Move features in JSON file to config YAML

ahmig · ahmig · commit 9b47abe7ad94 · 2025-09-11T16:41:00.000+02:00
diff --git a/config/config.yaml b/config/config.yaml
@@ -2,8 +2,6 @@ ALIGNMENT_REFERENCE:
   "NC_045512.2"
 PROBLEMATIC_VCF:
   "https://raw.githubusercontent.com/W-L/ProblematicSites_SARS-CoV2/da322c32004f7b16bdaa6a8ee7fd24d56e79d9dc/problematic_sites_sarsCov2.vcf"
-FEATURES_JSON:
-  "config/sarscov2_features.json"
 COORDINATES_JSON:
   "config/sarscov2_coordinates.json"
 GENETIC_CODE_JSON:
@@ -65,6 +63,19 @@ DEMIX:
 WINDOW:
   WIDTH: 1000
   STEP: 50
+GB_FEATURES:
+  "ORF1ab polyprotein": "ORF1ab"
+  "ORF1a polyprotein": "ORF1ab"
+  "surface glycoprotein": "S"
+  "ORF3a protein": "ORF3a"
+  "envelope protein": "E"
+  "membrane glycoprotein": "M"
+  "ORF6 protein": "ORF6"
+  "ORF7a protein": "ORF7"
+  "ORF7b": "ORF7"
+  "ORF8 protein": "ORF8"
+  "nucleocapsid phosphoprotein": "N"
+  "ORF10 protein": "ORF10"
 GISAID:
   CREDENTIALS: "config/gisaid.yaml"
   DATE_COLUMN: "CollectionDate"
diff --git a/config/sarscov2_features.json b/config/sarscov2_features.json
diff --git a/workflow/rules/evolution.smk b/workflow/rules/evolution.smk
@@ -1,13 +1,14 @@
 rule N_S_sites:
     threads: 1
     conda: "../envs/biopython.yaml"
+    params:
+        gb_features = config["GB_FEATURES"],
     input:
         fasta = OUTDIR/f"{OUTPUT_NAME}.ancestor.fasta",
         gb = OUTDIR/"reference.gb",
-        features = Path(config["FEATURES_JSON"]).resolve(),
-        genetic_code = Path(config["GENETIC_CODE_JSON"]).resolve()
+        genetic_code = Path(config["GENETIC_CODE_JSON"]).resolve(),
     output:
-        csv = temp(OUTDIR/f"{OUTPUT_NAME}.ancestor.N_S.sites.csv")
+        csv = temp(OUTDIR/f"{OUTPUT_NAME}.ancestor.N_S.sites.csv"),
     log:
         LOGDIR / "N_S_sites" / "log.txt"
     script:
diff --git a/workflow/rules/report.smk b/workflow/rules/report.smk
@@ -44,11 +44,11 @@ rule window:
     conda: "../envs/biopython.yaml"
     params:
         window = config["WINDOW"]["WIDTH"],
-        step = config["WINDOW"]["STEP"]
+        step = config["WINDOW"]["STEP"],
+        select_gb_features = config.get("GB_FEATURES", {}),  # if empty, uses all available features
     input:
         vcf = OUTDIR/f"{OUTPUT_NAME}.variants.tsv",
         gb = OUTDIR/"reference.gb",
-        features = config["FEATURES_JSON"]
     output:
         window_df = temp(OUTDIR/f"{OUTPUT_NAME}.window.csv"),
     log:
@@ -229,7 +229,7 @@ rule evo_plots:
     params: 
         design = config["PLOTS"]
     input: 
-        N_S = OUTDIR/f"{OUTPUT_NAME}.ancestor.N_S.sites.csv",
+        n_s_sites = OUTDIR/f"{OUTPUT_NAME}.ancestor.N_S.sites.csv",
         vcf =  OUTDIR/f"{OUTPUT_NAME}.variants.tsv",
         metadata = config["METADATA"]
     output:
diff --git a/workflow/scripts/N_S_sites_from_fasta.py b/workflow/scripts/N_S_sites_from_fasta.py
@@ -11,7 +11,9 @@
 
 def split_into_codons(seq: str) -> list:
     """Split the complete CDS feature in to a list of codons"""
-    codons = [seq[i:i + 3] for i in range(0, len(seq), 3) if "N" not in seq[i:i + 3]]
+    codons = [
+        seq[i:i + 3] for i in range(0, len(seq), 3) if "N" not in seq[i:i + 3]
+    ]
     return codons
 
 
@@ -29,7 +31,7 @@ def calculate_potential_changes(genetic_code: dict) -> dict:
         for codon_p in range(0, 3):
             nts = ["A", "G", "T", "C"]
             # Do not consider self substitutions, e.g. A->A
-            nts.remove(codon[codon_p]) 
+            nts.remove(codon[codon_p])
             for nt in nts:
                 codon_mutated = list(codon)
                 # Mutate the basepair
@@ -43,49 +45,55 @@ def calculate_potential_changes(genetic_code: dict) -> dict:
 
 
 def get_feature_codons(alignment: Gb2Alignment, annotation: list) -> dict:
-    dct = {key:alignment.ntSequences(key)[1].sequence for key in annotation}
-    return {key:split_into_codons(item) for key,item in dct.items()}
+    dct = {key: alignment.ntSequences(key)[1].sequence for key in annotation}
+    return {key: split_into_codons(item) for key, item in dct.items()}
 
 
-def get_df(codons: dict, genetic_code: dict) -> pd.DataFrame:
-    keys = []
+def calculate_ns_sites(codons: dict, genetic_code: dict) -> pd.DataFrame:
+    features = []
     N_sites = []
     S_sites = []
     values = calculate_potential_changes(genetic_code)
     for key, item in codons.items():
-        keys.append(key)
+        features.append(key)
         N = sum([values["N"][x] for x in item if x in values["N"].keys()])
         S = sum([values["S"][x] for x in item if x in values["S"].keys()])
         N_sites.append(N)
-        S_sites.append(S) 
-    return pd.DataFrame({"gene": keys, "N": N_sites, "S": S_sites})
+        S_sites.append(S)
+    return pd.DataFrame({"gene": features, "N": N_sites, "S": S_sites})
 
 
 def main():
 
-    logging.basicConfig(filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"], level=logging.INFO)
-    
+    logging.basicConfig(
+        filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"],
+        level=logging.INFO
+    )
+
     logging.info("Reading features")
-    with open(snakemake.input.features) as f:
-        feature_list = list(json.load(f).keys())
+    feature_list = list(snakemake.params.gb_features.keys())
 
     logging.info("Reading genetic code")
     with open(snakemake.input.genetic_code) as f:
         genetic_code = json.load(f)
-    
+
     logging.info("Create alignment object")
     features = Features(snakemake.input.gb)
-    seq = list(FastaReads(snakemake.input.fasta))[0]
+    fasta_reads = list(FastaReads(snakemake.input.fasta))
+    if len(fasta_reads) > 1:
+        logging.warning(
+            f"More than one record found in {snakemake.input.fasta}, selecting the first one")
+    seq = fasta_reads[0]
     aln = Gb2Alignment(seq, features)
 
-    logging.info("Splitting ancestral sequence into codons")
-    codons_dict = get_feature_codons(aln, feature_list)
+    logging.info("Splitting input sequence into codons")
+    codons = get_feature_codons(aln, feature_list)
 
     logging.info("Calculating synonymous and non synonymous sites")
-    df = get_df(codons_dict, genetic_code)
+    df = calculate_ns_sites(codons, genetic_code)
 
     logging.info("Saving results")
-    df.to_csv(snakemake.output.csv,index= False)
+    df.to_csv(snakemake.output.csv, index=False)
 
 
 if __name__ == "__main__":
diff --git a/workflow/scripts/report/NV_description.R b/workflow/scripts/report/NV_description.R
@@ -198,8 +198,8 @@ window_plot <- window %>%
   ggplot() +
   aes(
     x = position,
-    y = fractions,
-    color = gen
+    y = fraction,
+    color = feature
   ) +
   geom_point() +
   geom_line(
@@ -209,7 +209,7 @@ window_plot <- window %>%
   ) +
   scale_y_continuous(
     label = scales::percent,
-    limits = c(0, max(window$fractions) + 0.005)
+    limits = c(0, max(window$fraction) + 0.005)
   ) +
   xlim(c(0, 29903)) +
   scale_color_manual(values = GENE_PALETTE) +
@@ -235,7 +235,7 @@ window_plot_nsp <- window_plot +
     data = npc,
     aes(
       x = (summary_start + summary_end) / 2,
-      y = max(window$fractions) + 0.002,
+      y = max(window$fraction) + 0.002,
       label = summary_nsp
     ),
     inherit.aes = FALSE,
@@ -266,9 +266,8 @@ ggsave(
 
 # Zoom in in spike
 log_info("Plotting summary for variants in the spike")
-
 spike_pos <- window %>%
-  filter(gen == "S") %>%
+  filter(feature == "S") %>%
   pull(position)
 
 vcf_spike <- vcf %>%
@@ -278,12 +277,12 @@ vcf_spike <- vcf %>%
   )
 
 window_plot_spike <- window %>%
-  filter(gen == "S") %>%
+  filter(feature == "S") %>%
   ggplot() +
   aes(
     x = position,
-    y = fractions,
-    color = gen
+    y = fraction,
+    color = feature
   ) +
   geom_point() +
   geom_line(
@@ -293,7 +292,7 @@ window_plot_spike <- window %>%
   ) +
   scale_y_continuous(
     label = scales::percent,
-    limits = c(0, max(window$fractions) + 0.005)
+    limits = c(0, max(window$fraction) + 0.005)
   ) +
   xlim(c(min(spike_pos), max(spike_pos))) +
   scale_color_manual(
@@ -438,8 +437,8 @@ vcf %>%
 window %>%
   transmute(
     POS = position,
-    feature = gen,
-    prop_PolymorphicSites = fractions
+    feature = feature,
+    prop_PolymorphicSites = fraction
   ) %>%
   write.csv(snakemake@output[["table_1"]], row.names = FALSE)
 
diff --git a/workflow/scripts/report/evo_plots.R b/workflow/scripts/report/evo_plots.R
@@ -15,7 +15,7 @@ source(snakemake@params[["design"]])
 # Read inputs
 vcf <- read_delim(snakemake@input[["vcf"]])
 metadata <- read_delim(snakemake@input[["metadata"]])
-N_S_position <- read_delim(snakemake@input[["N_S"]])
+N_S_position <- read_delim(snakemake@input[["n_s_sites"]])
 
 # DATA PROCESSING
 # Create SNP variable and select useful variables
diff --git a/workflow/scripts/window.py b/workflow/scripts/window.py
@@ -6,46 +6,57 @@
 from gb2seq.features import Features
 
 
-def window_calculation(sites: set, window: int, step: int, coord: str) -> pd.DataFrame:
-    ft = Features(coord) # Create Features object to obtain annotations
-    positions = []
-    pct = []
-    genes = []
-    lim_sup = len(ft.reference.sequence) + 1
+def window_calculation(sites: set, window: int, step: int, gb_features: Features) -> pd.DataFrame:
+    positions, fractions, features = [], [], []
+    lim_sup = len(gb_features.reference.sequence) + 1
     for position in range(1, lim_sup, step):
-        if len(ft.getFeatureNames(position)) == 0:
-            genes.append("Intergenic")
+        if len(gb_features.getFeatureNames(position)) == 0:
+            features.append("Intergenic")
         else:
-            genes.append(list(ft.getFeatureNames(position))[0])
+            # If more than one feature on a site, include both (sorted lexicographically)
+            features.append("|".join(sorted(gb_features.getFeatureNames(position))))
         # Add percent (excluding initial and final positions)
         if position - window not in range(1, lim_sup):
-            pct.append(0)
+            fractions.append(0.0)
         else:
             # Calculate no. of polimorphisms in the window
-            num_snp = len([x for x in sites if x in range(position - window, position + 1)])
-            pct.append(num_snp / window)
+            num_snp = len([x for x in sites if x in range(
+                position - window, position + 1)])
+            fractions.append(num_snp / window)
         positions.append(position)
-    return pd.DataFrame({"position": positions, "fractions": pct, "gen": genes})
+    return pd.DataFrame({"position": positions, "fraction": fractions, "feature": features})
 
 
 def main():
 
-    logging.basicConfig(filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"], level=logging.INFO)
+    logging.basicConfig(
+        filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"],
+        level=logging.INFO
+    )
 
     logging.info("Reading input VCF")
     df = pd.read_table(snakemake.input.vcf)
     sites = set(df.POS)
 
-    logging.info("Reading features")
-    with open(snakemake.input.features) as f:
-        features_key = json.load(f)
+    logging.info("Reading genbank features")
+    features = Features(snakemake.input.gb)
+
+    logging.info("Calculating polimorphic sites sliding window")
+    windows = window_calculation(
+        sites,
+        snakemake.params.window,
+        snakemake.params.step,
+        features
+    )
+
+    if len(snakemake.params.select_gb_features) != 0:
+        logging.info("Filtering and renaming genbank features")
+        windows = windows[
+            windows.feature.isin(snakemake.params.select_gb_features.keys())
+        ].replace(snakemake.params.select_gb_features)
 
-    logging.info("Sliding window calculation of proportion of polimorphic sites")
-    frame = window_calculation(sites, snakemake.params.window, snakemake.params.step, snakemake.input.gb)
-    
     logging.info("Saving results")
-    frame.replace(features_key, inplace = True)
-    frame.to_csv(snakemake.output.window_df, index= False)
+    windows.to_csv(snakemake.output.window_df, index=False)
 
 
 if __name__ == "__main__":
