snakefile.annotate

#----------------------------------------------------
#  Setup
#----------------------------------------------------

configfile: workflow.basedir + "/Config/config.default.yml"

# import libraries
import os
import sys
import pandas as pd
from datetime import datetime

# import custom functions 
sys.path.insert(0, Path(workflow.basedir).parent.as_posix())
import Scripts.Functions.functions as func

now = datetime.now()
dt_string = now.strftime("%Y-%m-%d_%H-%M-%S")

# Get the names of query samples from the paths given in the query section of the config
samples = list(config['query_datasets'].keys())

# get consensus methods to run 
consensus_run = func.get_consensus_methods(config)

# get tools to run
tools_to_run = func.get_tools_to_run(config)
consensus_tools = func.get_consensus_tools(config)
print(consensus_tools)

# set defult parameters
func.set_reference_batch_parameters(config)
func.set_downsampling_parameters(config)
func.set_ontology_parameters(config, mode = 'annotation')
func.set_gene_conversion_parameters(config)
func.set_benchmark_directory(config, mode = 'annotation')
print(config)

#----------------------------------------------------
#  Final rule all 
#----------------------------------------------------

rule all:
  input:
    expand(config['output_dir'] + '/{sample}/report/{sample}.{consensus_type}.prediction_report.{reference}.' + dt_string + '.html',
    reference = config['references'],
    sample = samples,
    consensus_type = consensus_run)

#----------------------------------------------------
#  Preprocess
#----------------------------------------------------

rule preprocess:
  input:
    reference = lambda wildcards:config['references'][wildcards.reference]['expression'],
    query = config['query_datasets'].values()
  output:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    query = expand(config['output_dir'] + "/{sample}/{{reference}}/expression.csv", sample = samples),
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  log: 
    config['output_dir'] + "/model/{reference}/preprocess.log"
  params:
    basedir = {workflow.basedir},
    labfile = lambda wildcards:config['references'][wildcards.reference]['labels'],
    batch_path = lambda wildcards:config['references'][wildcards.reference]['batch'],
    out = config['output_dir'],
    convert_genes = lambda wildcards:config["references"][wildcards.reference]['convert_ref_mm_to_hg'],
    reference_name = "{reference}",
    query_names = samples,
    min_cells_per_cluster = lambda wildcards:config["references"][wildcards.reference]['min_cells_per_cluster'],
    downsample_value = lambda wildcards:config["references"][wildcards.reference]['downsample']['value'],
    downsample_per_class = lambda wildcards:config["references"][wildcards.reference]['downsample']['stratified'],
  shell:
    """
    Rscript {params.basedir}/Scripts/preprocess.R \
    {input.reference} \
    "{input.query}" \
    {params.out} \
    {params.convert_genes} \
    {params.labfile} \
    {params.reference_name} \
    "{params.query_names}" \
    {params.min_cells_per_cluster} \
    {params.downsample_value} \
    {params.downsample_per_class} \
    "{params.batch_path}" \
    &> {log} 
    """


#----------------------------------------------------
#  Ontology 
#----------------------------------------------------

rule ontology:
  input:
     labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
     ontology = config['output_dir'] + "/model/{reference}/ontology/ontology.csv"
  log:
     config['output_dir'] + "/model/{reference}/ontology.log"
  params:
     basedir = {workflow.basedir},
     out = config['output_dir'],
     reference_name = "{reference}",
     ontology_path = lambda wildcards:config["references"][wildcards.reference]["ontology"]["ontology_path"],
     ontology_column = lambda wildcards:config["references"][wildcards.reference]["ontology"]["ontology_column"]
  shell:
     """
      Rscript {params.basedir}/Scripts/ontology.R \
       {params.out} \
       {params.reference_name} \
       {input.labfile} \
       "{params.ontology_path}" \
       "{params.ontology_column}" \
      &> {log}
     """

#----------------------------------------------------
#  Consensus
#----------------------------------------------------

rule consensus:
  input:
    results = expand(config["output_dir"] + "/{{sample}}/{{reference}}/{tool}/{tool}_pred.csv",
                     tool=tools_to_run), 
    ontology = config['output_dir'] + "/model/{reference}/ontology/ontology.csv"
  output:
    prediction_summary = config['output_dir'] + "/{sample}/{reference}/{consensus_type}/Prediction_Summary_{ontology}.tsv"
  log: 
    config["output_dir"] + "/{sample}/{reference}/{consensus_type}/Gatherpreds.{ontology}.log"
  params:
    basedir = {workflow.basedir}, 
    tools = tools_to_run,
    consensus_tools = consensus_tools,
    consensus_type = "{consensus_type}",
    sample = config["output_dir"] + "/{sample}/{reference}/",
    min_agree = config["consensus"]["type"]["majority"]["min_agree"],
    metrics_file =  lambda wildcards:(f'{config["references"][wildcards.reference]["output_dir_benchmark"]}/'f'{wildcards.reference}/report/metrics_label.csv'),
    CAWPE_mode = config["consensus"]["type"]["CAWPE"]["mode"],
    alpha = config["consensus"]["type"]["CAWPE"]["alpha"],
    ontology_label = "{ontology}",
    accuracy_metric = config["consensus"]["type"]["CAWPE"]["accuracy_metric"]
  shell:
    """
    Rscript {params.basedir}/Scripts/calculate_consensus.R \
    {params.sample} \
    {output.prediction_summary} \
    "{params.tools}" \
    "{params.consensus_tools}" \
    "{params.consensus_type}" \
    "{params.min_agree}" \
    "{params.ontology_label}" \
    "{input.ontology}" \
    {params.metrics_file} \
    "{params.CAWPE_mode}" \
    "{params.alpha}" \
    "{params.accuracy_metric}" \
    &> {log}
    """  

#----------------------------------------------------
#  Knit Report 
#----------------------------------------------------

def ontology_function(wildcards):
    return [
        f"{config['output_dir']}/{wildcards.sample}/{wildcards.reference}/{wildcards.consensus_type}/Prediction_Summary_{ontology}.tsv"
        for ontology in config["references"][wildcards.reference]["ontology"]["ontology_column"]
    ]
    
rule knit_report:
  input: 
    pred = ontology_function,
    query = lambda wildcards:config['query_datasets'][wildcards.sample],
    ontology = config['output_dir'] + "/model/{reference}/ontology/ontology.csv"
  output: 
    report_path = config['output_dir'] + '/{sample}/report/{sample}.{consensus_type}.prediction_report.{reference}.' + dt_string + '.html'
  log:
    config['output_dir'] + "/{sample}/report/report.{consensus_type}.{reference}.log"
  params:
    basedir = {workflow.basedir},
    output_dir = config['output_dir'], 
    sample = "{sample}",
    tools = tools_to_run, 
    consensus_tools = consensus_tools,
    refs = "{reference}",
    ontology_columns = lambda wildcards:config["references"][wildcards.reference]["ontology"]["ontology_column"],
    marker_genes = config['marker_genes'],
    cons_type = "{consensus_type}"
  threads: 1
  resources:
  shell:
    """
    Rscript -e "rmarkdown::render(
            '{params.basedir}/Notebooks/annotate_report.Rmd',
            params = list(tools         = '{params.tools}',
                          cons_tools    = '{params.consensus_tools}',
                          refs          = '{params.refs}',
                          output_dir    = '{params.output_dir}',
                          consensus     = '{input.pred}',
                          cons_type     = '{params.cons_type}',
                          sample        = '{params.sample}',
                          marker_genes  = '{params.marker_genes}',
                          threads       = '{threads}',
                          query         = '{input.query}',
                          ontology_path = '{input.ontology}',
                          ontology_columns = '{params.ontology_columns}'),
            output_file = '{output.report_path}')" \
    &> {log}
    """

#----------------------------------------------------
#   SingleR 
#----------------------------------------------------

rule train_SingleR:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/SingleR/SingleR_model.Rda"
  params:
    basedir = {workflow.basedir}
  log: 
    config['output_dir'] + "/model/{reference}/SingleR/SingleR.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/SingleR/SingleR_train_benchmark.txt"
  threads: 
    config['SingleR']['threads']
  resources: 
  shell:
    """
    Rscript {params.basedir}/Scripts/SingleR/train_SingleR.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    &> {log}
    """

rule predict_SingleR:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/SingleR/SingleR_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/SingleR/SingleR_pred.csv"
  params:
    basedir = {workflow.basedir}
  log: 
    config['output_dir'] + "/{sample}/{reference}/SingleR/SingleR.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/SingleR/SingleR_predict_benchmark.txt"
  threads: 
    config['SingleR']['threads']
  resources: 
  shell:
    """
    Rscript {params.basedir}/Scripts/SingleR/predict_SingleR.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   scPred 
#----------------------------------------------------

rule train_scPred:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scPred_{classifier}/scPred_{classifier}.Rda"
  params:
    basedir = {workflow.basedir},
    classifier = "{classifier}"
  log: 
    config['output_dir'] + "/model/{reference}/scPred_{classifier}/scPred_{classifier}.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scPred_{classifier}/scPred_{classifier}_train_benchmark.txt"
  threads: 
    config['scPred']['threads']
  resources: 
  shell:
    """
    Rscript {params.basedir}/Scripts/scPred/train_scPred.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.classifier} \
    &> {log}
    """

rule predict_scPred:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scPred_{classifier}/scPred_{classifier}.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scPred_{classifier}/scPred_{classifier}_pred.csv"
  params:
    basedir = {workflow.basedir},
    classifier = "{classifier}"
  log: 
    config['output_dir'] + "/{sample}/{reference}/scPred_{classifier}/scPred_{classifier}.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scPred_{classifier}/scPred_{classifier}_predict_benchmark.txt"
  threads: 
    config['scPred']['threads']
  resources: 
  shell:
    """
    Rscript {params.basedir}/Scripts/scPred/predict_scPred.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    {params.classifier} \
    &> {log}
    """
    
rule gather_scPred:
  input:
    expand(config['output_dir'] + "/{{sample}}/{{reference}}/scPred_{classifier}/scPred_{classifier}_pred.csv",
    classifier = config['scPred']['classifier'])
    
#----------------------------------------------------
#   scClassify
#----------------------------------------------------

rule train_scClassify:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scClassify/scClassify_model.Rda"
  params:
    basedir = {workflow.basedir},
    topN = config['scClassify']['topN'], 
    weightsCal = config['scClassify']['weightsCal'],
    hopach_kmax = config['scClassify']['hopach_kmax']
  log:
    config['output_dir'] + "/model/{reference}/scClassify/scClassify.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scClassify/scClassify_train_benchmark.txt"
  threads: 
    config['scClassify']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/scClassify/train_scClassify.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.topN} \
    {params.weightsCal} \
    {params.hopach_kmax} \
    &> {log}
    """

rule predict_scClassify:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scClassify/scClassify_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scClassify/scClassify_pred.csv"
  params:
    basedir = {workflow.basedir},
    algorithm = config['scClassify']['algorithm'],
    similarity = config['scClassify']['similarity'],
    prob_threshold = config['scClassify']['prob_threshold'],
    cor_threshold_static = config['scClassify']['cor_threshold_static'],
    cor_threshold_high = config['scClassify']['cor_threshold_high']
  log:
    config['output_dir'] + "/{sample}/{reference}/scClassify/scClassify.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scClassify/scClassify_predict_benchmark.txt"
  threads: 
    config['scClassify']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/scClassify/predict_scClassify.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    {params.algorithm} \
    {params.similarity} \
    {params.prob_threshold} \
    {params.cor_threshold_static} \
    {params.cor_threshold_high} \
    &> {log}
    """

#----------------------------------------------------
#   SciBet
#----------------------------------------------------

rule train_SciBet:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/SciBet/SciBet_model.Rda"
  params:
    basedir = {workflow.basedir},
  log:
    config['output_dir'] + "/model/{reference}/SciBet/SciBet.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/SciBet/SciBet_train_benchmark.txt"
  threads: 
    config['SciBet']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/SciBet/train_SciBet.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    &> {log}
    """

rule predict_SciBet:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/SciBet/SciBet_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/SciBet/SciBet_pred.csv"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/SciBet/SciBet.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/SciBet/SciBet_predict_benchmark.txt"
  threads: 
    config['SciBet']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/SciBet/predict_SciBet.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   scHPL
#----------------------------------------------------

rule train_scHPL:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scHPL/scHPL_model.pkl"
  params:
    basedir = {workflow.basedir},
    classifier = config['scHPL']['classifier'],
    dimred = config['scHPL']['dimred']
  log:
    config['output_dir'] + "/model/{reference}/scHPL/scHPL.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scHPL/scHPL_train_benchmark.txt"
  threads: 
    config['scHPL']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/scHPL/train_scHPL.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {params.classifier} \
    {params.dimred} \
    &> {log}
    """

rule predict_scHPL:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scHPL/scHPL_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scHPL/scHPL_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['scHPL']['threshold']
  log:
    config['output_dir'] + "/{sample}/{reference}/scHPL/scHPL.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scHPL/scHPL_predict_benchmark.txt"
  threads: 
    config['scHPL']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/scHPL/predict_scHPL.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    &> {log}
    """

#----------------------------------------------------
#   SVM Linear
#----------------------------------------------------

rule train_SVMlinear:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/SVMlinear/SVMlinear_model.pkl"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/model/{reference}/SVMlinear/SVMlinear.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/SVMlinear/SVMlinear_train_benchmark.txt"
  threads: 
    config['SVMlinear']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/SVC/train_linearSVM.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    &> {log}
    """

rule predict_SVMlinear:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/SVMlinear/SVMlinear_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/SVMlinear/SVMlinear_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['SVMlinear']['threshold'],
    tool_name = config['SVMlinear']['classifier']
  log:
    config['output_dir'] + "/{sample}/{reference}/SVMlinear/SVMlinear.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/SVMlinear/SVMlinear_predict_benchmark.txt"
  threads: 
    config['SVMlinear']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/SVC/predict_SVM.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    {threads} \
    {params.tool_name} \
    &> {log}
    """

#----------------------------------------------------
#   SVC 
#----------------------------------------------------

rule train_SVC:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + "_model.pkl"
  params:
    basedir = {workflow.basedir},
    classifier = config['SVC']['classifier']
  log:
    config['output_dir'] + "/model/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + ".log"
  benchmark:
    config['output_dir'] + "/model/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + "_train_benchmark.txt"
  threads: 
    config['SVC']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/SVC/train_SVM.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {params.classifier} \
    {threads} \
    &> {log}
    """

rule predict_SVC:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + "_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + "_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['SVC']['threshold'],
    tool_name = config['SVC']['classifier']
  log:
    config['output_dir'] + "/{sample}/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + ".log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/SVC" + config['SVC']['classifier'] + "/SVC" + config['SVC']['classifier'] + "_predict_benchmark.txt"
  threads: 
    config['SVC']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/SVC/predict_SVM.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    {threads} \
    {params.tool_name} \
    &> {log}
    """

#----------------------------------------------------
#   singleCellNet
#----------------------------------------------------

rule train_singleCellNet:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/singleCellNet/singleCellNet_model.Rda"
  params:
    basedir = {workflow.basedir},
    nTrees = config['singleCellNet']['nTrees']
  log:
    config['output_dir'] + "/model/{reference}/singleCellNet/singleCellNet.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/singleCellNet/singleCellNet_train_benchmark.txt"
  threads: 
    config['singleCellNet']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/singleCellNet/train_singleCellNet.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.nTrees} \
    &> {log}
    """

rule predict_singleCellNet:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/singleCellNet/singleCellNet_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/singleCellNet/singleCellNet_pred.csv"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/singleCellNet/singleCellNet.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/singleCellNet/singleCellNet_predict_benchmark.txt"
  threads: 
    config['singleCellNet']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/singleCellNet/predict_singleCellNet.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   Correlation
#----------------------------------------------------

rule train_Correlation:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/Correlation/Correlation_model.Rda"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/model/{reference}/Correlation/Correlation.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/Correlation/Correlation_train_benchmark.txt"
  threads: 
    config['Correlation']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/Correlation/train_Correlation.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    &> {log}
    """

rule predict_Correlation:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/Correlation/Correlation_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/Correlation/Correlation_pred.csv"
  params:
      basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/Correlation/Correlation.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/Correlation/Correlation_predict_benchmark.txt"
  threads: 
    config['Correlation']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/Correlation/predict_Correlation.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   scLearn 
#----------------------------------------------------

rule train_scLearn: 
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scLearn/scLearn_model.Rda"
  params:
    basedir = {workflow.basedir},
    bootstrap_times = config['scLearn']['bootstrap_times']
  log:
    config['output_dir'] + "/model/{reference}/scLearn/scLearn.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scLearn/scLearn_train_benchmark.txt"
  threads: 
    config['scLearn']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/scLearn/train_scLearn.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.bootstrap_times} \
    &> {log}
    """

rule predict_scLearn:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scLearn/scLearn_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scLearn/scLearn_pred.csv"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/scLearn/scLearn.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scLearn/scLearn_train_benchmark.txt"
  threads: 
    config['scLearn']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/scLearn/predict_scLearn.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   ACTINN
#----------------------------------------------------

rule train_ACTINN:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/ACTINN/ACTINN_model.pkl"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/model/{reference}/ACTINN/ACTINN.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/ACTINN/ACTINN_predict_benchmark.txt"
  threads: 
    config['ACTINN']['threads']
  shell:
    """
    python {params.basedir}/Scripts/ACTINN/train_ACTINN.py \
           -trs {input.reference} \
           -trl {input.labfile} \
	         -mp {output.model} \
    &> {log}
    """

rule predict_ACTINN:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/ACTINN/ACTINN_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/ACTINN/ACTINN_pred.csv"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/ACTINN/ACTINN.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/ACTINN/ACTINN_predict_benchmark.txt"
  threads: 
    config['ACTINN']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/ACTINN/predict_ACTINN.py \
           -ts {input.query} \
           -mp {input.model} \
           -pp {output.pred} \
    &> {log}
    """

#----------------------------------------------------
#   scID 
#----------------------------------------------------

rule run_scID:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv",
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scID/scID_pred.csv"
  params:
    basedir = {workflow.basedir},
    estimate_weights_from_target = config['scID']['estimate_weights_from_target'],
    logFC = config['scID']['logFC']
  log:
    config['output_dir'] + "/{sample}/{reference}/scID/scID.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scID/scID_predict_benchmark.txt"
  threads: 
    config['scID']['threads']
  shell:
    """
    Rscript {params.basedir}/Scripts/scID/run_scID.R \
            {input.reference} \
	    {input.labfile} \
	    {input.query} \
	    {output.pred} \
	    {threads} \
      {params.estimate_weights_from_target} \
      {params.logFC} \
    &> {log}
    """

#----------------------------------------------------
#   scAnnotate
#----------------------------------------------------

rule run_scAnnotate:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv",
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scAnnotate/scAnnotate_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['scAnnotate']['threshold']
  log:
    config['output_dir'] + "/{sample}/{reference}/scAnnotate/scAnnotate.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scAnnotate/scAnnotate_predict_benchmark.txt"
  threads: 
    config['scAnnotate']['threads']
  shell:
    """
    Rscript {params.basedir}/Scripts/scAnnotate/run_scAnnotate.R \
            {input.reference} \
            {input.labfile} \
            {input.query} \
            {output.pred} \
            {threads} \
	    {params.threshold} \
    &> {log}
    """

#----------------------------------------------------
#   scNym
#----------------------------------------------------

rule run_scNym:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv",
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scNym/scNym_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['scNym']['threshold']
  log:
    config['output_dir'] + "/{sample}/{reference}/scNym/scNym.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scNym/scNym_predict_benchmark.txt"
  threads: 
    config['scNym']['threads']
  shell:
    """
    python {params.basedir}/Scripts/scNym/run_scNym.py \
            {input.reference} \
            {input.labfile} \
            {input.query} \
            {output.pred} \
            {params.threshold} \
    &> {log}
    """

#----------------------------------------------------
#   CellTypist 
#----------------------------------------------------

rule train_CellTypist:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/CellTypist/CellTypist_model.pkl"
  params:
    basedir = {workflow.basedir},
    feature_selection = config['CellTypist']['feature_selection']
  log:
    config['output_dir'] + "/model/{reference}/CellTypist/CellTypist.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/CellTypist/CellTypist_train_benchmark.txt"
  threads: 
    config['CellTypist']['threads']
  resources:
  shell:
    """
    python {params.basedir}/Scripts/CellTypist/train_CellTypist.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.feature_selection} \
    &> {log}
    """

rule predict_CellTypist:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/CellTypist/CellTypist_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/CellTypist/CellTypist_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['CellTypist']['threshold'],
    majority_voting = config['CellTypist']['majority_voting']
  log:
    config['output_dir'] + "/{sample}/{reference}/CellTypist/CellTypist.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/CellTypist/CellTypist_predict_benchmark.txt"
  threads: 
    config['CellTypist']['threads']
  resources:
  shell:
    """
    export CELLTYPIST_FOLDER='/root/.celltypist' && \
    python {params.basedir}/Scripts/CellTypist/predict_CellTypist.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    {threads} \
    {params.majority_voting} \
    &> {log}
    """

#----------------------------------------------------
#   Seurat
#----------------------------------------------------

rule train_Seurat:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/Seurat/Seurat_model.Rda"
  params:
    basedir = {workflow.basedir},
    nPC_computed = config['Seurat']['nPC_computed'],
    integration_method = config['Seurat']['integration_method']
  log:
    config['output_dir'] + "/model/{reference}/Seurat/Seurat.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/Seurat/Seurat_train_benchmark.txt"
  threads: 
    config['Seurat']['threads']
  resources:
  shell:
    """
    source /opt/conda/etc/profile.d/conda.sh && conda activate seurat5 && \
    Rscript {params.basedir}/Scripts/Seurat/train_Seurat.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.nPC_computed} \
    {params.integration_method} \
    &> {log} && \
    conda deactivate
    """

rule predict_Seurat:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/Seurat/Seurat_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/Seurat/Seurat_pred.csv"
  params:
    basedir = {workflow.basedir},
    nPC_used = config['Seurat']['nPC_used']
  log:
    config['output_dir'] + "/{sample}/{reference}/Seurat/Seurat.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/Seurat/Seurat_predict_benchmark.txt"
  threads: 
    config['Seurat']['threads']
  resources:
  shell:
    """
    source /opt/conda/etc/profile.d/conda.sh && conda activate seurat5 && \
    Rscript {params.basedir}/Scripts/Seurat/predict_Seurat.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    {params.nPC_used} \
    &> {log} && \
    conda deactivate
    """


#----------------------------------------------------
#   Symphony
#----------------------------------------------------

rule train_Symphony:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/Symphony/Symphony_model.Rda"
  params:
    basedir = {workflow.basedir},
    nPC_computed = config['Symphony']['nPC_computed']
  log:
    config['output_dir'] + "/model/{reference}/Symphony/Symphony.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/Symphony/Symphony_train_benchmark.txt"
  threads: 
    config['Symphony']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/Symphony/train_Symphony.R \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.nPC_computed} \
    &> {log}
    """

rule predict_Symphony:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/Symphony/Symphony_model.Rda"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/Symphony/Symphony_pred.csv"
  params:
    basedir = {workflow.basedir},
  log:
    config['output_dir'] + "/{sample}/{reference}/Symphony/Symphony.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/Symphony/Symphony_predict_benchmark.txt"
  threads: 
    config['Symphony']['threads']
  resources:
  shell:
    """
    Rscript {params.basedir}/Scripts/Symphony/predict_Symphony.R \
    {input.query} \
    {input.model} \
    {output.pred} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   scPoli
#----------------------------------------------------

rule train_scPoli:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scPoli/scPoli_model.pkl"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/model/{reference}/scPoli/scPoli.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scPoli/scPoli_train_benchmark.txt"
  shell:
    """
    python {params.basedir}/Scripts/scPoli/train_scPoli.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    &> {log}
    """

rule predict_scPoli:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scPoli/scPoli_model.pkl"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scPoli/scPoli_pred.csv"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/{sample}/{reference}/scPoli/scPoli.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scPoli/scPoli_predict_benchmark.txt"
  shell:
    """
    python {params.basedir}/Scripts/scPoli/predict_scPoli.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    &> {log}
    """

#----------------------------------------------------
#   scANVI
#----------------------------------------------------

rule train_scANVI:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/scANVI/model.pt"
  params:
    basedir = {workflow.basedir}
  log:
    config['output_dir'] + "/model/{reference}/scANVI/scANVI.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/scANVI/scANVI_train_benchmark.txt"
  shell:
    """
    python {params.basedir}/Scripts/scANVI/train_scANVI.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    &> {log}
    """

rule predict_scANVI:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/scANVI/model.pt"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/scANVI/scANVI_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['scANVI']['threshold']
  log:
    config['output_dir'] + "/{sample}/{reference}/scANVI/scANVI.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/scANVI/scANVI_predict_benchmark.txt"
  shell:
    """
    python {params.basedir}/Scripts/scANVI/predict_scANVI.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    &> {log}
    """

#----------------------------------------------------
#   CellBlast
#----------------------------------------------------

rule train_CellBlast:
  input:
    reference = config['output_dir'] + "/model/{reference}/expression.csv",
    labfile = config['output_dir'] + "/model/{reference}/downsampled_labels.csv"
  output:
    model = config['output_dir'] + "/model/{reference}/CellBlast/models/ref.h5ad"
  params:
    basedir = {workflow.basedir},
    n_models = config['CellBlast']['n_models']
  log:
    config['output_dir'] + "/model/{reference}/CellBlast/CellBlast.log"
  benchmark:
    config['output_dir'] + "/model/{reference}/CellBlast/CellBlast_train_benchmark.txt"
  threads: 
    config['CellBlast']['threads']
  shell:
    """
    python {params.basedir}/Scripts/CellBlast/train_CellBlast.py \
    {input.reference} \
    {input.labfile} \
    {output.model} \
    {threads} \
    {params.n_models} \
    &> {log}
    """

rule predict_CellBlast:
  input:
    query = config['output_dir'] + "/{sample}/{reference}/expression.csv",
    model = config['output_dir'] + "/model/{reference}/CellBlast/models/ref.h5ad"
  output:
    pred = config['output_dir'] + "/{sample}/{reference}/CellBlast/CellBlast_pred.csv"
  params:
    basedir = {workflow.basedir},
    threshold = config['CellBlast']['threshold']
  log:
    config['output_dir'] + "/{sample}/{reference}/CellBlast/CellBlast.log"
  benchmark:
    config['output_dir'] + "/{sample}/{reference}/CellBlast/CellBlast_predict_benchmark.txt"
  threads: 
    config['CellBlast']['threads']
  shell:
    """
    python {params.basedir}/Scripts/CellBlast/predict_CellBlast.py \
    {input.query} \
    {input.model} \
    {output.pred} \
    {params.threshold} \
    {threads} \
    &> {log}
    """

#----------------------------------------------------
#   The End 
#----------------------------------------------------