Make normalise work chunk by chunk

README.md

@@ -18,7 +18,7 @@ The append operation works by appending a VCF Zarr file to the store (which is a
 
 Note that only new _samples_ can be added - not new variants - so the store contains a fixed set of variants that must be known before it is created. This is usually not a limitation since the samples come from the same genotype array, or even from the same reference panel for imputed data.
 
-If the source VCFs have different (but typically overlapping) sets of variants, then they need to be harmonised with the full set of variants before being converted to VCF Zarr. This can be accomplished using bcftools merge or similar tools - it is not performed by vczstore itself.
+If the source VCFs have different (but typically overlapping) sets of variants, then they need to be harmonised with the full set of variants before being converted to VCF Zarr. This can be accomplished by running `vczstore normalise` before running append.
 
 The append operation will perform a check that the contig, position and allele (REF and ALT) fields all match before performing the update. It will fail if there is a mismatch, so samples with inconsistent variant sets cannot be appended. The check is strict - allele ordering must match exactly too.
 

tests/test_normalise.py

@@ -8,6 +8,7 @@
 from vczstore.normalise import (
     index_variants,
     normalise,
+    remap_genotypes,
     variant_alleles_are_equivalent,
 )
 
@@ -23,6 +24,8 @@ def make_vcz(
     variants_chunk_size=None,
     samples_chunk_size=None,
     call_genotype=None,
+    call_fields=None,
+    call_field_dims=None,
     ploidy=2,
 ):
 
@@ -97,6 +100,16 @@ def make_vcz(
             compressors=None,
             filters=None,
         )
+    if call_fields is not None:
+        for name, data in call_fields.items():
+            root.create_array(
+                name=f"call_{name}",
+                data=data,
+                chunks=(v_chunk, s_chunk, data.shape[2]),
+                dimension_names=call_field_dims[name],
+                compressors=None,
+                filters=None,
+            )
     return store
 
 
@@ -141,6 +154,28 @@ def test_variant_alleles_are_equivalent(
     assert_array_equal(updated, expected_updated)
 
 
+def test_remap_genotypes__no_op():
+    gt = np.array([[[0, 1]], [[1, 0]]], dtype=np.int8)
+    original = gt.copy()
+    remap_genotypes(gt, [], [])
+    assert_array_equal(gt, original)
+
+
+def test_remap_genotypes__single_variant():
+    # alleles [A, T, C] remapped to [A, C, T]: mapping is [0, 2, 1]
+    gt = np.array([[[0, 1]], [[2, 1]]], dtype=np.int8)
+    remap_genotypes(gt, [1], [np.array([0, 2, 1])])
+    assert_array_equal(gt, [[[0, 1]], [[1, 2]]])
+
+
+def test_remap_genotypes__multiple_variants():
+    gt = np.array([[[1, 2]], [[3, 0]]], dtype=np.int8)
+    # variant 0: vcz2=[A, C, T], vcz1=[A, T, C]: mapping [0, 2, 1]
+    # variant 1: vcz2=[A, G, T, C], vcz1=[A, T, C, G]: mapping [0, 3, 1, 2]
+    remap_genotypes(gt, [0, 1], [np.array([0, 2, 1]), np.array([0, 3, 1, 2])])
+    assert_array_equal(gt, [[[2, 1]], [[2, 0]]])
+
+
 def test_index_variants__success_subset():
     vcz1 = make_vcz([0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]])
     vcz2 = make_vcz([0, 0], [1, 3], [["A", "T"], ["T", "A"]])
@@ -207,7 +242,7 @@ def test_index_variants__new_allele():
         index_variants(vcz1, vcz2)
 
 
-@pytest.mark.parametrize("variants_chunk_size", [None, 5])
+@pytest.mark.parametrize("variants_chunk_size", [None, 1, 3, 4, 5, 10])
 def test_normalise(variants_chunk_size):
     vcz1 = make_vcz(
         variant_contig=[0, 0, 0, 0, 0, 0, 0, 0, 0],
@@ -266,3 +301,26 @@ def test_normalise(variants_chunk_size):
             [[0, 1]],  # remapped to vcz1 order
         ],
     )
+
+
+def test_normalise__other_call_fields_not_implemented():
+    vcz1 = make_vcz(
+        variant_contig=[0],
+        variant_position=[1],
+        alleles=[["A", "T", "G"]],
+    )
+
+    vcz2 = make_vcz(
+        variant_contig=[0],
+        variant_position=[1],
+        alleles=[["A", "G", "T"]],  # order different to vcz1
+        sample_id=["S1"],
+        call_genotype=[[[0, 2]]],
+        call_fields={"AD": np.array([[[10, -1, 20]]], dtype=np.int8)},
+        call_field_dims={"AD": ["variants", "samples", "alleles"]},
+    )
+
+    vcz2_norm = zarr.storage.MemoryStore()
+
+    with pytest.raises(NotImplementedError):
+        normalise(vcz1, vcz2, vcz2_norm)

vczstore/normalise.py

@@ -6,7 +6,7 @@
 from vcztools.retrieval import VczReader
 from vcztools.utils import array_dims, search
 
-from vczstore.utils import missing_val
+from vczstore.utils import missing_val, variant_chunk_slices
 
 
 def normalise(vcz1, vcz2, vcz2_norm):
@@ -35,6 +35,14 @@ def normalise(vcz1, vcz2, vcz2_norm):
         if var.startswith("call_"):
             # normalise genotype fields
             arr = root2[var]
+            for dim in array_dims(arr):
+                if dim in ("alleles", "alt_alleles", "genotypes") or dim.startswith(
+                    "local_"
+                ):
+                    raise NotImplementedError(
+                        f"Allele remapping not supported for dim {dim} in "
+                        f"variable {var}"
+                    )
             shape = (n_variants,) + arr.shape[1:]
             chunks = (variants_chunk_size,) + arr.chunks[1:]
             create_empty_group_array(
@@ -73,34 +81,57 @@ def normalise(vcz1, vcz2, vcz2_norm):
             arr = root1[var]
         norm_root[var][:] = arr[...]
 
+    # turn bool indexes into int array indexes
+    match_idx = np.where(index)[0]
+    remap_idx = np.where(remap_alleles)[0]
+
+    # find chunk boundaries
+    chunk_bounds = np.arange(0, n_variants, step=variants_chunk_size)
+    chunk_bounds = np.append(chunk_bounds, [n_variants])
+
+    # find chunk offsets for indexes
+    match_starts = np.searchsorted(match_idx, chunk_bounds)
+    remap_starts = np.searchsorted(remap_idx, chunk_bounds)
+
+    allele_mappings_list = list(allele_mappings.values())
+
     # copy variant-chunked arrays
-    # TODO: do this chunk-by-chunk
-    for var, arr in root2.arrays():
-        if array_dims(arr)[0] != "variants":
-            continue
-        # copy call arrays from vcz2, everything else from vcz1
-        if var.startswith("call_"):
-            arr = root2[var]
-            shape = (n_variants,) + arr.shape[1:]
-            data = np.full(shape, fill_value=missing_val(arr), dtype=arr.dtype)
-            data[index, ...] = arr[...]
-            if var == "call_genotype":
-                remap_genotypes(data, allele_mappings)
-            norm_root[var][:] = data
-        else:
-            arr = root1[var]
-            norm_root[var][:] = arr[...]
+    for i, v_sel in enumerate(variant_chunk_slices(root1)):
+        for var, arr in root2.arrays():
+            if array_dims(arr)[0] != "variants":
+                continue
+            # copy genotype fields from vcz2, everything else from vcz1
+            if var.startswith("call_"):
+                arr = root2[var]
+                chunk_n = v_sel.stop - v_sel.start
+                shape = (chunk_n,) + arr.shape[1:]
+                data = np.full(shape, fill_value=missing_val(arr), dtype=arr.dtype)
 
+                match_sl = slice(match_starts[i], match_starts[i + 1])
+                local_idx = match_idx[match_sl] - v_sel.start
+                data[local_idx, ...] = arr[match_sl, ...]
 
-def remap_genotypes(gt, allele_mappings):
+                if var == "call_genotype":
+                    remap_sl = slice(remap_starts[i], remap_starts[i + 1])
+                    local_remap_idx = remap_idx[remap_sl] - v_sel.start
+                    chunk_maps = allele_mappings_list[remap_sl]
+                    remap_genotypes(data, local_remap_idx, chunk_maps)
+
+                norm_root[var][v_sel] = data
+            else:
+                arr = root1[var]
+                norm_root[var][v_sel] = arr[v_sel, ...]
+
+
+def remap_genotypes(gt, indices, mappings):
     """Update a genotype array in-place by remapping allele indices.
 
-    allele_mappings is a dict of {variant_index: int_array} as returned by
-    index_variants.
+    indices and mappings are parallel arrays of variant positions and their
+    allele index remappings.
     """
     num_samples = gt.shape[1]
     ploidy = gt.shape[2]
-    for i, mapping in allele_mappings.items():
+    for i, mapping in zip(indices, mappings):
         for j in range(num_samples):
             for k in range(ploidy):
                 val = gt[i, j, k]
@@ -119,6 +150,10 @@ def index_variants(vcz1, vcz2):
             remapping.
         updated_allele_mappings: dict {variant_index: str_array} of the updated
             (merged) alleles for variants where vcz2 has extra alleles not in vcz1.
+
+    Note that the allele mappings are dicts which only contain sites where there
+    is a remapping. This is an efficient way to store allele mappings, since they
+    are rare, and are not known ahead of time.
     """
     root1 = zarr.open(vcz1, mode="r")
     root2 = zarr.open(vcz2, mode="r")