lars/minicalosim
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add bindings & plotting

Browse Source
This commit is contained in:
lars
2023-10-04 13:36:57 +02:00
parent 417053a583
commit f6db42ce09
8 changed files with 1270 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+1 -1
View File
@@ -23,7 +23,7 @@ RUN apt-get install -y dpkg-dev cmake g++ gcc binutils libx11-dev libxpm-dev lib
#RUN python3 --version && python3 -m ensurepip
RUN python3 -m pip install --upgrade pip
RUN python3 -m pip install pandas numpy matplotlib MarkupSafe wandb uproot setuptools awkward-pandas
RUN python3 -m pip install pandas numpy matplotlib MarkupSafe wandb uproot setuptools awkward-pandas plotly
RUN python3 -m pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# # # # GEANT
bind/G4Calo.py
+166 -22
View File
@@ -1,6 +1,8 @@
from minicalo import ConstructionWrapper
from minicalo import G4System as _G4System
import plotly.graph_objects as go
import numpy as np
import os
import subprocess
import uproot
@@ -34,25 +36,167 @@ class G4System(_G4System):
df = ttree["Hits;1"].arrays(library="pd")
return df
def displayEvent(self):
f_to_conv = max(
filter(lambda x: x.endswith(".prim"), os.listdir()), key=os.path.getctime
)
# Convert the .prim file to an eps graphic.
subprocess.run(["dawn", "-d", f_to_conv], stderr=subprocess.DEVNULL)
# Convert the eps graphic to png graphic.
subprocess.run(
[
"gs",
"-DEPSCrop",
"-dSAFER",
"-sDEVICE=png256",
"-r600",
"-o",
"event_raw.png",
f_to_conv.replace(".prim", ".eps"),
],
stdout=subprocess.DEVNULL,
)
subprocess.run(["convert", "event_raw.png", "-trim", "event.png"])
display(Image("event.png", width=500))
def displayEvent(self, particleSpec, minEnergy_GeV, maxEnergy_GeV=-1, sensor_width=np.array([50])):
# for loop over all layers
event = self.run_batch(1, particleSpec, minEnergy_GeV, maxEnergy_GeV)
to_plot = []
material_dict = {}
z0=0
# loop over materials
for layer_i, layer in enumerate(reversed(self.cw.getLayers())):
layer_i = len(self.cw.getLayers()) - layer_i -1
#
# plot layers
#
layer_hx = sensor_width / 2.
layer_hy = sensor_width / 2.
layer_z = layer.thickness
layer_material = layer.material
# add material to materials if it is not already in there
if layer_material not in material_dict.keys():
material_dict[layer_material] = {'name': layer_material,
'color': col_dict[layer_material],
'showlegend': False,
'flatshading': True,
'opacity': 0.2}
# add legend entry
to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0],
color=material_dict[layer_material]['color'],
showlegend=True, name=layer_material))
to_plot.append(go.Mesh3d(
# 8 vertices of a cube
x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * layer_hx,
y = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * layer_hy,
z = np.array([0, 0, 0, 0, -layer_z, -layer_z, -layer_z, -layer_z]) + z0,
**ijk_cube,
**material_dict[layer_material]
))
z0 += layer_z
#
# sensors
#
# if there are sensors in the current layer, add them
if layer_i in event['sensor_layer'].to_numpy():
is_in_layer = event['sensor_layer'].to_numpy() == layer_i
z = event['sensor_dz'].to_numpy()[is_in_layer]
n_sensors = len(z)
z=z[0]
xy_centers, hwidth = calculate_sensor_centers(n_sensors, sensor_width)
# loop over all sensors in current layer and add them to plot
for (x_center, y_center), energy in zip(xy_centers, event['sensor_energy'].to_numpy()[is_in_layer]):
to_plot.append(go.Mesh3d(
# 8 vertices of a cube
x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * hwidth + x_center,
y = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * hwidth + y_center,
z = np.array([0, 0, 0, 0, z, z, z, z]) + z0,
**ijk_cube,
flatshading=True,
color='black',
name='Sensor',
opacity= max(0.03, float(energy / event['total_dep_energy'].to_numpy())),
showlegend=False,
))
z0 += z
# add legend entry for sensors
to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0],
color='black', showlegend=True, name='Sensors'))
# add black-white colorbar for sensor hits
to_plot.append(go.Surface(
z=[[0, 0], [0, 0]],
colorscale=[[0, 'white'], [1, 'black']],
showscale=True,
cmin=0,
cmax=1,
colorbar=dict(
title='Fraction of total deposited Energy',
tickvals=[0, 1],
ticktext=['0', '1'],
ticks='outside',
ticklen=10,
),
))
#
# add red arrow for incoming particle
#
to_plot.append(
go.Scatter3d(
x=[0, 0],
y=[0, 0],
z=[-10, -2],
mode='lines+text',
line=dict(color='red', width=3), # You can change the color and width of the arrow
text=['Incoming ' + particleSpec],
textposition='bottom center',
hoverinfo='text',
showlegend=False,
))
#
# finally show plot
#
fig = go.Figure(data=[
*to_plot
])
fig.update_layout(legend=dict(x=0))
# add legend
fig.show()
#
# some helpers
#
def calculate_sensor_centers(X, square_size):
sensor_size = square_size / np.sqrt(X)
hsensor_size = sensor_size / 2.
centers = []
for j in range(int(np.sqrt(X))):
for i in range(int(np.sqrt(X))):
x_center = -25 + (i + 0.5) * sensor_size
y_center = -25 + (j + 0.5) * sensor_size
centers.append((x_center[0], y_center[0]))
return centers, hsensor_size[0]
ijk_cube = {
"i": [7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2],
"j": [3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3],
"k": [0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6],
}
col_dict = {
# active materials
"G4_POLYSTYRENE": 'red',
"G4_PLASTIC_SC_VINYLTOLUENE": 'blue',
"G4_BGO": 'green',
"G4_LSO": 'yellow',
"G4_LYSO": 'orange',
"G4_CESIUM_IODIDE": 'purple',
"G4_PbWO4": 'pink',
"G4_Si": 'brown',
# passive materials
"G4_Pb": 'darkgrey',
"G4_Fe": 'lightgrey',
"G4_W": 'grey',
"G4_Cu": 'dimgray',
"G4_BRASS": 'slategrey',
}
bind/bindings.cpp
+46 -8
View File
@@ -8,27 +8,65 @@
namespace py = pybind11;
template<class M>
void makeConstructionWrapper(M & m, std::string name){
py::class_<ConstructionWrapper>(m, name.data()).def(py::init())
.def("addLayer", &ConstructionWrapper::addLayer, py::arg("thickness"), py::arg("material"), py::arg("isActive")=true, py::arg("nx")=1, py::arg("ny")=-1)
.def("getXYWidth", &ConstructionWrapper::getXYWidth)
// bind overloaded getLayers function
.def("getLayers", (std::vector<Layer> & (ConstructionWrapper::*)()) &ConstructionWrapper::getLayers)
.def("getLayers", (const std::vector<Layer> & (ConstructionWrapper::*)() const) &ConstructionWrapper::getLayers)
.def("getNSensors", &ConstructionWrapper::getNSensors);
}
template<class M>
void makeG4System(M& m, std::string name){
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_gui", &G4System::run_gui)
.def("run_batch", &G4System::run_batch, py::arg("nEvents"), py::arg("partSpecies"), py::arg("minEnergy_GeV"), py::arg("maxEnergy_GeV"))
.def("run_visualize", &G4System::run_visualize, py::arg("partSpecies"), py::arg("minEnergy_GeV"), py::arg("maxEnergy_GeV"))
.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"))
.def("applyUICommand", &G4System::applyUICommand, py::arg("command"))
.def("printMaterial", &G4System::printMaterial, py::arg("name"));
.def("displayEvent", &G4System::displayEvent)
.def("printMaterial", &G4System::printMaterial, py::arg("name"))
.def("check", &G4System::check)
.def_readwrite("cw", &G4System::cw);
}
PYBIND11_MODULE(minicalo, m) {
// 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);
}
PYBIND11_MODULE(minicalo, m)
{
m.doc() = "pybind11 plugin"; // optional module docstring
makeConstructionWrapper(m, "ConstructionWrapper");
makeG4System(m, "G4System");
makeSensor(m, "Sensor");
makeLayer(m, "Layer");
}
bind/example.ipynb
+1045
View File
File diff suppressed because one or more lines are too long
include/ConstructionWrapper.hh
+3 -3
View File
@@ -68,9 +68,9 @@ public:
std::vector<Layer>& getLayers() ;
const std::vector<Layer>& getLayers() const;
double getXYWidth() const{
return xywidth;
}
double getXYWidth() const;
void resetSensorEnergies()const;//energies are mutable
int getNSensors()const{
int n_sensors = 0;
include/G4System.hh
+2
View File
@@ -43,6 +43,8 @@ void applyUICommand(const std::string& command){
UImanager->ApplyCommand(command);
}
ConstructionWrapper cw;
void displayEvent()const{}; //just a placeholder, this will be implemented in python
void printMaterial(const std::string& name)const;
src/ConstructionWrapper.cc
+5
View File
@@ -68,6 +68,11 @@ std::vector<Layer> & ConstructionWrapper::getLayers(){
return layers;
}
double ConstructionWrapper::getXYWidth() const{
return xywidth;
}
void ConstructionWrapper::resetSensorEnergies()const {
for(const auto & layer : layers){
for(const auto & sensor : layer.sensors){
src/G4System.cc
+2
View File
@@ -33,6 +33,8 @@ void G4System::init(ConstructionWrapper &CW){
//delete actionInitialization;
}
// save layers of CW to G4System class
G4System::cw = CW;
G4String session;