algorithms/onnx/CalorimeterParticleIDPreML.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/TensorCollection.h>
0010 #include <optional>
0011 #include <string>
0012 #include <string_view>
0013
0014 #include "algorithms/interfaces/WithPodConfig.h"
0015
0016 namespace eicrecon {
0017
0018 using CalorimeterParticleIDPreMLAlgorithm =
0019   algorithms::Algorithm<algorithms::Input<edm4eic::ClusterCollection,
0020                                                                                   std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>>,
0021                           algorithms::Output<edm4eic::TensorCollection,
0022                                                                                          std::optional<edm4eic::TensorCollection>>>;
0023
0024 class CalorimeterParticleIDPreML : public CalorimeterParticleIDPreMLAlgorithm,
0025                               public WithPodConfig<NoConfig> {
0026
0027 public:
0028   CalorimeterParticleIDPreML(std::string_view name)
0029       : CalorimeterParticleIDPreMLAlgorithm{name,
0030                                        {"inputClusters"},
0031                                        {"outputFeatureTensor", "outputTargetTensor"},
0032                                        ""} {
0033   }
0034
0035   void init() final;
0036   void process(const Input&, const Output&) const final;
0037 };
0038
0039 } // namespace eicrecon