feat: add fields parameter for VCF FORMAT output

Author: adilraza99

malariagen_data/anoph/to_vcf.py

@@ -11,6 +11,15 @@
 from . import plink_params
 from . import vcf_params
 
+# Supported FORMAT fields and their VCF header definitions.
+_VALID_FIELDS = {"GT", "GQ", "AD", "MQ"}
+_FORMAT_HEADERS = {
+    "GT": '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
+    "GQ": '##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">',
+    "AD": '##FORMAT=<ID=AD,Number=R,Type=Integer,Description="Allele Depth">',
+    "MQ": '##FORMAT=<ID=MQ,Number=1,Type=Integer,Description="Mapping Quality">',
+}
+
 
 class VcfExporter(
     AnophelesSnpData,
@@ -50,11 +59,23 @@ def snp_calls_to_vcf(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
         overwrite: plink_params.overwrite = False,
+        fields: vcf_params.vcf_fields = ("GT",),
     ) -> str:
         base_params._validate_sample_selection_params(
             sample_query=sample_query, sample_indices=sample_indices
         )
 
+        # Validate fields parameter.
+        fields = tuple(fields)
+        unknown = set(fields) - _VALID_FIELDS
+        if unknown:
+            raise ValueError(
+                f"Unknown FORMAT fields: {unknown}. "
+                f"Valid fields are: {sorted(_VALID_FIELDS)}"
+            )
+        if "GT" not in fields:
+            raise ValueError("GT must be included in fields.")
+
         if os.path.exists(output_path) and not overwrite:
             return output_path
 
@@ -74,14 +95,21 @@ def snp_calls_to_vcf(
         compress = output_path.endswith(".gz")
         opener = gzip.open if compress else open
 
+        # Determine which extra fields to include.
+        include_gq = "GQ" in fields
+        include_ad = "AD" in fields
+        include_mq = "MQ" in fields
+        format_str = ":".join(fields)
+
         with opener(output_path, "wt") as f:
             # Write VCF header.
             f.write("##fileformat=VCFv4.3\n")
             f.write(f"##fileDate={date.today().strftime('%Y%m%d')}\n")
             f.write("##source=malariagen_data\n")
             for contig in contigs:
                 f.write(f"##contig=<ID={contig}>\n")
-            f.write('##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\n')
+            for field in fields:
+                f.write(_FORMAT_HEADERS[field] + "\n")
             header_cols = [
                 "#CHROM",
                 "POS",
@@ -95,15 +123,36 @@ def snp_calls_to_vcf(
             ]
             f.write("\t".join(header_cols + list(sample_ids)) + "\n")
 
-            # Write records in chunks.
+            # Extract dask arrays.
             gt_data = ds["call_genotype"].data
             pos_data = ds["variant_position"].data
             contig_data = ds["variant_contig"].data
             allele_data = ds["variant_allele"].data
 
+            # Optional field arrays — may not exist in all datasets.
+            gq_data = None
+            ad_data = None
+            mq_data = None
+            if include_gq:
+                try:
+                    gq_data = ds["call_GQ"].data
+                except KeyError:
+                    pass
+            if include_ad:
+                try:
+                    ad_data = ds["call_AD"].data
+                except KeyError:
+                    pass
+            if include_mq:
+                try:
+                    mq_data = ds["call_MQ"].data
+                except KeyError:
+                    pass
+
             chunk_sizes = gt_data.chunks[0]
             offsets = np.cumsum((0,) + chunk_sizes)
 
+            # Write records in chunks.
             with self._spinner(f"Write VCF ({ds.sizes['variants']} variants)"):
                 for ci in range(len(chunk_sizes)):
                     start = offsets[ci]
@@ -113,6 +162,26 @@ def snp_calls_to_vcf(
                     contig_chunk = contig_data[start:stop].compute()
                     allele_chunk = allele_data[start:stop].compute()
 
+                    # Compute optional field chunks, handling missing data.
+                    gq_chunk = None
+                    ad_chunk = None
+                    mq_chunk = None
+                    if gq_data is not None:
+                        try:
+                            gq_chunk = gq_data[start:stop].compute()
+                        except (FileNotFoundError, KeyError):
+                            pass
+                    if ad_data is not None:
+                        try:
+                            ad_chunk = ad_data[start:stop].compute()
+                        except (FileNotFoundError, KeyError):
+                            pass
+                    if mq_data is not None:
+                        try:
+                            mq_chunk = mq_data[start:stop].compute()
+                        except (FileNotFoundError, KeyError):
+                            pass
+
                     for j in range(gt_chunk.shape[0]):
                         chrom = contigs[contig_chunk[j]]
                         pos = str(pos_chunk[j])
@@ -130,16 +199,47 @@ def snp_calls_to_vcf(
                         alt = ",".join(alt_alleles) if alt_alleles else "."
 
                         gt_row = gt_chunk[j]
-                        sample_fields = np.empty(gt_row.shape[0], dtype=object)
-                        for k in range(gt_row.shape[0]):
+                        n_samples = gt_row.shape[0]
+                        sample_fields = np.empty(n_samples, dtype=object)
+                        for k in range(n_samples):
+                            parts = []
+                            # GT (always present).
                             a0 = gt_row[k, 0]
                             a1 = gt_row[k, 1]
                             if a0 < 0 or a1 < 0:
-                                sample_fields[k] = "./."
+                                parts.append("./.")
                             else:
-                                sample_fields[k] = f"{a0}/{a1}"
-
-                        line = f"{chrom}\t{pos}\t.\t{ref}\t{alt}\t.\t.\t.\tGT\t"
+                                parts.append(f"{a0}/{a1}")
+                            # GQ.
+                            if include_gq:
+                                if gq_chunk is not None:
+                                    v = gq_chunk[j, k]
+                                    parts.append("." if v < 0 else str(v))
+                                else:
+                                    parts.append(".")
+                            # AD.
+                            if include_ad:
+                                if ad_chunk is not None:
+                                    ad_vals = ad_chunk[j, k]
+                                    parts.append(
+                                        ",".join(
+                                            "." if x < 0 else str(x) for x in ad_vals
+                                        )
+                                    )
+                                else:
+                                    parts.append(".")
+                            # MQ.
+                            if include_mq:
+                                if mq_chunk is not None:
+                                    v = mq_chunk[j, k]
+                                    parts.append("." if v < 0 else str(v))
+                                else:
+                                    parts.append(".")
+                            sample_fields[k] = ":".join(parts)
+
+                        line = (
+                            f"{chrom}\t{pos}\t.\t{ref}\t{alt}\t.\t.\t.\t{format_str}\t"
+                        )
                         line += "\t".join(sample_fields)
                         f.write(line + "\n")
 

malariagen_data/anoph/vcf_params.py

@@ -1,5 +1,7 @@
 """Parameters for VCF exporter functions."""
 
+from typing import Tuple
+
 from typing_extensions import Annotated, TypeAlias
 
 vcf_output_path: TypeAlias = Annotated[
@@ -9,3 +11,11 @@
     gzip compression.
     """,
 ]
+
+vcf_fields: TypeAlias = Annotated[
+    Tuple[str, ...],
+    """
+    FORMAT fields to include in the VCF output. Must include "GT".
+    Supported fields: "GT", "GQ", "AD", "MQ".
+    """,
+]

tests/anoph/test_vcf_exporter.py

@@ -148,3 +148,49 @@ def test_vcf_exporter_gzip(fixture, api: VcfExporter, tmp_path):
     with gzip.open(output_path, "rt") as f:
         first_line = f.readline()
     assert first_line.strip() == "##fileformat=VCFv4.3"
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_vcf_exporter_fields(fixture, api: VcfExporter, tmp_path):
+    region = api.contigs[0]
+
+    # Test with additional FORMAT fields.
+    output_path = str(tmp_path / "test_fields.vcf")
+    api.snp_calls_to_vcf(
+        output_path=output_path,
+        region=region,
+        fields=("GT", "GQ"),
+    )
+
+    with open(output_path) as f:
+        lines = f.readlines()
+
+    # Check FORMAT header lines.
+    format_lines = [line for line in lines if line.startswith("##FORMAT")]
+    assert len(format_lines) == 2
+    assert any("ID=GT" in line for line in format_lines)
+    assert any("ID=GQ" in line for line in format_lines)
+
+    # Check FORMAT column value.
+    data_lines = [line for line in lines if not line.startswith("#")]
+    assert len(data_lines) > 0
+    first_data = data_lines[0].strip().split("\t")
+    assert first_data[8] == "GT:GQ"
+
+    # Each sample field should have two colon-separated values.
+    for sample_val in first_data[9:]:
+        parts = sample_val.split(":")
+        assert len(parts) == 2, f"Expected GT:GQ, got {sample_val!r}"
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_vcf_exporter_fields_gt_required(fixture, api: VcfExporter, tmp_path):
+    region = api.contigs[0]
+    output_path = str(tmp_path / "test_no_gt.vcf")
+
+    with pytest.raises(ValueError, match="GT must be included"):
+        api.snp_calls_to_vcf(
+            output_path=output_path,
+            region=region,
+            fields=("GQ",),
+        )