From b9911791fe97bae18dadacad589e826d19126072 Mon Sep 17 00:00:00 2001
From: Niels Dekker <N.Dekker@lumc.nl>
Date: Sun, 4 Sep 2022 13:35:36 +0200
Subject: [PATCH] ENH: Add NumberOfWorkUnits to ElastixMain, used by
 CreateComponent

With this commit, the GenerateData() member function of both `TransformixFilter` and `ElastixRegistrationMethod` pass their NumberOfWorkUnits to their corresponding ElastixMain/TransformixMain object, which passed this NumberOfWorkUnits to each component that is an `itk::ProcessObject`.
---
 Core/Kernel/elxElastixMain.cxx             | 10 ++++++++++
 Core/Kernel/elxElastixMain.h               |  9 +++++++++
 Core/Main/itkElastixRegistrationMethod.hxx |  1 +
 Core/Main/itkTransformixFilter.hxx         |  1 +
 4 files changed, 21 insertions(+)

diff --git a/Core/Kernel/elxElastixMain.cxx b/Core/Kernel/elxElastixMain.cxx
index 27a6d5b6c..8bb057dad 100644
--- a/Core/Kernel/elxElastixMain.cxx
+++ b/Core/Kernel/elxElastixMain.cxx
@@ -730,6 +730,16 @@ ElastixMain::CreateComponent(const ComponentDescriptionType & name)
     itkExceptionMacro(<< "The following component could not be created: " << name);
   }
 
+  if (m_NumberOfWorkUnits > 0)
+  {
+    const auto processObject = dynamic_cast<itk::ProcessObject *>(component.GetPointer());
+
+    if (processObject)
+    {
+      processObject->SetNumberOfWorkUnits(m_NumberOfWorkUnits);
+    }
+  }
+
   return component;
 
 } // end CreateComponent()
diff --git a/Core/Kernel/elxElastixMain.h b/Core/Kernel/elxElastixMain.h
index 0327b95f7..dfbc9d201 100644
--- a/Core/Kernel/elxElastixMain.h
+++ b/Core/Kernel/elxElastixMain.h
@@ -313,6 +313,13 @@ class ElastixMain : public itk::Object
   virtual ParameterMapType
   GetTransformParametersMap() const;
 
+  void
+  SetNumberOfWorkUnits(const itk::ThreadIdType newValue)
+  {
+    m_NumberOfWorkUnits = newValue;
+  }
+
+
 protected:
   ElastixMain();
   ~ElastixMain() override;
@@ -399,6 +406,8 @@ class ElastixMain : public itk::Object
   ElastixMain(const Self &) = delete;
   void
   operator=(const Self &) = delete;
+
+  itk::ThreadIdType m_NumberOfWorkUnits{ 0 };
 };
 
 } // end namespace elastix
diff --git a/Core/Main/itkElastixRegistrationMethod.hxx b/Core/Main/itkElastixRegistrationMethod.hxx
index 556ff2c5e..b0e002c0e 100644
--- a/Core/Main/itkElastixRegistrationMethod.hxx
+++ b/Core/Main/itkElastixRegistrationMethod.hxx
@@ -227,6 +227,7 @@ ElastixRegistrationMethod<TFixedImage, TMovingImage>::GenerateData()
 
     // Create new instance of ElastixMain
     elx::DefaultConstruct<ElastixMainType> elastixMain;
+    elastixMain.SetNumberOfWorkUnits(this->GetNumberOfWorkUnits());
 
     // Set elastix levels
     elastixMain.SetElastixLevel(i);
diff --git a/Core/Main/itkTransformixFilter.hxx b/Core/Main/itkTransformixFilter.hxx
index 2c9bb4ff4..843429038 100644
--- a/Core/Main/itkTransformixFilter.hxx
+++ b/Core/Main/itkTransformixFilter.hxx
@@ -154,6 +154,7 @@ TransformixFilter<TMovingImage>::GenerateData()
 
   // Instantiate transformix
   elx::DefaultConstruct<TransformixMainType> transformixMain;
+  transformixMain.SetNumberOfWorkUnits(this->GetNumberOfWorkUnits());
 
   // Setup transformix for warping input image if given
   DataObjectContainerPointer inputImageContainer = nullptr;