BmnParticleEqualizer2D.h

Go to the documentation of this file.

// -------------------------------------------------------------------------
// -----                BmnParticleEqualizer2D header file               -----
// -----              Created 01/05/2024 by K. Gertsenberger           -----
// -------------------------------------------------------------------------
 
/** BmnParticleEqualizer2D.h
 ** @author Konstantin Gertsenberger <gertsen@jinr.ru>
 **
 ** The BmnParticleEqualizer2D add particles to equalize their number for all momentum intervals
 **
 ** The generator counts the particles already put on the stack by the other generators of the event and adds new ones
 ** to equalize their number, so it MUST be registered in FairPrimaryGenerator as the LAST generator (otherwise it
 ** would not see the particles of the generators registered after it). The requirement is verified in ReadEvent().
 **
 ** The kinematics of a new particle is sampled from the two dimensional sample histograms of its momentum interval:
 ** (P, Eta) gives the momentum and the pseudorapidity together, keeping their correlation, and then the azimuth angle
 ** is taken from the (Eta, Phi) histogram for the sampled Eta. The angles are therefore distributed as in the sample,
 ** i.e. inside the detector acceptance, and not uniformly.
 */
 
#ifndef BMNPARTICLEEQUALIZER2D_H
#define BMNPARTICLEEQUALIZER2D_H
 
#include "BmnGlobalTrack.h"
#include "BmnTrackMatch.h"
#include "FairGenerator.h"
#include "FairPrimaryGenerator.h"
#include "TClonesArray.h"
#include "TF1.h"
#include "TFile.h"
#include "TH1D.h"
#include "TH2D.h"
 
#include <map>
#include <memory>
#include <vector>
using namespace std;
 
class BmnParticleEqualizer2D : public FairGenerator
{
  public:
    /** Constructor with input histograms **/
    BmnParticleEqualizer2D(TString hist_file_name);
    /** Constructor with given PDG codes for particles to be equalized
    **@param pdg_codes Particles to be equalized (PDG codes)
    **@param min_p    Min value of momentum value
    **@param max_p    Max value of momentum value
    **@param intervals Interval count for momentum values to equalize particles at all the intervals independently
    **/
    BmnParticleEqualizer2D(TString hist_file_name,
                           vector<Int_t> pdg_codes,
                           Double_t min_p = 0.0,
                           Double_t max_p = 4.0,
                           Int_t intervals = 20);
 
    /** Destructor */
    virtual ~BmnParticleEqualizer2D() {}
 
    /** Initializer */
    Bool_t Init();
    /** Creates an event with equal number of particles at all momentum intervals
    **@param primGen   pointer to the FairPrimaryGenerator
    */
    virtual Bool_t ReadEvent(FairPrimaryGenerator* primGen);
 
    struct stParticleInfo
    {
        // Default 'stParticleInfo' constructor for storing particle distrubutons
        stParticleInfo(int pdg_key, Double_t min_p, Double_t max_p, Int_t interval_count);
        // 'stParticleInfo' destructor, empty becausse of using unique pointers
        ~stParticleInfo() {}
 
        vector<Int_t> fParticleNumber;   ///< PDG -> Particle Count (divided by momentum intervals -> vector)
        /// Part of the target number of the particles that has not been added yet, kept between the events so that a
        /// fractional target is reached on average (divided by momentum intervals -> vector)
        vector<Double_t> fParticleDebt;
        /// Relative reconstruction efficiency of the particle type, 1 for all the intervals if the efficiency map is
        /// absent in the file with the sample histograms (divided by momentum intervals -> vector)
        vector<Double_t> fEfficiency;
 
        unique_ptr<TH1D> fPHist;       ///< PDG -> P histogram (TH1D) for the particle type
        unique_ptr<TH1D> fPtHist;      ///< PDG -> Pt histogram (TH1D) for the particle type
        unique_ptr<TH1D> fYHist;       ///< PDG -> Y (rapidity) histogram (TH1D) for the particle type
        unique_ptr<TH1D> fEtaHist;     ///< PDG -> Eta (pseudorapidity) histogram (TH1D) for the particle type
        unique_ptr<TH1D> fThetaHist;   ///< PDG -> Theta (polar angle) histogram (TH1D) for the particle type
        unique_ptr<TH1D> fPhiHist;     ///< PDG -> Phi (azimuth angle) histogram (TH1D) for the particle type
 
        /// PDG -> (P, Eta) histogram (TH2D) of the whole momentum range, used to sample Eta for a given P when the
        /// momentum interval itself has too little statistics
        unique_ptr<TH2D> fPEtaHist;
        /// PDG -> (Eta, Phi) histogram (TH2D) of the whole momentum range, used as a fallback for the same reason
        unique_ptr<TH2D> fEtaPhiHist;
 
        vector<unique_ptr<TH1D>>
            fPHistVector;   ///< P histograms for the particle type (divided by momentum intervals -> vector)
        vector<unique_ptr<TH1D>>
            fEtaHistVector;   ///< Eta histograms for the particle type (divided by momentum intervals -> vector)
        vector<unique_ptr<TH1D>>
            fPhiHistVector;   ///< Phi histograms for the particle type (divided by momentum intervals -> vector)
        /// (P, Eta) histograms for the particle type (divided by momentum intervals -> vector). The Eta axis is
        /// identical to the one of 'fEtaPhiHistVector', so an Eta sampled here always has statistics there as well
        vector<unique_ptr<TH2D>> fPEtaHistVector;
        /// (Eta, Phi) histograms for the particle type (divided by momentum intervals -> vector)
        vector<unique_ptr<TH2D>> fEtaPhiHistVector;
        /// Rebinned copies of 'fPEtaHistVector' used to check whether the sample has statistics for a particle of the
        /// event, see SetSampleStatisticsCutRebin (divided by momentum intervals -> vector)
        vector<unique_ptr<TH2D>> fPEtaCutHistVector;
        /// The same rebinned copies of 'fEtaPhiHistVector' (divided by momentum intervals -> vector)
        vector<unique_ptr<TH2D>> fEtaPhiCutHistVector;
 
        /// Way the kinematics of a new particle is sampled in a momentum interval
        enum class EPSampling
        {
            kNone,           ///< the interval cannot be sampled at all, new particles are not added there
            kIntervalHist,   ///< P and Eta together from the (P, Eta) histogram of the interval itself
            kFitFunction     ///< P from the fitted function of the full P histogram (too little sample statistics)
        };
        vector<EPSampling>
            fPSampling;   ///< Momentum sampling way for the particle type (divided by momentum intervals -> vector)
 
        Double_t fY0;      ///< Mean rapidity
        Double_t fSigma;   ///< Sigma in the rapidity distribution
        Double_t fT;       ///< Temperature in the P distribution
    };
 
    /** Status bit the generator sets on every particle it adds. TObject keeps the bits 0-13 of fBits for itself, so
    ** 14 is the first one free for the user. CbmMCTrack copies these bits from TParticle, so the mark reaches the
    ** reconstructed (DST) data and the added particles can be told there from the ones of the physics generators:
    ** mc_track->TestBit(kAddedParticleBit) **/
    enum
    {
        kAddedParticleBit = BIT(14)
    };
 
    /** Set the target number of the particles of every type to be reached in every momentum interval of an event.
    ** Zero (default) keeps the original behaviour: the target of an interval is the maximum number among the particle
    ** types found in that interval of that event, so the numbers are equalized between the types but the momentum
    ** spectrum stays the one of the input events. A positive value makes the target the same for all the intervals and
    ** all the types, which gives a flat distribution both in the particle type and in the momentum - what is usually
    ** needed to train a PID. The value may be fractional (and normally is: it is the number of particles per event per
    ** interval, i.e. the wanted sample size divided by the number of events and by the number of intervals), the parts
    ** below one particle are accumulated between the events **/
    void SetTargetPerInterval(Double_t target) { fTargetPerInterval = target; }
    /** Set how much the two dimensional sample histograms are rebinned for the filtering of the event particles by
    ** the sample statistics (see SetSampleStatisticsCut). The filtering asks whether the sample has at least one entry
    ** in the cell a particle of the event falls into, so with the fine binning needed for a good sampling it becomes
    ** the more severe the smaller the sample is: measured on a sample of 100 tracks per momentum interval it discards
    ** about a half of the particles. The rebinned copies decouple the granularity of the filtering from the one of the
    ** sampling: the particles are still sampled from the full histograms, only the filtering uses the coarse ones.
    ** The factor has to divide the number of the bins of both axes; 1 makes the filtering use the full binning.
    ** Returns kFALSE (and keeps the previous value) if the factor cannot be applied **/
    Bool_t SetSampleStatisticsCutRebin(Int_t factor);
    /** Switch on the correction of the number of the added particles by the reconstruction efficiency (off by
    ** default). The efficiency map has to be present in the file with the sample histograms, see ProduceEfficiencyMap.
    ** When it is on, the number of the particles added to an interval is divided by the relative efficiency of that
    ** interval, compensating for the added particles lost at the reconstruction. 'min_efficiency' limits the resulting
    ** factor: the intervals with a lower efficiency are corrected as if it were equal to it **/
    Bool_t SetEfficiencyCorrection(Bool_t is_used, Double_t min_efficiency = 0.05);
 
    /** Set the cut applied to the tracks of the sample while producing the sample histograms (and while drawing the
    ** resulting distributions by ShowResultDistributions): 0 - all the primary particles are taken, 1 - the tracks
    ** without points in Silicon and GEM detectors are excluded, 2 (default) - the tracks without points in Silicon,
    ** GEM and ToF detectors are excluded. It defines the acceptance the sample histograms correspond to, so it has to
    ** be set before ProduceSampleHistograms() and cannot be changed for an already produced file. The reconstructed
    ** tracks counted by ShowResultDistributions and ProduceEfficiencyMap are selected in the same acceptance: the cut 2
    ** requires a ToF hit matched to the track **/
    static void SetSampleTrackCut(Int_t track_cut) { fSampleTrackCut = track_cut; }
    /** Switch the filtering of the event particles by the sample statistics on or off (on by default). When it is on,
    ** the particles of the event lying in the phase space regions where the sample histograms are empty are not
    ** counted: the generator cannot add such particles back, so counting them would leave the numbers unequalized **/
    void SetSampleStatisticsCut(Bool_t is_used) { fUseSampleStatisticsCut = is_used; }
 
    // static function to build the map of the reconstruction efficiency from the reconstructed (DST) files listed in
    // the given text file and to add it to the file with the sample histograms. Both the simulated and the
    // experimental DST files can be given: the particle type of a track is taken from the MC track matched to it when
    // it is available and is defined by the mass measured by the ToF detectors and by the sign of the charge
    // otherwise. The map is the number of the reconstructed tracks of a particle type in a momentum interval divided
    // by the number of the sample tracks in the same interval, normalized so that its maximum is 1, i.e. it is a
    // RELATIVE efficiency
    // return 0 if the map has been built successfully, a negative error code otherwise
    static int ProduceEfficiencyMap(TString input_list_file,
                                    TString histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
 
    // static function to build the map of the acceptance of the sample tracks from the simulated (SIM or DST) files
    // listed in the given text file and to add it to the file with the sample histograms. For every particle type and
    // momentum interval two (Eta, Phi) histograms of the primary particles are written: all of them
    // ('hAccAll<i>_pdg<code>') and the ones passing the sample track cut ('hAccPass<i>_pdg<code>', by default the ones
    // with points in Silicon, GEM and at least one of the ToF detectors, see SetSampleTrackCut). Their ratio is the
    // probability for a particle to get into the acceptance, the decays in flight included. The map is not used by
    // the generator yet
    // return 0 if the map has been built successfully, a negative error code otherwise
    static int ProduceAcceptanceMap(TString input_list_file,
                                    TString histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
 
    // static function to form sample histograms for equalizing the particle distributions
    // return 0 if the histograms have been produced successfully, a negative error code otherwise
    static int ProduceSampleHistograms(TString input_list_file,
                                       TString output_histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
 
    // fit P. Pt, Y, Eta histograms to generate the values according to the analytical functions
    static int FitPHistogram(unique_ptr<TH1D>& hP, int pdg_key, Double_t p_min, Double_t p_max);
    static int FitPtHistogram(unique_ptr<TH1D>& hPt,
                              int pdg_key,
                              Double_t pdg_mass,
                              Double_t pt_min,
                              Double_t pt_max,
                              Double_t& t_output);
    static int FitYHistogram(unique_ptr<TH1D>& hY, int pdg_key, Double_t& y0_output, Double_t& sigma_putput);
    static int FitEtaHistogram(unique_ptr<TH1D>& hEta, int pdg_key, Double_t eta_min, Double_t eta_max);
 
    /** additional functions to check results of fitting and particle equalizing **/
    // static function to fit the obtained sample histograms (to show whether the results are appropriate)
    static void FitSampleHistograms(TString input_histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
    // static function to show the obtained sample histograms with obtained fit (to check whether the results are
    // appropriate)
    static int ShowSampleHistograms(TString input_histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
    // draw histogram with multiplicity of different particle types for all the momentum intervals
    // as well as P, Pt, Y, Eta, Theta, Phi distributions after the equalizing
    // it work both for SIM and DST files listed in the given text file ('input_list_file' var)
    // return 0 if the distributions have been drawn successfully, a negative error code otherwise
    static int ShowResultDistributions(TString input_list_file);
 
    // none of the data members below is written to a ROOT file: they are all rebuilt from the file with the sample
    // histograms at every start, so they are marked as transient ('//!') to keep them out of the generated streamer
    // (the histograms are held by unique pointers, which cannot be streamed at all)
  private:
    static Int_t fSampleTrackCut;   ///< Cut applied to the tracks of the sample (see SetSampleTrackCut)
    Double_t fTargetPerInterval;    //! Target particle number per interval, 0 = maximum of the event (see the setter)
    Bool_t fUseEfficiencyCorrection;   //! Whether the added particle number is corrected by the efficiency
    Double_t fMinEfficiency;           //! Lower limit of the efficiency used for the correction
    Int_t fSampleStatisticsCutRebin;   //! Rebin factor of the histograms used to filter the particles of an event
    Bool_t fIsEfficiencyMapRead;       //! Whether the efficiency map has been found in the sample histogram file
    /// Whether the particles of the event are filtered by the sample statistics (see SetSampleStatisticsCut)
    Bool_t fUseSampleStatisticsCut;   //!
 
    static Double_t fPMin;         ///< Min value of momentum
    static Double_t fPMax;         ///< Max value of momentum
    static Int_t fIntervalCount;   ///< Interval count for momentum to equalize particles at all the intervals
    Double_t fIntervalStep;        //! Interval step
 
    // Whether correct input data for the generator is obtained. If false, the simulation run is stopped: running the
    // equalizing generator is a rare and critical case, so silently skipping it would spoil the whole produced data set
    Bool_t isActive;   //!
    // Particle PDG codes to be equalized
    static vector<Int_t> fParticles;
    // map with the full information on all the particle types
    map<Int_t, unique_ptr<stParticleInfo>> fParticleInfos;   //!
    // Maximum number of a particle type among all particles (divided by P intervals -> vector)
    vector<int> fMaxParticles;   //!
 
    // Get all sample histograms for all the particle types from the 'input_histo_file' file
    // and calculate fit parameters (if necessary) according to the given histograms
    Bool_t ReadSampleHistograms(TString input_histo_file = "$VMCWORKDIR/macro/recotools/particle_hists.root");
 
    // (re)build the rebinned copies of the two dimensional sample histograms used to filter the particles of an
    // event by the sample statistics, according to the current rebin factor
    // return kFALSE if the factor does not divide the number of the bins of some histogram (nothing is changed then)
    Bool_t BuildStatisticsCutHistograms();
 
    // verify that the momentum grid the sample histograms in the given file were produced with is exactly the one
    // this class is configured for. The histograms of the intervals are read by name only, so a shifted grid is not
    // detected by itself and would silently refer the values of an interval to a different momentum range
    static Bool_t CheckMomentumGrid(TFile* hist_file);
 
    // sample a random value along the Y axis of the 2D histogram from its conditional distribution for the given X
    // value, i.e. from the column of the X bin containing 'x'. If that column is empty, the X window is symmetrically
    // widened until some statistics is found, so the value is always taken from the nearest X region having it
    // return kFALSE (and do not change 'y') if the whole histogram is empty
    static Bool_t GetConditionalRandomY(const TH2D* hist, Double_t x, Double_t& y);
 
    // define PDG code of a particle by the given mass
    // among particle types set by 'vecParticlePDG' vector or in the TDatabasePDG if the vector is empty
    // tolerance: relative if tolerance > 0, absolute (in GeV) in case of tolerance <= 0
    // return PDG code or 0 if no corresponding particle was found
    static Int_t GetPDGCodeByMass(Double_t mass, Double_t tolerance = 1.e-6, const vector<int>& vecParticlePDG = {});
 
    // define the particle type of a reconstructed (DST) track. For the simulated data ('mc_tracks' is given) it is the
    // type of the related MC track: the one referenced by the global track itself (BmnGlobalTrack::GetRefIndex()) or,
    // if the reference is not set, the one matched to it ('global_match'). Otherwise, as well as for a simulated track
    // without any matching information, the type is defined by the mass measured by the ToF detectors and by the sign
    // of the charge, as the mass alone does not distinguish a particle from its antiparticle. 'tof_mass' is set to the
    // measured mass in the latter case and to -1 otherwise
    // return the PDG code or 0 if the type cannot be defined
    static Int_t GetGlobalTrackPDG(BmnGlobalTrack* global_track,
                                   BmnTrackMatch* global_match,
                                   TClonesArray* mc_tracks,
                                   Double_t& tof_mass);
 
    ClassDef(BmnParticleEqualizer2D, 1);
};
 
#endif