fix #59 bug for all opt functions

allel/opt/model.pyx

@@ -125,7 +125,7 @@ def haplotype_int8_count_alleles(cnp.int8_t[:, :] h not None,
             # iterate over haplotypes
             for j in range(n_haplotypes):
                 allele = h[i, j]
-                if allele >= 0:
+                if 0 <= allele <= max_allele:
                     ac[i, allele] += 1
 
     return np.asarray(ac)
@@ -155,7 +155,7 @@ def haplotype_int8_count_alleles_subpop(cnp.int8_t[:, :] h not None,
             for j in range(n_haplotypes):
                 idx = subpop[j]
                 allele = h[i, idx]
-                if allele >= 0:
+                if 0 <= allele <= max_allele:
                     ac[i, allele] += 1
 
     return np.asarray(ac)
@@ -218,7 +218,7 @@ def genotype_int8_count_alleles_masked(cnp.int8_t[:, :, :] g not None,
                     # iterate over alleles
                     for k in range(ploidy):
                         allele = g[i, j, k]
-                        if allele >= 0:
+                        if 0 <= allele <= max_allele:
                             ac[i, allele] += 1
 
     return np.asarray(ac)
@@ -250,7 +250,7 @@ def genotype_int8_count_alleles_subpop(cnp.int8_t[:, :, :] g not None,
                 idx = subpop[j]
                 for k in range(ploidy):
                     allele = g[i, idx, k]
-                    if allele >= 0:
+                    if 0 <= allele <= max_allele:
                         ac[i, allele] += 1
 
     return np.asarray(ac)
@@ -285,7 +285,7 @@ def genotype_int8_count_alleles_subpop_masked(cnp.int8_t[:, :, :] g not None,
                 if not mask[i, idx]:
                     for k in range(ploidy):
                         allele = g[i, idx, k]
-                        if allele >= 0:
+                        if 0 <= allele <= max_allele:
                             ac[i, allele] += 1
 
     return np.asarray(ac)

allel/test/test_model_api.py

@@ -987,7 +987,7 @@ def test_vstack(self):
 
 class HaplotypeArrayInterface(object):
 
-    def setup_instance(self, data):
+    def setup_instance(self, data, dtype=None):
         # to be implemented in sub-classes
         pass
 
@@ -1219,21 +1219,24 @@ def test_count_alleles(self):
                            [0, 2, 0],
                            [0, 0, 1],
                            [0, 0, 0]])
-        actual = self.setup_instance(haplotype_data).count_alleles()
-        aeq(expect, actual)
-        eq(4, actual.n_variants)
-        eq(3, actual.n_alleles)
+        for dtype in None, 'i1', 'i2':
+            h = self.setup_instance(haplotype_data, dtype=dtype)
+            actual = h.count_alleles()
+            aeq(expect, actual)
+            eq(4, actual.n_variants)
+            eq(3, actual.n_alleles)
 
     def test_count_alleles_subpop(self):
         expect = np.array([[1, 0, 0],
                            [0, 1, 0],
                            [0, 0, 1],
                            [0, 0, 0]])
-        h = self.setup_instance(haplotype_data)
-        actual = h.count_alleles(subpop=[0, 2])
-        aeq(expect, actual)
-        eq(4, actual.n_variants)
-        eq(3, actual.n_alleles)
+        for dtype in None, 'i1', 'i2':
+            h = self.setup_instance(haplotype_data, dtype=dtype)
+            actual = h.count_alleles(subpop=[0, 2])
+            aeq(expect, actual)
+            eq(4, actual.n_variants)
+            eq(3, actual.n_alleles)
 
     def test_count_alleles_subpops(self):
         expect_sub1 = np.array([[1, 0, 0],
@@ -1244,15 +1247,36 @@ def test_count_alleles_subpops(self):
                                 [0, 1, 0],
                                 [0, 0, 0],
                                 [0, 0, 0]])
-        h = self.setup_instance(haplotype_data)
-        subpops = {'sub1': [0, 2], 'sub2': [1, 2]}
-        actual = h.count_alleles_subpops(subpops=subpops)
-        aeq(expect_sub1, actual['sub1'])
-        aeq(expect_sub2, actual['sub2'])
-        eq(4, actual['sub1'].n_variants)
-        eq(3, actual['sub1'].n_alleles)
-        eq(4, actual['sub2'].n_variants)
-        eq(3, actual['sub2'].n_alleles)
+        for dtype in None, 'i1', 'i2':
+            h = self.setup_instance(haplotype_data, dtype=dtype)
+            subpops = {'sub1': [0, 2], 'sub2': [1, 2]}
+            actual = h.count_alleles_subpops(subpops=subpops)
+            aeq(expect_sub1, actual['sub1'])
+            aeq(expect_sub2, actual['sub2'])
+            eq(4, actual['sub1'].n_variants)
+            eq(3, actual['sub1'].n_alleles)
+            eq(4, actual['sub2'].n_variants)
+            eq(3, actual['sub2'].n_alleles)
+
+    def test_count_alleles_max_allele(self):
+        expect = np.array([[1, 1, 0],
+                           [0, 2, 0],
+                           [0, 0, 1],
+                           [0, 0, 0]])
+        for dtype in None, 'i1', 'i2':
+            h = self.setup_instance(haplotype_data, dtype=dtype)
+            actual = h.count_alleles()
+            eq(3, actual.n_alleles)
+            aeq(expect, actual)
+            actual = h.count_alleles(max_allele=2)
+            eq(3, actual.n_alleles)
+            aeq(expect, actual)
+            actual = h.count_alleles(max_allele=1)
+            eq(2, actual.n_alleles)
+            aeq(expect[:, :2], actual)
+            actual = h.count_alleles(max_allele=0)
+            eq(1, actual.n_alleles)
+            aeq(expect[:, :1], actual)
 
     def test_map_alleles(self):
 

allel/test/test_model_bcolz.py

@@ -203,8 +203,8 @@ class HaplotypeCArrayTests(HaplotypeArrayInterface, unittest.TestCase):
 
     _class = HaplotypeCArray
 
-    def setup_instance(self, data):
-        return HaplotypeCArray(data)
+    def setup_instance(self, data, dtype=None):
+        return HaplotypeCArray(data, dtype=dtype)
 
     def test_constructor(self):
 

allel/test/test_model_chunked.py

@@ -209,8 +209,9 @@ class HaplotypeChunkedArrayTests(HaplotypeArrayInterface, unittest.TestCase):
     def setUp(self):
         chunked.storage_registry['default'] = chunked.bcolzmem_storage
 
-    def setup_instance(self, data):
-        data = chunked.storage_registry['default'].array(data, chunklen=2)
+    def setup_instance(self, data, dtype=None):
+        data = chunked.storage_registry['default'].array(data, dtype=dtype,
+                                                         chunklen=2)
         return HaplotypeChunkedArray(data)
 
     def test_constructor(self):

allel/test/test_model_dask.py

@@ -29,8 +29,8 @@ class GenotypeDaskArrayTests(GenotypeArrayInterface, unittest.TestCase):
     _class = GenotypeDaskArray
 
     def setup_instance(self, data, dtype=None):
-        # ignore dtype
-        return GenotypeDaskArray.from_array(data, chunks=(2, 2, None))
+        return GenotypeDaskArray.from_array(np.asarray(data, dtype=dtype),
+                                            chunks=(2, 2, None))
 
     def test_constructor(self):
 
@@ -130,8 +130,9 @@ class HaplotypeDaskArrayTests(HaplotypeArrayInterface, unittest.TestCase):
 
     _class = HaplotypeDaskArray
 
-    def setup_instance(self, data):
-        return HaplotypeDaskArray.from_array(data, chunks=(2, 2))
+    def setup_instance(self, data, dtype=None):
+        return HaplotypeDaskArray.from_array(np.asarray(data, dtype=dtype),
+                                             chunks=(2, 2))
 
     def test_constructor(self):
 

allel/test/test_model_ndarray.py

@@ -187,8 +187,8 @@ class HaplotypeArrayTests(HaplotypeArrayInterface, unittest.TestCase):
 
     _class = HaplotypeArray
 
-    def setup_instance(self, data):
-        return HaplotypeArray(data)
+    def setup_instance(self, data, dtype=None):
+        return HaplotypeArray(data, dtype=dtype)
 
     def test_constructor(self):