Add tskit CLI

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)

bio2zarr/tskit.py

@@ -13,7 +13,7 @@ class TskitFormat(vcz.Source):
     def __init__(
         self,
         ts_path,
-        individuals_nodes,
+        individuals_nodes=None,
         sample_ids=None,
         contig_id=None,
         isolated_as_missing=False,
@@ -24,6 +24,13 @@ def __init__(
         self.isolated_as_missing = isolated_as_missing
 
         self.positions = self.ts.sites_position
+        if individuals_nodes is None:
+            if self.ts.num_individuals == 0:
+                raise ValueError(
+                    "No individuals found in the tree seqeuence, use individuals_nodes "
+                    "argument to specify the individuals nodes"
+                )
+            individuals_nodes = self.ts.individuals_nodes
 
         self._num_samples = individuals_nodes.shape[0]
         if self._num_samples < 1:
@@ -213,8 +220,8 @@ def generate_schema(
 def convert(
     ts_path,
     zarr_path,
-    individuals_nodes,
     *,
+    individuals_nodes=None,
     sample_ids=None,
     contig_id=None,
     isolated_as_missing=False,
@@ -225,7 +232,7 @@ def convert(
 ):
     tskit_format = TskitFormat(
         ts_path,
-        individuals_nodes,
+        individuals_nodes=individuals_nodes,
         sample_ids=sample_ids,
         contig_id=contig_id,
         isolated_as_missing=isolated_as_missing,

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.9"
 classifiers = [
@@ -52,6 +54,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 = [