lars/minicalosim
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

changed to slim pandas wrapper

Browse Source
This commit is contained in:
Jan Kieseler
2025-10-10 10:38:33 +02:00
parent d1155dc1e2
commit 13a913c068
4 changed files with 161 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bind/G4Calo.py
+137 -19
View File
@@ -1,4 +1,5 @@
import threading
import time
import warnings
import contextlib
import sys
@@ -24,8 +25,118 @@ 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(
@@ -47,20 +158,22 @@ def _run_mini_batch(
'temp_filename': tmp_fname,
'seed': seed}
# silent = False #DEBUG
path = "/work/jkiesele/minicalosim/bind/" #DEBUG
# path = ""
with TempFileManager() as tfm:
tfm.dump_input(to_process)
if silent:
subprocess.run(["G4Calo_exec.py", tfm.input_file, tfm.output_file],
subprocess.run([path+"G4Calo_exec.py", tfm.input_file, tfm.output_file],
stdout=subprocess.DEVNULL)#, stderr=subprocess.DEVNULL)
else:
subprocess.run(["G4Calo_exec.py", tfm.input_file, tfm.output_file])
subprocess.run([path+"G4Calo_exec.py", tfm.input_file, tfm.output_file])
data = tfm.read_output()
if return_geometry:
return data['df'], data['geometry']
return data['root_path'], data['geometry']
else:
return data['df']
return data['root_path']
def _run_mini_batch_wrapper(all_args):
return _run_mini_batch(*all_args)
@@ -100,6 +213,9 @@ def run_batch(
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)
@@ -145,31 +261,33 @@ def run_batch(
try:
if no_mp:
dfs = []
rp = []
for i in range(nCores):
dfs.append(_run_mini_batch(*args[i]))
rp.append(_run_mini_batch(*args[i]))
else:
with ThreadPoolExecutor(max_workers=nCores) as executor:
dfs = list(executor.map(_run_mini_batch_wrapper, args))
rp = list(executor.map(_run_mini_batch_wrapper, args))
except Exception as e:
print(e)
raise e
finally: #make sure to clean up the temporary files in case of an exception
for f in tmpfile:
if os.path.exists(f):
os.remove(f)
raise e
print('G4Calo: simulation finished, concatenating dataframes')
alldf = pd.concat(dfs)
alldf.reset_index(drop=True, inplace=True)
if len(filename):
alldf.to_pickle(filename)
return alldf
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,
bind/G4Calo_exec.py
+15 -34
View File
@@ -9,23 +9,6 @@ import glob
import pandas as pd
import os
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_batch(
geometry: GeometryDescriptor,
nEvents: int,
@@ -58,28 +41,26 @@ def run_batch(
try:
_G4System.run_batch(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"))
if file[0] != filename:
os.replace(file[0], filename)
print(f"G4Calo: output written to {filename}")
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
raise RuntimeError(f"Expected exactly one ROOT file, got {len(file)} - check multithreading. Matches: {file}")
except Exception as e:
raise e
finally:
#clean up output file
os.remove(filename)
return {'df':df, 'geometry':geometry}
# TO FIX: Geant4 adds "t<threadnumber>" to the filename, circumvent this for one thread, but this is not a good solution
#rename to original filename
root_path = filename
return {'root_path':root_path, 'geometry':geometry}
if __name__ == '__main__':
# do stuff
bind/example.py
+6 -4
View File
@@ -10,12 +10,14 @@ if __name__ == '__main__':
gd.addLayer(0.5, "G4_Pb", False)
gd.addLayer(1.,"G4_POLYSTYRENE",True,1)
df = run_batch(gd, 1000, 'gamma', 1)
df = run_batch(gd, 100, 'gamma', 1)
gd = GeometryDescriptor()
for _ in range(5):
for _ in range(30):
gd.addLayer(0.5, "G4_Pb", False)
gd.addLayer(1.,"G4_POLYSTYRENE",True,1)
df = run_batch(gd,10, 'gamma', 1)
#display_event(gd,"gamma", 2)
df = run_batch(gd,10, 'gamma', 40)
display_event(gd,"gamma", 2, outfile='event.html')