Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-09-07 08:28:12

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2026 Derek Anderson
0003 
0004 #include <algorithms/logger.h>
0005 #include <catch2/catch_test_macros.hpp>
0006 #include <edm4eic/CalorimeterHitCollection.h>
0007 #include <edm4eic/MCRecoCalorimeterHitLinkCollection.h>
0008 #include <edm4eic/ProtoClusterCollection.h>
0009 #include <edm4hep/CaloHitContributionCollection.h>
0010 #include <edm4hep/MCParticleCollection.h>
0011 #include <edm4hep/RawCalorimeterHitCollection.h>
0012 #include <edm4hep/SimCalorimeterHitCollection.h>
0013 #include <podio/RelationRange.h>
0014 #include <podio/detail/Link.h>
0015 #include <podio/detail/LinkCollectionImpl.h>
0016 #include <cmath>
0017 #include <cstddef>
0018 #include <deque>
0019 #include <functional>
0020 #include <limits>
0021 #include <memory>
0022 #include <set>
0023 #include <vector>
0024 
0025 #include "algorithms/calorimetry/CalorimeterTruthClustering.h"
0026 
0027 TEST_CASE("the CalorimeterTruthClustering algorithm runs", "[CalorimeterTruthClustering]") {
0028 
0029   eicrecon::CalorimeterTruthClustering algo_clustering("calorimeterTruthClustering");
0030   algo_clustering.level(algorithms::LogLevel::kDebug);
0031   algo_clustering.init();
0032 
0033   auto empty_rec_calo_hit_coll    = std::make_unique<edm4eic::CalorimeterHitCollection>();
0034   auto empty_mc_rec_hit_link_coll = std::make_unique<edm4eic::MCRecoCalorimeterHitLinkCollection>();
0035   auto empty_truth_clust_coll     = std::make_unique<edm4eic::ProtoClusterCollection>();
0036   algo_clustering.process({empty_rec_calo_hit_coll.get(), empty_mc_rec_hit_link_coll.get()},
0037                           {empty_truth_clust_coll.get()});
0038   REQUIRE(empty_truth_clust_coll->size() == 0);
0039 
0040   edm4hep::MutableMCParticle par_a = edm4hep::MutableMCParticle();
0041   par_a.setPDG(11);
0042   par_a.setCharge(-1.0);
0043   par_a.setMass(0.000511);
0044   par_a.setGeneratorStatus(1);
0045 
0046   edm4hep::MutableMCParticle par_b = edm4hep::MutableMCParticle();
0047   par_b.setPDG(22);
0048   par_b.setCharge(0.0);
0049   par_b.setMass(0.0);
0050   par_b.setGeneratorStatus(1);
0051 
0052   edm4hep::MutableMCParticle par_c = edm4hep::MutableMCParticle();
0053   par_c.setPDG(111);
0054   par_c.setCharge(0.0);
0055   par_c.setMass(0.134);
0056   par_c.setGeneratorStatus(1);
0057 
0058   edm4hep::MutableMCParticle par_c_c = edm4hep::MutableMCParticle();
0059   par_c_c.setPDG(22);
0060   par_c_c.setCharge(0.0);
0061   par_c_c.setMass(0.0);
0062   par_c_c.setGeneratorStatus(0);
0063   par_c_c.addToParents(par_c);
0064   par_c.addToDaughters(par_c_c);
0065 
0066   edm4hep::MutableMCParticle par_c_d = edm4hep::MutableMCParticle();
0067   par_c_d.setPDG(22);
0068   par_c_d.setCharge(0.0);
0069   par_c_d.setMass(0.0);
0070   par_c_d.setGeneratorStatus(0);
0071   par_c_d.addToParents(par_c);
0072   par_c.addToDaughters(par_c_d);
0073 
0074   auto mc_par_coll = std::make_unique<edm4hep::MCParticleCollection>();
0075   mc_par_coll->push_back(par_a);
0076   mc_par_coll->push_back(par_b);
0077   mc_par_coll->push_back(par_c);
0078   mc_par_coll->push_back(par_c_c);
0079   mc_par_coll->push_back(par_c_d);
0080 
0081   // define contributions
0082   //   - par A: 2.5 GeV --> hit A
0083   //   - par A: 0.5 GeV --> hit b
0084   //   - par B: 2.0 GeV --> hit b
0085   //   - par C: 0.5 GeV --> hit c (via decay photon)
0086   //   - par C: 0.5 GeV --> hit d (via decay photon)
0087   auto calo_hit_contrib_coll = std::make_unique<edm4hep::CaloHitContributionCollection>();
0088   auto contrib_a_a           = calo_hit_contrib_coll->create(11, 2.5);
0089   auto contrib_a_b           = calo_hit_contrib_coll->create(11, 0.5);
0090   auto contrib_b             = calo_hit_contrib_coll->create(22, 2.0);
0091   auto contrib_c_c           = calo_hit_contrib_coll->create(22, 0.5);
0092   auto contrib_c_d           = calo_hit_contrib_coll->create(22, 0.5);
0093   contrib_a_a.setParticle(par_a);
0094   contrib_a_b.setParticle(par_a);
0095   contrib_b.setParticle(par_b);
0096   contrib_c_c.setParticle(par_c_c);
0097   contrib_c_d.setParticle(par_c_d);
0098 
0099   auto sim_calo_hit_coll = std::make_unique<edm4hep::SimCalorimeterHitCollection>();
0100   auto sim_hit_a         = sim_calo_hit_coll->create(0, 2.5);
0101   auto sim_hit_b         = sim_calo_hit_coll->create(1, 2.5);
0102   auto sim_hit_c         = sim_calo_hit_coll->create(2, 0.5);
0103   auto sim_hit_d         = sim_calo_hit_coll->create(3, 0.5);
0104   sim_hit_a.addToContributions(contrib_a_a);
0105   sim_hit_b.addToContributions(contrib_a_b);
0106   sim_hit_b.addToContributions(contrib_b);
0107   sim_hit_c.addToContributions(contrib_c_c);
0108   sim_hit_d.addToContributions(contrib_c_d);
0109 
0110   auto raw_calo_hit_coll = std::make_unique<edm4hep::RawCalorimeterHitCollection>();
0111   auto raw_hit_a         = raw_calo_hit_coll->create(0, 250);
0112   auto raw_hit_b         = raw_calo_hit_coll->create(1, 250);
0113   auto raw_hit_c         = raw_calo_hit_coll->create(2, 50);
0114   auto raw_hit_d         = raw_calo_hit_coll->create(3, 50);
0115 
0116   auto mc_rec_hit_link_coll = std::make_unique<edm4eic::MCRecoCalorimeterHitLinkCollection>();
0117   auto mc_rec_hit_link_a    = mc_rec_hit_link_coll->create();
0118   auto mc_rec_hit_link_b    = mc_rec_hit_link_coll->create();
0119   auto mc_rec_hit_link_c    = mc_rec_hit_link_coll->create();
0120   auto mc_rec_hit_link_d    = mc_rec_hit_link_coll->create();
0121   mc_rec_hit_link_a.setFrom(raw_hit_a);
0122   mc_rec_hit_link_a.setTo(sim_hit_a);
0123   mc_rec_hit_link_b.setFrom(raw_hit_b);
0124   mc_rec_hit_link_b.setTo(sim_hit_b);
0125   mc_rec_hit_link_c.setFrom(raw_hit_c);
0126   mc_rec_hit_link_c.setTo(sim_hit_c);
0127   mc_rec_hit_link_d.setFrom(raw_hit_d);
0128   mc_rec_hit_link_d.setTo(sim_hit_d);
0129 
0130   auto rec_calo_hit_coll = std::make_unique<edm4eic::CalorimeterHitCollection>();
0131   auto rec_hit_a         = rec_calo_hit_coll->create(0, 2.5);
0132   auto rec_hit_b         = rec_calo_hit_coll->create(1, 2.5);
0133   auto rec_hit_c         = rec_calo_hit_coll->create(2, 0.5);
0134   auto rec_hit_d         = rec_calo_hit_coll->create(3, 0.5);
0135   rec_hit_a.setRawHit(raw_hit_a);
0136   rec_hit_b.setRawHit(raw_hit_b);
0137   rec_hit_c.setRawHit(raw_hit_c);
0138   rec_hit_d.setRawHit(raw_hit_d);
0139 
0140   // cluster rec hits based on truth info: should produce 3 clusters
0141   //   - clust A = {hit_a, hit_b}
0142   //   - clust B = {hit_b}
0143   //   - clust C = {hit_c, hit_d}
0144   auto truth_clust_coll = std::make_unique<edm4eic::ProtoClusterCollection>();
0145   algo_clustering.process({rec_calo_hit_coll.get(), mc_rec_hit_link_coll.get()},
0146                           {truth_clust_coll.get()});
0147   REQUIRE(truth_clust_coll->size() == 3);
0148 
0149   const std::set clust_a{0, 1};
0150   const std::set clust_b{1};
0151   const std::set clust_c{2, 3};
0152   for (const auto& clust : *truth_clust_coll) {
0153     for (const auto& hit : clust.getHits()) {
0154       const auto cell_id = hit.getCellID();
0155       switch (cell_id) {
0156       case 0:
0157         REQUIRE(clust_a.contains(cell_id));
0158         break;
0159       case 1:
0160         REQUIRE(clust_a.contains(cell_id));
0161         REQUIRE(clust_b.contains(cell_id));
0162         break;
0163       case 2:
0164         REQUIRE(clust_c.contains(cell_id));
0165         break;
0166       case 3:
0167         REQUIRE(clust_c.contains(cell_id));
0168         break;
0169       default:
0170         FAIL("Unknown cell ID encountered");
0171         break;
0172       }
0173     }
0174   }
0175 
0176   // weights are ratio of sim energy contributed by a particle to hit over
0177   // total sim energy of hit
0178   //   clust A --> {weight_a = 1.0, weight_b = 0.2}
0179   //   clust B --> {weight_b = 0.8}
0180   //   clust C --> {weight_c = 1.0, weight_d = 1.0}
0181   for (const auto& clust : *truth_clust_coll) {
0182     for (std::size_t ihit = 0; const auto& hit : clust.getHits()) {
0183       const auto cell_id = hit.getCellID();
0184       const float cell_w = clust.getWeights()[ihit];
0185       switch (cell_id) {
0186       case 0:
0187         REQUIRE(std::abs(cell_w - 1.0) < std::numeric_limits<float>::epsilon());
0188         break;
0189       case 1:
0190         // if parent cluster has 2 hits, then in cluster A
0191         // otherwise in cluster B
0192         if (clust.getHits().size() == 2) {
0193           REQUIRE(std::abs(cell_w - 0.2) < std::numeric_limits<float>::epsilon());
0194         } else {
0195           REQUIRE(std::abs(cell_w - 0.8) < std::numeric_limits<float>::epsilon());
0196         }
0197         break;
0198       case 2:
0199         REQUIRE(std::abs(cell_w - 1.0) < std::numeric_limits<float>::epsilon());
0200         break;
0201       case 3:
0202         REQUIRE(std::abs(cell_w - 1.0) < std::numeric_limits<float>::epsilon());
0203         break;
0204       default:
0205         FAIL("Unknown cell ID encountered");
0206         break;
0207       }
0208       ++ihit;
0209     }
0210   }
0211 }