diff --git a/bind/G4Calo.py b/bind/G4Calo.py index 6325597..47a591c 100644 --- a/bind/G4Calo.py +++ b/bind/G4Calo.py @@ -1,3 +1,14 @@ +import multiprocessing +import warnings + +def set_start_method(): + try: + multiprocessing.set_start_method('spawn', force=True) + except RuntimeError: + warnings.warn("Failed to set start method to 'spawn'. It may have already been set.") + +set_start_method() + from minicalo import GeometryDescriptor from minicalo import G4System as _G4System @@ -30,7 +41,7 @@ class __G4System(_G4System): save_file = len(filename) > 0 # filename without file ending(!) - filename = "_" + str(time.perf_counter_ns()) + ".root" + filename = "._" + str(time.perf_counter_ns()) + ".root" _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename) # TO FIX: Geant4 adds "t" to the filename, circumvent this for one thread, but this is not a good solution @@ -294,8 +305,6 @@ def index_out_of_bounds_workaround(tbranch): return df.reset_index(drop=True) -_s_G4System = __G4System()#singleton instance - def _run_mini_batch( cw : GeometryDescriptor, @@ -311,18 +320,23 @@ def _run_mini_batch( time.sleep(counter/1000) print(f'done sleeping {counter}') - from G4Calo import G4System + #this is now encapsuled + from G4Calo import __G4System + G4System = __G4System() G4System.init(cw) - df = _s_G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"") + df = G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"") return df -def run_batch(nEvents: int, +def run_batch( + gd : GeometryDescriptor, + nEvents: int, particleSpec: str, minEnergy_GeV: float, - maxEnergy_GeV: float = -1.0,): + maxEnergy_GeV: float = -1.0, + filename: str = ""): ''' splits the batch in jobs depending on how many cores are available and runs mini batches in parallel ''' @@ -342,7 +356,174 @@ def run_batch(nEvents: int, #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)]) + dfs = pool.starmap(_run_mini_batch, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, i) for i in range(nCores)]) return pd.concat(dfs) - \ No newline at end of file + +def _fill_event(gd : GeometryDescriptor, + particleSpec: str, + energy: float): + + from G4Calo import __G4System + G4System = __G4System() + G4System.init(gd) + G4System.run_batch(1, particleSpec, energy, energy,"") + return gd + +def display_event(gd : GeometryDescriptor, + particleSpec: str, + energy: float, + logE = False, renderer=None): + + #run _fill_event in forked mode using 1-core multiprocessing to avoid G4 singletons to interfere + with multiprocessing.Pool(1) as pool: + gd = pool.apply(_fill_event, (gd, particleSpec, energy)) + + #use gd to plot + # for loop over all layers + to_plot = [] + material_dict = {} + z0=0 + # sum up total deposited energy + total_dep_energy = 0 + for layer in gd.getLayers(): + for sensor in layer.sensors: + total_dep_energy += sensor.getEnergy() + if total_dep_energy == 0: + print("No energy deposited in calorimeter!") + total_dep_energy = 10**-8 # to avoid division by zero + # loop over materials + for layer in gd.getLayers(): + + layer_width = layer.nx * layer.sens_xwidth * 10. # in mm + + # + # plot layers + # + layer_hx = layer_width / 2. + layer_hy = layer_width / 2. + layer_z = layer.thickness * 10. # in mm + 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, + z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * layer_hy, + y = np.array([0, 0, 0, 0, layer_z, layer_z, layer_z, layer_z]) + z0, + **ijk_cube, + **material_dict[layer_material] + )) + # + # sensors + # + + # if there are sensors in the current layer, add them + if layer.sensors != []: + z = layer.sensors[0].getdz() + corr = 0. #layer.sensors[0].getX() - layer.sensors[0].getdx()/2. + layer_width/2. + # loop over all sensors in current layer and add them to plot + for sensor in layer.sensors: + x_center = sensor.getX() - corr + y_center = sensor.getY() - corr + hwidth = sensor.getdx() /2. + energy = sensor.getEnergy() + use_energy = float(energy / total_dep_energy) + if logE: + raise NotImplementedError + use_energy = np.log(use_energy+1.) # - np.log(total_dep_energy) + to_plot.append(go.Mesh3d( + # 8 vertices of a cube + x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * hwidth + x_center, + z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * hwidth + y_center, + y = np.array([0, 0, 0, 0, z, z, z, z]) + z0, + **ijk_cube, + flatshading=True, + color='black', + name='Sensor', + opacity= max(0.03, use_energy), + showlegend=False, + )) + + + z0 += layer_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]], + x=[[0, 0], [0, 0]], + y=[[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 + # + # Define the start and end points of the line + start_point = [0, 0, - z0*0.1] + end_point = [0, 0, - z0*0.25] + # Create the line trace + line_trace = go.Scatter3d( + x=[start_point[0], end_point[0]], + z=[start_point[1], end_point[1]], + y=[start_point[2], end_point[2]], + mode='lines', + line=dict(color='red', width=5), + name='Incoming particle', + showlegend=True, + ) + # Calculate the direction vector for the arrow + direction_vector = [(end_point[0] - start_point[0]), (end_point[1] - start_point[1]), (end_point[2] - start_point[2])] + # Create the arrowhead at the start point with the opposite direction + arrowhead_trace = go.Cone( + x=[start_point[0]], + z=[start_point[1]], + y=[start_point[2]], + u=[-direction_vector[0]], + w=[-direction_vector[1]], + v=[-direction_vector[2]], + sizemode='scaled', + sizeref=0.8, + showscale=False, + colorscale='Reds', + opacity=1.0, + anchor='tail', + ) + # Create the 3D scatter plot with both traces + to_plot.append(line_trace) + to_plot.append(arrowhead_trace) + # + # finally show plot + # + + fig = go.Figure(data=[ + *to_plot + ]) + fig.update_layout(legend=dict(x=0)) + # add legend + if renderer is not None: + fig.show(renderer=renderer) + else: + fig.show() \ No newline at end of file diff --git a/bind/bindings.cpp b/bind/bindings.cpp index 2c55687..4079dc9 100644 --- a/bind/bindings.cpp +++ b/bind/bindings.cpp @@ -33,8 +33,18 @@ PYBIND11_MODULE(minicalo, m) { .def("setNx", &Layer::setNx) .def("setNy", &Layer::setNy) .def("setIsActive", &Layer::setIsActive) + //direct accessors + .def_readwrite("nx", &Layer::nx) + .def_readwrite("ny", &Layer::ny) + .def_readwrite("sens_xwidth", &Layer::sens_xwidth) + .def_readwrite("sens_ywidth", &Layer::sens_ywidth) + .def_readwrite("thickness", &Layer::thickness) + .def_readwrite("material", &Layer::material) + .def_readwrite("isActive", &Layer::isActive) + .def("assignPhysicalVolume", &Layer::assignPhysicalVolume) .def("unAssign", &Layer::unAssign) + .def_readwrite("sensors", &Layer::sensors) .def(py::pickle( [](const Layer &l) { // __getstate__ return l.__getstate__(); @@ -46,7 +56,8 @@ PYBIND11_MODULE(minicalo, m) { py::class_(m, "GeometryDescriptor") .def(py::init<>()) - .def("addLayer", &GeometryDescriptor::addLayer) + //add these defaults: void addLayer(double thickness_cm, std::string material, bool isActive = true, int nx = 1, int ny = -1) + .def("addLayer", &GeometryDescriptor::addLayer, py::arg("thickness_cm"), py::arg("material"), py::arg("isActive") = true, py::arg("nx") = 1, py::arg("ny") = -1) .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers)) .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers, py::const_)) .def("getXYWidth", &GeometryDescriptor::getXYWidth) diff --git a/bind/example.py b/bind/example.py index 8a427b3..cc4081b 100644 --- a/bind/example.py +++ b/bind/example.py @@ -1,24 +1,21 @@ - -from G4Calo import GeometryDescriptor, G4System +import multiprocessing +from G4Calo import GeometryDescriptor, run_batch import sys -cw = GeometryDescriptor() +if __name__ == '__main__': + #multiprocessing.set_start_method('spawn') - -for _ in range(25): - cw.addLayer(0.5, "G4_Pb", False) - cw.addLayer(1.,"G4_POLYSTYRENE",True,1) - -#directly use G4System as a global singleton -G4System.init(cw) - - -df = G4System.run_batch(10, 'gamma', 1) -cw = GeometryDescriptor() - -for _ in range(5): - cw.addLayer(0.5, "G4_Pb", False) - cw.addLayer(1.,"G4_POLYSTYRENE",True,1) - -G4System.init(cw) -df = G4System.run_batch(10, 'gamma', 1) \ No newline at end of file + gd = GeometryDescriptor() + + for _ in range(25): + gd.addLayer(0.5, "G4_Pb", False) + gd.addLayer(1.,"G4_POLYSTYRENE",True,1) + + df = run_batch(gd, 1000, 'gamma', 1) + gd = GeometryDescriptor() + + for _ in range(5): + gd.addLayer(0.5, "G4_Pb", False) + gd.addLayer(1.,"G4_POLYSTYRENE",True,1) + + df = run_batch(gd,10, 'gamma', 1) \ No newline at end of file