BmnParticleEqualizer1D.h

Go to the documentation of this file.

// -------------------------------------------------------------------------
// -----                BmnParticleEqualizer1D header file               -----
// -----              Created 01/05/2024 by K. Gertsenberger           -----
// -------------------------------------------------------------------------
 
/** BmnParticleEqualizer1D.h
 ** @author Konstantin Gertsenberger <gertsen@jinr.ru>
 **
 ** The BmnParticleEqualizer1D 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 number of the particles of every type is equalized inside every momentum interval independently, up to the
 ** maximum number among the types found in the interval (the intervals themselves are not equalized with each other).
 **
 ** The momentum, the pseudorapidity and the azimuth angle of a new particle are sampled from the ONE dimensional
 ** sample histograms of its momentum interval, i.e. independently of each other: the correlations between the values
 ** inside the interval are not kept, use BmnParticleEqualizer2D if they matter. The angles are nevertheless distributed
 ** as in the sample, i.e. inside the detector acceptance, and not uniformly.
 **
 ** A file with the sample histograms produced by BmnParticleEqualizer2D can be read by this class as well (it is a
 ** superset), while a file produced by this class does not contain the two dimensional histograms needed by
 ** BmnParticleEqualizer2D.
 */
 
#ifndef BMNPARTICLEEQUALIZER1D_H
#define BMNPARTICLEEQUALIZER1D_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 <map>
#include <memory>
#include <vector>
using namespace std;
 
class BmnParticleEqualizer1D : public FairGenerator
{
  public:
    /** Constructor with input histograms **/
    BmnParticleEqualizer1D(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
    **/
    BmnParticleEqualizer1D(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 ~BmnParticleEqualizer1D() {}
 
    /** 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
 
        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)
 
        /// Way the momentum 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,   ///< from the P histogram of the interval itself (there is enough sample statistics)
            kFitFunction     ///< 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)
 
        /// Index of the momentum interval whose Eta histogram is used to sample the pseudorapidity for a given interval
        /// (the interval itself when it has enough statistics, otherwise the nearest one that has), -1 if there is none
        vector<Int_t> fEtaHistIndex;
        /// The same for the azimuth angle: index of the momentum interval whose Phi histogram is used
        vector<Int_t> fPhiHistIndex;
 
        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; }
    /** 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 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
    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");
 
    // 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);
 
    // for every momentum interval define the index of the interval histogram to be used to sample a value: the
    // interval itself if it has enough statistics, otherwise the nearest interval that has it (so the value stays
    // tied to a close momentum region instead of being taken from the whole momentum range)
    // the element is set to -1 if no interval has statistics at all
    static void DefineNearestFilledIntervals(const vector<unique_ptr<TH1D>>& hist_vector, vector<Int_t>& index_vector);
 
    // 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(BmnParticleEqualizer1D, 1);
};
 
#endif