From 5ff5caa48c499066764369bf813b1db0dfea5e78 Mon Sep 17 00:00:00 2001
From: Matt Cieslak <mattcieslak@gmail.com>
Date: Thu, 10 Oct 2024 20:18:27 -0400
Subject: [PATCH] [FIX] Constrain threading for intramodal template (#825)

Change threading for intramodal template
---
 qsiprep/interfaces/ants.py                   | 9 +++++++--
 qsiprep/workflows/dwi/intramodal_template.py | 1 +
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/qsiprep/interfaces/ants.py b/qsiprep/interfaces/ants.py
index 7069570d..e2216b10 100644
--- a/qsiprep/interfaces/ants.py
+++ b/qsiprep/interfaces/ants.py
@@ -4,6 +4,7 @@
 import os.path as op
 
 import nibabel as nb
+from nipype.interfaces.ants.base import ANTSCommand, ANTSCommandInputSpec
 from nipype.interfaces.base import (
     BaseInterfaceInputSpec,
     CommandLine,
@@ -22,7 +23,7 @@
 
 
 # Step 1 from DSI Studio, importing DICOM files or nifti
-class MultivariateTemplateConstruction2InputSpec(CommandLineInputSpec):
+class MultivariateTemplateConstruction2InputSpec(ANTSCommandInputSpec):
     dimension = traits.Enum(2, 3, 4, default=3, usedefault=True, argstr="-d %d")
     input_file = File(desc="txt or csv file with images", exists=True, position=-1)
     input_images = InputMultiObject(
@@ -93,7 +94,7 @@ class MultivariateTemplateConstruction2OutputSpec(TraitedSpec):
     iteration_templates = OutputMultiObject(File(exists=True))
 
 
-class MultivariateTemplateConstruction2(CommandLine):
+class MultivariateTemplateConstruction2(ANTSCommand):
     input_spec = MultivariateTemplateConstruction2InputSpec
     output_spec = MultivariateTemplateConstruction2OutputSpec
     _cmd = "antsMultivariateTemplateConstruction2.sh "
@@ -136,6 +137,10 @@ def _list_outputs(self):
 
         return outputs
 
+    def _num_threads_update(self):
+        # Parallelization is controlled by ants pexec
+        self.inputs.environ.update({"ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS": "1", "NSLOTS": "1"})
+
 
 class ImageMathInputSpec(BaseInterfaceInputSpec):
     in_file = File(exists=True, mandatory=True, position=3, argstr="%s")
diff --git a/qsiprep/workflows/dwi/intramodal_template.py b/qsiprep/workflows/dwi/intramodal_template.py
index 2df7fa1a..ef664884 100644
--- a/qsiprep/workflows/dwi/intramodal_template.py
+++ b/qsiprep/workflows/dwi/intramodal_template.py
@@ -119,6 +119,7 @@ def init_intramodal_template_wf(
             dimension=3, iteration_limit=num_iterations, **runtime_opts
         ),
         name="ants_mvtc2",
+        n_procs=omp_nthreads,
     )
 
     workflow.connect([