Merge pull request #120 from rest-for-physics/lobis-th2d

New generator TRestGeant4ParticleSourceCosmics

inc/TRestGeant4ParticleSourceCosmics.h

@@ -0,0 +1,33 @@
+
+#ifndef REST_TRESTGEANT4PARTICLESOURCECOSMICS_H
+#define REST_TRESTGEANT4PARTICLESOURCECOSMICS_H
+
+#include <TH2D.h>
+#include <TRandom3.h>
+#include <TRestGeant4ParticleSource.h>
+
+class TRestGeant4ParticleSourceCosmics : public TRestGeant4ParticleSource {
+   private:
+    std::set<std::string> fParticleNames;
+    std::string fFilename;
+    std::map<std::string, double> fParticleWeights;
+
+    std::map<std::string, TH2D*> fHistograms;
+    static std::mutex fMutex;
+    static std::unique_ptr<TRandom3> fRandom;
+
+   public:
+    void Update() override;
+    void InitFromConfigFile() override;
+
+    static void SetSeed(unsigned int seed);
+
+    TRestGeant4ParticleSourceCosmics();
+    ~TRestGeant4ParticleSourceCosmics() = default;
+
+    const char* GetName() const override { return "TRestGeant4ParticleSourceCosmics"; }
+
+    ClassDefOverride(TRestGeant4ParticleSourceCosmics, 1);
+};
+
+#endif  // REST_TRESTGEANT4PARTICLESOURCECOSMICS_H

src/TRestGeant4AnalysisProcess.cxx

@@ -130,7 +130,7 @@
 /// as the diagonal distance of a bounding box defined to contain all the
 /// hits that  produced an energy deposit.
 ///
-/// The following code ilustrates the addition of a primary event
+/// The following code illustrates the addition of a primary event
 /// observable.
 ///
 /// \code
@@ -403,8 +403,10 @@ TRestEvent* TRestGeant4AnalysisProcess::ProcessEvent(TRestEvent* inputEvent) {
     Double_t sensitiveVolumeEnergy = fOutputG4Event->GetEnergyInVolume(sensitiveVolumeName.Data());
     // Get time of the first hit in the sensitive volume
     double hitTime = std::numeric_limits<double>::infinity();
+    std::set<int> trackIdsInSensitiveVolume;
+
     for (const auto& track : fOutputG4Event->GetTracks()) {
-        const auto hits = track.GetHits();
+        const auto& hits = track.GetHits();
         for (size_t hitIndex = 0; hitIndex < hits.GetNumberOfHits(); hitIndex++) {
             const auto volumeName = hits.GetVolumeName(hitIndex);
             if (volumeName != sensitiveVolumeName) {
@@ -419,10 +421,65 @@ TRestEvent* TRestGeant4AnalysisProcess::ProcessEvent(TRestEvent* inputEvent) {
             if (time < hitTime) {
                 hitTime = time;
             }
+
+            trackIdsInSensitiveVolume.insert(track.GetTrackID());
         }
     }
+
     SetObservableValue("sensitiveVolumeFirstHitTime", hitTime);
 
+    std::set<int> trackParentsOfInterest;
+    for (const auto& trackIdInSensitiveVolume : trackIdsInSensitiveVolume) {
+        const auto track = fOutputG4Event->GetTrackByID(trackIdInSensitiveVolume);
+        TRestGeant4Track* parent = track;
+        bool found = false;
+        while (parent != nullptr && !found) {
+            // iterate over hits
+            const auto& hits = parent->GetHits();
+            if (hits.GetVolumeName(0) != sensitiveVolumeName) {
+                for (size_t hitIndex = 0; hitIndex < hits.GetNumberOfHits(); hitIndex++) {
+                    const auto volumeName = hits.GetVolumeName(hitIndex);
+                    if (volumeName != sensitiveVolumeName) {
+                        found = true;
+                        trackParentsOfInterest.insert(parent->GetTrackID());
+                        break;
+                    }
+                }
+            }
+
+            parent = parent->GetParentTrack();
+        }
+    }
+
+    if (!trackParentsOfInterest.empty()) {
+        const auto track = fOutputG4Event->GetTrackByID(*trackParentsOfInterest.begin());
+
+        const auto position = track->GetInitialPosition();
+        SetObservableValue("firstTrackInSensitivePositionX", position.X());
+        SetObservableValue("firstTrackInSensitivePositionY", position.Y());
+        SetObservableValue("firstTrackInSensitivePositionZ", position.Z());
+
+        const auto energy = track->GetInitialKineticEnergy();
+        SetObservableValue("firstTrackInSensitiveEnergy", energy);
+
+        const string particleName = track->GetParticleName().Data();
+        SetObservableValue("firstTrackInSensitiveParticle", particleName);
+
+        string parentParticleName;
+        if (track->GetParentTrack() != nullptr) {
+            parentParticleName = track->GetParentTrack()->GetParticleName().Data();
+        }
+        SetObservableValue("firstTrackInSensitiveParentParticle", parentParticleName);
+
+        const string creatorProcess = track->GetCreatorProcess().Data();
+        SetObservableValue("firstTrackInSensitiveCreatorProcess", creatorProcess);
+
+        const string volumeName = track->GetHits().GetVolumeName(0).Data();
+        SetObservableValue("firstTrackInSensitiveVolumeName", volumeName);
+    }
+
+    SetObservableValue("firstTrackInSensitiveOk", trackParentsOfInterest.size() == 1);
+
     if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Debug) {
         cout << "----------------------------" << endl;
         cout << "TRestGeant4Event : " << fOutputG4Event->GetID() << endl;

src/TRestGeant4Metadata.cxx

@@ -735,6 +735,7 @@
 #include <TRestGDMLParser.h>
 #include <TRestRun.h>
 
+#include "TRestGeant4ParticleSourceCosmics.h"
 #include "TRestGeant4PrimaryGeneratorInfo.h"
 
 using namespace std;
@@ -1065,6 +1066,10 @@ void TRestGeant4Metadata::ReadGenerator() {
         ReadParticleSource(source, sourceDefinition);
         AddParticleSource(source);
 
+        if (string(source->GetName()) == "TRestGeant4ParticleSourceCosmics") {
+            TRestGeant4ParticleSourceCosmics::SetSeed(GetSeed());
+        }
+
         sourceDefinition = GetNextElement(sourceDefinition);
     }
 

src/TRestGeant4ParticleSourceCosmics.cxx

@@ -0,0 +1,114 @@
+
+#include "TRestGeant4ParticleSourceCosmics.h"
+
+#include <TFile.h>
+#include <TH2D.h>
+
+using namespace std;
+
+mutex TRestGeant4ParticleSourceCosmics::fMutex;
+unique_ptr<TRandom3> TRestGeant4ParticleSourceCosmics::fRandom = nullptr;
+
+const map<string, string> geant4ParticleNames = {
+    {"neutron", "neutron"}, {"proton", "proton"}, {"gamma", "gamma"}, {"electron", "e-"}, {"muon", "mu-"}};
+
+TRestGeant4ParticleSourceCosmics::TRestGeant4ParticleSourceCosmics() = default;
+
+void TRestGeant4ParticleSourceCosmics::InitFromConfigFile() {
+    lock_guard<mutex> lock(fMutex);
+
+    cout << "TRestGeant4ParticleSourceCosmics::InitFromConfigFile" << endl;
+    fFilename = GetParameter("filename");
+    const auto particles = GetParameter("particles", "neutron,proton,gamma,electron,muon");
+
+    std::istringstream iss(particles);
+    std::string name;
+    while (std::getline(iss, name, ',')) {
+        fParticleNames.insert(name);
+    }
+
+    for (const auto& particle : fParticleNames) {
+        if (geant4ParticleNames.find(particle) == geant4ParticleNames.end()) {
+            cerr << "Particle name '" << particle << "' is not allowed." << endl;
+            exit(1);
+        }
+    }
+
+    auto file = TFile::Open(fFilename.c_str(), "READ");
+    if (file == nullptr) {
+        cerr << "File '" << fFilename << "' not found." << endl;
+        exit(1);
+    }
+    for (const auto& particle : fParticleNames) {
+        auto histogram = file->Get<TH2D>(string(particle + "_energy_zenith").c_str());
+        if (histogram == nullptr) {
+            cerr << "Histogram '" << particle << "' not found in file '" << fFilename << "'." << endl;
+            exit(1);
+        }
+        fHistograms[particle] = histogram;
+    }
+
+    double sum = 0;
+    for (const auto& particle : fParticleNames) {
+        auto hist = fHistograms.at(particle);
+        fParticleWeights[particle] += hist->GetEntries();
+        sum += fParticleWeights.at(particle);
+    }
+    for (auto& entry : fParticleWeights) {
+        entry.second /= sum;
+        cout << "TRestGeant4ParticleSourceCosmics::InitFromConfigFile: particle: " << entry.first
+             << " sampling weight: " << entry.second << endl;
+    }
+
+    // TODO: how to safely close the file?
+}
+
+void TRestGeant4ParticleSourceCosmics::Update() {
+    lock_guard<mutex> lock(fMutex);
+
+    RemoveParticles();
+
+    string particleName;
+    if (fParticleNames.size() == 1) {
+        particleName = *fParticleNames.begin();
+    } else {
+        const auto random = fRandom->Uniform(0, 1);
+        // choose a random particle based on the weights
+        double sum = 0;
+        for (const auto& entry : fParticleWeights) {
+            sum += entry.second;
+            if (random <= sum) {
+                particleName = entry.first;
+                break;
+            }
+        }
+    }
+
+    auto& hist = fHistograms.at(particleName);
+
+    double energy, zenith;
+    hist->GetRandom2(energy, zenith, fRandom.get());
+
+    particleName = geant4ParticleNames.at(particleName);
+
+    TRestGeant4Particle particle;
+    particle.SetParticleName(particleName);
+
+    particle.SetEnergy(energy * 1000);  // Convert from MeV to keV
+
+    double phi = fRandom->Uniform(0, 1) * TMath::TwoPi();
+    double zenithRad = zenith * TMath::DegToRad();
+
+    // direction towards -y (can be rotated later)
+    const TVector3 direction = {TMath::Sin(zenithRad) * TMath::Cos(phi), -TMath::Cos(zenithRad),
+                                TMath::Sin(zenithRad) * TMath::Sin(phi)};
+
+    particle.SetDirection(direction);
+
+    AddParticle(particle);
+}
+
+void TRestGeant4ParticleSourceCosmics::SetSeed(unsigned int seed) {
+    cout << "TRestGeant4ParticleSourceCosmics::SetSeed: " << seed << endl;
+    fRandom = std::make_unique<TRandom3>(seed);
+}