diff --git a/tests/test_ad_risk.py b/tests/test_ad_risk.py index e6cafab..016b752 100755 --- a/tests/test_ad_risk.py +++ b/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() \ No newline at end of file + unittest.main() diff --git a/tests/test_data_process.py b/tests/test_data_process.py new file mode 100644 index 0000000..bc9d7c4 --- /dev/null +++ b/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=", + '##FORMAT=', + "#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() diff --git a/utils/data_process.py b/utils/data_process.py index 082536c..cb168d5 100644 --- a/utils/data_process.py +++ b/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)