CIAassembly_pipeline

import pandas as pd
import os
from pathlib import Path
import sys

metadata = pd.read_table(
    (Path(workflow.basedir) / "config/units.tsv"), 
    dtype=str
    ).set_index(["sample"], drop=False)

def relative_to_basedir(path, basedir='basedir'):
    path = Path(path)
    if path.is_absolute():
        return str(path)
    else: 
        if basedir == 'basedir':
            base = Path(str(workflow.basedir))
        elif basedir == 'workdir':
            base = Path(os.getcwd())
        else:
            raise Exception('Must specify basedir or workdir.')
        return str((base / Path(path)).resolve())

# Convert relative paths to absolute path relative to workflow basedir
metadata['path'] = metadata['path'].map(relative_to_basedir)

# Add a workspace_path variable pointing to 
# data/file.fastq etc. This is relative to the workdir. 
metadata['workspace_path'] = (metadata['path']
    .map(lambda x: str(Path('data') / Path(x).name))
)

# add mapper attribute to metadata from config file
# either starlong or minimap2
metadata['mapper'] = (metadata['sample_type']
    .map(lambda x: config['mapper'][x])
)

def meta_attr(attrib='path', df=metadata):
    "Return a function taking wildcards and returning the "
    "given attribute, return the given in the metadata table."
    def inner(wc):
        ans = df.loc[str(wc.sample), attrib]
        return ans # df.loc[str(wc.sample), attrib]
    return inner

print("Starting CIA assembly with metadata object: {}".format(metadata), file=sys.stderr)

localrules: all, samtools_index

rule all:
    input:
        expand("cia_assemblies/{sample}.end.corrected.assembly.gtf", 
        sample=list(metadata['sample']))

def _find_annotation(wc):
    "If relative path, use relative to basedir"
    path = Path(config["annotation"][str(wc.annotation)])
    if path.is_absolute():
        return str(path)
    else: 
        return str(Path(workflow.basedir) / path)

rule symlink_annotation:
    input: 
        _find_annotation
    output: 
        "annotation/{annotation}"
    shell: 
        "ln -fs `readlink -f {input}` {output}"

def _find_path(wc):
    "Return the path of the file with the given basename wc.basename."
    ans = metadata[metadata['workspace_path'].map(lambda x: x.endswith(str(wc.basename)))]
    return ans['path'].values[0]


rule symlink_data: 
    input: 
        _find_path
    output: 
        "data/{basename}"
    shell:
        "ln -fs `readlink -f {input}` {output}"


rule minimap_idx:
    input: 
        target="annotation/genome.fasta"
    output: 
        "index/minimap2/genome.mmi"
    params:
        extra="-k 14"
    threads: 
        16
    wrapper: 
        "v1.20.0/bio/minimap2/index"

rule star_index:
    input:
        fasta="annotation/genome.fasta",
        gtf="annotation/transcriptome.gtf"
    output:
        directory("index/star"),
    threads: 16
    params:
        extra="--genomeSAindexNbases 12",
    log:
        "logs/star_index.log",
    wrapper:
        "v1.20.0/bio/star/index"

rule minimap_align:
    input:
        fastq=lambda wc: metadata.loc[str(wc.sample), 'workspace_path'],
        genome="annotation/genome.fasta",
        index="index/minimap2/genome.mmi"
    output:
        bam="mapping/minimap2/{sample}.bam",
    threads: 
        16
    benchmark: 
        "benchmarks/mapping/minimap2/{sample}_benchmark.txt"
    resources: 
        mem_mb=72000,
    conda: 
        "envs/minimap_samtools.yaml"
    shell: 
        "minimap2 -t {threads} "
        "-ax splice -uf "
        "{input.genome} "
        "{input.fastq} | "
        "samtools sort -@ {threads} -T $TMPDIR/{wildcards.sample} -o {output.bam} -; "

def _to_zcat(wc):
    "If sample ends with gz, then add zcat"
    sample_path = metadata.loc[str(wc.sample), 'workspace_path']
    if sample_path.endswith('.gz'):
        return '--readFilesCommand zcat'
    else:
        return ''

rule starlong_align: 
    input: 
        data=lambda wc: metadata.loc[str(wc.sample), 'workspace_path'],
        genome="index/star",
    output: 
        'mapping/starlong/{sample}.bam'
    conda: 
        "envs/star.yaml"
    threads: 16
    benchmark: 
        "benchmarks/mapping/starlong/{sample}_benchmark.txt"
    log: "logs/mapping/star/{sample}.log"
    resources: 
        mem_mb=72000,
    params: 
        zcat=_to_zcat,
        prefix='mapping/starlong/{sample}',
    script: 
        "scripts/starlong.py"

rule samtools_index:
    input:
        "{prefix}.bam",
    output:
        "{prefix}.bam.bai",
    log: 
        "logs/samtools/{prefix}.log"
    params:
        extra="",  
    threads: 4  
    wrapper:
        "v1.20.0/bio/samtools/index"

def _get_mapper_path(pattern="mapping/{mapper}/{{sample}}.bam"):
    "Given a path pattern, return a function taking wc and either "
    "mapping/minimap2/{sample}.bam or mapping/starlong/{sample}.bam"
    def inner(wc):
        mapper = metadata.loc[str(wc.sample), 'mapper']
        return pattern.format(mapper=mapper)
    return inner


rule bam2bed12minimap2:
    input: 
        bam2bed="src/flair/bin/bam2Bed12.py",
        bam=_get_mapper_path("mapping/{mapper}/{{sample}}.bam"),
        bai=_get_mapper_path("mapping/{mapper}/{{sample}}.bam.bai"),
        genome="annotation/genome.fasta",
    output:
        "bed/{sample}.classified.bed"
    threads: 16
    conda: 
        "envs/flair.yaml"
    shell: 
        "python {input.bam2bed} -i {input.bam} > {output}"

rule flair_correct:
    input: 
        bed="bed/{sample}.classified.bed",
        genome="annotation/genome.fasta",
        chrom_sizes="annotation/genome.chrom.sizes",
        gtf="annotation/transcriptome.gtf",
        flair="src/flair/flair.py"
    output:
        "assembly/flair/{sample}_all_corrected.psl"
    params:
        out_dir="assembly/flair/{sample}"
    threads: 16
    conda: 
        "envs/flair.yaml"
    shell: 
        "python {input.flair} correct "
        "--print_check "
        "-c {input.chrom_sizes} "
        "-t {threads} "
        "-q {input.bed} "
        "-g {input.genome} "
        "-f {input.gtf} "
        "-o {params.out_dir} "


def get_flair_cmdline(wc):
    sample_type = metadata.loc[str(wc.sample), "sample_type"]
    if sample_type not in config["sample_types"]:
        raise ValueError(
            "Error: {sample} has type {type} which is not one of "
            "allowed types {types}".format(
                sample=wc.sample, 
                type=sample_type,
                types=config["sample_types"]
            )
        )
    return config["params"]["flair_collapse_edp_tss"][sample_type]

rule flair_download: 
    output: 
        flair="src/flair/flair.py",
        bam2bed="src/flair/bin/bam2Bed12.py",
    params:
        outdir="src/flair",
        commit="d242d5ce6103158bd26b63f12e7598f9f70dae67"
    shell: 
        """
        if [ -d {params.outdir} ]; then
            rm -rf {params.outdir};
        fi;
        git clone https://github.com/BrooksLabUCSC/flair {params.outdir}
        cd {params.outdir}
        git reset --hard {params.commit}
        """

rule flair_collapse_edp_tss:
    input: 
        fastq=lambda wc: metadata.loc[str(wc.sample), 'workspace_path'],
        psl="assembly/flair/{sample}_all_corrected.psl",
        gtf="annotation/transcriptome.gtf",
        promoters="annotation/promoter.db.edp.bed",
        genome="annotation/genome.fasta",
        flair="src/flair/flair.py"
    output:
        "assembly/flair/{sample}_collapsed_annot.edp.tss/{sample}.isoforms.gtf"
    params:
        outDir="assembly/flair/{sample}_collapsed_annot.edp.tss/{sample}",
        tempdir="tmp/{sample}_flair_collapse",
        cmdline_params=get_flair_cmdline,
    threads: 20
    conda: 
        "envs/flair.yaml" # make sure includes minimap
    shell: 
        "mkdir -p {params.tempdir}; "
        "python {input.flair} collapse "
        "-g {input.genome} "
        "--temp_dir {params.tempdir} "
        "-r {input.fastq} "
        "-t {threads} "
        "-q {input.psl} "
        "-f {input.gtf} "
        "-p {input.promoters} "
        "-o {params.outDir} "
        "{params.cmdline_params}; " # i.e. -s 2 or -s 3
        "rm -rf {params.tempdir} "

rule sqanti_download: 
    output: "src/sqanti/sqanti3_qc.py"
    params:
        outdir="src/sqanti",
        commit=config["sqanti_git_hash"],
    shell: 
        """
        if [ -d "{params.outdir}" ]; then
            rm -rf {params.outdir};
        fi;
        git clone https://github.com/ConesaLab/SQANTI3 {params.outdir}
        cd {params.outdir}
        git reset --hard {params.commit}
        cd ../..
        """


# necessary since we install from an explicit environment
rule sqanti_env:
    output: "flags/sqanti_env.done"
    params: 
        sqanti_env=relative_to_basedir('envs/sqanti.txt')
    shell:
        """
        eval "$(conda shell.bash hook)"
        if ! `conda env list | grep cia-sqanti >/dev/null 2>/dev/null`; then
            conda create --file {params.sqanti_env} --name cia-sqanti
        fi;
        touch {output}
        """

rule sqanti_env_pip: 
    input: "flags/sqanti_env.done"
    output: touch("flags/sqanti_env_pip.done")
    conda: "cia-sqanti"
    shell: 
        """
        pip install bx-python==0.8.8
        pip install sklearn==0.0
        """

rule sqanti_env_cupcake: 
    input: "flags/sqanti_env.done"
    output: touch("flags/cupcake.done")
    conda: 'cia-sqanti'
    params:
        outdir="src/cDNA_Cupcake",
    shell:
        """
        if [ -d "{params.outdir}" ]; then \
            rm -rf {params.outdir}; \
        fi; 
        git clone https://github.com/Magdoll/cDNA_Cupcake.git {params.outdir}
        cd {params.outdir}
        python setup.py build
        python setup.py install
        cd ../..
        """

rule gtf_to_gene_pred: 
    input: 
        env="flags/sqanti_env.done",
        qc_script="src/sqanti/sqanti3_qc.py"
    output: "src/sqanti/utilities/gtfToGenePred"
    shell: 
        """
        wget http://hgdownload.cse.ucsc.edu/admin/exe/linux.x86_64/gtfToGenePred -O {output}
        chmod +x {output} 
        """

rule sqanti_qc_normal:
    input: 
        gtf="assembly/flair/{sample}_collapsed_annot.edp.tss/{sample}.isoforms.gtf",
        gtf_anno="annotation/transcriptome.gtf",
        genome="annotation/genome.fasta",
        polya_list="annotation/sqanti.polya.list",
        junctions_filtered="annotation/splice_junctions_filtered.tab",
        sqanti="src/sqanti/sqanti3_qc.py",
        flag="flags/sqanti_env_pip.done",
        cupcake_flag="flags/cupcake.done",
        gtftogenepred="src/sqanti/utilities/gtfToGenePred",
    output:
        gtf="assembly/sqanti/{sample}_collapsed_annot.edp.tss/{sample}.sqanti_corrected.gtf",
        classifxn="assembly/sqanti/{sample}_collapsed_annot.edp.tss/{sample}.sqanti_classification.txt"
    params:
        outDir="assembly/sqanti/{sample}_collapsed_annot.edp.tss/",
        prefix="{sample}.sqanti",
        cupcake=relative_to_basedir("src/cDNA_Cupcake", basedir='workdir'),
        internalPrimingWindow=50,
    resources: 
        mem_mb=40000
    benchmark:
        "benchmarks/sqanti/{sample}_benchmark.txt"
    threads: 16
    conda: 
        'cia-sqanti'
    shell: 
        "export PATH=$PATH:{params.cupcake}/sequence/; "
        "export PATH=$PATH:{params.cupcake}/rarefaction/; "
        "export PYTHONPATH={params.cupcake}/sequence/; "
        "echo $PYTHONPATH; "
        "if [ -d {params.outDir} ]; then "
        "   rm -rf {params.outDir}; "
        "fi; "
        "python {input.sqanti} "
        "--gtf {input.gtf} {input.gtf_anno}  {input.genome} "
        "-o {params.prefix} "
        "-d {params.outDir} "
        "--force_id_ignore "
        "-w {params.internalPrimingWindow} "
        "--polyA_motif_list {input.polya_list} "
        "-c {input.junctions_filtered} "
        "|| exit 0"

rule endCorrection:
    input: 
        endGuides="annotation/combined.rds.clusters.new.gff",
        annot="assembly/sqanti/{sample}_collapsed_annot.edp.tss/{sample}.sqanti_corrected.gtf",
        classifxn="assembly/sqanti/{sample}_collapsed_annot.edp.tss/{sample}.sqanti_classification.txt",
        txannot="annotation/transcriptome.gtf",
    output:
        "cia_assemblies/{sample}.end.corrected.assembly.gtf"
    params:
        outDir="cia_assemblies",
        R_script=relative_to_basedir('R/RDSFLAM.endGuidedCorrection.R'),
        R_source=relative_to_basedir('R/CIAmethods_source.R'),
    threads: 16
    resources:
        mem_mb=24000
    benchmark:
        "benchmarks/end_correction/{sample}_benchmark.txt"
    conda: 
        "envs/end_correction.yaml"
    shell: 
        "Rscript {params.R_script} "
        "-a {input.annot} "
        "-e {input.endGuides} "
        "-s {input.classifxn} "
        "-r {input.txannot} "
        "-o {params.outDir} "
        "-p {wildcards.sample} "
        "-c {params.R_source} "

# annotFile <-  "/data/processing1/alfonso/rerunAssembly/cdna_seq/cia_isoforms/cia_assemblies/heads_cia_assembly/cia_heads.all.seqs.gtf"
# LongEndsFile <-"data/cia_isoforms/clusters.rds/combined.rds.clusters.new.gff"
# sqantiClass <- "/data/processing1/alfonso/rerunAssembly/cdna_seq/cia_isoforms/cia_assemblies/heads_cia_assembly/sqanti/cia_heads.all.seqs.sqanti_classification.txt"
# ensAnnotFile <- "/data/repository/organisms/dm6_ensembl/ensembl/release-96/genes.gtf"