Dynamically pick call_genotype dtype

Author: benjeffery

bio2zarr/tskit.py

@@ -191,7 +191,7 @@ def generate_schema(
             vcz.ZarrArraySpec(
                 source=None,
                 name="call_genotype",
-                dtype="i1",
+                dtype=core.min_int_dtype(constants.INT_FILL, max_alleles - 1),
                 dimensions=["variants", "samples", "ploidy"],
                 description="Genotype for each variant and sample",
                 compressor=vcz.DEFAULT_ZARR_COMPRESSOR_GENOTYPES.get_config(),

tests/test_ts.py

@@ -395,3 +395,43 @@ def insert_branch_sites(ts, m=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)
+
+    def test_genotype_dtype_selection(self, tmp_path):
+        tables = tskit.TableCollection(sequence_length=100)
+        for _ in range(4):
+            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        mrca = tables.nodes.add_row(flags=0, time=1)
+        for i in range(4):
+            tables.edges.add_row(left=0, right=100, parent=mrca, child=i)
+        site_id = tables.sites.add_row(position=10, ancestral_state="A")
+        tables.mutations.add_row(site=site_id, node=0, derived_state="T")
+        tables.sort()
+        tree_sequence = tables.tree_sequence()
+        ts_path = tmp_path / "small_alleles.trees"
+        tree_sequence.dump(ts_path)
+
+        ind_nodes = np.array([[0, 1], [2, 3]])
+        format_obj = ts.TskitFormat(ts_path, ind_nodes)
+        schema = format_obj.generate_schema()
+        call_genotype_spec = next(s for s in schema.fields if s.name == "call_genotype")
+        assert call_genotype_spec.dtype == "i1"
+
+        tables = tskit.TableCollection(sequence_length=100)
+        for _ in range(4):
+            tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        mrca = tables.nodes.add_row(flags=0, time=1)
+        for i in range(4):
+            tables.edges.add_row(left=0, right=100, parent=mrca, child=i)
+        site_id = tables.sites.add_row(position=10, ancestral_state="A")
+        for i in range(32768):
+            tables.mutations.add_row(site=site_id, node=0, derived_state=f"ALLELE_{i}")
+
+        tables.sort()
+        tree_sequence = tables.tree_sequence()
+        ts_path = tmp_path / "large_alleles.trees"
+        tree_sequence.dump(ts_path)
+
+        format_obj = ts.TskitFormat(ts_path, ind_nodes)
+        schema = format_obj.generate_schema()
+        call_genotype_spec = next(s for s in schema.fields if s.name == "call_genotype")
+        assert call_genotype_spec.dtype == "i4"