Add tskit CLI

Author: benjeffery

bio2zarr/cli.py

@@ -10,6 +10,7 @@
 
 from . import icf as icf_mod
 from . import plink, provenance, vcf_utils
+from . import tskit as tskit_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(cls=NaturalOrderGroup, name="tskit2zarr")
+def tskit2zarr_main():
+    pass
+
+
+tskit2zarr_main.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,8 @@ def __init__(
         self.isolated_as_missing = isolated_as_missing
 
         self.positions = self.ts.sites_position
+        if individuals_nodes is None:
+            individuals_nodes = self.ts.individuals_nodes
 
         self._num_samples = individuals_nodes.shape[0]
         if self._num_samples < 1:
@@ -213,8 +215,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 +227,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

@@ -52,6 +52,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 = [