code/include/NewWave/WaveletEvent.hh

Go to the documentation of this file.

#ifndef NEWWAVE_WAVELET_EVENT_HH
#define NEWWAVE_WAVELET_EVENT_HH

#include "NewWave/WaveletBaseEvent.hh"
#include "NewWave/PixelDefinition.hh"
#include "NewWave/WaveletEngine.hh"
#include "NewWave/Exceptions.hh"
#include "NewWave/FrequencyBand.hh"
#include "NewWave/Utils.hh"
#include "NewWave/MomentumHelpers.hh"

#include <map>
#include <cmath>

#include <iostream>

namespace HepMC{
  class GenEvent;
}

namespace NewWave{
  
  using std::size_t;
  using std::map;
  using std::make_pair;
  
  template<typename T, typename momentum_type = Momentum<T> >
  class WaveletEvent : public WaveletBaseEvent{
    
  public:
    
    WaveletEvent(const T &particles,
                 const PixelDefinition &pixelDefn,
                 const WaveletEngine &engine): WaveletBaseEvent(),
    _pixelDefn(pixelDefn),
    _originalPixels(pixelDefn.makeEmptyPixelArray()),
    _engine(engine),
    _doInvert(true),
    _doParticles(true),
    _pileUpThreshold(0.),
    _doScale(true){
      
      _init(particles);
    }
    

    const WaveletCoefficients &coefficients()const{return _coefficients();}
    
    const FrequencyBands &frequencyBands()const{
      
      if(_frequencyBands.size() == _pixelDefn.nLevels() * _pixelDefn.nLevels()) return _frequencyBands;
      
      _frequencyBands.clear();
      _bandLookup.clear();
            
      for(const auto &c: _coefficients()){
        int hash = c.frequencyHash(_pixelDefn.nLevels());
        map<int, int>::const_iterator index = _bandLookup.find(hash);
        if(index==_bandLookup.end()){
          index = _bandLookup.insert(make_pair(hash, _frequencyBands.size())).first;
          _frequencyBands.push_back(FrequencyBand(c.yLevel(), c.phiLevel()));
        }
        
        _frequencyBands.at(index->second).addCoefficient(c);
      }
      
      return _frequencyBands;
    }
    
    const FrequencyBand &frequencyBand(const WaveletCoefficient &coeff)const{
      int hash = coeff.frequencyHash(_pixelDefn.nLevels());
      frequencyBands();
      map<int, int>::const_iterator index = _bandLookup.find(hash);
      if(index==_bandLookup.end()){
        throw InvalidFrequencyBand();
      }
      
      return _frequencyBands.at(index->second);
    }
    
    const PixelDefinition &pixelDefn()const{
      return _pixelDefn;
    }
    

    const PixelArray &pixels()const{
      
      if(!_doInvert) return _modifiedPixels;
      
      _modifiedPixels = _engine.inverseTransform(coefficients());
      _doInvert = false;
      return _modifiedPixels;
    }
    
    const PixelArray &originalPixels()const{
      return _originalPixels;
    }
    

    const T &particles()const{
      
      if(!_doInvert && !_doParticles) return _modifiedParticles;
      
      _ratio = pixels() / _originalPixels ;
      
      _modifiedParticles.clear();
      
      for(auto p: _originalParticles){
        size_t ybin   = _pixelDefn.yPixelIndex(momentum_type::rapidity(p));
        size_t phiBin = _pixelDefn.phiPixelIndex(momentum_type::phi(p));
        double ratio = _ratio.at(ybin).at(phiBin);
        
        momentum_type::update(_modifiedParticles, p, ratio, _pileUpThreshold, _doScale);
        
      }
      
      _doParticles = false;
      
      return _modifiedParticles;
    }
    
    

    void denoise(double noiseThreshold){
      
      filter([noiseThreshold](const WaveletCoefficient &c)->bool {return (fabs(c.value()) > noiseThreshold); });
      return;
    }


    void filter(const std::function<bool (const WaveletCoefficient&)> &select){
      
      for(WaveletCoefficient &coeff: _coefficients()){
        if(!select(coeff)){
          coeff.setValue(0.);
          _doInvert = true;
        }
      }
      return;
    }
    
    void scale(const std::function<double (const WaveletCoefficient&)> &scaler){
      for(WaveletCoefficient &coeff: _coefficients()){
        double s = scaler(coeff);
        coeff.setValue(s*coeff.value());
      }
      
      _doInvert = true;
      
      return;
    }
    
    
    void setPileUpThreshold(double threshold){
      if(threshold < 0.){
        throw PileUpThresholdException();
      }
      _pileUpThreshold = threshold;
      _doInvert = true;
      return;
    }
    

    void setScaleParticles(bool doScale){
      _doScale = doScale;
    }
    
  private:
    
    WaveletCoefficients &_coefficients()const{
      if(_coeffs.size() != 0) return _coeffs;
      _coeffs = _engine.transform(_originalPixels);
      for(auto &c: _coeffs){
        setEvent(c);
      }
      
      return _coeffs;
    }
    
    void _init(const T &input){
      
      for(auto p: input){
        if(_pixelDefn.covers(momentum_type::rapidity(p), momentum_type::phi(p))){
          _originalParticles.push_back(p);
          _addParticle(momentum_type::rapidity(p), momentum_type::phi(p), momentum_type::pT(p));
        }
      }
      
    }
    
    void _addParticle(double rapidity, double phi, double pT){
      size_t ybin   = _pixelDefn.yPixelIndex(rapidity);
      size_t phiBin = _pixelDefn.phiPixelIndex(phi);
      _originalPixels.at(ybin).at(phiBin) += pT;
      return;
    }
    
    HepMC::GenEvent *_genEvent()const;
    
    PixelDefinition _pixelDefn;
    
    mutable PixelArray _modifiedPixels;
    PixelArray _originalPixels;
    
    T _originalParticles;
    
    const WaveletEngine &_engine;
    
    mutable WaveletCoefficients _coeffs;
    
    mutable FrequencyBands _frequencyBands;
    mutable map<int, int> _bandLookup;
    
    mutable bool _doInvert;
    mutable bool _doParticles;
    
    double _pileUpThreshold;
    
    bool _doScale;
    
    mutable PixelArray _ratio;
    
    mutable T _modifiedParticles;
    
  };
  
  template<>
  HepMC::GenEvent* const &WaveletEvent<HepMC::GenEvent *>::particles()const;
  
  template<>
  void WaveletEvent<HepMC::GenEvent*>::_init(HepMC::GenEvent* const &input);
  
  
}

#endif