Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-17 08:07:33

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023, Christopher Dilks
0003 
0004 #include <catch2/catch_test_macros.hpp>
0005 #include <catch2/matchers/catch_matchers.hpp>
0006 #include <catch2/matchers/catch_matchers_floating_point.hpp>
0007 #include <edm4eic/CherenkovParticleIDCollection.h>
0008 #include <edm4eic/CherenkovParticleIDHypothesis.h>
0009 #include <edm4eic/TrackSegmentCollection.h>
0010 #include <edm4hep/Vector2f.h>
0011 #include <podio/RelationRange.h>
0012 #include <spdlog/common.h>
0013 #include <spdlog/logger.h>
0014 #include <spdlog/spdlog.h>
0015 #include <cmath>
0016 #include <gsl/pointers>
0017 #include <memory>
0018 #include <stdexcept>
0019 #include <string>
0020 #include <vector>
0021 
0022 #include "algorithms/pid/MergeParticleID.h"
0023 #include "algorithms/pid/MergeParticleIDConfig.h"
0024 
0025 TEST_CASE("the PID MergeParticleID algorithm runs", "[MergeParticleID]") {
0026 
0027   const float EPSILON = 1e-5;
0028 
0029   std::shared_ptr<spdlog::logger> logger = spdlog::default_logger()->clone("MergeParticleID");
0030   logger->set_level(spdlog::level::debug);
0031 
0032   // charged particle input data
0033   //----------------------------------------------------------
0034   /* just create two unique tracks; they don't have to be filled with data,
0035    * since `MergeParticleID` matches tracks by ID
0036    */
0037   auto tracks = std::make_unique<edm4eic::TrackSegmentCollection>();
0038   for (int i = 0; i < 2; i++) {
0039     tracks->create();
0040   }
0041 
0042   // Cherenkov PID inputs
0043   //----------------------------------------------------------
0044   /* 2 collections, each with 2 elements corresponding to 2 charged particles;
0045    * 2 or 3 possible PID hypotheses for each
0046    */
0047 
0048   edm4eic::MutableCherenkovParticleID pid;
0049   edm4eic::CherenkovParticleIDHypothesis pid_hyp;
0050 
0051   // collection 1 ------
0052   auto coll_cherenkov_1 = std::make_unique<edm4eic::CherenkovParticleIDCollection>();
0053 
0054   pid = coll_cherenkov_1->create();
0055   pid.setChargedParticle(tracks->at(0)); // track 0 first
0056   pid.setNpe(10);
0057   pid.setRefractiveIndex(1.05);
0058   pid.setPhotonEnergy(3e-9);
0059   pid_hyp.PDG    = 211; // pion hypothesis first
0060   pid_hyp.npe    = 10;
0061   pid_hyp.weight = 100;
0062   pid.addToHypotheses(pid_hyp);
0063   pid_hyp.PDG    = 321;
0064   pid_hyp.npe    = 8;
0065   pid_hyp.weight = 80;
0066   pid.addToHypotheses(pid_hyp);
0067   for (int i = 0; i <= pid.getNpe(); i++) {
0068     pid.addToThetaPhiPhotons(edm4hep::Vector2f{M_PI / 4, 0});
0069   }
0070 
0071   pid = coll_cherenkov_1->create();
0072   pid.setChargedParticle(tracks->at(1));
0073   pid.setNpe(11);
0074   pid.setRefractiveIndex(1.06);
0075   pid.setPhotonEnergy(4e-9);
0076   pid_hyp.PDG    = 321; // kaon hypothesis first
0077   pid_hyp.npe    = 10;
0078   pid_hyp.weight = 100;
0079   pid.addToHypotheses(pid_hyp);
0080   pid_hyp.PDG    = 211;
0081   pid_hyp.npe    = 10;
0082   pid_hyp.weight = 80;
0083   pid.addToHypotheses(pid_hyp);
0084   for (int i = 0; i <= pid.getNpe(); i++) {
0085     pid.addToThetaPhiPhotons(edm4hep::Vector2f{-M_PI / 4, 0});
0086   }
0087 
0088   // collection 2 ------
0089   auto coll_cherenkov_2 = std::make_unique<edm4eic::CherenkovParticleIDCollection>();
0090 
0091   pid = coll_cherenkov_2->create();
0092   pid.setChargedParticle(tracks->at(1)); // track 1 first
0093   pid.setNpe(4);
0094   pid.setRefractiveIndex(1.5);
0095   pid.setPhotonEnergy(3e-9);
0096   pid_hyp.PDG    = 211;
0097   pid_hyp.npe    = 3;
0098   pid_hyp.weight = 200;
0099   pid.addToHypotheses(pid_hyp);
0100   pid_hyp.PDG    = 321;
0101   pid_hyp.npe    = 2;
0102   pid_hyp.weight = 90;
0103   pid.addToHypotheses(pid_hyp);
0104   for (int i = 0; i <= pid.getNpe(); i++) {
0105     pid.addToThetaPhiPhotons(edm4hep::Vector2f{M_PI / 4, 0});
0106   }
0107 
0108   pid = coll_cherenkov_2->create();
0109   pid.setChargedParticle(tracks->at(0));
0110   pid.setNpe(6);
0111   pid.setRefractiveIndex(1.3);
0112   pid.setPhotonEnergy(4e-9);
0113   pid_hyp.PDG    = 321;
0114   pid_hyp.npe    = 4;
0115   pid_hyp.weight = 70;
0116   pid.addToHypotheses(pid_hyp);
0117   pid_hyp.PDG    = 211;
0118   pid_hyp.npe    = 1;
0119   pid_hyp.weight = 80;
0120   pid.addToHypotheses(pid_hyp);
0121   pid_hyp.PDG    = 11; // an additional hypothesis
0122   pid_hyp.npe    = 1;
0123   pid_hyp.weight = 60;
0124   pid.addToHypotheses(pid_hyp);
0125   for (int i = 0; i <= pid.getNpe(); i++) {
0126     pid.addToThetaPhiPhotons(edm4hep::Vector2f{-M_PI / 4, 0});
0127   }
0128 
0129   // Cherenkov PID tests
0130   //----------------------------------------------------------
0131 
0132   std::vector<gsl::not_null<const edm4eic::CherenkovParticleIDCollection*>> coll_cherenkov_list = {
0133       coll_cherenkov_1.get(), coll_cherenkov_2.get()};
0134 
0135   auto find_cherenkov_pid_for_track = [](const auto& coll, auto trk) {
0136     for (auto obj : *coll) {
0137       if (obj.getChargedParticle().id() == trk.id()) {
0138         return obj;
0139       }
0140     }
0141     FAIL("ERROR: cannot find CherenkovParticleID given track");
0142     if (coll->size() == 0) {
0143       throw std::runtime_error(
0144           "empty collection used in pid_MergeParticleID::find_cherenkov_pid_for_track");
0145     }
0146     return coll->at(0);
0147   };
0148 
0149   // additive weights
0150   SECTION("merge CherenkovParticleID: add hypothesis weights") {
0151 
0152     eicrecon::MergeParticleID algo("test");
0153     eicrecon::MergeParticleIDConfig cfg;
0154     cfg.mergeMode = eicrecon::MergeParticleIDConfig::kAddWeights;
0155     algo.applyConfig(cfg);
0156     algo.init();
0157 
0158     auto result = std::make_unique<edm4eic::CherenkovParticleIDCollection>();
0159     algo.process({coll_cherenkov_list}, {result.get()});
0160     auto pid_0 = find_cherenkov_pid_for_track(result, tracks->at(0));
0161     auto pid_1 = find_cherenkov_pid_for_track(result, tracks->at(1));
0162 
0163     REQUIRE_THAT(pid_0.getNpe(), Catch::Matchers::WithinAbs(10 + 6, EPSILON));
0164     REQUIRE_THAT(pid_1.getNpe(), Catch::Matchers::WithinAbs(11 + 4, EPSILON));
0165 
0166     REQUIRE_THAT(pid_0.getRefractiveIndex(),
0167                  Catch::Matchers::WithinAbs((10 * 1.05 + 6 * 1.3) / (10 + 6), EPSILON));
0168     REQUIRE_THAT(pid_1.getRefractiveIndex(),
0169                  Catch::Matchers::WithinAbs((11 * 1.06 + 4 * 1.5) / (11 + 4), EPSILON));
0170 
0171     REQUIRE_THAT(pid_0.getPhotonEnergy(),
0172                  Catch::Matchers::WithinAbs((10 * 3e-9 + 6 * 4e-9) / (10 + 6), EPSILON));
0173     REQUIRE_THAT(pid_1.getPhotonEnergy(),
0174                  Catch::Matchers::WithinAbs((11 * 4e-9 + 4 * 3e-9) / (11 + 4), EPSILON));
0175 
0176     REQUIRE(pid_0.hypotheses_size() == 3);
0177     for (auto hyp : pid_0.getHypotheses()) {
0178       switch (hyp.PDG) {
0179       case 211:
0180         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 1, EPSILON));
0181         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(100 + 80, EPSILON));
0182         break;
0183       case 321:
0184         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(8 + 4, EPSILON));
0185         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(80 + 70, EPSILON));
0186         break;
0187       case 11:
0188         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(1, EPSILON));
0189         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(60, EPSILON));
0190         break;
0191       default:
0192         FAIL("untested PDG hypothesis");
0193       }
0194     }
0195 
0196     REQUIRE(pid_1.hypotheses_size() == 2);
0197     for (auto hyp : pid_1.getHypotheses()) {
0198       switch (hyp.PDG) {
0199       case 211:
0200         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 3, EPSILON));
0201         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(80 + 200, EPSILON));
0202         break;
0203       case 321:
0204         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 2, EPSILON));
0205         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(100 + 90, EPSILON));
0206         break;
0207       default:
0208         FAIL("untested PDG hypothesis");
0209       }
0210     }
0211   }
0212 
0213   // multiplicative weights
0214   /* NOTE: the only difference from additive weights is the resulting weight values, but it
0215    * is still wise to re-test everything, so that any changes in the algorithm require corresponding
0216    * changes in these tests
0217    */
0218   SECTION("merge CherenkovParticleID: multiply hypothesis weights") {
0219 
0220     eicrecon::MergeParticleID algo("test");
0221     eicrecon::MergeParticleIDConfig cfg;
0222     cfg.mergeMode = eicrecon::MergeParticleIDConfig::kMultiplyWeights;
0223     algo.applyConfig(cfg);
0224     algo.init();
0225 
0226     auto result = std::make_unique<edm4eic::CherenkovParticleIDCollection>();
0227     algo.process({coll_cherenkov_list}, {result.get()});
0228     auto pid_0 = find_cherenkov_pid_for_track(result, tracks->at(0));
0229     auto pid_1 = find_cherenkov_pid_for_track(result, tracks->at(1));
0230 
0231     REQUIRE_THAT(pid_0.getNpe(), Catch::Matchers::WithinAbs(10 + 6, EPSILON));
0232     REQUIRE_THAT(pid_1.getNpe(), Catch::Matchers::WithinAbs(11 + 4, EPSILON));
0233 
0234     REQUIRE_THAT(pid_0.getRefractiveIndex(),
0235                  Catch::Matchers::WithinAbs((10 * 1.05 + 6 * 1.3) / (10 + 6), EPSILON));
0236     REQUIRE_THAT(pid_1.getRefractiveIndex(),
0237                  Catch::Matchers::WithinAbs((11 * 1.06 + 4 * 1.5) / (11 + 4), EPSILON));
0238 
0239     REQUIRE_THAT(pid_0.getPhotonEnergy(),
0240                  Catch::Matchers::WithinAbs((10 * 3e-9 + 6 * 4e-9) / (10 + 6), EPSILON));
0241     REQUIRE_THAT(pid_1.getPhotonEnergy(),
0242                  Catch::Matchers::WithinAbs((11 * 4e-9 + 4 * 3e-9) / (11 + 4), EPSILON));
0243 
0244     for (auto hyp : pid_0.getHypotheses()) {
0245       switch (hyp.PDG) {
0246       case 211:
0247         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 1, EPSILON));
0248         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(100 * 80, EPSILON));
0249         break;
0250       case 321:
0251         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(8 + 4, EPSILON));
0252         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(80 * 70, EPSILON));
0253         break;
0254       case 11:
0255         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(1, EPSILON));
0256         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(60, EPSILON));
0257         break;
0258       default:
0259         FAIL("untested PDG hypothesis");
0260       }
0261     }
0262 
0263     for (auto hyp : pid_1.getHypotheses()) {
0264       switch (hyp.PDG) {
0265       case 211:
0266         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 3, EPSILON));
0267         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(80 * 200, EPSILON));
0268         break;
0269       case 321:
0270         REQUIRE_THAT(hyp.npe, Catch::Matchers::WithinAbs(10 + 2, EPSILON));
0271         REQUIRE_THAT(hyp.weight, Catch::Matchers::WithinAbs(100 * 90, EPSILON));
0272         break;
0273       default:
0274         FAIL("untested PDG hypothesis");
0275       }
0276     }
0277   }
0278 }