diff --git a/bind/G4Calo.py b/bind/G4Calo.py index 77e5b75..c9f842b 100644 --- a/bind/G4Calo.py +++ b/bind/G4Calo.py @@ -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, diff --git a/bind/G4Calo_exec.py b/bind/G4Calo_exec.py index 9c9b678..8f64393 100755 --- a/bind/G4Calo_exec.py +++ b/bind/G4Calo_exec.py @@ -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" 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" 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 diff --git a/bind/example.py b/bind/example.py index 236f827..c95fbff 100644 --- a/bind/example.py +++ b/bind/example.py @@ -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) \ No newline at end of file + df = run_batch(gd,10, 'gamma', 40) + display_event(gd,"gamma", 2, outfile='event.html') \ No newline at end of file diff --git a/docker/Dockerfile b/docker/Dockerfile index 2cb44c1..42a4b32 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -97,7 +97,7 @@ RUN cd /root/minicalosim && git checkout $COMMIT && \ mkdir -p build && cd build && rm -rf * && cmake ../ && make -j4 -RUN cp /root/minicalosim/build/minicalo* /root/minicalosim/bind/G4Calo.py /root/minicalosim/bind/minicalo_tools.py /usr/local/lib/python3.10/dist-packages/ +RUN cp /root/minicalosim/build/minicalo* /root/minicalosim/bind/G4Calo.py /root/minicalosim/bind/minicalo_tools.py /root/minicalosim/bind/minipandas.py /usr/local/lib/python3.10/dist-packages/ RUN cp /root/minicalosim/bind/G4Calo_exec.py /usr/local/bin/ @@ -124,6 +124,8 @@ RUN fix-permissions /etc/jupyter/ # clean up RUN rm -rf /root/minicalosim +RUN pip3 install pyarrow + USER ${NB_UID} CMD ["start-notebook.sh"]