IParticlesLookUpTable.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef PARTICLESINCONETOOLS_IPARTICLELOOKUPTABLE_H
#define PARTICLESINCONETOOLS_IPARTICLELOOKUPTABLE_H
 
#include <limits>
#include <cmath>
#include "AthContainers/DataVector.h"
 
namespace xAOD {
 
 
  template<class T>
  class IParticlesLookUpTable {
  public:
  IParticlesLookUpTable( unsigned int nbins = 50 , float minPt=1e-3) :m_container(0), 
      m_nphiBins(nbins), 
      m_phiBinSize(m_2PI/m_nphiBins),
      m_minPt(minPt),
      m_phiBinnedLookUpTable()
      {} 
 
    void init( const DataVector<T>& particles );
      
    void clear();
 
    template<class O>
    bool iParticlesInCone( float eta, float phi, float dr, O& output ) const;
 
    bool isInitialized() const;
 
  private:
    static constexpr float m_2PI = 2*M_PI;
 
    float phiInRange(float phi) const { 
      while (phi >= M_PI) phi -= m_2PI;
      while (phi < -M_PI) phi += m_2PI;
      return phi;
    }
 
    int phiIndex(float phi) const { return (phi + M_PI)/m_phiBinSize; };
 
    void addEntry( int i, std::vector< const T* >& output ) const;
 
    void addEntry( int i, std::vector< ElementLink<DataVector<T> > >& output ) const;
 
    
    const DataVector<T>* m_container; //internal container
    int   m_nphiBins;    
    float m_phiBinSize;  
    float m_minPt;       
    std::vector< std::vector< int > > m_phiBinnedLookUpTable; 
  };
 
  template<class T>
  bool IParticlesLookUpTable<T>::isInitialized() const { return !m_phiBinnedLookUpTable.empty(); }
 
 
  template<class T>
  void IParticlesLookUpTable<T>::clear() {
    m_phiBinnedLookUpTable.clear();
  }
 
  template<class T>
  void IParticlesLookUpTable<T>::init( const DataVector<T>& particles ) {
    m_phiBinnedLookUpTable.resize(m_nphiBins);
    m_container = &particles;
    unsigned int size = particles.size();
    for( unsigned int i=0; i<size;++i ){
      if(particles[i]->pt()>m_minPt){//sanity check (mainly due to Truth)
    int index = phiIndex(particles[i]->phi());
    m_phiBinnedLookUpTable[index].push_back(i);
      }
    }
    //std::cout << " initializing lookup " << m_container->size() << std::endl;
    for( auto& vec : m_phiBinnedLookUpTable ) {
      if( vec.empty() ) continue;
      std::stable_sort(vec.begin(),vec.end(),[&](int i, int j) { return (*m_container)[i]->eta() < (*m_container)[j]->eta(); });
      // std::cout << " new phi slice " << vec.size() << " eta " ;
      // for( auto index : vec ){
      //   std::cout << "  " << (*m_container)[index]->eta();
      // }
      // std::cout << std::endl;
    }
  }
 
  template<class T> template<class O>
  bool IParticlesLookUpTable<T>::iParticlesInCone( float eta, float phi, float dr, O& output ) const {
 
    if( m_phiBinnedLookUpTable.empty() ) return false;
  
    auto compEta1 = [&](int i,double val ) { return (*m_container)[i]->eta() < val; };
    auto compEta2 = [&](double val,int i ) { return val < (*m_container)[i]->eta(); };
 
    int indexMin = phiIndex( phiInRange(phi-dr) );
    int indexMax = phiIndex( phiInRange(phi+dr) );
    //std::cout << " eta " << eta << " phi " << phi << " phi ranges " << indexMin << " " << indexMax << std::endl;
    // special treatment for boundary region
    std::vector< std::pair<int,int> > ranges;
    if( indexMin > indexMax ) {
      ranges.push_back( std::make_pair(0,indexMax) );
      ranges.push_back( std::make_pair(indexMin,m_nphiBins-1) );      
    }else{
      ranges.push_back( std::make_pair(indexMin,indexMax) );
    }
    
    float dr2Cut = dr*dr;
 
    // loop over ranges 
    for( auto& range : ranges ){
      indexMin = range.first;
      indexMax = range.second;
      for( ; indexMin <= indexMax; ++indexMin ){
        // get iterators for iparticles to be included 
        const std::vector< int >& tps = m_phiBinnedLookUpTable[indexMin];
        auto it_min = std::lower_bound (tps.begin(),tps.end(),eta-dr,compEta1  );
        auto it_max = std::upper_bound (it_min,tps.end(),eta+dr,compEta2 );
        //std::cout << " new range, entries in bounds " << std::distance(it_max,it_min) << std::endl;
        // add iparticles in cone
        for( ;it_min!=it_max;++it_min ){
          const T* entry = (*m_container)[*it_min];
          float deta = eta- entry->eta();
          float dphi = phiInRange(phi-entry->phi());
          float dr2 = deta*deta + dphi*dphi;
          //std::cout << "   new entry: " << *it_min << " eta,phi " << entry->eta() << " phi " << entry->phi() << " dr " << sqrt(dr2) << std::endl;
          if( dr2 < dr2Cut )  addEntry( *it_min, output );
        }
      }
    }
    return true;
  }
 
  template<class T>
  void IParticlesLookUpTable<T>::addEntry( int i, std::vector< const T* >& output ) const {
     output.push_back( (*m_container)[i]);
  }
 
  template<class T>
  void IParticlesLookUpTable<T>::addEntry( int i, std::vector< ElementLink<DataVector<T> > >& output ) const {
    output.push_back( ElementLink<DataVector<T> >(*m_container,i) );
  }
 
}   // end of namespace
 
#endif