SPTCSelectorXBestPerFamily.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 <vector>
#include <numeric>
 
#include <tracking/spacePointCreation/SpacePointTrackCand.h>
#include <tracking/trackFindingVXD/trackQualityEstimators/QualityEstimatorBase.h>
#include <tracking/trackFindingVXD/trackQualityEstimators/QualityEstimatorTripletFit.h>
#include <tracking/trackFindingVXD/trackQualityEstimators/QualityEstimatorCircleFit.h>
#include <tracking/trackFindingVXD/trackQualityEstimators/QualityEstimatorRiemannHelixFit.h>
#include <tracking/trackFindingVXD/trackQualityEstimators/QualityEstimatorRandom.h>
 
 
namespace Belle2 {
  class SPTCSelectorXBestPerFamily {
 
  public:
 
    SPTCSelectorXBestPerFamily(unsigned short xBest = 5, const std::string& estimationMethod = std::string("tripletFit")):
      m_xBest(xBest)
    {
      if (estimationMethod == "tripletFit") {
        m_estimator = std::make_unique<QualityEstimatorTripletFit>();
      } else if (estimationMethod == "circleFit") {
        m_estimator = std::make_unique<QualityEstimatorCircleFit>();
      } else if (estimationMethod == "helixFit") {
        m_estimator = std::make_unique<QualityEstimatorRiemannHelixFit>();
      } else if (estimationMethod == "random") {
        m_estimator = std::make_unique<QualityEstimatorRandom>();
      }
      B2ASSERT("QualityEstimator could not be initialized with method: " << estimationMethod, m_estimator);
    };
 
    ~SPTCSelectorXBestPerFamily() = default;
 
    void prepareSelector(unsigned short nFamilies)
    {
      m_bestPaths.clear();
      m_familyToIndex.clear();
      m_bestPaths.reserve(nFamilies);
      m_familyToIndex.resize(nFamilies, -1);
 
      m_currentIndex = 0;
    }
 
    void testNewSPTC(SpacePointTrackCand& sptc)
    {
      sptc.setQualityIndicator(m_estimator->estimateQuality(sptc.getSortedHits()));
      short family = sptc.getFamily();
 
      if (m_familyToIndex.at(family) == -1) {
        m_bestPaths.emplace_back(std::multiset<SpacePointTrackCand> { sptc });
        m_familyToIndex.at(family) = m_currentIndex;
        m_currentIndex++;
        return;
      }
 
      auto& currentSet = m_bestPaths.at(m_familyToIndex[family]);
 
      if (currentSet.size() == m_xBest) {
        std::multiset<SpacePointTrackCand>::iterator iter = currentSet.begin();
        if (sptc.getQualityIndicator() < iter->getQualityIndicator()) {
          return;
        }
        currentSet.erase(iter);
      }
      currentSet.emplace_hint(currentSet.cbegin(), sptc);
    }
 
    std::vector<SpacePointTrackCand> returnSelection() const
    {
      std::vector<SpacePointTrackCand> jointBestPaths;
      jointBestPaths.reserve(std::accumulate(m_bestPaths.begin(), m_bestPaths.end(), 0,
      [](int a, auto b) { return a + b.size(); }));
 
      for (auto&& set : m_bestPaths) {
        jointBestPaths.insert(jointBestPaths.end(), set.begin(), set.end());
      }
 
      return jointBestPaths;
    }
 
    void setMagneticFieldForQE(double bFieldZ)
    {
      m_estimator->setMagneticFieldStrength(bFieldZ);
    }
 
  private:
 
    std::unique_ptr<QualityEstimatorBase> m_estimator;
 
    std::vector<std::multiset<SpacePointTrackCand> > m_bestPaths;
 
    std::vector<short> m_familyToIndex;
 
    unsigned short m_currentIndex = 0;
 
    unsigned short m_xBest;
  };
}