751 lines
24 KiB
Python
751 lines
24 KiB
Python
"""
|
|
============================================================
|
|
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,
|
|
return_steps: bool = False):
|
|
"""
|
|
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.
|
|
return_steps : bool, optional
|
|
If True, also return a MiniFrame of step-level data (one row per
|
|
Geant4 step across all events). Returns a tuple (hits, steps).
|
|
|
|
Returns
|
|
-------
|
|
MiniFrame or tuple[MiniFrame, MiniFrame]
|
|
If return_steps is False (default): event-level MiniFrame with one
|
|
row per event. If return_steps is True: (hits, steps) where steps
|
|
has one row per Geant4 step with columns event_id, track_id,
|
|
step_no, pdg, pre_x/y/z, pre_E, post_x/y/z, post_E.
|
|
|
|
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
|
|
steps_df = None
|
|
try:
|
|
df = _assemble_results_to_mini_df(rp)
|
|
if return_steps:
|
|
steps_df = _assemble_results_to_mini_df(rp, tree_name="Steps")
|
|
finally:
|
|
for f in tmpfile:
|
|
if os.path.exists(f):
|
|
os.remove(f)
|
|
print('G4Calo: concatenation finished after {:.2f} seconds'.format(sw.elapsed()))
|
|
if return_steps:
|
|
return df, steps_df
|
|
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) |