Normalise variants with allele remapping

tests/test_normalise.py

@@ -1,115 +1,190 @@
-import pathlib
 import re
 
 import numpy as np
 import pytest
 import zarr
+from numpy.testing import assert_array_equal
 
-from vczstore.zarr_impl import index_variants, normalise, variant_alleles_are_equivalent
+from vczstore.normalise import (
+    index_variants,
+    normalise,
+    variant_alleles_are_equivalent,
+)
+
+
+# TODO: replace this with make_vcz in vcztools/bio2zarr
+# see https://github.com/sgkit-dev/vczstore/issues/37
+def make_vcz(
+    variant_contig,
+    variant_position,
+    alleles,
+    *,
+    sample_id=None,
+    variants_chunk_size=None,
+    samples_chunk_size=None,
+    call_genotype=None,
+    ploidy=2,
+):
 
-from .utils import compare_vcf_and_vcz, convert_vcf_to_vcz
+    variant_contig = np.asarray(variant_contig, dtype=np.int32)
+    variant_position = np.asarray(variant_position, dtype=np.int32)
+    num_variants = variant_contig.shape[0]
+    max_alleles = max(len(a) for a in alleles) if len(alleles) > 0 else 1
+    allele_array = np.full((num_variants, max_alleles), "", dtype="<U16")
+    for i, row in enumerate(alleles):
+        for j, a in enumerate(row):
+            allele_array[i, j] = a
 
+    num_samples = len(sample_id) if sample_id is not None else 0
+
+    v_chunk = variants_chunk_size if variants_chunk_size is not None else num_variants
+    v_chunk = max(v_chunk, 1)
+    s_chunk = (
+        samples_chunk_size if samples_chunk_size is not None else max(num_samples, 1)
+    )
 
-def create_variant_only_vcz(variant_contig, variant_position, variant_allele):
     store = zarr.storage.MemoryStore()
     root = zarr.create_group(store=store)
     root.create_array(
         name="contig_id",
         data=np.unique(variant_contig).astype(str),
         dimension_names=["contigs"],
+        compressors=None,
+        filters=None,
     )
     root.create_array(
         name="variant_contig",
-        data=np.array(variant_contig, dtype=np.int32),
+        data=variant_contig,
+        chunks=(v_chunk,),
         dimension_names=["variants"],
+        compressors=None,
+        filters=None,
     )
     root.create_array(
         name="variant_position",
-        data=np.array(variant_position, dtype=np.int32),
+        data=variant_position,
+        chunks=(v_chunk,),
         dimension_names=["variants"],
+        compressors=None,
+        filters=None,
     )
     root.create_array(
         name="variant_allele",
-        data=np.array(variant_allele, dtype="T"),
+        data=allele_array,
+        chunks=(v_chunk, max_alleles),
         dimension_names=["variants", "alleles"],
+        compressors=None,
+        filters=None,
     )
+    if sample_id is None:
+        sample_id_arr = np.array([], dtype="<U64")
+    else:
+        sample_id_arr = np.asarray(sample_id, dtype="<U64")
     root.create_array(
-        name="sample_id", data=np.array([], dtype="T"), dimension_names=["samples"]
+        name="sample_id",
+        data=sample_id_arr,
+        dimension_names=["samples"],
+        compressors=None,
+        filters=None,
     )
+    if call_genotype is not None:
+        call_genotype = np.asarray(call_genotype, dtype=np.int8)
+        root.create_array(
+            name="call_genotype",
+            data=call_genotype,
+            chunks=(v_chunk, s_chunk, ploidy),
+            dimension_names=["variants", "samples", "ploidy"],
+            compressors=None,
+            filters=None,
+        )
     return store
 
 
 @pytest.mark.parametrize(
-    ("variant_allele1", "variant_allele2", "expected"),
+    (
+        "variant_allele1",
+        "variant_allele2",
+        "expected_matched",
+        "expected_mapping",
+        "expected_updated",
+    ),
     [
-        ([], [], True),
-        (["A"], ["A"], True),
-        (["A", "T"], ["A", "T"], True),
-        (["A", ""], ["A"], True),
-        (["A", "."], ["A"], True),
-        (["A", "T", "."], ["A", "T", ".", "."], True),
-        (["A", "T", "."], ["A", "T", "", ""], True),
-        (["."], ["A"], False),
-        (["A", ".", "T"], ["A", "T"], False),
-        (["A", ".", "T"], ["A", "T", "."], False),
+        (["A"], ["A"], True, None, None),
+        (["A", "T"], ["A", "T"], True, None, None),
+        (["A", ""], ["A"], True, None, None),
+        (["A", "."], ["A"], True, None, None),
+        (["A", "T", "."], ["A", "T", ".", "."], True, None, None),
+        (["A", "T", "."], ["A", "T", "", ""], True, None, None),
+        # overlapping alt alleles
+        (["A", "C", "T"], ["A", "T"], True, [0, 2], None),
+        # new alt alleles
+        (["A", "T"], ["A", "T", "C"], True, [0, 1, 2], ["A", "T", "C"]),
+        (["A", "T"], ["A", "C", "T"], True, [0, 2, 1], ["A", "T", "C"]),
+        # indels
+        (["AT", "A"], ["A", "AT"], False, None, None),
+        # different alt allele
+        (["A", "C"], ["A", "T"], False, None, None),
     ],
 )
-def test_variant_alleles_are_equivalent(variant_allele1, variant_allele2, expected):
-    assert variant_alleles_are_equivalent(variant_allele1, variant_allele2) == expected
+def test_variant_alleles_are_equivalent(
+    variant_allele1,
+    variant_allele2,
+    expected_matched,
+    expected_mapping,
+    expected_updated,
+):
+    matched, mapping, updated = variant_alleles_are_equivalent(
+        np.array(variant_allele1), np.array(variant_allele2)
+    )
+    assert matched == expected_matched
+    assert_array_equal(mapping, expected_mapping)
+    assert_array_equal(updated, expected_updated)
 
 
 def test_index_variants__success_subset():
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]]
-    )
-    vcz2 = create_variant_only_vcz([0, 0], [1, 3], [["A", "T"], ["T", "A"]])
-    np.testing.assert_array_equal(index_variants(vcz1, vcz2), [True, False, True])
+    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"]])
+    index, *_ = index_variants(vcz1, vcz2)
+    assert_array_equal(index, [True, False, True])
 
 
 def test_index_variants__success_repeated_site():
     # example where multi-allelic sites are split to biallelic
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [1, 1, 3], [["A", "T"], ["A", "C"], ["T", "A"]]
-    )
-    vcz2 = create_variant_only_vcz([0, 0], [1, 1], [["A", "T"], ["A", "C"]])
-    np.testing.assert_array_equal(index_variants(vcz1, vcz2), [True, True, False])
+    vcz1 = make_vcz([0, 0, 0], [1, 1, 3], [["A", "T"], ["A", "C"], ["T", "A"]])
+    vcz2 = make_vcz([0, 0], [1, 1], [["A", "T"], ["A", "C"]])
+    index, *_ = index_variants(vcz1, vcz2)
+    assert_array_equal(index, [True, True, False])
 
 
 def test_index_variants__order_mismatch():
     # example where multi-allelic sites are split to biallelic
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [1, 1, 3], [["A", "T"], ["A", "C"], ["T", "A"]]
-    )
+    vcz1 = make_vcz([0, 0, 0], [1, 1, 3], [["A", "T"], ["A", "C"], ["T", "A"]])
     # note that the allele ordering is different for contig 0, position 1
-    vcz2 = create_variant_only_vcz([0, 0], [1, 1], [["A", "C"], ["A", "T"]])
+    vcz2 = make_vcz([0, 0], [1, 1], [["A", "C"], ["A", "T"]])
     with pytest.raises(
         match=re.escape(
-            "Variant not in (or out of order in) first vcz: "
+            "Variant in vcz2 not found in vcz1 (or vcz2 is out of order): "
             "variant_contig=0, variant_position=1, variant_allele=['A', 'T']"
         )
     ):
         index_variants(vcz1, vcz2)
 
 
 def test_index_variants__new_variant():
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [2, 3, 4], [["A", "T"], ["C", "G"], ["T", "A"]]
-    )
-    vcz2 = create_variant_only_vcz([0, 0, 0], [1, 2], [["A", "."], ["A", "T"]])
+    vcz1 = make_vcz([0, 0, 0], [2, 3, 4], [["A", "T"], ["C", "G"], ["T", "A"]])
+    vcz2 = make_vcz([0, 0, 0], [1, 2], [["A", "."], ["A", "T"]])
     with pytest.raises(
         match=re.escape(
-            "Variant not in (or out of order in) first vcz: "
+            "Variant in vcz2 not found in vcz1 (or vcz2 is out of order): "
             "variant_contig=0, variant_position=1, variant_allele=['A', '.']"
         )
     ):
         index_variants(vcz1, vcz2)
 
 
 def test_index_variants__new_variant_at_end():
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]]
-    )
-    vcz2 = create_variant_only_vcz([0, 0, 0], [1, 4], [["A", "T"], ["G", "A"]])
+    vcz1 = make_vcz([0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]])
+    vcz2 = make_vcz([0, 0, 0], [1, 4], [["A", "T"], ["G", "A"]])
     with pytest.raises(
         match=re.escape(
             "Variant not in first vcz: "
@@ -120,11 +195,9 @@ def test_index_variants__new_variant_at_end():
 
 
 def test_index_variants__new_allele():
-    vcz1 = create_variant_only_vcz(
-        [0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]]
-    )
+    vcz1 = make_vcz([0, 0, 0], [1, 2, 3], [["A", "T"], ["C", "G"], ["T", "A"]])
     # variant at contig 0, position 3 has different alleles
-    vcz2 = create_variant_only_vcz([0, 0], [1, 3], [["A", "T"], ["T", "G"]])
+    vcz2 = make_vcz([0, 0], [1, 3], [["A", "T"], ["T", "G"]])
     with pytest.raises(
         match=re.escape(
             "Variant not in first vcz: "
@@ -134,21 +207,62 @@ def test_index_variants__new_allele():
         index_variants(vcz1, vcz2)
 
 
-# bcftools merge tests/data/vcf/alleles-variants.vcf.gz \
-#  tests/data/vcf/alleles-1.vcf.gz \
-#  -o tests/data/vcf/alleles-1-norm.vcf.gz -W=csi
-def test_normalise(tmp_path):
-    vcz0 = convert_vcf_to_vcz("alleles-variants.vcf.gz", tmp_path, ploidy=2)
-    vcz1 = convert_vcf_to_vcz("alleles-1.vcf.gz", tmp_path)
+@pytest.mark.parametrize("variants_chunk_size", [None, 5])
+def test_normalise(variants_chunk_size):
+    vcz1 = make_vcz(
+        variant_contig=[0, 0, 0, 0, 0, 0, 0, 0, 0],
+        variant_position=[1, 2, 3, 4, 4, 5, 5, 6, 7],
+        alleles=[
+            ["A", "T"],
+            ["A", "C"],
+            ["A", "G"],
+            ["A", "C"],
+            ["A", "G", "T"],
+            ["A", "G"],
+            ["A", "T", "C"],
+            ["A", "G", "T"],
+            ["A", "T", "G"],
+        ],
+        variants_chunk_size=variants_chunk_size,
+    )
+
+    vcz2 = make_vcz(
+        variant_contig=[0, 0, 0, 0, 0, 0],
+        variant_position=[1, 2, 4, 5, 6, 7],
+        alleles=[
+            ["A", "T"],
+            ["A", "C"],
+            ["A", "C"],
+            ["A", "T", "C"],
+            ["A", "G"],
+            ["A", "G", "T"],  # order different to vcz1
+        ],
+        sample_id=["S1"],
+        call_genotype=[[[0, 0]], [[0, 1]], [[0, 0]], [[1, 1]], [[0, 1]], [[0, 2]]],
+    )
+
+    vcz2_norm = zarr.storage.MemoryStore()
 
-    vcz1_norm = pathlib.Path(tmp_path) / "alleles-1-norm.vcz"
+    normalise(vcz1, vcz2, vcz2_norm)
 
-    normalise(vcz0, vcz1, vcz1_norm)
+    root1 = zarr.open(vcz1)
+    root_norm = zarr.open(vcz2_norm)
 
-    compare_vcf_and_vcz(
-        tmp_path,
-        "view --no-version",
-        "alleles-1-norm.vcf.gz",
-        "view --no-version",
-        vcz1_norm,
+    assert_array_equal(root_norm["variant_contig"][:], root1["variant_contig"][:])
+    assert_array_equal(root_norm["variant_position"][:], root1["variant_position"][:])
+    assert_array_equal(root_norm["variant_allele"][:], root1["variant_allele"][:])
+    assert_array_equal(root_norm["sample_id"][:], ["S1"])
+    assert_array_equal(
+        root_norm["call_genotype"][:],
+        [
+            [[0, 0]],
+            [[0, 1]],
+            [[-1, -1]],
+            [[0, 0]],
+            [[-1, -1]],
+            [[-1, -1]],
+            [[1, 1]],
+            [[0, 1]],
+            [[0, 1]],  # remapped to vcz1 order
+        ],
     )

vczstore/cli.py

@@ -61,7 +61,7 @@ def append(vcz1, vcz2, impl, zarr_backend_storage):
 @click.argument("vcz2_norm", type=click.Path())
 def normalise(vcz1, vcz2, vcz2_norm):
     """Normalise variants in vcz2 with respect to vcz1 and write to vcz2_norm"""
-    from vczstore.zarr_impl import normalise as normalise_function
+    from vczstore.normalise import normalise as normalise_function
 
     normalise_function(vcz1, vcz2, vcz2_norm)