Add tskit CLI

Author: benjeffery

bio2zarr/cli.py

@@ -630,3 +630,82 @@ 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(
+    "--ploidy",
+    type=int,
+    default=2,
+    help="Ploidy to use when tree sequence has no individuals (default: 2)",
+)
+@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,
+    ploidy,
+    contig_id,
+    isolated_as_missing,
+    variants_chunk_size,
+    samples_chunk_size,
+    verbose,
+    progress,
+    worker_processes,
+    force,
+):
+    """
+    Convert a tree sequence file to vcfzarr format.
+
+    TS_PATH is the path to the tree sequence file.
+    ZARR_PATH is the path for the output zarr directory.
+    """
+    setup_logging(verbose)
+    check_overwrite_dir(zarr_path, force)
+
+    import tskit
+
+    from bio2zarr import tskit as tskit_mod
+
+    ts = tskit.load(ts_path)
+
+    if ts.num_individuals == 0:
+        individuals_nodes = ts.split_samples_by_ploidy(ploidy)
+    else:
+        individuals_nodes = ts.individuals_nodes
+
+    tskit_mod.convert(
+        ts_path,
+        zarr_path,
+        individuals_nodes,
+        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():
+    """
+    Convert tree sequence file(s) to the vcfzarr format.
+
+    See the online documentation at https://sgkit-dev.github.io/bio2zarr/
+    for more information.
+    """
+
+
+tskit2zarr_main.add_command(convert_tskit)

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 = [