AMSTrdMCTrack.h

Go to the documentation of this file.

#pragma once
 
#include "TRDConstants.h"
#include "EventGenTRDConsts.h"
 
#include <TGraph2D.h>
#include <TH3D.h>
 
#include <cmath>
#include <vector>
#include <iostream>
 
 
namespace AMSTrdMCTrack
{
    //______________________________________________________________________________________
    enum Coordinate
    {
        X,
        Y,
        Z
    };
    
    
    //______________________________________________________________________________________
    struct Cluster
    {
        double x;
        double y;
        double z;
    };
    
    
    //______________________________________________________________________________________
    struct TRDHitCluster
    {
        double x;
        double y;
        double z;
        int Layer;
        int Ladder;
        int Tube;
        double Edep;
    };
    
    
    //______________________________________________________________________________________
    struct Track
    {
        int trkID;
        int g3PID;
        std::vector<Cluster> clusters;
    };
    
    
    //______________________________________________________________________________________
    struct TRDTrack
    {
        int trkID;
        int g3PID;
        std::vector<TRDHitCluster> clusters;
    };
    
    
    //______________________________________________________________________________________
    template <typename T>
    double GetDistance(const T& P1, const T& P2)
    {
        return std::sqrt((P2.x - P1.x) * (P2.x - P1.x) + 
                         (P2.y - P1.y) * (P2.y - P1.y) + 
                         (P2.z - P1.z) * (P2.z - P1.z));
    }
    
    
    //______________________________________________________________________________________
    bool IsClusterInsideTRD(const TH2F * h2D_TRDMap_XZ,
                            const TH2F * h2D_TRDMap_YZ,
                            const TH2F * h2D_TRDMap_XY,
                            double x,
                            double y,
                            double z)
    {
        static TH2D * histogramXZ = nullptr;
        static TH2D * histogramYZ = nullptr;
 
        if (!histogramXZ || !histogramYZ)
        {
            histogramXZ = new TH2D("histogramXZ", ";X (cm);Z (cm)", 18, EventGenTRDConsts::SIM_TRD_XMIN, EventGenTRDConsts::SIM_TRD_XMAX, 20, EventGenTRDConsts::SIM_TRD_ZMIN, EventGenTRDConsts::SIM_TRD_ZMAX);
            histogramYZ = new TH2D("histogramYZ", ";Y (cm);Z (cm)", 18, EventGenTRDConsts::SIM_TRD_YMIN, EventGenTRDConsts::SIM_TRD_YMAX, 20, EventGenTRDConsts::SIM_TRD_ZMIN, EventGenTRDConsts::SIM_TRD_ZMAX);
            std::cerr<<"\n\n[INFO] Static histogramXZ and histogramYZ were created in AMSTrdMCTrack::IsClusterInsideTRD().\n";
        }
        
        for (int x_bin = 1; x_bin <= 18; ++x_bin) {       
            for (int z_bin = 1; z_bin <= 20; ++z_bin) {   
                histogramXZ->SetBinContent(x_bin, z_bin, 0);
                if (4 < z_bin && z_bin <= 16)             
                    histogramXZ->SetBinContent(x_bin, z_bin, 1);
                if ((x_bin == 1 || x_bin == 18) && z_bin <= 12)
                    histogramXZ->SetBinContent(x_bin, z_bin, 0);
            }
        }
        
        for (int y_bin = 1; y_bin <= 18; ++y_bin) {       
            for (int z_bin = 1; z_bin <= 20; ++z_bin) {   
                histogramYZ->SetBinContent(y_bin, z_bin, 0);
                if (z_bin <= 4 || z_bin > 16)             
                    histogramYZ->SetBinContent(y_bin, z_bin, 1);
                if ((y_bin <= 2 || y_bin >= 17) && z_bin <= 4)
                    histogramYZ->SetBinContent(y_bin, z_bin, 0);
            }
        }
        
        
        int bin_XZ = histogramXZ->FindBin(x, z);
        if (histogramXZ->GetBinContent(bin_XZ) <= 0)
        {
            // return false;   /// this isn't working right now
        }
        
        int bin_YZ = histogramYZ->FindBin(y, z);
        if (histogramYZ->GetBinContent(bin_YZ) <= 0)
        {
            // return false;   /// this isn't working right now
        }
        
        
        
        
        if (!h2D_TRDMap_XZ || !h2D_TRDMap_YZ || !h2D_TRDMap_XY)
        {
            std::cerr << "\n\n[ERROR]: Unable to open XZ, YZ and XY histograms.\n\n";
            throw std::runtime_error("\n\nh2D_TRDMap histograms not found\n\n");
        }
 
        int bin_XY = const_cast<TH2F*>(h2D_TRDMap_XY)->FindBin(x, y);
        if (h2D_TRDMap_XY->GetBinContent(bin_XY) <= 0)
        {
            return false;
        }
 
        bin_XZ = const_cast<TH2F*>(h2D_TRDMap_XZ)->FindBin(x, z);
        if (h2D_TRDMap_XZ->GetBinContent(bin_XZ) <= 0)
        {
            return false;
        }
 
        bin_YZ = const_cast<TH2F*>(h2D_TRDMap_YZ)->FindBin(y, z);
        if (h2D_TRDMap_YZ->GetBinContent(bin_YZ) <= 0)
        {
            return false;
        }
 
        return true;
    }
    
    
    //______________________________________________________________________________________
    void printTRDTrackInfo(const std::vector<TRDTrack>& vTracks)
    {
        for (const auto& track : vTracks)
        {
            printf("\ntrkID:  %5d, g3PID:  %5d", track.trkID, track.g3PID);
            
            for (const auto& cluster : track.clusters)
                printf("\n    (x, y, z, Edep):  %8.3f, %8.3f, %8.3f, %8.4f    |    Estimated (Layer, Ladder):  %5d, %5d.", cluster.x, cluster.y, cluster.z, cluster.Edep, cluster.Layer, cluster.Ladder);
        }
    }
    
    
    //______________________________________________________________________________________
    std::vector<double> GetAllPosValues(const std::vector<Track>& allTracks, Coordinate coord)
    {
        std::vector<double> allValues;
 
        for (const Track& track : allTracks)
        {
            for (const Cluster& cluster : track.clusters)
            {
                switch (coord)
                {
                    case X:
                        allValues.push_back(cluster.x);
                        break;
                    case Y:
                        allValues.push_back(cluster.y);
                        break;
                    case Z:
                        allValues.push_back(cluster.z);
                        break;
                }
            }
        }
 
        return allValues;
    }
    
    
    //______________________________________________________________________________________
    std::vector<double> GetTrkPosValues(const Track& track, Coordinate coord)
    {
        std::vector<double> values;
 
        for (const Cluster& cluster : track.clusters) {
            switch (coord) {
                case X:
                    values.push_back(cluster.x);
                    break;
                case Y:
                    values.push_back(cluster.y);
                    break;
                case Z:
                    values.push_back(cluster.z);
                    break;
            }
        }
 
        return values;
    }
    
    
    TGraph2D * SetTrackClustersToGraph(const AMSTrdMCTrack::Track &track)
    {
        
        TGraph2D *graph = new TGraph2D();
        for (const auto &cluster : track.clusters)
            graph->SetPoint(graph->GetN(), cluster.x, cluster.y, cluster.z);
        
        graph->SetMarkerColor(kRed);  // Set cluster color to red
        graph->SetMarkerStyle(20);    // Set marker style
        return graph;
    }
 
    
    TH3D * SetTrackClustersInTH3D(const AMSTrdMCTrack::Track &track)
    {
        
        TH3D* hist = new TH3D("clustersHist", "Clusters;X;Y;Z", 100, TRD_XYMIN, TRD_XYMAX, 100, TRD_XYMIN, TRD_XYMAX, nZBins, TRD_ZMIN, TRD_ZMAX);
 
        for (const AMSTrdMCTrack::Cluster &cluster : track.clusters)
            hist->Fill(cluster.x, cluster.y, cluster.z);
 
        return hist;
    }
}