CbmRichMCbmDenoiseCnn.cxx

Go to the documentation of this file.

/* Copyright (C) 2024-2025 UGiessen, Giessen
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Martin Beyer [committer] */
 
#if HAVE_ONNXRUNTIME
 
#include "CbmRichMCbmDenoiseCnn.h"
 
#include "CbmDigiManager.h"
#include "CbmEvent.h"
#include "CbmRichDetectorData.h"
#include "CbmRichGeoHandler.h"
#include "CbmRichHit.h"
 
#include <Logger.h>
 
#include <TClonesArray.h>
#include <TStopwatch.h>
 
#include <iostream>
#include <sstream>
#include <vector>
 
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
 
void CbmRichMCbmDenoiseCnn::Init()
{
  fOrtEnv = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, GetName());
 
  fOrtSessionOptions = std::make_unique<Ort::SessionOptions>();
  // Thread numbers need to be explicitly set. With the current Ort version it can't detect
  // the numbers on the cluster, should be fixed in a later version.
  fOrtSessionOptions->SetIntraOpNumThreads(1);
  fOrtSessionOptions->SetInterOpNumThreads(1);
  fOrtSessionOptions->SetExecutionMode(ORT_SEQUENTIAL);
  fOrtSessionOptions->SetGraphOptimizationLevel(ORT_ENABLE_ALL);
 
  fOrtSession = std::make_unique<Ort::Session>(*fOrtEnv, fOnnxFilePath.c_str(), *fOrtSessionOptions);
 
  fOrtRunOptions    = std::make_unique<Ort::RunOptions>(nullptr);
  fOrtAllocatorInfo = std::make_unique<Ort::MemoryInfo>(Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU));
  fOrtInput         = std::make_unique<Ort::Value>(Ort::Value::CreateTensor<float>(
    *fOrtAllocatorInfo, fInput.data(), fInput.size(), fInputShape.data(), fInputShape.size()));
 
  fCbmRichGeoHandler = &CbmRichGeoHandler::GetInstance();
}
 
void CbmRichMCbmDenoiseCnn::Process(CbmEvent* event, const TClonesArray* richHits)
{
  int nHits = event ? event->GetNofData(ECbmDataType::kRichHit) : richHits->GetEntriesFast();
  std::vector<int> timeWindowHitIndices;
  for (int iHit = 0; iHit < nHits; iHit++) {  // Sliding time window loop
    timeWindowHitIndices.clear();
    int seedIndex       = event ? event->GetIndex(ECbmDataType::kRichHit, iHit) : iHit;
    CbmRichHit* seedHit = static_cast<CbmRichHit*>(richHits->At(seedIndex));
    if (!seedHit) continue;
    const auto gridIdxSeed = AddressToGridIndex(seedHit->GetAddress());
    if (gridIdxSeed < 0 || gridIdxSeed >= static_cast<int>(fInput.size())) {
      LOG(error) << GetName() << "::Process: Invalid grid index for seed hit: " << gridIdxSeed
                 << ". Skipping seed hit.";
      continue;
    }
    timeWindowHitIndices.push_back(seedIndex);
    int hitsInTimeWindow = 1;
    for (int jHit = iHit + 1; jHit < nHits; jHit++) {  // Search for hits in time window
      int jHitIndex   = event ? event->GetIndex(ECbmDataType::kRichHit, jHit) : jHit;
      CbmRichHit* hit = static_cast<CbmRichHit*>(richHits->At(jHitIndex));
      if (!hit) continue;
      const auto gridIdxHit = AddressToGridIndex(hit->GetAddress());
      if (gridIdxHit < 0 || gridIdxHit >= static_cast<int>(fInput.size())) {
        LOG(error) << GetName() << "::Process: Invalid grid index for hit: " << gridIdxHit << ". Skipping hit.";
        continue;
      }
      double dt = hit->GetTime() - seedHit->GetTime();
      if (dt < fTimeWindowLength) {
        hitsInTimeWindow++;
        timeWindowHitIndices.push_back(jHitIndex);
      }
      else {
        break;
      }
    }
 
    if (hitsInTimeWindow >= fMinHitsInTimeWindow) {
      ProcessTimeWindow(timeWindowHitIndices, seedHit->GetTime(), richHits);
      iHit += hitsInTimeWindow - 1;  // Move to last hit inside time window
    }
    else {
      seedHit->SetIsNoiseNN(true);
    }
  }
}
 
void CbmRichMCbmDenoiseCnn::ProcessTimeWindow(const std::vector<int>& timeWindowHitIndices, const double& seedHitTime,
                                              const TClonesArray* richHits)
{
  fInput = {};  // Reset all input values to 0.0
  std::vector<int> gridIndices;
  gridIndices.reserve(50);
  for (const auto& hitIndex : timeWindowHitIndices) {
    CbmRichHit* hit    = static_cast<CbmRichHit*>(richHits->At(hitIndex));
    const auto gridIdx = AddressToGridIndex(hit->GetAddress());
    if (gridIdx < 0 || gridIdx >= static_cast<int>(fInput.size())) {
      // This should never happen since those hits are already filtered.
      // Keep it included for safety, throwing fatal.
      LOG(fatal) << GetName() << "::ProcessTimeWindow: Invalid grid index: " << gridIdx;
      continue;
    }
    gridIndices.push_back(gridIdx);
    // Shift time by 1ns to distinguish empty pixels from seed hit time
    fInput[gridIdx] = static_cast<float>((hit->GetTime() - seedHitTime + 1.0) / (fTimeWindowLength + 1.0));
  }
 
  const auto output = Inference(gridIndices);
 
  for (std::size_t i = 0; i < timeWindowHitIndices.size(); i++) {
    CbmRichHit* hit = static_cast<CbmRichHit*>(richHits->At(timeWindowHitIndices[i]));
    bool isNoise    = output[i] < fClassificationThreshold;
    hit->SetIsNoiseNN(isNoise);
  }
}
 
std::vector<float> CbmRichMCbmDenoiseCnn::Inference(const std::vector<int>& gridIndices)
{
  auto output = fOrtSession->Run(*fOrtRunOptions.get(), fInputNames.data(), fOrtInput.get(), 1, fOutputNames.data(), 1);
  float* intarr = output.front().GetTensorMutableData<float>();
 
  std::vector<float> out(gridIndices.size());
  std::transform(gridIndices.begin(), gridIndices.end(), out.begin(),
                 [intarr](int gridIdx) { return intarr[gridIdx]; });
  return out;
}
 
int CbmRichMCbmDenoiseCnn::AddressToGridIndex(int address)
{
  CbmRichPixelData* pixel_data = fCbmRichGeoHandler->GetPixelDataByAddress(address);
 
  // Calculate local X [0,7],Y [0,7] indices of one MAPMT
  int pmtUID  = pixel_data->fPmtId;
  int pmtIndX = (pmtUID >> 4) & 0xF;  // index ascending from -x to x
  int pmtIndY = pmtUID & 0xF;         // index ascending from y to -y
 
  // Calculate global X [0,31],Y [0,72] indices
  int globalIndX = 8 * pmtIndX + pixel_data->fPixelId % 8;
  int globalIndY = 8 * pmtIndY + pixel_data->fPixelId / 8;
 
  return 32 * globalIndY + globalIndX;
}
 
#endif  // HAVE_ONNXRUNTIME