EIC code displayed by LXR master/​EICrecon/​src/​algorithms/​calorimetry/​TrackClusterMergeSplitter.cc	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2024-09-27 07:02:57

 // SPDX-License-Identifier: LGPL-3.0-or-later
 // Copyright (C) 2024 Derek Anderson
 
 #include <edm4eic/CalorimeterHit.h>
 #include <edm4hep/Vector3f.h>
 #include <edm4hep/utils/vector_utils.h>
 #include <fmt/core.h>
 #include <podio/ObjectID.h>
 #include <podio/RelationRange.h>
 #include <stdint.h>
 #include <cmath>
 #include <cstddef>
 #include <gsl/pointers>
 
 // algorithm definition
 #include "TrackClusterMergeSplitter.h"
 #include "algorithms/calorimetry/TrackClusterMergeSplitterConfig.h"
 
 
 
 namespace eicrecon {
 
   // --------------------------------------------------------------------------
   //! Initialize algorithm
   // --------------------------------------------------------------------------
   void TrackClusterMergeSplitter::init(const dd4hep::Detector* detector) {
 
     // grab detector id
     m_idCalo = detector -> constant<int>(m_cfg.idCalo);
     debug("Collecting projections to detector with system id {}", m_idCalo);
 
   }  // end 'init(dd4hep::Detector*)'
 
 
 
   // --------------------------------------------------------------------------
   //! Process inputs
   // --------------------------------------------------------------------------
   /*! Primary algorithm call: algorithm ingests a collection
    *  protoclusters and a collection of track projections.
    *  It then decides to merge or split protoclusters according
    *  to the following algorithm:
    *    1. Identify all tracks projections pointing to the
    *       specified calorimeter.
    *    2. Match relevant track projections to protoclusters
    *       based on distance between projection and the energy-
    *       weighted barycenter of the protocluster;
    *    3. For each cluster-track pair:
    *       i.  Calculate the significance of the pair's
    *           E/p relative to the provided mean E/p and
    *           its RMS; and
    *       ii. If the significance is less than the
    *           significance specified by `minSigCut`,
    *           merge all clusters within `drAdd`.
    *    4. Create a protocluster for each merged cluster
    *       and copy all unused protoclusters into output.
    *       - If multiple tracks point to the same merged
    *         cluster, create a new protocluster for each
    *         projection with hit weighted relative to
    *         the track momentum.
    */
   void TrackClusterMergeSplitter::process(
     const TrackClusterMergeSplitter::Input& input,
     const TrackClusterMergeSplitter::Output& output
   ) const {
 
     // grab inputs/outputs
     const auto [in_protoclusters, in_projections] = input;
     auto [out_protoclusters] = output;
 
     // exit if no clusters in collection
     if (in_protoclusters->size() == 0) {
       debug("No proto-clusters in input collection.");
       return;
     }
 
     // ------------------------------------------------------------------------
     // 1. Identify projections to calorimeter
     // ------------------------------------------------------------------------
     VecProj vecProject;
     get_projections(in_projections, vecProject);
 
     // ------------------------------------------------------------------------
     // 2. Match relevant projections to clusters
     // ------------------------------------------------------------------------
     MapToVecProj mapProjToSplit;
     if (vecProject.size() == 0) {
       debug("No projections to match clusters to.");
       return;
     } else {
       match_clusters_to_tracks(in_protoclusters, vecProject, mapProjToSplit);
     }
 
     // ------------------------------------------------------------------------
     // 3. Loop over projection-cluster pairs to check if merging is needed
     // ------------------------------------------------------------------------
     SetClust setUsedClust;
     MapToVecClust mapClustToMerge;
     for (auto& [clustSeed, vecMatchProj] : mapProjToSplit) {
 
       // at this point, track-cluster matches are 1-to-1
       // so grab matched track
       auto projSeed = vecMatchProj.front();
 
       // skip if cluster is already used
       if (setUsedClust.count(clustSeed)) {
         continue;
       }
 
       // add cluster to list and flag as used
       mapClustToMerge[clustSeed].push_back( clustSeed );
       setUsedClust.insert( clustSeed );
 
       // grab cluster energy and projection momentum
       const float eClustSeed = get_cluster_energy(clustSeed);
       const float eProjSeed = m_cfg.avgEP * edm4hep::utils::magnitude(projSeed.momentum);
 
       // ----------------------------------------------------------------------
       // 3.i. Calculate significance
       // ----------------------------------------------------------------------
       const float sigSeed = (eClustSeed - eProjSeed) / m_cfg.sigEP;
       trace("Seed energy = {}, expected energy = {}, significance = {}", eClustSeed, eProjSeed, sigSeed);
 
       // ----------------------------------------------------------------------
       // 3.ii. If significance is above threshold, do nothing.
       //       Otherwise identify clusters to merge.
       // ----------------------------------------------------------------------
       if (sigSeed > m_cfg.minSigCut) {
         continue;
       }
 
       // get eta, phi of seed
       const auto posSeed = get_cluster_position(clustSeed);
       const float etaSeed = edm4hep::utils::eta(posSeed);
       const float phiSeed = edm4hep::utils::angleAzimuthal(posSeed);
 
       // loop over other clusters
       float eClustSum = eClustSeed;
       float sigSum = sigSeed;
       for (auto in_cluster : *in_protoclusters) {
 
         // ignore used clusters
         if (setUsedClust.count(in_cluster)) {
           continue;
         }
 
         // get eta, phi of cluster
         const auto posClust = get_cluster_position(in_cluster);
         const float etaClust = edm4hep::utils::eta(posClust);
         const float phiClust = edm4hep::utils::angleAzimuthal(posClust);
 
         // get distance to seed
         const float drToSeed = std::hypot(
           etaSeed - etaClust,
           std::remainder(phiSeed - phiClust, 2. * M_PI)
         );
 
         // --------------------------------------------------------------------
         // If inside merging-window, add to list of clusters to merge
         // --------------------------------------------------------------------
         if (drToSeed > m_cfg.drAdd) {
           continue;
         } else {
           mapClustToMerge[clustSeed].push_back( in_cluster );
           setUsedClust.insert( in_cluster );
         }
 
         // --------------------------------------------------------------------
         // if picked up cluster w/ matched track, add projection to list
         // --------------------------------------------------------------------
         if (mapProjToSplit.count(in_cluster)) {
           vecMatchProj.insert(
             vecMatchProj.end(),
             mapProjToSplit[in_cluster].begin(),
             mapProjToSplit[in_cluster].end()
           );
         }
 
         // increment sums and output debugging
         eClustSum += get_cluster_energy(in_cluster);
         sigSum = (eClustSum - eProjSeed) / m_cfg.sigEP;
         trace(
           "{} clusters to merge: current sum = {}, current significance = {}, {} track(s) pointing to merged cluster",
           mapClustToMerge[clustSeed].size(),
           eClustSum,
           sigSum,
           vecMatchProj.size()
         );
       }  // end cluster loop
     }  // end matched cluster-projection loop
 
     // ------------------------------------------------------------------------
     // 4. Create an output protocluster for each merged cluster and for
     //    each track pointing to merged cluster
     // ------------------------------------------------------------------------
     for (auto& [clustSeed, vecClustToMerge] : mapClustToMerge) {
       merge_and_split_clusters(
         vecClustToMerge,
         mapProjToSplit[clustSeed],
         out_protoclusters
       );
     }  // end clusters to merge loop
 
     // ------------------------------------------------------------------------
     // copy unused clusters to output
     // ------------------------------------------------------------------------
     for (auto in_cluster : *in_protoclusters) {
 
       // ignore used clusters
       if (setUsedClust.count(in_cluster)) {
         continue;
       }
 
       // copy cluster and add to output collection
       edm4eic::MutableProtoCluster out_cluster = in_cluster.clone();
       out_protoclusters->push_back(out_cluster);
       trace("Copied input cluster {} onto output cluster {}",
         in_cluster.getObjectID().index,
         out_cluster.getObjectID().index
       );
 
     }  // end cluster loop
 
   }  // end 'process(Input&, Output&)'
 
 
 
   // --------------------------------------------------------------------------
   //! Collect projections pointing to calorimeter
   // --------------------------------------------------------------------------
   void TrackClusterMergeSplitter::get_projections(
     const edm4eic::TrackSegmentCollection* projections,
     VecProj& relevant_projects
   ) const {
 
     // return if projections are empty
     if (projections->size() == 0) {
       debug("No projections in input collection.");
       return;
     }
 
     // collect projections
     for (auto project : *projections) {
       for (auto point : project.getPoints()) {
         if (
           (point.system  == m_idCalo) &&
           (point.surface == 1)
         ) {
           relevant_projects.push_back(point);
         }
       }  // end point loop
     }  // end projection loop
     debug("Collected relevant projections: {} to process", relevant_projects.size());
 
   }  // end 'get_projections(edm4eic::CalorimeterHit&, edm4eic::TrackSegmentCollection&, VecTrkPoint&)'
 
 
 
   // --------------------------------------------------------------------------
   //! Match clusters to track projections
   // --------------------------------------------------------------------------
   /*! FIXME remove this once cluster-track matching has been centralized
    */
   void TrackClusterMergeSplitter::match_clusters_to_tracks(
     const edm4eic::ProtoClusterCollection* clusters,
     const VecProj& projections,
     MapToVecProj& matches
   ) const {
 
 
     // loop over relevant projections
     for (uint32_t iProject = 0; iProject < projections.size(); ++iProject) {
 
       // grab projection
       auto project = projections[iProject];
 
       // get eta, phi of projection
       const float etaProj = edm4hep::utils::eta(project.position);
       const float phiProj = edm4hep::utils::angleAzimuthal(project.position);
 
       // to store matched cluster
       edm4eic::ProtoCluster match;
 
       // find closest cluster
       bool foundMatch = false;
       float dMatch = m_cfg.drAdd;
       for (auto cluster : *clusters) {
 
         // get eta, phi of cluster
         const auto  posClust = get_cluster_position(cluster);
         const float etaClust = edm4hep::utils::eta(posClust);
         const float phiClust = edm4hep::utils::angleAzimuthal(posClust);
 
         // calculate distance to centroid
         const float dist = std::hypot(
           etaProj - etaClust,
           std::remainder(phiProj - phiClust, 2. * M_PI)
         );
 
         // if closer, set match to current projection
         if (dist <= dMatch) {
           foundMatch = true;
           dMatch = dist;
           match = cluster;
         }
       }  // end cluster loop
 
       // record match if found
       if (foundMatch) {
         matches[match].push_back(project);
         trace("Matched cluster to track projection: eta-phi distance = {}", dMatch);
       }
     }  // end cluster loop
     debug("Finished matching clusters to track projections: {} matches", matches.size());
 
   }  // end 'match_clusters_to_tracks(edm4eic::ClusterCollection*, VecTrkPoint&, MapToVecProj&)'
 
 
 
   // --------------------------------------------------------------------------
   //! Merge identified clusters and split if needed
   // --------------------------------------------------------------------------
   /*! If multiple tracks are pointing to merged cluster, a new protocluster
    *  is created for each track w/ hits weighted by its distance to the track
    *  and the track's momentum.
    */
   void TrackClusterMergeSplitter::merge_and_split_clusters(
     const VecClust& to_merge,
     const VecProj& to_split,
     edm4eic::ProtoClusterCollection* out_protoclusters
   ) const {
 
     // if only 1 matched track, no need to split
     if (to_split.size() == 1) {
       edm4eic::MutableProtoCluster new_clust = out_protoclusters->create();
       for (const auto& old_clust : to_merge) {
         for (const auto& hit : old_clust.getHits()) {
           new_clust.addToHits( hit );
           new_clust.addToWeights( 1. );
         }
         trace("Merged input cluster {} into output cluster {}", old_clust.getObjectID().index, new_clust.getObjectID().index);
       }
       return;
     }
 
     // otherwise split merged cluster for each matched track
     std::vector<edm4eic::MutableProtoCluster> new_clusters;
     for (const auto& proj : to_split) {
       new_clusters.push_back( out_protoclusters->create() );
     }
     trace("Splitting merged cluster across {} tracks", to_split.size());
 
     // loop over all hits from all clusters to merge
     std::vector<float> weights( to_split.size(), 1. );
     for (const auto& old_clust : to_merge) {
       for (const auto& hit : old_clust.getHits()) {
 
         // calculate hit's weight for each track
         for (std::size_t iProj = 0; const auto& proj : to_split) {
 
           // get track eta, phi
           const float etaProj = edm4hep::utils::eta(proj.position);
           const float phiProj = edm4hep::utils::angleAzimuthal(proj.position);
 
           // get hit eta, phi
           const float etaHit = edm4hep::utils::eta(hit.getPosition());
           const float phiHit = edm4hep::utils::angleAzimuthal(hit.getPosition());
 
           // get track momentum, distance to hit
           const float mom  = edm4hep::utils::magnitude(proj.momentum);
           const float dist = std::hypot(
             etaHit - etaProj,
             std::remainder(phiHit - phiProj, 2. * M_PI)
           );
 
           // set weight
           weights[iProj] = std::exp(-1. * dist / m_cfg.transverseEnergyProfileScale) * mom;
           ++iProj;
         }
 
         // normalize weights
         float wTotal = 0.;
         for (const float weight : weights) {
           wTotal += weight;
         }
         for (float& weight : weights) {
           weight /= wTotal;
         }
 
         // add hit to each split merged cluster w/ relevant weight
         for (std::size_t iProj = 0; auto& new_clust : new_clusters) {
           new_clust.addToHits( hit );
           new_clust.addToWeights( weights[iProj] );
         }
 
       }  // end hits to merge loop
     }  // end clusters to merge loop
 
   }  // end 'merge_and_split_clusters(VecClust&, VecProj&, edm4eic::MutableCluster&)'
 
 
 
   // --------------------------------------------------------------------------
   //! Grab current energy of protocluster
   // --------------------------------------------------------------------------
   float TrackClusterMergeSplitter::get_cluster_energy(const edm4eic::ProtoCluster& clust) const {
 
     float eClust = 0.;
     for (auto hit : clust.getHits()) {
       eClust += hit.getEnergy();
     }
     return eClust / m_cfg.sampFrac;
 
   }  // end 'get_cluster_energy(edm4eic::ProtoCluster&)'
 
 
 
   // --------------------------------------------------------------------------
   //! Get current center of protocluster
   // --------------------------------------------------------------------------
   edm4hep::Vector3f TrackClusterMergeSplitter::get_cluster_position(const edm4eic::ProtoCluster& clust) const {
 
     // grab total energy
     const float eClust = get_cluster_energy(clust) * m_cfg.sampFrac;
 
     // calculate energy-weighted center
     float wTotal = 0.;
     edm4hep::Vector3f position(0., 0., 0.);
     for (auto hit : clust.getHits()) {
 
       // calculate weight
       float weight = hit.getEnergy() / eClust;
       wTotal += weight;
 
       // update cluster position
       position = position + (hit.getPosition() * weight);
     }
 
     float norm = 1.;
     if (wTotal == 0.) {
       warning("Total weight of 0 in position calculation!");
     } else {
       norm = wTotal;
     }
     return position / norm;
 
   }  // end 'get_cluster_position(edm4eic::ProtoCluster&)'
 
 }  // end eicrecon namespace

This page was automatically generated by the 2.3.7 LXR engine. The LXR team