Count states not alleles

Author: benjeffery

bio2zarr/tskit.py

@@ -118,8 +118,11 @@ def generate_schema(
 
         # Determine max number of alleles
         max_alleles = 0
-        for variant in self.ts.variants():
-            max_alleles = max(max_alleles, len(variant.alleles))
+        for site in self.ts.sites():
+            states = {site.ancestral_state}
+            for mut in site.mutations:
+                states.add(mut.derived_state)
+            max_alleles = max(len(states), max_alleles)
 
         logging.info(f"Scanned tskit with {n} samples and {m} variants")
         logging.info(