From 6a4f6ac1f9da9d69e7fd47d6aadd23c07448f195 Mon Sep 17 00:00:00 2001
From: jscgh <j.caley@unsw.edu.au>
Date: Fri, 29 Nov 2024 12:24:32 +1100
Subject: [PATCH] Working with multiqc enabled

---
 main.nf                         | 14 +++++++-------
 modules/local/run_helixfold3.nf |  5 ++---
 workflows/helixfold3.nf         | 25 ++++++++-----------------
 3 files changed, 17 insertions(+), 27 deletions(-)

diff --git a/main.nf b/main.nf
index 017aa3a0..7608de18 100644
--- a/main.nf
+++ b/main.nf
@@ -27,8 +27,8 @@ if (params.mode.toLowerCase().split(",").contains("esmfold")) {
     include { PREPARE_ESMFOLD_DBS } from './subworkflows/local/prepare_esmfold_dbs'
     include { ESMFOLD             } from './workflows/esmfold'
 }
-if (params.mode == "helixfold3") {
-    include { PREPARE_HELIXFOLD3_DBS     } from './subworkflows/local/prepare_helixfold3_dbs'
+if (params.mode.toLowerCase().split(",").contains("helixfold3")) {
+    include { PREPARE_HELIXFOLD3_DBS    } from './subworkflows/local/prepare_helixfold3_dbs'
     include { HELIXFOLD3                } from './workflows/helixfold3'
 }
 
@@ -213,7 +213,7 @@ workflow NFCORE_PROTEINFOLD {
     //
     // WORKFLOW: Run helixfold3
     //
-    if(params.mode == "helixfold3") {
+    if(requested_modes.contains("helixfold3")) {
         //
         // SUBWORKFLOW: Prepare helixfold3 DBs
         //
@@ -266,10 +266,10 @@ workflow NFCORE_PROTEINFOLD {
             PREPARE_HELIXFOLD3_DBS.out.helixfold3_init_models,
             PREPARE_HELIXFOLD3_DBS.out.helixfold3_maxit_src
         )
-        ch_helixfold3_top_ranked_pdb = helixfold3.out.top_ranked_pdb
-        ch_multiqc                = ch_multiqc.mix(helixfold3.out.multiqc_report.collect())
-        ch_versions               = ch_versions.mix(helixfold3.out.versions)
-        ch_report_input           = ch_report_input.mix(helixfold3.out.pdb_msa)
+        ch_helixfold3_top_ranked_pdb    = HELIXFOLD3.out.top_ranked_pdb
+        ch_multiqc                      = ch_multiqc.mix(HELIXFOLD3.out.multiqc_report.collect())
+        ch_versions                     = ch_versions.mix(HELIXFOLD3.out.versions)
+        ch_report_input                 = ch_report_input.mix(HELIXFOLD3.out.pdb_msa)
     }
 
     //
diff --git a/modules/local/run_helixfold3.nf b/modules/local/run_helixfold3.nf
index f75629a0..8c941af6 100644
--- a/modules/local/run_helixfold3.nf
+++ b/modules/local/run_helixfold3.nf
@@ -30,9 +30,9 @@ process RUN_HELIXFOLD3 {
 
     output:
     path ("${fasta.baseName}*")
-    tuple val(meta), path ("${meta.id}_helixfold3.cif")     , emit: main_cif
-    tuple val(meta), path ("${meta.id}_helixfold3.pdb")     , emit: main_pdb
+    tuple val(meta), path ("${meta.id}_helixfold3.pdb")     , emit: top_ranked_pdb
     tuple val(meta), path ("ranked*pdb")                    , emit: pdb
+    tuple val(meta), path ("${meta.id}_helixfold3.cif")     , emit: main_cif
     tuple val(meta), path ("*_mqc.tsv")                     , emit: multiqc
     path "versions.yml", emit: versions
 
@@ -74,7 +74,6 @@ process RUN_HELIXFOLD3 {
         --model_name allatom_demo \
         --init_model "./init_models/HelixFold3-240814.pdparams" \
         --infer_times 4 \
-        --diff_batch_size 1 \
         --logging_level "ERROR" \
         --precision "bf16"
 
diff --git a/workflows/helixfold3.nf b/workflows/helixfold3.nf
index c24fc622..1d18e7bd 100644
--- a/workflows/helixfold3.nf
+++ b/workflows/helixfold3.nf
@@ -65,7 +65,7 @@ workflow HELIXFOLD3 {
         ch_helixfold3_maxit_src
     )
 
-    RUN_HELIXFOLD3_ALL_ATOM
+    RUN_HELIXFOLD3
         .out
         .multiqc
         .map { it[1] }
@@ -73,19 +73,9 @@ workflow HELIXFOLD3 {
         .map { [ [ "model": "helixfold3" ], it.flatten() ] }
         .set { ch_multiqc_report }
 
-    ch_pdb            = ch_pdb.mix(RUN_HELIXFOLD3_ALL_ATOM.out.pdb)
-    ch_top_ranked_pdb = ch_top_ranked_pdb.mix(RUN_HELIXFOLD3_ALL_ATOM.out.top_ranked_pdb)
-    ch_versions       = ch_versions.mix(RUN_HELIXFOLD3_ALL_ATOM.out.versions)
-
-    RUN_HELIXFOLD3_ALL_ATOM
-        .out
-        .pdb
-        .combine(ch_dummy_file)
-        .map {
-            it[0]["model"] = "helixfold3"
-            it
-        }
-        .set { ch_pdb_msa }
+    ch_pdb            = ch_pdb.mix(RUN_HELIXFOLD3.out.pdb)
+    ch_top_ranked_pdb = ch_top_ranked_pdb.mix(RUN_HELIXFOLD3.out.top_ranked_pdb)
+    ch_versions       = ch_versions.mix(RUN_HELIXFOLD3.out.versions)
 
     ch_top_ranked_pdb
         .map { [ it[0]["id"], it[0], it[1] ] }
@@ -100,9 +90,10 @@ workflow HELIXFOLD3 {
         .set { ch_pdb_msa }
 
     emit:
-    top_ranked_pdb = ch_top_ranked_pdb // channel: [ id, /path/to/*.pdb ]
-    multiqc_report = ch_multiqc_report // channel: /path/to/multiqc_report.html
-    versions       = ch_versions       // channel: [ path(versions.yml) ]
+    pdb_msa        = ch_pdb_msa          // channel: [ meta, /path/to/*.pdb, dummy_file ]
+    top_ranked_pdb = ch_top_ranked_pdb   // channel: [ id, /path/to/*.pdb ]
+    multiqc_report = ch_multiqc_report   // channel: /path/to/multiqc_report.html
+    versions       = ch_versions         // channel: [ path(versions.yml) ]
 }
 
 /*