Add tskit CLI

Author: benjeffery

bio2zarr/__main__.py

@@ -17,6 +17,7 @@ def bio2zarr():
 bio2zarr.add_command(cli.vcf2zarr_main)
 bio2zarr.add_command(cli.plink2zarr)
 bio2zarr.add_command(cli.vcfpartition)
+bio2zarr.add_command(cli.tskit2zarr)
 
 if __name__ == "__main__":
     bio2zarr()

bio2zarr/cli.py

@@ -9,6 +9,7 @@
 import tabulate
 
 from . import plink, provenance, vcf_utils
+from . import tskit as tskit_mod
 from . import vcf as vcf_mod
 
 logger = logging.getLogger(__name__)
@@ -630,3 +631,52 @@ def vcfpartition(vcfs, verbose, num_partitions, partition_size):
         )
         for region in regions:
             click.echo(f"{region}\t{vcf_path}")
+
+
+@click.command(name="convert")
+@click.argument("ts_path", type=click.Path(exists=True))
+@click.argument("zarr_path", type=click.Path())
+@click.option("--contig-id", type=str, help="Contig/chromosome ID (default: '1')")
+@click.option(
+    "--isolated-as-missing", is_flag=True, help="Treat isolated nodes as missing"
+)
+@variants_chunk_size
+@samples_chunk_size
+@verbose
+@progress
+@worker_processes
+@force
+def convert_tskit(
+    ts_path,
+    zarr_path,
+    contig_id,
+    isolated_as_missing,
+    variants_chunk_size,
+    samples_chunk_size,
+    verbose,
+    progress,
+    worker_processes,
+    force,
+):
+    setup_logging(verbose)
+    check_overwrite_dir(zarr_path, force)
+
+    tskit_mod.convert(
+        ts_path,
+        zarr_path,
+        contig_id=contig_id,
+        isolated_as_missing=isolated_as_missing,
+        variants_chunk_size=variants_chunk_size,
+        samples_chunk_size=samples_chunk_size,
+        worker_processes=worker_processes,
+        show_progress=progress,
+    )
+
+
+@version
+@click.group()
+def tskit2zarr():
+    pass
+
+
+tskit2zarr.add_command(convert_tskit)

pyproject.toml

@@ -25,6 +25,8 @@ dependencies = [
   # colouredlogs pulls in humanfriendly",
   "cyvcf2",
   "bed_reader",
+  # TODO Using dev version of tskit for CI, FIXME before release
+  "tskit @ git+https://github.com/tskit-dev/tskit.git@main#subdirectory=python",
 ]
 requires-python = ">=3.10"
 classifiers = [
@@ -51,6 +53,7 @@ documentation = "https://sgkit-dev.github.io/bio2zarr/"
 [project.scripts]
 vcf2zarr = "bio2zarr.cli:vcf2zarr_main"
 vcfpartition = "bio2zarr.cli:vcfpartition"
+tskit2zarr = "bio2zarr.cli:tskit2zarr_main"
 
 [project.optional-dependencies]
 dev = [

tests/data/ts/example.trees

@@ -61,6 +61,15 @@
     local_alleles=False,
 )
 
+DEFAULT_TSKIT_CONVERT_ARGS = dict(
+    contig_id=None,
+    isolated_as_missing=False,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    show_progress=True,
+    worker_processes=1,
+)
+
 DEFAULT_PLINK_CONVERT_ARGS = dict(
     variants_chunk_size=None,
     samples_chunk_size=None,
@@ -635,6 +644,116 @@ def test_vcf_convert_overwrite_zarr_confirm_yes(self, mocked, tmp_path, response
             (self.vcf_path,), str(zarr_path), **DEFAULT_CONVERT_ARGS
         )
 
+    @pytest.mark.parametrize(("progress", "flag"), [(True, "-P"), (False, "-Q")])
+    @mock.patch("bio2zarr.tskit.convert")
+    def test_convert_tskit(self, mocked, tmp_path, progress, flag):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path} {flag}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+        args = dict(DEFAULT_TSKIT_CONVERT_ARGS)
+        args["show_progress"] = progress
+        mocked.assert_called_once_with(
+            ts_path,
+            str(zarr_path),
+            **args,
+        )
+
+    @pytest.mark.parametrize("response", ["y", "Y", "yes"])
+    @mock.patch("bio2zarr.tskit.convert")
+    def test_tskit_convert_overwrite_zarr_confirm_yes(self, mocked, tmp_path, response):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        zarr_path.mkdir()
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path}",
+            catch_exceptions=False,
+            input=response,
+        )
+        assert result.exit_code == 0
+        assert f"Do you want to overwrite {zarr_path}" in result.stdout
+        assert len(result.stderr) == 0
+        mocked.assert_called_once_with(
+            ts_path,
+            str(zarr_path),
+            **DEFAULT_TSKIT_CONVERT_ARGS,
+        )
+
+    @pytest.mark.parametrize("response", ["n", "N", "No"])
+    @mock.patch("bio2zarr.tskit.convert")
+    def test_tskit_convert_overwrite_zarr_confirm_no(self, mocked, tmp_path, response):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        zarr_path.mkdir()
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path}",
+            catch_exceptions=False,
+            input=response,
+        )
+        assert result.exit_code == 1
+        assert "Aborted" in result.stderr
+        mocked.assert_not_called()
+
+    @pytest.mark.parametrize("force_arg", ["-f", "--force"])
+    @mock.patch("bio2zarr.tskit.convert")
+    def test_tskit_convert_overwrite_zarr_force(self, mocked, tmp_path, force_arg):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        zarr_path.mkdir()
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path} {force_arg}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+        mocked.assert_called_once_with(
+            ts_path,
+            str(zarr_path),
+            **DEFAULT_TSKIT_CONVERT_ARGS,
+        )
+
+    @mock.patch("bio2zarr.tskit.convert")
+    def test_tskit_convert_with_options(self, mocked, tmp_path):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path} --contig-id chr1 "
+            "--isolated-as-missing -l 100 -w 50 -p 4",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        assert len(result.stdout) == 0
+        assert len(result.stderr) == 0
+
+        expected_args = dict(DEFAULT_TSKIT_CONVERT_ARGS)
+        expected_args["contig_id"] = "chr1"
+        expected_args["isolated_as_missing"] = True
+        expected_args["variants_chunk_size"] = 100
+        expected_args["samples_chunk_size"] = 50
+        expected_args["worker_processes"] = 4
+
+        mocked.assert_called_once_with(
+            ts_path,
+            str(zarr_path),
+            **expected_args,
+        )
+
 
 class TestVcfEndToEnd:
     vcf_path = "tests/data/vcf/sample.vcf.gz"
@@ -908,10 +1027,36 @@ def test_part_size_multiple_vcfs(self):
 
 
 @pytest.mark.parametrize(
-    "cmd", [main.bio2zarr, cli.vcf2zarr_main, cli.plink2zarr, cli.vcfpartition]
+    "cmd",
+    [
+        main.bio2zarr,
+        cli.vcf2zarr_main,
+        cli.plink2zarr,
+        cli.vcfpartition,
+        cli.tskit2zarr,
+    ],
 )
 def test_version(cmd):
     runner = ct.CliRunner()
     result = runner.invoke(cmd, ["--version"], catch_exceptions=False)
     s = f"version {provenance.__version__}\n"
     assert result.stdout.endswith(s)
+
+
+class TestTskitEndToEnd:
+    def test_convert(self, tmp_path):
+        ts_path = "tests/data/ts/example.trees"
+        zarr_path = tmp_path / "zarr"
+        runner = ct.CliRunner()
+        result = runner.invoke(
+            cli.tskit2zarr,
+            f"convert {ts_path} {zarr_path}",
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        result = runner.invoke(
+            cli.vcf2zarr_main, f"inspect {zarr_path}", catch_exceptions=False
+        )
+        assert result.exit_code == 0
+        # Arbitrary check
+        assert "variant_position" in result.stdout

tests/test_cli.py

@@ -237,7 +237,7 @@ def test_examples(self, chunk_size, size, start, stop):
         # It works in CI on Linux, but it'll probably break at some point.
         # It's also necessary to update these numbers each time a new data
         # file gets added
-        ("tests/data", 5030777),
+        ("tests/data", 5045029),
         ("tests/data/vcf", 5018640),
         ("tests/data/vcf/sample.vcf.gz", 1089),
     ],

tests/test_core.py

@@ -40,29 +40,54 @@ def test_simple_tree_sequence(self, tmp_path):
                 tmp_path / "test.trees", zarr_path, ind_nodes, show_progress=False
             )
             zroot = zarr.open(zarr_path, mode="r")
-            assert zroot["variant_position"].shape == (3,)
-            assert list(zroot["variant_position"][:]) == [10, 20, 30]
+            pos = zroot["variant_position"][:]
+            assert pos.shape == (3,)
+            assert pos.dtype == np.int8
+            assert np.array_equal(pos, [10, 20, 30])
 
             alleles = zroot["variant_allele"][:]
+            assert alleles.shape == (3, 2)
+            assert alleles.dtype == "O"
             assert np.array_equal(alleles, [["A", "T"], ["C", "G"], ["G", "A"]])
 
             genotypes = zroot["call_genotype"][:]
+            assert genotypes.shape == (3, 2, 2)
+            assert genotypes.dtype == np.int8
             assert np.array_equal(
                 genotypes, [[[1, 1], [0, 0]], [[0, 0], [1, 1]], [[1, 0], [0, 0]]]
             )
 
             phased = zroot["call_genotype_phased"][:]
+            assert phased.shape == (3, 2)
+            assert phased.dtype == np.bool
             assert np.all(phased)
 
             contigs = zroot["contig_id"][:]
+            assert contigs.shape == (1,)
+            assert contigs.dtype == "O"
             assert np.array_equal(contigs, ["1"])
 
             contig = zroot["variant_contig"][:]
+            assert contig.shape == (3,)
+            assert contig.dtype == np.int8
             assert np.array_equal(contig, [0, 0, 0])
 
             samples = zroot["sample_id"][:]
+            assert samples.shape == (2,)
+            assert samples.dtype == "O"
             assert np.array_equal(samples, ["tsk_0", "tsk_1"])
 
+            assert set(zroot.array_keys()) == {
+                "variant_position",
+                "variant_allele",
+                "call_genotype",
+                "call_genotype_phased",
+                "call_genotype_mask",
+                "contig_id",
+                "variant_contig",
+                "sample_id",
+            }
+
 
 class TestTskitFormat:
     """Unit tests for TskitFormat without using full conversion."""