Implement metagenomic mode

main.nf

@@ -44,6 +44,7 @@ include { DOWNLOAD_BAKTA_DB } from './modules/helpers.nf'
 include { SORF_EXTRA } from './modules/find_sorf_extra.nf'
 include { EXTEND_OR_GENERATE_AUXILIARY_DB } from './modules/generate_auxiliary_db.nf'
 include { EXTEND_ANNOTATIONS } from './modules/extend_annotations.nf'
+include { ASSEMBLE_MAGS } from './subworkflows/assemble_metagenomic_contigs.nf'
 
 
 /*
@@ -70,42 +71,57 @@ workflow {
         bakta_db = DOWNLOAD_BAKTA_DB(params.bakta_db_type)
     }
 
-    infiles = Channel.fromPath("${params.indir}/*${params.infile_extension}")
+    if (params.meta) {
 
-    infiles
+        // the MAG assembly process returns 
+        ch_asm = ASSEMBLE_MAGS("${params.indir}")
+        // TODO: refactor this. This channel is needed for correct combination of feature pickles (CDS, optional RNA)
+        // but it is redundant with the infiles channel
+    } else {
+        infiles = Channel.fromPath("${params.indir}/*${params.infile_extension}")
+        ch_asm = infiles.map { asm -> tuple(sampleIdFromName(asm.name), asm) }
+    }
+
+    ch_asm
         .combine(bakta_db)
         .set { infiles_and_bakta_db }
 
-    ch_asm = infiles.map { asm -> tuple(sampleIdFromName(asm.name), asm) }
+    // ch_asm.view() // DEBUG
+
 
     //-----------------------------
     // CDS prediction
     //-----------------------------
     cds_outputs = FIND_CDSS(infiles_and_bakta_db)
-    all_cds_outputs = cds_outputs.collect()
 
-    cds_pkl_list_ch = all_cds_outputs
-        .flatten()
-        .filter { it.name.endsWith('.pkl') }
-        .collect()
+    cds_outputs
+        .multiMap { sample_id, seqs_faa, cds_pkl ->
+            ch_cds_pkl: cds_pkl
+            ch_cds_faa: seqs_faa
+        }
+        .set { cds_outputs_ch }
 
-    ch_cds_pkl = cds_pkl_list_ch.flatten()
+    // cds_outputs_ch.ch_cds_pkl.view() // DEBUG
 
     //-----------------------------
     // Cluster + annotate
     //-----------------------------
-    CLUSTER_PROTEOME(cds_outputs)
+    CLUSTER_PROTEOME(cds_outputs_ch.ch_cds_faa)
     CLUSTER_PROTEOME.out.set { cluster_out_ch }  
 
-    // TODO: use multiMap (https://docs.seqera.io/nextflow/reference/operator#multimap)
-    rep_proteins_ch = cluster_out_ch.map { all_seqs, clustering_tsv, rep_seq -> rep_seq }
-    clustering_tsv_ch = cluster_out_ch.map { all_seqs, clustering_tsv, rep_seq -> clustering_tsv }
-    all_seqs_ch = cluster_out_ch.map { all_seqs, clustering_tsv, rep_seq -> all_seqs }   
+    cluster_out_ch
+        .multiMap { all_seqs, clustering_tsv, rep_seq ->
+            all_seqs_ch: all_seqs
+            clustering_tsv_ch: clustering_tsv
+            rep_proteins_ch: rep_seq
+        }
+        .set { cluster_out }
 
-    rep_proteins_ch
+    cluster_out.rep_proteins_ch
         .combine(bakta_db)
         .set { rep_proteins_and_dbs }
 
+    // optional expert systems that can be used for annotation
     hmm_ch = params.user_hmms ? Channel.of(file(params.user_hmms)) : []
     prots_ch = params.user_proteins ? Channel.of(file(params.user_proteins)) : []
 
@@ -137,8 +153,8 @@ workflow {
     // consider switching to explicit parameter definition in the config
     if( (params.mmseqs_args ?: '') != '--min-seq-id 1.0 -c 1.0 --alignment-mode 3' ) {
         EXTEND_ANNOTATIONS(
-            clustering_tsv_ch,
-            all_seqs_ch,
+            cluster_out.clustering_tsv_ch,
+            cluster_out.all_seqs_ch,
             bulk_annotations
         )
         bulk_ann_final_ch = EXTEND_ANNOTATIONS.out.bulk_annotations_extended
@@ -153,42 +169,47 @@ workflow {
     // NOTE: cache is not utilised if channel values are collected in a different order
     // TODO: sort collected values in cds_pkl_list_ch?
     MERGE_ANNOTATIONS(
-        cds_pkl_list_ch,
+        cds_outputs_ch.ch_cds_pkl.collect(),
         bulk_ann_final_ch
     )
 
     DETECT_PSEUDOGENES(MERGE_ANNOTATIONS.out.annotated_pickles, bakta_db, bakta_db_type)
     DETECT_PSEUDOGENES.out.annotated_samples_updated
         .set { ch_cds_annot_pkl }
 
+    ch_cds_keyed = ch_cds_annot_pkl.map { p -> tuple(sampleIdFromName(p.name), p) }
+
     //-----------------------------
     // RNA prediction
     //-----------------------------
-    rna_outputs = FIND_RNAS(infiles_and_bakta_db)
-
-    ch_rna_pkl = rna_outputs
-        .flatten()
-        .filter { it.name.endsWith('.pkl') }
+    if ( !params.skip_rna ) {
+        ch_rna_keyed = FIND_RNAS(infiles_and_bakta_db)
+        ch_sorf_in = ch_cds_keyed
+            .join(ch_rna_keyed)
+            .join(ch_asm)
+            .map { sid, cds_pkl, rna_pkl, asm -> tuple(sid, asm, [cds_pkl, rna_pkl]) }
+            .combine(bakta_db)
+    } else {
+        ch_sorf_in = ch_cds_keyed
+            .join(ch_asm)
+            .map { sid, cds_pkl, asm -> tuple(sid, asm, [cds_pkl]) }
+            .combine(bakta_db)
+
+        // ch_cds_keyed.view() // DEBUG
+        // SORF_EXTRA_NO_RNA(ch_sorf_in)
+    }
+
+    // ch_sorf_in.view() // DEBUG
 
     //-----------------------------
     // SORF extra search
     //-----------------------------
-    ch_cds_keyed = ch_cds_annot_pkl.map { p -> tuple(sampleIdFromName(p.name), p) }
-    ch_rna_keyed = ch_rna_pkl.map { p -> tuple(sampleIdFromName(p.name), p) }
-
-    ch_sorf_in = ch_cds_keyed
-        .join(ch_rna_keyed)
-        .join(ch_asm)
-        .map { sid, cds_pkl, rna_pkl, asm -> tuple(sid, asm, cds_pkl, rna_pkl) }
-        .combine(bakta_db)
-
     SORF_EXTRA(ch_sorf_in)
+    gff3_anno = SORF_EXTRA.out.gff3_annotations
 
     // TODO: refactor branching
     if ( params.auxiliary_db && (!file(params.auxiliary_db).exists() || params.extend_auxdb) ) {
-        SORF_EXTRA
-            .out
-            .gff3_annotations
+        gff3_anno
             .map { sample_id, anno_gff3, anno_pkl -> anno_pkl }
             .collect()
             .set { final_pkl_anno }

modules/assemble_metagenomic_contigs.nf

@@ -0,0 +1,56 @@
+process ASSEMBLE_CONTIGS {
+    tag "metagenomic"
+    label "metagenomic"
+    cpus { params.metaspades_threads }
+    memory { params.metaspades_mem_cap }
+    publishDir params.outdir, enabled: params.save_intermediate, mode: 'copy'
+
+    input:
+    tuple val(sample_id), path(sample_fwd), path(sample_rev) // paired-end reads assumbed by default
+
+    output:
+    tuple val(sample_id), path("megahit_out_${sample_id}/final.contigs.fa")
+
+    script:
+    // """
+    // metaspades.py -1 ${sample_fwd} -2 ${sample_rev} \
+    // -o metaspades_output \
+    // -t ${params.metaspades_threads} \
+    // -m ${params.metaspades_mem_cap}
+    // """
+    """
+    megahit -1 ${sample_fwd} -2 ${sample_rev} \
+    -o megahit_out_${sample_id} \
+    -t ${params.metaspades_threads} \
+    -m 17179869184
+    """
+}
+
+
+process ASSEMBLE_CONTIGS_METASPADES {
+    tag "metagenomic"
+    label "metagenomic"
+    cpus { params.metaspades_threads }
+    memory { params.metaspades_mem_cap }
+    publishDir params.outdir, enabled: params.save_intermediate, mode: 'copy'
+
+    input:
+    tuple val(sample_id), path(sample_fwd), path(sample_rev) // paired-end reads assumbed by default
+
+    output:
+    tuple val(sample_id), path("metaspades_out_${sample_id}/final.contigs.fa")
+
+    script:
+    // """
+    // metaspades.py -1 ${sample_fwd} -2 ${sample_rev} \
+    // -o metaspades_output \
+    // -t ${params.metaspades_threads} \
+    // -m ${params.metaspades_mem_cap}
+    // """
+    """
+    megahit -1 ${sample_fwd} -2 ${sample_rev} \
+    -o metaspades_out_${sample_id} \
+    -t ${params.metaspades_threads} \
+    -m 17179869184
+    """
+}

modules/find_cds.nf

@@ -5,17 +5,18 @@ process FIND_CDS {
     publishDir params.outdir, enabled: params.save_intermediate, mode: 'copy'
 
     input:
-    tuple path(assembly), path(bakta_db)
+    tuple val(sample_id), path(assembly), path(bakta_db)
 
     output:
-    tuple path("CDSS_bakta/${output_prefix}.cds-only.faa"), path("CDSS_bakta/${output_prefix}.cds-only.pkl")
+    tuple val(sample_id), path("CDSS_bakta/${sample_id}.cds-only.faa"), path("CDSS_bakta/${sample_id}.cds-only.pkl")
 
     script:
-    output_prefix = assembly.getBaseName()
+    meta = params.meta ? "--meta" : "" // Bakta (Pyrodigal) metagenome mode
     """
     bakta --db ${bakta_db} --cds-only  \
+    ${meta} \
     --output CDSS_bakta  \
-    --prefix ${output_prefix} \
+    --prefix ${sample_id} \
     --threads ${task.cpus} \
     ${assembly}
     """

modules/find_rnas.nf

@@ -5,13 +5,12 @@ process FIND_RNAS {
     publishDir params.outdir, enabled: params.save_intermediate, mode: 'copy'
 
     input:
-    tuple path(assembly), path(bakta_db)
+    tuple val(sample_id), path(assembly), path(bakta_db)
 
     output:
-    path("RNAS_bakta/${output_prefix}.rna-only.pkl")
+    tuple val(sample_id), path("RNAS_bakta/${sample_id}.rna-only.pkl")
 
     script:
-    output_prefix = assembly.getBaseName()
     """
     export TMPDIR=\$(mktemp -d)
     echo \$TMPDIR
@@ -25,7 +24,7 @@ process FIND_RNAS {
     bakta --db ${bakta_db} --rna-only  \
     --verbose \
     --output RNAS_bakta  \
-    --prefix ${output_prefix} \
+    --prefix ${sample_id} \
     --threads ${task.cpus} \
     ${assembly}
     """

modules/find_sorf_extra.nf

@@ -22,21 +22,67 @@ process SORF_EXTRA {
     )
 
     input:
-    tuple val(sample_id), path(assembly), path(cds_pkl), path(rna_pkl), path(bakta_db)
+    tuple val(sample_id), path(assembly), path(feature_pkls), path(bakta_db)
 
     output:
     tuple val(sample_id), path("${sample_id}/${sample_id}.sorf-extra.gff3"), path("${sample_id}/${sample_id}.sorf-extra.pkl"), emit: gff3_annotations
 
     script:
     def compliant = params.compliant ? "--compliant" : ""
+    def feature_pkls_param = feature_pkls.collect { "--feature-pickle ${it}" }.join(' ')
     """
     bakta --db ${bakta_db} --sorf-extra  \
-        --cds-pickle ${cds_pkl} \
-        --rna-pickle ${rna_pkl} \
+        ${feature_pkls_param} \
         --output ${sample_id}  \
         --prefix ${sample_id} \
         --threads ${task.cpus} \
         ${compliant} \
         ${assembly}
     """
-}
+}
+
+
+// // the same process but without the RNA features input
+// // TODO: thinks of a better design to implement optional input file
+// // maybe something like this would do: https://nextflow-io.github.io/patterns/optional-input/
+// process SORF_EXTRA_NO_RNA {
+//     tag { sample_id }
+//     label "sorf_extra_search"
+//     label 'bakta'
+
+//     // TODO: allow specifying output path for the gff3 file in Bakta to avoid doing this
+//     publishDir (
+//         params.outdir, 
+//         mode: 'copy', 
+//         saveAs: { filename -> 
+//             def name = file(filename).name
+//             if (name.endsWith(".sorf-extra.gff3")) {
+//                 def base = name.replaceFirst(/\.sorf-extra\.gff3$/, '')
+//                 return "${base}.gff3"
+//             }
+//             // // TODO: save final pickle objects 
+//             // else if (name.endsWith(".sorf-extra.pkl")) {
+//             //     return name
+//             // }
+//             return null
+//         }
+//     )
+
+//     input:
+//     tuple val(sample_id), path(assembly), path(cds_pkl), path(bakta_db)
+
+//     output:
+//     tuple val(sample_id), path("${sample_id}/${sample_id}.sorf-extra.gff3"), path("${sample_id}/${sample_id}.sorf-extra.pkl"), emit: gff3_annotations
+
+//     script:
+//     def compliant = params.compliant ? "--compliant" : ""
+//     """
+//     bakta --db ${bakta_db} --sorf-extra  \
+//         --feature-pickle ${cds_pkl} \
+//         --output ${sample_id}  \
+//         --prefix ${sample_id} \
+//         --threads ${task.cpus} \
+//         ${compliant} \
+//         ${assembly}
+//     """
+// }

nextflow.config

@@ -42,11 +42,25 @@ params {
 
     // TODO: add debug flag that sets save_intermediate to `true` and turns on extra logging
 
+    // metagenome mode
+    sample_1_suffix = '_1' // or 'fwd'
+    sample_2_suffix = '_2'
+
+    // metaSPAdes params
+    metaspades_threads = 16
+    metaspades_mem_cap = 32 // GB
+
     // Bakta settings
     bakta_db_dir = get_bakta_db_dir() // TODO: download bakta db if not found under this path
     bakta_db_type = "light" // `light|full`
     bakta_db = get_bakta_db_path()
 
+    // CDS search params
+    meta = false // assemble MAGs and run Bakta in metagenomic mode
+
+    // RNA prediction params
+    skip_rna = false
+
     bakta_args = "" // TODO: define parameters explicitly
     compliant = false
 

subworkflows/assemble_metagenomic_contigs.nf

@@ -0,0 +1,20 @@
+include { ASSEMBLE_CONTIGS } from '../modules/assemble_metagenomic_contigs.nf'
+
+workflow ASSEMBLE_MAGS {
+    take:
+    indir // tuple val(sample_id), path(assembly)
+
+    main:
+    reads_str = "${params.indir}/*{${params.sample_1_suffix},${params.sample_2_suffix}}*${params.infile_extension}"
+    println "${reads_str} DEBUG reads_str" 
+
+    read_pairs_ch = Channel.fromFilePairs(reads_str, checkIfExists: true)
+    read_pairs_ch.map { meta, reads -> [meta, reads[0], reads[1]] }
+        .set {read_pairs_ch}
+    // read_pairs_ch.view() // DEBUG
+
+    contigs = ASSEMBLE_CONTIGS(read_pairs_ch)
+
+    emit:
+    contigs
+}