Update calculation of allele frequency-weighted distances

Replace Python script with PathoGenOmics/afwdist, a faster, lightweight tool

Author: ahmig

workflow/envs/afwdist.yaml

@@ -0,0 +1,5 @@
+channels:
+  - bioconda
+  - conda-forge
+dependencies:
+  - afwdist==1.0.0

workflow/rules/distances.smk

@@ -1,17 +1,52 @@
-rule weighted_distances:
-    threads: 1
+rule extract_afwdist_variants:
     conda: "../envs/biopython.yaml"
     params:
-        samples = expand("{sample}", sample = iter_samples()),
-        mask_class = ["mask"]
+        sample_col = "SAMPLE",
+        position_col = "POS",
+        sequence_col = "ALT",
+        frequency_col = "ALT_FREQ",
+        mask_class = ["mask"],
     input:
-        tsv = OUTDIR/f"{OUTPUT_NAME}.variants.tsv",
-        vcf = lambda wildcards: select_problematic_vcf(),
+        variants = OUTDIR/f"{OUTPUT_NAME}.variants.tsv",
+        mask_vcf = lambda wildcards: select_problematic_vcf(),
         ancestor = OUTDIR/f"{OUTPUT_NAME}.ancestor.fasta",
-        reference = OUTDIR/"reference.fasta"
+        reference = OUTDIR/"reference.fasta",
     output:
-        distances = REPORT_DIR_TABLES/f"distances.csv"
+        variants = temp(OUTDIR/f"{OUTPUT_NAME}.variants.afwdist.csv"),
     log:
-        LOGDIR / "weighted_distances" / "log.txt"
+        LOGDIR/"extract_afwdist_variants"/"log.txt"
     script:
-        "../scripts/weighted_distances.py"
+        "../scripts/extract_afwdist_variants.py"
+
+
+rule afwdist_weighted_distances:
+    conda: "../envs/afwdist.yaml"
+    params:
+        extra_args = ""
+    input:
+        variants = OUTDIR/f"{OUTPUT_NAME}.variants.afwdist.csv",
+        reference = OUTDIR/f"{OUTPUT_NAME}.ancestor.fasta",
+    output:
+        distances = temp(REPORT_DIR_TABLES/"distances.raw.csv"),
+    log:
+        LOGDIR/"afwdist_weighted_distances"/"log.txt"
+    shell:
+        "afwdist "
+        "-i {input.variants:q} "
+        "-r {input.reference:q} "
+        "-o {output.distances:q} "
+        "{params.extra_args} >{log:q} 2>&1"
+
+
+rule format_afwdist_results:
+    conda: "../envs/biopython.yaml"
+    params:
+        samples = sorted(iter_samples()),
+    input:
+        distances = REPORT_DIR_TABLES/"distances.raw.csv",
+    output:
+        distances = REPORT_DIR_TABLES/"distances.csv",
+    log:
+        LOGDIR/"format_afwdist_results"/"log.txt"
+    script:
+        "../scripts/format_afwdist_results.py"

workflow/scripts/extract_afwdist_variants.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+
+import logging
+from typing import List
+
+import pandas as pd
+from Bio import SeqIO
+from Bio.SeqRecord import SeqRecord
+from Bio.Seq import Seq
+
+
+def read_monofasta(path: str) -> SeqRecord:
+    fasta = SeqIO.parse(path, "fasta")
+    record = next(fasta)
+    if next(fasta, None) is not None:
+        logging.warning(f"There are unread records left in '{path}'")
+    return record
+
+
+def read_masked_sites(vcf_path: str, mask_classes: List[str]) -> List[int]:
+    """
+    Parse a VCF containing positions for masking. Assumes the VCF file is
+    formatted as in:
+    github.com/W-L/ProblematicSites_SARS-CoV2/blob/master/problematic_sites_sarsCov2.vcf
+    with a "mask" or "caution" recommendation in column 7.
+    Masked sites are specified with params.
+    """
+    vcf = pd.read_csv(
+        vcf_path,
+        sep="\s+",
+        comment="#",
+        names=("CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO")
+    )
+    return vcf.loc[vcf.FILTER.isin(mask_classes), "POS"].tolist()
+
+
+def build_ancestor_variant_table(ancestor: Seq, reference: Seq, reference_name: str, masked_positions: List[int]) -> pd.DataFrame:
+    pos = []
+    alt = []
+    for i in range(1, len(ancestor) + 1):
+        if i not in masked_positions and ancestor[i-1] != reference[i-1]:
+            pos.append(i)
+            alt.append(reference[i-1])
+    df = pd.DataFrame({snakemake.params.position_col: pos, snakemake.params.sequence_col: alt})
+    df[snakemake.params.frequency_col] = 1  # As a reference genome, we assume all positions have fixed alleles
+    df[snakemake.params.sample_col] = reference_name
+    return df
+
+
+if __name__ == "__main__":
+
+    logging.basicConfig(filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"], level=logging.INFO)
+
+    colnames = {
+        snakemake.params.sample_col: "sample",
+        snakemake.params.position_col: "position",
+        snakemake.params.sequence_col: "sequence",
+        snakemake.params.frequency_col: "frequency"
+    }
+
+    logging.info("Reading input tables")
+    # Variants
+    variants = pd.read_table(snakemake.input.variants, sep="\t")
+    logging.info(f"Read {len(variants)} variant records")
+    # VCF with sites to mask
+    masked_sites = read_masked_sites(snakemake.input.mask_vcf, snakemake.params.mask_class)
+    logging.info(f"Read {len(masked_sites)} masked positions")
+
+    logging.info("Reading input FASTA files")
+    # Case ancestor
+    ancestor = read_monofasta(snakemake.input.ancestor)
+    logging.info(f"Ancestor: '{ancestor.description}', length={len(ancestor.seq)}")
+    # Alignment reference
+    reference = read_monofasta(snakemake.input.reference)
+    logging.info(f"Reference: '{reference.description}', length={len(reference.seq)}")
+
+    logging.info("Processing ancestor variants")
+    ancestor_table = build_ancestor_variant_table(ancestor.seq, reference.seq, reference.id, masked_sites)
+    logging.info(f"Ancestor has {len(ancestor_table)} variants")
+    all_variants = pd.concat([variants, ancestor_table], ignore_index=True)
+    logging.info(f"Combined table has {len(all_variants)} variants")
+
+    logging.info("Renaming and selecting columns")
+    output = all_variants.rename(columns=colnames)[list(colnames.values())]
+    logging.info("Filtering sites")
+    output = output[~output.position.isin(masked_sites)]
+    logging.info(f"There are {len(output)} rows left")
+    
+    logging.info("Writing results")
+    output.to_csv(snakemake.output.variants, index=False)

workflow/scripts/format_afwdist_results.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+
+import logging
+import pandas as pd
+
+
+if __name__ == "__main__":
+
+    logging.basicConfig(filename=snakemake.log[0], format=snakemake.config["LOG_PY_FMT"], level=logging.INFO)
+
+    logging.info("Read pairwise distances")
+    df = pd.read_csv(snakemake.input.distances)
+
+    logging.info("Initializing formatted output")
+    output = pd.DataFrame(
+        columns=snakemake.params.samples,
+        index=snakemake.params.samples,
+        dtype="float64"
+    )
+
+    logging.info("Filling table")
+    for i, row in df.iterrows():
+        output.loc[row.sample_m, row.sample_n] = row.distance
+        output.loc[row.sample_n, row.sample_m] = row.distance
+    
+    logging.info("Writing formatted results")
+    output.to_csv(snakemake.output.distances)