algorithms/calorimetry/TrackClusterMergeSplitter.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Derek Anderson
0003
0004 #pragma once
0005
0006 #include <DD4hep/Detector.h>
0007 #include <algorithms/algorithm.h>
0008 #include <edm4eic/ProtoClusterCollection.h>
0009 #include <edm4eic/TrackPoint.h>
0010 #include <edm4eic/TrackSegmentCollection.h>
0011 #include <edm4hep/Vector3f.h>
0012 #include <podio/ObjectID.h>
0013 #include <algorithm>
0014 #include <map>
0015 #include <set>
0016 #include <string>
0017 #include <string_view>
0018 #include <vector>
0019
0020 #include "TrackClusterMergeSplitterConfig.h"
0021 // for algorithm configuration
0022 #include "algorithms/interfaces/WithPodConfig.h"
0023
0024
0025
0026 namespace eicrecon {
0027
0028   // --------------------------------------------------------------------------
0029   //! Comparator struct for protoclusters
0030   // --------------------------------------------------------------------------
0031   /*! Organizes protoclusters by their ObjectID's in decreasing collection
0032    *  ID first, and second by decreasing index second.
0033    */
0034   struct CompareProto {
0035
0036     bool operator() (const edm4eic::ProtoCluster& lhs, const edm4eic::ProtoCluster& rhs) const {
0037       if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
0038         return (lhs.getObjectID().index < rhs.getObjectID().index);
0039       } else {
0040         return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
0041       }
0042     }
0043
0044   };  // end CompareObjectID
0045
0046
0047
0048   // --------------------------------------------------------------------------
0049   //! Convenience types
0050   // --------------------------------------------------------------------------
0051   using VecProj = std::vector<edm4eic::TrackPoint>;
0052   using VecClust = std::vector<edm4eic::ProtoCluster>;
0053   using SetClust = std::set<edm4eic::ProtoCluster, CompareProto>;
0054   using MapToVecProj = std::map<edm4eic::ProtoCluster, VecProj, CompareProto>;
0055   using MapToVecClust = std::map<edm4eic::ProtoCluster, VecClust, CompareProto>;
0056
0057
0058
0059   // --------------------------------------------------------------------------
0060   //! Algorithm input/output
0061   // --------------------------------------------------------------------------
0062   using TrackClusterMergeSplitterAlgorithm = algorithms::Algorithm<
0063     algorithms::Input<
0064       edm4eic::ProtoClusterCollection,
0065       edm4eic::TrackSegmentCollection
0066     >,
0067     algorithms::Output<
0068       edm4eic::ProtoClusterCollection
0069     >
0070   >;
0071
0072
0073
0074   // --------------------------------------------------------------------------
0075   //! Track-Based Cluster Merger/Splitter
0076   // --------------------------------------------------------------------------
0077   /*! An algorithm which takes a collection of proto-clusters, matches
0078    *  track projections, and then decides to merge or split those proto-
0079    *  clusters based on average E/p from simulations.
0080    *
0081    *  Heavily inspired by Eur. Phys. J. C (2017) 77:466
0082    */
0083   class TrackClusterMergeSplitter :
0084     public TrackClusterMergeSplitterAlgorithm,
0085     public WithPodConfig<TrackClusterMergeSplitterConfig>
0086   {
0087
0088     public:
0089
0090       // ctor
0091       TrackClusterMergeSplitter(std::string_view name) :
0092         TrackClusterMergeSplitterAlgorithm {
0093           name,
0094           {"InputProtoClusterCollection", "InputTrackProjections"},
0095           {"OutputProtoClusterCollection"},
0096           "Merges or splits clusters based on tracks projected to them."
0097         } {}
0098
0099       // public methods
0100       void init(const dd4hep::Detector* detector);
0101       void process (const Input&, const Output&) const final;
0102
0103     private:
0104
0105       // private methods
0106       void get_projections(const edm4eic::TrackSegmentCollection* projections, VecProj& relevant_projects) const;
0107       void match_clusters_to_tracks(const edm4eic::ProtoClusterCollection* clusters, const VecProj& projections, MapToVecProj& matches) const;
0108       void merge_and_split_clusters(const VecClust& to_merge, const VecProj& to_split, edm4eic::ProtoClusterCollection* out_protoclusters) const;
0109       float get_cluster_energy(const edm4eic::ProtoCluster& clust) const;
0110       edm4hep::Vector3f get_cluster_position(const edm4eic::ProtoCluster& clust) const;
0111
0112       // calorimeter id
0113       int m_idCalo {0};
0114
0115   };  // end TrackClusterMergeSplitter
0116
0117 }  // end eicrecon namespace