changed to be serialisable for multiprocessing

CMakeLists.txt

@@ -4,8 +4,6 @@
 cmake_minimum_required(VERSION 3.16...3.21)
 project(minicalo)
 
-
-
 #----------------------------------------------------------------------------
 # Find Geant4 package, activating all available UI and Vis drivers by default
 # You can set WITH_GEANT4_UIVIS to OFF via the command line or ccmake/cmake-gui
@@ -33,19 +31,32 @@ file(GLOB sources ${PROJECT_SOURCE_DIR}/src/*.cc)
 file(GLOB headers ${PROJECT_SOURCE_DIR}/include/*.hh)
 
 #----------------------------------------------------------------------------
-# Add the executable, and link it to the Geant4 libraries
+# Find Python and pybind11
+#
+find_package(Python3 COMPONENTS Development Interpreter REQUIRED)
+add_subdirectory(lib/pybind11)
+
+# Add the include directories for pybind11, Python, and the project headers
+include_directories(
+  ${PROJECT_SOURCE_DIR}/include
+  ${Python3_INCLUDE_DIRS}
+  lib/pybind11/include
+)
+
+# Collect binding sources and headers
+file(GLOB_RECURSE BIND_SOURCES "${PROJECT_SOURCE_DIR}/bind/*.cpp" "${PROJECT_SOURCE_DIR}/bind/*.cc")
+
+# Add the pybind11 module
+pybind11_add_module(minicalo ${sources} ${BIND_SOURCES})
+
+# Link the Geant4 and Python libraries with the pybind11 module
+target_link_libraries(minicalo PUBLIC ${Geant4_LIBRARIES} ${Python3_LIBRARIES})
+
+#----------------------------------------------------------------------------
+# Add the executable, and link it to the Geant4 and Python libraries
 #
 add_executable(exampleB4a exampleB4a.cc ${sources} ${headers})
-target_link_libraries(exampleB4a ${Geant4_LIBRARIES})
-
-
-## needs to be added later
-
-file(GLOB_RECURSE SOURCES "${PROJECT_SOURCE_DIR}/src/*.cc" )
-add_subdirectory(lib/pybind11)
-pybind11_add_module(minicalo ${SOURCES} "${PROJECT_SOURCE_DIR}/bind/bindings.cpp")
-target_include_directories(minicalo PUBLIC lib/pybind11/include)
-target_link_libraries(minicalo PUBLIC ${Geant4_LIBRARIES})
+target_link_libraries(exampleB4a ${Geant4_LIBRARIES} ${Python3_LIBRARIES})
 
 #----------------------------------------------------------------------------
 # Copy all scripts to the build directory, i.e. the directory in which we
@@ -62,14 +73,14 @@ set(EXAMPLEB4A_SCRIPTS
   run1.mac
   run2.mac
   vis.mac
-  )
+)
 
 foreach(_script ${EXAMPLEB4A_SCRIPTS})
   configure_file(
     ${PROJECT_SOURCE_DIR}/${_script}
     ${PROJECT_BINARY_DIR}/${_script}
     COPYONLY
-    )
+  )
 endforeach()
 
 #----------------------------------------------------------------------------

bind/G4Calo.py

@@ -31,7 +31,7 @@ class __G4System(_G4System):
         save_file = len(filename) > 0
         # filename without file ending(!)
         filename = "_" + str(time.perf_counter_ns()) + ".root"
-        _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename=filename)
+        _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename)
         
         # TO FIX: Geant4 adds "t<threadnumber>" to the filename, circumvent this for one thread, but this is not a good solution
         file = glob.glob(filename.replace(".root",  "*.root"))
@@ -294,4 +294,55 @@ def index_out_of_bounds_workaround(tbranch):
     return df.reset_index(drop=True)
 
 
-G4System = __G4System()#singleton instance
+_s_G4System = __G4System()#singleton instance
+
+
+def _run_mini_batch(
+        cw : GeometryDescriptor,
+        nEvents: int,
+        particleSpec: str,
+        minEnergy_GeV: float,
+        maxEnergy_GeV: float = -1.0,
+        counter : int = 0):
+
+    
+    print(f"Running mini batch {counter} with {nEvents} events")
+    import time
+    time.sleep(counter/1000)
+    print(f'done sleeping {counter}')
+
+    from  G4Calo import G4System
+    G4System.init(cw)
+
+    df = _s_G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"")
+    return df
+
+
+
+def run_batch(nEvents: int,
+        particleSpec: str,
+        minEnergy_GeV: float,
+        maxEnergy_GeV: float = -1.0,):
+    '''
+    splits the batch in jobs depending on how many cores are available and runs mini batches in parallel
+    '''
+    assert nEvents > 0
+
+    nCores = multiprocessing.cpu_count()
+    #make sure to adjust cores such that at least 200 events are run per core
+    nCores = min(nCores, nEvents // 200 + 1)
+
+    print(f"Running on {nCores} cores")
+    nEventsPerCore = nEvents // nCores
+    print(f"Running {nEventsPerCore} events per core")
+    nEventsLastCore = nEvents - nEventsPerCore * (nCores - 1)
+    print(f"Running {nEventsLastCore} events on last core")
+    
+    nevents = [nEventsPerCore if i < nCores - 1 else nEventsLastCore for i in range(nCores)]
+
+    #use a multiprocessing pool to run the mini batches in parallel
+    with multiprocessing.Pool(nCores) as pool:
+        dfs = pool.starmap(_run_mini_batch, [(cw, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, i) for i in range(nCores)])
+
+    return pd.concat(dfs)
+        

bind/bindings.cpp

@@ -1,79 +1,78 @@
-
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
-#include <pybind11/operators.h>
-
 #include "GeometryDescriptor.hh"
 #include "G4System.hh"
 
 namespace py = pybind11;
 
+PYBIND11_MODULE(minicalo, m) {
+    py::class_<Sensor>(m, "Sensor")
+        .def(py::init<>())
+        .def("getEnergy", &Sensor::getEnergy)
+        .def("getPos", &Sensor::getPos)
+        .def("getSize", &Sensor::getSize)
+        .def("getX", &Sensor::getX)
+        .def("getY", &Sensor::getY)
+        .def("getZ", &Sensor::getZ)
+        .def("getdx", &Sensor::getdx)
+        .def("getdy", &Sensor::getdy)
+        .def("getdz", &Sensor::getdz)
+        .def(py::pickle(
+            [](const Sensor &s) { // __getstate__
+                return s.__getstate__();
+            },
+            [](py::tuple t) { // __setstate__
+                return Sensor::__setstate__(t);
+            }
+        ));
 
-template<class M>
-void makeGeometryDescriptor(M & m, std::string name){
+    py::class_<Layer>(m, "Layer")
+        .def(py::init<>())
+        .def("setThickness", &Layer::setThickness)
+        .def("setMaterial", &Layer::setMaterial)
+        .def("setNx", &Layer::setNx)
+        .def("setNy", &Layer::setNy)
+        .def("setIsActive", &Layer::setIsActive)
+        .def("assignPhysicalVolume", &Layer::assignPhysicalVolume)
+        .def("unAssign", &Layer::unAssign)
+        .def(py::pickle(
+            [](const Layer &l) { // __getstate__
+                return l.__getstate__();
+            },
+            [](py::tuple t) { // __setstate__
+                return Layer::__setstate__(t);
+            }
+        ));
 
-    py::class_<GeometryDescriptor>(m, name.data()).def(py::init())
-    .def("addLayer", &GeometryDescriptor::addLayer, py::arg("thickness"), py::arg("material"), py::arg("isActive")=true, py::arg("nx")=1, py::arg("ny")=-1)
-    .def("getXYWidth", &GeometryDescriptor::getXYWidth)
-    // bind overloaded getLayers function
-    .def("getLayers", (std::vector<Layer> & (GeometryDescriptor::*)()) &GeometryDescriptor::getLayers)
-    .def("getLayers", (const std::vector<Layer> & (GeometryDescriptor::*)() const) &GeometryDescriptor::getLayers)
-    .def("isAssigned", &GeometryDescriptor::isAssigned)
-    .def("getNSensors", &GeometryDescriptor::getNSensors);
+    py::class_<GeometryDescriptor>(m, "GeometryDescriptor")
+        .def(py::init<>())
+        .def("addLayer", &GeometryDescriptor::addLayer)
+        .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers))
+        .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers, py::const_))
+        .def("getXYWidth", &GeometryDescriptor::getXYWidth)
+        .def("resetSensorEnergies", &GeometryDescriptor::resetSensorEnergies)
+        .def("getNSensors", &GeometryDescriptor::getNSensors)
+        .def("printSensorEnergies", &GeometryDescriptor::printSensorEnergies)
+        .def("isAssigned", &GeometryDescriptor::isAssigned)
+        .def("unAssign", &GeometryDescriptor::unAssign)
+        .def("isEmpty", &GeometryDescriptor::isEmpty)
+        .def(py::pickle(
+            [](const GeometryDescriptor &g) { // __getstate__
+                return g.__getstate__();
+            },
+            [](py::tuple t) { // __setstate__
+                return GeometryDescriptor::__setstate__(t);
+            }
+        ));
 
+    //now for G4System
+    py::class_<G4System>(m, "G4System")
+        .def(py::init<>())
+        .def("init", &G4System::init)
+        .def("run_visualize", &G4System::run_visualize)
+        .def("run_gui", &G4System::run_gui)
+        .def("run_batch", &G4System::run_batch)
+        .def("applyUICommand", &G4System::applyUICommand)
+        .def("displayEvent", &G4System::displayEvent)
+        .def("printMaterial", &G4System::printMaterial);
 }
-
-
-template <class M>
-void makeG4System(M &m, std::string name)
-{
-    py::class_<G4System>(m, name.data()).def(py::init())
-    .def("init", &G4System::init, py::arg("cw")).def("run_visualize", &G4System::run_visualize, py::arg("partSpecies"), py::arg("minEnergy_GeV"), py::arg("maxEnergy_GeV"))
-    .def("run_gui", &G4System::run_gui)
-    .def("run_batch", &G4System::run_batch, py::arg("nEvents"), py::arg("partSpecies"), py::arg("minEnergy_GeV"), py::arg("maxEnergy_GeV"), py::arg("filename")="_1234567890_Hits.root")
-    .def("applyUICommand", &G4System::applyUICommand, py::arg("command"))
-    .def("displayEvent", &G4System::displayEvent)
-   // .def("printMaterial", &G4System::printMaterial, py::arg("name"))
-    //.def("check", &G4System::check)
-    .def("getGeometryDescriptor", &G4System::getGeometryDescriptor);
-}
-
-// create bindings for Layer class
-template<class M>
-void makeLayer(M &m, std::string name){
-    py::class_<Layer>(m, name.data()).def(py::init())
-    .def_readwrite("thickness", &Layer::thickness)
-    .def_readwrite("material", &Layer::material)
-    .def_readwrite("nx", &Layer::nx)
-    .def_readwrite("ny", &Layer::ny)
-    .def_readwrite("isActive", &Layer::isActive)
-    .def_readwrite("sens_xwidth", &Layer::sens_xwidth)
-    .def_readwrite("sens_ywidth", &Layer::sens_ywidth)
-    .def_readwrite("sensors", &Layer::sensors);
-}
-
-// create bindings for sensor class
-template<class M>
-void makeSensor(M &m, std::string name){
-    py::class_<Sensor>(m, name.data()).def(py::init())
-    .def("getEnergy", &Sensor::getEnergy)
-    .def("getPos", &Sensor::getPos)
-    .def("getSize", &Sensor::getSize)
-    .def("getX", &Sensor::getX)
-    .def("getY", &Sensor::getY)
-    .def("getZ", &Sensor::getZ)
-    .def("getdx", &Sensor::getdx)
-    .def("getdy", &Sensor::getdy)
-    .def("getdz", &Sensor::getdz);
-}
-
-
-
-PYBIND11_MODULE(minicalo, m)
-{
-    m.doc() = "pybind11 plugin"; // optional module docstring
-    makeGeometryDescriptor(m, "GeometryDescriptor");
-    makeG4System(m, "G4System");
-    makeSensor(m, "Sensor");
-    makeLayer(m, "Layer");
-}

include/G4System.hh

@@ -22,7 +22,8 @@ class G4System{
   friend class GeometryDescriptor;
 public:
 
-G4System(bool Gui=false):gui(Gui){};
+G4System():gui(false){};
+G4System(bool Gui):gui(Gui){};
 ~G4System(){
     if(visManager != nullptr){
     delete visManager;

include/GeometryDescriptor.hh

@@ -1,64 +1,77 @@
-
 #ifndef GeometryDescriptor_HH
 #define GeometryDescriptor_HH
 
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
 #include "G4ThreeVector.hh"
 #include <vector>
 #include <string>
-//#include "G4VPhysicalVolume.hh"
-class G4VPhysicalVolume;
+#include "G4VPhysicalVolume.hh"
 class G4System;
 
-class Sensor{
+class Sensor {
 public:
-    Sensor(){};
-    ~Sensor(){};
+    Sensor() {};
+    ~Sensor() {};
 
-    const G4double getEnergy()const{
+    const G4double getEnergy() const {
         return energy;
     }
-    G4ThreeVector getPos()const{
+    G4ThreeVector getPos() const {
         return position;
     }
-    G4ThreeVector getSize()const{
+    G4ThreeVector getSize() const {
         return size;
     }
 
-    double getX()const{
+    double getX() const {
         return position.x();
-    }  
+    }
 
-    double getY()const{
+    double getY() const {
         return position.y();
     }
 
-    double getZ()const{
+    double getZ() const {
         return position.z();
     }
 
-    double getdx()const{
+    double getdx() const {
         return size.x();
     }
-    
-    double getdy()const{
+
+    double getdy() const {
         return size.y();
     }
 
-    double getdz()const{
+    double getdz() const {
         return size.z();
     }
 
     G4ThreeVector position;
     G4ThreeVector size;
     mutable G4double energy;
+
+    pybind11::tuple __getstate__() const {
+        return pybind11::make_tuple(position.x(), position.y(), position.z(), size.x(), size.y(), size.z(), energy);
+    }
+
+    static Sensor __setstate__(pybind11::tuple t) {
+        if (t.size() != 7) throw std::runtime_error("Invalid state!");
+        Sensor sensor;
+        sensor.position = G4ThreeVector(t[0].cast<double>(), t[1].cast<double>(), t[2].cast<double>());
+        sensor.size = G4ThreeVector(t[3].cast<double>(), t[4].cast<double>(), t[5].cast<double>());
+        sensor.energy = t[6].cast<G4double>();
+        return sensor;
+    }
 };
 
-class Layer{
-    
+class Layer {
 public:
-    Layer();
-    ~Layer(){};
-    
+    Layer() : thickness(0),sens_xwidth(0), sens_ywidth(0),  material(""), nx(1), ny(1), isActive(false), physicalVolume(nullptr) {};
+    ~Layer() {};
+
     void setThickness(double thickness_cm);
     void setMaterial(std::string material);
     void setNx(int nx);
@@ -66,7 +79,7 @@ public:
     void setIsActive(bool isActive);
     void assignPhysicalVolume(G4VPhysicalVolume* physicalVolume);
 
-    void unAssign(){
+    void unAssign() {
         physicalVolume = nullptr;
         sensors.clear();
     }
@@ -83,68 +96,107 @@ public:
     std::string name;
 
     std::vector<Sensor> sensors;
+
+    pybind11::tuple __getstate__() const {
+        return pybind11::make_tuple(thickness, sens_xwidth, sens_ywidth, material, nx, ny, isActive, sensors);
+    }
+
+    static Layer __setstate__(pybind11::tuple t) {
+        if (t.size() != 8) throw std::runtime_error("Invalid state!");
+        Layer layer;
+        layer.thickness = t[0].cast<double>();
+        layer.sens_xwidth = t[1].cast<double>();
+        layer.sens_ywidth = t[2].cast<double>();
+        layer.material = t[3].cast<std::string>();
+        layer.nx = t[4].cast<int>();
+        layer.ny = t[5].cast<int>();
+        layer.isActive = t[6].cast<bool>();
+        layer.sensors = t[7].cast<std::vector<Sensor>>();
+        layer.physicalVolume = nullptr; // Reset pointer
+        return layer;
+    }
 };
 
-
-class GeometryDescriptor{
-
+class GeometryDescriptor {
 public:
-  GeometryDescriptor(double xy_width=50);
-  ~GeometryDescriptor();
+    GeometryDescriptor() : xywidth(50), g4system(nullptr) {};
+    ~GeometryDescriptor();
 
-  void addLayer(double thickness_cm, std::string material, bool isActive=true, int nx=1, int ny=-1);
-  
-   std::vector<Layer>& getLayers() ;
-   const std::vector<Layer>& getLayers() const;
+    void addLayer(double thickness_cm, std::string material, bool isActive = true, int nx = 1, int ny = -1);
 
-   double getXYWidth() const;
+    std::vector<Layer>& getLayers() {
+        return layers;
+    }
+    const std::vector<Layer>& getLayers() const {
+        return layers;
+    }
 
+    double getXYWidth() const {
+        return xywidth;
+    }
 
-    void resetSensorEnergies()const;//energies are mutable
-    int getNSensors()const{ 
+    void resetSensorEnergies() const; // energies are mutable
+    int getNSensors() const {
         int n_sensors = 0;
-        for(const auto& layer: layers){
+        for (const auto& layer : layers) {
             n_sensors += layer.sensors.size();
         }
         return n_sensors;
     }
 
-     Layer* getLayerByVolume(G4VPhysicalVolume* volume);
+    Layer* getLayerByVolume(G4VPhysicalVolume* volume);
+    const Layer* getLayerByVolume(G4VPhysicalVolume* volume) const;
+    Sensor* getSensorByVolume(G4VPhysicalVolume* volume);
+    const Sensor* getSensorByVolume(G4VPhysicalVolume* volume) const;
+    void printSensorEnergies() const;
 
-     const Layer* getLayerByVolume(G4VPhysicalVolume* volume)const;
-
-     Sensor* getSensorByVolume(G4VPhysicalVolume* volume);
-
-     const Sensor* getSensorByVolume(G4VPhysicalVolume* volume)const;
-
-    void printSensorEnergies()const;
-
-    bool isAssigned()const{
-        if(layers.size() == 0){
+    bool isAssigned() const {
+        if (layers.empty()) {
             return false;
-        }
-        else{
+        } else {
             return layers[0].physicalVolume != nullptr;
         }
     }
 
-    void unAssign(){
-        for(auto& layer: layers){
+    void unAssign() {
+        for (auto& layer : layers) {
             layer.unAssign();
         }
     }
 
-    bool isEmpty()const{
-        return layers.size() == 0;
+    bool isEmpty() const {
+        return layers.empty();
     }
 
-    void setG4System(G4System* g4system){
+    void setG4System(G4System* g4system) {
         this->g4system = g4system;
     }
-private:
+
+    pybind11::tuple __getstate__() const {
+        pybind11::list layer_list;
+        for (const auto& layer : layers) {
+            layer_list.append(layer.__getstate__());
+        }
+        return pybind11::make_tuple(xywidth, layer_list);
+    }
+
+    static GeometryDescriptor __setstate__(pybind11::tuple t) {
+        if (t.size() != 2) throw std::runtime_error("Invalid state!");
+        GeometryDescriptor geom;
+        geom.xywidth = t[0].cast<double>();
+
+        pybind11::list layer_list = t[1].cast<pybind11::list>();
+        for (auto item : layer_list) {
+            geom.layers.push_back(Layer::__setstate__(item.cast<pybind11::tuple>()));
+        }
+        geom.g4system = nullptr; // Reset pointer
+        return geom;
+    }
+
+//private:
     double xywidth;
     std::vector<Layer> layers;
-    G4System * g4system;
+    G4System* g4system;
 };
 
-#endif
+#endif

src/GeometryDescriptor.cc

@@ -4,14 +4,6 @@
 #include "G4VPhysicalVolume.hh"
 #include "G4System.hh"
 
-Layer::Layer(){
-    thickness = 0;
-    material = "";
-    nx = 1;
-    ny = 1;
-    isActive = false;
-    physicalVolume = nullptr;
-}
 
 void Layer::setThickness(double thickness){
     this->thickness = thickness;
@@ -37,10 +29,6 @@ void Layer::assignPhysicalVolume(G4VPhysicalVolume* physicalVolume){
     this->physicalVolume = physicalVolume;
 }
 
-GeometryDescriptor::GeometryDescriptor(double xy_width){
-    xywidth = xy_width;
-}
-
 GeometryDescriptor::~GeometryDescriptor() {
     if(g4system != nullptr){
         if(this == g4system->assigned_cw){ //unassign
@@ -68,18 +56,6 @@ void GeometryDescriptor::addLayer(double thickness, std::string material, bool i
     layers.push_back(layer);
 }
 
-const std::vector<Layer> & GeometryDescriptor::getLayers() const{
-    return layers;
-}
-
-std::vector<Layer> & GeometryDescriptor::getLayers(){
-    return layers;
-}
-
-
-double GeometryDescriptor::getXYWidth() const{
-       return xywidth;
-}
 
 void GeometryDescriptor::resetSensorEnergies()const {
     for(const auto & layer : layers){