Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-07-05 08:15:25

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