eic-opticks/src/GPURaytrace.h

0001 #include <filesystem>
0002 #include <fstream>
0003 #include <iostream>
0004
0005 #include "G4AutoLock.hh"
0006 #include "G4BooleanSolid.hh"
0007 #include "G4Cerenkov.hh"
0008 #include "G4Electron.hh"
0009 #include "G4Event.hh"
0010 #include "G4GDMLParser.hh"
0011 #include "G4LogicalVolumeStore.hh"
0012 #include "G4OpBoundaryProcess.hh"
0013 #include "G4OpticalPhoton.hh"
0014 #include "G4PhysicalConstants.hh"
0015 #include "G4PrimaryParticle.hh"
0016 #include "G4PrimaryVertex.hh"
0017 #include "G4RunManager.hh"
0018 #include "G4RunManagerFactory.hh"
0019 #include "G4SDManager.hh"
0020 #include "G4Scintillation.hh"
0021 #include "G4SubtractionSolid.hh"
0022 #include "G4SystemOfUnits.hh"
0023 #include "G4ThreeVector.hh"
0024 #include "G4Track.hh"
0025 #include "G4TrackStatus.hh"
0026 #include "G4UserEventAction.hh"
0027 #include "G4UserRunAction.hh"
0028 #include "G4UserSteppingAction.hh"
0029 #include "G4UserTrackingAction.hh"
0030 #include "G4VPhysicalVolume.hh"
0031 #include "G4VProcess.hh"
0032 #include "G4VUserDetectorConstruction.hh"
0033 #include "G4VUserPrimaryGeneratorAction.hh"
0034
0035 #include "g4cx/G4CXOpticks.hh"
0036 #include "sysrap/NP.hh"
0037 #include "sysrap/SEvt.hh"
0038 #include "sysrap/STrackInfo.h"
0039 #include "sysrap/spho.h"
0040 #include "sysrap/sphoton.h"
0041 #include "u4/U4.hh"
0042 #include "u4/U4Random.hh"
0043 #include "u4/U4StepPoint.hh"
0044 #include "u4/U4Touchable.h"
0045 #include "u4/U4Track.h"
0046
0047 namespace
0048 {
0049 G4Mutex genstep_mutex = G4MUTEX_INITIALIZER;
0050 }
0051
0052 bool IsSubtractionSolid(G4VSolid *solid)
0053 {
0054     if (!solid)
0055         return false;
0056
0057     // Check if the solid is directly a G4SubtractionSolid
0058     if (dynamic_cast<G4SubtractionSolid *>(solid))
0059         return true;
0060
0061     // If the solid is a Boolean solid, check its constituent solids
0062     G4BooleanSolid *booleanSolid = dynamic_cast<G4BooleanSolid *>(solid);
0063     if (booleanSolid)
0064     {
0065         G4VSolid *solidA = booleanSolid->GetConstituentSolid(0);
0066         G4VSolid *solidB = booleanSolid->GetConstituentSolid(1);
0067
0068         // Recursively check the constituent solids
0069         if (IsSubtractionSolid(solidA) || IsSubtractionSolid(solidB))
0070             return true;
0071     }
0072
0073     // For other solid types, return false
0074     return false;
0075 }
0076
0077 std::string str_tolower(std::string s)
0078 {
0079     std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
0080     return s;
0081 }
0082
0083 struct PhotonHit : public G4VHit
0084 {
0085     PhotonHit() = default;
0086
0087     PhotonHit(unsigned id, G4double energy, G4double time, G4ThreeVector position, G4ThreeVector direction,
0088               G4ThreeVector polarization)
0089         : fid(id), fenergy(energy), ftime(time), fposition(position), fdirection(direction), fpolarization(polarization)
0090     {
0091     }
0092
0093     // Copy constructor
0094     PhotonHit(const PhotonHit &right)
0095         : G4VHit(right), fid(right.fid), fenergy(right.fenergy), ftime(right.ftime), fposition(right.fposition),
0096           fdirection(right.fdirection), fpolarization(right.fpolarization)
0097     {
0098     }
0099
0100     // Assignment operator
0101     const PhotonHit &operator=(const PhotonHit &right)
0102     {
0103         if (this != &right)
0104         {
0105             G4VHit::operator=(right);
0106             fid = right.fid;
0107             fenergy = right.fenergy;
0108             ftime = right.ftime;
0109             fposition = right.fposition;
0110             fdirection = right.fdirection;
0111             fpolarization = right.fpolarization;
0112         }
0113         return *this;
0114     }
0115
0116     // Equality operator
0117     G4bool operator==(const PhotonHit &right) const
0118     {
0119         return (this == &right);
0120     }
0121
0122     // Print method
0123     void Print() override
0124     {
0125         G4cout << "Detector id: " << fid << " energy: " << fenergy << " nm" << " time: " << ftime << " ns"
0126                << " position: " << fposition << " direction: " << fdirection << " polarization: " << fpolarization
0127                << G4endl;
0128     }
0129
0130     // Member variables
0131     G4int fid{0};
0132     G4double fenergy{0};
0133     G4double ftime{0};
0134     G4ThreeVector fposition{0, 0, 0};
0135     G4ThreeVector fdirection{0, 0, 0};
0136     G4ThreeVector fpolarization{0, 0, 0};
0137 };
0138
0139 using PhotonHitsCollection = G4THitsCollection<PhotonHit>;
0140
0141 struct PhotonSD : public G4VSensitiveDetector
0142 {
0143     PhotonSD(G4String name) : G4VSensitiveDetector(name), fHCID(-1)
0144     {
0145         G4String HCname = name + "_HC";
0146         collectionName.insert(HCname);
0147         G4cout << collectionName.size() << "   PhotonSD name:  " << name << " collection Name: " << HCname << G4endl;
0148     }
0149
0150     void Initialize(G4HCofThisEvent *hce) override
0151     {
0152         fPhotonHitsCollection = new PhotonHitsCollection(SensitiveDetectorName, collectionName[0]);
0153         if (fHCID < 0)
0154         {
0155             // G4cout << "PhotonSD::Initialize:  " << SensitiveDetectorName << "   " << collectionName[0] << G4endl;
0156             fHCID = G4SDManager::GetSDMpointer()->GetCollectionID(collectionName[0]);
0157         }
0158         hce->AddHitsCollection(fHCID, fPhotonHitsCollection);
0159     }
0160
0161     G4bool ProcessHits(G4Step *aStep, G4TouchableHistory *) override
0162     {
0163         G4Track *theTrack = aStep->GetTrack();
0164         if (theTrack->GetDefinition() != G4OpticalPhoton::OpticalPhotonDefinition())
0165             return false;
0166
0167         G4double theEnergy = theTrack->GetTotalEnergy() / CLHEP::eV;
0168
0169         // Create a new hit (CopyNr is set to 0 as DetectorID is omitted)
0170         PhotonHit *newHit = new PhotonHit(
0171             0, // CopyNr set to 0
0172             theEnergy, theTrack->GetGlobalTime(), aStep->GetPostStepPoint()->GetPosition(),
0173             aStep->GetPostStepPoint()->GetMomentumDirection(), aStep->GetPostStepPoint()->GetPolarization());
0174
0175         fPhotonHitsCollection->insert(newHit);
0176         theTrack->SetTrackStatus(fStopAndKill);
0177         return true;
0178     }
0179
0180     void EndOfEvent(G4HCofThisEvent *) override
0181     {
0182         G4int NbHits = fPhotonHitsCollection->entries();
0183         G4cout << "PhotonSD::EndOfEvent Number of PhotonHits: " << NbHits << G4endl;
0184     }
0185
0186     void AddOpticksHits()
0187     {
0188         SEvt *sev = SEvt::Get_EGPU();
0189         unsigned int num_hits = sev->GetNumHit(0);
0190
0191         for (int idx = 0; idx < int(num_hits); idx++)
0192         {
0193             sphoton hit;
0194             sev->getHit(hit, idx);
0195             G4ThreeVector position = G4ThreeVector(hit.pos.x, hit.pos.y, hit.pos.z);
0196             G4ThreeVector direction = G4ThreeVector(hit.mom.x, hit.mom.y, hit.mom.z);
0197             G4ThreeVector polarization = G4ThreeVector(hit.pol.x, hit.pol.y, hit.pol.z);
0198             int theCreationProcessid;
0199             if (OpticksPhoton::HasCerenkovFlag(hit.flagmask))
0200             {
0201                 theCreationProcessid = 0;
0202             }
0203             else if (OpticksPhoton::HasScintillationFlag(hit.flagmask))
0204             {
0205                 theCreationProcessid = 1;
0206             }
0207             else
0208             {
0209                 theCreationProcessid = -1;
0210             }
0211             std::cout << hit.wavelength << " " << position << " " << direction << " " << polarization << std::endl;
0212
0213             PhotonHit *newHit = new PhotonHit(0, hit.wavelength, hit.time, position, direction, polarization);
0214             fPhotonHitsCollection->insert(newHit);
0215         }
0216     }
0217
0218   private:
0219     PhotonHitsCollection *fPhotonHitsCollection{nullptr};
0220     G4int fHCID;
0221 };
0222
0223 struct DetectorConstruction : G4VUserDetectorConstruction
0224 {
0225     DetectorConstruction(std::filesystem::path gdml_file) : gdml_file_(gdml_file)
0226     {
0227     }
0228
0229     G4VPhysicalVolume *Construct() override
0230     {
0231         parser_.Read(gdml_file_.string(), false);
0232         G4VPhysicalVolume *world = parser_.GetWorldVolume();
0233
0234         G4CXOpticks::SetGeometry(world);
0235         G4LogicalVolumeStore *lvStore = G4LogicalVolumeStore::GetInstance();
0236
0237         static G4VisAttributes invisibleVisAttr(false);
0238
0239         // Check if the store is not empty
0240         if (lvStore && !lvStore->empty())
0241         {
0242             // Iterate over all logical volumes in the store
0243             for (auto &logicalVolume : *lvStore)
0244             {
0245                 G4VSolid *solid = logicalVolume->GetSolid();
0246
0247                 // Check if the solid uses subtraction
0248                 if (IsSubtractionSolid(solid))
0249                 {
0250                     // Assign the invisible visual attributes to the logical volume
0251                     logicalVolume->SetVisAttributes(&invisibleVisAttr);
0252
0253                     // Optionally, print out the name of the logical volume
0254                     G4cout << "Hiding logical volume: " << logicalVolume->GetName() << G4endl;
0255                 }
0256             }
0257         }
0258
0259         return world;
0260     }
0261
0262     void ConstructSDandField() override
0263     {
0264         G4cout << "ConstructSDandField is called." << G4endl;
0265         G4SDManager *SDman = G4SDManager::GetSDMpointer();
0266
0267         const G4GDMLAuxMapType *auxmap = parser_.GetAuxMap();
0268         for (auto const &[logVol, listType] : *auxmap)
0269         {
0270             for (auto const &auxtype : listType)
0271             {
0272                 if (auxtype.type == "SensDet")
0273                 {
0274                     G4cout << "Attaching sensitive detector to logical volume: " << logVol->GetName() << G4endl;
0275                     G4String name = logVol->GetName() + "_PhotonDetector";
0276                     PhotonSD *aPhotonSD = new PhotonSD(name);
0277                     SDman->AddNewDetector(aPhotonSD);
0278                     logVol->SetSensitiveDetector(aPhotonSD);
0279                 }
0280             }
0281         }
0282     }
0283
0284   private:
0285     std::filesystem::path gdml_file_;
0286     G4GDMLParser parser_;
0287 };
0288
0289 struct PrimaryGenerator : G4VUserPrimaryGeneratorAction
0290 {
0291     SEvt *sev;
0292
0293     PrimaryGenerator(SEvt *sev) : sev(sev)
0294     {
0295     }
0296
0297     void GeneratePrimaries(G4Event *event) override
0298     {
0299         G4ThreeVector position_mm(0.0 * m, 0.0 * m, 0.0 * m);
0300         G4double time_ns = 0;
0301         G4ThreeVector direction(0, 0.2, 0.8);
0302         G4double wavelength_nm = 0.1;
0303
0304         G4PrimaryVertex *vertex = new G4PrimaryVertex(position_mm, time_ns);
0305         G4PrimaryParticle *particle = new G4PrimaryParticle(G4Electron::Definition());
0306         particle->SetKineticEnergy(5 * GeV);
0307         particle->SetMomentumDirection(direction);
0308         vertex->SetPrimary(particle);
0309         event->AddPrimaryVertex(vertex);
0310     }
0311 };
0312
0313 struct EventAction : G4UserEventAction
0314 {
0315     SEvt *sev;
0316     G4int fTotalG4Hits{0};
0317
0318     EventAction(SEvt *sev) : sev(sev)
0319     {
0320     }
0321
0322     void BeginOfEventAction(const G4Event *event) override
0323     {
0324     }
0325
0326     void EndOfEventAction(const G4Event *event) override
0327     {
0328         G4HCofThisEvent *hce = event->GetHCofThisEvent();
0329         if (hce)
0330         {
0331             for (G4int i = 0; i < hce->GetNumberOfCollections(); i++)
0332             {
0333                 G4VHitsCollection *hc = hce->GetHC(i);
0334                 if (hc)
0335                 {
0336                     fTotalG4Hits += hc->GetSize();
0337                 }
0338             }
0339         }
0340     }
0341
0342     G4int GetTotalG4Hits() const
0343     {
0344         return fTotalG4Hits;
0345     }
0346 };
0347
0348 struct RunAction : G4UserRunAction
0349 {
0350     EventAction *fEventAction;
0351
0352     RunAction(EventAction *eventAction) : fEventAction(eventAction)
0353     {
0354     }
0355
0356     void BeginOfRunAction(const G4Run *run) override
0357     {
0358     }
0359
0360     void EndOfRunAction(const G4Run *run) override
0361     {
0362         if (G4Threading::IsMasterThread())
0363         {
0364             G4CXOpticks *gx = G4CXOpticks::Get();
0365
0366             auto start = std::chrono::high_resolution_clock::now();
0367             gx->simulate(0, false);
0368             cudaDeviceSynchronize();
0369             auto end = std::chrono::high_resolution_clock::now();
0370             // Compute duration
0371             std::chrono::duration<double> elapsed = end - start;
0372             std::cout << "Simulation time: " << elapsed.count() << " seconds" << std::endl;
0373
0374             // unsigned int num_hits = SEvt::GetNumHit(EGPU);
0375             SEvt *sev = SEvt::Get_EGPU();
0376             unsigned int num_hits = sev->GetNumHit(0);
0377
0378             std::cout << "Opticks: NumCollected:  " << sev->GetNumGenstepFromGenstep(0) << std::endl;
0379             std::cout << "Opticks: NumCollected:  " << sev->GetNumPhotonCollected(0) << std::endl;
0380             std::cout << "Opticks: NumHits:  " << num_hits << std::endl;
0381             std::cout << "Geant4: NumHits:  " << fEventAction->GetTotalG4Hits() << std::endl;
0382             std::ofstream outFile("opticks_hits_output.txt");
0383             if (!outFile.is_open())
0384             {
0385                 std::cerr << "Error opening output file!" << std::endl;
0386                 return;
0387             }
0388
0389             for (int idx = 0; idx < int(num_hits); idx++)
0390             {
0391                 sphoton hit;
0392                 sev->getHit(hit, idx);
0393                 G4ThreeVector position = G4ThreeVector(hit.pos.x, hit.pos.y, hit.pos.z);
0394                 G4ThreeVector direction = G4ThreeVector(hit.mom.x, hit.mom.y, hit.mom.z);
0395                 G4ThreeVector polarization = G4ThreeVector(hit.pol.x, hit.pol.y, hit.pol.z);
0396                 int theCreationProcessid;
0397                 if (OpticksPhoton::HasCerenkovFlag(hit.flagmask))
0398                 {
0399                     theCreationProcessid = 0;
0400                 }
0401                 else if (OpticksPhoton::HasScintillationFlag(hit.flagmask))
0402                 {
0403                     theCreationProcessid = 1;
0404                 }
0405                 else
0406                 {
0407                     theCreationProcessid = -1;
0408                 }
0409                 //    std::cout << "Adding hit from Opticks:" << hit.wavelength << " " << position << " " << direction
0410                 //    << "
0411                 //    "
0412                 //              << polarization << std::endl;
0413                 outFile << hit.time << " " << hit.wavelength << "  " << "(" << position.x() << ", " << position.y()
0414                         << ", " << position.z() << ")  " << "(" << direction.x() << ", " << direction.y() << ", "
0415                         << direction.z() << ")  " << "(" << polarization.x() << ", " << polarization.y() << ", "
0416                         << polarization.z() << ")  " << "CreationProcessID=" << theCreationProcessid << std::endl;
0417             }
0418
0419             outFile.close();
0420         }
0421     }
0422 };
0423
0424 struct SteppingAction : G4UserSteppingAction
0425 {
0426     SEvt *sev;
0427
0428     SteppingAction(SEvt *sev) : sev(sev)
0429     {
0430     }
0431
0432     void UserSteppingAction(const G4Step *aStep)
0433     {
0434         G4Track *aTrack;
0435         G4int fNumPhotons = 0;
0436
0437         G4StepPoint *preStep = aStep->GetPostStepPoint();
0438         G4VPhysicalVolume *volume = preStep->GetPhysicalVolume();
0439
0440         if (aStep->GetTrack()->GetDefinition() == G4OpticalPhoton::OpticalPhotonDefinition())
0441         {
0442             // Kill if step count exceeds 10000 to avoid reflection forever
0443             if (aStep->GetTrack()->GetCurrentStepNumber() > 10000)
0444             {
0445                 aStep->GetTrack()->SetTrackStatus(fStopAndKill);
0446             }
0447         }
0448
0449         if (volume && volume->GetName() == "MirrorPyramid")
0450         {
0451             aTrack = aStep->GetTrack();
0452             if (aTrack->GetDefinition() != G4OpticalPhoton::OpticalPhotonDefinition())
0453             {
0454                 aTrack->SetTrackStatus(fStopAndKill);
0455             }
0456         }
0457
0458         G4SteppingManager *fpSteppingManager =
0459             G4EventManager::GetEventManager()->GetTrackingManager()->GetSteppingManager();
0460         G4StepStatus stepStatus = fpSteppingManager->GetfStepStatus();
0461
0462         if (stepStatus != fAtRestDoItProc)
0463         {
0464             G4ProcessVector *procPost = fpSteppingManager->GetfPostStepDoItVector();
0465             size_t MAXofPostStepLoops = fpSteppingManager->GetMAXofPostStepLoops();
0466             for (size_t i3 = 0; i3 < MAXofPostStepLoops; i3++)
0467             {
0468                 if ((*procPost)[i3]->GetProcessName() == "Cerenkov")
0469                 {
0470                     aTrack = aStep->GetTrack();
0471                     const G4DynamicParticle *aParticle = aTrack->GetDynamicParticle();
0472                     G4double charge = aParticle->GetDefinition()->GetPDGCharge();
0473                     const G4Material *aMaterial = aTrack->GetMaterial();
0474                     G4MaterialPropertiesTable *MPT = aMaterial->GetMaterialPropertiesTable();
0475
0476                     G4MaterialPropertyVector *Rindex = MPT->GetProperty(kRINDEX);
0477                     if (!Rindex || Rindex->GetVectorLength() == 0)
0478                     {
0479                         G4cout << "WARNING: Material has no valid RINDEX data. Skipping Cerenkov calculation."
0480                                << G4endl;
0481                         return;
0482                     }
0483
0484                     G4Cerenkov *proc = (G4Cerenkov *)(*procPost)[i3];
0485                     fNumPhotons = proc->GetNumPhotons();
0486
0487                     G4AutoLock lock(&genstep_mutex); // <-- Mutex is locked here
0488
0489                     if (fNumPhotons > 0)
0490                     {
0491                         G4double Pmin = Rindex->Energy(0);
0492                         G4double Pmax = Rindex->GetMaxEnergy();
0493                         G4double nMax = Rindex->GetMaxValue();
0494                         G4double beta1 = aStep->GetPreStepPoint()->GetBeta();
0495                         G4double beta2 = aStep->GetPostStepPoint()->GetBeta();
0496                         G4double beta = (beta1 + beta2) * 0.5;
0497                         G4double BetaInverse = 1. / beta;
0498                         G4double maxCos = BetaInverse / nMax;
0499                         G4double maxSin2 = (1.0 - maxCos) * (1.0 + maxCos);
0500                         G4double MeanNumberOfPhotons1 =
0501                             proc->GetAverageNumberOfPhotons(charge, beta1, aMaterial, Rindex);
0502                         G4double MeanNumberOfPhotons2 =
0503                             proc->GetAverageNumberOfPhotons(charge, beta2, aMaterial, Rindex);
0504
0505                         U4::CollectGenstep_G4Cerenkov_modified(aTrack, aStep, fNumPhotons, BetaInverse, Pmin, Pmax,
0506                                                                maxCos, maxSin2, MeanNumberOfPhotons1,
0507                                                                MeanNumberOfPhotons2);
0508                     }
0509                 }
0510                 if ((*procPost)[i3]->GetProcessName() == "Scintillation")
0511                 {
0512                     G4Scintillation *proc1 = (G4Scintillation *)(*procPost)[i3];
0513                     fNumPhotons = proc1->GetNumPhotons();
0514                     if (fNumPhotons > 0)
0515                     {
0516                         aTrack = aStep->GetTrack();
0517                         const G4Material *aMaterial = aTrack->GetMaterial();
0518                         G4MaterialPropertiesTable *MPT = aMaterial->GetMaterialPropertiesTable();
0519                         if (!MPT || !MPT->ConstPropertyExists(kSCINTILLATIONTIMECONSTANT1))
0520                         {
0521                             G4cout << "WARNING: Material has no valid SCINTILLATIONTIMECONSTANT1 data. Skipping "
0522                                       "Scintillation calculation."
0523                                    << G4endl;
0524                             return;
0525                         }
0526                         G4double SCINTILLATIONTIMECONSTANT1 = MPT->GetConstProperty(kSCINTILLATIONTIMECONSTANT1);
0527
0528                         U4::CollectGenstep_DsG4Scintillation_r4695(aTrack, aStep, fNumPhotons, 1,
0529                                                                    SCINTILLATIONTIMECONSTANT1);
0530                     }
0531                 }
0532             }
0533         }
0534     }
0535 };
0536
0537 struct TrackingAction : G4UserTrackingAction
0538 {
0539     const G4Track *transient_fSuspend_track = nullptr;
0540     SEvt *sev;
0541
0542     TrackingAction(SEvt *sev) : sev(sev)
0543     {
0544     }
0545
0546     void PreUserTrackingAction_Optical_FabricateLabel(const G4Track *track)
0547     {
0548     }
0549
0550     void PreUserTrackingAction(const G4Track *track) override
0551     {
0552     }
0553
0554     void PostUserTrackingAction(const G4Track *track) override
0555     {
0556     }
0557 };
0558
0559 struct G4App
0560 {
0561     G4App(std::filesystem::path gdml_file)
0562         : sev(SEvt::CreateOrReuse_EGPU()), det_cons_(new DetectorConstruction(gdml_file)),
0563           prim_gen_(new PrimaryGenerator(sev)), event_act_(new EventAction(sev)), run_act_(new RunAction(event_act_)),
0564           stepping_(new SteppingAction(sev)),
0565
0566           tracking_(new TrackingAction(sev))
0567     {
0568     }
0569
0570     //~G4App(){ G4CXOpticks::Finalize();}
0571
0572     // Create "global" event
0573     SEvt *sev;
0574
0575     G4VUserDetectorConstruction *det_cons_;
0576     G4VUserPrimaryGeneratorAction *prim_gen_;
0577     EventAction *event_act_;
0578     RunAction *run_act_;
0579     SteppingAction *stepping_;
0580     TrackingAction *tracking_;
0581 };