shuaizhou / parser_vcf.py

import os import re import gzip from typing import Tuple, List import pandas as pd class QC: _DP_RE = re.compile(r"(?:^|;)\s*DP=(\d+(?:\.\d+)?)") def __init__( self, uid: str, library_num: int, output_qc: str, input_vcf: str, basic_reads: str, bamstats: str, depth: int ): self.uid = uid self.library_num = library_num self.output_qc = output_qc self.input_vcf = input_vcf self.basic_reads = basic_reads self.bamstats = bamstats self.depth = depth def _gt_to_bases(self, gt: str, ref: str, alt: str) -> str | None: if gt is None: return None gt = str(gt) if gt in {".", "./.", ".|."}: return None alts = [] if alt in {None, ".", ""} else str(alt).split(",") alleles = [ref] + alts # 0->REF, 1..->ALT if "/" in gt: parts = gt.split("/") elif "|" in gt: parts = gt.split("|") else: parts = [gt] # haploid like "1" bases = [] for p in parts: if p in {".", ""}: bases.append(".") continue try: idx = int(p) except ValueError: return None if 0 <= idx < len(alleles): bases.append(alleles[idx]) else: return None return "/".join(bases) @classmethod def dp_difference(cls, info: str, dp: int) -> float: """ INFO列里的DP 与 sample列里的DP 差值（INFO - sample） """ if info is None: return float(0 - dp) m = cls._DP_RE.search(str(info)) if m: return float(m.group(1)) - dp return float(0 - dp) def parse_vcf_row(self, row) -> pd.Series: """ row: tuple/list-like, 对应VCF的 0..9 列: 0 CHROM,1 POS,2 ID,3 REF,4 ALT,5 QUAL,6 FILTER,7 INFO,8 FORMAT,9 SAMPLE """ chr_id = row[0] position = row[1] ref_base = row[3] alt_base = row[4] fmt_keys = str(row[8]).split(":") sample_vals = str(row[9]).split(":") fmt = dict(zip(fmt_keys, sample_vals)) data = {"A": 0, "T": 0, "G": 0, "C": 0} gt = fmt.get("GT", None) gt_base = self._gt_to_bases(gt, ref_base, alt_base) # AD ad = fmt.get("AD", None) ad_sum = 0 if ad not in {None, ".", ""}: try: ad_int = [int(i) for i in str(ad).split(",")] ad_sum = sum(ad_int) bases = [ref_base] if alt_base in {None, ".", ""} else [ref_base] + str(alt_base).split(",") for b, n in zip(bases, ad_int): if b in data: # only A/T/G/C data[b] = n except ValueError: ad = None ad_sum = 0 else: ad = None # DP dp = 0 v = fmt.get("DP", None) if v not in {None, ".", ""}: try: dp = int(v) except ValueError: dp = 0 dp_diff = self.dp_difference(row[7], dp) # GQ / RGQ gq = None for k in ("GQ", "RGQ"): v = fmt.get(k, None) if v not in {None, ".", ""}: try: gq = float(v) except ValueError: gq = None break # PL pls = fmt.get("PL", None) if pls in {".", ""}: pls = None if ad_sum > 0: afA = data["A"] / ad_sum afT = data["T"] / ad_sum afG = data["G"] / ad_sum afC = data["C"] / ad_sum else: afA = afT = afG = afC = 0.0 return pd.Series( [ chr_id, position, data["A"], data["T"], data["G"], data["C"], afA, afT, afG, afC, gt, gt_base, ad, ad_sum, dp, dp_diff, gq, pls ], index=[ "chr", "pos", "A", "T", "G", "C", "A_frac", "T_frac", "G_frac", "C_frac", "GT", "GT_base", "AD", "AD_sum", "DP", "DP_diff", "GQ", "PL" ] ) def read_vcf(self, input_vcf: str) -> Tuple[pd.DataFrame, List[str], List[str]]: """ 返回: (vcf_df, vcf_header(##), header_line(#CHROM...)) """ is_gzipped = input_vcf.endswith(".gz") open_func = gzip.open if is_gzipped else open mode = "rt" if is_gzipped else "r" vcf_header: List[str] = [] header_line: List[str] = [] # 逐行读header，避免 readlines() 占内存 with open_func(input_vcf, mode) as f: for line in f: if line.startswith("##"): vcf_header.append(line.rstrip("\n")) continue if line.startswith("#CHROM"): header_line = line.strip().split() break # 主体用 pandas 读（comment='#' 会跳过所有header行） vcf_df = pd.read_csv( input_vcf, comment="#", sep="\t", header=None, compression="infer", dtype={0: str, 1: int, 3: str, 4: str, 7: str, 8: str, 9: str}, low_memory=False, ) return vcf_df, vcf_header, header_line def parse_vcf_df(self, vcf_df: pd.DataFrame) -> pd.DataFrame: """ 只解析，不过滤。比 iterrows() 更快。 """ out = [] # itertuples 更快；name=None 返回普通tuple for row in vcf_df.itertuples(index=False, name=None): out.append(self.parse_vcf_row(row)) return pd.DataFrame(out) def main(self) -> pd.DataFrame: # logger.info(f"starting: {self.input_vcf}") vcf_df, vcf_header, header_line = self.read_vcf(self.input_vcf) parsed = self.parse_vcf_df(vcf_df) # logger.info(f"done: {self.input_vcf}") return parsed