From d27a0e9d76527eaefcdd13670f0fdebbbc860852 Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Tue, 8 Jul 2025 15:30:29 -0400 Subject: [PATCH 1/3] CalorimeterParticleID: use track-cluster matches to determine p in E/p --- .../onnx/CalorimeterParticleIDPostML.cc | 13 +++++++++++-- .../onnx/CalorimeterParticleIDPostML.h | 3 +++ .../onnx/CalorimeterParticleIDPreML.cc | 19 +++++++++---------- .../onnx/CalorimeterParticleIDPreML.h | 4 +++- src/detectors/EEMC/EEMC.cc | 13 ++++++++----- .../CalorimeterParticleIDPostML_factory.h | 6 ++++-- .../CalorimeterParticleIDPreML_factory.h | 3 ++- src/global/reco/reco.cc | 6 +++--- 8 files changed, 43 insertions(+), 24 deletions(-) diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.cc b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc index e9e0bd20c6..cb3569ad24 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPostML.cc +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc @@ -20,8 +20,8 @@ void CalorimeterParticleIDPostML::init() { void CalorimeterParticleIDPostML::process(const CalorimeterParticleIDPostML::Input& input, const CalorimeterParticleIDPostML::Output& output) const { - const auto [in_clusters, in_assocs, prediction_tensors] = input; - auto [out_clusters, out_assocs, out_particle_ids] = output; + const auto [in_clusters, in_track_matches, in_assocs, prediction_tensors] = input; + auto [out_clusters, out_track_matches, out_assocs, out_particle_ids] = output; if (prediction_tensors->size() != 1) { error("Expected to find a single tensor, found {}", prediction_tensors->size()); @@ -79,6 +79,15 @@ void CalorimeterParticleIDPostML::process(const CalorimeterParticleIDPostML::Inp prob_electron // float likelihood )); + // propagate track matches + for (auto in_track_match : *in_track_matches) { + if (in_track_match.getCluster() == in_cluster) { + auto out_track_match = in_track_match.clone(); + out_track_match.setCluster(out_cluster); + out_track_matches->push_back(out_track_match); + } + } + // propagate associations for (auto in_assoc : *in_assocs) { if (in_assoc.getRec() == in_cluster) { diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.h b/src/algorithms/onnx/CalorimeterParticleIDPostML.h index 3dd3c68c3c..f7059a5f97 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPostML.h +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -18,9 +19,11 @@ namespace eicrecon { using CalorimeterParticleIDPostMLAlgorithm = algorithms::Algorithm< algorithms::Input, edm4eic::TensorCollection>, algorithms::Output, edm4hep::ParticleIDCollection>>; diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc index 745ea2f085..886d2cd3e0 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc @@ -26,7 +26,7 @@ void CalorimeterParticleIDPreML::init() { void CalorimeterParticleIDPreML::process(const CalorimeterParticleIDPreML::Input& input, const CalorimeterParticleIDPreML::Output& output) const { - const auto [clusters, cluster_assocs] = input; + const auto [clusters, track_matches, cluster_assocs] = input; auto [feature_tensors, target_tensors] = output; edm4eic::MutableTensor feature_tensor = feature_tensors->create(); @@ -45,19 +45,18 @@ void CalorimeterParticleIDPreML::process(const CalorimeterParticleIDPreML::Input for (edm4eic::Cluster cluster : *clusters) { double momentum = NAN; { - // FIXME: use track momentum once matching to tracks becomes available - edm4eic::MCRecoClusterParticleAssociation best_assoc; - for (auto assoc : *cluster_assocs) { - if (assoc.getRec() == cluster) { - if ((not best_assoc.isAvailable()) || (assoc.getWeight() > best_assoc.getWeight())) { - best_assoc = assoc; + auto best_match = edm4eic::TrackClusterMatch::makeEmpty(); + for (auto match : *track_matches) { + if (match.getCluster() == cluster) { + if ((not best_match.isAvailable()) || (match.getWeight() > best_match.getWeight())) { + best_match = match; } } } - if (best_assoc.isAvailable()) { - momentum = edm4hep::utils::magnitude(best_assoc.getSim().getMomentum()); + if (best_match.isAvailable()) { + momentum = edm4hep::utils::magnitude(best_match.getTrack().getMomentum()); } else { - warning("Can't find association for cluster. Skipping..."); + trace("Can't find a match for the cluster. Skipping..."); continue; } } diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.h b/src/algorithms/onnx/CalorimeterParticleIDPreML.h index 9ae4858d66..f0964da0e2 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPreML.h +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -17,6 +18,7 @@ namespace eicrecon { using CalorimeterParticleIDPreMLAlgorithm = algorithms::Algorithm< algorithms::Input>, algorithms::Output>>; @@ -26,7 +28,7 @@ class CalorimeterParticleIDPreML : public CalorimeterParticleIDPreMLAlgorithm, public: CalorimeterParticleIDPreML(std::string_view name) : CalorimeterParticleIDPreMLAlgorithm{ - name, {"inputClusters"}, {"outputFeatureTensor", "outputTargetTensor"}, ""} {} + name, {"inputClusters", "inputTrackClusterMatches"}, {"outputFeatureTensor", "outputTargetTensor"}, ""} {} void init() final; void process(const Input&, const Output&) const final; diff --git a/src/detectors/EEMC/EEMC.cc b/src/detectors/EEMC/EEMC.cc index f03b906771..24cd49bdbb 100644 --- a/src/detectors/EEMC/EEMC.cc +++ b/src/detectors/EEMC/EEMC.cc @@ -150,7 +150,7 @@ void InitPlugin(JApplication* app) { #endif #if EDM4EIC_VERSION_MAJOR >= 8 {"EcalEndcapNClustersWithoutPIDAndShapes", // edm4eic::Cluster - "EcalEndcapNClusterAssociationsWithoutPIDAndShapes"}, // edm4eic::MCRecoClusterParticleAssociation + "EcalEndcapNClusterWithoutPIDAssociationsAndShapes"}, // edm4eic::MCRecoClusterParticleAssociation #else {"EcalEndcapNClustersWithoutShapes", // edm4eic::Cluster "EcalEndcapNClusterAssociationsWithoutShapes"}, // edm4eic::MCRecoClusterParticleAssociation @@ -168,8 +168,8 @@ void InitPlugin(JApplication* app) { #if EDM4EIC_VERSION_MAJOR >= 8 "EcalEndcapNClustersWithoutPID", {"EcalEndcapNClustersWithoutPIDAndShapes", - "EcalEndcapNClusterAssociationsWithoutPIDAndShapes"}, - {"EcalEndcapNClustersWithoutPID", "EcalEndcapNClusterAssociationsWithoutPID"}, + "EcalEndcapNClusterWithoutPIDAssociationsAndShapes"}, + {"EcalEndcapNClustersWithoutPID", "EcalEndcapNClusterWithoutPIDAssociations"}, #else "EcalEndcapNClusters", {"EcalEndcapNClustersWithoutShapes", "EcalEndcapNClusterAssociationsWithoutShapes"}, @@ -196,7 +196,8 @@ void InitPlugin(JApplication* app) { "EcalEndcapNParticleIDPreML", { "EcalEndcapNClustersWithoutPID", - "EcalEndcapNClusterAssociationsWithoutPID", + "EcalEndcapNTrackClusterWithoutPIDMatches", + "EcalEndcapNClusterWithoutPIDAssociations", }, { "EcalEndcapNParticleIDInput_features", @@ -220,11 +221,13 @@ void InitPlugin(JApplication* app) { "EcalEndcapNParticleIDPostML", { "EcalEndcapNClustersWithoutPID", - "EcalEndcapNClusterAssociationsWithoutPID", + "EcalEndcapNTrackClusterWithoutPIDMatches", + "EcalEndcapNClusterWithoutPIDAssociations", "EcalEndcapNParticleIDOutput_probability_tensor", }, { "EcalEndcapNClusters", + "EcalEndcapNTrackClusterMatches", "EcalEndcapNClusterAssociations", "EcalEndcapNClusterParticleIDs", }, diff --git a/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h index 399a3af72f..a82df7472a 100644 --- a/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h +++ b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h @@ -19,10 +19,12 @@ class CalorimeterParticleIDPostML_factory std::unique_ptr m_algo; PodioInput m_cluster_input{this}; + PodioInput m_track_cluster_matches_input{this}; PodioInput m_cluster_assoc_input{this}; PodioInput m_prediction_tensor_input{this}; PodioOutput m_cluster_output{this}; + PodioOutput m_track_cluster_matches_output{this}; PodioOutput m_cluster_assoc_output{this}; PodioOutput m_particle_id_output{this}; @@ -38,8 +40,8 @@ class CalorimeterParticleIDPostML_factory void Process(int32_t /* run_number */, uint64_t /* event_number */) { m_algo->process( - {m_cluster_input(), m_cluster_assoc_input(), m_prediction_tensor_input()}, - {m_cluster_output().get(), m_cluster_assoc_output().get(), m_particle_id_output().get()}); + {m_cluster_input(), m_track_cluster_matches_input(), m_cluster_assoc_input(), m_prediction_tensor_input()}, + {m_cluster_output().get(), m_track_cluster_matches_output().get(), m_cluster_assoc_output().get(), m_particle_id_output().get()}); } }; diff --git a/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h b/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h index fb5fe9e161..99f4da0c3d 100644 --- a/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h +++ b/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h @@ -19,6 +19,7 @@ class CalorimeterParticleIDPreML_factory std::unique_ptr m_algo; PodioInput m_cluster_input{this}; + PodioInput m_track_cluster_matches_input{this}; PodioInput m_cluster_assoc_input{this}; PodioOutput m_feature_tensor_output{this}; @@ -35,7 +36,7 @@ class CalorimeterParticleIDPreML_factory void ChangeRun(int32_t /* run_number */) {} void Process(int32_t /* run_number */, uint64_t /* event_number */) { - m_algo->process({m_cluster_input(), m_cluster_assoc_input()}, + m_algo->process({m_cluster_input(), m_track_cluster_matches_input(), m_cluster_assoc_input()}, {m_feature_tensor_output().get(), m_target_tensor_output().get()}); } }; diff --git a/src/global/reco/reco.cc b/src/global/reco/reco.cc index 4e75c56c95..7a52209b16 100644 --- a/src/global/reco/reco.cc +++ b/src/global/reco/reco.cc @@ -209,12 +209,12 @@ void InitPlugin(JApplication* app) { // Backward app->Add(new JOmniFactoryGeneratorT( - "EcalEndcapNBarrelTrackClusterMatches", - {"CalorimeterTrackProjections", "EcalEndcapNClusters"}, {"EcalEndcapNTrackClusterMatches"}, + "EcalEndcapNTrackClusterMatches", + {"CalorimeterTrackProjections", "EcalEndcapNClustersWithoutPID"}, {"EcalEndcapNTrackClusterWithoutPIDMatches"}, {.calo_id = "EcalEndcapN_ID"}, app)); app->Add(new JOmniFactoryGeneratorT( - "HcalEndcapNBarrelTrackClusterMatches", + "HcalEndcapNTrackClusterMatches", {"CalorimeterTrackProjections", "HcalEndcapNClusters"}, {"HcalEndcapNTrackClusterMatches"}, {.calo_id = "HcalEndcapN_ID"}, app)); From 19447d55183cd71c1bbc7f048375360848fdc7a3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 8 Jul 2025 19:32:55 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/algorithms/onnx/CalorimeterParticleIDPostML.cc | 2 +- src/algorithms/onnx/CalorimeterParticleIDPostML.h | 6 ++---- src/algorithms/onnx/CalorimeterParticleIDPreML.cc | 4 ++-- src/algorithms/onnx/CalorimeterParticleIDPreML.h | 9 +++++---- .../calorimetry/CalorimeterParticleIDPostML_factory.h | 7 ++++--- src/global/reco/reco.cc | 9 ++++----- 6 files changed, 18 insertions(+), 19 deletions(-) diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.cc b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc index cb3569ad24..6bb0ab4599 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPostML.cc +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc @@ -21,7 +21,7 @@ void CalorimeterParticleIDPostML::process(const CalorimeterParticleIDPostML::Inp const CalorimeterParticleIDPostML::Output& output) const { const auto [in_clusters, in_track_matches, in_assocs, prediction_tensors] = input; - auto [out_clusters, out_track_matches, out_assocs, out_particle_ids] = output; + auto [out_clusters, out_track_matches, out_assocs, out_particle_ids] = output; if (prediction_tensors->size() != 1) { error("Expected to find a single tensor, found {}", prediction_tensors->size()); diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.h b/src/algorithms/onnx/CalorimeterParticleIDPostML.h index f7059a5f97..2ddf73c111 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPostML.h +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.h @@ -18,12 +18,10 @@ namespace eicrecon { using CalorimeterParticleIDPostMLAlgorithm = algorithms::Algorithm< - algorithms::Input, edm4eic::TensorCollection>, - algorithms::Output, edm4hep::ParticleIDCollection>>; diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc index 886d2cd3e0..33d759818e 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc @@ -26,8 +26,8 @@ void CalorimeterParticleIDPreML::init() { void CalorimeterParticleIDPreML::process(const CalorimeterParticleIDPreML::Input& input, const CalorimeterParticleIDPreML::Output& output) const { - const auto [clusters, track_matches, cluster_assocs] = input; - auto [feature_tensors, target_tensors] = output; + const auto [clusters, track_matches, cluster_assocs] = input; + auto [feature_tensors, target_tensors] = output; edm4eic::MutableTensor feature_tensor = feature_tensors->create(); feature_tensor.addToShape(clusters->size()); diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.h b/src/algorithms/onnx/CalorimeterParticleIDPreML.h index f0964da0e2..bbdb340e10 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPreML.h +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.h @@ -17,8 +17,7 @@ namespace eicrecon { using CalorimeterParticleIDPreMLAlgorithm = algorithms::Algorithm< - algorithms::Input>, algorithms::Output>>; @@ -27,8 +26,10 @@ class CalorimeterParticleIDPreML : public CalorimeterParticleIDPreMLAlgorithm, public: CalorimeterParticleIDPreML(std::string_view name) - : CalorimeterParticleIDPreMLAlgorithm{ - name, {"inputClusters", "inputTrackClusterMatches"}, {"outputFeatureTensor", "outputTargetTensor"}, ""} {} + : CalorimeterParticleIDPreMLAlgorithm{name, + {"inputClusters", "inputTrackClusterMatches"}, + {"outputFeatureTensor", "outputTargetTensor"}, + ""} {} void init() final; void process(const Input&, const Output&) const final; diff --git a/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h index a82df7472a..90926a6047 100644 --- a/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h +++ b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h @@ -39,9 +39,10 @@ class CalorimeterParticleIDPostML_factory void ChangeRun(int32_t /* run_number */) {} void Process(int32_t /* run_number */, uint64_t /* event_number */) { - m_algo->process( - {m_cluster_input(), m_track_cluster_matches_input(), m_cluster_assoc_input(), m_prediction_tensor_input()}, - {m_cluster_output().get(), m_track_cluster_matches_output().get(), m_cluster_assoc_output().get(), m_particle_id_output().get()}); + m_algo->process({m_cluster_input(), m_track_cluster_matches_input(), m_cluster_assoc_input(), + m_prediction_tensor_input()}, + {m_cluster_output().get(), m_track_cluster_matches_output().get(), + m_cluster_assoc_output().get(), m_particle_id_output().get()}); } }; diff --git a/src/global/reco/reco.cc b/src/global/reco/reco.cc index 7a52209b16..3b6c28d956 100644 --- a/src/global/reco/reco.cc +++ b/src/global/reco/reco.cc @@ -210,13 +210,12 @@ void InitPlugin(JApplication* app) { // Backward app->Add(new JOmniFactoryGeneratorT( "EcalEndcapNTrackClusterMatches", - {"CalorimeterTrackProjections", "EcalEndcapNClustersWithoutPID"}, {"EcalEndcapNTrackClusterWithoutPIDMatches"}, - {.calo_id = "EcalEndcapN_ID"}, app)); + {"CalorimeterTrackProjections", "EcalEndcapNClustersWithoutPID"}, + {"EcalEndcapNTrackClusterWithoutPIDMatches"}, {.calo_id = "EcalEndcapN_ID"}, app)); app->Add(new JOmniFactoryGeneratorT( - "HcalEndcapNTrackClusterMatches", - {"CalorimeterTrackProjections", "HcalEndcapNClusters"}, {"HcalEndcapNTrackClusterMatches"}, - {.calo_id = "HcalEndcapN_ID"}, app)); + "HcalEndcapNTrackClusterMatches", {"CalorimeterTrackProjections", "HcalEndcapNClusters"}, + {"HcalEndcapNTrackClusterMatches"}, {.calo_id = "HcalEndcapN_ID"}, app)); #endif // EDM4EIC_VERSION_MAJOR >= 8 From 678b8634faabf740c38f50c2eb6b48c1e20139e6 Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Mon, 14 Jul 2025 17:21:46 -0400 Subject: [PATCH 3/3] IWYU --- src/algorithms/onnx/CalorimeterParticleIDPreML.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc index 33d759818e..bdf239a62b 100644 --- a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc @@ -1,19 +1,19 @@ // SPDX-License-Identifier: LGPL-3.0-or-later -// Copyright (C) 2024 Dmitry Kalinkin +// Copyright (C) 2024 - 2025 Dmitry Kalinkin #include #if EDM4EIC_VERSION_MAJOR >= 8 -#include -#include #include +#include #include #include #include #include -#include - +#include +#include #include +#include #include "CalorimeterParticleIDPreML.h"