changed to be serialisable for multiprocessing

This commit is contained in:
Jan Kieseler
2024-07-23 14:01:20 +02:00
parent da97ad4613
commit e849d0bac5
6 changed files with 257 additions and 167 deletions
+67 -68
View File
@@ -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");
}