Restrict bcftools consensus to region-filtered VCFs

tests/test_ad_risk.py

@@ -50,6 +50,7 @@ def test_1(self) -> None:
         model_outputs = self.vcf_processor.predict(self.model, self.checkpoint_path, self.trainer, self.dataloader, self.vcf_dataset)
         gene_tissue_embeds = model_outputs['embeddings'].iloc[0]
         preds = self.adrisk(gene_tissue_embeds)
+        print(f"TestADrisk_pred={preds[0]:.8f}")
         self.assertAlmostEqual(preds[0], 0.66763765, places=1) # tested on h100, dev coreweave on 10/31/25
 
 
@@ -68,8 +69,10 @@ def test_1(self) -> None:
         Test the ADrisk prediction pipeline with an input VCF file
         """
         preds = self.adrisk(self.vcf_path, self.gene_ids, self.tissue_ids)
+        print(f"TestADriskFromVCF_pred0={preds.iloc[0].ad_risk.item():.8f}")
+        print(f"TestADriskFromVCF_pred1={preds.iloc[1].ad_risk.item():.8f}")
         self.assertAlmostEqual(preds.iloc[0].ad_risk.item(), 0.66763765, places=1) # tested on h100, dev coreweave on 10/31/25
 
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()

tests/test_data_process.py

@@ -0,0 +1,86 @@
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+import pandas as pd
+
+CURRENT_PATH = Path(__file__).parent
+sys.path.insert(0, str(CURRENT_PATH.parent))
+
+from utils.data_process import ExtractSeqFromBed
+
+
+class TestExtractSeqFromBed(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tmp_dir = Path(tempfile.mkdtemp(prefix="variantformer-data-process-"))
+        cls.ref_fasta = cls.tmp_dir / "ref.fa"
+        cls.ref_fasta.write_text(">chr1\nACGTACGTACGTACGTACGT\n", encoding="ascii")
+        subprocess.run(["samtools", "faidx", str(cls.ref_fasta)], check=True)
+
+        raw_vcf = cls.tmp_dir / "variants.vcf"
+        raw_vcf.write_text(
+            "\n".join(
+                [
+                    "##fileformat=VCFv4.2",
+                    "##contig=<ID=chr1,length=20>",
+                    '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">',
+                    "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tSAMPLE",
+                    "chr1\t5\t.\tA\tG\t.\tPASS\t.\tGT\t1/1",
+                    "chr1\t12\t.\tT\tC\t.\tPASS\t.\tGT\t1/1",
+                    "",
+                ]
+            ),
+            encoding="ascii",
+        )
+        cls.vcf_path = cls.tmp_dir / "variants.vcf.gz"
+        subprocess.run(
+            [
+                "bcftools",
+                "view",
+                "-Oz",
+                "--write-index",
+                "-o",
+                str(cls.vcf_path),
+                str(raw_vcf),
+            ],
+            check=True,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        shutil.rmtree(cls.tmp_dir, ignore_errors=True)
+
+    def test_apply_bcftools_consensus_respects_region_coordinates(self):
+        extractor = ExtractSeqFromBed(neighbour_hood=0, ref_fasta=str(self.ref_fasta))
+        region = SimpleNamespace(chrom="chr1", start=4, end=5, cCRE="cre-1")
+
+        mutated_seq, mutations = extractor.apply_bcftools_consensus(
+            region, str(self.vcf_path), str(self.ref_fasta)
+        )
+
+        self.assertEqual(mutated_seq, "G")
+        self.assertEqual(mutations, 1)
+
+    def test_process_subject_multiple_regions(self):
+        extractor = ExtractSeqFromBed(neighbour_hood=0, ref_fasta=str(self.ref_fasta))
+        bed_regions = pd.DataFrame(
+            [
+                {"chrom": "chr1", "start": 4, "end": 5, "cCRE": "cre-1"},
+                {"chrom": "chr1", "start": 11, "end": 12, "cCRE": "cre-2"},
+            ]
+        )
+
+        result = extractor.process_subject(str(self.vcf_path), bed_regions)
+
+        self.assertEqual(result["sequence"].tolist(), ["G", "C"])
+        self.assertEqual(result["start_cre"].tolist(), [4, 11])
+        self.assertEqual(result["end_cre"].tolist(), [5, 12])
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/data_process.py

@@ -1,5 +1,6 @@
 import os
 import subprocess
+import tempfile
 from multiprocessing import Pool
 import multiprocessing as mp
 from concurrent.futures import ProcessPoolExecutor
@@ -14,6 +15,54 @@ def __init__(
         self.ref_fasta = ref_fasta
         self.upstream_neighbour_hood = upstream_neighbour_hood
 
+    def _build_consensus_args(self, vcf_file: str, variant_type: str = None):
+        if variant_type == "SNP":
+            return [
+                "bcftools",
+                "consensus",
+                "-H",
+                "I",
+                "-e",
+                'ALT~\"<.*>\" || TYPE!=\"snp\"',
+                vcf_file,
+            ]
+        return [
+            "bcftools",
+            "consensus",
+            "-H",
+            "I",
+            "-e",
+            'ALT~\"<.*>\"',
+            vcf_file,
+        ]
+
+    def _create_filtered_vcf(self, vcf_file: str, region_str: str):
+        with tempfile.NamedTemporaryFile(suffix=".vcf.gz", delete=False) as tmp_vcf:
+            filtered_vcf = tmp_vcf.name
+        cmd_view = [
+            "bcftools",
+            "view",
+            "-r",
+            region_str,
+            "--regions-overlap",
+            "1",
+            "-Oz",
+            "--write-index",
+            "-o",
+            filtered_vcf,
+            vcf_file,
+        ]
+        result = subprocess.run(cmd_view, capture_output=True, text=True)
+        if result.returncode != 0:
+            self._cleanup_filtered_vcf(filtered_vcf)
+            raise ValueError(f"Error running bcftools view: {result.stderr}")
+        return filtered_vcf
+
+    def _cleanup_filtered_vcf(self, filtered_vcf: str):
+        for path in (filtered_vcf, f"{filtered_vcf}.csi", f"{filtered_vcf}.tbi"):
+            if os.path.exists(path):
+                os.remove(path)
+
     def apply_bcftools_consensus(
         self, region, vcf_file, reference_fasta, variant_type: str = None
     ):
@@ -35,40 +84,36 @@ def apply_bcftools_consensus(
             else:
                 mutated_seq = "".join(result_ref.stdout.strip().split("\n")[1:])
                 return mutated_seq, 0
-        # If vcf_file is not None, run bcftools consensus
-        # Command to extract the reference sequence and apply mutations
-        if variant_type == "SNP":
-            bcftools_args = [
-                "bcftools",
-                "consensus",
-                "-H",
-                "I",
-                "-e",
-                'ALT~\"<.*>\" || TYPE!=\"snp\"',
-                vcf_file,
-            ]
-        else:
-            bcftools_args = [
-                "bcftools",
-                "consensus",
-                "-H",
-                "I",
-                "-e",
-                'ALT~\"<.*>\"',
-                vcf_file,
-            ]
+        try:
+            filtered_vcf = self._create_filtered_vcf(vcf_file, region_str)
+        except ValueError as exc:
+            print(region_str)
+            print(f"\n{exc}")
+            print("Falling back to ref genome")
+            result_ref = subprocess.run(cmd_ref, capture_output=True, text=True)
+            if result_ref.returncode != 0:
+                print(f"\nError running samtools faidx: {result_ref.stderr}")
+                return None, 0
+            mutated_seq = "".join(result_ref.stdout.strip().split("\n")[1:])
+            return mutated_seq, 0
 
         # Use piped commands without shell=True
-        samtools_process = subprocess.Popen(
-            cmd_ref, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
-        )
-        result = subprocess.run(
-            bcftools_args, stdin=samtools_process.stdout, capture_output=True, text=True
-        )
-        samtools_process.stdout.close()
-        samtools_stderr = samtools_process.stderr.read()
-        samtools_process.stderr.close()
-        samtools_process.wait()
+        try:
+            samtools_process = subprocess.Popen(
+                cmd_ref, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+            )
+            result = subprocess.run(
+                self._build_consensus_args(filtered_vcf, variant_type=variant_type),
+                stdin=samtools_process.stdout,
+                capture_output=True,
+                text=True,
+            )
+            samtools_process.stdout.close()
+            samtools_stderr = samtools_process.stderr.read()
+            samtools_process.stderr.close()
+            samtools_process.wait()
+        finally:
+            self._cleanup_filtered_vcf(filtered_vcf)
 
         # If bcftools consensus fails, return the reference sequence
         if result.returncode != 0:
@@ -412,39 +457,37 @@ def apply_bcftools_consensus_to_gene(
             else:
                 mutated_seq = "".join(result_ref.stdout.strip().split("\n")[1:])
                 return mutated_seq
-        # If vcf_file is not None, run bcftools consensus
-        if variant_type == "SNP":
-            bcftools_args = [
-                "bcftools",
-                "consensus",
-                "-H",
-                "I",
-                "-e",
-                'ALT~\"<.*>\" || TYPE!=\"snp\"',
-                vcf_file,
-            ]
-        else:
-            bcftools_args = [
-                "bcftools",
-                "consensus",
-                "-H",
-                "I",
-                "-e",
-                'ALT~\"<.*>\"',
-                vcf_file,
-            ]
+        try:
+            filtered_vcf = self._create_filtered_vcf(vcf_file, region_str)
+        except ValueError as exc:
+            print(region_str)
+            print(f"\n{exc}")
+            print("Falling back to reference")
+            result_ref = subprocess.run(cmd_ref, capture_output=True, text=True)
+            if result_ref.returncode != 0:
+                raise ValueError(
+                    f"Error running bcftools consensus: {result_ref.stderr}"
+                )
+            mutated_seq = "".join(result_ref.stdout.strip().split("\n")[1:])
+            return mutated_seq
 
         # Use piped commands without shell=True
-        samtools_process = subprocess.Popen(
-            cmd_ref, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
-        )
-        result = subprocess.run(
-            bcftools_args, stdin=samtools_process.stdout, capture_output=True, text=True
-        )
-        samtools_process.stdout.close()
-        samtools_stderr = samtools_process.stderr.read()
-        samtools_process.stderr.close()
-        samtools_process.wait()
+        try:
+            samtools_process = subprocess.Popen(
+                cmd_ref, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+            )
+            result = subprocess.run(
+                self._build_consensus_args(filtered_vcf, variant_type=variant_type),
+                stdin=samtools_process.stdout,
+                capture_output=True,
+                text=True,
+            )
+            samtools_process.stdout.close()
+            samtools_stderr = samtools_process.stderr.read()
+            samtools_process.stderr.close()
+            samtools_process.wait()
+        finally:
+            self._cleanup_filtered_vcf(filtered_vcf)
 
         if result.returncode != 0:
             print(region_str)