episodic 0.0.1__py3-none-any.whl

episodic/__about__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2023-present Wytamma Wirth <wytamma.wirth@me.com>
+#
+# SPDX-License-Identifier: MIT
+__version__ = "0.0.1"

episodic/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023-present Wytamma Wirth <wytamma.wirth@me.com>
+#
+# SPDX-License-Identifier: MIT

episodic/__main__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: 2023-present Wytamma Wirth <wytamma.wirth@me.com>
+#
+# SPDX-License-Identifier: MIT
+import sys
+
+if __name__ == "__main__":
+    from episodic.cli import episodic
+
+    sys.exit(episodic())

episodic/cli/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: 2023-present Wytamma Wirth <wytamma.wirth@me.com>
+#
+# SPDX-License-Identifier: MIT
+from pathlib import Path
+
+from snk.cli import CLI
+
+episodic = CLI(pipeline_dir_path = Path(__file__).parent.parent)

episodic/snk.yaml

@@ -0,0 +1,125 @@
+logo: Episodic
+tagline: A pipeline for fitting and testing Fixed Local Clock (FLC) molecular clock models for episodic evolution.
+annotations:
+  alignment:
+    type: Path
+    help: "Path to a fasta file containing the alignment."
+    required: true
+  group:
+    type: List[str]
+    help: "Group sequences containing this value (in the header) to define a FLC. Can specify multiple."
+    required: true
+  date_delimiter:
+    type: str
+    help: "Delimiter to use to split the date from the rest of the header."
+    required: false
+    default: '"|"'
+  date_index:
+    type: int
+    help: "Index of the date in the header. 0-based."
+    required: false
+    default: -1
+  newick:
+    type: Path
+    help: "Path to a newick tree file. If provided topology will be fixed."
+    required: false
+  clock:
+    type: List[str]
+    help: "Clock models to test. Options are 'strict', 'relaxed', and 'flc-stem'. Can specify multiple."
+    required: false
+    default: ['strict', 'relaxed', 'flc-stem']
+  chain_length:
+    type: int
+    help: "Length of the chain."
+    required: false
+    default: 10000000
+  samples:
+    type: int
+    help: "Number of samples to draw from the chain."
+    required: false
+    default: 10000
+  rate_gamma_prior_shape:
+    type: float
+    help: "Shape parameter for the gamma prior on the rate."
+    required: false
+    default: 0.5
+  rate_gamma_prior_scale:
+    type: float
+    help: "Scale parameter for the gamma prior on the rate."
+    required: false
+    default: 0.1
+  outdir:
+    type: Path
+    help: "Path to the output directory."
+    required: false
+    default: results
+  dated:
+    type: bool
+    help: Create a timestamped output subdirectory
+    required: false
+    default: false
+  duplicates:
+    type: int
+    help: "Number of duplicate runs. To test for convergence."
+    required: false
+    default: 2
+  trees:
+    type: bool
+    help: "If false, the trees will not be saved."
+    required: false
+    default: true
+  mcc_tree:
+    heights:
+      type: List[str]
+      help: "Height to use for the MCC tree. Can specify multiple. 'mean' (default), 'median', 'keep' or 'ca'"
+      required: false
+      default: ['mean']
+  marginal_likelihood:
+    estimate:
+      type: bool
+      help: "If true, the PS/SS will be used to estimate the marginal likelihood."
+      required: false
+      default: True
+    paths:
+      type: int
+      help: "Number of paths to use for the marginal likelihood estimation."
+      required: false
+      default: 100
+    chain_length:
+      type: int
+      help: "Length of the chain for the marginal likelihood estimation."
+      required: false
+      default: 1000000
+    duplicates:
+      type: int
+      help: "Number of duplicate MLE runs."
+      required: false
+      default: 3
+  fit_clocks:
+    type: bool
+    help: "If true, BEAST will be used to fit the clocks. Use --no-fit-clocks to only run the MLE analysis."
+    required: false
+    default: true
+  beast:
+    threads:
+      type: int
+      help: "Number of threads to use for BEAST."
+      required: false
+      default: 4
+    resources:
+      runtime: 
+        help: "Runtime in minutes."
+        default: 10080 # 7 days
+      mem_mb: 
+        help: "Memory to request."
+        default: 400M
+      partition: 
+        help: "Partition to submit to."
+        default: gpu-a100
+      gres: 
+        help: "GPU resource to request."
+        default: gpu:1
+    envmodules:
+      default:
+        - "GCC/11.3.0"
+        - "beagle-lib/3.1.2-CUDA-11.7.0"

episodic/workflow/Snakefile

@@ -0,0 +1,107 @@
+from utils import decimal_year_to_date
+from scripts.populate_beast_template import taxa_from_fasta
+
+
+SNAKE_DIR = Path(workflow.basedir)
+TEMPLATE_DIR = SNAKE_DIR / "templates"
+SCRIPT_DIR = SNAKE_DIR / "scripts"
+
+OUT_DIR=Path(config["outdir"])
+if config["dated"]:
+    # create timestamped output directory
+    from datetime import datetime
+    OUT_DIR = OUT_DIR / datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
+CLOCK_DIR = OUT_DIR / "clocks"
+
+MLE: bool = config["marginal_likelihood"].get("estimate")
+fit_clocks: bool = config["fit_clocks"]
+rate_gamma_prior_scale=config["rate_gamma_prior_scale"]
+rate_gamma_prior_shape=config["rate_gamma_prior_shape"]
+clocks = expand("{clock}_{rate_gamma_prior_shape}_{rate_gamma_prior_scale}", clock=config["clock"], rate_gamma_prior_shape=rate_gamma_prior_shape, rate_gamma_prior_scale=rate_gamma_prior_scale)
+flc_clocks = [c for c in clocks if "flc" in c]
+relaxed_clocks = [c for c in clocks if "relaxed" in c]
+alignment_path=config["alignment"]
+date_delimiter=config["date_delimiter"]
+date_index=config["date_index"]
+
+duplicates = range(1, config["duplicates"] + 1) if fit_clocks else []
+mle_duplicates = range(1, config["marginal_likelihood"].get("duplicates") + 1) if MLE else []
+
+ALL_LOG_FILES = expand(CLOCK_DIR / "{clock}" / "{clock}_{duplicate}" / "{clock}_{duplicate}.log", clock=clocks, duplicate=duplicates)
+PER_CLOCK_LOG_FILES = lambda wildcards: [CLOCK_DIR / wildcards.clock / f"{wildcards.clock}_{duplicate}" / f"{wildcards.clock}_{duplicate}.log" for duplicate in duplicates]
+
+TAXA = taxa_from_fasta(
+    alignment_path,
+    date_delimiter=date_delimiter,
+    date_index=date_index,
+)
+
+most_recent_sampling_date = decimal_year_to_date(min(TAXA, key=lambda taxa: taxa.date).date)
+
+print(f"Running Episodic with {len(TAXA)} taxa")
+print(f"Most recent sampling date: {most_recent_sampling_date}")
+
+include: "rules/beast.smk"
+include: "rules/config.smk"
+include: "rules/report.smk"
+include: "rules/tree.smk"
+
+
+
+
+CLOCK_FILES = []
+
+CLOCK_FILES.extend(
+    [
+        expand(CLOCK_DIR / "{clock}" / "{clock}_{duplicate}" / "{clock}_{duplicate}_trace_plots", clock=clocks, duplicate=duplicates),
+        expand(CLOCK_DIR / "{clock}" / "{clock}-summary.csv", clock=clocks),
+    ]
+)
+CLOCK_FILES.extend(
+    [CLOCK_DIR / f"{clock}" / f"{clock}-violin.svg" for clock in flc_clocks],
+)
+CLOCK_FILES.extend(
+    [CLOCK_DIR / f"{clock}" / f"{clock}-odds.csv" for clock in flc_clocks],
+)
+CLOCK_FILES.extend(
+    [CLOCK_DIR / "clocks-violin.svg",
+    CLOCK_DIR / "clocks-trace.svg",]
+)
+if config.get("trees"):
+    CLOCK_FILES.extend(
+        expand(
+            CLOCK_DIR / "{clock}" / "{clock}_{duplicate}" / "{clock}_{duplicate}.mcc.{heights}.{ext}",
+            clock=clocks,
+            duplicate=duplicates,
+            heights=config["mcc_tree"].get("heights", "mean"),
+            ext=["nwk", "svg"]
+        ),
+    )
+    CLOCK_FILES.extend(
+        expand(
+            CLOCK_DIR / "{clock}" / "{clock}_{duplicate}" / "{clock}_{duplicate}.rate_quantiles.{ext}", 
+            clock=clocks, 
+            duplicate=duplicates, 
+            ext=["csv", "svg"]
+        )
+    )
+
+
+
+OUTPUT_FILES = [
+    OUT_DIR / "config.yaml",
+]
+
+if fit_clocks:
+    OUTPUT_FILES.extend(
+        CLOCK_FILES
+    )
+
+if MLE:
+    OUTPUT_FILES.append(
+        OUT_DIR / "mle" / "mle.svg"
+    )
+
+rule all:
+    input: 
+        *OUTPUT_FILES

episodic/workflow/envs/arviz.yml

@@ -0,0 +1,12 @@
+name: arviz
+channels:
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip
+  - pip:
+    - arviz
+    - pandas
+    - typer
+    - bokeh
+    - holoviews

episodic/workflow/envs/beast.yml

@@ -0,0 +1,6 @@
+channels:
+  - bioconda
+  - conda-forge
+dependencies:
+  - beast==1.10.4
+  - beagle-lib<4

episodic/workflow/envs/ggtree.yml

@@ -0,0 +1,11 @@
+name: ggtree
+channels:
+  - conda-forge
+  - bioconda
+  - defaults
+dependencies:
+  - r-base
+  - bioconductor-treeio
+  - bioconductor-ggtree
+  - r-optparse
+  - r-svglite

episodic/workflow/envs/phylo.yml

@@ -0,0 +1,12 @@
+name: phylo
+channels:
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip
+  - pip:
+    - dendropy
+    - numpy
+    - typer
+    - scipy
+    - matplotlib

episodic/workflow/envs/plot_traces.yml

@@ -0,0 +1,11 @@
+channels:
+  - conda-forge
+dependencies:
+  - python=3.9
+  - numpy
+  - typer
+  - pandas
+  - plotly
+  - pip
+  - pip:
+    - kaleido

episodic/workflow/envs/python.yml

@@ -0,0 +1,12 @@
+channels:
+  - conda-forge
+  - bioconda
+  - defaults
+dependencies:
+  - python==3.9
+  - pip
+  - pip:
+    - pandas
+    - seaborn
+    - matplotlib
+    - typer

episodic/workflow/profiles/slurm/config.yaml

@@ -0,0 +1,36 @@
+# Note that we use `threads` here as SLURM `--cpus-per-task`.
+cluster: >-
+  mkdir -p logs && sbatch
+  $(if [[ '{resources.account}' ]]; then echo '-A {resources.account}'; fi)
+  $(if [[ '{resources.partition}' ]]; then echo '-p {resources.partition}'; fi)
+  --parsable
+  --time={resources.runtime}
+  --mem={resources.mem_mb}
+  -N {resources.nodes}
+  -n {resources.tasks_per_node}
+  -c {threads}
+  -o logs/{rule}-{wildcards}.out -e logs/{rule}-{wildcards}.err
+  $(if [[ '{resources.qos}' ]]; then echo '-q {resources.qos}'; fi)
+  $(if [[ '{resources.gres}' ]]; then echo '--gres={resources.gres}'; fi)
+  {resources.extra}
+default-resources:
+  - account=''
+  - partition=''
+  - runtime=15
+  - mem_mb=4000
+  - nodes=1
+  - tasks_per_node=1
+  - qos='covid19'
+  - gres=''
+  - extra=''
+cluster-cancel: "scancel {jobid}"
+jobs: 50
+use-conda: true
+use-envmodules: true
+printshellcmds: true
+rerun-incomplete: true
+keep-going: true
+local-cores: 1
+max-jobs-per-second: 10
+max-status-checks-per-second: 1
+latency-wait: 30

episodic/workflow/profiles/slurm/slurm-cancle

@@ -0,0 +1,3 @@
+#!/bin/sh
+
+scancel $(echo $* | grep -Eo '[0-9]+')

episodic/workflow/rules/beast.smk

@@ -0,0 +1,76 @@
+
+rule create_beast_xml:
+    input:
+        alignment = config["alignment"],
+    output:
+        beast_XML_file = CLOCK_DIR / "{clock}" / "{name}" / "{name}.xml",
+    params:
+        template = TEMPLATE_DIR / "beast_xml_template.jinja",
+        date_delimiter = "\|" if config.get("date_delimiter") == "|" else config.get("date_delimiter"),
+        date_index = config.get("date_index", -1),
+        groups = " ".join(config.get("group")),
+        clock = lambda wildcards: wildcards.clock.split("_")[0],
+        rate_gamma_prior_shape = config.get("rate_gamma_prior_shape"),
+        rate_gamma_prior_scale = config.get("rate_gamma_prior_scale"),
+        chain_length = config.get("chain_length"),
+        samples = config.get("samples"),
+        mle = lambda wildcards: "--mle" if "mle" in wildcards.name else "",
+        mle_chain_length = f"--mle-chain-length {config['marginal_likelihood'].get('chain_length')}",
+        mle_path_steps = f"--mle-path-steps {config['marginal_likelihood'].get('paths')}",
+        no_trace = lambda wildcards: "--no-trace" if "mle" in wildcards.name else "",
+        no_trees = lambda wildcards: "--no-trees" if not config.get("trees") or "mle" in wildcards.name else "",
+        fixed_tree = f'--fixed-tree {config.get("fixed_tree")}'  if config.get("fixed_tree") else "",
+    shell:
+        """
+        python {SCRIPT_DIR}/populate_beast_template.py \
+            {params.template} \
+            --output {output.beast_XML_file} \
+            --alignment {input.alignment} \
+            --date-delimiter {params.date_delimiter} \
+            --date-index {params.date_index} \
+            --groups {params.groups} \
+            --clock {params.clock} \
+            --rate-gamma-prior-shape {params.rate_gamma_prior_shape} \
+            --rate-gamma-prior-scale {params.rate_gamma_prior_scale} \
+            --chain-length {params.chain_length} \
+            --samples {params.samples} \
+            {params.mle} \
+            {params.mle_chain_length} \
+            {params.mle_path_steps} \
+            {params.no_trace} \
+            {params.no_trees} \
+            {params.fixed_tree}
+        """
+
+TREES = {"beast_trees_file": CLOCK_DIR / "{clock}" / "{name}" / "{name}.trees"} if config.get("trees") else {}
+
+rule run_beast:
+    input:
+        beast_XML_file = rules.create_beast_xml.output.beast_XML_file,
+    output:
+        beast_stdout_file = CLOCK_DIR / "{clock}" / "{name}" / "{name}.stdout",
+        beast_log_file = CLOCK_DIR / "{clock}" / "{name}" / "{name}.log",
+        **TREES,
+    threads: config["beast"].get("threads")
+    resources:
+        **config["beast"].get("resources", {}),
+    envmodules:
+        *config["beast"].get("envmodules", []),
+    conda:
+        "../envs/beast.yml"
+    shell:
+        """
+        beast -working -overwrite -beagle_GPU -threads {threads} {input.beast_XML_file} > {output.beast_stdout_file}
+        """
+
+MLE_OUT_DIR = OUT_DIR / "mle" / "{clock}"
+
+use rule create_beast_xml as create_mle_xml with:
+    output:
+        beast_XML_file = MLE_OUT_DIR / "{name}.xml",
+
+use rule run_beast as run_mle_beast with:
+    input:
+        beast_XML_file = rules.create_mle_xml.output.beast_XML_file,
+    output:
+        beast_stdout_file = MLE_OUT_DIR / "{name}.stdout",

episodic/workflow/rules/config.smk

@@ -0,0 +1,10 @@
+# convert config to yaml
+import yaml
+
+rule write_config_to_file:
+    output: OUT_DIR / "config.yaml"
+    params: yaml_config=yaml.dump(config)
+    shell:
+        """
+        echo "{params.yaml_config}" > {output}
+        """

episodic/workflow/rules/report.smk

@@ -0,0 +1,95 @@
+MLE_DIR = OUT_DIR / "mle"
+
+rule extract_mle:
+    input:
+        expand(OUT_DIR / "mle" / "{clock}" / "{clock}_mle_{duplicate}.stdout", clock=clocks, duplicate=mle_duplicates),
+    output:
+        MLE_DIR / "mle.svg",
+    conda:
+        "../envs/python.yml"
+    shell:
+        """
+        python {SCRIPT_DIR}/extract_mle.py {MLE_DIR}
+        """
+
+rule plot_traces:
+    """
+    Makes trace plots from the beast log file.
+    """
+    input:
+        rules.run_beast.output.beast_log_file,
+    output:
+        directory(CLOCK_DIR / "{clock}" / "{name}" / "{name}_trace_plots/"),
+    conda:
+        "../envs/plot_traces.yml"
+    shell:
+        """
+        ${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/plot_traces.py {input} {output}
+        """
+
+rule summary:
+    """
+    Makes combined summarys from the beast log file.
+    """
+    input:
+        PER_CLOCK_LOG_FILES,
+    output:
+        posterior_svg=CLOCK_DIR / "{clock}" / "{clock}-summary.csv",
+    params:
+        output=lambda wildcards: CLOCK_DIR / f"{wildcards.clock}" / f"{wildcards.clock}-summary.csv",
+    conda:
+        "../envs/arviz.yml"
+    shell:
+        """
+        ${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/arviz_output.py summary {input} {params.output}
+        """
+
+rule plot_flc_rates:
+    """
+    Makes plots from the flc clock files.
+    """
+    input:
+        lambda wildcards: [CLOCK_DIR / wildcards.clock / f"{wildcards.clock}_{duplicate}" / f"{wildcards.clock}_{duplicate}.log" for duplicate in duplicates if "flc" in wildcards.clock],
+    output:
+        rate_svg=CLOCK_DIR / "{clock}" / "{clock}-violin.svg" 
+    params:
+        directory=lambda wildcards: CLOCK_DIR / f"{wildcards.clock}",
+        gamma_shape=rate_gamma_prior_shape,
+        gamma_scale=rate_gamma_prior_scale,
+    conda:
+        "../envs/arviz.yml"
+    shell:
+        """
+        ${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/arviz_output.py rates {input} {params.directory} --gamma-shape {params.gamma_shape} --gamma-scale {params.gamma_scale}
+        """
+
+use rule plot_flc_rates as plot_rates with:
+    input:
+        expand(CLOCK_DIR / "{clock}" / "{clock}_{duplicate}" / "{clock}_{duplicate}.log", clock=clocks, duplicate=duplicates),
+    output:
+        clocks_violin=CLOCK_DIR / "clocks-violin.svg",
+        clocks_trace=CLOCK_DIR / "clocks-trace.svg",
+    params:
+        directory=CLOCK_DIR,
+        gamma_shape=rate_gamma_prior_shape,
+        gamma_scale=rate_gamma_prior_scale,
+
+rule calculate_odds:
+    input:
+        lambda wildcards: [CLOCK_DIR / wildcards.clock / f"{wildcards.clock}_{duplicate}" / f"{wildcards.clock}_{duplicate}.log" for duplicate in duplicates if "flc" in wildcards.clock],
+    output:
+        rate_svg=CLOCK_DIR / "{clock}" / "{clock}-odds.csv" 
+    params:
+        directory=lambda wildcards: CLOCK_DIR / f"{wildcards.clock}",
+        gamma_shape=rate_gamma_prior_shape,
+        gamma_scale=rate_gamma_prior_scale,
+    conda:
+        "../envs/python.yml"
+    shell:
+        """
+        ${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/calculate_odds.py {input} {output} --gamma-shape {params.gamma_shape} --gamma-scale {params.gamma_scale}
+        for file in {input}
+        do
+            ${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/calculate_odds.py $file ${{file%.log}}-odds.csv --gamma-shape {params.gamma_shape} --gamma-scale {params.gamma_scale}
+        done
+        """

episodic/workflow/rules/tree.smk

@@ -0,0 +1,68 @@
+
+rule max_clade_credibility_tree:
+    """
+    Makes trace plots from the beast log file.
+    """
+    input:
+        rules.run_beast.output.beast_trees_file,
+    output:
+        CLOCK_DIR / "{clock}" / "{name}" / "{name}.mcc.{heights}.nexus",
+    params:
+        burnin = int(int(config['samples']) * 0.1),
+    conda:
+        "../envs/beast.yml"
+    shell:
+        """
+        treeannotator -burninTrees {params.burnin} -heights {wildcards.heights} {input} {output}
+        """
+
+
+rule max_clade_credibility_tree_newick:
+    """
+    Makes trace plots from the beast log file.
+    """
+    input:
+        rules.max_clade_credibility_tree.output,
+    output:
+        CLOCK_DIR / "{clock}" / "{name}" / "{name}.mcc.{heights}.nwk",
+    conda:
+        "../envs/phylo.yml"
+    shell:
+        "${{CONDA_PREFIX}}/bin/python {SCRIPT_DIR}/tree_converter.py {input} {output} --node-label posterior"
+
+
+rule max_clade_credibility_tree_render:
+    """
+    Renders the MCC tree in SVG format.
+    """
+    input:
+        rules.max_clade_credibility_tree.output,
+    output:
+        CLOCK_DIR / "{clock}" / "{name}" / "{name}.mcc.{heights}.svg",
+    params:
+        mrsd = most_recent_sampling_date,
+    conda:
+        "../envs/ggtree.yml"
+    shell:
+        "${{CONDA_PREFIX}}/bin/RScript {SCRIPT_DIR}/plotMCCtree.R --input {input} --output {output} --mrsd {params.mrsd}"
+
+
+rule rate_quantile_analysis:
+    input:
+        rules.run_beast.output.beast_trees_file,
+    output:
+        csv = CLOCK_DIR / "{clock}" / "{name}" / "{name}.rate_quantiles.csv",
+        svg = CLOCK_DIR / "{clock}" / "{name}" / "{name}.rate_quantiles.svg",
+    params:
+        groups = " ".join(f"-g {group}" for group in config['group']),
+    conda:
+        "../envs/phylo.yml"
+    shell:
+        """
+        python {SCRIPT_DIR}/phylo_rate_quantile_analysis.py \
+          {input} \
+          --output-csv {output.csv} \
+          --output-plot {output.svg} \
+          {params.groups}
+        """
+