Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-06-29 07:05:30

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Wouter Deconinck
0003 
0004 #include <algorithm>
0005 #include <cmath>
0006 #include <fmt/format.h>
0007 
0008 #include "Gaudi/Algorithm.h"
0009 #include "GaudiAlg/GaudiTool.h"
0010 #include "GaudiAlg/Producer.h"
0011 #include "GaudiAlg/Transformer.h"
0012 #include "GaudiKernel/RndmGenerators.h"
0013 
0014 #include "JugBase/DataHandle.h"
0015 
0016 // Event Model related classes
0017 #include "edm4hep/MCParticleCollection.h"
0018 #include "edm4eic/MCRecoParticleAssociationCollection.h"
0019 #include "edm4eic/ReconstructedParticleCollection.h"
0020 #include "edm4hep/utils/vector_utils.h"
0021 
0022 namespace {
0023 enum DetectorTags { kTagB0 = 1, kTagRP = 2, kTagOMD = 3, kTagZDC = 4 };
0024 }
0025 
0026 namespace Jug::Fast {
0027 
0028 class SmearedFarForwardParticles : public GaudiAlgorithm {
0029 private:
0030   DataHandle<edm4hep::MCParticleCollection> m_inputMCParticles{"inputMCParticles", Gaudi::DataHandle::Reader, this};
0031   DataHandle<edm4eic::ReconstructedParticleCollection> m_outputParticles{"SmearedFarForwardParticles",
0032                                                                       Gaudi::DataHandle::Writer, this};
0033   DataHandle<edm4eic::MCRecoParticleAssociationCollection> m_outputAssocCollection{"MCRecoParticleAssociation",
0034                                                                                 Gaudi::DataHandle::Writer, this};
0035 
0036   Gaudi::Property<bool> m_enableZDC{this, "enableZDC", true};
0037   Gaudi::Property<bool> m_enableB0{this, "enableB0", true};
0038   Gaudi::Property<bool> m_enableRP{this, "enableRP", true};
0039   Gaudi::Property<bool> m_enableOMD{this, "enableOMD", true};
0040 
0041   // Beam energy, only used to determine the RP/OMD momentum ranges
0042   Gaudi::Property<double> m_ionBeamEnergy{this, "ionBeamEnergy", 0.};
0043   // RP default to 10-on-100 setting
0044   // Pz > 60% of beam energy (60% x 100GeV = 60GeV)
0045   // theta from 0.2mrad -> 5mrad
0046   Gaudi::Property<double> m_thetaMinRP{this, "thetaMinRP", 0.2e-3};
0047   Gaudi::Property<double> m_thetaMaxRP{this, "thetaMaxRP", 5e-3};
0048   Gaudi::Property<double> m_pMinRigidityRP{this, "pMinRigidityRP", 0.60};
0049   // B0
0050   Gaudi::Property<double> m_thetaMinB0{this, "thetaMinB0", 6.0e-3};
0051   Gaudi::Property<double> m_thetaMaxB0{this, "thetaMaxB0", 20.0e-3};
0052   // OMD default to 10-on-100 setting
0053   // 25% < P/Ebeam < 60% of beam energy (25% x 100GeV = 25GeV and 60% x 100GeV = 60GeV)
0054   // Angles both given for the small angle full-acceptance part,
0055   // and for the larger angle part where we only measure |phi| > rad
0056   Gaudi::Property<double> m_thetaMinFullOMD{this, "thetaMinFullOMD", 0.};
0057   Gaudi::Property<double> m_thetaMaxFullOMD{this, "thetaMaxFullOMD", 2e-3};
0058   Gaudi::Property<double> m_thetaMinPartialOMD{this, "thetaMinPartialOMD", 2.0e-3};
0059   Gaudi::Property<double> m_thetaMaxPartialOMD{this, "thetaMaxPartialOMD", 5.0e-3};
0060   Gaudi::Property<double> m_pMinRigidityOMD{this, "pMinRigidityOMD", 0.25};
0061   Gaudi::Property<double> m_pMaxRigidityOMD{this, "pMaxRigidityOMD", 0.60};
0062 
0063   // Crossing angle, set to -25mrad
0064   Gaudi::Property<double> m_crossingAngle{this, "crossingAngle",
0065                                           -0.025}; //-0.025}; -- causes double rotation with afterburner
0066 
0067   Rndm::Numbers m_gaussDist;
0068 
0069   using RecPart = edm4eic::MutableReconstructedParticle;
0070   using Assoc   = edm4eic::MutableMCRecoParticleAssociation;
0071   using RecData = std::pair<RecPart, Assoc>;
0072 
0073 public:
0074   SmearedFarForwardParticles(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0075     declareProperty("inputMCParticles", m_inputMCParticles, "MCParticles");
0076     declareProperty("outputParticles", m_outputParticles, "ReconstructedParticles");
0077     declareProperty("outputAssociations", m_outputAssocCollection, "MCRecoParticleAssociation");
0078   }
0079   StatusCode initialize() override {
0080     if (GaudiAlgorithm::initialize().isFailure()) {
0081       return StatusCode::FAILURE;
0082     }
0083     IRndmGenSvc* randSvc = svc<IRndmGenSvc>("RndmGenSvc", true);
0084     // use 0 for mean and 1 for standard deviation. Can rescale appropriately for the
0085     // different subsystems
0086     StatusCode sc = m_gaussDist.initialize(randSvc, Rndm::Gauss(0.0, 1.0));
0087     if (!sc.isSuccess()) {
0088       return StatusCode::FAILURE;
0089     }
0090     return StatusCode::SUCCESS;
0091   }
0092   StatusCode execute() override {
0093     const auto& mc = *(m_inputMCParticles.get());
0094     auto& rc       = *(m_outputParticles.createAndPut());
0095     auto& assoc    = *(m_outputAssocCollection.createAndPut());
0096 
0097     double ionBeamEnergy = 0;
0098     if (m_ionBeamEnergy > 0) {
0099       ionBeamEnergy = m_ionBeamEnergy;
0100     } else {
0101       for (const auto& part : mc) {
0102         if (part.getGeneratorStatus() == 4 && part.getPDG() == 2212) {
0103           auto E = part.getEnergy();
0104           if (33 < E && E < 50) {
0105             ionBeamEnergy = 41;
0106           } else if (80 < E && E < 120) {
0107             ionBeamEnergy = 100;
0108           } else if (220 < E && E < 330) {
0109             ionBeamEnergy = 275;
0110           } else {
0111             warning() << "Ion beam energy " << E << " not a standard setting." << endmsg;
0112             ionBeamEnergy = E;
0113           }
0114           break;
0115         }
0116       }
0117       if (ionBeamEnergy == 0) {
0118         warning() << "No incoming ion beam; using 100 GeV ion beam energy." << endmsg;
0119         ionBeamEnergy = 100;
0120       }
0121     }
0122 
0123     std::vector<std::vector<RecData>> rc_parts;
0124     if (m_enableZDC) {
0125       rc_parts.push_back(zdc(mc, ionBeamEnergy));
0126     }
0127     if (m_enableRP) {
0128       rc_parts.push_back(rp(mc, ionBeamEnergy));
0129     }
0130     if (m_enableB0) {
0131       rc_parts.push_back(b0(mc, ionBeamEnergy));
0132     }
0133     if (m_enableOMD) {
0134       rc_parts.push_back(omd(mc, ionBeamEnergy));
0135     }
0136     for (const auto& det : rc_parts) {
0137       for (const auto& [part, link] : det) {
0138         rc.push_back(part);
0139         assoc.push_back(link);
0140       }
0141     }
0142     return StatusCode::SUCCESS;
0143   }
0144 
0145 private:
0146   // ZDC smearing as in eic_smear
0147   // https://github.com/eic/eicsmeardetectors/blob/9a1831dd97bf517b80a06043b9ee4bfb96b483d8/SmearMatrixDetector_0_1_FF.cxx#L224
0148   std::vector<RecData> zdc(const edm4hep::MCParticleCollection& mc, const double /* ionBeamEnergy */) {
0149     std::vector<RecData> rc;
0150     for (const auto& part : mc) {
0151       if (part.getGeneratorStatus() > 1) {
0152         if (msgLevel(MSG::DEBUG)) {
0153           debug() << "ignoring particle with generatorStatus = " << part.getGeneratorStatus() << endmsg;
0154         }
0155         continue;
0156       }
0157       // only detect neutrons and photons
0158       const auto mom_ion = rotateLabToIonDirection(part.getMomentum());
0159       if (part.getPDG() != 2112 && part.getPDG() != 22) {
0160         continue;
0161       }
0162       // only 0-->4.5 mrad
0163       const double mom_ion_theta = edm4hep::utils::anglePolar(mom_ion);
0164       const double mom_ion_phi   = edm4hep::utils::angleAzimuthal(mom_ion);
0165       if (mom_ion_theta > 4.5 / 1000.) {
0166         continue;
0167       }
0168 
0169       double conTerm = 0.05;  // default 5%
0170       double stoTerm = 0.5;   // default 50%
0171       double angTerm = 0.003; // 3mrad
0172 
0173       if (part.getPDG() == 2112) {
0174         conTerm = 0.05;                 // default 5%
0175         stoTerm = 0.5;                  // default 50%
0176         angTerm = 0.003;                // 3mrad
0177       } else if (part.getPDG() == 22) { // EMCAL expected to have slightly better performance
0178         conTerm = 0.03;                 // default 3%
0179         stoTerm = 0.10;                 // default 10% for WSciFi
0180         angTerm = 0.001;                // 1mrad is the detault for the block size
0181       }
0182 
0183       // explicit double precision due to E*E - m*m
0184       const double E    = part.getEnergy();
0185       const double dE   = sqrt((conTerm * E) * (conTerm * E) + stoTerm * stoTerm * E) * m_gaussDist(); // 50%/SqrtE + 5%
0186       const double Es   = E + dE;
0187       const double th   = mom_ion_theta;
0188       const double dth  = (angTerm / sqrt(E)) * m_gaussDist();
0189       const double ths  = th + dth;
0190       const double phi  = mom_ion_phi;
0191       const double dphi = 0;
0192       const double phis = phi + dphi;
0193       const double moms = sqrt(Es * Es - part.getMass() * part.getMass());
0194       // now cast back into float
0195       const auto mom3s_ion = edm4hep::utils::sphericalToVector(moms, ths, phis);
0196       const auto mom3s     = rotateIonToLabDirection(mom3s_ion);
0197       RecPart rec_part;
0198       rec_part.setType(kTagZDC);
0199       rec_part.setEnergy(static_cast<float>(Es));
0200       rec_part.setMomentum({mom3s.x, mom3s.y, mom3s.z});
0201       rec_part.setReferencePoint({static_cast<float>(part.getVertex().x), static_cast<float>(part.getVertex().y),
0202                                   static_cast<float>(part.getVertex().z)});
0203       rec_part.setCharge(static_cast<int16_t>(part.getCharge()));
0204       rec_part.setMass(static_cast<float>(part.getMass()));
0205       rec_part.setGoodnessOfPID(1.);
0206       rec_part.setPDG(part.getPDG());
0207       Assoc assoc;
0208       assoc.setRecID(rec_part.getObjectID().index);
0209       assoc.setSimID(part.getObjectID().index);
0210       assoc.setWeight(1.);
0211       assoc.setRec(rec_part);
0212       //assoc.setSim(part);
0213 
0214       // rec_part.mcID();
0215       rc.emplace_back(rec_part, assoc);
0216 
0217       if (msgLevel(MSG::DEBUG)) {
0218         const auto& part_p    = part.getMomentum();
0219         const auto part_p_mag = std::hypot(part_p.x, part_p.y, part_p.z);
0220         debug()
0221             << fmt::format(
0222                    "Found ZDC particle: {}, Etrue: {}, Emeas: {}, ptrue: {}, pmeas: {}, theta_true: {}, theta_meas: {}",
0223                    part.getPDG(), E, rec_part.getEnergy(), part_p_mag, edm4hep::utils::magnitude(rec_part.getMomentum()), th,
0224                    edm4hep::utils::anglePolar(rec_part.getMomentum()))
0225             << endmsg;
0226       }
0227     }
0228     return rc;
0229   }
0230   // Fast B0 as in
0231   // https://github.com/eic/eicsmeardetectors/blob/9a1831dd97bf517b80a06043b9ee4bfb96b483d8/SmearMatrixDetector_0_1_FF.cxx#L254
0232   std::vector<RecData> b0(const edm4hep::MCParticleCollection& mc, const double /* ionBeamEnergy */) {
0233     std::vector<RecData> rc;
0234     for (const auto& part : mc) {
0235       if (part.getGeneratorStatus() > 1) {
0236         if (msgLevel(MSG::DEBUG)) {
0237           debug() << "ignoring particle with getGeneratorStatus = " << part.getGeneratorStatus() << endmsg;
0238         }
0239         continue;
0240       }
0241       // only detect charged hadrons and photons
0242       if (part.getPDG() != 2212 && part.getPDG() != -2212 && part.getPDG() != 211 && part.getPDG() != -211 &&
0243           part.getPDG() != 321 && part.getPDG() != -321 && part.getPDG() != 22) {
0244         continue;
0245       }
0246       // only 6-->20 mrad
0247       const auto mom_ion = removeCrossingAngle(part.getMomentum()); // rotateLabToIonDirection(part.getMomentum());
0248       const auto mom_ion_theta = edm4hep::utils::anglePolar(mom_ion);
0249       if (mom_ion_theta < m_thetaMinB0 || mom_ion_theta > m_thetaMaxB0) {
0250         continue;
0251       }
0252       auto [rc_part, assoc] = smearMomentum(part);
0253       // we don't detect photon energy, just its angles and presence
0254       if (part.getPDG() == 22) {
0255         rc_part.setMomentum({0, 0, 0});
0256         rc_part.setEnergy(0);
0257       }
0258       rc_part.setType(kTagB0);
0259       rc.emplace_back(rc_part, assoc);
0260       if (msgLevel(MSG::DEBUG)) {
0261         const auto& part_p      = part.getMomentum();
0262         const auto part_p_pt    = edm4hep::utils::magnitudeTransverse(part_p);
0263         const auto part_p_mag   = edm4hep::utils::magnitude(part_p);
0264         const auto part_p_theta = edm4hep::utils::anglePolar(part_p);
0265         debug() << fmt::format("Found B0 particle: {}, ptrue: {}, pmeas: {}, pttrue: {}, ptmeas: {}, theta_true: {}, "
0266                                "theta_meas: {}",
0267                                part.getPDG(), part_p_mag, edm4hep::utils::magnitude(rc_part.momentum()), part_p_pt,
0268                                edm4hep::utils::magnitudeTransverse(rc_part.momentum()), part_p_theta,
0269                                edm4hep::utils::anglePolar(rc_part.momentum()))
0270                 << endmsg;
0271       }
0272     }
0273 
0274     return rc;
0275   }
0276 
0277   std::vector<RecData> rp(const edm4hep::MCParticleCollection& mc, const double ionBeamEnergy) {
0278     std::vector<RecData> rc;
0279     for (const auto& part : mc) {
0280       if (part.getGeneratorStatus() > 1) {
0281         if (msgLevel(MSG::DEBUG)) {
0282           debug() << "ignoring particle with getGeneratorStatus = " << part.getGeneratorStatus() << endmsg;
0283         }
0284         continue;
0285       }
0286       // only detect protons
0287       if (part.getPDG() != 2212) {
0288         continue;
0289       }
0290       const auto mom_ion = removeCrossingAngle(part.getMomentum()); // rotateLabToIonDirection(part.getMomentum());
0291       const auto mom_ion_theta = edm4hep::utils::anglePolar(mom_ion);
0292       if (mom_ion_theta < m_thetaMinRP || mom_ion_theta > m_thetaMaxRP ||
0293           mom_ion.z < m_pMinRigidityRP * ionBeamEnergy) {
0294         continue;
0295       }
0296       auto [rc_part, assoc] = smearMomentum(part);
0297       rc_part.setType(kTagRP);
0298       rc.emplace_back(rc_part, assoc);
0299       if (msgLevel(MSG::DEBUG)) {
0300         const auto& part_p      = part.getMomentum();
0301         const auto part_p_pt    = edm4hep::utils::magnitudeTransverse(part_p);
0302         const auto part_p_mag   = edm4hep::utils::magnitude(part_p);
0303         const auto part_p_theta = edm4hep::utils::anglePolar(part_p);
0304         debug() << fmt::format("Found RP particle: {}, ptrue: {}, pmeas: {}, pttrue: {}, ptmeas: {}, theta_true: {}, "
0305                                "theta_meas: {}",
0306                                part.getPDG(), part_p_mag, edm4hep::utils::magnitude(rc_part.momentum()), part_p_pt,
0307                                edm4hep::utils::magnitudeTransverse(rc_part.momentum()), part_p_theta,
0308                                edm4hep::utils::anglePolar(rc_part.momentum()))
0309                 << endmsg;
0310       }
0311     }
0312     return rc;
0313   }
0314 
0315   std::vector<RecData> omd(const edm4hep::MCParticleCollection& mc, const double ionBeamEnergy) {
0316     std::vector<RecData> rc;
0317     for (const auto& part : mc) {
0318       if (part.getGeneratorStatus() > 1) {
0319         if (msgLevel(MSG::DEBUG)) {
0320           debug() << "ignoring particle with getGeneratorStatus = " << part.getGeneratorStatus() << endmsg;
0321         }
0322         continue;
0323       }
0324       // only detect protons
0325       if (part.getPDG() != 2212) {
0326         continue;
0327       }
0328       const auto mom_ion = removeCrossingAngle(part.getMomentum()); // rotateLabToIonDirection(part.getMomentum());
0329       if (mom_ion.z < m_pMinRigidityOMD * ionBeamEnergy || mom_ion.z > m_pMaxRigidityOMD * ionBeamEnergy) {
0330         continue;
0331       }
0332       auto [rc_part, assoc] = smearMomentum(part);
0333       rc_part.setType(kTagOMD);
0334       rc.emplace_back(rc_part, assoc);
0335       if (msgLevel(MSG::DEBUG)) {
0336         const auto& part_p      = part.getMomentum();
0337         const auto part_p_pt    = edm4hep::utils::magnitudeTransverse(part_p);
0338         const auto part_p_mag   = edm4hep::utils::magnitude(part_p);
0339         const auto part_p_theta = edm4hep::utils::anglePolar(part_p);
0340         debug() << fmt::format("Found OMD particle: {}, ptrue: {}, pmeas: {}, pttrue: {}, ptmeas: {}, theta_true: {}, "
0341                                "theta_meas: {}",
0342                                part.getPDG(), part_p_mag, edm4hep::utils::magnitude(rc_part.momentum()), part_p_pt,
0343                                edm4hep::utils::magnitudeTransverse(rc_part.momentum()), part_p_theta,
0344                                edm4hep::utils::anglePolar(rc_part.momentum()))
0345                 << endmsg;
0346       }
0347     }
0348     return rc;
0349   }
0350 
0351   // all momentum smearing in EIC-smear for the far-forward region uses
0352   // the same 2 relations for P and Pt smearing (B0, RP, OMD)
0353   RecData smearMomentum(const edm4hep::MCParticle& part) {
0354     const auto mom_ion = rotateLabToIonDirection(part.getMomentum());
0355     const double p     = std::hypot(mom_ion.x, mom_ion.y, mom_ion.z);
0356     const double dp    = (0.025 * p) * m_gaussDist();
0357     const double ps    = p + dp;
0358 
0359     // const double pt  = std::hypot(mom_ion.x, mom_ion.y);
0360     // const double dpt = (0.03 * pt) * m_gaussDist();
0361     // just apply relative smearing on px and py
0362     const double dpxs = (0.03 * mom_ion.x) * m_gaussDist(); //+ (1 + dpt / pt);
0363     const double dpys = (0.03 * mom_ion.y) * m_gaussDist(); //+ (1 + dpt / pt);
0364 
0365     const double pxs = mom_ion.x + dpxs;
0366     const double pys = mom_ion.y + dpys;
0367 
0368     // now get pz
0369     const double pzs = sqrt(ps * ps - pxs * pxs - pys * pys);
0370 
0371     // And build our 3-vector
0372     const edm4hep::Vector3f psmear_ion{static_cast<float>(pxs), static_cast<float>(pys), static_cast<float>(pzs)};
0373     const auto psmear = rotateIonToLabDirection(psmear_ion);
0374     edm4eic::MutableReconstructedParticle rec_part;
0375     rec_part.setType(-1);
0376     rec_part.setEnergy(std::hypot(ps, part.getMass()));
0377     rec_part.setMomentum({psmear.x, psmear.y, psmear.z});
0378     rec_part.setReferencePoint({static_cast<float>(part.getVertex().x), static_cast<float>(part.getVertex().y),
0379                                 static_cast<float>(part.getVertex().z)});
0380     rec_part.setCharge(static_cast<int16_t>(part.getCharge()));
0381     rec_part.setMass(static_cast<float>(part.getMass()));
0382     rec_part.setGoodnessOfPID(1); // perfect PID
0383     rec_part.setPDG(part.getPDG());
0384     Assoc assoc;
0385     assoc.setRecID(rec_part.getObjectID().index);
0386     assoc.setSimID(part.getObjectID().index);
0387     assoc.setWeight(1.);
0388     assoc.setRec(rec_part);
0389     //assoc.setSim(part);
0390 
0391     return {rec_part, assoc};
0392   }
0393 
0394   // Rotate 25mrad about the y-axis
0395   edm4hep::Vector3f rotateLabToIonDirection(const edm4hep::Vector3f& vec) const {
0396     const auto sth = sin(-m_crossingAngle);
0397     const auto cth = cos(-m_crossingAngle);
0398     return {static_cast<float>(cth * vec.x + sth * vec.z), static_cast<float>(vec.y),
0399             static_cast<float>(-sth * vec.x + cth * vec.z)};
0400   }
0401 
0402   edm4hep::Vector3f rotateIonToLabDirection(const edm4hep::Vector3f& vec) const {
0403     const auto sth = sin(m_crossingAngle);
0404     const auto cth = cos(m_crossingAngle);
0405     return {static_cast<float>(cth * vec.x + sth * vec.z), static_cast<float>(vec.y),
0406             static_cast<float>(-sth * vec.x + cth * vec.z)};
0407   }
0408 
0409   edm4hep::Vector3f removeCrossingAngle(const edm4hep::Vector3f& vec) const {
0410     const auto sth = std::sin(-m_crossingAngle);
0411     const auto cth = std::cos(-m_crossingAngle);
0412     return {static_cast<float>(cth * vec.x + sth * vec.z), static_cast<float>(vec.y),
0413             static_cast<float>(-sth * vec.x + cth * vec.z)};
0414   }
0415 };
0416 
0417 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0418 DECLARE_COMPONENT(SmearedFarForwardParticles)
0419 
0420 } // namespace Jug::Fast