algorithms/calorimetry/ImagingClusterReco.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Chao Peng, Wouter Deconinck, David Lawrence
0003
0004 /*
0005  *  Reconstruct the cluster/layer info for imaging calorimeter
0006  *  Logarithmic weighting is used to describe energy deposit in transverse direction
0007  *
0008  *  Author: Chao Peng (ANL), 06/02/2021
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 // Event Model related classes
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   /** Imaging cluster reconstruction.
0050    *
0051    *  Reconstruct the cluster/layer info for imaging calorimeter
0052    *  Logarithmic weighting is used to describe energy deposit in transverse direction
0053    *
0054    *  \ingroup reco
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             // get cluster and associated layers
0082             auto cl = reconstruct_cluster(pcl);
0083             auto cl_layers = reconstruct_cluster_layers(pcl);
0084
0085             // Get cluster direction from the layer profile
0086             auto [theta, phi] = fit_track(cl_layers);
0087             cl.setIntrinsicTheta(theta);
0088             cl.setIntrinsicPhi(phi);
0089             // no error on the intrinsic direction TODO
0090
0091             // store layer and clusters on the datastore
0092             for (const auto& layer: cl_layers) {
0093                 layers->push_back(layer);
0094                 cl.addToClusters(layer);
0095             }
0096             clusters->push_back(cl);
0097
0098             // If mcHits are available, associate cluster with MCParticle
0099             if (mchits->size() > 0) {
0100
0101                 // 1. find pclhit with the largest energy deposition
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                 // 2. find mchit with same CellID
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                     // break if no matching hit found for this CellID
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                 // 3. find mchit's MCParticle
0126                 const auto &mcp = mchit->getContributions(0).getParticle();
0127
0128                 // set association
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         // debug output
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         // using map to have hits sorted by layer
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 //            if (layer_map.count(lid) == 0) {
0159 //                std::vector<std::pair<const edm4eic::CalorimeterHit, float>> v;
0160 //                layer_map[lid] = {};
0161 //            }
0162             layer_map[lid].push_back({hit, weights[i]});
0163         }
0164
0165         // create layers
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         // Calculate averages
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         // positionError not set
0200         // Intrinsic direction meaningless in a cluster layer --> not set
0201
0202         // Calculate radius as the standard deviation of the hits versus the cluster center
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         // TODO Skewedness
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             // energy weighting for the other variables
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         // shower radius estimate (eta-phi plane)
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         // Skewedness not calculated TODO
0261
0262         // Optionally store the MC truth associated with the first hit in this cluster
0263         // FIXME no connection between cluster and truth in edm4hep
0264         // if (mcHits) {
0265         //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0266         //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0267         //}
0268
0269         return cluster;
0270     }
0271
0272     std::pair<double /* polar */, double /* azimuthal */> 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         // cannot fit
0283         if (nrows < 2) {
0284             return {};
0285         }
0286
0287         mean_pos = mean_pos / nrows;
0288         // fill position data
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         // theta and phi
0304         return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0305     }
0306 };
0307
0308 } // namespace eicrecon