SineSubtract.h

#ifndef _SINE_SUBTRACT_H
#define _SINE_SUBTRACT_H

#ifdef DONT_HAVE_MINUIT2
#warning "Need Minuit2 for SineSubtract." 
#else


/**************************************************************************************
 * \class FFTtools::SineSubtract 
 *
 * \brief This class contains an independent implementation of the sine
 * subtraction CW removal algorithm, suggested by Andres/ Steph
 * (ANITA Elog 621).
 *
 * Cosmin Deaconu <cozzyd@kicp.uchicago.edu> is responsible for this mess. 
 *  
 * This, SineSubtract,  is the controlling class. Related classes are
 * SineFitter, which does the minimization, SineFitter::SineFitFn, which
 * provides the objective function and its derivative, SineFitterLimits, which
 * encodes parameter limits, and SineSubtractResult, which stores stuff. 
 *
 * Many options are available, some of which may even be documented. For
 * example, the part of the waveform used to train sine subtraction may be
 * limited (e.g. if you wanted to just use the non-trigger portion).  Multiple
 * waveforms may be fit for at once (a single frequency is fit for in each
 * iteration, with the amplitude and phase allowed to vary per waveform). It is
 * also possible to enable storage of intermediate graphs and power spectra at
 * each step.  This is useful for diagnostics and making pretty plots, but
 * incurs some overhead, so is not recommended for production use.
 * Gratuitously, there are methods to generate plots and even Beamer slides. 
 *
 * A sketch of the algorithm is: 
 *
 *    - Rectify waveform and estimate power spectrum (according to the power spectrum option). 
 *
 *    - Find peak, fmax, in power spectrum by finding max magnitude/(1+nfails)^nfails_exp) , 
 *      according to the peak finding option. See below for meaning of nfails. 
 *
 *    - Find sine ( A (sin wt + phi)) that when subtracted minimizes "power"
 *    (avg(v_i^2)); initialize with w = 2pi fmax, A = mag, phase from fft.
 *    Uninterpolated graph is used for this! 
 *
 *    - If power reduced by more than min_power_reduction, subtract 
 *
 *    - If power not reduced by more than min_power reduction, increment nfails
 *    for that bin 
 *
 *    - Repeat until max_iter_without_reduction iterations with not enough
 *    improvement   
 *
 * 
 * --- 
 *
 *  A lot of time has been spent profiling and trying to make this run as fast
 *  as possible. This results in code that is more difficult to understand than
 *  a naive implementation would allow. Any suggestions for performance
 *  enhancements are awelcome. 
 * 
 *  If ENABLE_VECTORIZE is defined, custom vectorized code (using Agner Fog's
 *  VCL, http://www.agner.org/optimize/#vectorclass) is used for the objective
 *  function and gradient.  At least until compilers get smart enough to
 *  autovectorize stuff better, this is substantially faster, but obviously
 *  only makes sense in a processor with SIMD instructions.  256-bit floating
 *  point registers are used (e.g. AVX) but the VCL will gracefully fallback to
 *  128-bit registers at compile time should your processor be too old. At some
 *  point AVX-512 will appear, so that might require a few changes. 
 *
 *  **WARNING**: You will not get the same "answers" in vectorized mode. This
 *  is expected, due to usage of different math functions and the general
 *  non-associativity of of floating point math. 
 *
 * ---
 *
 *  A number of additional compile-time options influence the behavior: 
 *
 *  - If SINE_SUBTRACT_USE_FLOATS is defined, the fitter will use floats
 *  instead of doubles. In principle, this should be twice the vectorization,
*  but in practice, the overhead of copying and casting seems to be
*  problematic. 
 *
 *  - If SINE_SUBTRACT_FORCE_ALIGNED is defined, the fitter will copy the
 *  doubles into an aligned array. This might make a performance difference on
 *  old processors, but seems to not help for the most part. 
 *  
 *  - If SINE_SUBTRACT_PROFILE is defined, the time spent calling subtractCW is
 *  printed. 
 *
 *  - if SINE_SUBTRACT_DONT_HORIZONTAL_ADD is defined, horizontal addition is
 *  disabled.  You probably don't want to do this, I added this when I was
 *  investigating some strange numerical behavior. 
 *
 *  See macros/testArtificialSubtract.C to see SineSubtract in action.  
 *
 * --- 
 *
 * TODO list: 
 *
 *  - Better  tuning of default step sizes / limits ? 
 *
 *  - Minuit2 adds quite a bit of overhead, and is one of the things preventing
 *  this from running even faster
 *
 ********************************************************************************************/ 

#include "Minuit2/Minuit2Minimizer.h" 
#include "FFTtools.h" 
#include "TObject.h"

#include <vector>

class TGraph; 
class TCanvas; 

namespace FFTtools
{

   class SineFitterLimits
   {
     public: 

       double maxA_relative_to_guess; 

       double minA_relative_to_guess; 

       double max_n_df_relative_to_guess; 

       double phase_start_error; 

       double amp_start_error; 

       double freq_start_error; 



       SineFitterLimits(double maxA_relative_to_guess_ = 4, double minA_relative_to_guess_ = 0.25, 
                        double max_n_df_relative_to_guess_ = 1, double phase_start_error_ = TMath::Pi() / 16, 
                        double amp_start_error_ = 0.1, double freq_start_error_  = 0.1) : 
                               maxA_relative_to_guess(maxA_relative_to_guess_), 
                               minA_relative_to_guess(minA_relative_to_guess_),
                               max_n_df_relative_to_guess(max_n_df_relative_to_guess_) , 
                               phase_start_error(phase_start_error_) , 
                               amp_start_error(amp_start_error_) , 
                               freq_start_error(freq_start_error_)  
                             {;} 

     
             
   }; 

    class SineFitter
    {

      public: 
        SineFitter(); 

        virtual ~SineFitter(); 

        void setGuess(double f, int ntrace, const double *ph, double amp); 

        void doFit(int ntrace, const int * nsamples, const double ** x, const double **y, const double * w = 0, const double ** env = 0); 

        double getFreq() const {return freq;} 

        const double* getPhase() const {return &phase[0];} 

        const double* getAmp() const {return &amp[0];} 

        double getFreqErr() const {return freq_err;} 

        const double *getPhaseErr() const {return &phase_err[0];} 

        const double *getAmpErr() const {return &amp_err[0];} 

        double getPower() const {return min.MinValue(); } 

        void setVerbose(bool v) { verbose=v; } 

        int getStatus() const { return min.Status(); }

        void setLimitOptions(const SineFitterLimits * lims) { limits = *lims; }

        ROOT::Minuit2::Minuit2Minimizer * minimizer() { return &min; }

        TGraph* getEvalRecordGraph(int i = -1){
          i = i < 0 ? grEvalRecords.size() - 1 : i;
          if(i >= 0 && i < grEvalRecords.size()){
            return grEvalRecords[i];
          }
          return NULL;
        }
        size_t nEvalRecords(){return grEvalRecords.size();}
        void SetDoEvalRecord(bool doEvalRecord){fDoEvalRecord = doEvalRecord;}
        Bool_t GetDoEvalRecord() const {return fDoEvalRecord;}
        void deleteEvalRecords();

      private:
        ROOT::Minuit2::Minuit2Minimizer min; 
        bool verbose; 
        double eval(const double *p); 
        double freq;
        std::vector<double> phase; 
        std::vector<double> amp; 
        double freq_err;
        std::vector<double> phase_err;
        std::vector<double> amp_err;
        bool fDoEvalRecord;
        std::vector<TGraph*> grEvalRecords;
      
        double freq_factor; 
        double min_amp_mean, max_amp_mean;  
        double min_amp_guess, max_amp_guess; 
        SineFitterLimits limits; 



#ifndef __CINT__ //CINT is stupid about subclasses 

        class SineFitFn : public ROOT::Math::IGradientFunctionMultiDim 
        {
          public:

            SineFitFn(SineFitter* parent = NULL); 

            virtual ~SineFitFn(); 

            void setXY(int ntraces, const int * nsamples, const double **xx, const double **yy, const double * w = 0, const double ** env = 0); 

            virtual double DoEval(const double *p) const; 

            virtual double DoDerivative(const double *p, unsigned int coord) const; 

            unsigned NDim() const { return 1+2*nt; } 

            virtual ROOT::Math::IBaseFunctionMultiDim * Clone() const;
          
          private:
            SineFitter* fContainer; // pointer to parent
          
            int *ns;
            int nt; 
#ifdef SINE_SUBTRACT_USE_FLOATS
            float **x; 
            float **y; 
            float **env; 
            const double **xp, **yp, ***envp; 
#elif defined(SINE_SUBTRACT_FORCE_ALIGNED)
            double **x; 
            double **y; 
            double **env; 
            const double **xp, **yp, **envp; 
#else
            const double ** x; 
            const double ** y; 
            const double ** env; 
#endif 
            const double * wgt; 

          
        } f;
#endif

    }; 

    /* This class stores the minimization result. It would be a struct if CINT weren't
     * stupid about structs */
    class SineSubtractResult
    {

      public: 

      void clear(); 

      void append(const SineSubtractResult * r); 

      virtual ~SineSubtractResult(){;} 


      std::vector<double> powers; 

      std::vector<std::vector<double> >phases; 

      std::vector<double> freqs; 

      std::vector<std::vector<double> > amps; 

      std::vector<std::vector<double > > phases_errs; 

      std::vector<double> freqs_errs; 

      std::vector<std::vector<double> >amps_errs; 


      ClassDef(SineSubtractResult,3); 
    }; 


    class SineSubtract 
    {

      public: 

        SineSubtract(int max_iter_without_reduction = 3, double min_power_reduction = 0.05, bool store = false); 


        SineSubtract(const TGraph * freq_dependent_min_power_reduction, int max_iter_without_reduction = 3, bool store = false); 


        virtual ~SineSubtract(); 

        /* Subtract CW from a single trace. Sine Subtraction can handle both evenly spaced and unevely spaced waveforms. 
         * For evenly spaced waveforms, you should pass dt <= 0, otherwise you should pass the nominal average. 
         * 
         * The input graph is not touched. A new graph is allocated for the subtraction. 
         *
         * @param g The input TGraph to subtract from 
         * @param dt The nominal sample rate for uneven waveforms. If <=0, then the graph is assumed to be even and dt is computed from the first time step. 
         * @return The subtracted Graph.  
         */
        TGraph * subtractCW(const TGraph * g, double dt, const SineSubtractResult* result = NULL);


        void subtractCW(int ng, TGraph ** g, double dt, const double * w = 0, const SineSubtractResult* result = NULL); 

        void setFreqLimits(double min, double max) { fmin.clear(); fmin.push_back(min); fmax.clear(); fmax.push_back(max); } 

        void setFreqLimits(int nbands, const double * min, const double * max) { fmin.clear(); fmax.clear(); for (int i = 0; i < nbands;i++) { fmin.push_back(min[i]); fmax.push_back(max[i]); }; } 

        void unsetFreqLimits() { fmin.clear(); fmax.clear(); high_factor = 1;  } 


        void setVerbose(bool v) {verbose = v; fitter.setVerbose(v); }


        void setTraceLimits(int min, int max) { tmin = min; tmax = max; }



        enum PowerSpectrumEstimator
        {
          FFT, 
          LOMBSCARGLE 
        }; 


        enum FFT_Params
        {
          FFT_NPAD, 
          FFT_NPARAMS
        }; 

        enum LOMBSCARGLE_Params
        {
          LS_OVERSAMPLE_FACTOR, 
          LS_NPARAMS
        }; 


        void setNFailsExponent(double exp = 0.5) { nfail_exponent = exp ; }


        void setPowerSpectrumEstimator(PowerSpectrumEstimator estimator, const double * params = 0); 


        enum PeakFindingOption
        {
          GLOBALMAX, 
          NEIGHBORFACTOR, 
          TSPECTRUM, 
          SAVGOLSUB 
        }; 


        enum NEIGHBORFACTOR_Params
        {
          NF_NEIGHBOR_FACTOR, 
          NF_NPARAMS
        }; 

        /* These are basically based on the 
         * options to TSpectrum::searchHighRes
         */ 
        enum TSPECTRUM_Params
        {
          TS_SIGMA, 
          TS_THRESHOLD, 
          TS_NDECONV_ITERATIONS, //default = 3
          TS_AVERAGE_WINDOW, //default = 3
          TS_NPARAMS
        }; 

        enum SAVGOLSUB_Params
        {
          SGS_ORDER, 
          SGS_WIDTH, 
          SGS_NPARAMS
        }; 

        void setPeakFindingOption(PeakFindingOption peak, const double * params = 0); 

        enum EnvelopeOption
        {
          ENV_NONE, 
          ENV_HILBERT, //fit to hilbert envelope
          ENV_RMS, //fit to sliding RMS envelope 
          ENV_PEAK //fit to sliding RMS envelope 
        };  

        enum ENV_HILBERT_Params
        {
          ENV_HILBERT_FIT_ORDER, //default = 3, if < 0, no fit (just take values) 
          ENV_HILBERT_NPARAMS

        }; 

        enum ENV_RMS_Params
        {
          ENV_RMS_FIT_ORDER, //default = 3, if < 0, no fit (just take envelope values) 
          ENV_RMS_WINDOW, //default = 5
          ENV_RMS_NPARAMS
        }; 

       enum ENV_PEAK_Params
        {
          ENV_PEAK_FIT_ORDER, //default = 3, if < 0, no fit (just take envelope values) 
          ENV_PEAK_MINDISTANCE, //default = 5 (appropriate for ~200 MHz min freq) 
          ENV_PEAK_NPARAMS
        }; 

        void setEnvelopeOption(EnvelopeOption env, const double * params = 0); 

        const std::vector<double> * getPowerSequence() const { return &r.powers; } 

        const std::vector<std::vector<TGraph*> > * getStoredGraphs() const { return &gs; }

        const std::vector<TGraph*> * getStoredSpectra() const { return &spectra; }


        size_t nStoredGraphsInChannel() const { return gs[0].size(); } 

        size_t nChannels() const { return gs.size(); } 

        TGraph* storedGraph(int i, int c) { return gs[c][i]; } 
        TGraph* storedEnvelope(int i, int c) { return env_gs[c][i]; } 

        size_t nStoredSpectra() const { return spectra.size(); } 

        TGraph* storedSpectra(int i) { return spectra[i]; } 

        int getNSines() const; 
        
        /* Return the number of traces fit for */
        int getNGraphs() const { return r.phases.size(); } 

        const double * getPhases(int g) const { return &r.phases[g][0]; }

        const double * getFreqs() const { return &r.freqs[0]; }

        const double * getAmps(int g) const { return &r.amps[g][0]; }

        const double * getPhaseErrs(int g) const { return &r.phases_errs[g][0]; }

        const double * getFreqErrs() const { return &r.freqs_errs[0]; }

        const double * getAmpErrs(int g) const { return &r.amps_errs[g][0]; }

        const SineSubtractResult * getResult() const { return &r; } 

        /* Set the maximum number of failed iterations before giving up. */
        void setMaxIterationsWithoutReduction(int max) { maxiter = max;  } 

        void setAbsoluteMaxIterations(int max) { abs_maxiter = max; } 

        void setMaxSuccessfulIterations(int max) { max_successful_iter = max; } 

        void setMinPowerReduction(double min) {min_power_reduction = min; g_min_power = 0;}; 

        /* Set the frequency-dependent threshold before giving up. */
        void setMinPowerReduction(const TGraph *g) { g_min_power= g; }  


        void setStore(bool s) { store = s; }

        void setFitterLimitOptions( const SineFitterLimits * lims) { fitter.setLimitOptions(lims); }

        //these are super auxilliary and could be moved out (they only need public methods anyway) 
        
        void makePlots(TCanvas * cpower = 0,TCanvas *cspectra = 0, int ncols = 4) const; 

        void makeSlides(const char * title = "SineSubtract", const char * prefix = "sinsub", const char * outdir = ".", const char* format = "png", bool standalone = true) const; 

        SineFitter * sineFitter() { return &fitter; } 

      private: 
        int abs_maxiter; 
        int maxiter; 
        double nfail_exponent; 
        int max_successful_iter; 
        int tmin, tmax; 
        std::vector<double> fmin; 
        std::vector<double> fmax; 
        double min_power_reduction; 
        std::vector<std::vector<TGraph*> > gs; 
        std::vector<std::vector<TGraph*> > env_gs; 
        std::vector<TGraph*> spectra;
        SineSubtractResult r; 
        bool store; 
        double high_factor; 

        PowerSpectrumEstimator power_estimator;
        double * power_estimator_params; 
        PeakFindingOption peak_option; 
        double * peak_option_params; 


        EnvelopeOption envelope_option; 
        double * envelope_option_params; 

        const TGraph * g_min_power; 

        bool verbose; 
        void reset(); 
        unsigned id; 

        int findMaxFreq(int Nfreq, const double * freq, const double * mag, const int * nfails) const; 
        bool allowedFreq(double f, double df) const; 

        SineFitter fitter; 
    }; 

}

#endif
#endif