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