tr_he_analysis.cxx

// Copyright CERN and copyright holders of ALICE O2. This software is
// distributed under the terms of the GNU General Public License v3 (GPL
// Version 3), copied verbatim in the file "COPYING".
//
// See http://alice-o2.web.cern.ch/license for full licensing information.
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
//

#include "PWGLF/DataModel/LFNucleiTables.h"
#include <TLorentzVector.h>
#include <TF1.h>
#include "ReconstructionDataFormats/Track.h"
#include "Framework/runDataProcessing.h"
#include "Framework/AnalysisTask.h"
#include "Framework/AnalysisDataModel.h"
#include "Framework/ASoAHelpers.h"
#include "Framework/HistogramRegistry.h"
#include "Common/Core/trackUtilities.h"
#include "Common/DataModel/Multiplicity.h"
#include "Common/DataModel/PIDResponse.h"
#include "Common/DataModel/EventSelection.h"
#include "Common/DataModel/Centrality.h"
#include "Common/DataModel/TrackSelectionTables.h"
#include "Common/Core/trackUtilities.h"
#include "Common/CCDB/EventSelectionParams.h"
#include "PWGLF/DataModel/LFParticleIdentification.h"
#include "ReconstructionDataFormats/PID.h"
#include "Common/Core/PID/TPCPIDResponse.h"

namespace o2::aod
{
namespace h3_data
{
DECLARE_SOA_COLUMN(tPt, tpt, double);
DECLARE_SOA_COLUMN(tY, ty, double);
DECLARE_SOA_COLUMN(tEta, teta, double);
DECLARE_SOA_COLUMN(tPhi, tphi, double);
DECLARE_SOA_COLUMN(tPx, tpx, double);
DECLARE_SOA_COLUMN(tPy, tpy, double);
DECLARE_SOA_COLUMN(tPz, tpz, double);
DECLARE_SOA_COLUMN(tE, te, double);
//DECLARE_SOA_COLUMN(tKink, tkink, int);
DECLARE_SOA_COLUMN(tCharge, tcharge, int);

DECLARE_SOA_COLUMN(tP, tp, double);
DECLARE_SOA_COLUMN(tH3DEdx, th3dedx, double);
DECLARE_SOA_COLUMN(tnSigmaTPC, tnsigmatpc, double);
//DECLARE_SOA_COLUMN(tH3Sigma, th3sigma, double);
DECLARE_SOA_COLUMN(tTOFSignalH3, ttofsignalh3, double);

DECLARE_SOA_COLUMN(tDcaXY, tdcaxy, double);
DECLARE_SOA_COLUMN(tDcaZ, tdcaz, double);
DECLARE_SOA_COLUMN(tSigmaYX, tsigmayx, double);
DECLARE_SOA_COLUMN(tSigmaXYZ, tsigmaxyz, double);
DECLARE_SOA_COLUMN(tSigmaZ, tsigmaz, double);

DECLARE_SOA_COLUMN(tnTPCcluster, tntpccluster, int);
DECLARE_SOA_COLUMN(tnITScluster, tnitscluster, int);
}    
DECLARE_SOA_TABLE(H3Data, "H3Data", "h3_data",
                h3_data::tPt, h3_data::tY, h3_data::tEta, h3_data::tPhi, h3_data::tPx, h3_data::tPy, h3_data::tPz,
                h3_data::tE, h3_data::tCharge, h3_data::tP,
                h3_data::tH3DEdx, h3_data::tnSigmaTPC, h3_data::tTOFSignalH3, h3_data::tDcaXY, h3_data::tDcaZ, h3_data::tSigmaYX,
                h3_data::tSigmaXYZ, h3_data::tSigmaZ, h3_data::tnTPCcluster, h3_data::tnITScluster
                );
}
namespace{
    const int nBetheParams = 6; 
    const int nParticles = 2; 
    static const std::vector<std::string> particleNames{"triton", "helion"};
    static const std::vector<int> particlePdgCodes{o2::constants::physics::kTriton, o2::constants::physics::kHelium3};
    static const std::vector<double> particleMasses{o2::constants::physics::MassTriton, o2::constants::physics::MassHelium3};  
    static const std::vector<int> particleCharge{1, 2};
    static const std::vector<std::string> betheBlochParNames{"p0", "p1", "p2", "p3", "p4", "resolution"};    
    constexpr double betheBlochDefault[nParticles][nBetheParams]{
    {0.313129,   181.664226, 2779397163087.684082, 2.130773,    29.609643, 0.09},   // triton
    {70.584685,  3.196364,   0.133878,             2.731736,    1.675617,  0.09}}; //Helion 
}
using namespace o2;
using namespace o2::framework;
using namespace o2::framework::expressions;
using TracksFull = soa::Join<aod::TracksIU, aod::TracksExtra, aod::TracksCovIU, o2::aod::TracksDCA, aod::pidTPCPi, aod::pidTOFmass, aod::pidTOFbeta, aod::pidTPCLfFullTr>;

class Particle {
public:
    TString name;
    int pdgCode;
    double mass;
    int charge;
    double resolution;
    std::vector<double> betheParams;

    Particle(const std::string& name_, int pdgCode_, double mass_, int charge_,
             const o2::framework::LabeledArray<double>& bethe) {
        name = TString(name_);
        pdgCode = pdgCode_;
        mass = mass_;
        charge = charge_;

        resolution = bethe.get(name, "resolution"); // Access the "resolution" parameter

        // Access other Bethe-Bloch parameters
        betheParams.clear();
        for (int i = 0; i < 5; ++i) {
            betheParams.push_back(bethe.get(name, i));
        }
    }
};

struct tr_he_analysis {
    Produces<o2::aod::H3Data> h3_data;
    //Produces<o2::aod::He3_data> He3_table;
    //OutputObj<o2::framework::TableBuilder<H3_data>> H3_table{"H3_table"};
    //OutputObj<o2::framework::TableBuilder<He3_data>> He3_table{"He3_table"};
    HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject};
    std::vector<Particle> particles;
    o2::framework::LabeledArray<double> bethe;
    Configurable<bool> enableTr{"enableTr", true, "Flag to enable triton analysis."};
    Configurable<bool> enableHe{"enableHe", true, "Flag to enable helium-3 analysis."};

    ConfigurableAxis binsPt{"binsPt", {VARIABLE_WIDTH, 0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.425, 0.45, 0.475, 0.5, 0.5125, 0.525, 0.5375, 0.55, 0.5625, 0.575, 0.5875, 0.6, 0.6125, 0.625, 0.6375, 0.65, 0.6625, 0.675, 0.6875, 0.7, 0.7125, 0.725, 0.7375, 0.75, 0.7625, 0.775, 0.7875, 0.8, 0.8125, 0.825, 0.8375, 0.85, 0.8625, 0.875, 0.8875, 0.9, 0.9125, 0.925, 0.9375, 0.95, 0.9625, 0.975, 0.9875, 1.0, 1.0125, 1.025, 1.0375, 1.05, 1.0625, 1.075, 1.0875, 1.1, 1.1125, 1.125, 1.1375, 1.15, 1.1625, 1.175, 1.1875, 1.2, 1.2125, 1.225, 1.2375, 1.25, 1.2625, 1.275, 1.2875, 1.3, 1.3125, 1.325, 1.3375, 1.35, 1.3625, 1.375, 1.3875, 1.4, 1.4125, 1.425, 1.4375, 1.45, 1.4625, 1.475, 1.4875, 1.5, 1.5125, 1.525, 1.5375, 1.55, 1.5625, 1.575, 1.5875, 1.6, 1.6125, 1.625, 1.6375, 1.65, 1.6625, 1.675, 1.6875, 1.7, 1.7125, 1.725, 1.7375, 1.75, 1.7625, 1.775, 1.7875, 1.8, 1.8125, 1.825, 1.8375, 1.85, 1.8625, 1.875, 1.8875, 1.9, 1.9125, 1.925, 1.9375, 1.95, 1.9625, 1.975, 1.9875, 2.0, 2.0625, 2.125, 2.1875, 2.25, 2.3125, 2.375, 2.4375, 2.5, 2.625, 2.75, 2.875, 3.0, 3.25, 3.5, 3.75, 4.0, 4.5, 5.0, 6.0, 7.0, 8.0}, ""};
    ConfigurableAxis binsPtHe{"binsPtTr", {VARIABLE_WIDTH, 1.0, 1.25, 1.50, 1.75, 2.0, 2.25, 2.50, 2.75, 3.0, 3.25, 3.50, 3.75, 4.0, 4.50, 5.0, 6.0, 7.0, 8.0}, ""};
    ConfigurableAxis binsPtZHe{"binsPtZTr", {VARIABLE_WIDTH, 0.5, 0.625, 0.75, 0.875, 1.0, 1.125, 1.25, 1.375, 1.5, 1.625, 1.75, 1.875, 2.0, 2.25, 2.5, 3.0, 3.5, 4.0}, ""};
    ConfigurableAxis binsPtTr{"binsPtHe", {VARIABLE_WIDTH, 1.0, 1.25, 1.50, 1.75, 2.0, 2.25, 2.50, 2.75, 3.0, 3.25, 3.50, 3.75, 4.0, 4.50, 5.0, 6.0, 7.0, 8.0}, ""};
    ConfigurableAxis binsPtZTr{"binsPtZHe", {VARIABLE_WIDTH, 0.5, 0.625, 0.75, 0.875, 1.0, 1.125, 1.25, 1.375, 1.5, 1.625, 1.75, 1.875, 2.0, 2.25, 2.5, 3.0, 3.5, 4.0}, ""};
    ConfigurableAxis binsdEdx{"binsdEdx", {600, 0.f, 3000.f}, ""};
    ConfigurableAxis binsBeta{"binsBeta", {120, 0.0, 1.2}, ""};
    ConfigurableAxis binsDCA{"binsDCA", {400, -1.f, 1.f}, ""};
    ConfigurableAxis binsSigmaTPC{"binsSigmaTPC", {1000, -100, 100}, ""};
    ConfigurableAxis binsSigmaTOF{"binsSigmaTOF", {1000, -100, 100}, ""};
    ConfigurableAxis binsMassTr{"binsMassTr", {250, -2.5, 2.5f}, ""};
    ConfigurableAxis binsMassHe{"binsMassHe", {300, -3., 3.f}, ""};
    // Set the event selection cuts
    struct : ConfigurableGroup {
        Configurable<bool> useSel8{"useSel8", true, "Use Sel8 for run3 Event Selection"};
        Configurable<bool> TVXtrigger{"TVXtrigger", false, "Use TVX for Event Selection (default w/ Sel8)"};
        Configurable<bool> removeTFBorder{"removeTFBorder", false, "Remove TimeFrame border (default w/ Sel8)"};
        Configurable<bool> removeITSROFBorder{"removeITSROFBorder", false, "Remove ITS Read-Out Frame border (default w/ Sel8)"};
    } evselOptions;

    Configurable<bool> cfgBetheParam{"cfgBetheParam", true, "Using own or built in bethe parametrization"};
    // Set the multiplity event limits
    Configurable<float> cfgLowMultCut{"cfgLowMultCut", 0.0f, "Accepted multiplicity percentage lower limit"};
    Configurable<float> cfgHighMultCut{"cfgHighMultCut", 100.0f, "Accepted multiplicity percentage higher limit"};

    // Set the z-vertex event cut limits
    Configurable<float> cfgHighCutVertex{"cfgHighCutVertex", 10.0f, "Accepted z-vertex upper limit"};
    Configurable<float> cfgLowCutVertex{"cfgLowCutVertex", -10.0f, "Accepted z-vertex lower limit"};

    // Set the quality cuts for tracks
    Configurable<bool> rejectFakeTracks{"rejectFakeTracks", false, "Flag to reject ITS-TPC fake tracks (for MC)"};
    Configurable<float> cfgCutITSClusters{"cfgCutITSClusters", -1.f, "Minimum number of ITS clusters"};
    Configurable<float> cfgCutTPCXRows{"cfgCutTPCXRows", -1.f, "Minimum number of crossed TPC rows"};
    Configurable<float> cfgCutTPCClusters{"cfgCutTPCClusters", 40.f, "Minimum number of found TPC clusters"};
    Configurable<int> nITSLayer{"nITSLayer", 0, "ITS Layer (0-6)"};

    // Set the kinematic and PID cuts for tracks
    struct : ConfigurableGroup {
        Configurable<float> pCut{"pCut", 0.3f, "Value of the p selection for spectra (default 0.3)"};
        Configurable<float> etaCut{"etaCut", 0.8f, "Value of the eta selection for spectra (default 0.8)"};
        Configurable<float> yLowCut{"yLowCut", -1.0f, "Value of the low rapidity selection for spectra (default -1.0)"};
        Configurable<float> yHighCut{"yHighCut", 1.0f, "Value of the high rapidity selection for spectra (default 1.0)"};
    } kinemOptions;

    struct : ConfigurableGroup {
        Configurable<float> nsigmaTPCTr{"nsigmaTPCTr", 5.f, "Value of the Nsigma TPC cut for tritons"};
        Configurable<float> nsigmaTPCHe{"nsigmaTPCHe", 5.f, "Value of the Nsigma TPC cut for helium-3"};
    } nsigmaTPCvar;

    void init(o2::framework::InitContext&){
        const AxisSpec pAxis{binsPt, "#it{p} (GeV/#it{c})"};
        const AxisSpec ptAxis{binsPt, "#it{p}_{T} (GeV/#it{c})"};
        const AxisSpec ptHeAxis{binsPtHe, "#it{p}_{T} (GeV/#it{c})"};
        const AxisSpec ptZHeAxis{binsPtZHe, "#it{p}_{T}/z (GeV/#it{c})"};
        const AxisSpec ptTrAxis{binsPtTr, "#it{p}_{T} (GeV/#it{c})"};
        const AxisSpec ptZTrAxis{binsPtZTr, "#it{p}_{T}/z (GeV/#it{c})"};
        const AxisSpec dedxAxis{binsdEdx, "d#it{E}/d#it{x} A.U."};
        const AxisSpec betaAxis{binsBeta, "TOF #beta"};
        const AxisSpec dcaxyAxis{binsDCA, "DCAxy (cm)"};
        const AxisSpec dcazAxis{binsDCA, "DCAz (cm)"};
        const AxisSpec massTrAxis{binsMassTr, ""};
        const AxisSpec massHeAxis{binsMassHe, ""};
        const AxisSpec SigmaTPCAxis{binsSigmaTPC, ""};
        const AxisSpec SigmaTOFAxis{binsSigmaTOF, ""};
    
        histos.add<TH1>("histogram/h1pT", "Track #it{p}_{T}; #it{p}_{T} (GeV/#it{c}); counts", HistType::kTH1F, {{500, 0., 10.}});
        histos.add<TH1>("histogram/h1p", "Track momentum; p (GeV/#it{c}); counts", HistType::kTH1F, {{500, 0., 10.}});
        histos.add<TH2>("histogram/h2TPCsignVsTPCmomentum", "TPC <-dE/dX> vs #it{p}/Z; Signed #it{p} (GeV/#it{c}); TPC <-dE/dx> (a.u.)", HistType::kTH2F, {{400, -8.f, 8.f}, {dedxAxis}});
        histos.add<TH2>("histogram/h2TOFbetaVsP", "TOF #beta vs #it{p}/Z; Signed #it{p} (GeV/#it{c}); TOF #beta", HistType::kTH2F, {{250, -5.f, 5.f}, {betaAxis}});
        histos.add<TH1>("event/eventSelection", "eventSelection", HistType::kTH1D, {{7, -0.5, 6.5}});
        auto h = histos.get<TH1>(HIST("event/eventSelection"));
        h->GetXaxis()->SetBinLabel(1, "Total");
        h->GetXaxis()->SetBinLabel(2, "TVX trigger cut");
        h->GetXaxis()->SetBinLabel(3, "TF border cut");
        h->GetXaxis()->SetBinLabel(4, "ITS ROF cut");
        h->GetXaxis()->SetBinLabel(5, "TVX + TF + ITS ROF");
        h->GetXaxis()->SetBinLabel(6, "Sel8 cut");
        h->GetXaxis()->SetBinLabel(7, "Z-vert Cut");
        double flatBetheBloch[nParticles * nBetheParams];
        for (int i = 0; i < nParticles; ++i) {
            for (int j = 0; j < nBetheParams; ++j) {
                flatBetheBloch[i * nBetheParams + j] = betheBlochDefault[i][j];
            }
        }

        o2::framework::LabeledArray<double> bethe(
        flatBetheBloch, nParticles, nBetheParams,
        particleNames, betheBlochParNames
        );
    }
    void fillHistoTable(TracksFull const& tracks, soa::Join<aod::Collisions, aod::EvSels>::iterator const& event){

        std::vector<Particle> particles;
        for (size_t i = 0; i < particleNames.size(); ++i) {
            particles.emplace_back(
            particleNames[i], particlePdgCodes[i], particleMasses[i], particleCharge[i], bethe);
        }
        Particle triton(particleNames[0], particlePdgCodes[0], particleMasses[0], particleCharge[0], bethe);
        Particle helion(particleNames[1], particlePdgCodes[1], particleMasses[1], particleCharge[1], bethe);

        std::cout << "Particle 1: " << triton.name << ", PDG Code: " << triton.pdgCode
                  << ", Mass: " << triton.mass << ", Charge: " << triton.charge
                  << ", Resolution: " << triton.resolution << std::endl;
        std::cout << "Particle 2: " << helion.name << ", PDG Code: " << helion.pdgCode
                  << ", Mass: " << helion.mass << ", Charge: " << helion.charge
              << ", Resolution: " << helion.resolution << std::endl;

        histos.fill(HIST("event/eventSelection"), 0);
        if ((event.selection_bit(aod::evsel::kNoITSROFrameBorder)) && (event.selection_bit(aod::evsel::kNoTimeFrameBorder)) && (event.selection_bit(aod::evsel::kIsTriggerTVX))) {
            histos.fill(HIST("event/eventSelection"), 4);
        }
        if (evselOptions.useSel8 && !event.sel8())
            return;
        histos.fill(HIST("event/eventSelection"), 5);
        if (event.posZ() < cfgLowCutVertex || event.posZ() > cfgHighCutVertex)
            return;
        histos.fill(HIST("event/eventSelection"), 6);
        if (cfgBetheParam)
        {
            for (auto& track : tracks) {
                if (TMath::Abs(track.tpcInnerParam()) < kinemOptions.pCut)
                    continue;
                if (TMath::Abs(track.eta()) < kinemOptions.etaCut)
                    continue;
                for (size_t i = 0; i < particles.size(); i++){
                    if (track.tpcNClsFound() < cfgCutTPCClusters)
                        continue;
                    if (track.itsNCls() < cfgCutITSClusters)
                        continue;
                    if (track.tpcNClsCrossedRows() < cfgCutTPCXRows)
                        continue;
                    histos.fill(HIST("tracks/h1pT"), track.pt());
                    histos.fill(HIST("tracks/h1p"), track.p());
                    histos.fill(HIST("tracks/h2TPCsignVsTPCmomentum"), track.tpcInnerParam() / (1.f * track.sign()), track.tpcSignal());
                    histos.fill(HIST("tracks/h2TOFbetaVsP"), track.p() / (1.f * track.sign()), track.beta());
                    if (std::abs(getTPCnSigma(track, particles.at(i))) < nsigmaTPCvar.nsigmaTPCTr){
                        double tPt = track.pt();
                        double tY = track.y();
                        double tEta = track.eta();
                        double tPhi = track.phi();
                        double tPx = track.px();
                        double tPy = track.py();
                        double tPz = track.pz();
                        double tE = track.energy(track.mass());
                        //int tKink = track.kink();
                        int tCharge = track.sign();

                        double tP = track.p();
                        double tH3DEdx = track.tpcSignal();
                        double tnSigmaTPC = track.tpcNSigmaTr();
                        double tTOFSignalH3 = track.mass() * track.mass();

                        double tDcaXY = track.dcaXY();
                        double tDcaZ = track.dcaZ();
                        double tSigmaYX = track.sigmaY();
                        double tSigmaXYZ = track.sigmaSnp();
                        double tSigmaZ = track.sigmaZ();

                        int tnTPCcluster = track.tpcNClsFound(); 
			            int tnITScluster = track.itsNCls();

                        h3_data(tPt, tY, tEta, tPhi, tPx, tPy, tPz, tE, tCharge, tP, tH3DEdx, tnSigmaTPC, tTOFSignalH3, tDcaXY, tDcaZ, tSigmaXYZ, tSigmaYX, tSigmaZ, tnTPCcluster, tnITScluster);
                    }
                }
            }
        }    
        else if (!cfgBetheParam)
        {
            for (auto& track : tracks) {
                if (TMath::Abs(track.tpcInnerParam()) < kinemOptions.pCut)
                    continue;
                if (TMath::Abs(track.eta()) < kinemOptions.etaCut)
                    continue;
                if (track.tpcNClsFound() < cfgCutTPCClusters)
                    continue;
                if (track.itsNCls() < cfgCutITSClusters)
                    continue;
                if (track.tpcNClsCrossedRows() < cfgCutTPCXRows)
                    continue;
                histos.fill(HIST("tracks/h1pT"), track.pt());
                histos.fill(HIST("tracks/h1p"), track.p());
                histos.fill(HIST("tracks/h2TPCsignVsTPCmomentum"), track.tpcInnerParam() / (1.f * track.sign()), track.tpcSignal());
                histos.fill(HIST("tracks/h2TOFbetaVsP"), track.p() / (1.f * track.sign()), track.beta());
                if (std::abs(track.tpcNSigmaTr()) < nsigmaTPCvar.nsigmaTPCTr){
                    double tPt = track.pt();
                    double tY = track.y();
                    double tEta = track.eta();
                    double tPhi = track.phi();
                    double tPx = track.px();
                    double tPy = track.py();
                    double tPz = track.pz();
                    double tE = track.energy(track.mass());
                    //int tKink = track.kink();
                    int tCharge = track.sign();

                    double tP = track.p();
                    double tH3DEdx = track.tpcSignal();
                    double tnSigmaTPC = track.tpcNSigmaTr();
                    double tTOFSignalH3 = track.mass() * track.mass();

                    double tDcaXY = track.dcaXY();
                    double tDcaZ = track.dcaZ();
                    double tSigmaYX = track.sigmaY();
                    double tSigmaXYZ = track.sigmaSnp();
                    double tSigmaZ = track.sigmaZ();

                    int tnTPCcluster = track.tpcNClsFound(); 
			        int tnITScluster = track.itsNCls();

                    h3_data(tPt, tY, tEta, tPhi, tPx, tPy, tPz, tE, tCharge, tP, tH3DEdx, tnSigmaTPC, tTOFSignalH3, tDcaXY, tDcaZ, tSigmaXYZ, tSigmaYX, tSigmaZ, tnTPCcluster, tnITScluster);
                }
            }
        }    
    }

    template <class T>
    float getTPCnSigma(T const& track, Particle const& particle) {
	const float rigidity = track.tpcInnerParam();
    if (!track.hasTPC()) return -999; 
      
	double expBethe{tpc::BetheBlochAleph(static_cast<double>(particle.charge*rigidity  / particle.mass), particle.betheParams[0], particle.betheParams[1], particle.betheParams[2], particle.betheParams[3], particle.betheParams[4])}; 
	double expSigma{expBethe * particle.resolution};
	float sigmaTPC = static_cast<float>((track.tpcSignal() - expBethe) / expSigma);
	return sigmaTPC;     		
    }  

};  

WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
{
  return WorkflowSpec{
    adaptAnalysisTask<tr_he_analysis>(cfgc),
    };
}