""" ============================================================ G4Calo Public Interface Functions ============================================================ This module provides two main functions intended for student use: - run_batch(...) - display_event(...) They allow running Geant4 calorimeter simulations and visualising individual events using a pandas-like, lightweight data structure (MiniFrame) built for teaching and analysis. ------------------------------------------------------------ Function: run_batch(...) ------------------------------------------------------------ Purpose ------- Run a full Geant4 simulation batch with automatic parallelisation and return the results as a lightweight, pandas-like table (MiniFrame). Overview -------- `run_batch` distributes a requested number of Geant4 events across available CPU cores, runs all mini-batches in parallel, collects their temporary ROOT output files, converts them to numpy arrays, and finally assembles everything into a single MiniFrame for analysis or visualisation. Parameters ---------- gd : GeometryDescriptor The calorimeter geometry to simulate. nEvents : int Total number of events to simulate. Must be > 0. particleSpec : str or list[str] Particle type(s) to generate, e.g. "gamma" or ["e-", "pi+"]. minEnergy_GeV : float Minimum kinetic energy in GeV for the primary particles. maxEnergy_GeV : float, optional Maximum kinetic energy in GeV. If negative (default), identical to minEnergy_GeV. filename : str, optional Optional output file name for storing the resulting MiniFrame using pickle format. If empty, no file is written. no_mp : bool, optional If True, disables multiprocessing and runs all events sequentially. Useful for debugging. manual_seed : int, optional Optional random-seed override. If < 0 (default), an internal per-core seed sequence is generated and persisted between runs. Returns ------- MiniFrame A table-like object containing one row per simulated event. Each row stores scalar quantities (e.g. true_energy, total_dep_energy) and fixed-size arrays for each sensor-level variable (sensor_energy, sensor_x, sensor_y, ...). Notes ----- - Automatically chooses the number of cores so that each core processes at least ~80 events. - Temporary ROOT files are written to /dev/shm or /tmp and removed after merging. - Prints timing information for simulation and concatenation steps. - For deterministic results, specify manual_seed. Example ------- >>> df = run_batch(gd, nEvents=1000, particleSpec="gamma", minEnergy_GeV=1.0) >>> print(df.describe()) >>> df.to_pickle("results.pkl") ------------------------------------------------------------ Function: display_event(...) ------------------------------------------------------------ Purpose ------- Visualise the energy deposition of a single simulated event in 3D using Plotly. Overview -------- `display_event` performs a one-event Geant4 simulation with the given geometry, retrieves deposited sensor energies, and renders a 3D view of the calorimeter including materials, layers, sensors, and the incoming particle direction. Parameters ---------- gd : GeometryDescriptor The calorimeter geometry descriptor (modified in-place). particleSpec : str or list[str] Particle type(s) to simulate, e.g. "e-", "pi+", "gamma". energy : float Particle kinetic energy in GeV. logE : bool, optional If True, use logarithmic scaling for sensor opacity (currently not implemented). renderer : str, optional Optional Plotly renderer ("browser", "notebook", "png", ...). If omitted, uses the default interactive renderer. seed : int, optional Optional random seed for reproducibility. outfile : str, optional If provided, saves the figure as an interactive HTML file instead of displaying it. Returns ------- None — the function displays or saves the visualisation. What the Plot Shows ------------------- - Coloured cubes for each material layer (colour legend included) - Black semi-transparent cubes for sensors, opacity ∝ deposited energy - Red arrow showing incoming particle direction - Axes labelled in mm (x, y, z following HEP conventions) - Interactive rotation and zoom supported Example ------- >>> display_event(gd, particleSpec="gamma", energy=1.0, renderer="browser") Typical Workflow ---------------- >>> df = run_batch(gd, 200, "e-", 1.0) >>> print(df.describe()) >>> display_event(gd, "e-", 1.0, renderer="browser") These two functions form the primary interface for student exercises: `run_batch` to simulate, and `display_event` to visualise results. All heavy lifting (parallelisation, I/O, and conversion) is handled internally so that focus remains on physics interpretation. """ import threading import time import warnings import contextlib import sys from minicalo_tools import TempFileManager from concurrent.futures import ThreadPoolExecutor 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 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): """ Run a full Geant4 simulation batch with automatic parallelisation. Distributes the requested number of events across CPU cores, runs each mini-batch in parallel, collects temporary ROOT files, converts them to numpy arrays, and merges everything into a lightweight, pandas-like MiniFrame. Parameters ---------- gd : GeometryDescriptor Calorimeter geometry to simulate. nEvents : int Total number of events to simulate. particleSpec : str or list[str] Particle type(s), e.g. "gamma" or ["e-", "pi+"]. minEnergy_GeV : float Minimum kinetic energy in GeV. maxEnergy_GeV : float, optional Maximum kinetic energy in GeV. If negative, equals minEnergy_GeV. filename : str, optional Optional pickle filename for saving results. no_mp : bool, optional Disable multiprocessing (run single-threaded) for debugging. manual_seed : int, optional Optional random seed. If < 0, seeds are generated automatically. Returns ------- MiniFrame Table-like structure with one row per event, containing both scalar values and fixed-size sensor arrays. Notes ----- - Automatically chooses the number of cores to balance workload. - Temporary files are written to /dev/shm or /tmp and removed after merging. - Prints timing information for both simulation and data assembly. - For reproducibility, set manual_seed. Example ------- >>> df = run_batch(gd, 1000, "gamma", 1.0) >>> print(df.describe()) >>> df.to_pickle("results.pkl") """ 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 = ""): """ Visualise a single simulated event in 3D using Plotly. Runs a one-event Geant4 simulation with the provided geometry, extracts the deposited sensor energies, and renders an interactive 3D view of the calorimeter with materials, layers, and energy deposition. Parameters ---------- gd : GeometryDescriptor The calorimeter geometry (modified in place by the simulation). particleSpec : str or list[str] Particle type(s), e.g. "e-", "pi+", "gamma". energy : float Particle kinetic energy in GeV. logE : bool, optional Use logarithmic opacity scaling (currently not implemented). renderer : str, optional Plotly renderer ("browser", "notebook", etc.). Default: interactive. seed : int, optional Random seed for reproducibility. outfile : str, optional If set, saves the figure as an interactive HTML file. Returns ------- None — displays or saves a 3D Plotly figure. Example ------- >>> display_event(gd, "gamma", 1.0, renderer="browser") The plot shows: - Coloured cubes for material layers - Semi-transparent black cubes for sensors (opacity ∝ deposited energy) - Red arrow for incoming particle direction - Interactive rotation and zoom """ 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)