Add tests of tskit source

Author: benjeffery

bio2zarr/ts.py

@@ -95,7 +95,7 @@ def _make_sample_mapping(self, ploidy):
                 ts.num_samples // ploidy, ploidy, dtype=np.int32
             )
             self.max_ploidy = ploidy
-            self.sample_ids = np.arange(ts.num_samples, dtype=np.int32)
+            self.sample_ids = (ts.nodes_flags & tskit.NODE_IS_SAMPLE).nonzero()[0]
 
         self._num_samples = len(self.individual_ploidies)
 
@@ -125,33 +125,30 @@ def iter_alleles(self, start, stop, num_alleles):
             yield alleles
 
     def iter_alleles_and_genotypes(self, start, stop, shape, num_alleles):
-        gt = np.zeros(shape, dtype=np.int8)
-        phased = np.zeros(shape[:-1], dtype=bool)
+        # In tskit, all genotypes are considered phased
+        phased = np.ones(shape[:-1], dtype=bool)
 
         for variant in self.ts.variants(
             samples=self.sample_ids,
             isolated_as_missing=self.isolated_as_missing,
             left=self.positions[start],
             right=self.positions[stop] if stop < self.num_records else None,
         ):
+            gt = np.full(shape, constants.INT_FILL, dtype=np.int8)
             alleles = np.full(num_alleles, constants.STR_FILL, dtype="O")
             for i, allele in enumerate(variant.alleles):
+                # None is returned by tskit in the case of a missing allele
+                if allele is None:
+                    continue
                 assert i < num_alleles
                 alleles[i] = allele
 
             genotypes = variant.genotypes
             sample_index = 0
             for i, ploidy in enumerate(self.individual_ploidies):
                 for j in range(ploidy):
-                    if j < self.max_ploidy:  # Only fill up to max_ploidy
-                        try:
-                            gt[i, j] = genotypes[sample_index + j]
-                        except IndexError:
-                            # This can happen if the ploidy varies between individuals
-                            gt[i, j] = -2  # Fill value
-
-                # In tskit, all genotypes are considered phased
-                phased[i] = True
+                    if j < self.max_ploidy:
+                        gt[i, j] = genotypes[sample_index + j]
                 sample_index += ploidy
 
             yield alleles, (gt, phased)

tests/test_ts.py

@@ -2,6 +2,7 @@
 import tempfile
 
 import numpy as np
+import pytest
 import tskit
 import zarr
 
@@ -58,3 +59,332 @@ def test_simple_tree_sequence(self, tmp_path):
 
             samples = zroot["sample_id"][:]
             assert np.array_equal(samples, ["tsk_0", "tsk_1"])
+
+
+class TestTskitFormat:
+    """Unit tests for TskitFormat without using full conversion."""
+
+    @pytest.fixture()
+    def simple_ts(self, tmp_path):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=0, time=1)  # MRCA for 0,1
+        tables.nodes.add_row(flags=0, time=1)  # MRCA for 2,3
+        tables.edges.add_row(left=0, right=100, parent=4, child=0)
+        tables.edges.add_row(left=0, right=100, parent=4, child=1)
+        tables.edges.add_row(left=0, right=100, parent=5, child=2)
+        tables.edges.add_row(left=0, right=100, parent=5, child=3)
+        site_id = tables.sites.add_row(position=10, ancestral_state="A")
+        tables.mutations.add_row(site=site_id, node=4, derived_state="T")
+        site_id = tables.sites.add_row(position=20, ancestral_state="C")
+        tables.mutations.add_row(site=site_id, node=5, derived_state="G")
+        site_id = tables.sites.add_row(position=30, ancestral_state="G")
+        tables.mutations.add_row(site=site_id, node=0, derived_state="A")
+        tables.sort()
+        tree_sequence = tables.tree_sequence()
+        ts_path = tmp_path / "test.trees"
+        tree_sequence.dump(ts_path)
+        return ts_path, tree_sequence
+
+    @pytest.fixture()
+    def no_individuals_ts(self, tmp_path):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.nodes.add_row(flags=0, time=1)  # MRCA for 0,1
+        tables.nodes.add_row(flags=0, time=1)  # MRCA for 2,3
+        tables.edges.add_row(left=0, right=100, parent=4, child=0)
+        tables.edges.add_row(left=0, right=100, parent=4, child=1)
+        tables.edges.add_row(left=0, right=100, parent=5, child=2)
+        tables.edges.add_row(left=0, right=100, parent=5, child=3)
+        site_id = tables.sites.add_row(position=10, ancestral_state="A")
+        tables.mutations.add_row(site=site_id, node=4, derived_state="T")
+        site_id = tables.sites.add_row(position=20, ancestral_state="C")
+        tables.mutations.add_row(site=site_id, node=5, derived_state="G")
+        tables.sort()
+        tree_sequence = tables.tree_sequence()
+        ts_path = tmp_path / "no_individuals.trees"
+        tree_sequence.dump(ts_path)
+        return ts_path, tree_sequence
+
+    def test_initialization(self, simple_ts):
+        ts_path, tree_sequence = simple_ts
+
+        # Test with default parameters
+        format_obj = ts.TskitFormat(ts_path)
+        assert format_obj.path == ts_path
+        assert format_obj.ts.num_sites == tree_sequence.num_sites
+        assert format_obj.contig_id == "1"
+        assert not format_obj.isolated_as_missing
+
+        # Test with custom parameters
+        format_obj = ts.TskitFormat(ts_path, contig_id="chr1", isolated_as_missing=True)
+        assert format_obj.contig_id == "chr1"
+        assert format_obj.isolated_as_missing
+        assert format_obj.path == ts_path
+
+    def test_basic_properties(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+
+        assert format_obj.num_records == format_obj.ts.num_sites
+        assert format_obj.num_samples == 2  # Two individuals
+        assert len(format_obj.samples) == 2
+        assert format_obj.samples[0].id == "tsk_0"
+        assert format_obj.samples[1].id == "tsk_1"
+
+        assert format_obj.root_attrs == {}
+
+        contigs = format_obj.contigs
+        assert len(contigs) == 1
+        assert contigs[0].id == "1"
+
+    def test_sample_mapping_with_individuals(self, simple_ts):
+        ts_path, _ = simple_ts
+
+        format_obj = ts.TskitFormat(ts_path)
+        assert format_obj.num_samples == 2
+        assert format_obj.max_ploidy == 2
+        assert format_obj.individual_ploidies == [2, 2]
+
+        # Should raise error if ploidy specified with individuals
+        with pytest.raises(
+            ValueError, match="Cannot specify ploidy when individuals are present"
+        ):
+            ts.TskitFormat(ts_path, ploidy=2)
+
+    def test_sample_mapping_without_individuals(self, no_individuals_ts):
+        ts_path, tree_sequence = no_individuals_ts
+
+        # Default ploidy should be 1
+        format_obj = ts.TskitFormat(ts_path)
+        assert format_obj.num_samples == 4
+        assert format_obj.max_ploidy == 1
+        assert list(format_obj.individual_ploidies) == [1, 1, 1, 1]
+
+        # Explicitly set ploidy to 2
+        format_obj = ts.TskitFormat(ts_path, ploidy=2)
+        assert format_obj.num_samples == 2
+        assert format_obj.max_ploidy == 2
+        assert list(format_obj.individual_ploidies) == [2, 2]
+
+        with pytest.raises(ValueError, match="Ploidy must be >= 1"):
+            ts.TskitFormat(ts_path, ploidy=0)
+
+        with pytest.raises(ValueError, match="Sample size must be divisible by ploidy"):
+            ts.TskitFormat(ts_path, ploidy=3)
+
+    def test_schema_generation(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+
+        schema = format_obj.generate_schema()
+        assert schema.dimensions["variants"].size == 3
+        assert schema.dimensions["samples"].size == 2
+        assert schema.dimensions["ploidy"].size == 2
+        assert schema.dimensions["alleles"].size == 2  # A/T, C/G, G/A -> max is 2
+        field_names = [field.name for field in schema.fields]
+        assert "variant_position" in field_names
+        assert "variant_allele" in field_names
+        assert "variant_contig" in field_names
+        assert "call_genotype" in field_names
+        assert "call_genotype_phased" in field_names
+        assert "call_genotype_mask" in field_names
+        schema = format_obj.generate_schema(
+            variants_chunk_size=10, samples_chunk_size=5
+        )
+        assert schema.dimensions["variants"].chunk_size == 10
+        assert schema.dimensions["samples"].chunk_size == 5
+
+    def test_iter_contig(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+        contig_indices = list(format_obj.iter_contig(1, 3))
+        assert contig_indices == [0, 0]
+
+    def test_iter_field(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+        positions = list(format_obj.iter_field("position", None, 0, 3))
+        assert positions == [10, 20, 30]
+        positions = list(format_obj.iter_field("position", None, 1, 3))
+        assert positions == [20, 30]
+        with pytest.raises(ValueError, match="Unknown field"):
+            list(format_obj.iter_field("unknown_field", None, 0, 3))
+
+    def test_iter_alleles(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+        alleles_list = list(format_obj.iter_alleles(0, 3, 2))
+
+        expected_alleles = np.array([["A", "T"], ["C", "G"], ["G", "A"]])
+        assert len(alleles_list) == 3
+        assert np.array_equal(alleles_list, expected_alleles)
+
+        # Test with different start/stop
+        alleles_list = list(format_obj.iter_alleles(1, 3, 2))
+        assert len(alleles_list) == 2
+        expected_alleles = np.array([["C", "G"], ["G", "A"]])
+        assert np.array_equal(alleles_list, expected_alleles)
+
+    def test_iter_alleles_and_genotypes(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+
+        shape = (2, 2)  # (num_samples, max_ploidy)
+        results = list(format_obj.iter_alleles_and_genotypes(0, 3, shape, 2))
+
+        assert len(results) == 3
+
+        for i, (alleles, (gt, phased)) in enumerate(results):
+            if i == 0:
+                assert tuple(alleles) == ("A", "T")
+                assert np.array_equal(gt, [[1, 1], [0, 0]])
+            elif i == 1:
+                assert tuple(alleles) == ("C", "G")
+                assert np.array_equal(gt, [[0, 0], [1, 1]])
+            elif i == 2:
+                assert tuple(alleles) == ("G", "A")
+                assert np.array_equal(gt, [[1, 0], [0, 0]])
+            assert np.all(phased)
+
+    def test_partial_iter_alleles_and_genotypes(self, simple_ts):
+        ts_path, _ = simple_ts
+        format_obj = ts.TskitFormat(ts_path)
+
+        shape = (2, 2)
+        results = list(format_obj.iter_alleles_and_genotypes(1, 3, shape, 2))
+        assert len(results) == 2
+        alleles, (gt, phased) = results[0]
+        assert tuple(alleles) == ("C", "G")
+        assert np.array_equal(gt, [[0, 0], [1, 1]])
+        assert np.all(phased)
+
+        alleles, (gt, phased) = results[1]
+        assert tuple(alleles) == ("G", "A")
+        assert np.array_equal(gt, [[1, 0], [0, 0]])
+        assert np.all(phased)
+
+    def test_variable_ploidy(self, tmp_path):
+        # Create a tree sequence with mixed ploidy
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=0)
+        tables.individuals.add_row(flags=0, location=(0, 0), metadata=b"")
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0, individual=1)
+        tables.nodes.add_row(flags=0, time=1)
+        tables.nodes.add_row(flags=0, time=1)
+        tables.edges.add_row(left=0, right=100, parent=5, child=0)
+        tables.edges.add_row(left=0, right=100, parent=5, child=1)
+        tables.edges.add_row(left=0, right=100, parent=6, child=2)
+        tables.edges.add_row(left=0, right=100, parent=6, child=3)
+        tables.edges.add_row(left=0, right=100, parent=6, child=4)
+
+        # Add a site with a mutation at individual 0, node 0
+        site_id = tables.sites.add_row(position=10, ancestral_state="A")
+        tables.mutations.add_row(site=site_id, node=0, derived_state="T")
+
+        # Add another site with mutation at individual 1, node 3
+        # (middle chromosome of triploid)
+        site_id = tables.sites.add_row(position=20, ancestral_state="C")
+        tables.mutations.add_row(site=site_id, node=3, derived_state="G")
+
+        tables.sort()
+        tree_sequence = tables.tree_sequence()
+        ts_path = tmp_path / "mixed_ploidy.trees"
+        tree_sequence.dump(ts_path)
+
+        format_obj = ts.TskitFormat(ts_path)
+
+        assert format_obj.max_ploidy == 3
+        assert format_obj.individual_ploidies == [2, 3]
+
+        shape = (2, 3)  # (num_samples, max_ploidy)
+        results = list(format_obj.iter_alleles_and_genotypes(0, 2, shape, 2))
+
+        assert len(results) == 2
+
+        alleles, (gt, phased) = results[0]
+        assert tuple(alleles) == ("A", "T")
+        # First site - derived on 1st of diploid
+        expected_gt = np.array([[1, 0, -2], [0, 0, 0]])
+        assert np.array_equal(gt, expected_gt)
+        assert np.all(phased)
+
+        alleles, (gt, phased) = results[1]
+        assert tuple(alleles) == ("C", "G")
+        # Second site - derived on 2nd of triploid
+        expected_gt = np.array([[0, 0, -2], [0, 1, 0]])
+        assert np.array_equal(gt, expected_gt)
+        assert np.all(phased)
+
+        # Check the fill value
+        assert gt[0, 2] == -2
+
+    def test_isolated_as_missing(self, tmp_path):
+        def insert_branch_sites(ts, m=1):
+            if m == 0:
+                return ts
+            tables = ts.dump_tables()
+            tables.sites.clear()
+            tables.mutations.clear()
+            for tree in ts.trees():
+                left, right = tree.interval
+                delta = (right - left) / (m * len(list(tree.nodes())))
+                x = left
+                for u in tree.nodes():
+                    if tree.parent(u) != tskit.NULL:
+                        for _ in range(m):
+                            site = tables.sites.add_row(position=x, ancestral_state="0")
+                            tables.mutations.add_row(
+                                site=site, node=u, derived_state="1"
+                            )
+                            x += delta
+            return tables.tree_sequence()
+
+        tables = tskit.Tree.generate_balanced(2, span=10).tree_sequence.dump_tables()
+        # This also tests sample nodes that are not a single block at
+        # the start of the nodes table.
+        tables.nodes.add_row(time=0, flags=tskit.NODE_IS_SAMPLE)
+        tree_sequence = insert_branch_sites(tables.tree_sequence())
+        print(tree_sequence.tables)
+
+        ts_path = tmp_path / "isolated_sample.trees"
+        tree_sequence.dump(ts_path)
+
+        format_obj_default = ts.TskitFormat(ts_path, isolated_as_missing=False)
+        shape = (3, 1)  # (num_samples, max_ploidy)
+        results_default = list(
+            format_obj_default.iter_alleles_and_genotypes(0, 1, shape, 2)
+        )
+
+        assert len(results_default) == 1
+        alleles, (gt_default, phased) = results_default[0]
+        assert tuple(alleles) == ("0", "1")
+
+        # Sample 2 should have the ancestral state (0) when isolated_as_missing=False
+        expected_gt_default = np.array([[1], [0], [0]])
+        assert np.array_equal(gt_default, expected_gt_default)
+
+        format_obj_missing = ts.TskitFormat(ts_path, isolated_as_missing=True)
+        results_missing = list(
+            format_obj_missing.iter_alleles_and_genotypes(0, 1, shape, 2)
+        )
+
+        assert len(results_missing) == 1
+        alleles, (gt_missing, phased) = results_missing[0]
+        assert tuple(alleles) == ("0", "1")
+
+        # Individual 2 should have missing values (-1) when isolated_as_missing=True
+        expected_gt_missing = np.array([[1], [0], [-1]])
+        assert np.array_equal(gt_missing, expected_gt_missing)