From 8a7606ab303a0cab83072df3ae718fc8c88654fc Mon Sep 17 00:00:00 2001
From: Oscar Esteban <code@oscaresteban.es>
Date: Sun, 18 Aug 2024 17:32:50 +0200
Subject: [PATCH] enh: revise number of processes required by nodes (memory
 issues in DWI pipeline)

---
 mriqc/workflows/diffusion/base.py   | 24 ++++++++++++------------
 mriqc/workflows/diffusion/output.py |  5 +++++
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/mriqc/workflows/diffusion/base.py b/mriqc/workflows/diffusion/base.py
index 2155b74a1..0053dc0d8 100644
--- a/mriqc/workflows/diffusion/base.py
+++ b/mriqc/workflows/diffusion/base.py
@@ -131,21 +131,21 @@ def dmri_qc_workflow(name='dwiMRIQC'):
     get_lowb = pe.Node(
         ExtractOrientations(),
         name='get_lowb',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     # Generate B0 reference
     dwi_ref = pe.Node(
         RobustAverage(mc_method=None),
         name='dwi_ref',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     hmc_b0 = pe.Node(
         Volreg(args='-Fourier -twopass', zpad=4, outputtype='NIFTI_GZ'),
         name='hmc_b0',
         mem_gb=3.0,
-        n_procs=config.nipype.omp_nthreads,
+        n_procs=config.nipype.nprocs,
     )
 
     # Calculate brainmask
@@ -164,13 +164,13 @@ def dmri_qc_workflow(name='dwiMRIQC'):
     averages = pe.MapNode(
         WeightedStat(),
         name='averages',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
         iterfield=['in_weights'],
     )
     stddev = pe.MapNode(
         WeightedStat(stat='std'),
         name='stddev',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
         iterfield=['in_weights'],
     )
 
@@ -180,38 +180,38 @@ def dmri_qc_workflow(name='dwiMRIQC'):
             nthreads=config.nipype.omp_nthreads,
         ),
         name='dwidenoise',
-        n_procs=config.nipype.omp_nthreads,
+        n_procs=config.nipype.nprocs,
     )
     drift = pe.Node(
         CorrectSignalDrift(),
         name='drift',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
     sp_mask = pe.Node(
         SpikingVoxelsMask(),
         name='sp_mask',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     # Fit DTI/DKI model
     dwimodel = pe.Node(
         DiffusionModel(),
         name='dwimodel',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=config.nipype.nprocs,
     )
 
     # Calculate CC mask
     cc_mask = pe.Node(
         CCSegmentation(),
         name='cc_mask',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     # Run PIESNO noise estimation
     piesno = pe.Node(
         PIESNO(),
         name='piesno',
-        n_procs=max(1, config.nipype.omp_nthreads // 2),
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     # EPI to MNI registration
@@ -483,7 +483,7 @@ def hmc_workflow(name='dMRI_HMC'):
         Volreg(args='-Fourier -twopass', zpad=4, outputtype='NIFTI_GZ'),
         name='motion_correct',
         mem_gb=3.0,
-        n_procs=config.nipype.omp_nthreads,
+        n_procs=config.nipype.nprocs,
     )
 
     bvec_rot = pe.Node(RotateVectors(), name='bvec_rot')
diff --git a/mriqc/workflows/diffusion/output.py b/mriqc/workflows/diffusion/output.py
index 87d073a7c..43c68b38e 100644
--- a/mriqc/workflows/diffusion/output.py
+++ b/mriqc/workflows/diffusion/output.py
@@ -81,10 +81,12 @@ def init_dwi_report_wf(name='dwi_report_wf'):
     mosaic_fa = pe.Node(
         PlotMosaic(cmap='Greys_r'),
         name='mosaic_fa',
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
     mosaic_md = pe.Node(
         PlotMosaic(cmap='Greys_r'),
         name='mosaic_md',
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     mosaic_snr = pe.MapNode(
@@ -97,6 +99,7 @@ def init_dwi_report_wf(name='dwi_report_wf'):
         ),
         name='mosaic_snr',
         iterfield=['before', 'after'],
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     mosaic_noise = pe.MapNode(
@@ -106,6 +109,7 @@ def init_dwi_report_wf(name='dwi_report_wf'):
         ),
         name='mosaic_noise',
         iterfield=['in_file'],
+        n_procs=max(1, config.nipype.nprocs // 2),
     )
 
     if config.workflow.species.lower() in ('rat', 'mouse'):
@@ -187,6 +191,7 @@ def _gen_entity(inlist):
     plot_heatmap = pe.Node(
         DWIHeatmap(scalarmap_label='Shell-wise Fractional Anisotropy (FA)'),
         name='plot_heatmap',
+        n_procs=config.nipype.nprocs,
     )
     ds_report_hm = pe.Node(
         DerivativesDataSink(