Parallel execution for GetCurrentClusterFit

MuonColliderSoft · Oct 17, 2024 · 64560fb · 64560fb
1 parent 4e3fea9
commit 64560fb
Show file tree

Hide file tree

Showing 2 changed files with 99 additions and 2 deletions.
diff --git a/include/ParConeClusteringAlgorithm.h b/include/ParConeClusteringAlgorithm.h
@@ -22,6 +22,12 @@
 using lc_content::KDTreeLinkerAlgo;
 using lc_content::KDTreeNodeInfoT;
 
+#ifdef TBB_ENABLED
+#include "tbb/tbb.h"
+#include "Pandora/AlgorithmHeaders.h"
+#include "Pandora/StatusCodes.h"
+#endif
+
 /**
  *  @brief  ConeClusteringAlgorithm class
  */
@@ -236,6 +242,35 @@ class ParConeClusteringAlgorithm : public pandora::Algorithm
     float           m_mipTrackChi2Cut;              ///< Max value of fit chi2 for track seeded cluster to retain its IsMipTrack status
 
     unsigned int    m_firstLayer;                   ///< cache the pseudo layer at IP
+
+#ifdef TBB_ENABLED
+
+    mutable tbb::queuing_mutex p_mutex;
+
+    class CurrentClusterFit
+    {
+        using Cluster = pandora::Cluster;
+        using ClusterVector = std::vector<const Cluster*>;
+    public:
+        CurrentClusterFit(const ParConeClusteringAlgorithm& algo, const ClusterVector& c_vector, ClusterFitResultMap& c_fitMap) :
+            coneAlgorithm(algo),
+            clusterVector(c_vector),
+            clusterFitResultMap(c_fitMap),
+            p_result(pandora::STATUS_CODE_SUCCESS)
+        {}
+        virtual ~CurrentClusterFit() {}
+
+        void operator()(const tbb::blocked_range<std::size_t>& crange) const;
+
+    private:
+        const ParConeClusteringAlgorithm& coneAlgorithm;
+        const ClusterVector& clusterVector;
+        ClusterFitResultMap& clusterFitResultMap;
+        mutable pandora::StatusCode p_result;
+    };
+
+#endif
+
 };
 
 class ParConeClusteringAlgorithmFactory : public pandora::AlgorithmFactory

diff --git a/src/ParConeClusteringAlgorithm.cc b/src/ParConeClusteringAlgorithm.cc
@@ -198,7 +198,10 @@ StatusCode ParConeClusteringAlgorithm::GetCurrentClusterFitResults(const Cluster
 {
     if (!clusterFitResultMap.empty())
         return STATUS_CODE_INVALID_PARAMETER;
-
+#ifdef TBB_ENABLED
+    tbb::parallel_for(tbb::blocked_range<std::size_t>(0, clusterVector.size()),
+            CurrentClusterFit(*this, clusterVector, clusterFitResultMap));
+#else
     for (ClusterVector::const_iterator iter = clusterVector.begin(); iter != clusterVector.end(); ++iter)
     {
         const Cluster *const pCluster = *iter;
@@ -244,7 +247,7 @@ StatusCode ParConeClusteringAlgorithm::GetCurrentClusterFitResults(const Cluster
         if (!clusterFitResultMap.insert(ClusterFitResultMap::value_type(pCluster, clusterFitResult)).second)
             return STATUS_CODE_FAILURE;
     }
-
+#endif
     return STATUS_CODE_SUCCESS;
 }
 
@@ -947,3 +950,62 @@ StatusCode ParConeClusteringAlgorithm::ReadSettings(const TiXmlHandle xmlHandle)
 
     return STATUS_CODE_SUCCESS;
 }
+
+
+
+
+#ifdef TBB_ENABLED
+void ParConeClusteringAlgorithm::CurrentClusterFit::operator()(const tbb::blocked_range<std::size_t>& crange) const
+{
+    for (std::size_t idx = crange.begin(); idx != crange.end(); ++idx)
+    {
+        const Cluster *const pCluster = clusterVector[idx];
+        ClusterFitResult clusterFitResult;
+
+        if (pCluster->GetNCaloHits() > 1)
+        {
+            const unsigned int innerLayer(pCluster->GetInnerPseudoLayer());
+            const unsigned int outerLayer(pCluster->GetOuterPseudoLayer());
+            const unsigned int nLayersSpanned(outerLayer - innerLayer);
+
+            if (nLayersSpanned > coneAlgorithm.m_nLayersSpannedForFit)
+            {
+                unsigned int nLayersToFit(coneAlgorithm.m_nLayersToFit);
+
+                if (pCluster->GetMipFraction() - coneAlgorithm.m_nLayersToFitLowMipCut < std::numeric_limits<float>::epsilon())
+                    nLayersToFit *= coneAlgorithm.m_nLayersToFitLowMipMultiplier;
+
+                const unsigned int startLayer( (nLayersSpanned > nLayersToFit) ? (outerLayer - nLayersToFit) : innerLayer);
+                (void) ClusterFitHelper::FitLayerCentroids(pCluster, startLayer, outerLayer, clusterFitResult);
+
+                if (clusterFitResult.IsFitSuccessful())
+                {
+                    const float dotProduct(clusterFitResult.GetDirection().GetDotProduct(pCluster->GetInitialDirection()));
+                    const float chi2(clusterFitResult.GetChi2());
+
+                    if (((dotProduct < coneAlgorithm.m_fitSuccessDotProductCut1) &&
+                            (chi2 > coneAlgorithm.m_fitSuccessChi2Cut1)) ||
+                        ((dotProduct < coneAlgorithm.m_fitSuccessDotProductCut2) &&
+                                (chi2 > coneAlgorithm.m_fitSuccessChi2Cut2)) )
+                    {
+                        clusterFitResult.SetSuccessFlag(false);
+                    }
+                }
+            }
+            else if (nLayersSpanned > coneAlgorithm.m_nLayersSpannedForApproxFit)
+            {
+                const CartesianVector centroidChange(pCluster->GetCentroid(outerLayer) - pCluster->GetCentroid(innerLayer));
+                clusterFitResult.Reset();
+                clusterFitResult.SetDirection(centroidChange.GetUnitVector());
+                clusterFitResult.SetSuccessFlag(true);
+            }
+        }
+
+        {
+            tbb::queuing_mutex::scoped_lock lock(coneAlgorithm.p_mutex);
+            if (!clusterFitResultMap.insert(ClusterFitResultMap::value_type(pCluster, clusterFitResult)).second)
+                p_result = STATUS_CODE_FAILURE;
+        }
+    }
+}
+#endif