Jan Kieseler
529 lines
17 KiB
Python
529 lines
17 KiB
Python
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
|
|
|
|
import plotly.graph_objects as go
|
|
import pandas as pd
|
|
import numpy as np
|
|
import os
|
|
import glob
|
|
import subprocess
|
|
import uproot
|
|
import time
|
|
|
|
from IPython.display import Image, display
|
|
|
|
|
|
class __G4System(_G4System):
|
|
|
|
|
|
def run_batch(
|
|
self,
|
|
nEvents: int,
|
|
particleSpec: str,
|
|
minEnergy_GeV: float,
|
|
maxEnergy_GeV: float = -1.0,
|
|
filename: str = "",
|
|
):
|
|
|
|
if maxEnergy_GeV < 0:
|
|
maxEnergy_GeV = minEnergy_GeV
|
|
|
|
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)
|
|
|
|
# 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"))
|
|
print(file)
|
|
if len(file)>1:
|
|
raise Exception("More than one file found! - Check mulithreading")
|
|
else:
|
|
filename = file[0]
|
|
|
|
# conversion from root to pandas dataframe
|
|
|
|
# TO FIX: circumvent index out of bounds error
|
|
ttree = uproot.open(filename)
|
|
try:
|
|
df = ttree["Hits;1"].arrays(library="pd")
|
|
except:
|
|
df = index_out_of_bounds_workaround(ttree["Hits;1"])
|
|
|
|
# save to pickle and delete root file
|
|
if save_file:
|
|
df.to_pickle(filename.replace(".root", ".pkl"))
|
|
os.system("rm "+filename)
|
|
|
|
return df
|
|
|
|
|
|
def displayEvent(self, logE = False, renderer=None):
|
|
# for loop over all layers
|
|
to_plot = []
|
|
material_dict = {}
|
|
z0=0
|
|
|
|
# sum up total deposited energy
|
|
total_dep_energy = 0
|
|
for layer in self.getGeometryDescriptor().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 self.getGeometryDescriptor().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()
|
|
|
|
|
|
|
|
|
|
|
|
#
|
|
# 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 = -square_size/2 + (i + 0.5) * sensor_size
|
|
y_center = -square_size/2 + (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',
|
|
}
|
|
|
|
def index_out_of_bounds_workaround(tbranch):
|
|
# ugly workaround for index out of bounds error
|
|
dfs = []
|
|
entries = tbranch['sensor_energy'].num_entries
|
|
|
|
# check each entry
|
|
for entry in range(entries):
|
|
try:
|
|
df = tbranch.arrays(entry_start=entry,entry_stop=entry+1, library='pd')
|
|
dfs.append(df)
|
|
except:
|
|
continue
|
|
|
|
df = pd.concat(dfs)
|
|
return df.reset_index(drop=True)
|
|
|
|
|
|
|
|
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}')
|
|
|
|
#this is now encapsuled
|
|
from G4Calo import __G4System
|
|
G4System = __G4System()
|
|
G4System.init(cw)
|
|
|
|
df = G4System.run_batch(nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV,"")
|
|
return df
|
|
|
|
|
|
|
|
def run_batch(
|
|
gd : GeometryDescriptor,
|
|
nEvents: int,
|
|
particleSpec: str,
|
|
minEnergy_GeV: float,
|
|
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
|
|
'''
|
|
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, [(gd, nevents[i], particleSpec, minEnergy_GeV, maxEnergy_GeV, i) for i in range(nCores)])
|
|
|
|
return pd.concat(dfs)
|
|
|
|
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() |