Jan Kieseler
94 lines
2.9 KiB
Python
Executable File
94 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
from minicalo import GeometryDescriptor
|
|
from minicalo import G4System
|
|
from minicalo_tools import TempFileManager
|
|
import uproot
|
|
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,
|
|
particleSpec ,
|
|
minEnergy_GeV: float,
|
|
maxEnergy_GeV: float,
|
|
temp_filename: str,
|
|
seed: int
|
|
):
|
|
|
|
if maxEnergy_GeV < 0:
|
|
maxEnergy_GeV = minEnergy_GeV
|
|
|
|
#check if particleSpec is a list of strings or a single string and assert
|
|
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"
|
|
if isinstance(particleSpec, str):
|
|
particleSpec = [particleSpec]
|
|
|
|
# filename without file ending(!)
|
|
if not temp_filename.endswith(".root"):
|
|
filename = temp_filename+ ".root"
|
|
else:
|
|
filename = temp_filename
|
|
|
|
_G4System = G4System() # create instance of G4System only here
|
|
_G4System.init(geometry, seed)
|
|
|
|
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 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
|
|
except Exception as e:
|
|
raise e
|
|
finally:
|
|
os.remove(filename)
|
|
|
|
return {'df':df, 'geometry':geometry}
|
|
|
|
if __name__ == '__main__':
|
|
# do stuff
|
|
# get the first two arguments that correspond to input and output file
|
|
import sys
|
|
assert len(sys.argv) == 3, "Please provide input and output file"
|
|
|
|
io = TempFileManager.slave_init(sys.argv[1], sys.argv[2])
|
|
input_data = io.read_input()
|
|
output_data = run_batch(**input_data)
|
|
io.dump_output(output_data)
|