eic-opticks/src/GPUPhotonSourceMinimal.h

0001 #include <filesystem>
0002 #include <fstream>
0003 #include <vector>
0004
0005 #include "G4Event.hh"
0006 #include "G4GDMLParser.hh"
0007 #include "G4OpticalPhoton.hh"
0008 #include "G4PhysicalConstants.hh"
0009 #include "G4PrimaryParticle.hh"
0010 #include "G4PrimaryVertex.hh"
0011 #include "G4SystemOfUnits.hh"
0012 #include "G4ThreeVector.hh"
0013 #include "G4UserEventAction.hh"
0014 #include "G4UserRunAction.hh"
0015 #include "G4VPhysicalVolume.hh"
0016 #include "G4VUserDetectorConstruction.hh"
0017 #include "G4VUserPrimaryGeneratorAction.hh"
0018
0019 #include "g4cx/G4CXOpticks.hh"
0020 #include "sysrap/NP.hh"
0021 #include "sysrap/SEvt.hh"
0022 #include "sysrap/sphoton.h"
0023
0024 #include "config.h"
0025 #include "torch.h"
0026
0027 struct DetectorConstruction : G4VUserDetectorConstruction
0028 {
0029     DetectorConstruction(std::filesystem::path gdml_file) : gdml_file_(gdml_file)
0030     {
0031     }
0032
0033     G4VPhysicalVolume *Construct() override
0034     {
0035         parser_.Read(gdml_file_.string(), false);
0036         G4VPhysicalVolume *world = parser_.GetWorldVolume();
0037
0038         G4CXOpticks::SetGeometry(world);
0039
0040         return world;
0041     }
0042
0043   private:
0044     std::filesystem::path gdml_file_;
0045     G4GDMLParser parser_;
0046 };
0047
0048 struct PrimaryGenerator : G4VUserPrimaryGeneratorAction
0049 {
0050     gphox::Config cfg;
0051     SEvt *sev;
0052
0053     PrimaryGenerator(const gphox::Config &cfg, SEvt *sev) : cfg(cfg), sev(sev)
0054     {
0055     }
0056
0057     void GeneratePrimaries(G4Event *event) override
0058     {
0059         std::vector<sphoton> sphotons = generate_photons(cfg.torch);
0060
0061         size_t num_floats = sphotons.size() * 4 * 4;
0062         float *data = reinterpret_cast<float *>(sphotons.data());
0063         NP *photons = NP::MakeFromValues<float>(data, num_floats);
0064
0065         photons->reshape({static_cast<int64_t>(sphotons.size()), 4, 4});
0066
0067         for (const sphoton &p : sphotons)
0068         {
0069             G4ThreeVector position_mm(p.pos.x, p.pos.y, p.pos.z);
0070             G4double time_ns = p.time;
0071             G4ThreeVector direction(p.mom.x, p.mom.y, p.mom.z);
0072             G4double wavelength_nm = p.wavelength;
0073             G4ThreeVector polarization(p.pol.x, p.pol.y, p.pol.z);
0074
0075             G4PrimaryVertex *vertex = new G4PrimaryVertex(position_mm, time_ns);
0076             G4double kineticEnergy = h_Planck * c_light / (wavelength_nm * nm);
0077
0078             G4PrimaryParticle *particle = new G4PrimaryParticle(G4OpticalPhoton::Definition());
0079             particle->SetKineticEnergy(kineticEnergy);
0080             particle->SetMomentumDirection(direction);
0081             particle->SetPolarization(polarization);
0082
0083             vertex->SetPrimary(particle);
0084             event->AddPrimaryVertex(vertex);
0085         }
0086
0087         sev->SetInputPhoton(photons);
0088     }
0089 };
0090
0091 struct EventAction : G4UserEventAction
0092 {
0093     SEvt *sev;
0094     unsigned int fTotalOpticksHits{0};
0095
0096     EventAction(SEvt *sev) : sev(sev)
0097     {
0098     }
0099
0100     void BeginOfEventAction(const G4Event *event) override
0101     {
0102         sev->beginOfEvent(event->GetEventID());
0103     }
0104
0105     void EndOfEventAction(const G4Event *event) override
0106     {
0107         int eventID = event->GetEventID();
0108         sev->addEventConfigArray();
0109         sev->gather();
0110         sev->endOfEvent(eventID);
0111
0112         // GPU-based simulation
0113         G4CXOpticks *gx = G4CXOpticks::Get();
0114
0115         gx->simulate(eventID, false);
0116         cudaDeviceSynchronize();
0117
0118         SEvt *sev_gpu = SEvt::Get_EGPU();
0119         unsigned int num_hits = sev_gpu->GetNumHit(0);
0120
0121         std::cout << "Opticks: NumHits:  " << num_hits << std::endl;
0122         fTotalOpticksHits += num_hits;
0123
0124         std::ofstream outFile("opticks_hits_output.txt");
0125         if (!outFile.is_open())
0126         {
0127             std::cerr << "Error opening output file!" << std::endl;
0128         }
0129         else
0130         {
0131             for (int idx = 0; idx < int(num_hits); idx++)
0132             {
0133                 sphoton hit;
0134                 sev_gpu->getHit(hit, idx);
0135                 G4ThreeVector position = G4ThreeVector(hit.pos.x, hit.pos.y, hit.pos.z);
0136                 G4ThreeVector direction = G4ThreeVector(hit.mom.x, hit.mom.y, hit.mom.z);
0137                 G4ThreeVector polarization = G4ThreeVector(hit.pol.x, hit.pol.y, hit.pol.z);
0138                 outFile << hit.time << " " << hit.wavelength << "  " << "(" << position.x() << ", " << position.y()
0139                         << ", " << position.z() << ")  " << "(" << direction.x() << ", " << direction.y() << ", "
0140                         << direction.z() << ")  " << "(" << polarization.x() << ", " << polarization.y() << ", "
0141                         << polarization.z() << ")" << std::endl;
0142             }
0143             outFile.close();
0144         }
0145
0146         gx->reset(eventID);
0147     }
0148
0149     unsigned int GetTotalOpticksHits() const
0150     {
0151         return fTotalOpticksHits;
0152     }
0153 };
0154
0155 struct RunAction : G4UserRunAction
0156 {
0157     EventAction *fEventAction;
0158
0159     RunAction(EventAction *eventAction) : fEventAction(eventAction)
0160     {
0161     }
0162
0163     void BeginOfRunAction(const G4Run *) override
0164     {
0165     }
0166
0167     void EndOfRunAction(const G4Run *) override
0168     {
0169         std::cout << "Opticks: NumHits:  " << fEventAction->GetTotalOpticksHits() << std::endl;
0170     }
0171 };
0172
0173 struct G4App
0174 {
0175     G4App(const gphox::Config &cfg, std::filesystem::path gdml_file)
0176         : sev(SEvt::CreateOrReuse_ECPU()), det_cons_(new DetectorConstruction(gdml_file)),
0177           prim_gen_(new PrimaryGenerator(cfg, sev)), event_act_(new EventAction(sev)),
0178           run_act_(new RunAction(event_act_))
0179     {
0180     }
0181
0182     SEvt *sev;
0183
0184     G4VUserDetectorConstruction *det_cons_;
0185     G4VUserPrimaryGeneratorAction *prim_gen_;
0186     EventAction *event_act_;
0187     RunAction *run_act_;
0188 };