Add read_vcfzarr

Author: tomwhite

requirements-dev.txt

@@ -3,4 +3,5 @@ pre-commit
 pytest
 pytest-cov
 hypothesis
-statsmodels
+pytest-datadir
+statsmodels

requirements.txt

@@ -1,4 +1,4 @@
+dask[array]
 numpy
 xarray
-dask[array]
 scipy

setup.cfg

@@ -54,19 +54,19 @@ ignore =
 profile = black
 default_section = THIRDPARTY
 known_first_party = sgkit
-known_third_party = dask,hypothesis,numpy,pandas,pytest,setuptools,xarray
+known_third_party = dask,hypothesis,numpy,pandas,pytest,setuptools,xarray,zarr
 multi_line_output = 3
 include_trailing_comma = True
 force_grid_wrap = 0
 use_parentheses = True
 line_length = 88
 
+[mypy-dask.*]
+ignore_missing_imports = True
 [mypy-numpy.*]
 ignore_missing_imports = True
 [mypy-pandas.*]
 ignore_missing_imports = True
-[mypy-dask.*]
-ignore_missing_imports = True
 [mypy-pytest.*]
 ignore_missing_imports = True
 [mypy-statsmodels.*]
@@ -75,5 +75,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-setuptools]
 ignore_missing_imports = True
+[mypy-zarr]
+ignore_missing_imports = True
 [mypy-sgkit.tests.*]
+disallow_untyped_calls = False
 disallow_untyped_defs = False

sgkit/__init__.py

@@ -6,6 +6,7 @@
     create_genotype_call_dataset,
     create_genotype_dosage_dataset,
 )
+from .io.vcfzarr_reader import read_vcfzarr
 from .stats.aggregation import count_alleles
 from .stats.association import gwas_linear_regression
 
@@ -18,4 +19,5 @@
     "count_alleles",
     "create_genotype_dosage_dataset",
     "gwas_linear_regression",
+    "read_vcfzarr",
 ]

sgkit/api.py

@@ -33,7 +33,7 @@ def create_genotype_call_dataset(
         The (index of the) contig for each variant.
     variant_position : array_like, int
         The reference position of the variant.
-    variant_alleles : array_like, S1
+    variant_alleles : array_like, zero-terminated bytes, e.g. "S1"
         The possible alleles for the variant.
     sample_id : array_like, str
         The unique identifier of the sample.

sgkit/io/__init__.py

@@ -0,0 +1,77 @@
+from typing import Any
+
+import dask.array as da
+import numpy as np
+import xarray as xr
+import zarr
+
+from ..api import DIM_VARIANT, create_genotype_call_dataset
+from ..typing import ArrayLike, PathType
+from ..utils import encode_array
+
+
+def read_vcfzarr(path: PathType) -> xr.Dataset:
+    """Read a VCF Zarr file.
+
+    Loads VCF variant, sample, and genotype data as Dask arrays within a Dataset
+    from a Zarr file created using scikit-allel's `vcf_to_zarr` function.
+
+    Since `vcf_to_zarr` does not preserve phasing information, there is no
+    `call/genotype_phased` variable in the resulting dataset.
+
+    Parameters
+    ----------
+    path : PathType
+        Path to the Zarr file.
+
+    Returns
+    -------
+    xr.Dataset
+        The dataset of genotype calls, created using `create_genotype_call_dataset`.
+    """
+
+    vcfzarr = zarr.open_group(str(path), mode="r")
+
+    # Index the contig names
+    variants_chrom = da.from_zarr(vcfzarr["variants/CHROM"]).astype(str)
+    variant_contig, variant_contig_names = encode_array(variants_chrom.compute())
+    variant_contig = variant_contig.astype("int16")
+    variant_contig_names = list(variant_contig_names)
+
+    # For variant alleles, combine REF and ALT into a single array
+    # and calculate the number of alleles so we can set the dtype correctly
+    variants_ref = da.from_zarr(vcfzarr["variants/REF"])
+    variants_alt = da.from_zarr(vcfzarr["variants/ALT"])
+
+    def max_str_len(arr: ArrayLike) -> Any:
+        return arr.map_blocks(
+            lambda s: np.char.str_len(s.astype(str)), dtype=np.int8
+        ).max()
+
+    max_allele_length = max(
+        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
+    )
+    variant_alleles = variants_ref_alt.astype(f"S{max_allele_length}")
+
+    variants_id = da.from_zarr(vcfzarr["variants/ID"]).astype(str)
+
+    ds = create_genotype_call_dataset(
+        variant_contig_names=variant_contig_names,
+        variant_contig=variant_contig,
+        variant_position=da.from_zarr(vcfzarr["variants/POS"]),
+        variant_alleles=variant_alleles,
+        sample_id=da.from_zarr(vcfzarr["samples"]).astype(str),
+        call_genotype=da.from_zarr(vcfzarr["calldata/GT"]),
+        variant_id=variants_id,
+    )
+
+    # Add a mask for variant ID
+    ds["variant/id_mask"] = (
+        [DIM_VARIANT],
+        variants_id == ".",
+    )
+
+    return ds

sgkit/io/vcfzarr_reader.py

@@ -0,0 +1,31 @@
+##fileformat=VCFv4.0
+##fileDate=20090805
+##source=myImputationProgramV3.1
+##reference=1000GenomesPilot-NCBI36
+##phasing=partial
+##INFO=<ID=NS,Number=1,Type=Integer,Description="Number of Samples With Data">
+##INFO=<ID=AN,Number=1,Type=Integer,Description="Total number of alleles in called genotypes">
+##INFO=<ID=AC,Number=.,Type=Integer,Description="Allele count in genotypes, for each ALT allele, in the same order as listed">
+##INFO=<ID=DP,Number=1,Type=Integer,Description="Total Depth">
+##INFO=<ID=AF,Number=.,Type=Float,Description="Allele Frequency">
+##INFO=<ID=AA,Number=1,Type=String,Description="Ancestral Allele">
+##INFO=<ID=DB,Number=0,Type=Flag,Description="dbSNP membership, build 129">
+##INFO=<ID=H2,Number=0,Type=Flag,Description="HapMap2 membership">
+##FILTER=<ID=s50,Description="Less than 50% of samples have data">
+##FILTER=<ID=q10,Description="Quality below 10">
+##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
+##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">
+##FORMAT=<ID=DP,Number=1,Type=Integer,Description="Read Depth">
+##FORMAT=<ID=HQ,Number=2,Type=Integer,Description="Haplotype Quality">
+##ALT=<ID=DEL:ME:ALU,Description="Deletion of ALU element">
+##ALT=<ID=CNV,Description="Copy number variable region">
+#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	NA00001	NA00002	NA00003
+19	111	.	A	C	9.6	.	.	GT:HQ	0|0:10,15	0|0:10,10	0/1:3,3
+19	112	.	A	G	10	.	.	GT:HQ	0|0:10,10	0|0:10,10	0/1:3,3
+20	14370	rs6054257	G	A	29	PASS	NS=3;DP=14;AF=0.5;DB;H2	GT:GQ:DP:HQ	0|0:48:1:51,51	1|0:48:8:51,51	1/1:43:5:.,.
+20	17330	.	T	A	3	q10	NS=3;DP=11;AF=0.017	GT:GQ:DP:HQ	0|0:49:3:58,50	0|1:3:5:65,3	0/0:41:3:.,.
+20	1110696	rs6040355	A	G,T	67	PASS	NS=2;DP=10;AF=0.333,0.667;AA=T;DB	GT:GQ:DP:HQ	1|2:21:6:23,27	2|1:2:0:18,2	2/2:35:4:.,.
+20	1230237	.	T	.	47	PASS	NS=3;DP=13;AA=T	GT:GQ:DP:HQ	0|0:54:.:56,60	0|0:48:4:51,51	0/0:61:2:.,.
+20	1234567	microsat1	G	GA,GAC	50	PASS	NS=3;DP=9;AA=G;AN=6;AC=3,1	GT:GQ:DP	0/1:.:4	0/2:17:2	./.:40:3
+20	1235237	.	T	.	.	.	.	GT	0/0	0|0	./.
+X	10	rsTest	AC	A,ATG,C	10	PASS	.	GT	0	0/1	0|2

sgkit/tests/data/sample.vcf

@@ -0,0 +1,64 @@
+import numpy as np
+from numpy.testing import assert_array_equal
+
+from sgkit import read_vcfzarr
+
+
+def test_read_vcfzarr(shared_datadir):
+    # The file sample.vcf.zarr.zip was created by running the following
+    # in a python session with the scikit-allel package installed.
+    #
+    # import allel
+    # allel.vcf_to_zarr("sample.vcf", "sample.vcf.zarr.zip")
+
+    path = shared_datadir / "sample.vcf.zarr.zip"
+    ds = read_vcfzarr(path)
+
+    assert ds.attrs["contigs"] == ["19", "20", "X"]
+    assert_array_equal(ds["variant/contig"], [0, 0, 1, 1, 1, 1, 1, 1, 2])
+    assert_array_equal(
+        ds["variant/position"],
+        [111, 112, 14370, 17330, 1110696, 1230237, 1234567, 1235237, 10],
+    )
+    assert_array_equal(
+        ds["variant/alleles"],
+        [
+            [b"A", b"C", b"", b""],
+            [b"A", b"G", b"", b""],
+            [b"G", b"A", b"", b""],
+            [b"T", b"A", b"", b""],
+            [b"A", b"G", b"T", b""],
+            [b"T", b"", b"", b""],
+            [b"G", b"GA", b"GAC", b""],
+            [b"T", b"", b"", b""],
+            [b"AC", b"A", b"ATG", b"C"],
+        ],
+    )
+    assert_array_equal(
+        ds["variant/id"],
+        [".", ".", "rs6054257", ".", "rs6040355", ".", "microsat1", ".", "rsTest"],
+    )
+    assert_array_equal(
+        ds["variant/id_mask"],
+        [True, True, False, True, False, True, False, True, False],
+    )
+
+    assert_array_equal(ds["sample/id"], ["NA00001", "NA00002", "NA00003"])
+
+    call_genotype = np.array(
+        [
+            [[0, 0], [0, 0], [0, 1]],
+            [[0, 0], [0, 0], [0, 1]],
+            [[0, 0], [1, 0], [1, 1]],
+            [[0, 0], [0, 1], [0, 0]],
+            [[1, 2], [2, 1], [2, 2]],
+            [[0, 0], [0, 0], [0, 0]],
+            [[0, 1], [0, 2], [-1, -1]],
+            [[0, 0], [0, 0], [-1, -1]],
+            [[0, -1], [0, 1], [0, 2]],
+        ],
+        dtype="i1",
+    )
+    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

sgkit/tests/data/sample.vcf.zarr.zip

@@ -1,7 +1,9 @@
+from pathlib import Path
 from typing import Any, Union
 
 import dask.array as da
 import numpy as np
 
-DType = Any
 ArrayLike = Union[np.ndarray, da.Array]
+DType = Any
+PathType = Union[str, Path]