ECLCRFinderModule.cc

/**************************************************************************
 * 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.                  *
 **************************************************************************/
 
/* Own header. */
#include <ecl/modules/eclCRFinder/ECLCRFinderModule.h>
 
/* ECL headeers. */
#include <ecl/dataobjects/ECLCalDigit.h>
#include <ecl/dataobjects/ECLConnectedRegion.h>
#include <ecl/geometry/ECLNeighbours.h>
 
/* Basf2 headers. */
#include <framework/gearbox/Unit.h>
#include <framework/logging/Logger.h>
#include <mdst/dataobjects/EventLevelClusteringInfo.h>
 
/* C++ headers. */
#include <algorithm>
#include <iostream>
#include <array>
 
// NAMESPACE(S)
using namespace Belle2;
 
//-----------------------------------------------------------------
//                 Register the Module
//-----------------------------------------------------------------
REG_MODULE(ECLCRFinder);
REG_MODULE(ECLCRFinderPureCsI);
 
//-----------------------------------------------------------------
//                 Implementation
//-----------------------------------------------------------------
ECLCRFinderModule::ECLCRFinderModule() : Module(), m_eclCalDigits(eclCalDigitArrayName()),
  m_eclConnectedRegions(eclConnectedRegionArrayName()), m_eventLevelClusteringInfo(eventLevelClusteringInfoName())
{
  // Set description
  setDescription("ECLCRFinderModule");
 
  // Parallel processing certification.
  setPropertyFlags(c_ParallelProcessingCertified);
 
  // Add module parameters.
  addParam("energyCut0", m_energyCut[0], "Seed energy cut.", 20.0 * Belle2::Unit::MeV);
  addParam("energyCut1", m_energyCut[1], "Growth energy cut.", 20.0 * Belle2::Unit::MeV);
  addParam("energyCut2", m_energyCut[2], "Digit energy cut.", 0.5 * Belle2::Unit::MeV);
  addParam("timeCut0", m_timeCut[0], "Seed time cut (negative values for residual cut).", 400.);
  addParam("timeCut1", m_timeCut[1], "Growth time cut (negative values for residual cut).", 400.);
  addParam("timeCut2", m_timeCut[2], "Digit time cut (negative values for residual cut).", 99999.);
  addParam("timeCut0maxEnergy", m_timeCut_maxEnergy[0], "Time cut is only applied below this energy for seed crystals.",
           99999.0 * Belle2::Unit::MeV);
  addParam("timeCut1maxEnergy", m_timeCut_maxEnergy[1], "Time cut is only applied below this energy for growth crystals.",
           99999.0 * Belle2::Unit::MeV);
  addParam("timeCut2maxEnergy", m_timeCut_maxEnergy[2], "Time cut is only applied below this energy for digits.",
           0.0 * Belle2::Unit::MeV);
  addParam("mapType0", m_mapType[0], "Map type for seed crystals.", std::string("N"));
  addParam("mapType1", m_mapType[1], "Map type for growth crystals.",  std::string("N"));
  addParam("mapPar0", m_mapPar[0],
           "Map parameter for seed crystals (radius (type=R), integer (for type=N) or fraction (for type=MC)).", 1.0);
  addParam("mapPar1", m_mapPar[1],
           "Map parameter for growth crystals (radius (type=R), integer (for type=N) or fraction (for type=MC)).", 1.0);
  addParam("skipFailedTimeFitDigits", m_skipFailedTimeFitDigits, "Digits with failed fits are skipped when checking timing cuts.", 0);
  addParam("useParametersFromDatabase", m_useParametersFromDatabase, "get energy and time cuts from payload", true);
 
}
 
ECLCRFinderModule::~ECLCRFinderModule()
{
  ;
}
 
void ECLCRFinderModule::initialize()
{
  B2DEBUG(200, "ECLCRFinderModule::initialize()");
 
  // Register dataobjects.
  m_eclCalDigits.isRequired(eclCalDigitArrayName());
  m_eclConnectedRegions.registerInDataStore(eclConnectedRegionArrayName());
  m_eventLevelClusteringInfo.isRequired(eventLevelClusteringInfoName());
 
  // Register relations.
  m_eclConnectedRegions.registerRelationTo(m_eclCalDigits);
 
  // Initialize neighbour maps.
  m_neighbourMaps.resize(2);
  m_neighbourMaps[0] = new ECL::ECLNeighbours(m_mapType[0], m_mapPar[0]);
  m_neighbourMaps[1] = new ECL::ECLNeighbours(m_mapType[1], m_mapPar[1]);
 
  // Resize the vectors
  m_cellIdToCheckVec.resize(8737); 
  m_cellIdToSeedVec.resize(8737); 
  m_cellIdToGrowthVec.resize(8737); 
  m_cellIdToDigitVec.resize(8737); 
  m_cellIdToTempCRIdVec.resize(8737); 
  m_calDigitStoreArrPosition.resize(8737);
 
  m_adj.resize(8737);
  m_visited.resize(8737);
}
 
//-----------------------------------------------------------------
//..By default, get the energy and time thresholds from the
//  eclClusteringParameters dbOject
void ECLCRFinderModule::beginRun()
{
  if (m_useParametersFromDatabase and m_eclClusteringParameters.hasChanged()) {
    std::array<double, 3> CRF_energyCut = m_eclClusteringParameters->getCRFEnergyCut();
    m_energyCut[0] = CRF_energyCut[0];
    m_energyCut[1] = CRF_energyCut[1];
    m_energyCut[2] = CRF_energyCut[2];
 
    std::array<double, 3> CRF_timeCut = m_eclClusteringParameters->getCRFTimeCut();
    m_timeCut[0] = CRF_timeCut[0];
    m_timeCut[1] = CRF_timeCut[1];
    m_timeCut[2] = CRF_timeCut[2];
 
    std::array<double, 3> CRF_timeCutMaxEnergy = m_eclClusteringParameters->getCRFTimeCutMaxEnergy();
    m_timeCut_maxEnergy[0] = CRF_timeCutMaxEnergy[0];
    m_timeCut_maxEnergy[1] = CRF_timeCutMaxEnergy[1];
    m_timeCut_maxEnergy[2] = CRF_timeCutMaxEnergy[2];
  }
 
  // Check user inputs: [2]: digit, [1]: growth, [0]: seed
  // overall energy thresholds
  if (std::isless(m_energyCut[0], m_energyCut[1])) B2FATAL("ECLCRFinderModule::beginRun(): m_energyCut[0]=" << m_energyCut[0] <<
                                                             " must be larger or equal than m_energyCut[1]=" << m_energyCut[1]);
  if (std::isless(m_energyCut[1], m_energyCut[2])) B2FATAL("ECLCRFinderModule::beginRun(): m_energyCut[1]=" << m_energyCut[1] <<
                                                             " must be larger or equal than m_energyCut[2]=" << m_energyCut[2]);
  // timing threshold (can depend on energy, but we make the check here even stronger by checking that the timing is looser without checking the timing energy range)
  if (std::isgreater(m_timeCut[0], m_timeCut[1]))
    B2FATAL("ECLCRFinderModule::beginRun(): m_timeCut[0] must be less or equal than m_timeCut[1].");
  if (std::isgreater(m_timeCut[1], m_timeCut[2]))
    B2FATAL("ECLCRFinderModule::beginRun(): m_timeCut[1] must be less or equal than m_timeCut[2].");
}
 
 
//-----------------------------------------------------------------
 
void ECLCRFinderModule::event()
{
  B2DEBUG(200, "ECLCRFinderModule::event()");
 
  // Reset the vector(s).
  std::fill(m_cellIdToCheckVec.begin(), m_cellIdToCheckVec.end(), 0);
  std::fill(m_cellIdToSeedVec.begin(), m_cellIdToSeedVec.end(), 0);
  std::fill(m_cellIdToGrowthVec.begin(), m_cellIdToGrowthVec.end(), 0);
  std::fill(m_cellIdToDigitVec.begin(), m_cellIdToDigitVec.end(), 0);
  std::fill(m_cellIdToTempCRIdVec.begin(), m_cellIdToTempCRIdVec.end(), 0);
 
  // Fill a vector that can be used to map cellid -> store array position
  std::fill(m_calDigitStoreArrPosition.begin(), m_calDigitStoreArrPosition.end(), -1);
  for (int i = 0; i < m_eclCalDigits.getEntries(); i++) {
    m_calDigitStoreArrPosition[m_eclCalDigits[i]->getCellId()] = i;
  }
 
  // Clear the map(s).
  m_cellIdToTempCRIdMap.clear();
 
  //-------------------------------------------------------
  // fill digits into maps
  for (const auto& eclCalDigit : m_eclCalDigits) {
    const double energy         = eclCalDigit.getEnergy();
    const double time           = eclCalDigit.getTime();
    const double timeresolution = eclCalDigit.getTimeResolution();
    const int cellid            = eclCalDigit.getCellId();
    const bool fitfailed        = eclCalDigit.isFailedFit();
 
    double timeresidual = 999.;
    if (!fitfailed and fabs(timeresolution) > 1e-9) {
      timeresidual = time / timeresolution;
    }
 
    // Negative timecut is interpreted as cut on time residual, positive cut as cut on the time!
    // Start filling all crystals to a map. Growth and seed crystals are strict subsets.
    if (std::isgreaterequal(energy, m_energyCut[2])) {
      if (fitfailed > 0 and m_skipFailedTimeFitDigits > 0) continue;
      if (!fitfailed
          and energy < m_timeCut_maxEnergy[2]) { //check timing cuts only if we have a good fit and if the energy is below threshold
        if (m_timeCut[2] > 1e-9 and fabs(time) > m_timeCut[2]) continue;
        if (m_timeCut[2] < -1e-9  and fabs(timeresidual) > fabs(m_timeCut[2])) continue;
      }
      m_cellIdToDigitVec[cellid] = 1;
      B2DEBUG(250, "ECLCRFinderModule::event(), adding 'all digit' cellid = " << cellid << " " << energy << " " << time << " " <<
              timeresidual);
 
      // check growth only if they already passed the digit check
      if (std::isgreaterequal(energy, m_energyCut[1])) {
        if (!fitfailed
            and energy < m_timeCut_maxEnergy[1]) { //check timing cuts only if we have a good fit and if the energy is below threshold
          if (m_timeCut[1] > 1e-9 and fabs(time) > m_timeCut[1]) continue;
          if (m_timeCut[1] < -1e-9  and fabs(timeresidual) > fabs(m_timeCut[1])) continue;
        }
        m_cellIdToGrowthVec[cellid] = 1;
        B2DEBUG(250, "ECLCRFinderModule::event(), adding 'growth digit' cellid = " << cellid << " " << energy << " " << time << " " <<
                timeresidual);
 
 
        // check seed only if they already passed the growth check
        if (std::isgreaterequal(energy, m_energyCut[0])) {
          if (!fitfailed
              and energy < m_timeCut_maxEnergy[0]) { //check timing cuts only if we have a good fit and if the energy is below threshold
            if (m_timeCut[0] > 1e-9 and fabs(time) > m_timeCut[0]) continue;
            if (m_timeCut[0] < -1e-9  and fabs(timeresidual) > fabs(m_timeCut[0])) continue;
          }
          m_cellIdToSeedVec[cellid] = 1;
          B2DEBUG(250, "ECLCRFinderModule::event(), adding 'seed digit' cellid = " << cellid << " " << energy << " " << time << " " <<
                  timeresidual);
        } // end seed
      } //end growth
    }// end digit
  }//end filling maps
 
  // we start with seed crystals A and attach all growth crystals B
  std::vector<std::vector<int>> connectedRegions_AB = getConnectedRegions(m_cellIdToSeedVec, m_cellIdToGrowthVec, 0);
  std::vector<int> connectedRegions_AB_flattened = flattenVector(connectedRegions_AB);
  std::vector<int> AB = oneHotVector(connectedRegions_AB_flattened, m_cellIdToSeedVec.size());
 
  // Check if any of the growth crystals could grow to other growth crystals
  std::vector<std::vector<int>> connectedRegions_ABB = getConnectedRegions(AB, m_cellIdToGrowthVec, 0);
  std::vector<int> connectedRegions_ABB_flattened = flattenVector(connectedRegions_ABB);
  std::vector<int> ABB = oneHotVector(connectedRegions_ABB_flattened, AB.size());
 
  // and finally: attach all normal digits
  std::vector<std::vector<int>> connectedRegions_ABBC = getConnectedRegions(ABB, m_cellIdToDigitVec, 0);
 
  //final step: merge all CRs that share at least one crystal
  std::vector<std::set<int>> connectedRegionsMerged_ABBC_sets = mergeVectorsUsingBFSTraversal(connectedRegions_ABBC);
 
  // Create CRs and add relations to digits.
  unsigned int connectedRegionID = 0;
  for (const auto& xcr : connectedRegionsMerged_ABBC_sets) {
 
    // Append to store array
    const auto aCR = m_eclConnectedRegions.appendNew();
 
    // Set CR ID
    aCR->setCRId(connectedRegionID);
    connectedRegionID++;
 
    // Add all digits
    for (int x : xcr) {
      const int pos = m_calDigitStoreArrPosition[x];
      aCR->addRelationTo(m_eclCalDigits[pos], 1.0);
    }
  }
 
}
 
void ECLCRFinderModule::endRun()
{
  B2DEBUG(200, "ECLCRFinderModule::endRun()");
}
 
 
void ECLCRFinderModule::terminate()
{
  B2DEBUG(200, "ECLCRFinderModule::terminate()");
  for (unsigned int i = 0; i < m_neighbourMaps.size(); i++) {
    if (m_neighbourMaps[i]) delete m_neighbourMaps[i];
  }
 
}
 
std::vector<int> ECLCRFinderModule::flattenVector(std::vector<std::vector<int>>& A)
{
  std::vector<int> C;
  for (const auto& B : A) {
    C.insert(C.end(), B.begin(), B.end());
  }
  std::sort(C.begin(), C.end());
  C.erase(std::unique(C.begin(), C.end()), C.end());
  return C;
}
 
std::vector<int> ECLCRFinderModule::oneHotVector(std::vector<int>& A, const int n)
{
  std::vector<int> C(n, 0);
  for (int x : A) {
    if (x >= 0 && x < n) {
      C[x] = 1;
    }
  }
  return C;
}
 
std::vector<std::set<int>> ECLCRFinderModule::mergeVectorsUsingBFSTraversal(std::vector<std::vector<int>>& A)
{
  std::vector<std::set<int>> output;
 
  // 1. Data Structures for lookup
  // We use a fixed-size adjacency list for all crystals (Max ID 8736)
  std::vector<int> relevantNodes;
  relevantNodes.reserve(A.size() * 4);
 
  for (int node = 0; node < 8737; ++node) {
    m_adj[node].clear();
    m_visited[node] = false;
  }
 
  // 2. Build the Graph
  // If a vector contains {1, 2, 3}, we mark them as connected neighbors.
  for (const auto& group : A) {
    if (group.empty()) continue;
 
    int root = group[0];
    relevantNodes.push_back(root);
 
    for (size_t i = 1; i < group.size(); ++i) {
      int neighbor = group[i];
      relevantNodes.push_back(neighbor);
 
      // Add bi-directional edges
      m_adj[root].push_back(neighbor);
      m_adj[neighbor].push_back(root);
    }
  }
 
  // 3. Find Connected Components (BFS)
  // We only iterate over nodes that actually appeared in the event.
  for (int startNode : relevantNodes) {
    if (m_visited[startNode]) continue;
 
    // Start a new cluster
    std::set<int> component;
    std::vector<int> q; // BFS Queue
 
    m_visited[startNode] = true;
    q.push_back(startNode);
    component.insert(startNode);
 
    int head = 0;
    while (head < static_cast<int>(q.size())) {
      int u = q[head++];
 
      for (int v : m_adj[u]) {
        if (!m_visited[v]) {
          m_visited[v] = true;
          component.insert(v);
          q.push_back(v);
        }
      }
    }
 
    // Move the finished component to output
    output.push_back(std::move(component));
  }
 
  return output;
}
 
std::vector<std::vector<int>> ECLCRFinderModule::getConnectedRegions(const std::vector<int>& A, const std::vector<int>& B,
                           const int maptype)
{
  std::vector<std::vector<int>> connectedRegions;
 
  // We iterate 0..8737 for the seeds (A), as A is a sparse vector mapped to cellID
  for (unsigned int i = 0; i < A.size(); ++i) {
    if (A[i] > 0) {
      std::vector<int> region;
      // Reserve with neighbor count to prevent reallocations
      region.reserve(21);
      region.push_back(i);
 
      // Instead of looping j = 0 to 8737, we only loop over actual geometric neighbors.
      const auto& neighbors = m_neighbourMaps[maptype]->getNeighbours(i);
 
      for (int neighborID : neighbors) {
        // We only care if this specific neighbor is active in vector B
        if (neighborID < static_cast<int>(B.size()) && B[neighborID] > 0) {
          region.push_back(neighborID);
        }
      }
 
      std::sort(region.begin(), region.end());
      region.erase(std::unique(region.begin(), region.end()), region.end());
 
      connectedRegions.push_back(region);
    }
  }
 
  return connectedRegions;
}