MMSeqs GPU

conf/base.config

@@ -52,7 +52,7 @@ process {
     withLabel:process_high_memory {
         memory = { 200.GB * task.attempt }
     }
-    withLabled:process_gpu {
+    withLabel:process_gpu {
         accelerator = 1
     }
     withLabel:error_ignore {

main.nf

@@ -458,7 +458,9 @@ workflow NFCORE_PROTEINFOLD {
             PREPARE_BOLTZ_DBS.out.boltz_model,
             PREPARE_COLABFOLD_DBS.out.colabfold_db,
             PREPARE_COLABFOLD_DBS.out.uniref30,
-            params.boltz_use_msa_server
+            params.boltz_use_msa_server,
+            true,
+            params.colabfold_db_gpu
         )
         ch_multiqc                  = ch_multiqc.mix(BOLTZ.out.multiqc_report)
         ch_versions                 = ch_versions.mix(BOLTZ.out.versions)

modules/local/mmseqs_gpu/main.nf

@@ -0,0 +1,56 @@
+process PREPARE_MMSEQS_DB {
+    tag "$meta.id"
+    label "process_medium"
+
+    container "/srv/scratch/sbf/containers/mmseqs-gpu.sif"
+
+    input:
+    tuple val(meta), path(fasta)
+
+    output:
+    tuple val(meta), path("querydb")
+
+    script:
+    """
+    mkdir "querydb"
+    /app/mmseqs/bin/mmseqs createdb "${fasta}" ./querydb/${meta.id}
+    """
+}
+
+process SEARCH_MMSEQS_GPU {
+    tag "$meta.id"
+    label "process_medium"
+    label "process_gpu"
+
+    container "/srv/scratch/sbf/containers/mmseqs-gpu.sif"
+
+    input:
+    tuple val(meta), path(fasta), path(querydb)
+    path ("mmseqs-gpu")
+
+    output:
+    tuple val(meta), path ("colabfold_${meta.id}_hits.a3m")
+
+    script:
+    """
+    mkdir "msas"
+    /app/mmseqs/bin/mmseqs search \
+        --gpu 1 \
+        ./${querydb}/${meta.id} \
+        ./mmseqs-gpu/colabfold_envdb_202108_db \
+        ./msas/colabfold_${meta.id} \
+        ./tmp
+
+    # Convert to a3m files
+    mkdir "msa_inter"
+    mkdir "tmp_out"
+    /app/mmseqs/bin/mmseqs result2msa \
+        ./querydb/${meta.id} \
+        ./mmseqs-gpu/colabfold_envdb_202108_db \
+        ./msas/colabfold_${meta.id} \
+        ./msa_inter/colabfold_${meta.id}
+    /app/mmseqs/bin/mmseqs unpackdb \
+        ./msa_inter/colabfold_${meta.id} ./tmp_out/colabfold_${meta.id}
+    mv ./tmp_out/colabfold_${meta.id}/0 ./colabfold_${meta.id}_hits.a3m
+    """
+}

modules/local/run_boltz_mmseqs_gpu/main.nf

@@ -0,0 +1,73 @@
+/*
+ * Run Boltz
+ */
+process RUN_BOLTZ_MMSEQS_GPU {
+    tag "$meta.id"
+    label 'process_medium'
+    label 'process_gpu'
+
+    container "nf-core/proteinfold_boltz:dev"
+
+    input:
+    tuple val(meta), path(yaml_samplesheet)
+    path (files)
+    path ('boltz1_conf.ckpt')
+    path ('ccd.pkl')
+
+    output:
+    tuple val(meta), path ("boltz_results_*/processed/msa/*.npz")               , emit: msa
+    tuple val(meta), path ("boltz_results_*/processed/structures/*.npz")        , emit: structures
+    tuple val(meta), path ("boltz_results_*/predictions/*/confidence*.json")    , emit: confidence
+    tuple val(meta), path ("${meta.id}_plddt_mqc.tsv")                          , emit: multiqc
+    tuple val(meta), path ("*boltz.pdb")                                        , emit: pdb
+    tuple val(meta), path ("boltz_results_*/predictions/*/plddt_*model_0.npz")  , emit: plddt
+    tuple val(meta), path ("boltz_results_*/predictions/*/pae_*model_0.npz")    , emit: pae
+
+    path "versions.yml", emit: versions
+
+    when:
+    task.ext.when == null || task.ext.when
+
+    script:
+    // Exit if running this module with -profile conda / -profile mamba
+    if (workflow.profile.tokenize(',').intersect(['conda', 'mamba']).size() >= 1) {
+        error("Local RUN_BOLTZ module does not support Conda. Please use Docker / Singularity / Podman instead.")
+    }
+    def version = "0.4.1"
+    def args = task.ext.args ?: ''
+
+    """
+    boltz predict "${yaml_samplesheet}" --cache ./ --write_full_pae --output_format pdb
+    cp boltz_results_*/predictions/*/*.pdb ./${meta.id}_boltz.pdb
+
+    echo -e Atom_serial_number"\\t"Atom_name"\\t"Residue_name"\\t"Residue_sequence_number"\\t"pLDDT > ${meta.id}_plddt_mqc.tsv
+    awk '{print \$2"\\t"\$3"\\t"\$4"\\t"\$6"\\t"\$11}' boltz_results_*/predictions/*/*.pdb | grep -v 'N/A' | uniq >> ${meta.id}_plddt_mqc.tsv
+
+    cat <<-END_VERSIONS > versions.yml
+    "${task.process}":
+        boltz: $version
+    END_VERSIONS
+    """
+
+    stub:
+    def version = "0.4.1"
+    """
+    mkdir -p boltz_results_${meta.id}/processed/msa/
+    mkdir -p boltz_results_${meta.id}/processed/structures/
+    mkdir -p boltz_results_${meta.id}/predictions/${meta.id}/
+
+    touch ${meta.id}_boltz.pdb
+    touch boltz_results_${meta.id}/processed/msa/${meta.id}.npz
+    touch boltz_results_${meta.id}/processed/structures/${meta.id}.npz
+    touch boltz_results_${meta.id}/predictions/${meta.id}/confidence_${meta.id}.json
+    touch boltz_results_${meta.id}/predictions/${meta.id}/${meta.id}.pdb
+    touch boltz_results_${meta.id}/predictions/${meta.id}/plddt_${meta.id}_model_0.npz
+    touch boltz_results_${meta.id}/predictions/${meta.id}/pae_${meta.id}_model_0.npz
+    touch ${meta.id}_plddt_mqc.tsv
+
+    cat <<-END_VERSIONS > versions.yml
+    "${task.process}":
+        boltz: $version
+    END_VERSIONS
+    """
+}

modules/local/samplesheetutils/main.nf

@@ -0,0 +1,21 @@
+process SAMPLESHEET_BOLTZ_MSA {
+    tag "$meta.id"
+    label "process_medium"
+
+    container "/srv/scratch/sbf/containers/samplesheet-utils-1.3.sif"
+
+    input:
+    tuple val(meta), path(fasta), path(msa)
+
+    output:
+    tuple val(meta), path("samplesheet.yaml"), emit: formatted_yaml
+
+    script:
+    """
+    create-samplesheet \
+        --directory ./ \
+        --msa-dir ./ \
+        --output-file samplesheet.yaml \
+        --yaml
+    """
+}

workflows/boltz.nf

@@ -22,6 +22,9 @@ include { BOLTZ_FASTA } from '../modules/local/data_convertor/boltz_fasta'
 include { SPLIT_MSA } from '../modules/local/msa_manager/split_msa'
 include { MMSEQS_COLABFOLDSEARCH } from '../modules/local/mmseqs_colabfoldsearch'
 include { MULTIFASTA_TO_CSV      } from '../modules/local/multifasta_to_csv'
+include { PREPARE_MMSEQS_DB      } from '../modules/local/mmseqs_gpu'
+include { SEARCH_MMSEQS_GPU      } from '../modules/local/mmseqs_gpu'
+include { SAMPLESHEET_BOLTZ_MSA  } from '../modules/local/samplesheetutils/'
 //
 // SUBWORKFLOW: Consisting entirely of nf-core/modules
 //
@@ -51,6 +54,8 @@ workflow BOLTZ {
     ch_colabfold_db // channel: [ path(colabfold_db) ]
     ch_uniref30     // channel: [ path(uniref30) ]
     msa_server
+    mmseqs_gpu_msa // If true, run normal pipeline
+    ch_colabfold_db_gpu
 
     main:
     ch_samplesheet.join(
@@ -99,16 +104,38 @@ workflow BOLTZ {
         .set{ch_prepare_fasta}
     }
 
-    BOLTZ_FASTA(
-            ch_prepare_fasta
-        )
+    if (!mmseqs_gpu_msa) {
+        BOLTZ_FASTA(
+                ch_prepare_fasta
+            )
 
-    RUN_BOLTZ(
-        BOLTZ_FASTA.out.formatted_fasta.map{[it[0], it[1]]},
-        BOLTZ_FASTA.out.formatted_fasta.map{it[2]},
-        ch_boltz_model,
-        ch_boltz_ccd
-    )
+        RUN_BOLTZ(
+            BOLTZ_FASTA.out.formatted_fasta.map{[it[0], it[1]]},
+            BOLTZ_FASTA.out.formatted_fasta.map{it[2]},
+            ch_boltz_model,
+            ch_boltz_ccd
+        )
+    } else {
+        PREPARE_MMSEQS_DB(
+            ch_samplesheet
+        )
+        PREPARE_MMSEQS_DB.out.view()
+        ch_samplesheet.join(PREPARE_MMSEQS_DB.out).view()
+        SEARCH_MMSEQS_GPU(
+            ch_samplesheet.join(PREPARE_MMSEQS_DB.out),
+            ch_colabfold_db_gpu,
+        )
+        SEARCH_MMSEQS_GPU.out.view()
+        SAMPLESHEET_BOLTZ_MSA(
+            ch_samplesheet.join(SEARCH_MMSEQS_GPU.out)
+        )
+        RUN_BOLTZ(
+            SAMPLESHEET_BOLTZ_MSA.out,
+            [],
+            ch_boltz_model,
+            ch_boltz_ccd
+        )
+    }
 
     RUN_BOLTZ
         .out