multigene_LOF_search.snake

"""Snakemake file."""
import logging
import logzero
from logzero import logger as log

import os
import inspect

from pathlib import Path

import ruamel.yaml as yaml

import pandas as pd
import numpy as np

from matplotlib import pyplot as plt
import seaborn as sns
sns.set_style("whitegrid")

import munch

from veoibd_synapse.rules import pathify_by_key_ends, SnakeRule, SnakeRun, recode_graph, apply_template
from veoibd_synapse.misc import process_config

from veoibd_synapse import misc

# Metadata
__author__ = "Gus Dunn"
__email__ = "w.gus.dunn@gmail.com"


#### COMMON RUN STUFF ####
original_config_path = Path(workflow.overwrite_configfile)
ORIGINAL_CONFIG_AS_STRING = original_config_path.read_text()

cfg = process_config(str(original_config_path))
cfg = pathify_by_key_ends(cfg)

# Decide whether to output debug lvl msgs
if cfg.COMMON.DEBUG == True:
    logzero.loglevel(logging.DEBUG)
else:
    logzero.loglevel(logging.INFO)


log.debug("Initializing common run stuff.")
RUN = SnakeRun(cfg=cfg, snakefile=Path(inspect.getfile(inspect.currentframe())))


# Save current copy of the configfile
SAVED_CONFIG = RUN.out_dir / "{NAME}.yaml".format(NAME=RUN.name)
SAVED_CONFIG.parent.mkdir(parents=True, exist_ok=True)
with SAVED_CONFIG.open('w') as cnf_out:
    cnf_out.write(ORIGINAL_CONFIG_AS_STRING)

PRE = []
ALL = []
DEBUG = []

# add specific useful stuff to RUN
RUN.globals.input_vcfs = [vcf.parts[-1].rstrip('.vcf.gz') for vcf in cfg.VCF_CHECK.IN.VCF_DIR.glob("*.vcf.gz")]

log.debug("BEGIN defining rules.")
############ BEGIN PIPELINE RULES ############
# ------------------------- #
#### VCF_CHECK ####
VCF_CHECK = SnakeRule(run=RUN, name="VCF_CHECK")

# input
VCF_CHECK.i.subjects_vcf = str(VCF_CHECK.IN.VCF_DIR / "{vcf}.vcf.gz"),

# output
VCF_CHECK.o.vcf_checked_sentinel = str(RUN.interim_dir / "{vcf}.vcf.has_been_checked")

VCF_CHECK.o.vcf_checked_sentinel_expd = expand(VCF_CHECK.o.vcf_checked_sentinel,
                                               vcf=RUN.globals.input_vcfs)

# ---
rule VCF_CHECK:
    log:
        path=str(VCF_CHECK.log)

    input:
        subjects_vcf=VCF_CHECK.i.subjects_vcf,

    output:
        subjects_checked_vcf=VCF_CHECK.o.vcf_checked_sentinel,

    run:
        snpsift = """SnpSift -Xmx4g vcfCheck {input.subjects_vcf} 2> {output.subjects_checked_vcf} 2> {log.path}"""
        shell(snpsift)

ALL.append(VCF_CHECK.o.vcf_checked_sentinel_expd)




# ------------------------- #
#### FILTER_GTF ####
FILTER_GTF = SnakeRule(run=RUN, name="FILTER_GTF")

# params
FILTER_GTF.p.a_patterns = FILTER_GTF.PARAMS.A_PATTERNS
FILTER_GTF.p.b_pattern = FILTER_GTF.PARAMS.B_PATTERN

# input
FILTER_GTF.i.gtf = str(FILTER_GTF.IN.GTF)

# output
FILTER_GTF.o.temp_gtf = str(FILTER_GTF.out_dir / "temp_gtf.gtf")
FILTER_GTF.o.filtered_gtf = str(FILTER_GTF.out_dir / "filtered_gtf.gtf")

# ---
rule FILTER_GTF:
    log:
        path=str(FILTER_GTF.log)

    params:
        a_patterns=FILTER_GTF.p.a_patterns,
        b_pattern=FILTER_GTF.p.b_pattern,

    input:
        gtf=FILTER_GTF.i.gtf,

    output:
        temp_gtf=temp(FILTER_GTF.o.temp_gtf),
        filtered_gtf=FILTER_GTF.o.filtered_gtf,

    run:
        shell("rm -f {log.path}")

        cmd_tmplt = """grep -P '{pattern}' {{input.gtf}} >> {{output.temp_gtf}} 2>> {{log.path}}"""

        for pattern in params.a_patterns:
            shell(cmd_tmplt.format(pattern=pattern))

        shell("""grep -P "{pattern}" {{output.temp_gtf}} > {{output.filtered_gtf}} 2>> {{log.path}}""".format(pattern=params.b_pattern))


ALL.append(rules.FILTER_GTF.output)


# ------------------------- #
#### FILTER_SUBJECTS_VCFS ####
FILTER_SUBJECTS_VCFS = SnakeRule(run=RUN, name="FILTER_SUBJECTS_VCFS")

# input
FILTER_SUBJECTS_VCFS.i.subjects_vcf = str(VCF_CHECK.IN.VCF_DIR / "{vcf}.vcf.gz"),
FILTER_SUBJECTS_VCFS.i.vcf_checked_sentinel_expd = VCF_CHECK.o.vcf_checked_sentinel_expd,
FILTER_SUBJECTS_VCFS.i.filtered_gtf = FILTER_GTF.o.filtered_gtf

# output
FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf = str(FILTER_SUBJECTS_VCFS.out_dir / "{vcf}.filtered.vcf")

FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf_expd = expand(FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf,
                                                           vcf=RUN.globals.input_vcfs)


# ---
rule FILTER_SUBJECTS_VCFS:
    log:
        path=str(FILTER_SUBJECTS_VCFS.log)

    input:
        subjects_vcf=FILTER_SUBJECTS_VCFS.i.subjects_vcf,
        vcf_checked_sentinel_expd=FILTER_SUBJECTS_VCFS.i.vcf_checked_sentinel_expd,
        filtered_gtf=FILTER_SUBJECTS_VCFS.i.filtered_gtf,

    output:
        subjects_filtered_vcf=FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf,

    run:
        shell("vt view -H {input.subjects_vcf} > {output.subjects_filtered_vcf} 2> {log.path}")

        shell("bedtools intersect -u -a {input.subjects_vcf} -b {input.filtered_gtf} >> "
              " {output.subjects_filtered_vcf} 2>> {log.path}")

ALL.append(FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf_expd)



# ------------------------- #
#### SNPEFF ####
SNPEFF = SnakeRule(run=RUN, name="SNPEFF")

# params
SNPEFF.p.genome_db = SNPEFF.PARAMS.GENOME_DB

# input
SNPEFF.i.subjects_filtered_vcf = FILTER_SUBJECTS_VCFS.o.subjects_filtered_vcf

# output
SNPEFF.o.stats = str(SNPEFF.out_dir / "{vcf}.html")
SNPEFF.o.snpeff_vcf = str(SNPEFF.out_dir / "{vcf}.snpeff.vcf")
# SNPEFF.o.output_1 = str(SNPEFF.out_dir / "{something}.ext")

SNPEFF.o.stats_expd = expand(SNPEFF.o.stats,
                             vcf=RUN.globals.input_vcfs)

SNPEFF.o.snpeff_vcf_expd = expand(SNPEFF.o.snpeff_vcf,
                                  vcf=RUN.globals.input_vcfs)

# ---
rule SNPEFF:
    log:
        path=str(SNPEFF.log)

    params:
        genome_db=SNPEFF.p.genome_db,

    input:
        subjects_filtered_vcf=SNPEFF.i.subjects_filtered_vcf,

    output:
        stats=SNPEFF.o.stats,
        snpeff_vcf=SNPEFF.o.snpeff_vcf,

    shell:
        "snpEff -Xmx4g -v -stats {output.stats} {params.genome_db} {input.subjects_filtered_vcf} > {output.snpeff_vcf} "
        "&> {log.path} "

ALL.append(SNPEFF.o.stats_expd)
ALL.append(SNPEFF.o.snpeff_vcf_expd)



# ------------------------- #
#### SNPSIFT_ANNOTATE ####
SNPSIFT_ANNOTATE = SnakeRule(run=RUN, name="SNPSIFT_ANNOTATE")

# params
# SNPSIFT_ANNOTATE.p.param_1 = SNPSIFT_ANNOTATE.PARAMS.param_1

# input
SNPSIFT_ANNOTATE.i.snpeff_vcf = SNPEFF.o.snpeff_vcf
SNPSIFT_ANNOTATE.i.dbsnp = SNPSIFT_ANNOTATE.IN.DBSNP

# output
SNPSIFT_ANNOTATE.o.annotated_vcf = str(SNPSIFT_ANNOTATE.out_dir / "{vcf}.snpeff.annotated.vcf")

SNPSIFT_ANNOTATE.o.annotated_vcf_expd = expand(SNPSIFT_ANNOTATE.o.annotated_vcf,
                                               vcf=RUN.globals.input_vcfs)

# ---
rule SNPSIFT_ANNOTATE:
    log:
        path=str(SNPSIFT_ANNOTATE.log)

    # params:
    #     param_1=SNPSIFT_ANNOTATE.p.param_1,

    input:
        snpeff_vcf=SNPSIFT_ANNOTATE.i.snpeff_vcf,
        dbsnp=SNPSIFT_ANNOTATE.i.dbsnp,

    output:
        annotated_vcf=SNPSIFT_ANNOTATE.o.annotated_vcf,

    shell:
        "SnpSift -Xmx4g annotate {input.dbsnp} {input.snpeff_vcf} > {output.annotated_vcf} "
        "&> {log.path} "

ALL.append(SNPSIFT_ANNOTATE.o.annotated_vcf_expd)



# ------------------------- #
#### SNPSIFT ####
SNPSIFT = SnakeRule(run=RUN, name="SNPSIFT")

# params
SNPSIFT.p.filter_str = SNPSIFT.PARAMS.FILTER_STR

# input
SNPSIFT.i.snpeff_vcf = SNPSIFT_ANNOTATE.o.annotated_vcf

# output
SNPSIFT.o.snpsift_vcf = str(SNPSIFT.out_dir / "{vcf}.snpsift.vcf")

SNPSIFT.o.snpsift_vcf_expd = expand(SNPSIFT.o.snpsift_vcf,
                                   vcf=RUN.globals.input_vcfs)

# ---
rule SNPSIFT:
    log:
        path=str(SNPSIFT.log)

    params:
        filter_str=SNPSIFT.p.filter_str,

    input:
        snpeff_vcf=SNPSIFT.i.snpeff_vcf,

    output:
        snpsift_vcf=SNPSIFT.o.snpsift_vcf,

    shell:
        """cat {input.snpeff_vcf} | SnpSift -Xmx4g filter "{params.filter_str}" > {output.snpsift_vcf}"""
        "&> {log.path} "

ALL.append(SNPSIFT.o.snpsift_vcf_expd)




# ------------------------- #


#### ALL ####
# ---
rule all:
    input: ALL


# ------------------------- #
#### DRAW_RULE_GRAPH ####
DRAW_RULE_GRAPH = SnakeRule(run=RUN, name="DRAW_RULE_GRAPH")

# params
DRAW_RULE_GRAPH.p.pretty_names = RUN.pretty_names

# input

# output
DRAW_RULE_GRAPH.o.rule_graph_dot = str(DRAW_RULE_GRAPH.out_dir / "rule_graph.dot")
DRAW_RULE_GRAPH.o.recoded_rule_graph_dot = str(DRAW_RULE_GRAPH.out_dir / "recoded_rule_graph.dot")
DRAW_RULE_GRAPH.o.recoded_rule_graph_svg = str(DRAW_RULE_GRAPH.out_dir / "recoded_rule_graph.svg")

# ---
rule draw_rule_graph:
    log:
        path=str(DRAW_RULE_GRAPH.log)

    params:
        pretty_names=DRAW_RULE_GRAPH.p.pretty_names,

    input:
        Snakefile=str(RUN.snakefile.absolute()),
        config=str(SAVED_CONFIG),

    output:
        rule_graph_dot=DRAW_RULE_GRAPH.o.rule_graph_dot,
        recoded_rule_graph_dot=DRAW_RULE_GRAPH.o.recoded_rule_graph_dot,
        recoded_rule_graph_svg=DRAW_RULE_GRAPH.o.recoded_rule_graph_svg,

    run:
        rule_name = cfg.COMMON.DRAW_RULE
        shell("snakemake -p -s {input.Snakefile}  --configfile {input.config} "+rule_name+" --rulegraph > {output.rule_graph_dot}")

        recode_graph(dot=output.rule_graph_dot,
                     new_dot=output.recoded_rule_graph_dot,
                     pretty_names=RUN.pretty_names,
                     rules_to_drop=['save_run_config',rule_name],
                     color="#50D0FF",
                     use_pretty_names=False)

        shell("dot -Tsvg {output.recoded_rule_graph_dot} -o {output.recoded_rule_graph_svg} -v ; echo ''")


# ------------------------- #
#### DRAW_DAG_GRAPH ####
DRAW_DAG_GRAPH = SnakeRule(run=RUN, name="DRAW_DAG_GRAPH")

# params
DRAW_DAG_GRAPH.p.pretty_names = RUN.pretty_names

# input

# output
DRAW_DAG_GRAPH.o.dag_graph_dot = str(DRAW_DAG_GRAPH.out_dir / "dag_graph.dot")
DRAW_DAG_GRAPH.o.recoded_dag_graph_dot = str(DRAW_DAG_GRAPH.out_dir / "recoded_dag_graph.dot")
DRAW_DAG_GRAPH.o.recoded_dag_graph_svg = str(DRAW_DAG_GRAPH.out_dir / "recoded_dag_graph.svg")

# ---
rule draw_dag_graph:
    log:
        path=str(DRAW_DAG_GRAPH.log)

    params:
        pretty_names=DRAW_DAG_GRAPH.p.pretty_names,

    input:
        Snakefile=str(RUN.snakefile.absolute()),
        config=str(SAVED_CONFIG),

    output:
        dag_graph_dot=DRAW_DAG_GRAPH.o.dag_graph_dot,
        recoded_dag_graph_dot=DRAW_DAG_GRAPH.o.recoded_dag_graph_dot,
        recoded_dag_graph_svg=DRAW_DAG_GRAPH.o.recoded_dag_graph_svg,

    run:
        rule_name = cfg.COMMON.DRAW_RULE
        shell("snakemake -p -s {input.Snakefile}  --configfile {input.config} "+rule_name+" --dag > {output.dag_graph_dot}")

        recode_graph(dot=output.dag_graph_dot,
                     new_dot=output.recoded_dag_graph_dot,
                     pretty_names=RUN.pretty_names,
                     rules_to_drop=['save_run_config',rule_name],
                     color="#50D0FF",
                     use_pretty_names=False)

        shell("dot -Tsvg {output.recoded_dag_graph_dot} -o {output.recoded_dag_graph_svg} -v ; echo ''")