changed to slim pandas wrapper

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')
+    print('G4Calo: simulation finished after {:.2f} seconds, concatenating outputs'.format(sw.elapsed()))
-
+    sw.reset()
-    alldf = pd.concat(dfs)
+    sw.start()
-    alldf.reset_index(drop=True, inplace=True)
+    df = None
-
+    try:
-    if len(filename):
+        df = _assemble_results_to_mini_df(rp)
-        alldf.to_pickle(filename)
+    finally: #make sure to delete temp files
-    
+        for f in tmpfile:
-    return alldf
+            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

@@ -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")
+            raise RuntimeError(f"Expected exactly one ROOT file, got {len(file)} - check multithreading. Matches: {file}")
-        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
     except Exception as e:
-        raise e
+        #clean up output file 
-    finally:
         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

@@ -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)
+    df = run_batch(gd,10, 'gamma', 40)
-    #display_event(gd,"gamma", 2)
+    display_event(gd,"gamma", 2, outfile='event.html')

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"]