Jan Kieseler
526 lines
17 KiB
Python
526 lines
17 KiB
Python
import threading
|
|
import time
|
|
import warnings
|
|
import contextlib
|
|
import sys
|
|
from minicalo_tools import TempFileManager
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
'''
|
|
give G4Calo package an executable that takes input and output json files names
|
|
|
|
This calls the actual Geant stuff.
|
|
The input are all the arguments and a geometry descriptor
|
|
The output is the data frame as well as the changed geometry descriptor
|
|
|
|
The user only interacts with a wrapper that calls this executable through a system call
|
|
'''
|
|
|
|
|
|
from minicalo import GeometryDescriptor
|
|
|
|
import plotly.graph_objects as go
|
|
import pandas as pd
|
|
import numpy as np
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import awkward as ak
|
|
import uproot
|
|
import glob
|
|
|
|
import pyarrow as pa
|
|
|
|
from IPython.display import Image, display
|
|
from minipandas import MiniFrame
|
|
|
|
|
|
#### helpers #####
|
|
|
|
class Stopwatch(object):
|
|
def __init__(self):
|
|
self.start_time = None
|
|
|
|
def start(self):
|
|
self.start_time = time.time()
|
|
def print_elapsed(self):
|
|
elapsed_time = self.elapsed()
|
|
print(f"Elapsed time: {elapsed_time:.2f} seconds")
|
|
return elapsed_time
|
|
|
|
def elapsed(self):
|
|
if self.start_time is None:
|
|
raise RuntimeError("Stopwatch has not been started.")
|
|
self.end_time = time.time()
|
|
elapsed_time = self.end_time - self.start_time
|
|
return elapsed_time
|
|
def reset(self):
|
|
self.start_time = None
|
|
|
|
def _inspect_schema(first_path: str, tree_name: str = "Hits"):
|
|
"""
|
|
Open one ROOT file and infer:
|
|
- scalar_cols: list[str]
|
|
- array_cols: list[str]
|
|
- array_sizes: dict[col, int]
|
|
- dtypes_scalar: dict[col, np.dtype]
|
|
- dtypes_array: dict[col, np.dtype]
|
|
- n0: number of entries in this file
|
|
"""
|
|
with uproot.open(first_path) as f:
|
|
t = f[tree_name]
|
|
arr0 = t.arrays(library="ak")
|
|
|
|
fields = ak.fields(arr0)
|
|
|
|
def _ndim_from_type(a):
|
|
"""recursively go into the array until we have a scalar.
|
|
This works since we know the arrays are regular"""
|
|
if isinstance(a, ak.Array):
|
|
if(len(a) == 0):
|
|
raise RuntimeError("Cannot inspect schema, empty array encountered")
|
|
return 1 + _ndim_from_type(a[0])
|
|
else:
|
|
return 0
|
|
|
|
# list-like vs scalar by type inspection (no heavy checks)
|
|
array_cols = [c for c in fields if _ndim_from_type(arr0[c]) > 1]
|
|
scalar_cols = [c for c in fields if _ndim_from_type(arr0[c]) == 1]
|
|
|
|
# fixed sizes (assumed regular) + dtypes
|
|
array_sizes = {c: int(ak.num(arr0[c], axis=1)[0]) for c in array_cols}
|
|
dtypes_scalar = {c: np.asarray(arr0[c]).dtype for c in scalar_cols}
|
|
dtypes_array = {c: ak.to_numpy(arr0[c]).dtype for c in array_cols}
|
|
|
|
n0 = len(arr0)
|
|
return scalar_cols, array_cols, array_sizes, dtypes_scalar, dtypes_array, n0
|
|
|
|
def _count_total_entries(root_paths, tree_name: str = "Hits") -> int:
|
|
total = 0
|
|
for p in root_paths:
|
|
with uproot.open(p) as f:
|
|
total += int(f[tree_name].num_entries)
|
|
return total
|
|
|
|
def _assemble_results_to_mini_df(root_paths, tree_name: str = "Hits"):
|
|
"""
|
|
Returns:
|
|
scalars: dict[col, np.ndarray] # shape (N,)
|
|
arrays: dict[col, np.ndarray] # shape (N, M[col])
|
|
"""
|
|
# 1) learn schema & sizes from first file
|
|
scalar_cols, array_cols, array_sizes, dtypes_scalar, dtypes_array, _ = \
|
|
_inspect_schema(root_paths[0], tree_name)
|
|
|
|
# 2) total events
|
|
total_N = _count_total_entries(root_paths, tree_name)
|
|
|
|
# 3) preallocate
|
|
scalars = {c: np.empty(total_N, dtype=dtypes_scalar[c]) for c in scalar_cols}
|
|
arrays = {c: np.empty((total_N, array_sizes[c]), dtype=dtypes_array[c]) for c in array_cols}
|
|
|
|
# 4) fill by slices
|
|
pos = 0
|
|
for p in root_paths:
|
|
with uproot.open(p) as f:
|
|
arr = f[tree_name].arrays(library="ak")
|
|
n = len(arr)
|
|
sl = slice(pos, pos + n)
|
|
|
|
for c in scalar_cols:
|
|
scalars[c][sl] = np.asarray(arr[c])
|
|
for c in array_cols:
|
|
arrays[c][sl, :] = ak.to_numpy(arr[c]) # (n, fixed_size)
|
|
|
|
pos += n
|
|
|
|
#merge the dicts
|
|
merged = {**scalars, **arrays}
|
|
return MiniFrame(merged)
|
|
|
|
|
|
def _run_mini_batch(
|
|
cw : GeometryDescriptor,
|
|
nEvents: int,
|
|
particleSpec ,
|
|
minEnergy_GeV: float,
|
|
maxEnergy_GeV: float = -1.0,
|
|
seed : int = 0,
|
|
tmp_fname : str = "",
|
|
silent : bool = True,
|
|
return_geometry : bool = False):
|
|
|
|
to_process = {'geometry': cw,
|
|
'nEvents': nEvents,
|
|
'particleSpec': particleSpec,
|
|
'minEnergy_GeV': minEnergy_GeV,
|
|
'maxEnergy_GeV': maxEnergy_GeV,
|
|
'temp_filename': tmp_fname,
|
|
'seed': seed}
|
|
|
|
|
|
path = "/work/jkiesele/minicalosim/bind/" #DEBUG
|
|
# path = ""
|
|
|
|
with TempFileManager() as tfm:
|
|
tfm.dump_input(to_process)
|
|
if silent:
|
|
subprocess.run([path+"G4Calo_exec.py", tfm.input_file, tfm.output_file],
|
|
stdout=subprocess.DEVNULL)#, stderr=subprocess.DEVNULL)
|
|
else:
|
|
subprocess.run([path+"G4Calo_exec.py", tfm.input_file, tfm.output_file])
|
|
data = tfm.read_output()
|
|
if return_geometry:
|
|
return data['root_path'], data['geometry']
|
|
else:
|
|
return data['root_path']
|
|
|
|
def _run_mini_batch_wrapper(all_args):
|
|
return _run_mini_batch(*all_args)
|
|
|
|
def create_temp_root_fname():
|
|
#check if we can put the temporary files in /dev/shm, otherwise use /tmp
|
|
if os.path.exists('/dev/shm'):
|
|
tmpdir = '/dev/shm'
|
|
else:
|
|
tmpdir = '/tmp'
|
|
#create a temp file name
|
|
tmpfilename = tempfile.NamedTemporaryFile(delete=False, dir=tmpdir, suffix=".root").name
|
|
#delete file as we only need the name
|
|
os.remove(tmpfilename)
|
|
return tmpfilename
|
|
|
|
def run_batch(
|
|
gd : GeometryDescriptor,
|
|
nEvents: int,
|
|
particleSpec ,
|
|
minEnergy_GeV: float,
|
|
maxEnergy_GeV: float = -1.0,
|
|
filename: str = "",
|
|
no_mp: bool = False,
|
|
manual_seed : int = -1):
|
|
'''
|
|
splits the batch in jobs depending on how many cores are available and runs mini batches in parallel
|
|
this will create threads.
|
|
If no_mp is set to True, it will run the batch in a single thread
|
|
The seeding is not thread safe (unless provided by hand)
|
|
'''
|
|
assert nEvents > 0
|
|
assert minEnergy_GeV > 0
|
|
|
|
if not isinstance(particleSpec, str):
|
|
assert isinstance(particleSpec, list), "particleSpec must be a string or a list of strings"
|
|
for p in particleSpec:
|
|
assert isinstance(p, str), "particleSpec must be a string or a list of strings"
|
|
|
|
sw = Stopwatch()
|
|
sw.start()
|
|
|
|
nCores = os.cpu_count()
|
|
#make sure to adjust cores such that at least 80 events are run per core
|
|
nCores = min(nCores, nEvents // 80 + 1)
|
|
|
|
if no_mp:
|
|
nCores = 1
|
|
|
|
print(f"G4Calo: 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)]
|
|
|
|
if manual_seed >= 0:
|
|
seed = manual_seed
|
|
else:
|
|
if os.path.exists(os.path.expanduser("~/.g4calo_seeds.txt")):
|
|
with open(os.path.expanduser("~/.g4calo_seeds.txt"), "r") as f:
|
|
seed = int(f.read())
|
|
else:
|
|
seed = 1
|
|
with open(os.path.expanduser("~/.g4calo_seeds.txt"), "w") as f:
|
|
f.write(str(seed + nCores))
|
|
|
|
print(f"G4Calo: Batch seed: {seed}")
|
|
|
|
|
|
|
|
#create a list of temporary file names, one for each core, no full path, just the names.
|
|
tmpfile = [f"{create_temp_root_fname()}_g4calo_{i}.root" for i in range(nCores)]
|
|
|
|
args = [(gd,
|
|
nevents[i],
|
|
particleSpec,
|
|
minEnergy_GeV,
|
|
maxEnergy_GeV,
|
|
seed+i,
|
|
tmpfile[i],
|
|
True,
|
|
False) for i in range(nCores)]
|
|
|
|
try:
|
|
if no_mp:
|
|
rp = []
|
|
for i in range(nCores):
|
|
rp.append(_run_mini_batch(*args[i]))
|
|
else:
|
|
|
|
with ThreadPoolExecutor(max_workers=nCores) as executor:
|
|
rp = list(executor.map(_run_mini_batch_wrapper, args))
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
for f in tmpfile:
|
|
if os.path.exists(f):
|
|
os.remove(f)
|
|
raise e
|
|
|
|
print('G4Calo: simulation finished after {:.2f} seconds, concatenating outputs'.format(sw.elapsed()))
|
|
sw.reset()
|
|
sw.start()
|
|
df = None
|
|
try:
|
|
df = _assemble_results_to_mini_df(rp)
|
|
finally: #make sure to delete temp files
|
|
for f in tmpfile:
|
|
if os.path.exists(f):
|
|
os.remove(f)
|
|
print('G4Calo: concatenation finished after {:.2f} seconds'.format(sw.elapsed()))
|
|
return df
|
|
|
|
def _fill_event(gd : GeometryDescriptor,
|
|
particleSpec,
|
|
energy: float,
|
|
seed: int = -1):
|
|
|
|
return _run_mini_batch(gd, 1, particleSpec, energy, return_geometry=True, seed=seed, tmp_fname=create_temp_root_fname())[1]
|
|
|
|
|
|
#
|
|
## below only for visualization
|
|
# 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 display_event(gd : GeometryDescriptor,
|
|
particleSpec,
|
|
energy: float,
|
|
logE = False, renderer=None, seed = -1, outfile : str = ""):
|
|
|
|
gd = _fill_event(gd, particleSpec, energy)
|
|
|
|
print('creating plot')
|
|
#use gd to plot
|
|
# for loop over all layers
|
|
to_plot = []
|
|
material_dict = {}
|
|
z0=0
|
|
# sum up total deposited energy, this is just a normalization factor
|
|
total_dep_energy = 0
|
|
max_dep_energy = 0
|
|
for layer in gd.getLayers():
|
|
for sensor in layer.sensors:
|
|
total_dep_energy += sensor.getEnergy()
|
|
max_dep_energy = max(max_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
|
|
if max_dep_energy == 0:
|
|
print("No energy deposited in any sensor!")
|
|
max_dep_energy = 10**-8
|
|
|
|
# loop over layers, invert order
|
|
for layer in gd.getLayers()[::-1]:
|
|
|
|
layer_width = layer.nx * layer.sens_xwidth * 10. # in mm
|
|
|
|
#
|
|
# plot layers
|
|
#
|
|
layer_hx = layer_width / 2.
|
|
layer_hy = layer_width / 2.
|
|
layer_dz = layer.thickness * 10. # in mm
|
|
layer_material = layer.material
|
|
z0 = layer.getZ()
|
|
|
|
# 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': min(0.2 * 20. / len(gd.getLayers()) + 1e-3, 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_dz, layer_dz, layer_dz, layer_dz]) + z0,
|
|
y = (layer_dz / 2) * np.array([-1, -1, -1, -1, 1, 1, 1, 1]) + 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 / max_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 = (layer_dz / 2) * np.array([-1, -1, -1, -1, 1, 1, 1, 1]) + z0,
|
|
#y = np.array([0, 0, 0, 0, z, z, z, z]) + z0,
|
|
**ijk_cube,
|
|
flatshading=True,
|
|
color='black',
|
|
name='Sensor',
|
|
opacity= use_energy, #max(0.03, use_energy),
|
|
showlegend=False,
|
|
))
|
|
|
|
|
|
# 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='Energy/max(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, gd.getLayers()[0].getZ()-20.]
|
|
end_point = [0, 0, gd.getLayers()[0].getZ()-10.]
|
|
# 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,
|
|
)
|
|
if (not isinstance(particleSpec, list)) or len(particleSpec) == 1:
|
|
# 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=[end_point[0]],
|
|
z=[end_point[1]],
|
|
y=[end_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))
|
|
|
|
#rotate standard view point
|
|
# use layer_width as it is constant for all layers and the largest dimension
|
|
fig.update_layout(scene_camera=dict(eye=dict(x=1 * 1.5, y=-1.5, z=1 * 1.5)))
|
|
#name the axes in the HEP way, so y and z switch names
|
|
fig.update_layout(scene=dict(xaxis_title='x [mm]', yaxis_title='z [mm]', zaxis_title='y [mm]'))
|
|
|
|
if len(outfile) > 0:
|
|
fig.write_html(outfile)
|
|
|
|
else:
|
|
if renderer is not None:
|
|
fig.show(renderer=renderer)
|
|
else:
|
|
fig.show()
|
|
|
|
|
|
|
|
def to_numpy(df_column, dtype=None):
|
|
return np.array(df_column.tolist(), dtype=dtype) |