Skip to content

Commit fc072ce

Browse files
lobislobis
authored
Merge pull request #129 from rest-for-physics/lobis-cosmic-revamp
Update cosmic sampler
2 parents f1c00d2 + 4b9e3ad commit fc072ce

File tree

3 files changed

+133
-50
lines changed

3 files changed

+133
-50
lines changed

‎include/Application.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ class Application {
4848
private:
4949
SimulationManager fSimulationManager;
5050

51-
void ValidateOutputFile(const std::string& outputFile) const;
51+
static void ValidateOutputFile(const std::string& outputFile);
5252
};
5353

5454
#endif // REST_APPLICATION_H

‎src/Application.cxx‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
#ifndef GEANT4_WITHOUT_G4RunManagerFactory
1616
#include <G4RunManagerFactory.hh>
1717
#endif
18+
#include <TRestGeant4ParticleSourceCosmics.h>
19+
1820
#include <G4UImanager.hh>
1921
#include <G4VSteppingVerbose.hh>
2022
#include <cstdlib>
@@ -26,7 +28,6 @@
2628
#include "PhysicsList.h"
2729
#include "PrimaryGeneratorAction.h"
2830
#include "RunAction.h"
29-
#include "SimulationManager.h"
3031
#include "SteppingAction.h"
3132
#include "SteppingVerbose.h"
3233

@@ -492,10 +493,31 @@ void Application::Run(const CommandLineOptions::Options& options) {
492493

493494
const auto nEntries = run->GetEntries();
494495

496+
if (metadata->GetNumberOfSources() == 1 &&
497+
string(metadata->GetParticleSource(0)->GetName()) == "TRestGeant4ParticleSourceCosmics") {
498+
auto source = dynamic_cast<TRestGeant4ParticleSourceCosmics*>(metadata->GetParticleSource(0));
499+
auto histogramsTransformed = source->GetHistogramsTransformed();
500+
double totalParticlesPerUnitTimePerSurface = 0;
501+
for (const auto& [name, histogram] : histogramsTransformed) {
502+
totalParticlesPerUnitTimePerSurface += histogram->Integral();
503+
}
504+
const double nParticlesLaunched = metadata->GetNumberOfEvents();
505+
const double equivalentSurface =
506+
metadata->GetGeant4PrimaryGeneratorInfo().GetSpatialGeneratorCosmicSurfaceTermCm2();
507+
const double time = nParticlesLaunched / (totalParticlesPerUnitTimePerSurface * equivalentSurface);
508+
cout << "Total particles launched: " << nParticlesLaunched << endl;
509+
cout << "Total particles per second per cm2: " << totalParticlesPerUnitTimePerSurface << endl;
510+
cout << "Equivalent surface: " << equivalentSurface << " cm2" << endl;
511+
cout << "Total time to launch all particles: " << time << " s" << endl;
512+
cout << "Counts per second: " << double(run->GetEntries()) / time << endl;
513+
cout << "Counts per second (wall time): "
514+
<< double(run->GetEntries()) / fSimulationManager.GetElapsedTime() << endl;
515+
metadata->SetSimulationTime(time);
516+
}
517+
495518
run->UpdateOutputFile();
496519
run->CloseFile();
497520

498-
metadata->PrintMetadata();
499521
run->PrintMetadata();
500522

501523
const auto nEventsAtEnd =
@@ -505,13 +527,14 @@ void Application::Run(const CommandLineOptions::Options& options) {
505527
ValidateOutputFile(filename);
506528

507529
cout << "\n\t- Total simulation time is " << ToTimeStringLong(fSimulationManager.GetElapsedTime()) << ", "
508-
<< nEventsAtEnd << " processed events (" << nEventsAtEnd / fSimulationManager.GetElapsedTime()
509-
<< " per second) and " << nEntries << " events saved to output file ("
510-
<< nEntries / fSimulationManager.GetElapsedTime() << " per second)" << endl;
530+
<< nEventsAtEnd << " processed events ("
531+
<< double(nEventsAtEnd) / fSimulationManager.GetElapsedTime() << " per second) and " << nEntries
532+
<< " events saved to output file (" << double(nEntries) / fSimulationManager.GetElapsedTime()
533+
<< " per second)" << endl;
511534
cout << "\t- Output file: " << filename << endl << endl;
512535
}
513536

514-
void Application::ValidateOutputFile(const string& filename) const {
537+
void Application::ValidateOutputFile(const string& filename) {
515538
bool error = false;
516539

517540
const auto run = TRestRun(filename);

‎src/PrimaryGeneratorAction.cxx‎

Lines changed: 103 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ std::mutex PrimaryGeneratorAction::fPrimaryGenerationMutex;
4343
TF1* PrimaryGeneratorAction::fEnergyDistributionFunction = nullptr;
4444
TF1* PrimaryGeneratorAction::fAngularDistributionFunction = nullptr;
4545
TF2* PrimaryGeneratorAction::fEnergyAndAngularDistributionFunction = nullptr;
46+
TH2D* energyAndAngularDistributionHistogram = nullptr;
4647

4748
PrimaryGeneratorAction::PrimaryGeneratorAction(SimulationManager* simulationManager)
4849
: G4VUserPrimaryGeneratorAction(), fSimulationManager(simulationManager) {
@@ -230,6 +231,30 @@ void PrimaryGeneratorAction::SetGeneratorSpatialDensity(TString str) {
230231
fGeneratorSpatialDensityFunction = new TF3("GeneratorDistFunc", str);
231232
}
232233

234+
TVector2 PointOnUnitDisk() {
235+
double r = TMath::Sqrt(G4UniformRand());
236+
double theta = G4UniformRand() * 2 * TMath::Pi();
237+
return {r * TMath::Cos(theta), r * TMath::Sin(theta)};
238+
}
239+
240+
pair<bool, TVector3> IntersectionLineSphere(const TVector3& lineOrigin, const TVector3& lineDirection) {
241+
// sphere origin is always (0,0)
242+
// return the first intersection point
243+
// https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
244+
const double a = lineDirection.Dot(lineDirection);
245+
const double b = 2 * lineDirection.Dot(lineOrigin);
246+
const double c = lineOrigin.Dot(lineOrigin) - 1;
247+
248+
const double discriminant = b * b - 4 * a * c;
249+
if (discriminant < 0) {
250+
return {false, {0, 0, 0}};
251+
}
252+
const double t1 = (-b + sqrt(discriminant)) / (2 * a);
253+
const double t2 = (-b - sqrt(discriminant)) / (2 * a);
254+
const double t = min(t1, t2);
255+
return {true, lineOrigin + t * lineDirection};
256+
}
257+
233258
void PrimaryGeneratorAction::GeneratePrimaries(G4Event* event) {
234259
lock_guard<mutex> lock(fPrimaryGenerationMutex);
235260
auto simulationManager = fSimulationManager;
@@ -261,29 +286,61 @@ void PrimaryGeneratorAction::GeneratePrimaries(G4Event* event) {
261286
G4ParticleTable::GetParticleTable()->FindParticle(particle.GetParticleName().Data()));
262287

263288
if (fCosmicCircumscribedSphereRadius == 0.) {
264-
// radius in mm
265289
fCosmicCircumscribedSphereRadius = fSimulationManager->GetRestMetadata()
266290
->GetGeant4PrimaryGeneratorInfo()
267-
.GetSpatialGeneratorCosmicRadius();
291+
.GetSpatialGeneratorCosmicRadius(); // radius in mm
292+
}
293+
const G4ThreeVector referenceDirection = {0, -1, 0};
294+
const auto& direction = particle.GetMomentumDirection();
295+
const double zenith = TMath::ACos(
296+
direction.Dot({referenceDirection.x(), referenceDirection.y(), referenceDirection.z()}));
297+
298+
const TVector2 positionOnDisk = PointOnUnitDisk();
299+
const TVector2 positionOnEllipse = {positionOnDisk.X() / TMath::Cos(zenith) + TMath::Tan(zenith),
300+
positionOnDisk.Y()};
301+
302+
double phi = TVector2(direction.X(), direction.Z()).Phi();
303+
const TVector2 positionOnEllipseRotated = {
304+
positionOnEllipse.X() * TMath::Cos(phi) - positionOnEllipse.Y() * TMath::Sin(phi),
305+
positionOnEllipse.X() * TMath::Sin(phi) + positionOnEllipse.Y() * TMath::Cos(phi),
306+
};
307+
const TVector3 positionOrigin = {
308+
// this may be wrong, or how phi is derived from direction
309+
-1 * positionOnEllipseRotated.X(),
310+
1.0,
311+
-1 * positionOnEllipseRotated.Y(),
312+
};
313+
const auto [intersectionFlag, intersection] = IntersectionLineSphere(positionOrigin, direction);
314+
if (!intersectionFlag) {
315+
cerr << "PrimaryGeneratorAction: cosmic generator intersection not found (this should never "
316+
"happen)"
317+
<< endl;
318+
exit(1);
268319
}
269320

270-
const auto position = ComputeCosmicPosition(fParticleGun.GetParticleMomentumDirection(),
271-
fCosmicCircumscribedSphereRadius);
272-
273-
fParticleGun.SetParticlePosition(position);
274321
fParticleGun.SetParticleEnergy(particle.GetEnergy() * keV);
275-
fParticleGun.SetParticleMomentumDirection({particle.GetMomentumDirection().X(),
276-
particle.GetMomentumDirection().Y(),
277-
particle.GetMomentumDirection().Z()});
322+
323+
G4ThreeVector particlePosition = {
324+
intersection.X() * fCosmicCircumscribedSphereRadius,
325+
intersection.Y() * fCosmicCircumscribedSphereRadius,
326+
intersection.Z() * fCosmicCircumscribedSphereRadius,
327+
};
328+
G4ThreeVector sourceDirection = {source->GetDirection().X(), source->GetDirection().Y(),
329+
source->GetDirection().Z()};
330+
G4ThreeVector particleDirection = {direction.X(), direction.Y(), direction.Z()};
331+
332+
const double angleBetweenDirections = sourceDirection.angle(referenceDirection);
333+
if (angleBetweenDirections > 0) {
334+
const G4ThreeVector cross = referenceDirection.cross(sourceDirection);
335+
particleDirection.rotate(angleBetweenDirections, cross);
336+
particlePosition.rotate(angleBetweenDirections, cross);
337+
}
338+
339+
fParticleGun.SetParticleMomentumDirection(particleDirection);
340+
fParticleGun.SetParticlePosition(particlePosition);
341+
278342
fParticleGun.GeneratePrimaryVertex(event);
279343

280-
/*
281-
cout << "PrimaryGeneratorAction - INFO: cosmic particle generated: "
282-
<< fParticleGun.GetParticleDefinition()->GetParticleName()
283-
<< " energy: " << fParticleGun.GetParticleEnergy() / keV << " keV"
284-
<< " position: " << fParticleGun.GetParticlePosition() / mm << " mm"
285-
<< " direction: " << fParticleGun.GetParticleMomentumDirection() << endl;
286-
*/
287344
return;
288345
}
289346

@@ -317,21 +374,17 @@ void PrimaryGeneratorAction::GeneratePrimaries(G4Event* event) {
317374

318375
for (int i = 0; i < restG4Metadata->GetNumberOfSources(); i++) {
319376
vector<TRestGeant4Particle> particles = restG4Metadata->GetParticleSource(i)->GetParticles();
320-
// std::cout << "Source : " << i << std::endl;
321377
for (const auto& p : particles) {
322-
// ParticleDefinition should be always declared first (after position).
323378
SetParticleDefinition(i, p);
324379
SetParticleEnergyAndDirection(i, p);
325380

326-
// p.Print();
327-
328381
if (spatialGeneratorTypeEnum == SpatialGeneratorTypes::COSMIC) {
329382
const auto position = ComputeCosmicPosition(fParticleGun.GetParticleMomentumDirection(),
330383
fCosmicCircumscribedSphereRadius);
331384
fParticleGun.SetParticlePosition(position);
332385
}
333386

334-
if (spatialGeneratorTypeEnum == SpatialGeneratorTypes::SOURCE) {
387+
else if (spatialGeneratorTypeEnum == SpatialGeneratorTypes::SOURCE) {
335388
G4ThreeVector position = {p.GetOrigin().X(), p.GetOrigin().Y(), p.GetOrigin().Z()};
336389
fParticleGun.SetParticlePosition(position);
337390
}
@@ -822,37 +875,44 @@ void PrimaryGeneratorAction::SetParticleEnergyAndDirection(Int_t particleSourceI
822875
restG4Metadata->GetParticleSource(particleSourceIndex)->GetEnergyDistributionType().Data();
823876
const auto energyDistTypeEnum = StringToEnergyDistributionTypes(energyDistTypeName);
824877

825-
if (energyDistTypeEnum != EnergyDistributionTypes::FORMULA2 &&
826-
angularDistTypeEnum != AngularDistributionTypes::FORMULA2) {
827-
// when not using 'FORMULA2'
878+
if ((energyDistTypeEnum != EnergyDistributionTypes::FORMULA2 &&
879+
angularDistTypeEnum != AngularDistributionTypes::FORMULA2) &&
880+
(energyDistTypeEnum != EnergyDistributionTypes::TH2D &&
881+
angularDistTypeEnum != AngularDistributionTypes::TH2D)) {
828882
SetParticleEnergy(particleSourceIndex, particle);
829883
SetParticleDirection(particleSourceIndex, particle);
830884
return;
831885
}
832886

833-
// this function should only be invoked when both energy and angular dist are "formula2"
834-
if (energyDistTypeEnum != EnergyDistributionTypes::FORMULA2 ||
835-
angularDistTypeEnum != AngularDistributionTypes::FORMULA2) {
836-
cout << "PrimaryGeneratorAction::SetParticleEnergyAndDirection - this method should only invoked "
837-
"when both angular and energy dist are 'formula2'"
887+
double energy, angle;
888+
if (energyDistTypeEnum == EnergyDistributionTypes::FORMULA2 &&
889+
angularDistTypeEnum == AngularDistributionTypes::FORMULA2) {
890+
{
891+
lock_guard<mutex> lock(fDistributionFormulaMutex);
892+
fEnergyAndAngularDistributionFunction->GetRandom2(energy, angle, fRandom);
893+
energy *= keV;
894+
}
895+
} else if (energyDistTypeEnum == EnergyDistributionTypes::TH2D &&
896+
angularDistTypeEnum == AngularDistributionTypes::TH2D) {
897+
if (energyAndAngularDistributionHistogram == nullptr) {
898+
const auto filename = source->GetEnergyDistributionFilename();
899+
const auto name = source->GetEnergyDistributionNameInFile();
900+
cout << "Loading energy and angular distribution from file " << filename << " with name " << name
901+
<< endl;
902+
TFile* file = TFile::Open(filename);
903+
energyAndAngularDistributionHistogram =
904+
file->Get<TH2D>(name); // it's the same for both angular and energy
905+
}
906+
energyAndAngularDistributionHistogram->GetRandom2(energy, angle);
907+
energy *= MeV; // energy in these histograms is in MeV. TODO: parse energy from axis label
908+
angle *= TMath::DegToRad();
909+
} else {
910+
cout << "Energy/Angular distribution type 'formula2' or 'TH2D' should be used on both energy and "
911+
"angular"
838912
<< endl;
839913
exit(1);
840914
}
841915

842-
if (fEnergyAndAngularDistributionFunction == nullptr) {
843-
RESTError << "PrimaryGeneratorAction::SetParticleEnergyAndDirection - energy and angular "
844-
"distribution function is not set"
845-
<< RESTendl;
846-
exit(1);
847-
}
848-
849-
double energy, angle;
850-
{
851-
lock_guard<mutex> lock(fDistributionFormulaMutex);
852-
fEnergyAndAngularDistributionFunction->GetRandom2(energy, angle, fRandom);
853-
}
854-
energy *= keV;
855-
856916
G4ThreeVector direction = {source->GetDirection().X(), source->GetDirection().Y(),
857917
source->GetDirection().Z()};
858918

0 commit comments

Comments
 (0)