Modifications to support AmBe sim ROOT files as Primary generator

include/WCSimAmBePrimaryReader.hh

@@ -0,0 +1,93 @@
+#ifndef WCSIMAMBEPRIMARYREADER_HH
+#define WCSIMAMBEPRIMARYREADER_HH
+
+#include <string>
+#include <vector>
+
+#include "Math/Vector4D.h"
+
+class TFile;
+class TTree;
+
+struct WCSimAmBeParticle {
+  int track_id;
+  int parent_id;
+  int pdg;
+  ROOT::Math::XYZTVector position;
+  ROOT::Math::XYZTVector momentum;
+  std::string process;
+
+  WCSimAmBeParticle()
+      : track_id(-1),
+        parent_id(-1),
+        pdg(0),
+        position(),
+        momentum(),
+        process("") {}
+};
+
+struct WCSimAmBeEvent {
+  int rank;
+  int thread_id;
+  int event_id;
+  std::vector<WCSimAmBeParticle> particles;
+
+  WCSimAmBeEvent() : rank(-1), thread_id(-1), event_id(-1), particles() {}
+
+  void Clear() {
+    rank = -1;
+    thread_id = -1;
+    event_id = -1;
+    particles.clear();
+  }
+};
+
+class WCSimAmBePrimaryReader {
+ public:
+  WCSimAmBePrimaryReader();
+  virtual ~WCSimAmBePrimaryReader();
+
+  bool Open(const std::string& filename,
+            const std::string& treeName = "EmergingParticles");
+  void Close();
+
+  bool IsOpen() const;
+  bool HasTree() const;
+
+  bool NextEvent(WCSimAmBeEvent& event);
+  void Reset();
+
+  long long GetEntries() const;
+  long long GetCurrentEntry() const;
+  std::string GetFileName() const;
+  std::string GetTreeName() const;
+
+ private:
+  bool SetupBranches();
+  void ClearBranchPointers();
+  bool ValidateCurrentEntry() const;
+
+ private:
+  TFile* fFile;
+  TTree* fTree;
+
+  std::string fFileName;
+  std::string fTreeName;
+  long long fCurrentEntry;
+  long long fNEntries;
+
+  // Event-level branches
+  int fbranchEmergingRank;
+  int fbranchEmergingThreadId;
+  int fbranchEmergingEventId;
+
+  // Particle-level branches
+  std::vector<int>* fbranchEmergingId;
+  std::vector<int>* fbranchEmergingParentId;
+  std::vector<int>* fbranchEmergingPDG;
+  std::vector<ROOT::Math::XYZTVector>* fbranchEmergingPos;
+  std::vector<ROOT::Math::XYZTVector>* fbranchEmergingP;
+  std::vector<std::string>* fbranchEmergingProcess;
+};
+
+#endif

include/WCSimPrimaryGeneratorAction.hh

@@ -20,6 +20,8 @@
 #include "Framework/Interaction/Interaction.h"
 #endif
 #include "WCSimRootOptions.hh"
+#include "G4ThreeVector.hh"
+
 
 class WCSimDetectorConstruction;
 class G4ParticleGun;
@@ -28,6 +30,9 @@ class G4Event;
 class WCSimPrimaryGeneratorMessenger;
 class G4Generator;
 
+class WCSimAmBePrimaryReader;
+
+
 class WCSimPrimaryGeneratorAction : public G4VUserPrimaryGeneratorAction
 {
 public:
@@ -78,6 +83,12 @@ public:
 
   void SaveOptionsToOutput(WCSimRootOptions * wcopt);
 
+
+  // For AmBe source sim
+  G4bool OpenAmBePrimaryFile(const G4String& fileName);
+  void SetAmBePositionOffset(const G4ThreeVector& v) { amBePositionOffset = v; }
+  G4ThreeVector GetAmBePositionOffset() const { return amBePositionOffset; }
+
 private:
   WCSimDetectorConstruction*      myDetector;
   G4ParticleGun*                  particleGun;
@@ -140,7 +151,12 @@ private:
 	G4String neutrinosDirectory;
 	G4bool loadNewPrimaries;
 	G4int primariesoffset;
-
+
+    // AmBe external ROOT-input mode
+  G4bool useAmBeRootInput;
+  G4String amBeInputFileName;
+  WCSimAmBePrimaryReader* amBeReader;
+  G4ThreeVector amBePositionOffset;
 public:
 
   inline void SetMulineEvtGenerator(G4bool choice) { useMulineEvt = choice; }
@@ -159,6 +175,12 @@ public:
   inline void SetGPSEvtGenerator(G4bool choice) { useGPSEvt = choice; }
   inline G4bool IsUsingGPSEvtGenerator()  { return useGPSEvt; }
 
+  inline void SetAmBeRootGenerator(G4bool choice) { useAmBeRootInput = choice; }
+  inline G4bool IsUsingAmBeRootGenerator() const { return useAmBeRootInput; }
+
+  inline void SetAmBeInputFileName(const G4String& fileName) { amBeInputFileName = fileName; }
+  inline G4String GetAmBeInputFileName() const { return amBeInputFileName; }
+
   inline void OpenVectorFile(G4String fileName) 
   {
     if ( inputFile.is_open() ) 

include/WCSimPrimaryGeneratorMessenger.hh

@@ -5,6 +5,7 @@ class WCSimPrimaryGeneratorAction;
 class G4UIdirectory;
 class G4UIcmdWithAString;
 class G4UIcmdWithAnInteger;
+class G4UIcmdWith3VectorAndUnit;
 
 #include "G4UImessenger.hh"
 #include "globals.hh"
@@ -29,6 +30,10 @@ class WCSimPrimaryGeneratorMessenger: public G4UImessenger
   G4UIcmdWithAString* primariesfileDirectoryCmd;
   G4UIcmdWithAString* neutrinosfileDirectoryCmd;
   G4UIcmdWithAnInteger* primariesStartEventCmd;
+
+  // For AmBe simulation
+  G4UIcmdWithAString* ambeFileCmd;
+  G4UIcmdWith3VectorAndUnit* ambeOffsetCmd;
 
 };
 

src/WCSimAmBePrimaryReader.cc

@@ -0,0 +1,217 @@
+#include "WCSimAmBePrimaryReader.hh"
+
+#include <iostream>
+
+// ROOT
+#include "TFile.h"
+#include "TTree.h"
+
+WCSimAmBePrimaryReader::WCSimAmBePrimaryReader()
+    : fFile(0),
+      fTree(0),
+      fFileName(""),
+      fTreeName("EmergingParticles"),
+      fCurrentEntry(0),
+      fNEntries(0),
+      fbranchEmergingRank(-1),
+      fbranchEmergingThreadId(-1),
+      fbranchEmergingEventId(-1),
+      fbranchEmergingId(0),
+      fbranchEmergingParentId(0),
+      fbranchEmergingPDG(0),
+      fbranchEmergingPos(0),
+      fbranchEmergingP(0),
+      fbranchEmergingProcess(0) {}
+
+WCSimAmBePrimaryReader::~WCSimAmBePrimaryReader() {
+  Close();
+}
+
+bool WCSimAmBePrimaryReader::Open(const std::string& filename,
+                                  const std::string& treeName) {
+  Close();
+
+  fFileName = filename;
+  fTreeName = treeName;
+  fCurrentEntry = 0;
+  fNEntries = 0;
+
+  fFile = TFile::Open(filename.c_str(), "READ");
+  if (!fFile || fFile->IsZombie()) {
+    std::cerr << "WCSimAmBePrimaryReader::Open(): failed to open file: "
+              << filename << std::endl;
+    Close();
+    return false;
+  }
+
+  fTree = dynamic_cast<TTree*>(fFile->Get(treeName.c_str()));
+  if (!fTree) {
+    std::cerr << "WCSimAmBePrimaryReader::Open(): failed to find tree '"
+              << treeName << "' in file: " << filename << std::endl;
+    Close();
+    return false;
+  }
+
+  if (!SetupBranches()) {
+    std::cerr << "WCSimAmBePrimaryReader::Open(): failed to set branch addresses"
+              << std::endl;
+    Close();
+    return false;
+  }
+
+  fNEntries = static_cast<long long>(fTree->GetEntries());
+
+  std::cout << "WCSimAmBePrimaryReader: opened file " << fFileName
+            << " with tree " << fTreeName
+            << " containing " << fNEntries << " entries." << std::endl;
+
+  return true;
+}
+
+void WCSimAmBePrimaryReader::Close() {
+  ClearBranchPointers();
+
+  if (fFile) {
+    fFile->Close();
+    delete fFile;
+    fFile = 0;
+  }
+
+  fTree = 0;
+  fFileName = "";
+  fTreeName = "EmergingParticles";
+  fCurrentEntry = 0;
+  fNEntries = 0;
+  fbranchEmergingRank = -1;
+  fbranchEmergingThreadId = -1;
+  fbranchEmergingEventId = -1;
+}
+
+bool WCSimAmBePrimaryReader::IsOpen() const {
+  return (fFile != 0);
+}
+
+bool WCSimAmBePrimaryReader::HasTree() const {
+  return (fTree != 0);
+}
+
+void WCSimAmBePrimaryReader::Reset() {
+  fCurrentEntry = 0;
+}
+
+long long WCSimAmBePrimaryReader::GetEntries() const {
+  return fNEntries;
+}
+
+long long WCSimAmBePrimaryReader::GetCurrentEntry() const {
+  return fCurrentEntry;
+}
+
+std::string WCSimAmBePrimaryReader::GetFileName() const {
+  return fFileName;
+}
+
+std::string WCSimAmBePrimaryReader::GetTreeName() const {
+  return fTreeName;
+}
+
+bool WCSimAmBePrimaryReader::SetupBranches() {
+  if (!fTree) return false;
+
+  ClearBranchPointers();
+
+  fTree->SetBranchAddress("Rank", &fbranchEmergingRank);
+  fTree->SetBranchAddress("Thread", &fbranchEmergingThreadId);
+  fTree->SetBranchAddress("EventId", &fbranchEmergingEventId);
+  fTree->SetBranchAddress("TrackId", &fbranchEmergingId);
+  fTree->SetBranchAddress("ParentId", &fbranchEmergingParentId);
+  fTree->SetBranchAddress("PDG", &fbranchEmergingPDG);
+  fTree->SetBranchAddress("Vertex", &fbranchEmergingPos);
+  fTree->SetBranchAddress("Momentum", &fbranchEmergingP);
+  fTree->SetBranchAddress("Process", &fbranchEmergingProcess);
+
+  return true;
+}
+
+void WCSimAmBePrimaryReader::ClearBranchPointers() {
+  fbranchEmergingId = 0;
+  fbranchEmergingParentId = 0;
+  fbranchEmergingPDG = 0;
+  fbranchEmergingPos = 0;
+  fbranchEmergingP = 0;
+  fbranchEmergingProcess = 0;
+}
+
+bool WCSimAmBePrimaryReader::ValidateCurrentEntry() const {
+  if (!fbranchEmergingId ||
+      !fbranchEmergingParentId ||
+      !fbranchEmergingPDG ||
+      !fbranchEmergingPos ||
+      !fbranchEmergingP ||
+      !fbranchEmergingProcess) {
+    std::cerr << "WCSimAmBePrimaryReader::ValidateCurrentEntry(): "
+              << "one or more branch pointers are null." << std::endl;
+    return false;
+  }
+
+  const std::size_t n = fbranchEmergingPDG->size();
+
+  if (fbranchEmergingId->size() != n ||
+      fbranchEmergingParentId->size() != n ||
+      fbranchEmergingPos->size() != n ||
+      fbranchEmergingP->size() != n ||
+      fbranchEmergingProcess->size() != n) {
+    std::cerr << "WCSimAmBePrimaryReader::ValidateCurrentEntry(): "
+              << "inconsistent vector sizes in EventId "
+              << fbranchEmergingEventId << std::endl;
+    return false;
+  }
+
+  return true;
+}
+
+bool WCSimAmBePrimaryReader::NextEvent(WCSimAmBeEvent& event) {
+  event.Clear();
+
+  if (!fTree) {
+    std::cerr << "WCSimAmBePrimaryReader::NextEvent(): no tree is loaded."
+              << std::endl;
+    return false;
+  }
+
+  if (fCurrentEntry >= fNEntries) {
+    return false;
+  }
+
+  Long64_t bytesRead = fTree->GetEntry(fCurrentEntry);
+  if (bytesRead <= 0) {
+    std::cerr << "WCSimAmBePrimaryReader::NextEvent(): failed to read entry "
+              << fCurrentEntry << std::endl;
+    return false;
+  }
+
+  if (!ValidateCurrentEntry()) {
+    return false;
+  }
+
+  event.rank = fbranchEmergingRank;
+  event.thread_id = fbranchEmergingThreadId;
+  event.event_id = fbranchEmergingEventId;
+
+  event.particles.reserve(fbranchEmergingPDG->size());
+
+  for (std::size_t i = 0; i < fbranchEmergingPDG->size(); ++i) {
+    WCSimAmBeParticle particle;
+    particle.track_id = fbranchEmergingId->at(i);
+    particle.parent_id = fbranchEmergingParentId->at(i);
+    particle.pdg = fbranchEmergingPDG->at(i);
+    particle.position = fbranchEmergingPos->at(i);
+    particle.momentum = fbranchEmergingP->at(i);
+    particle.process = fbranchEmergingProcess->at(i);
+
+    event.particles.push_back(particle);
+  }
+
+  ++fCurrentEntry;
+  return true;
+}