fix: mask sites before dN and dS calculation

Author: ahmig

workflow/envs/bedtools.yaml

@@ -0,0 +1,5 @@
+channels:
+  - conda-forge
+  - bioconda
+dependencies:
+  - bedtools==2.31.1

workflow/rules/evolution.smk

@@ -14,13 +14,27 @@ rule filter_genbank_features:
         "../scripts/filter_genbank_features.py"
 
 
+rule build_problematic_bed:
+    conda:
+        "../envs/bedtools.yaml"
+    input:
+        vcf = lambda wildcards: select_problematic_vcf(),
+    output:
+        bed = temp(OUTDIR / "sites_masked.bed"),
+    log:
+        LOGDIR / "build_problematic_bed" / "log.txt",
+    shell:
+        "bedtools merge -i {input.vcf} >{output.bed} 2>{log}"
+
+
 rule n_s_sites:
     threads: 1
     conda: "../envs/biopython.yaml"
     params:
         gb_qualifier_display = "gene",
     input:
         fasta = OUTDIR/f"{OUTPUT_NAME}.ancestor.fasta",
+        masked = OUTDIR / "sites_masked.bed",
         gb = OUTDIR/"reference.cds.gb",
         genetic_code = Path(config["GENETIC_CODE_JSON"]).resolve(),
     output:
@@ -35,6 +49,7 @@ rule calculate_dnds:
     conda: "../envs/renv.yaml"
     input: 
         n_s_sites = OUTDIR/f"{OUTPUT_NAME}.ancestor.n_s.sites.csv",
+        masked = OUTDIR / "sites_masked.bed",
         variants =  OUTDIR/f"{OUTPUT_NAME}.variants.tsv",
         metadata = config["METADATA"]
     output:

workflow/scripts/calculate_dnds.R

@@ -8,6 +8,7 @@ sink(log, type = "output")
 library(dplyr)
 library(readr)
 library(tidyr)
+library(purrr)
 library(logger)
 
 log_threshold(INFO)
@@ -37,6 +38,21 @@ metadata <- read_delim(snakemake@input[["metadata"]]) %>%
   select(ID, interval) %>%
   rename(SAMPLE = ID)
 
+log_info("Reading masked sites BED")
+masked <- read_tsv(
+  snakemake@input$masked,
+  col_names = c("chrom", "start", "end"),
+  col_types = "cii",
+  comment = "#"
+) %>%
+  mutate(POS = map2(start + 1, end, seq.int)) %>%
+  unnest(POS) %>%
+  pull(POS)
+
+log_info("Filtering variants on masked sites")
+variants <- variants %>%
+  filter(!POS %in% masked, !is.na(SYNONYMOUS))
+
 log_debug("Adding metadata to variants table")
 variants <- left_join(variants, metadata)
 

workflow/scripts/n_s_sites_from_fasta.py

@@ -7,19 +7,50 @@
 
 import pandas as pd
 from Bio import SeqIO
-from Bio.SeqRecord import SeqRecord
+from Bio.SeqRecord import SeqRecord, SeqFeature
 from Bio.Seq import Seq
 
 
 NTS = ("A", "C", "G", "T")
+MASK_CHR = "N"
 
 
 def read_monofasta(path: str) -> SeqRecord:
     return SeqIO.read(path, format="fasta")
 
 
+def read_masked_positions(bed_path: str) -> set:
+    masked = set()
+    with open(bed_path) as fh:
+        for line in fh:
+            if line.startswith("#") or not line.strip():
+                continue
+            parts = line.strip().split("\t")
+            start, end = int(parts[1]), int(parts[2])
+            masked.update(range(start, end))
+    return masked
+
+
+def mask_feature(feature: SeqFeature, record: SeqRecord, masked_positions: set) -> SeqRecord:
+    extracted = feature.extract(record)
+    seq_chars = list(str(extracted.seq).upper())
+    # Handle reverse strand and compound locations
+    genomic_positions = []
+    for part in feature.location.parts:
+        positions = list(range(int(part.start), int(part.end)))
+        if part.strand == -1:
+            positions.reverse()
+        genomic_positions.extend(positions)
+    # Mask sites on sequence
+    for i, pos in enumerate(genomic_positions):
+        if seq_chars[i] not in NTS or pos in masked_positions:
+            seq_chars[i] = MASK_CHR
+    extracted.seq = Seq("".join(seq_chars))
+    return extracted
+
+
 def split_into_codons(seq: Seq) -> list:
-    """Split the complete CDS feature in to a list of codons"""
+    """Split the complete CDS feature in to a list of codons (if ACGT only)"""
     return [
         seq[i:i + 3] for i in range(0, len(seq)-2, 3) if all(char in NTS for char in seq[i:i + 3])
     ]
@@ -79,6 +110,9 @@ def main():
     with open(snakemake.input.genetic_code) as f:
         genetic_code = json.load(f)
 
+    logging.info("Reading masked sites BED")
+    masked = read_masked_positions(snakemake.input.masked)
+
     logging.info("Reading GenBank file")
     gb = SeqIO.read(snakemake.input.gb, format="gb")
 
@@ -95,7 +129,7 @@ def main():
             logging.warning(f"Identifier '{identifier}' is already among coding records and will not be replaced by feature at {feature.location}")
         else:
             logging.debug(f"Adding feature")
-            coding_records[identifier] = feature.extract(record)
+            coding_records[identifier] = mask_feature(feature, record, masked)
 
     logging.info(f"Splitting {len(coding_records)} records into codons")
     codons = get_feature_codons(coding_records)