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

File indexing completed on 2026-06-17 07:50:49

 // SPDX-License-Identifier: LGPL-3.0-or-later
 // Copyright (C) 2026 Sebouh Paul, Baptiste Fraisse
 
 #include <Evaluator/DD4hepUnits.h>
 #include <TVector3.h>
 #include <edm4eic/Cluster.h>
 #include <edm4eic/ReconstructedParticleCollection.h>
 #include <edm4eic/Vertex.h>
 #include <edm4hep/Vector3f.h>
 #include <stdint.h>
 #include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstddef>
 #include <stdexcept>
 #include <tuple>
 #include <vector>
 
 #include "LambdaReconstruction.h"
 #include "TLorentzVector.h"
 
 namespace eicrecon {
 
 void LambdaReconstruction::init() {
   try {
     m_zMax = m_detector->constant<double>(m_cfg.offsetPositionName);
   } catch (std::runtime_error&) {
     m_zMax = 35800;
     trace("Failed to get {} from the detector, using default value of {}", m_cfg.offsetPositionName,
           m_zMax);
   }
 }
 
 // reconstruction machinery from n+g+g
 
 bool LambdaReconstruction::reconstruct_from_triplet(
     const edm4eic::ReconstructedParticle& n_in, const edm4eic::ReconstructedParticle& g1_in,
     const edm4eic::ReconstructedParticle& g2_in,
     edm4eic::ReconstructedParticleCollection* out_lambdas,
     edm4eic::ReconstructedParticleCollection* out_decay_products) const {
   static const double m_neutron = m_particleSvc.particle(2112).mass;
   static const double m_pi0     = m_particleSvc.particle(111).mass;
   static const double m_lambda  = m_particleSvc.particle(3122).mass;
 
   const double En = n_in.getEnergy();
   const double E1 = g1_in.getEnergy();
   const double E2 = g2_in.getEnergy();
 
   if (En <= m_neutron) {
     return false;
   }
   const double pn = std::sqrt(std::max(0.0, En * En - m_neutron * m_neutron));
 
   TVector3 xn;
   TVector3 x1;
   TVector3 x2;
 
   if (n_in.clusters_size() == 0 || g1_in.clusters_size() == 0 || g2_in.clusters_size() == 0) {
     return false;
   }
 
   const auto rn  = n_in.getClusters(0).getPosition();
   const auto rg1 = g1_in.getClusters(0).getPosition();
   const auto rg2 = g2_in.getClusters(0).getPosition();
 
   xn.SetXYZ(rn.x * dd4hep::mm, rn.y * dd4hep::mm, rn.z * dd4hep::mm);
   x1.SetXYZ(rg1.x * dd4hep::mm, rg1.y * dd4hep::mm, rg1.z * dd4hep::mm);
   x2.SetXYZ(rg2.x * dd4hep::mm, rg2.y * dd4hep::mm, rg2.z * dd4hep::mm);
 
   xn.RotateY(-m_cfg.globalToProtonRotation);
   x1.RotateY(-m_cfg.globalToProtonRotation);
   x2.RotateY(-m_cfg.globalToProtonRotation);
 
   TVector3 vtx(0, 0, 0);
   double f  = 0.0;
   double df = 0.5;
 
   const double theta_open_expected = 2 * std::asin(m_pi0 / (2 * std::sqrt(E1 * E2)));
 
   TLorentzVector n;
   TLorentzVector g1;
   TLorentzVector g2;
   TLorentzVector lambda;
 
   for (int i = 0; i < m_cfg.iterations; i++) {
     n      = {pn * (xn - vtx).Unit(), En};
     g1     = {E1 * (x1 - vtx).Unit(), E1};
     g2     = {E2 * (x2 - vtx).Unit(), E2};
     lambda = n + g1 + g2;
 
     const double theta_open = g1.Angle(g2.Vect());
     if (theta_open > theta_open_expected) {
       f -= df;
     } else if (theta_open < theta_open_expected) {
       f += df;
     }
 
     vtx = lambda.Vect() * (f * m_zMax / lambda.Z());
     df /= 2.0;
   }
 
   const double mass_rec = lambda.M();
   if (std::abs(mass_rec - m_lambda) > m_cfg.lambdaMassWindow * m_lambda) {
     return false;
   }
 
   // rotate everything back to lab coordinates
   vtx.RotateY(m_cfg.globalToProtonRotation);
   lambda.RotateY(m_cfg.globalToProtonRotation);
   n.RotateY(m_cfg.globalToProtonRotation);
   g1.RotateY(m_cfg.globalToProtonRotation);
   g2.RotateY(m_cfg.globalToProtonRotation);
 
   // boost decay products to Lambda rest frame
   auto b = -lambda.BoostVector();
   n.Boost(b);
   g1.Boost(b);
   g2.Boost(b);
 
   // vertex back to EDM units (mm)
   vtx = vtx * (1.0 / dd4hep::mm);
 
   // saving reco lambda
   auto rec_lambda = out_lambdas->create();
   rec_lambda.setPDG(3122);
   rec_lambda.setEnergy(lambda.E());
   rec_lambda.setMomentum({static_cast<float>(lambda.X()), static_cast<float>(lambda.Y()),
                           static_cast<float>(lambda.Z())});
   rec_lambda.setReferencePoint(
       {static_cast<float>(vtx.X()), static_cast<float>(vtx.Y()), static_cast<float>(vtx.Z())});
   rec_lambda.setCharge(0);
   rec_lambda.setMass(mass_rec);
 
   // --- cm neutron
   auto neutron_cm = out_decay_products->create();
   neutron_cm.setPDG(2112);
   neutron_cm.setEnergy(n.E());
   neutron_cm.setMomentum(
       {static_cast<float>(n.X()), static_cast<float>(n.Y()), static_cast<float>(n.Z())});
   neutron_cm.setReferencePoint(
       {static_cast<float>(vtx.X()), static_cast<float>(vtx.Y()), static_cast<float>(vtx.Z())});
   neutron_cm.setCharge(0);
   neutron_cm.setMass(m_neutron);
 
   // --- cm gammas
   auto gamma1_cm = out_decay_products->create();
   gamma1_cm.setPDG(22);
   gamma1_cm.setEnergy(g1.E());
   gamma1_cm.setMomentum(
       {static_cast<float>(g1.X()), static_cast<float>(g1.Y()), static_cast<float>(g1.Z())});
   gamma1_cm.setReferencePoint(
       {static_cast<float>(vtx.X()), static_cast<float>(vtx.Y()), static_cast<float>(vtx.Z())});
   gamma1_cm.setCharge(0);
   gamma1_cm.setMass(0);
 
   auto gamma2_cm = out_decay_products->create();
   gamma2_cm.setPDG(22);
   gamma2_cm.setEnergy(g2.E());
   gamma2_cm.setMomentum(
       {static_cast<float>(g2.X()), static_cast<float>(g2.Y()), static_cast<float>(g2.Z())});
   gamma2_cm.setReferencePoint(
       {static_cast<float>(vtx.X()), static_cast<float>(vtx.Y()), static_cast<float>(vtx.Z())});
   gamma2_cm.setCharge(0);
   gamma2_cm.setMass(0);
 
   rec_lambda.addToParticles(n_in);
   rec_lambda.addToParticles(g1_in);
   rec_lambda.addToParticles(g2_in);
 
   neutron_cm.addToParticles(n_in);
   gamma1_cm.addToParticles(g1_in);
   gamma2_cm.addToParticles(g2_in);
 
   return true;
 }
 
 void LambdaReconstruction::process(const LambdaReconstruction::Input& input,
                                    const LambdaReconstruction::Output& output) const {
   const auto [neutralsHcal, neutralsB0, neutralsEcalEndcapP, neutralsLFHCAL] = input;
   auto [out_lambdas, out_decay_products]                                     = output;
 
   const double m_pi0    = m_particleSvc.particle(111).mass;
   const double m_lambda = m_particleSvc.particle(3122).mass;
 
   // --------------------------------------------------------------------------
   // Step 1: collect neutral candidates by broad detector category
   //
   // We keep two categories for ranking purposes:
   //   - ZDC-preferred neutrals
   //   - other far-forward neutrals
   // while centralizing the detector-specific collection handling in a single
   // local table to avoid duplicating hardcoded loops throughout the algorithm.
   // --------------------------------------------------------------------------
 
   std::vector<edm4eic::ReconstructedParticle> neutrons_zdc;
   std::vector<edm4eic::ReconstructedParticle> neutrons_other;
   std::vector<edm4eic::ReconstructedParticle> gammas_zdc;
   std::vector<edm4eic::ReconstructedParticle> gammas_other;
 
   struct NeutralInputDescription {
     const edm4eic::ReconstructedParticleCollection* coll;
     bool use_gamma;
     bool use_neutron;
     bool is_zdc;
   };
 
   const std::array<NeutralInputDescription, 4> input_descs{{
       {neutralsHcal, true, true, true},
       {neutralsB0, true, false, false},
       {neutralsEcalEndcapP, true, false, false},
       {neutralsLFHCAL, false, true, false},
   }};
 
   for (const auto& desc : input_descs) {
     for (const auto& part : *desc.coll) {
       if (desc.use_neutron && part.getPDG() == 2112) {
         (desc.is_zdc ? neutrons_zdc : neutrons_other).push_back(part);
       }
       if (desc.use_gamma && part.getPDG() == 22) {
         (desc.is_zdc ? gammas_zdc : gammas_other).push_back(part);
       }
     }
   }
 
   if (neutrons_zdc.empty() && neutrons_other.empty()) {
     debug("No neutrons in input collection");
     return;
   }
 
   // --------------------------------------------------------------------------
   // Step 2: build a unified photon pool while retaining whether each photon
   // comes from the ZDC-preferred category.
   // --------------------------------------------------------------------------
 
   using ParticleT = edm4eic::ReconstructedParticle;
 
   std::vector<ParticleT> gamma_pool;
   std::vector<bool> gamma_is_zdc;
 
   gamma_pool.reserve(gammas_zdc.size() + gammas_other.size());
   gamma_is_zdc.reserve(gammas_zdc.size() + gammas_other.size());
 
   for (const auto& g : gammas_zdc) {
     gamma_pool.push_back(g);
     gamma_is_zdc.push_back(1);
   }
   for (const auto& g : gammas_other) {
     gamma_pool.push_back(g);
     gamma_is_zdc.push_back(0);
   }
 
   if (gamma_pool.size() < 2) {
     debug("Less than 2 gammas found ({})", gamma_pool.size());
     return;
   }
 
   // --------------------------------------------------------------------------
   // Invariant-mass helpers
   // --------------------------------------------------------------------------
 
   auto invMass2 = [](const ParticleT& a, const ParticleT& b) {
     const auto pa = a.getMomentum();
     const auto pb = b.getMomentum();
 
     const double Ea = a.getEnergy();
     const double Eb = b.getEnergy();
 
     const double px = pa.x + pb.x;
     const double py = pa.y + pb.y;
     const double pz = pa.z + pb.z;
     const double E  = Ea + Eb;
 
     return E * E - (px * px + py * py + pz * pz);
   };
 
   auto invMass2_3 = [](const ParticleT& a, const ParticleT& b, const ParticleT& c) {
     const auto pa = a.getMomentum();
     const auto pb = b.getMomentum();
     const auto pc = c.getMomentum();
 
     const double Ea = a.getEnergy();
     const double Eb = b.getEnergy();
     const double Ec = c.getEnergy();
 
     const double px = pa.x + pb.x + pc.x;
     const double py = pa.y + pb.y + pc.y;
     const double pz = pa.z + pb.z + pc.z;
     const double E  = Ea + Eb + Ec;
 
     return E * E - (px * px + py * py + pz * pz);
   };
 
   // --------------------------------------------------------------------------
   // Step 3: build pi0 candidates from photon pairs
   // --------------------------------------------------------------------------
 
   enum class GammaCategory : uint8_t { TwoZDC = 0, OneZDC = 1, ZeroZDC = 2 };
 
   struct Pi0Pair {
     int i;
     int j;
     GammaCategory cat;
     double dmpi0_cand;
   };
 
   std::vector<Pi0Pair> pi0_pairs;
   pi0_pairs.reserve(gamma_pool.size() * (gamma_pool.size() - 1) / 2);
 
   for (size_t i = 0; i < gamma_pool.size(); ++i) {
     for (size_t j = i + 1; j < gamma_pool.size(); ++j) {
       const auto& g1 = gamma_pool[i];
       const auto& g2 = gamma_pool[j];
 
       const double m2 = invMass2(g1, g2);
       if (m2 <= 0.0) {
         continue;
       }
 
       const double m  = std::sqrt(m2);
       const double dm = std::abs(m - m_pi0);
       if (dm > m_cfg.pi0Window * m_pi0) {
         continue;
       }
 
       const bool z1 = gamma_is_zdc[i];
       const bool z2 = gamma_is_zdc[j];
 
       const GammaCategory cat = (z1 && z2)   ? GammaCategory::TwoZDC
                                 : (z1 || z2) ? GammaCategory::OneZDC
                                              : GammaCategory::ZeroZDC;
 
       pi0_pairs.push_back({static_cast<int>(i), static_cast<int>(j), cat, dm});
     }
   }
 
   if (pi0_pairs.empty()) {
     return;
   }
 
   // --------------------------------------------------------------------------
   // Step 4: build Lambda candidate pool from pi0-neutron combinations
   // --------------------------------------------------------------------------
 
   enum class NeutronCategory : uint8_t { ZDC = 0, Other = 1 };
 
   struct LambdaCandidate {
     int g_i;
     int g_j;
     int n_idx;
     NeutronCategory n_cat;
     GammaCategory g_cat;
     double chi2;
     double E;
     double pz;
   };
 
   std::vector<LambdaCandidate> cands;
   cands.reserve(64); // NB <= 4^3 candidates when searching for 3 daughters across 4 calos
 
   auto try_neutrons = [&](const auto& neutron_list, NeutronCategory n_cat) {
     for (const auto& pp : pi0_pairs) {
       const auto& g1 = gamma_pool[pp.i];
       const auto& g2 = gamma_pool[pp.j];
 
       for (size_t in = 0; in < neutron_list.size(); ++in) {
         const auto& n = neutron_list[in];
 
         const double m2L = invMass2_3(n, g1, g2);
         if (m2L <= 0.0) {
           continue;
         }
 
         const double mL = std::sqrt(m2L);
         const double dL = mL - m_lambda;
 
         if (std::abs(dL) > m_cfg.lambdaMassWindow * m_lambda) {
           continue;
         }
 
         const double pi0_term    = pp.dmpi0_cand / (m_cfg.pi0Window * m_pi0);
         const double lambda_term = dL / (m_cfg.lambdaMassWindow * m_lambda);
 
         const double chi2 = pi0_term * pi0_term + lambda_term * lambda_term;
 
         const double E = n.getEnergy() + g1.getEnergy() + g2.getEnergy();
 
         const auto pn  = n.getMomentum();
         const auto pg1 = g1.getMomentum();
         const auto pg2 = g2.getMomentum();
 
         const double pz = pn.z + pg1.z + pg2.z;
 
         cands.push_back({pp.i, pp.j, static_cast<int>(in), n_cat, pp.cat, chi2, E, pz});
       }
     }
   };
 
   if (!neutrons_zdc.empty()) {
     try_neutrons(neutrons_zdc, NeutronCategory::ZDC);
   }
   if (!neutrons_other.empty()) {
     try_neutrons(neutrons_other, NeutronCategory::Other);
   }
 
   if (cands.empty()) {
     return;
   }
 
   // --------------------------------------------------------------------------
   // Step 5: rank candidates and keep the best one
   // Preference order:
   //   1. ZDC neutron over other neutron
   //   2. more ZDC photons in the pi0 candidate
   //   3. lower combined mass-based chi2
   //   4. larger forward momentum pz
   //   5. larger total energy
   // --------------------------------------------------------------------------
 
   auto better = [&](const LambdaCandidate& a, const LambdaCandidate& b) -> bool {
     if (a.n_cat != b.n_cat) {
       return static_cast<int>(a.n_cat) < static_cast<int>(b.n_cat);
     }
     if (a.g_cat != b.g_cat) {
       return static_cast<int>(a.g_cat) < static_cast<int>(b.g_cat);
     }
     if (a.chi2 != b.chi2) {
       return a.chi2 < b.chi2;
     }
     if (a.pz != b.pz) {
       return a.pz > b.pz;
     }
     return a.E > b.E;
   };
 
   int best_k = 0;
   for (int k = 1; k < static_cast<int>(cands.size()); ++k) {
     if (better(cands[k], cands[best_k])) {
       best_k = k;
     }
   }
 
   const auto& best = cands[best_k];
 
   const auto& g1 = gamma_pool[best.g_i];
   const auto& g2 = gamma_pool[best.g_j];
   const auto& n =
       (best.n_cat == NeutronCategory::ZDC) ? neutrons_zdc[best.n_idx] : neutrons_other[best.n_idx];
 
   reconstruct_from_triplet(n, g1, g2, out_lambdas, out_decay_products);
 }
 
 } // namespace eicrecon

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