Update TrackClusterMergeSplitter to output track-cluster associations (PFA0)

src/algorithms/calorimetry/TrackClusterMergeSplitter.h

@@ -5,13 +5,24 @@
 
 #include <DD4hep/Detector.h>
 #include <algorithms/algorithm.h>
-#include <edm4eic/ProtoClusterCollection.h>
+#include <edm4eic/CalorimeterHit.h>
+#include <edm4eic/ClusterCollection.h>
+#include <edm4eic/EDM4eicVersion.h>
+#include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
+#if EDM4EIC_VERSION_MAJOR >= 8
+#include <edm4eic/TrackClusterMatchCollection.h>
+#endif
+#include <edm4eic/Track.h>
 #include <edm4eic/TrackPoint.h>
 #include <edm4eic/TrackSegmentCollection.h>
-#include <edm4hep/Vector3f.h>
-#include <podio/ObjectID.h>
+#if EDM4EIC_VERSION_MAJOR >= 7
+#include <edm4eic/MCRecoCalorimeterHitAssociationCollection.h>
+#else
+#include <edm4hep/SimCalorimeterHitCollection.h>
+#endif
 #include <algorithm>
 #include <map>
+#include <optional>
 #include <set>
 #include <string>
 #include <string_view>
@@ -26,14 +37,14 @@
 namespace eicrecon {
 
   // --------------------------------------------------------------------------
-  //! Comparator struct for protoclusters
+  //! Comparator struct for object IDs
   // --------------------------------------------------------------------------
-  /*! Organizes protoclusters by their ObjectID's in decreasing collection
+  /*! Organizes objects by their ObjectID's in decreasing collection
    *  ID first, and second by decreasing index second.
    */
-  struct CompareProto {
+  template <typename T> struct CompareObjectID {
 
-    bool operator() (const edm4eic::ProtoCluster& lhs, const edm4eic::ProtoCluster& rhs) const {
+    bool operator() (const T& lhs, const T& rhs) const {
       if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
         return (lhs.getObjectID().index < rhs.getObjectID().index);
       } else {
@@ -43,16 +54,24 @@ namespace eicrecon {
 
   };  // end CompareObjectID
 
+  // specialization for clusters, hits
+  using CompareClust = CompareObjectID<edm4eic::Cluster>;
+  using CompareHit = CompareObjectID<edm4eic::CalorimeterHit>;
+
 
 
   // --------------------------------------------------------------------------
   //! Convenience types
   // --------------------------------------------------------------------------
+  using VecTrk = std::vector<edm4eic::Track>;
   using VecProj = std::vector<edm4eic::TrackPoint>;
-  using VecClust = std::vector<edm4eic::ProtoCluster>;
-  using SetClust = std::set<edm4eic::ProtoCluster, CompareProto>;
-  using MapToVecProj = std::map<edm4eic::ProtoCluster, VecProj, CompareProto>;
-  using MapToVecClust = std::map<edm4eic::ProtoCluster, VecClust, CompareProto>;
+  using VecClust = std::vector<edm4eic::Cluster>;
+  using SetClust = std::set<edm4eic::Cluster, CompareClust>;
+  using MapToVecTrk = std::map<edm4eic::Cluster, VecTrk, CompareClust>;
+  using MapToVecProj = std::map<edm4eic::Cluster, VecProj, CompareClust>;
+  using MapToVecClust = std::map<edm4eic::Cluster, VecClust, CompareClust>;
+  using MapToWeight = std::map<edm4eic::CalorimeterHit, double, CompareHit>;
+  using VecWeights = std::vector<MapToWeight>;
 
 
 
@@ -61,11 +80,21 @@ namespace eicrecon {
   // --------------------------------------------------------------------------
   using TrackClusterMergeSplitterAlgorithm = algorithms::Algorithm<
     algorithms::Input<
-      edm4eic::ProtoClusterCollection,
-      edm4eic::TrackSegmentCollection
+      edm4eic::ClusterCollection,
+      edm4eic::TrackSegmentCollection,
+      std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>,
+#if EDM4EIC_VERSION_MAJOR >= 7
+      std::optional<edm4eic::MCRecoCalorimeterHitAssociationCollection>
+#else
+      std::optional<edm4hep::SimCalorimeterHitCollection>
+#endif
     >,
     algorithms::Output<
-      edm4eic::ProtoClusterCollection
+      edm4eic::ClusterCollection,
+#if EDM4EIC_VERSION_MAJOR >= 8
+      edm4eic::TrackClusterMatchCollection,
+#endif
+      std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>
     >
   >;
 
@@ -74,8 +103,8 @@ namespace eicrecon {
   // --------------------------------------------------------------------------
   //! Track-Based Cluster Merger/Splitter
   // --------------------------------------------------------------------------
-  /*! An algorithm which takes a collection of proto-clusters, matches
-   *  track projections, and then decides to merge or split those proto-
+  /*! An algorithm which takes a collection of clusters, matches
+   *  track projections, and then decides to merge or split those
    *  clusters based on average E/p from simulations.
    *
    *  Heavily inspired by Eur. Phys. J. C (2017) 77:466
@@ -91,8 +120,16 @@ namespace eicrecon {
       TrackClusterMergeSplitter(std::string_view name) :
         TrackClusterMergeSplitterAlgorithm {
           name,
-          {"InputProtoClusterCollection", "InputTrackProjections"},
-          {"OutputProtoClusterCollection"},
+#if EDM4EIC_VERSION_MAJOR >= 7
+          {"InputClusterCollection", "InputTrackProjections", "InputMCClusterAssociations", "InputMCHitAssociations"},
+#else
+          {"InputClusterCollection", "InputTrackProjections", "InputMCClusterAssociations", "InputMCHits"},
+#endif
+#if EDM4EIC_VERSION_MAJOR >= 8
+          {"OutputProtoClusterCollection", "OutputTrackClusterMatches", "OutputMCClusterAssociations"},
+#else
+          {"OutputProtoClusterCollection", "OutputMCClusterAssociations"},
+#endif
           "Merges or splits clusters based on tracks projected to them."
         } {}
 
@@ -103,11 +140,15 @@ namespace eicrecon {
     private:
 
       // private methods
-      void get_projections(const edm4eic::TrackSegmentCollection* projections, VecProj& relevant_projects) const;
-      void match_clusters_to_tracks(const edm4eic::ProtoClusterCollection* clusters, const VecProj& projections, MapToVecProj& matches) const;
-      void merge_and_split_clusters(const VecClust& to_merge, const VecProj& to_split, edm4eic::ProtoClusterCollection* out_protoclusters) const;
-      float get_cluster_energy(const edm4eic::ProtoCluster& clust) const;
-      edm4hep::Vector3f get_cluster_position(const edm4eic::ProtoCluster& clust) const;
+      void get_projections(const edm4eic::TrackSegmentCollection* projections, VecProj& relevant_projects, VecTrk& relevant_trks) const;
+      void match_clusters_to_tracks(const edm4eic::ClusterCollection* clusters, const VecProj& projections, const VecTrk& tracks, MapToVecProj& matched_projects, MapToVecTrk& matched_tracks) const;
+      void merge_and_split_clusters(const VecClust& to_merge, const VecProj& to_split, std::vector<edm4eic::MutableCluster>& new_clusters) const;
+      void make_cluster(const VecClust& old_clusts, edm4eic::MutableCluster& new_clust, std::optional<MapToWeight> split_weights = std::nullopt) const;
+#if EDM4EIC_VERSION_MAJOR >= 7
+      void collect_associations(const edm4eic::MutableCluster& new_clust, const VecClust& old_clusts, const edm4eic::MCRecoClusterParticleAssociationCollection* old_clust_assocs, const edm4eic::MCRecoCalorimeterHitAssociationCollection* old_hit_assocs, edm4eic::MCRecoClusterParticleAssociationCollection* new_clust_assocs) const;
+#else
+      void collect_associations(const edm4eic::MutableCluster& new_clust, const VecClust& old_clusts, const edm4eic::MCRecoClusterParticleAssociationCollection* old_clust_assocs, const edm4hep::SimCalorimeterHitCollection* old_sim_hits, edm4eic::MCRecoClusterParticleAssociationCollection* new_clust_assocs) const;
+#endif
 
       // calorimeter id
       int m_idCalo {0};

src/detectors/BHCAL/BHCAL.cc

@@ -194,10 +194,20 @@ extern "C" {
 
         app->Add(
           new JOmniFactoryGeneratorT<TrackClusterMergeSplitter_factory>(
-            "HcalBarrelSplitMergeProtoClusters",
-            {"HcalBarrelIslandProtoClusters",
-             "CalorimeterTrackProjections"},
-            {"HcalBarrelSplitMergeProtoClusters"},
+            "HcalBarrelSplitMergeClustersWithoutShapes",
+            {"HcalBarrelClusters",
+             "CalorimeterTrackProjections",
+             "HcalBarrelClusterAssociations",
+#if EDM4EIC_VERSION_MAJOR >= 7
+             "HcalBarrelRawHitAssociations"},
+#else
+             "HcalBarrelHits"},
+#endif
+            {"HcalBarrelSplitMergeClustersWithoutShapes",
+#if EDM4EIC_VERSION_MAJOR >= 8
+             "HcalBarrelTrackSplitMergeClusterMatches",
+#endif
+             "HcalBarrelSplitMergeClusterAssociationsWithoutShapes"},
             {
               .idCalo = "HcalBarrel_ID",
               .minSigCut = -2.0,
@@ -211,34 +221,19 @@ extern "C" {
           )
         );
 
-        app->Add(
-          new JOmniFactoryGeneratorT<CalorimeterClusterRecoCoG_factory>(
-             "HcalBarrelSplitMergeClustersWithoutShapes",
-            {"HcalBarrelSplitMergeProtoClusters", // edm4eic::ProtoClusterCollection
-             "HcalBarrelHits"},                          // edm4hep::SimCalorimeterHitCollection
-            {"HcalBarrelSplitMergeClustersWithoutShapes", // edm4eic::Cluster
-             "HcalBarrelSplitMergeClusterAssociationsWithoutShapes"}, // edm4eic::MCRecoClusterParticleAssociation
-            {
-              .energyWeight = "log",
-              .sampFrac = 1.0,
-              .logWeightBase = 6.2,
-              .enableEtaBounds = false
-            },
-            app   // TODO: Remove me once fixed
-          )
-        );
-
         app->Add(
           new JOmniFactoryGeneratorT<CalorimeterClusterShape_factory>(
             "HcalBarrelSplitMergeClusters",
             {"HcalBarrelSplitMergeClustersWithoutShapes",
              "HcalBarrelSplitMergeClusterAssociationsWithoutShapes"},
             {"HcalBarrelSplitMergeClusters",
              "HcalBarrelSplitMergeClusterAssociations"},
-            {},
+            {
+              .energyWeight = "log",
+              .logWeightBase = 6.2
+            },
             app
           )
         );
-
     }
 }

src/detectors/EEMC/EEMC.cc

@@ -185,25 +185,6 @@ extern "C" {
           )
         );
 
-        app->Add(
-          new JOmniFactoryGeneratorT<TrackClusterMergeSplitter_factory>(
-            "EcalEndcapNSplitMergeProtoClusters",
-            {"EcalEndcapNIslandProtoClusters",
-             "CalorimeterTrackProjections"},
-            {"EcalEndcapNSplitMergeProtoClusters"},
-            {
-              .idCalo = "EcalEndcapN_ID",
-              .minSigCut = -1.0,
-              .avgEP = 1.0,
-              .sigEP = 0.10,
-              .drAdd = 0.08,
-              .sampFrac = 1.0,
-              .transverseEnergyProfileScale = 1.0
-            },
-            app   // TODO: remove me once fixed
-          )
-        );
-
 #if EDM4EIC_VERSION_MAJOR >= 8
         app->Add(new JOmniFactoryGeneratorT<CalorimeterParticleIDPreML_factory>(
             "EcalEndcapNParticleIDPreML",
@@ -248,19 +229,31 @@ extern "C" {
 #endif
 
         app->Add(
-          new JOmniFactoryGeneratorT<CalorimeterClusterRecoCoG_factory>(
-             "EcalEndcapNSplitMergeClustersWithoutShapes",
-            {"EcalEndcapNSplitMergeProtoClusters", // edm4eic::ProtoClusterCollection
-             "EcalEndcapNHits"}, // edm4hep::SimCalorimeterHitCollection
-            {"EcalEndcapNSplitMergeClustersWithoutShapes", // edm4eic::Cluster
-             "EcalEndcapNSplitMergeClusterAssociationsWithoutShapes"}, // edm4eic::MCRecoClusterParticleAssociation
+          new JOmniFactoryGeneratorT<TrackClusterMergeSplitter_factory>(
+            "EcalEndcapNSplitMergeClustersWithoutShapes",
+            {"EcalEndcapNClustersWithoutPID",
+             "CalorimeterTrackProjections",
+             "EcalEndcapNClusterAssociations",
+#if EDM4EIC_VERSION_MAJOR >= 7
+             "EcalEndcapNRawHitAssociations"},
+#else
+             "EcalEndcapNHits"},
+#endif
+            {"EcalEndcapNSplitMergeClustersWithoutShapes",
+#if EDM4EIC_VERSION_MAJOR >= 8
+             "EcalEndcapNTrackSplitMergeClusterMatches",
+#endif
+             "EcalEndcapNSplitMergeClusterAssociationsWithoutShapes"},
             {
-              .energyWeight = "log",
+              .idCalo = "EcalEndcapN_ID",
+              .minSigCut = -1.0,
+              .avgEP = 1.0,
+              .sigEP = 0.10,
+              .drAdd = 0.08,
               .sampFrac = 1.0,
-              .logWeightBase = 3.6,
-              .enableEtaBounds = false
+              .transverseEnergyProfileScale = 1.0
             },
-            app   // TODO: Remove me once fixed
+            app   // TODO: remove me once fixed
           )
         );
 

src/detectors/EHCAL/EHCAL.cc

@@ -163,10 +163,20 @@ extern "C" {
         );
         app->Add(
           new JOmniFactoryGeneratorT<TrackClusterMergeSplitter_factory>(
-            "HcalEndcapNSplitMergeProtoClusters",
-            {"HcalEndcapNIslandProtoClusters",
-             "CalorimeterTrackProjections"},
-            {"HcalEndcapNSplitMergeProtoClusters"},
+            "HcalEndcapNSplitMergeClustersWithoutShapes",
+            {"HcalEndcapNClusters",
+             "CalorimeterTrackProjections",
+             "HcalEndcapNClusterAssociations",
+#if EDM4EIC_VERSION_MAJOR >= 7
+             "HcalEndcapNRawHitAssociations"},
+#else
+             "HcalEndcapNHits"},
+#endif
+            {"HcalEndcapNSplitMergeClustersWithoutShapes",
+#if EDM4EIC_VERSION_MAJOR >= 8
+             "HcalEndcapNTrackSplitMergeClusterMatches",
+#endif
+             "HcalEndcapNSplitMergeClusterAssociationsWithoutShapes"},
             {
               .idCalo = "HcalEndcapN_ID",
               .minSigCut = -2.0,
@@ -179,22 +189,6 @@ extern "C" {
             app   // TODO: remove me once fixed
           )
         );
-        app->Add(
-          new JOmniFactoryGeneratorT<CalorimeterClusterRecoCoG_factory>(
-             "HcalEndcapNSplitMergeClustersWithoutShapes",
-            {"HcalEndcapNSplitMergeProtoClusters", // edm4eic::ProtoClusterCollection
-             "HcalEndcapNHits"}, // edm4hep::SimCalorimeterHitCollection
-            {"HcalEndcapNSplitMergeClustersWithoutShapes", // edm4eic::Cluster
-             "HcalEndcapNSplitMergeClusterAssociationsWithoutShapes"}, // edm4eic::MCRecoClusterParticleAssociation
-            {
-              .energyWeight = "log",
-              .sampFrac = 1.0,
-              .logWeightBase = 6.2,
-              .enableEtaBounds = false
-            },
-            app   // TODO: Remove me once fixed
-          )
-        );
         app->Add(
           new JOmniFactoryGeneratorT<CalorimeterClusterShape_factory>(
             "HcalEndcapNSplitMergeClusters",

src/detectors/FEMC/FEMC.cc

@@ -160,10 +160,20 @@ extern "C" {
 
         app->Add(
           new JOmniFactoryGeneratorT<TrackClusterMergeSplitter_factory>(
-            "EcalEndcapPSplitMergeProtoClusters",
-            {"EcalEndcapPIslandProtoClusters",
-             "CalorimeterTrackProjections"},
-            {"EcalEndcapPSplitMergeProtoClusters"},
+            "EcalEndcapPSplitMergeClustersWithoutShapes",
+            {"EcalEndcapPClusters",
+             "CalorimeterTrackProjections",
+             "EcalEndcapPClusterAssociations",
+#if EDM4EIC_VERSION_MAJOR >= 7
+             "EcalEndcapPRawHitAssociations"},
+#else
+             "EcalEndcapPHits"},
+#endif
+            {"EcalEndcapPSplitMergeClustersWithoutShapes",
+#if EDM4EIC_VERSION_MAJOR >= 8
+             "EcalEndcapPTrackSplitMergeClusterMatches",
+#endif
+             "EcalEndcapPSplitMergeClusterAssociationsWithoutShapes"},
             {
               .idCalo = "EcalEndcapP_ID",
               .minSigCut = -2.0,
@@ -177,23 +187,6 @@ extern "C" {
           )
         );
 
-        app->Add(
-          new JOmniFactoryGeneratorT<CalorimeterClusterRecoCoG_factory>(
-             "EcalEndcapPSplitMergeClustersWithoutShapes",
-            {"EcalEndcapPSplitMergeProtoClusters", // edm4eic::ProtoClusterCollection
-             "EcalEndcapPHits"},                          // edm4hep::SimCalorimeterHitCollection
-            {"EcalEndcapPSplitMergeClustersWithoutShapes", // edm4eic::Cluster
-             "EcalEndcapPSplitMergeClusterAssociationsWithoutShapes"}, // edm4eic::MCRecoClusterParticleAssociation
-            {
-              .energyWeight = "log",
-              .sampFrac = 1.0,
-              .logWeightBase = 3.6,
-              .enableEtaBounds = false
-            },
-            app   // TODO: Remove me once fixed
-          )
-        );
-
         app->Add(
           new JOmniFactoryGeneratorT<CalorimeterClusterShape_factory>(
             "EcalEndcapPSplitMergeClusters",