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 <algorithms/algorithm.h>
0007 #include <algorithms/geo.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 namespace eicrecon {
0025
0026 // --------------------------------------------------------------------------
0027 //! Comparator struct for protoclusters
0028 // --------------------------------------------------------------------------
0029 /*! Organizes protoclusters by their ObjectID's in decreasing collection
0030    *  ID first, and second by decreasing index second.
0031    */
0032 struct CompareProto {
0033
0034   bool operator()(const edm4eic::ProtoCluster& lhs, const edm4eic::ProtoCluster& rhs) const {
0035     if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
0036       return (lhs.getObjectID().index < rhs.getObjectID().index);
0037     } else {
0038       return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
0039     }
0040   }
0041
0042 }; // end CompareObjectID
0043
0044 // --------------------------------------------------------------------------
0045 //! Convenience types
0046 // --------------------------------------------------------------------------
0047 using VecProj       = std::vector<edm4eic::TrackPoint>;
0048 using VecClust      = std::vector<edm4eic::ProtoCluster>;
0049 using SetClust      = std::set<edm4eic::ProtoCluster, CompareProto>;
0050 using MapToVecProj  = std::map<edm4eic::ProtoCluster, VecProj, CompareProto>;
0051 using MapToVecClust = std::map<edm4eic::ProtoCluster, VecClust, CompareProto>;
0052
0053 // --------------------------------------------------------------------------
0054 //! Algorithm input/output
0055 // --------------------------------------------------------------------------
0056 using TrackClusterMergeSplitterAlgorithm = algorithms::Algorithm<
0057     algorithms::Input<edm4eic::ProtoClusterCollection, edm4eic::TrackSegmentCollection>,
0058     algorithms::Output<edm4eic::ProtoClusterCollection>>;
0059
0060 // --------------------------------------------------------------------------
0061 //! Track-Based Cluster Merger/Splitter
0062 // --------------------------------------------------------------------------
0063 /*! An algorithm which takes a collection of proto-clusters, matches
0064    *  track projections, and then decides to merge or split those proto-
0065    *  clusters based on average E/p from simulations.
0066    *
0067    *  Heavily inspired by Eur. Phys. J. C (2017) 77:466
0068    */
0069 class TrackClusterMergeSplitter : public TrackClusterMergeSplitterAlgorithm,
0070                                   public WithPodConfig<TrackClusterMergeSplitterConfig> {
0071
0072 public:
0073   // ctor
0074   TrackClusterMergeSplitter(std::string_view name)
0075       : TrackClusterMergeSplitterAlgorithm{
0076             name,
0077             {"InputProtoClusterCollection", "InputTrackProjections"},
0078             {"OutputProtoClusterCollection"},
0079             "Merges or splits clusters based on tracks projected to them."} {}
0080
0081   // public methods
0082   void init();
0083   void process(const Input&, const Output&) const final;
0084
0085 private:
0086   const algorithms::GeoSvc& m_geo = algorithms::GeoSvc::instance();
0087
0088   // private methods
0089   void get_projections(const edm4eic::TrackSegmentCollection* projections,
0090                        VecProj& relevant_projects) const;
0091   void match_clusters_to_tracks(const edm4eic::ProtoClusterCollection* clusters,
0092                                 const VecProj& projections, MapToVecProj& matches) const;
0093   void merge_and_split_clusters(const VecClust& to_merge, const VecProj& to_split,
0094                                 edm4eic::ProtoClusterCollection* out_protoclusters) const;
0095   float get_cluster_energy(const edm4eic::ProtoCluster& clust) const;
0096   edm4hep::Vector3f get_cluster_position(const edm4eic::ProtoCluster& clust) const;
0097
0098   // calorimeter id
0099   unsigned int m_idCalo{0};
0100
0101 }; // end TrackClusterMergeSplitter
0102
0103 } // namespace eicrecon