EventGeneratorTRD.h

Go to the documentation of this file.

#pragma once
 
#include <TH1F.h>
#include <TROOT.h>
#include <TError.h>
#include <TCanvas.h>
#include <TProfile.h>
#include <TRandom3.h>
 
#include "AMSTrdMCTrack.h"
#include "EventGenTRDConsts.h"
 
#include <cmath>
#include <vector>
#include <iostream>
 
 
namespace EventGeneratorTRD
{
    //______________________________________________________________________________________
    std::vector<AMSTrdMCTrack::Track> GenerateTRDEvent(const TH1F * hFittedTRDtrksinMC, const bool& debugMode = false, const double& SIGMA_TRACK_CLUSTERS = 1)
    {
        TRandom3 * r3 = new TRandom3();
        r3->SetSeed(0);
        
        AMSTrdMCTrack::Cluster P0 = {
            r3->Uniform(EventGenTRDConsts::SIM_TRD_XYMIN, EventGenTRDConsts::SIM_TRD_XYMAX),
            r3->Uniform(EventGenTRDConsts::SIM_TRD_XYMIN, EventGenTRDConsts::SIM_TRD_XYMAX),
            r3->Uniform(EventGenTRDConsts::SIM_TRD_ZMIN,  EventGenTRDConsts::SIM_TRD_ZMAX)
        }; 
        
        // int N = r3->Integer(51);                  /// Generates integer between 0 and 50
        int N = hFittedTRDtrksinMC->GetRandom();     
        if (debugMode)
            N = 5;
        if (N == 1) N = 2; 
        
        
        std::vector<AMSTrdMCTrack::Track> vTracks(N);
        unsigned int nSimClustersInEvent = 0;
        
 
        for(int i = 0; i < N; i++)
        {
            vTracks[i].trkID = i+1;  
            
            if(i == 0) 
                vTracks[i].g3PID = 14;
            else
                vTracks[i].g3PID = r3->Integer(101) - 50;
            
            AMSTrdMCTrack::Cluster Pn = {
                r3->Uniform(EventGenTRDConsts::SIM_TRD_XYMIN, EventGenTRDConsts::SIM_TRD_XYMAX),
                r3->Uniform(EventGenTRDConsts::SIM_TRD_XYMIN, EventGenTRDConsts::SIM_TRD_XYMAX),
                r3->Uniform(EventGenTRDConsts::SIM_TRD_ZMIN,  EventGenTRDConsts::SIM_TRD_ZMAX)
            };
            
            
            int M = r3->Integer(99) + 2;  
 
            double dx = (Pn.x - P0.x) / M;
            double dy = (Pn.y - P0.y) / M;
            double dz = (Pn.z - P0.z) / M;
 
            for(int j = 0; j < M; j++)
            {
                ++nSimClustersInEvent;
                
                vTracks[i].clusters.push_back(
                {
                    P0.x + j * dx + r3->Gaus(0, SIGMA_TRACK_CLUSTERS),
                    P0.y + j * dy + r3->Gaus(0, SIGMA_TRACK_CLUSTERS),
                    P0.z + j * dz + r3->Gaus(0, SIGMA_TRACK_CLUSTERS)
                });
            }
        }    
        
        
        if (debugMode)
            printf("Generated %d tracks in sim event, containing %d sim clusters.", N, nSimClustersInEvent);
       
        
        delete r3;
        return vTracks;
    }
    
    
    //______________________________________________________________________________________
    std::vector<int> EstimateLayerLadderTube(double x, double y, double z, const bool debugMode = false)
    {
        static const TH2F * const TRDModuleMap_XZ = new TH2F ("TRDModuleMap_XZ", ";x (cm);z (cm);", 18, EventGenTRDConsts::SIM_TRD_XMIN, EventGenTRDConsts::SIM_TRD_XMAX, 20, EventGenTRDConsts::SIM_TRD_ZMIN, EventGenTRDConsts::SIM_TRD_ZMAX);  
        static const TH2F * const TRDModuleMap_YZ = new TH2F ("TRDModuleMap_YZ", ";y (cm);z (cm);", 18, EventGenTRDConsts::SIM_TRD_YMIN, EventGenTRDConsts::SIM_TRD_YMAX, 20, EventGenTRDConsts::SIM_TRD_ZMIN, EventGenTRDConsts::SIM_TRD_ZMAX);  
        
        int layerXZ = TRDModuleMap_XZ->GetYaxis()->FindBin(z) - 1;
        int layerYZ = TRDModuleMap_YZ->GetYaxis()->FindBin(z) - 1;
 
        if (layerXZ != layerYZ)
        {
            std::cerr<<"\n[ERROR] z: "<<z<<". layerXZ: "<<layerXZ<<" and layerYZ: "<<layerYZ<<"\n";
            throw std::runtime_error("Layer mismatch between XZ and YZ maps");
        }
        int layer = layerXZ;
 
        int ladder = 0;
        double leftEdge, rightEdge;
        double pos = 0;
 
        if (layer >= 0 && layer <= 3) {
            ladder = TRDModuleMap_YZ->GetXaxis()->FindBin(y) - 3;
            pos = y;
            leftEdge = TRDModuleMap_YZ->GetXaxis()->GetBinLowEdge(ladder + 3);
            rightEdge = TRDModuleMap_YZ->GetXaxis()->GetBinUpEdge(ladder + 3);
        } else if (layer >= 4 && layer <= 11) {
            ladder = TRDModuleMap_XZ->GetXaxis()->FindBin(x) - 2;
            pos = x;
            leftEdge = TRDModuleMap_XZ->GetXaxis()->GetBinLowEdge(ladder + 2);
            rightEdge = TRDModuleMap_XZ->GetXaxis()->GetBinUpEdge(ladder + 2);
        } else if (layer >= 12 && layer <= 15) {
            ladder = TRDModuleMap_XZ->GetXaxis()->FindBin(x) - 1;
            pos = x;
            leftEdge = TRDModuleMap_XZ->GetXaxis()->GetBinLowEdge(ladder + 1);
            rightEdge = TRDModuleMap_XZ->GetXaxis()->GetBinUpEdge(ladder + 1);
        } else if (layer >= 16 && layer <= 19) {
            ladder = TRDModuleMap_YZ->GetXaxis()->FindBin(y) - 1;
            pos = y;
            leftEdge = TRDModuleMap_YZ->GetXaxis()->GetBinLowEdge(ladder + 1);
            rightEdge = TRDModuleMap_YZ->GetXaxis()->GetBinUpEdge(ladder + 1);
        } else {
            throw std::runtime_error("Layer index out of allowed range");
        }
 
        double partitionWidth = (rightEdge - leftEdge) / 16;
        int tube = static_cast<int>((pos - leftEdge) / partitionWidth);
        //tube = (tube + 0) % 16;  /// This adds 0/3/etc/ to tube and then uses modulo 16 to handle overflow. Found to be not needed.
        
        if (tube < 0 || tube >= 16)
        {
            std::cerr<<"\n\nz: "<<z<<". leftEdge: "<<leftEdge<<" and rightEdge: "<<rightEdge<<"\n";
            throw std::runtime_error("Tube calculation out of range");
        }
        
        if (ladder < 0 || ladder > 17)
        {
            if (debugMode)
                std::cerr<<"\n\nSkipping cluster x: "<<x<<", y: "<<y<<", z: "<<z<<".  Layer, Ladder = "<<layer<<", "<<ladder<<"\n";
            // throw std::runtime_error("Ladder calculation out of range");
        }
 
        std::vector<int> result = {layer, ladder, tube};
        return result;
    }
    
    
    //______________________________________________________________________________________
    double GetRandomSampledEdep(const TH1F * hTRDEkinMaxZPrMC, const TH2F * h2D_TRDEkinMaxZPrMC_Edep, const bool debugMode = false)
    {
        if (!hTRDEkinMaxZPrMC || !h2D_TRDEkinMaxZPrMC_Edep)
            throw std::runtime_error("\n\nUnable to open Ekin and Edep histograms\n\n");
 
        double sampledEkin = hTRDEkinMaxZPrMC->GetRandom();
 
        int binX = h2D_TRDEkinMaxZPrMC_Edep->GetXaxis()->FindBin(sampledEkin);
 
        TH1D * hProjectionY = h2D_TRDEkinMaxZPrMC_Edep->ProjectionY("_py", binX, binX, "e"); // "e" option to use bin errors for weighting
 
        double sampledEdep = hProjectionY->GetRandom();
 
        delete hProjectionY;
        
        if (debugMode)
            printf("\nFrom GetRandomSampledEdep(): Sampled Ekin = %8.3f,  binX = %8d,  Sampled Edep = %8.3f.", sampledEkin, binX, sampledEdep);
 
        return sampledEdep;
    }
    
    
    //______________________________________________________________________________________
    double generateNearEdge(double min, double max, TRandom3& r3)
    {
        
        double width = (max - min) * 0.02;        
        if (r3.Uniform() < 0.5)
            return r3.Uniform(min, min + width);  
        else
            return r3.Uniform(max - width, max);  
    }
    
    
    //______________________________________________________________________________________
    AMSTrdMCTrack::Cluster GenerateClusterNearTRDEdge(TRandom3& r3)
    {
        
        AMSTrdMCTrack::Cluster cluster;
        cluster.z = r3.Uniform(EventGenTRDConsts::SIM_TRD_ZMIN + 10, EventGenTRDConsts::SIM_TRD_ZMAX - 10);
        
        if (cluster.z <= 96.6)                                  
        {
            cluster.x = r3.Uniform(-77.0, 77.0);
            cluster.y = generateNearEdge(-72.0, 72.0, r3);
        }
        
        else if (cluster.z > 96.6 && cluster.z <= 119.8)        
        {
            cluster.x = generateNearEdge(-81.0, 81.0, r3);
            cluster.y = r3.Uniform(-84.0, 84.0);
        }
        
        else if (cluster.z > 119.8 && cluster.z <= 130.0)       
        {
            cluster.x = generateNearEdge(-91.0, 91.0, r3);
            cluster.y = r3.Uniform(-92.0, 92.0);
        }
        
        else if (cluster.z > 130.0 && cluster.z <= 143.0)       
        {
            cluster.x = r3.Uniform(-97.0, 97.0);
            cluster.y = generateNearEdge(-92.0, 92.0, r3);
        }
        else
            throw std::runtime_error("generated cluster.z value from r3.Uniform > 143.0!");
        
        return cluster;
    }
    
    
    //______________________________________________________________________________________
    AMSTrdMCTrack::Cluster GenerateClusterAtMinDistance(const AMSTrdMCTrack::Cluster& P0, TRandom3& r3)
    {
        const AMSTrdMCTrack::Cluster fixedPoint = {r3.Gaus(0.0, 20),
                                                   r3.Gaus(0.0, 20),
                                                   r3.Gaus(P0.z, 20)};  //r3.Uniform(70, 160)
        AMSTrdMCTrack::Cluster Pn;
        
        double dirX = fixedPoint.x - P0.x;
        double dirY = fixedPoint.y - P0.y;
        double dirZ = fixedPoint.z - P0.z;
        
        double length = std::sqrt(dirX * dirX + dirY * dirY + dirZ * dirZ);
        dirX /= length;
        dirY /= length;
        dirZ /= length;
        
        // double minDistance = minDistance + r3.Uniform(0.0, 10.0);  // Ensure it is at least minDistance away
        const double minDistance = EventGenTRDConsts::MIN_SIDE_TRACK_LENGTH + r3.Gaus(20, 5);
 
        int N = 1;
        double distance = 0.0;
        while (distance < 2 * minDistance)
        {
            Pn.x = P0.x + N * minDistance * dirX;
            Pn.y = P0.y + N * minDistance * dirY;
            Pn.z = P0.z + N * minDistance * dirZ;
            
            distance = AMSTrdMCTrack::GetDistance(P0, Pn);
            N++;
        }
 
        return Pn;
    }
    
    
    //______________________________________________________________________________________
    AMSTrdMCTrack::Cluster GenerateClusterOppositeP0(const AMSTrdMCTrack::Cluster& P0, TRandom3& r3)
    {
        AMSTrdMCTrack::Cluster Pn;
        Pn.z = 2 * 110 - P0.z;  
        
        if (r3.Uniform() < 0.3333)
        {
            Pn.x = r3.Gaus(-1*P0.x, 20);
            Pn.y = r3.Gaus( 1*P0.y, 20);
        }
        else if (r3.Uniform() >= 0.3333 && r3.Uniform() < 0.6666)
        {
            Pn.x = r3.Gaus( 1*P0.x, 20);
            Pn.y = r3.Gaus(-1*P0.y, 20);
        }
        else
        {
            Pn.x = r3.Gaus(0, 40);
            Pn.y = r3.Gaus(0, 40);
        }
        
        return Pn;
    }
    
    
    
    
    //______________________________________________________________________________________
    std::vector<AMSTrdMCTrack::TRDTrack> GenerateTRDEvent(const TH1F * hTRDEkinMaxZPrMC,
                                                          const TH2F * h2D_TRDEkinMaxZPrMC_Edep,
                                                          const TH2F * h2D_TRDMap_XZ,
                                                          const TH2F * h2D_TRDMap_YZ,
                                                          const TH2F * h2D_TRDMap_XY,
                                                          const int nFakeTracksPerEvent = 1,
                                                          const bool debugMode = false,
                                                          TH1D * hQA_nAttemptsPerFakeEvent = NULL, 
                                                          TH2D * hQA_FakeAngleVSLength = NULL,
                                                          double SIGMA_TRACK_CLUSTERS = 0.001)
    {
        if (!hTRDEkinMaxZPrMC || !h2D_TRDEkinMaxZPrMC_Edep || !h2D_TRDMap_XZ || !h2D_TRDMap_YZ || !h2D_TRDMap_XY)
            throw std::runtime_error("std::vector<AMSTrdMCTrack::TRDTrack> GenerateTRDEvent(): Histograms not be loaded.");
        
        TRandom3 r3;
        r3.SetSeed(0);
 
        std::vector<AMSTrdMCTrack::TRDTrack> vTracks;
        vTracks.reserve(nFakeTracksPerEvent);
 
        unsigned int nSimClustersInEvent = 0;
        int attempts = 0;
 
        AMSTrdMCTrack::Cluster P0 = GenerateClusterNearTRDEdge(r3);
        if (debugMode)
            printf("\nP0 = (%8.3f, %8.3f, %8.3f) cm", P0.x, P0.y, P0.z);
        
 
        while (int(vTracks.size()) < nFakeTracksPerEvent && attempts < EventGenTRDConsts::MAX_ATTEMPTS_PER_EVENT * nFakeTracksPerEvent)
        {
            attempts++;
            if (debugMode)
                printf("\nnAttepts to find Pn = %d", attempts);
 
            
            // AMSTrdMCTrack::Cluster Pn = GenerateClusterNearTRDEdge(r3);     /// --> Is working okay, but perhaps with some issues.
            // AMSTrdMCTrack::Cluster Pn = GenerateClusterOppositeP0(P0, r3);  /// --> Is not working efficiently. Too make events with no clusters.
            AMSTrdMCTrack::Cluster Pn = GenerateClusterAtMinDistance(P0, r3);
 
            double dx = Pn.x - P0.x;
            double dy = Pn.y - P0.y;
            double dz = Pn.z - P0.z;
            double magnitude = AMSTrdMCTrack::GetDistance(P0, Pn);
            double cosTheta = fabs(dz / magnitude);
 
            double radians = acos(cosTheta);  // acos() returns the angle in radians
            double degrees = radians * (180.0 / M_PI);
            if (debugMode)
            {   
                printf("\nCheck Pn = (%8.3f, %8.3f, %8.3f) cm", Pn.x, Pn.y, Pn.z);
                printf("\nAngle = %5.2f radians or %6.1f degrees.\nLength = %5.1f cm", radians, degrees, magnitude);
            }
            
            if (radians <= EventGenTRDConsts::MIN_Z_ANGLE_THETA)
                continue; 
            if (radians > EventGenTRDConsts::MAX_Z_ANGLE_THETA)
                continue;
            if (magnitude < EventGenTRDConsts::MIN_SIDE_TRACK_LENGTH)
                continue;
            
            if (debugMode)
                printf("\nInitial good line found! Set Pn = (%8.3f, %8.3f, %8.3f) cm", Pn.x, Pn.y, Pn.z);
            
            AMSTrdMCTrack::TRDTrack track;
            track.trkID = vTracks.size() + 1;
            track.g3PID = (track.trkID == 1) ? 14 : r3.Integer(101) - 50;
 
            int M = r3.Gaus(20) + 20;  
            
            
            for (int j = 0; j < M; j++)
            {
                double x = P0.x + j * dx / M + r3.Gaus(0, SIGMA_TRACK_CLUSTERS);
                double y = P0.y + j * dy / M + r3.Gaus(0, SIGMA_TRACK_CLUSTERS);
                double z = P0.z + j * dz / M + r3.Gaus(0, SIGMA_TRACK_CLUSTERS);
 
                if (x < EventGenTRDConsts::SIM_TRD_XMIN || x > EventGenTRDConsts::SIM_TRD_XMAX ||
                    y < EventGenTRDConsts::SIM_TRD_YMIN || y > EventGenTRDConsts::SIM_TRD_YMAX ||
                    z < EventGenTRDConsts::SIM_TRD_ZMIN || z > EventGenTRDConsts::SIM_TRD_ZMAX)
                    continue;  
                
                bool isInside = AMSTrdMCTrack::IsClusterInsideTRD(h2D_TRDMap_XZ, h2D_TRDMap_YZ, h2D_TRDMap_XY, x, y, z);
                if (isInside == false)
                    continue;  
                
                std::vector<int> estLLT = EstimateLayerLadderTube(x, y, z, debugMode);
                if (estLLT[1] < 0 || estLLT[1] > 17)
                    continue;  
                
                double Edep = GetRandomSampledEdep(hTRDEkinMaxZPrMC, h2D_TRDEkinMaxZPrMC_Edep);
                
                track.clusters.push_back({x, y, z, estLLT[0], estLLT[1], estLLT[2], Edep});  
                ++nSimClustersInEvent;
            }
 
            if (!track.clusters.empty())
                vTracks.push_back(track);
            
            if (track.clusters.size() >= 2)
            {
                double distance = AMSTrdMCTrack::GetDistance(track.clusters.front(), track.clusters.back());
                
                if (debugMode)
                    printf("\nActual length of fake track: %8.2f cm.", distance);
                
                if (hQA_FakeAngleVSLength)
                    hQA_FakeAngleVSLength->Fill(degrees, distance);
            }
        }
        
        
        if (debugMode)
        {
            printf("\nGenerated %d fake track(s) in this event, containing %d synthetic/fake clusters.", int(vTracks.size()), nSimClustersInEvent);
            AMSTrdMCTrack::printTRDTrackInfo(vTracks);
            printf("\n");
        }
        
        if (hQA_nAttemptsPerFakeEvent)
            hQA_nAttemptsPerFakeEvent->Fill(attempts);
        
        static int nCalls = 0;
        nCalls++;
        if (vTracks.empty())
            printf("\n\n[WARNING]: For nCall = %d, vTracks generated by GenerateTRDEvent() is empty. | nAttempts = %d\n", nCalls, attempts);
 
        return vTracks;
    }
    
    
    
    
    //______________________________________________________________________________________
    TH2F * GetModuleHistogramFakeTracks(const std::vector<AMSTrdMCTrack::TRDTrack>& fake_tracks, const bool debugMode = false, TH2F * hTRDEventMap_XZ = NULL, TH2F * hTRDEventMap_YZ = NULL)
    {
        if (fake_tracks.empty())
        {
            std::cerr<<"\n[WARNING]: In GetModuleHistogramFakeTracks(), provided std::vector<AMSTrdMCTrack::TRDTrack> is empty.\n\n";
            return nullptr;
        }
        
        TH2F * hTRDModuleMap = new TH2F("hTRDModuleMap_FakeTracks", ";Ladder;Layer;Edep", 18, 0, 18, 20, 0, 20);
 
        // Fill the histogram using the tracks data
        for (const auto& track : fake_tracks)
        {
            for (const auto& cluster : track.clusters)
            {
                int ladder_hist = cluster.Ladder;
                int layer      = cluster.Layer;
                
                // Determine ladder value based on layer
                if (layer >= 0 && layer <= 3)
                    ladder_hist = ladder_hist + 2;
                else if (layer >= 4 && layer <= 11)
                    ladder_hist = ladder_hist + 1;
                else if (layer >= 12 && layer <= 15)
                    ladder_hist = ladder_hist + 0;
                else if (layer >= 16 && layer <= 19)
                    ladder_hist = ladder_hist + 0;
                else
                    throw std::runtime_error("Layer index out of allowed range 0-19");
                
                hTRDModuleMap->Fill(ladder_hist, layer, cluster.Edep);
                
                if (hTRDEventMap_XZ) hTRDEventMap_XZ->Fill(cluster.x, cluster.z, cluster.Edep);
                if (hTRDEventMap_YZ) hTRDEventMap_YZ->Fill(cluster.y, cluster.z, cluster.Edep);
            }
        }
        
        return hTRDModuleMap;
    }
}