Fixup documentation and add some info logging

Some tests

Author: jeromekelleher

bio2zarr/tskit.py

@@ -28,7 +28,12 @@ def __init__(
         else:
             # input 'ts' is a path.
             self._path = ts
+            logger.info(f"Loading from {ts}")
             self.ts = tskit.load(ts)
+        logger.info(
+            f"Input has {self.ts.num_individuals} individuals and "
+            f"{self.ts.num_sites} sites"
+        )
 
         self.contig_id = contig_id if contig_id is not None else "1"
         self.isolated_as_missing = isolated_as_missing
@@ -42,6 +47,7 @@ def __init__(
         sample_ids = model_mapping.individuals_name
 
         self._num_samples = individuals_nodes.shape[0]
+        logger.info(f"Converting for {self._num_samples} samples")
         if self._num_samples < 1:
             raise ValueError("individuals_nodes must have at least one sample")
         self.max_ploidy = individuals_nodes.shape[1]
@@ -100,6 +106,7 @@ def iter_field(self, field_name, shape, start, stop):
     def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
         # All genotypes in tskit are considered phased
         phased = np.ones(shape[:-1], dtype=bool)
+        logger.debug(f"Getting genotpes start={start} stop={stop}")
 
         for variant in self.ts.variants(
             isolated_as_missing=self.isolated_as_missing,

docs/tskit2zarr/python_api.md

@@ -12,13 +12,14 @@ This will convert the [tskit](https://tskit.dev) tree sequence stored
 at ``ts_path`` to VCF Zarr stored at ``vcz_path`` using 8 worker processes.
 The details of how we map from the
 tskit {ref}`tskit:sec_data_model` to VCF Zarr are taken care of by
-TreeSequence.map_to_vcf_model method, which is called with no
+{meth}`tskit.TreeSequence.map_to_vcf_model`
+method, which is called with no
 parameters by default if the ``model_mapping`` parameter to
 {func}`~bio2zarr.tskit.convert` is not specified.
 
 For more control over the properties of the output, for example
 to pick a specific subset of individuals, you can use
-TreeSequence.map_to_vcf_model
+{meth}`~tskit.TreeSequence.map_to_vcf_model`
 to return the required mapping:
 
 ```python

tests/test_tskit.py

@@ -518,3 +518,37 @@ def test_against_tskit_vcf_output(ts, tmp_path):
         )
     )
     xt.assert_equal(ds1, ds2)
+
+
+def assert_ts_ds_equal(ts, ds, ploidy=2):
+    assert ds.sizes["ploidy"] == ploidy
+    assert ds.sizes["variants"] == ts.num_sites
+    assert ds.sizes["samples"] == ts.num_individuals
+    # Msprime guarantees that this will be true.
+    nt.assert_array_equal(
+        ts.genotype_matrix().reshape((ts.num_sites, ts.num_individuals, ploidy)),
+        ds.call_genotype.values,
+    )
+    nt.assert_array_equal(
+        ds.call_genotype_phased.values,
+        np.ones((ts.num_sites, ts.num_individuals), dtype=bool),
+    )
+    # Specialised for the limited form of mutations used here
+    nt.assert_equal(
+        ds.variant_allele[:, 0].values, [site.ancestral_state for site in ts.sites()]
+    )
+    nt.assert_equal(
+        ds.variant_allele[:, 1].values,
+        [mutation.derived_state for mutation in ts.mutations()],
+    )
+    nt.assert_equal(ds.variant_position, ts.sites_position)
+
+
+@pytest.mark.parametrize("worker_processes", [0, 1, 2, 15])
+def test_workers(tmp_path, worker_processes):
+    ts = msprime.sim_ancestry(10, sequence_length=1000, random_seed=42)
+    ts = add_mutations(ts)
+    out = tmp_path / "tskit.zarr"
+    tsk.convert(ts, out, worker_processes=worker_processes)
+    ds = sg.load_dataset(out)
+    assert_ts_ds_equal(ts, ds)