Allow variant/ALT to be 1D

Author: tomwhite

sgkit/io/vcfzarr_reader.py

@@ -10,6 +10,12 @@
 from ..utils import encode_array
 
 
+def _ensure_2d(arr: ArrayLike) -> ArrayLike:
+    if arr.ndim == 1:
+        arr = arr.reshape(-1, 1)
+    return arr
+
+
 def read_vcfzarr(path: PathType) -> xr.Dataset:
     """Read a VCF Zarr file.
 
@@ -52,7 +58,7 @@ def max_str_len(arr: ArrayLike) -> Any:
         da.compute(max_str_len(variants_ref), max_str_len(variants_alt))
     )
     variants_ref_alt = da.concatenate(
-        [variants_ref.reshape(-1, 1), variants_alt], axis=1
+        [_ensure_2d(variants_ref), _ensure_2d(variants_alt)], axis=1
     )
     variant_alleles = variants_ref_alt.astype(f"S{max_allele_length}")
 

sgkit/tests/test_vcfzarr_reader.py

@@ -2,6 +2,7 @@
 from numpy.testing import assert_array_equal
 
 from sgkit import read_vcfzarr
+from sgkit.io.vcfzarr_reader import _ensure_2d
 
 
 def test_read_vcfzarr(shared_datadir):
@@ -62,3 +63,8 @@ def test_read_vcfzarr(shared_datadir):
     assert_array_equal(ds["call/genotype"], call_genotype)
     assert_array_equal(ds["call/genotype_mask"], call_genotype < 0)
     assert "call/genotype_phased" not in ds
+
+
+def test_ensure_2d():
+    assert_array_equal(_ensure_2d(np.array([0, 2, 1])), np.array([[0], [2], [1]]))
+    assert_array_equal(_ensure_2d(np.array([[0], [2], [1]])), np.array([[0], [2], [1]]))