CATFinderModule.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
#include <framework/core/Module.h>
 
#include <framework/datastore/StoreArray.h>
#include <mva/methods/ONNX.h>
#include <tracking/trackingUtilities/eventdata/hits/CDCWireHit.h>
#include <tracking/trackingUtilities/rootification/StoreWrappedObjPtr.h>
 
namespace Belle2 {
 
  class CDCHit;
  class RecoHitInformation;
  class RecoTrack;

   * @class CATFinderModule
   * @brief Belle II module for the CDC AI Track Finder (CATFinder) using a Graph Neural Network (GNN) and condensation clustering.
   *
   * The CATFinderModule processes Central Drift Chamber (CDC) wire hits and reconstructs charged particle tracks
   * by applying a Graph Neural Network (GNN) followed by condensation-based clustering and spatial hit ordering.
   *
   * The module workflow consists of:
   * - **Preprocessing** ``CDCWireHit`` objects to construct the input feature tensor.
   * - **GNN inference** using a trained model loaded from a weight file.
   * - **Postprocessing** GNN outputs to extract condensation points, assign hits, and build `RecoTrack` objects.
   *
   * The GNN is run via the ONNX interface.
   */
  class CATFinderModule : public Module {
 
  public:
    CATFinderModule();

    void initialize() override;

    void beginRun() override;

    void event() override;
 
  private:

    std::string m_CDCRecoTracksName;

    std::string m_catFinderWeightfileName = "CATFinderWeightFile";

    StoreArray<CDCHit> m_CDCHits;

    TrackingUtilities::StoreWrappedObjPtr<std::vector<TrackingUtilities::CDCWireHit>> m_wireHitVector{"CDCWireHitVector"};

    StoreArray<RecoTrack> m_CDCRecoTracks;

    StoreArray<RecoHitInformation> m_recoHitInformations;

    float m_tdcOffset = 0.;

    float m_tdcScale = 0.;

    float m_adcClip = 0.;

    float m_slayerScale = 0.;

    float m_clayerScale = 0.;

    float m_layerScale = 0.;

    float m_spatialCoordinatesScale = 0.;

    unsigned int m_nInputFeatures = 0;

    unsigned int m_latentSpaceNDim = 0;

    float m_tBeta = 0.;

    float m_tDistance = 0.;

    float m_maxRadius = 0.;

    unsigned int m_minNumberHits = 0.;

    std::string m_inputTFeaturesName;

    std::string m_outputTBetaName;

    std::string m_outputTCoordinatesName;

    std::string m_outputTMomentumName;

    std::string m_outputTVertexName;

    std::string m_outputTChargeName;

    std::unique_ptr<MVA::ONNX::Session> m_session;
  };

}