Added chunking of allele_mapping array so that da.map blocks works

Removed numba decorator for the apply_allele_mapping function, for now
malariagen · Aug 8, 2024 · 5e502bf · 5e502bf
1 parent ab6f28c
commit 5e502bf
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/malariagen_data/anoph/snp_data.py b/malariagen_data/anoph/snp_data.py
@@ -1627,15 +1627,23 @@ def biallelic_snp_calls(
             # Store alleles, transformed.
             variant_allele = ds_bi["variant_allele"].data
             variant_allele = variant_allele.rechunk((variant_allele.chunks[0], -1))
-            # variant_allele_out = da.map_blocks(
-            #    lambda block: apply_allele_mapping(block, allele_mapping, max_allele=1),
-            #    variant_allele,
-            #    dtype=variant_allele.dtype,
-            #    chunks=(variant_allele.chunks[0], [2]),
-            # )
-            variant_allele_out = apply_allele_mapping(
-                variant_allele.compute(), allele_mapping, max_allele=1
+
+            # Chunk allele mapping according to same variant_allele.
+            allele_mapping_chunked = da.from_array(
+                allele_mapping, chunks=variant_allele.chunks
             )
+
+            # Apply allele mapping blockwise to variant_allele.
+            variant_allele_out = da.map_blocks(
+                lambda allele, map: apply_allele_mapping(allele, map, max_allele=1),
+                variant_allele,
+                allele_mapping_chunked,
+                dtype=variant_allele.dtype,
+                chunks=(variant_allele.chunks[0], [2]),
+            )
+            # variant_allele_out = apply_allele_mapping(
+            #    variant_allele.compute(), allele_mapping, max_allele=1
+            # )
             data_vars["variant_allele"] = ("variants", "alleles"), variant_allele_out
 
             # Store allele counts, transformed, so we don't have to recompute.

diff --git a/malariagen_data/util.py b/malariagen_data/util.py
@@ -1348,7 +1348,7 @@ def trim_alleles(ac):
     return mapping
 
 
-@numba.njit
+# @numba.njit
 def apply_allele_mapping(x, mapping, max_allele):
     """Transform an array x, where the columns correspond to alleles,
     according to an allele mapping.
@@ -1360,9 +1360,7 @@ def apply_allele_mapping(x, mapping, max_allele):
     n_sites = x.shape[0]
     n_alleles = x.shape[1]
     assert mapping.shape[0] == n_sites
-    assert (
-        mapping.shape[1] == n_alleles
-    )  # these are not the same, work out what's going on - try running code with debugger? or print statementsd
+    assert mapping.shape[1] == n_alleles
 
     # Create output array.
     out = np.empty(shape=(n_sites, max_allele + 1), dtype=x.dtype)