Correct assertions in ImpData

python/tests/beagle_numba.py

@@ -50,6 +50,8 @@
 __VERSION__ = "0.0.0"
 __DATE__ = "20XXXXXX"
 
+_ACGT_ALLELES = [0, 1, 2, 3, tskit.MISSING_DATA]
+
 
 @dataclass(frozen=True)
 class GeneticMap:
@@ -105,20 +107,29 @@ def __post_init__(self):
         assert self.alleles.shape[0] / 2 == len(
             self.individual_names
         ), "Number of query haplotypes is not equal to twice the number of individuals."
-        assert (
-            len(self.site_pos) == self.alleles.shape[1]
-        ), "Number of sites in alleles is not equal to the number of site positions."
-        assert (
-            len(self.site_pos) == self.alleles.shape[1]
+        assert len(self.site_pos) == len(
+            self.alts
         ), "Number of sites in refs is not equal to the number of site positions."
+        assert len(self.site_pos) == len(
+            self.refs
+        ), "Number of sites in alts is not equal to the number of site positions."
         assert (
             len(self.site_pos) == self.alleles.shape[1]
-        ), "Number of sites in alts is not equal to the number of site positions."
+        ), "Number of sites in alleles is not equal to the number of site positions."
         assert (
             self.alleles.shape == self.allele_probs.shape
         ), "Dimensions in alleles and allele probabilities don't match."
         for i in range(len(self.alleles)):
             assert np.all(np.isin(self.alleles[:, i], [self.refs[i], self.alts[i]]))
+        assert np.all(
+            np.isin(np.unique(self.refs), _ACGT_ALLELES)
+        ), "Unrecognized alleles are in REF alleles."
+        assert np.all(
+            np.isin(np.unique(self.alts), _ACGT_ALLELES)
+        ), "Unrecognized alleles are in ALT alleles."
+        assert np.all(
+            np.isin(np.unique(self.alleles), _ACGT_ALLELES)
+        ), "Unrecognized alleles are in alleles."
 
     @property
     def num_sites(self):