Fixup genotypes iterator and call_LAA/call_LAD

Author: jeromekelleher

bio2zarr/core.py

@@ -63,6 +63,27 @@ def chunk_aligned_slices(z, n, max_chunks=None):
     return slices
 
 
+def first_dim_slice_iter(z, start, stop):
+    """
+    Efficiently iterate over the specified slice of the first dimension of the zarr
+    array z.
+    """
+    chunk_size = z.chunks[0]
+    first_chunk = start // chunk_size
+    last_chunk = (stop // chunk_size) + (stop % chunk_size != 0)
+    for chunk in range(first_chunk, last_chunk):
+        Z = z.blocks[chunk]
+        chunk_start = chunk * chunk_size
+        chunk_stop = chunk_start + chunk_size
+        slice_start = None
+        if start > chunk_start:
+            slice_start = start - chunk_start
+        slice_stop = None
+        if stop < chunk_stop:
+            slice_stop = stop - chunk_start
+        yield from Z[slice_start:slice_stop]
+
+
 def du(path):
     """
     Return the total bytes stored at this path.

bio2zarr/vcf2zarr/vcz.py

@@ -550,15 +550,20 @@ def compute_laa_field(genotypes):
 
 
 def compute_lad_field(ad, laa):
-    lad = np.full((ad.shape[0], 2), -2, dtype=ad.dtype)
-    ref_ref = np.where((laa[:, 0] == -2) & (laa[:, 1] == -2))[0]
-    lad[ref_ref, 0] = ad[ref_ref, 0]
-    ref_alt = np.where((laa[:, 0] != -2) & (laa[:, 1] == -2))[0]
-    lad[ref_alt, 0] = ad[ref_alt, 0]
-    lad[ref_alt, 1] = ad[ref_alt, laa[ref_alt, 0]]
-    alt_alt = np.where((laa[:, 0] != -2) & (laa[:, 1] != -2))[0]
-    lad[alt_alt, 0] = ad[alt_alt, laa[alt_alt, 0]]
-    lad[alt_alt, 1] = ad[alt_alt, laa[alt_alt, 1]]
+    try:
+        lad = np.full((ad.shape[0], 2), -2, dtype=ad.dtype)
+        ref_ref = np.where((laa[:, 0] == -2) & (laa[:, 1] == -2))[0]
+        lad[ref_ref, 0] = ad[ref_ref, 0]
+        ref_alt = np.where((laa[:, 0] != -2) & (laa[:, 1] == -2))[0]
+        lad[ref_alt, 0] = ad[ref_alt, 0]
+        lad[ref_alt, 1] = ad[ref_alt, laa[ref_alt, 0]]
+        alt_alt = np.where((laa[:, 0] != -2) & (laa[:, 1] != -2))[0]
+        lad[alt_alt, 0] = ad[alt_alt, laa[alt_alt, 0]]
+        lad[alt_alt, 1] = ad[alt_alt, laa[alt_alt, 1]]
+    except Exception as e:
+        print("ad = ", ad)
+        print("laa = ", laa)
+        raise e
     return lad
 
 
@@ -875,21 +880,22 @@ def encode_local_alleles_partition(self, partition_index):
         call_AD_source = self.icf.fields["FORMAT/AD"].iter_values(
             partition.start, partition.stop
         )
-
         gt_array = zarr.open_array(
             store=self.wip_partition_array_path(partition_index, "call_genotype"),
             mode="r",
         )
-        for chunk_index in range(gt_array.cdata_shape[0]):
-            for genotypes in gt_array.blocks[chunk_index]:
-                laa = compute_laa_field(genotypes)
-                j = call_LAA.next_buffer_row()
-                call_LAA.buff[j] = laa
-
-                ad = next(call_AD_source)
-                j = call_LAD.next_buffer_row()
-                lad = compute_lad_field(ad, laa)
-                call_LAD.buff[j] = lad
+        for genotypes in core.first_dim_slice_iter(
+            gt_array, partition.start, partition.stop
+        ):
+            laa = compute_laa_field(genotypes)
+            j = call_LAA.next_buffer_row()
+            call_LAA.buff[j] = laa
+
+            ad = next(call_AD_source)
+            k = call_LAD.next_buffer_row()
+            assert j == k
+            lad = compute_lad_field(ad, laa)
+            call_LAD.buff[j] = lad
 
         call_LAA.flush()
         self.finalise_partition_array(partition_index, "call_LAA")

tests/test_core.py

@@ -194,6 +194,41 @@ def test_5_chunk_1(self, n, expected):
         assert result == expected
 
 
+class TestFirstDimSliceIter:
+    @pytest.mark.parametrize("chunk_size", [1, 3, 4, 5])
+    @pytest.mark.parametrize(
+        ("size", "start", "stop"),
+        [
+            (10, 0, 4),
+            (10, 0, 8),
+            (10, 0, 10),
+            (10, 4, 4),
+            (10, 4, 8),
+            (10, 4, 10),
+            (10, 0, 5),
+            (10, 0, 3),
+            (10, 0, 9),
+            (10, 1, 5),
+            (10, 1, 1),
+            (10, 1, 2),
+            (10, 1, 3),
+            (10, 1, 4),
+            (10, 1, 10),
+            (10, 5, 5),
+            (10, 5, 6),
+            (10, 5, 7),
+            (5, 0, 5),
+            (5, 1, 1),
+            (5, 1, 3),
+        ],
+    )
+    def test_examples(self, chunk_size, size, start, stop):
+        a = np.arange(size, dtype=int)
+        z = zarr.empty(size, chunks=chunk_size, dtype=int)
+        z[:] = a
+        assert list(core.first_dim_slice_iter(z, start, stop)) == list(a[start:stop])
+
+
 @pytest.mark.skipif(sys.platform != "linux", reason="Only valid on Linux")
 @pytest.mark.parametrize(
     ("path", "expected"),