File indexing completed on 2024-09-27 07:02:57
0001
0002
0003
0004
0005
0006
0007
0008
0009
0010
0011 #pragma once
0012
0013 #include <Eigen/Dense>
0014 #include <algorithm>
0015
0016 #include <algorithms/algorithm.h>
0017 #include <DDRec/CellIDPositionConverter.h>
0018 #include <DDRec/Surface.h>
0019 #include <DDRec/SurfaceManager.h>
0020
0021 #include "algorithms/calorimetry/ClusterTypes.h"
0022
0023
0024 #include <edm4hep/MCParticleCollection.h>
0025 #include <edm4hep/SimCalorimeterHitCollection.h>
0026 #include <edm4hep/utils/vector_utils.h>
0027 #include <edm4eic/CalorimeterHitCollection.h>
0028 #include <edm4eic/ClusterCollection.h>
0029 #include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
0030 #include <edm4eic/ProtoClusterCollection.h>
0031
0032 #include "algorithms/interfaces/WithPodConfig.h"
0033 #include "ImagingClusterRecoConfig.h"
0034
0035 namespace eicrecon {
0036
0037 using ImagingClusterRecoAlgorithm = algorithms::Algorithm<
0038 algorithms::Input<
0039 edm4eic::ProtoClusterCollection,
0040 edm4hep::SimCalorimeterHitCollection
0041 >,
0042 algorithms::Output<
0043 edm4eic::ClusterCollection,
0044 edm4eic::MCRecoClusterParticleAssociationCollection,
0045 edm4eic::ClusterCollection
0046 >
0047 >;
0048
0049
0050
0051
0052
0053
0054
0055
0056 class ImagingClusterReco
0057 : public ImagingClusterRecoAlgorithm,
0058 public WithPodConfig<ImagingClusterRecoConfig> {
0059
0060 public:
0061 ImagingClusterReco(std::string_view name)
0062 : ImagingClusterRecoAlgorithm{name,
0063 {"inputProtoClusterCollection", "mcHits"},
0064 {"outputClusterCollection", "outputClusterAssociations", "outputLayerCollection"},
0065 "Reconstruct the cluster/layer info for imaging calorimeter."} {}
0066
0067 public:
0068
0069 void init() { }
0070
0071 void process(const Input& input, const Output& output) const final {
0072
0073 const auto [proto, mchits] = input;
0074 auto [clusters, associations, layers] = output;
0075
0076 for (const auto& pcl: *proto) {
0077 if (!pcl.getHits().empty() && !pcl.getHits(0).isAvailable()) {
0078 warning("Protocluster hit relation is invalid, skipping protocluster");
0079 continue;
0080 }
0081
0082 auto cl = reconstruct_cluster(pcl);
0083 auto cl_layers = reconstruct_cluster_layers(pcl);
0084
0085
0086 auto [theta, phi] = fit_track(cl_layers);
0087 cl.setIntrinsicTheta(theta);
0088 cl.setIntrinsicPhi(phi);
0089
0090
0091
0092 for (const auto& layer: cl_layers) {
0093 layers->push_back(layer);
0094 cl.addToClusters(layer);
0095 }
0096 clusters->push_back(cl);
0097
0098
0099 if (mchits->size() > 0) {
0100
0101
0102 auto pclhits = pcl.getHits();
0103 auto pclhit = std::max_element(
0104 pclhits.begin(),
0105 pclhits.end(),
0106 [](const auto &pclhit1, const auto &pclhit2) {
0107 return pclhit1.getEnergy() < pclhit2.getEnergy();
0108 }
0109 );
0110
0111
0112 const edm4hep::SimCalorimeterHit* mchit = nullptr;
0113 for (auto h : *mchits) {
0114 if (h.getCellID() == pclhit->getCellID()) {
0115 mchit = &h;
0116 break;
0117 }
0118 }
0119 if( !mchit ){
0120
0121 warning("Proto-cluster has highest energy in CellID {}, but no mc hit with that CellID was found.", pclhit->getCellID());
0122 break;
0123 }
0124
0125
0126 const auto &mcp = mchit->getContributions(0).getParticle();
0127
0128
0129 auto clusterassoc = associations->create();
0130 clusterassoc.setRecID(cl.getObjectID().index);
0131 clusterassoc.setSimID(mcp.getObjectID().index);
0132 clusterassoc.setWeight(1.0);
0133 clusterassoc.setRec(cl);
0134 clusterassoc.setSim(mcp);
0135 }
0136
0137 }
0138
0139
0140 for (const auto& cl: *clusters) {
0141 debug("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.getObjectID().index,
0142 cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0143 cl.getIntrinsicPhi() / M_PI * 180.
0144 );
0145 }
0146 }
0147
0148 private:
0149
0150 static std::vector<edm4eic::MutableCluster> reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) {
0151 const auto& hits = pcl.getHits();
0152 const auto& weights = pcl.getWeights();
0153
0154 std::map<int, std::vector<std::pair<const edm4eic::CalorimeterHit, float>>> layer_map;
0155 for (unsigned i = 0; i < hits.size(); ++i) {
0156 const auto hit = hits[i];
0157 auto lid = hit.getLayer();
0158
0159
0160
0161
0162 layer_map[lid].push_back({hit, weights[i]});
0163 }
0164
0165
0166 std::vector<edm4eic::MutableCluster> cl_layers;
0167 for (const auto &[lid, layer_hits]: layer_map) {
0168 auto layer = reconstruct_layer(layer_hits);
0169 cl_layers.push_back(layer);
0170 }
0171 return cl_layers;
0172 }
0173
0174 static edm4eic::MutableCluster reconstruct_layer(const std::vector<std::pair<const edm4eic::CalorimeterHit, float>>& hits) {
0175 edm4eic::MutableCluster layer;
0176 layer.setType(Jug::Reco::ClusterType::kClusterSlice);
0177
0178 double energy{0};
0179 double energyError{0};
0180 double time{0};
0181 double timeError{0};
0182 double sumOfWeights{0};
0183 auto pos = layer.getPosition();
0184 for (const auto &[hit, weight]: hits) {
0185 energy += hit.getEnergy() * weight;
0186 energyError += std::pow(hit.getEnergyError() * weight, 2);
0187 time += hit.getTime() * weight;
0188 timeError += std::pow(hit.getTimeError() * weight, 2);
0189 pos = pos + hit.getPosition() * weight;
0190 sumOfWeights += weight;
0191 layer.addToHits(hit);
0192 }
0193 layer.setEnergy(energy);
0194 layer.setEnergyError(std::sqrt(energyError));
0195 layer.setTime(time / sumOfWeights);
0196 layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0197 layer.setNhits(hits.size());
0198 layer.setPosition(pos / sumOfWeights);
0199
0200
0201
0202
0203 double radius = 0.;
0204 for (const auto &[hit, weight]: hits) {
0205 radius += std::pow(edm4hep::utils::magnitude(hit.getPosition() - layer.getPosition()), 2);
0206 }
0207 layer.addToShapeParameters(std::sqrt(radius / layer.getNhits()));
0208
0209
0210 return layer;
0211 }
0212
0213 static edm4eic::MutableCluster reconstruct_cluster(const edm4eic::ProtoCluster& pcl) {
0214 edm4eic::MutableCluster cluster;
0215
0216 const auto& hits = pcl.getHits();
0217 const auto& weights = pcl.getWeights();
0218
0219 cluster.setType(Jug::Reco::ClusterType::kCluster3D);
0220 double energy = 0.;
0221 double energyError = 0.;
0222 double time = 0.;
0223 double timeError = 0.;
0224 double meta = 0.;
0225 double mphi = 0.;
0226 double r = 9999 * dd4hep::cm;
0227 for (unsigned i = 0; i < hits.size(); ++i) {
0228 const auto &hit = hits[i];
0229 const auto &weight = weights[i];
0230 energy += hit.getEnergy() * weight;
0231 energyError += std::pow(hit.getEnergyError() * weight, 2);
0232
0233 const double energyWeight = hit.getEnergy() * weight;
0234 time += hit.getTime() * energyWeight;
0235 timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0236 meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0237 mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0238 r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0239 cluster.addToHits(hit);
0240 }
0241 cluster.setEnergy(energy);
0242 cluster.setEnergyError(std::sqrt(energyError));
0243 cluster.setTime(time / energy);
0244 cluster.setTimeError(std::sqrt(timeError) / energy);
0245 cluster.setNhits(hits.size());
0246 cluster.setPosition(edm4hep::utils::sphericalToVector(r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0247
0248
0249 double radius = 0.;
0250 for (const auto &hit: hits) {
0251 radius += std::pow(
0252 std::hypot(
0253 (edm4hep::utils::eta(hit.getPosition()) - edm4hep::utils::eta(cluster.getPosition())),
0254 (edm4hep::utils::angleAzimuthal(hit.getPosition()) - edm4hep::utils::angleAzimuthal(cluster.getPosition()))
0255 ),
0256 2.0
0257 );
0258 }
0259 cluster.addToShapeParameters(std::sqrt(radius / cluster.getNhits()));
0260
0261
0262
0263
0264
0265
0266
0267
0268
0269 return cluster;
0270 }
0271
0272 std::pair<double , double > fit_track(const std::vector<edm4eic::MutableCluster> &layers) const {
0273 int nrows = 0;
0274 decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0275 for (const auto &layer: layers) {
0276 if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0277 mean_pos = mean_pos + layer.getPosition();
0278 nrows += 1;
0279 }
0280 }
0281
0282
0283 if (nrows < 2) {
0284 return {};
0285 }
0286
0287 mean_pos = mean_pos / nrows;
0288
0289 Eigen::MatrixXd pos(nrows, 3);
0290 int ir = 0;
0291 for (const auto &layer: layers) {
0292 if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0293 auto delta = layer.getPosition() - mean_pos;
0294 pos(ir, 0) = delta.x;
0295 pos(ir, 1) = delta.y;
0296 pos(ir, 2) = delta.z;
0297 ir += 1;
0298 }
0299 }
0300
0301 Eigen::JacobiSVD <Eigen::MatrixXd> svd(pos, Eigen::ComputeThinU | Eigen::ComputeThinV);
0302 const auto dir = svd.matrixV().col(0);
0303
0304 return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0305 }
0306 };
0307
0308 }