algorithms/onnx/CalorimeterParticleIDPostML.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 Dmitry Kalinkin
0003
0004 #pragma once
0005
0006 #include <algorithms/algorithm.h>
0007 #include <edm4eic/ClusterCollection.h>
0008 #include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
0009 #include <edm4eic/MCRecoClusterParticleLinkCollection.h>
0010 #include <edm4eic/TensorCollection.h>
0011 #include <edm4hep/ParticleIDCollection.h>
0012 #include <optional>
0013 #include <string>
0014 #include <string_view>
0015
0016 #include "algorithms/interfaces/WithPodConfig.h"
0017
0018 namespace eicrecon {
0019
0020 using CalorimeterParticleIDPostMLAlgorithm = algorithms::Algorithm<
0021     algorithms::Input<edm4eic::ClusterCollection,
0022                       std::optional<edm4eic::MCRecoClusterParticleLinkCollection>,
0023                       edm4eic::TensorCollection>,
0024     algorithms::Output<edm4eic::ClusterCollection,
0025                        std::optional<edm4eic::MCRecoClusterParticleLinkCollection>,
0026                        std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>,
0027                        edm4hep::ParticleIDCollection>>;
0028
0029 class CalorimeterParticleIDPostML : public CalorimeterParticleIDPostMLAlgorithm,
0030                                     public WithPodConfig<NoConfig> {
0031
0032 public:
0033   CalorimeterParticleIDPostML(std::string_view name)
0034       : CalorimeterParticleIDPostMLAlgorithm{
0035             name,
0036             {"inputClusters", "inputClusterLinks", "inputPredictionsTensor"},
0037             {"outputClusters", "outputClusterLinks", "outputClusterAssociations",
0038              "outputParticleIDs"},
0039             ""} {}
0040
0041   void init() final;
0042   void process(const Input&, const Output&) const final;
0043 };
0044
0045 } // namespace eicrecon