Snakefile-sorting-hat

import yaml
import glob
from snakemake.utils import min_version


# set minimum snakemake version
min_version("7.32.4")

metadata = {}

def get_cluster_config_file():
    """
    Returns the rule-specific cluster (lsf or slurm) resource settings file to be copied to
    acc working directory before the snakemake call
    """
    if 'cluster_config' in config:
        return config['cluster_config']
    else:
        return None


def fill_metadata(acc):
    with open(f"{acc}/{acc}.metadata_summary.yaml", 'r') as meta:
        metadata[acc] = yaml.load(meta, Loader=yaml.FullLoader)

def get_log_handler_script():
    if 'loghandler' in config:
        return f"--log-handler-script {config['loghandler']} &>/dev/null"
    return ""

def get_tool(acc):
    """
    If a tool is given in the config, then that tool is preferred,
    otherwise it is determined by the experiment type.
    """
    if 'tool' in config:
        return config['tool']
    if acc not in metadata:
        fill_metadata(acc)
    if 'experiment_type' not in metadata[acc]:
        print(f"Experiment type not found for {acc}")
        return None
    if 'differential' in metadata[acc]['experiment_type']:
        return "all-diff"
    if 'baseline' in metadata[acc]['experiment_type']:
        return "all-baseline"
    print(f"Could not determine tool for {acc}")
    return None


def get_organism(acc):
    if acc not in metadata:
        fill_metadata(acc)
    return metadata[acc]['organism']

def get_exp_type(acc):
    if acc not in metadata:
        fill_metadata(acc)
    return metadata[acc]['experiment_type']

def get_assay_contrast_line(acc):
    if acc not in metadata:
        fill_metadata(acc)
    if 'contrasts' in metadata[acc]:
        contrast_ids = "::".join(metadata[acc]['contrasts'].keys())
        contrast_labels = "&&".join(metadata[acc]['contrasts'].values())
        return f"contrast_ids={contrast_ids} contrast_labels=\"{contrast_labels}\""
    elif 'assays' in metadata[acc]:
        assay_ids = "::".join(metadata[acc]['assays'].keys())
        assay_labels = "&&".join(metadata[acc]['assays'].values())
        return f"assay_ids={assay_ids} assay_labels=\"{assay_labels}\""
    return None

def get_gff_path(acc):
    if acc not in metadata:
        fill_metadata(acc)
    return metadata[acc]['gff']

def get_skip_steps_file():
    if 'skip_steps_file' in config:
        return config['skip_steps_file']
    else:
        return None

def get_run_deconv_file():
    if 'run_deconv_file' in config:
        return config['run_deconv_file']
    else:
        return None

def get_meta_config():
    if 'atlas_meta_config' in config:
        return config['atlas_meta_config']
    else:
        return None

def get_user_species(wildcards):
    if 'species' in config:
        if get_organism( wildcards['accession'] ) in config['species'].split(":"):
            return 'True'
        else:
            return 'False'
    else:
        return 'True'

def get_db_params(param):
    """Return config parameters to establish connection with isl db."""
    if param in config:
        return config[param]
    else:
        sys.exit(1)

def read_skip_accessions_file(analysis_goal):
    """
    List of accessions that should be skipped from reprocessing
    or recalculations
    """
    if analysis_goal == 'reprocess':
        if 'skip_accessions_reproc_file' in config:
            with open(config['skip_accessions_reproc_file'], 'r') as stream:
                try:
                    skip_accessions_reprocessing=yaml.safe_load(stream)
                except yaml.YAMLError as exc:
                    print(exc)
            return skip_accessions_reprocessing['skips']['accession']
        else:
            return None
    elif analysis_goal == 'recalculations':
        print(f"NOTE - skip accessions not implemented for recalculations")
        return None
    else:
        print(f"WARNING - analysis_goal {analysis_goal} is not reprocess nor recalculations")
        return None



def get_outputs():
    import os
    ACC=[]
    if 'accessions' in config:
        ACC=config['accessions'].split(":")
    else:
        ACC=glob.glob('E-*')

    if 'goal' in config:
        if config['goal'] == 'recalculations':

            if 'atlas_exps' in config:
                # filter out ACC based on accessions that are available in atlas_exps
                valid_acc_dir = next(os.walk( config['atlas_exps'] ))[1]
                # intersection of the two lists
                pass_ACC = list(set(ACC) & set(valid_acc_dir))
                # filter E-PROT-* as no recalculations are needed for these. 
                # In any case, Snakemake-recalculations will exit with error for proteomics experiments
                nonprot_ACC = list(filter(lambda k: not 'E-PROT' in k, pass_ACC))
                return expand("{accession}/{accession}.recalculations.done", accession=nonprot_ACC)
            else:
                # filter E-PROT-* as no recalculations are needed for these. 
                # In any case, Snakemake-recalculations will exit with error for proteomics experiments
                nonprot_ACC = list(filter(lambda k: not 'E-PROT' in k, ACC))
                return expand("{accession}/{accession}.recalculations.done", accession=nonprot_ACC)
        elif config['goal'] == 'reprocess':
            # filter accessions that should not be reprocessed
            skip_accessions_reprocessing = read_skip_accessions_file('reprocess')
            if skip_accessions_reprocessing != None:
                print(f"Skipping reprocessing for the following accessions: {skip_accessions_reprocessing}")
                pass_ACC = [x for x in ACC if x not in skip_accessions_reprocessing ]
            if pass_ACC:
                return expand("{accession}/{accession}.reprocess.done", accession=pass_ACC )
            else:
                sys.exit("No accessions to reprocess. Exiting...")
        else:
            return None
    else:
        return None


localrules: all

wildcard_constraints:
    accession = "E-\D+-\d+"

rule all:
    input:
        required_outputs=get_outputs()


rule get_experiment_metadata:
    conda:
        "envs/perl-atlas-modules.yaml"
    log: "{accession}/logs/get_experiment_metadata.log"
    params:
        atlas_meta_config = config['atlas_meta_config'],
        atlas_prod = config['atlas_prod'],
        atlas_exps = config['atlas_exps'],
        check_species_file=config['check_sp_file'],
        prot_magetab_files=config['prot_magetabfiles']
    input:
        config_file = "{accession}/{accession}-configuration.xml",
        gtf_dir = config['gtf_dir']
    output:
        metadata = "{accession}/{accession}.metadata_summary.yaml"
    shell:
        """
        set -e 
        mkdir -p {wildcards.accession}/logs
        exec &> "{log}"
        export ATLAS_META_CONFIG="{params.atlas_meta_config}"
        export ATLAS_PROD="{params.atlas_prod}"
        export ATLAS_EXPS="{params.atlas_exps}"
        source "{workflow.basedir}/atlas-bash-util/generic_routines.sh"

        # proteomics baseline experiments do not need to include exp type assays/contrasts in the summary yaml
        prot_exp_type='experimentType="proteomics_baseline'
        if grep -q "$prot_exp_type" {input.config_file} ; then
            rm -f {output.metadata}
            echo "---" >> {output.metadata}
            echo "experiment_type: "$(grep "$prot_exp_type" {input.config_file}  | cut -d'"' -f 2) >> {output.metadata}
        else
            perl {workflow.basedir}/bin/get_exptype_contrasts_assays.pl {wildcards.accession} {input.config_file} {output.metadata}
        fi
        # we may avoid this adding proteomics_baseline_dia to AtlasSiteConfig.yml (perl-atlas-modules)

        if [ -e "{params.atlas_exps}/{wildcards.accession}/{wildcards.accession}.condensed-sdrf.tsv" ]; then
            echo "Trying with condensed in atlas_exps..."
            organism=$(get_organism_from_condensed_sdrf "{params.atlas_exps}/{wildcards.accession}/{wildcards.accession}.condensed-sdrf.tsv")
        elif [ -e "{wildcards.accession}/{wildcards.accession}.condensed-sdrf.tsv" ]; then
            echo "Trying with condensed in local experiment dir..."
            organism=$(get_organism_from_condensed_sdrf "{wildcards.accession}/{wildcards.accession}.condensed-sdrf.tsv")
        elif [ -e "{wildcards.accession}/{wildcards.accession}-configuration.xml" ]; then
            echo "Trying with perl (produces organism with capitals and spaces)"
            organism=$(perl {workflow.basedir}/bin/get_experiment_info.pl --experiment {wildcards.accession} --xmlfile {wildcards.accession}/{wildcards.accession}-configuration.xml --organism | sed 's/ /_/g' | tr '[:upper:]' '[:lower:]')
        else
            echo "Nothing worked...Exiting with error"
            >&2 echo "Can't retrieve organism: neither {wildcards.accession}.condensed-sdrf.tsv nor {wildcards.accession}/{wildcards.accession}-configuration.xml not found "
            exit 1
        fi
        # fix species name
        while read line; do
            echo "$line"
            DIR=$line
        done < <(perl {workflow.basedir}/bin/bulk_recalculations_check_species_names.pl $(echo $organism) {params.check_species_file}  )
        organism_check=$(echo $DIR)

        ##echo $organism "check-> "$organism_check >> organisms.txt
        organism=$organism_check

        gff_file=$( find {input.gtf_dir}/$organism -iname "$organism.*.gff3" | head -n1)
        if [ -z "$gff_file" ]; then
            gff_file=$( find {input.gtf_dir}/$organism -iname "$organism_*.gff3" | head -n1)
        fi
	    echo "GTF file location sourced"

        prot_exp_type_all='experimentType="proteomics_'
        # this applies to all proteomics experiments
        if grep -q "$prot_exp_type_all" {input.config_file} ; then
            idf={params.prot_magetab_files}/{wildcards.accession}/{wildcards.accession}.idf.txt
            sdrf={params.prot_magetab_files}/{wildcards.accession}/{wildcards.accession}.sdrf.txt
        else
            magetabfiles=$( perl {workflow.basedir}/bin/get_experiment_info.pl --experiment {wildcards.accession} --xmlfile {wildcards.accession}/{wildcards.accession}-configuration.xml --magetabfiles --not-strict)
            echo "MAGETAB location sourced"
            idf=$( echo $magetabfiles | awk -F',' '{{ print $1 }}' )
            sdrf=$( echo $magetabfiles | awk -F',' '{{ print $2 }}' )
        fi

        if [ -z "$gff_file" ] || [ -z "$idf" ] || [ -z "$sdrf" ]; then
            echo "GFF: $gff_file"
            echo "IDF: $idf"
            echo "SDRF: $sdrf"
            echo "GFF, IDF or SDRF not found. YAML not created...Exiting with error"
            exit 1
        else
            echo "organism: $organism" >> {output.metadata}
            echo "gff: $gff_file" >> {output.metadata}
            echo "idf: $idf" >> {output.metadata}
            echo "sdrf: $sdrf" >> {output.metadata}
        fi
        """

rule produce_reprocess_call:
    input:
        metadata = "{accession}/{accession}.metadata_summary.yaml",
        working_directory = "{accession}"
    log: "{accession}/logs/produce_reprocess_call.log"
    params:
        bioentities_properties=config['bioentities_properties'],
        skip_steps_file=get_skip_steps_file(),
        run_deconv_file=get_run_deconv_file(),
        methods_base=config['methods_base'],
        methods_dif=config['methods_dif'],
        zooma_exclusions=config['zooma_exclusions'],
        isl_dir=config['isl_dir'],
        isl_genomes=config['isl_genomes'],
        irap_versions=config['irap_versions'],
        irap_container=config['irap_container'],
        tmp_dir=config['tmp_dir'],
        atlas_exps = config['atlas_exps'],
        priv_stat_file = config['priv_stat_file'],
        cluster_config = get_cluster_config_file(),
        log_handler_script=get_log_handler_script(),
        goal=config['goal'],
        atlasmeta_config=get_meta_config(),
        allowed_organism=get_user_species,
        oracle_home=get_db_params('oracle_home'),
        python_user=get_db_params('python_user'),
        python_connect_string=get_db_params('python_connect_string'),
        python_password=get_db_params('python_password'),
        deconv_ref=config['deconv_ref']

    output:
        touch("{accession}/{accession}.reprocess.done")
    run:
        command_line=f"""
        set -e # snakemake on the cluster doesn't stop on error when --keep-going is set
        mkdir -p {wildcards.accession}/logs
        exec &> {log}
        pushd {input.working_directory}
        if [ -f {params.cluster_config} ]; then
            if [ ! -f "cluster.yaml" ]; then
                ln -sf {params.cluster_config} cluster.yaml
            fi
        else
            echo "cluster resource settings file does not exist"
            exit 1
        fi

        if [[ "{params.allowed_organism}" == "True" ]]; then

            snakemake --use-singularity --restart-times 1 --latency-wait 100 {config['sm_options']} --config \
                accession={wildcards.accession} \
                goal={params.goal} \
                tool={get_tool(wildcards.accession)} \
                metadata_summary=../{input.metadata} \
                bioentities_properties={params.bioentities_properties} \
                skip_steps_file={params.skip_steps_file} \
                run_deconv_file={params.run_deconv_file} \
                methods_base={params.methods_base} \
                methods_dif={params.methods_dif} \
                atlas_exps={params.atlas_exps} \
                priv_stat_file={params.priv_stat_file} \
                zooma_exclusions={params.zooma_exclusions} \
                isl_dir={params.isl_dir} \
                isl_genomes={params.isl_genomes} \
                irap_versions={params.irap_versions} \
                irap_container={params.irap_container} \
                tmp_dir={params.tmp_dir} \
                atlas_meta_config={params.atlasmeta_config} \
                oracle_home={params.oracle_home} \
                python_user={params.python_user} \
                python_connect_string={params.python_connect_string} \
                python_password={params.python_password} \
                deconv_ref={params.deconv_ref} \
                -s {workflow.basedir}/Snakefile-reprocess {params.log_handler_script}

        else
            echo "WARNING: Organism not allowed for reprocessing" 
        fi
        popd
        """
        #print(command_line)
        shell(command_line)


rule produce_recalculations_call:
    input:
        metadata = "{accession}/{accession}.metadata_summary.yaml",
        working_directory = "{accession}"
    log: "{accession}/logs/produce_recalculations_call.log"
    params:
        bioentities_properties=config['bioentities_properties'],
        skip_steps_file=get_skip_steps_file(),
        cluster_config = get_cluster_config_file(),
        run_deconv_file=get_run_deconv_file(),
        log_handler_script=get_log_handler_script(),
        goal=config['goal'],
        atlas_exps = config['atlas_exps'],
        priv_stat_file = config['priv_stat_file'],
        tmp_dir=config['tmp_dir'],
        allowed_organism=get_user_species,
        oracle_home=get_db_params('oracle_home'),
        python_user=get_db_params('python_user'),
        python_connect_string=get_db_params('python_connect_string'),
        python_password=get_db_params('python_password'),
        deconv_ref=config['deconv_ref'],
        irap_versions=config['irap_versions'],
        irap_container=config['irap_container'],
        methods_base=config['methods_base'],
        methods_dif=config['methods_dif'],
        isl_dir=config['isl_dir'],
        isl_genomes=config['isl_genomes']
	
    output:
        touch("{accession}/{accession}.recalculations.done")
    run:
        command_line=f"""
        set -e # snakemake on the cluster doesn't stop on error when --keep-going is set
        mkdir -p {wildcards.accession}/logs
        exec &> {log}
        pushd {input.working_directory}
        if [ -f {params.cluster_config} ]; then
            if [ ! -f "cluster.yaml" ]; then
                ln -sf {params.cluster_config} cluster.yaml
            fi
        else
            echo "cluster resource settings file does not exist"
            exit 1	    
        fi

        if [[ "{params.allowed_organism}" == "True" ]]; then

            snakemake --restart-times 4 --latency-wait 250 {config['sm_options']} --config \
                accession={wildcards.accession} \
                goal={params.goal} \
                tool={get_tool(wildcards.accession)} \
                metadata_summary=../{input.metadata} \
                bioentities_properties={params.bioentities_properties} \
                skip_steps_file={params.skip_steps_file} \
		run_deconv_file={params.run_deconv_file} \
		atlas_exps={params.atlas_exps} \
		priv_stat_file={params.priv_stat_file} \
	       	tmp_dir={params.tmp_dir} \
                oracle_home={params.oracle_home} \
                python_user={params.python_user} \
                python_connect_string={params.python_connect_string} \
                python_password={params.python_password} \
		deconv_ref={params.deconv_ref} \
		irap_versions={params.irap_versions} \
                irap_container={params.irap_container} \
		methods_base={params.methods_base} \
                methods_dif={params.methods_dif} \
		isl_dir={params.isl_dir} \
                isl_genomes={params.isl_genomes} \
                -s {workflow.basedir}/Snakefile-recalculations {params.log_handler_script}
                
        else
            echo "WARNING: Organism not allowed for reprocessing" 
        fi
        popd
        """
        #print(command_line)
        shell(command_line)