Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-11-07 09:30:55

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2026 Chun Yuen Tsang, Minho Kim
0003 
0004 #include "CalorimeterCALOROCReco.h"
0005 
0006 #include <DD4hep/Alignments.h>
0007 #include <DD4hep/IDDescriptor.h>
0008 #include <DD4hep/Objects.h>
0009 #include <DD4hep/Readout.h>
0010 #include <DD4hep/Segmentations.h>
0011 #include <DD4hep/Shapes.h>
0012 #include <DD4hep/VolumeManager.h>
0013 #include <DD4hep/Volumes.h>
0014 #include <DD4hep/detail/SegmentationsInterna.h>
0015 #include <DDSegmentation/BitFieldCoder.h>
0016 #include <DDSegmentation/MultiSegmentation.h>
0017 #include <DDSegmentation/Segmentation.h>
0018 #include <Evaluator/DD4hepUnits.h>
0019 #include <Math/GenVector/Cartesian3D.h>
0020 #include <Math/GenVector/DisplacementVector3D.h>
0021 #include <edm4eic/CALOROC1BSample.h>
0022 #include <edm4eic/unit_system.h>
0023 #include <edm4hep/CaloHitContribution.h>
0024 #include <edm4hep/MCParticle.h>
0025 #include <edm4hep/Vector3f.h>
0026 #include <fmt/ranges.h>
0027 #include <podio/ObjectID.h>
0028 #include <podio/RelationRange.h>
0029 #include <podio/detail/Link.h>
0030 #include <podio/detail/LinkCollectionImpl.h>
0031 #include <stdint.h>
0032 #include <algorithm>
0033 #include <cctype>
0034 #include <cmath>
0035 #include <fstream>
0036 #include <functional>
0037 #include <initializer_list>
0038 #include <iterator>
0039 #include <memory>
0040 #include <sstream>
0041 #include <string>
0042 #include <tuple>
0043 #include <unordered_map>
0044 #include <utility>
0045 #include <vector>
0046 
0047 #include "algorithms/calorimetry/CalorimeterCALOROCRecoConfig.h"
0048 
0049 using namespace dd4hep;
0050 
0051 namespace eicrecon {
0052 
0053 double CalorimeterCALOROCReco::_energyCor(double referencePos, double energy, double z) const {
0054   double length = std::abs(referencePos - z);
0055   double factor =
0056       m_cfg.attenuationParameters[0] * std::exp(-length / m_cfg.attenuationParameters[1]) +
0057       (1 - m_cfg.attenuationParameters[0]) * std::exp(-length / m_cfg.attenuationParameters[2]);
0058   return energy / factor;
0059 }
0060 
0061 void CalorimeterCALOROCReco::init() {
0062 
0063   if (m_cfg.attenuationReferencePositionNamePos.empty() ||
0064       m_cfg.attenuationReferencePositionNameNeg.empty()) {
0065     error("You MUST provide the name of reference positions for attenuation correction");
0066   }
0067 
0068   m_reference_z_p = m_geo.detector()->constant<double>(m_cfg.attenuationReferencePositionNamePos) *
0069                     edm4eic::unit::mm / dd4hep::mm;
0070   m_reference_z_n = m_geo.detector()->constant<double>(m_cfg.attenuationReferencePositionNameNeg) *
0071                     edm4eic::unit::mm / dd4hep::mm;
0072 
0073   // do not get the layer/sector ID if no readout class provided
0074   if (m_cfg.readout.empty()) {
0075     error("You MUST provide the name of the readout to decode CellID.");
0076   }
0077 
0078   // First, try and get the IDDescriptor. This will throw an exception if it fails.
0079   try {
0080     id_spec = m_detector->readout(m_cfg.readout).idSpec();
0081   } catch (...) {
0082     warning("Failed to get idSpec for {}", m_cfg.readout);
0083     return;
0084   }
0085   // Next, try and get the readout fields. This will throw a different exception.
0086   try {
0087     id_dec = id_spec.decoder();
0088     if (!m_cfg.sectorField.empty()) {
0089       sector_idx = id_dec->index(m_cfg.sectorField);
0090       debug("Find sector field {}, index = {}", m_cfg.sectorField, sector_idx);
0091     }
0092     if (!m_cfg.layerField.empty()) {
0093       layer_idx = id_dec->index(m_cfg.layerField);
0094       debug("Find layer field {}, index = {}", m_cfg.layerField, layer_idx);
0095     }
0096     if (!m_cfg.maskPosFields.empty()) {
0097       std::size_t tmp_mask = 0;
0098       for (auto& field : m_cfg.maskPosFields) {
0099         tmp_mask |= id_spec.field(field)->mask();
0100       }
0101       // assign this mask if all fields succeed
0102       gpos_mask = tmp_mask;
0103     }
0104   } catch (...) {
0105     if (id_dec == nullptr) {
0106       warning("Failed to load ID decoder for {}", m_cfg.readout);
0107       std::stringstream readouts;
0108       for (auto r : m_detector->readouts()) {
0109         readouts << "\"" << r.first << "\", ";
0110       }
0111       warning("Available readouts: {}", readouts.str());
0112     } else {
0113       warning("Failed to find field index for {}.", m_cfg.readout);
0114       if (!m_cfg.sectorField.empty()) {
0115         warning(" -- looking for sector field \"{}\".", m_cfg.sectorField);
0116       }
0117       if (!m_cfg.layerField.empty()) {
0118         warning(" -- looking for layer field  \"{}\".", m_cfg.layerField);
0119       }
0120       if (!m_cfg.maskPosFields.empty()) {
0121         warning(" -- looking for masking fields  \"{}\".", fmt::join(m_cfg.maskPosFields, ", "));
0122       }
0123       std::stringstream fields;
0124       for (auto field : id_spec.decoder()->fields()) {
0125         fields << "\"" << field.name() << "\", ";
0126       }
0127       warning("Available fields: {}", fields.str());
0128       warning("n.b. The local position, sector id and layer id will not be correct for this.");
0129       warning("Position masking may not be applied.");
0130       warning("however, the position, energy, and time values should still be good.");
0131     }
0132 
0133     return;
0134   }
0135 
0136   id_spec = m_detector->readout(m_cfg.readout).idSpec();
0137 
0138   // local detector name has higher priority
0139   if (!m_cfg.localDetElement.empty()) {
0140     try {
0141       m_local = m_detector->detector(m_cfg.localDetElement);
0142     } catch (...) {
0143       error("failed to load local coordinate system from DetElement {}", m_cfg.localDetElement);
0144       return;
0145     }
0146   } else {
0147     std::vector<std::pair<std::string, int>> fields;
0148     for (auto f : m_cfg.localDetFields) {
0149       fields.emplace_back(f, 0);
0150     }
0151     local_mask = id_spec.get_mask(fields);
0152     // use all fields if nothing provided
0153     if (fields.empty()) {
0154       local_mask = ~static_cast<decltype(local_mask)>(0);
0155     }
0156   }
0157 
0158   // Get the field indices and field-dependent conversion factors if necessary
0159   if (m_cfg.edep_to_npe_fields.empty() || m_cfg.edep_to_npe_filename.empty()) {
0160     error("You MUST provide edep_to_npe files and filename for layer-by-layer calibration.");
0161   }
0162 
0163   for (const auto& field : m_cfg.edep_to_npe_fields) {
0164     auto field_idx = id_dec->index(field);
0165     m_field_idxs.push_back(field_idx);
0166   }
0167 
0168   std::ifstream infile(m_cfg.edep_to_npe_filename);
0169   if (!infile) {
0170     error("Unable to open LUT file: {}", m_cfg.edep_to_npe_filename);
0171   }
0172   std::string line;
0173   while (std::getline(infile, line)) {
0174     std::istringstream iss(line);
0175     std::vector<int> keys;
0176     for (std::size_t i = 0; i < m_cfg.edep_to_npe_fields.size(); i++) {
0177       int value;
0178       iss >> value;
0179       keys.push_back(value);
0180     }
0181     double factor;
0182     if (!(iss >> factor)) {
0183       error("Malformed LUT file: {}", m_cfg.edep_to_npe_filename);
0184     }
0185     m_edep_to_npe_lut[keys] = factor;
0186   }
0187 }
0188 
0189 double CalorimeterCALOROCReco::_sumADC(const edm4eic::RawCALOROCHit& ADC) const {
0190   double sum = 0;
0191 
0192   // check high gain ADC first. If it saturates, we switch to lowGainADC
0193   for (const auto& bSample : ADC.getBSamples()) {
0194     if (bSample.highGainADC >= m_cfg.highGainDR) {
0195       sum += static_cast<double>(bSample.lowGainADC);
0196     } else {
0197       sum += static_cast<double>(bSample.highGainADC) / m_cfg.gainRatio;
0198     }
0199   }
0200   return sum;
0201 }
0202 
0203 double CalorimeterCALOROCReco::_toa(const edm4eic::RawCALOROCHit& ADC) const {
0204   // find the relative toa
0205   int samplePhase        = ADC.getSamplePhase();
0206   int timeStamp          = ADC.getTimeStamp();
0207   const auto& bSamples   = ADC.getBSamples();
0208   size_t idx_toa         = 0;
0209   uint16_t timeOfArrival = 0;
0210   for (; idx_toa < bSamples.size(); ++idx_toa) {
0211     timeOfArrival = bSamples.at(idx_toa).timeOfArrival;
0212     if (timeOfArrival > 0) {
0213       break;
0214     }
0215   }
0216 
0217   return samplePhase * (25. / 1042.) + (timeStamp + idx_toa) * 25. - timeOfArrival * (25. / 1024.);
0218 }
0219 
0220 double CalorimeterCALOROCReco::_timeWalkCorrection(double toa, double lowGainADC) const {
0221   if (static_cast<double>(lowGainADC) - m_cfg.timeWalkCorrectionParameters[2] > 0) {
0222     return toa - (m_cfg.timeWalkCorrectionParameters[1] *
0223                       pow(static_cast<double>(lowGainADC) - m_cfg.timeWalkCorrectionParameters[2],
0224                           m_cfg.timeWalkCorrectionParameters[3]) +
0225                   m_cfg.timeWalkCorrectionParameters[0]);
0226   } else {
0227     return toa;
0228   }
0229 }
0230 
0231 void CalorimeterCALOROCReco::process(const CalorimeterCALOROCReco::Input& input,
0232                                      const CalorimeterCALOROCReco::Output& output) const {
0233 
0234   const auto [npeHitsP, ADCPs, npeHitsN, ADCNs]       = input;
0235   auto [recohits, rawhits, rawhitsLink, rawhitsAssoc] = output;
0236 
0237   // match NpeHits and ADC hits by cellID
0238   std::unordered_map<dd4hep::rec::CellID, size_t> cellID2NpeHitNID, cellID2ADCNID, cellID2ADCPID;
0239   for (size_t i = 0; i < npeHitsN->size(); ++i) {
0240     const auto& hit                   = npeHitsN->at(i);
0241     cellID2NpeHitNID[hit.getCellID()] = i;
0242   }
0243 
0244   for (size_t i = 0; i < ADCPs->size(); ++i) {
0245     const auto& ADCP                = ADCPs->at(i);
0246     cellID2ADCPID[ADCP.getCellID()] = i;
0247   }
0248 
0249   for (size_t i = 0; i < ADCNs->size(); ++i) {
0250     const auto& ADCN                = ADCNs->at(i);
0251     cellID2ADCNID[ADCN.getCellID()] = i;
0252   }
0253 
0254   for (const auto& npeHitP : *npeHitsP) {
0255     auto cellID = npeHitP.getCellID();
0256 
0257     // ignore data point if the hit is not showing up on all ADC and NpeHit collections
0258     auto it = cellID2NpeHitNID.find(cellID);
0259     if (it == cellID2NpeHitNID.end()) {
0260       continue;
0261     }
0262     const auto& npeHitN = npeHitsN->at(it->second);
0263     it                  = cellID2ADCPID.find(cellID);
0264     if (it == cellID2ADCPID.end()) {
0265       continue;
0266     }
0267     const auto& ADCP = ADCPs->at(it->second);
0268     it               = cellID2ADCNID.find(cellID);
0269     if (it == cellID2ADCNID.end()) {
0270       continue;
0271     }
0272     const auto& ADCN = ADCNs->at(it->second);
0273 
0274     // get layer and sector ID
0275     const int lid = id_dec != nullptr && !m_cfg.layerField.empty()
0276                         ? static_cast<int>(id_dec->get(cellID, layer_idx))
0277                         : -1;
0278     const int sid = id_dec != nullptr && !m_cfg.sectorField.empty()
0279                         ? static_cast<int>(id_dec->get(cellID, sector_idx))
0280                         : -1;
0281 
0282     auto tP = _toa(ADCP);
0283     auto tN = _toa(ADCN);
0284 
0285     // get position of the hit;
0286     double zpos;
0287 
0288     if (m_cfg.useNpeHitPos) {
0289       zpos = static_cast<double>(npeHitP.getPosition().z);
0290     } else {
0291       if (m_cfg.timeWalkCor) {
0292         tP = _timeWalkCorrection(tP, this->_sumADC(ADCP));
0293         tN = _timeWalkCorrection(tN, this->_sumADC(ADCN));
0294       }
0295       auto dt = tN - tP;
0296       zpos    = dt * m_cfg.lightSpeedParameters[0] + m_cfg.lightSpeedParameters[1];
0297     }
0298     double time = 0.5 * (tP + tN);
0299     // bound zpos by sipm location
0300     zpos = std::max(m_reference_z_n, std::min(m_reference_z_p, zpos));
0301 
0302     // fetch edep to npe factor on this layer
0303     std::vector<int> key;
0304     key.reserve(m_field_idxs.size());
0305     for (const auto idx : m_field_idxs) {
0306       key.push_back(static_cast<int>(id_dec->get(cellID, idx)));
0307     }
0308 
0309     const auto lut_it = m_edep_to_npe_lut.find(key);
0310     if (lut_it == m_edep_to_npe_lut.end() || lut_it->second == 0.0) {
0311       warning("Missing/invalid edep->npe factor for key {} (cellID={})", fmt::join(key, ","),
0312               cellID);
0313       continue;
0314     }
0315     const double eDep2NpeFactor = lut_it->second;
0316     // find averaged energy
0317     double npeP, npeN;
0318     // get values from both sides in a loop
0319     for (bool NSide : {true, false}) {
0320       auto& ADC = NSide ? ADCN : ADCP;
0321       auto& npe = NSide ? npeN : npeP;
0322       switch (m_cfg.proxy_type) {
0323       case CalorimeterCALOROCRecoConfig::ProxyType::sum:
0324         npe = _sumADC(ADC);
0325         break;
0326       case CalorimeterCALOROCRecoConfig::ProxyType::templateFit:
0327         error("Proxy type not implemented.");
0328         break;
0329       case CalorimeterCALOROCRecoConfig::ProxyType::simpson:
0330         error("Proxy type not implemented.");
0331         break;
0332       default:
0333         error("Proxy type not implemented.");
0334       }
0335     }
0336 
0337     // perform layer-by-layer energy correction
0338     double chargeP = npeP / eDep2NpeFactor;
0339     double chargeN = npeN / eDep2NpeFactor;
0340 
0341     // attenuation correction
0342     double corEP = _energyCor(m_reference_z_p, chargeP, zpos) * m_cfg.slope + m_cfg.intercept;
0343     double corEN = _energyCor(m_reference_z_n, chargeN, zpos) * m_cfg.slope + m_cfg.intercept;
0344     double corE  = std::sqrt(corEP * corEN);
0345 
0346     // raw hits for all pulses
0347     edm4hep::MutableRawCalorimeterHit rawhit;
0348     rawhit.setCellID(cellID);
0349     rawhit.setAmplitude(npeP);
0350     rawhit.setTimeStamp(time);
0351 
0352     double edep = 0;
0353 
0354     // link to parents
0355     std::unordered_map<podio::ObjectID, edm4eic::MutableMCRecoCalorimeterHitLink> links_staging;
0356     std::unordered_map<podio::ObjectID, edm4eic::MutableMCRecoCalorimeterHitAssociation>
0357         rawassocs_staging;
0358 
0359     for (bool NSide : {true, false}) {
0360       const auto& npeHit = NSide ? npeHitN : npeHitP;
0361       for (const auto& contrib : npeHit.getContributions()) {
0362         if (links_staging.find(contrib.getParticle().getObjectID()) != links_staging.end()) {
0363           continue;
0364         }
0365 
0366         edep += contrib.getEnergy();
0367         edm4eic::MutableMCRecoCalorimeterHitAssociation assoc;
0368         assoc.setRawHit(rawhit);
0369 
0370         edm4eic::MutableMCRecoCalorimeterHitLink link;
0371         link.setFrom(rawhit);
0372 
0373         assoc.setSimHit(npeHit);
0374         assoc.setWeight(contrib.getEnergy());
0375 
0376         link.setTo(npeHit);
0377         link.setWeight(contrib.getEnergy());
0378 
0379         rawassocs_staging[contrib.getParticle().getObjectID()] = assoc;
0380         links_staging[contrib.getParticle().getObjectID()]     = link;
0381       }
0382     }
0383 
0384     if (edep > 0.0) {
0385       for (auto& [key, link] : links_staging) {
0386         auto newLink = rawhitsLink->create();
0387         newLink.setTo(link.getTo());
0388         newLink.setFrom(link.getFrom());
0389         newLink.setWeight(link.getWeight() / edep);
0390       }
0391       for (auto& [key, assoc] : rawassocs_staging) {
0392         assoc.setWeight(assoc.getWeight() / edep);
0393         rawhitsAssoc->push_back(assoc);
0394       }
0395     }
0396 
0397     dd4hep::DetElement local;
0398     dd4hep::Position gpos;
0399     try {
0400       // global positions
0401       gpos = m_converter->position(cellID);
0402 
0403       // masked position (look for a mother volume)
0404       if (gpos_mask != 0) {
0405         auto mpos = m_converter->position(cellID & ~gpos_mask);
0406         // replace corresponding coords
0407         for (const char& c : m_cfg.maskPos) {
0408           switch (std::tolower(c)) {
0409           case 'x':
0410             gpos.SetX(mpos.X());
0411             break;
0412           case 'y':
0413             gpos.SetY(mpos.Y());
0414             break;
0415           case 'z':
0416             gpos.SetZ(mpos.Z());
0417             break;
0418           default:
0419             break;
0420           }
0421         }
0422       }
0423       gpos.SetZ(zpos * dd4hep::mm);
0424 
0425       // local positions
0426       if (m_cfg.localDetElement.empty()) {
0427         auto volman = m_detector->volumeManager();
0428         local       = volman.lookupDetElement(cellID & local_mask);
0429       } else {
0430         local = m_local;
0431       }
0432     } catch (...) {
0433       continue;
0434     }
0435 
0436     const auto lpos = local.nominal().worldToLocal(gpos);
0437     std::vector<double> cdim;
0438     // get segmentation dimensions
0439 
0440     const dd4hep::DDSegmentation::Segmentation* segmentation =
0441         m_converter->findReadout(local).segmentation()->segmentation;
0442     auto segmentation_type = segmentation->type();
0443 
0444     while (segmentation_type == "MultiSegmentation") {
0445       const auto* multi_segmentation =
0446           dynamic_cast<const dd4hep::DDSegmentation::MultiSegmentation*>(segmentation);
0447       const dd4hep::DDSegmentation::Segmentation& sub_segmentation =
0448           multi_segmentation->subsegmentation(cellID);
0449 
0450       segmentation      = &sub_segmentation;
0451       segmentation_type = segmentation->type();
0452     }
0453 
0454     if (segmentation_type == "CartesianGridXY" || segmentation_type == "HexGridXY" ||
0455         segmentation_type == "CartesianGridXYStaggered") {
0456       auto cell_dim = m_converter->cellDimensions(cellID);
0457       cdim.resize(3);
0458       cdim[0] = cell_dim[0];
0459       cdim[1] = cell_dim[1];
0460       debug("Using segmentation for cell dimensions: {}", fmt::join(cdim, ", "));
0461     } else if (segmentation_type == "CartesianStripZ") {
0462       auto cell_dim = m_converter->cellDimensions(cellID);
0463       cdim.resize(3);
0464       cdim[2] = cell_dim[0];
0465       debug("Using segmentation for cell dimensions: {}", fmt::join(cdim, ", "));
0466     } else {
0467       if ((segmentation_type != "NoSegmentation") && (!warned_unsupported_segmentation)) {
0468         warning("Unsupported segmentation type \"{}\"", segmentation_type);
0469         warned_unsupported_segmentation = true;
0470       }
0471 
0472       // Using bounding box instead of actual solid so the dimensions are always in dim_x, dim_y, dim_z
0473       cdim =
0474           m_converter->findContext(cellID)->volumePlacement().volume().boundingBox().dimensions();
0475       std::ranges::transform(cdim, cdim.begin(), [](auto&& PH1) {
0476         return std::multiplies<double>()(std::forward<decltype(PH1)>(PH1), 2);
0477       });
0478       debug("Using bounding box for cell dimensions: {}", fmt::join(cdim, ", "));
0479     }
0480 
0481     //create constant vectors for passing to hit initializer list
0482     //FIXME: needs to come from the geometry service/converter
0483     const decltype(edm4eic::CalorimeterHitData::position) position(
0484         gpos.x() / dd4hep::mm, gpos.y() / dd4hep::mm, gpos.z() / dd4hep::mm);
0485     const decltype(edm4eic::CalorimeterHitData::dimension) dimension(
0486         cdim.at(0) / dd4hep::mm, cdim.at(1) / dd4hep::mm, cdim.at(2) / dd4hep::mm);
0487     const decltype(edm4eic::CalorimeterHitData::local) local_position(
0488         lpos.x() / dd4hep::mm, lpos.y() / dd4hep::mm, lpos.z() / dd4hep::mm);
0489 
0490     auto recohit =
0491         recohits->create(cellID, corE, 0, time, 0, position, dimension, sid, lid, local_position);
0492     recohit.setRawHit(rawhit);
0493     rawhits->push_back(rawhit);
0494   }
0495 }
0496 
0497 } // namespace eicrecon