Jan Kieseler
75 lines
2.3 KiB
Python
Executable File
75 lines
2.3 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 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)
|
|
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 RuntimeError(f"Expected exactly one ROOT file, got {len(file)} - check multithreading. Matches: {file}")
|
|
except Exception as e:
|
|
#clean up output file
|
|
os.remove(filename)
|
|
|
|
|
|
# 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
|
|
# 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)
|