Skip to content
Snippets Groups Projects
Commit 44400973 authored by Johannes Schumann's avatar Johannes Schumann
Browse files

Merge branch 'multifileoutput' into 'master' 

Multifile output

See merge request !29
parents a76fd88c c45f5705
No related branches found
No related tags found
1 merge request!29Multifile output
Pipeline #26242 passed
Stage: test
Stage: coverage
Stage: doc
Stage: deploy
Hide whitespace changes
Inline Side-by-side
Showing
with 117 additions and 97 deletions
km3buu/output.py
+ 117
− 97
View file @ 44400973
...@@ -608,6 +608,7 @@ class GiBUUOutput: ...@@ -608,6 +608,7 @@ class GiBUUOutput:
def write_detector_file(gibuu_output, def write_detector_file(gibuu_output,
ofile="gibuu.offline.root", ofile="gibuu.offline.root",
no_files=1,
geometry=CanVolume(), geometry=CanVolume(),
livetime=3.156e7, livetime=3.156e7,
propagate_tau=True): # pragma: no cover propagate_tau=True): # pragma: no cover
...@@ -620,6 +621,8 @@ def write_detector_file(gibuu_output, ...@@ -620,6 +621,8 @@ def write_detector_file(gibuu_output,
Output object which wraps the information from the GiBUU output files Output object which wraps the information from the GiBUU output files
ofile: str ofile: str
Output filename Output filename
no_files: int (default: 1)
Number of output files written
geometry: DetectorVolume geometry: DetectorVolume
The detector geometry which should be used The detector geometry which should be used
livetime: float livetime: float
...@@ -628,12 +631,6 @@ def write_detector_file(gibuu_output, ...@@ -628,12 +631,6 @@ def write_detector_file(gibuu_output,
if not isinstance(geometry, DetectorVolume): if not isinstance(geometry, DetectorVolume):
raise TypeError("Geometry needs to be a DetectorVolume") raise TypeError("Geometry needs to be a DetectorVolume")
evt = ROOT.Evt()
outfile = ROOT.TFile.Open(ofile, "RECREATE")
tree = ROOT.TTree("E", "KM3NeT Evt Tree")
tree.Branch("Evt", evt, 32000, 4)
mc_trk_id = 0
def add_particles(particle_data, pos_offset, rotation, start_mc_trk_id, def add_particles(particle_data, pos_offset, rotation, start_mc_trk_id,
timestamp, status): timestamp, status):
nonlocal evt nonlocal evt
...@@ -678,6 +675,10 @@ def write_detector_file(gibuu_output, ...@@ -678,6 +675,10 @@ def write_detector_file(gibuu_output,
sec_lep_type *= -1 sec_lep_type *= -1
event_data = gibuu_output.arrays event_data = gibuu_output.arrays
if no_files > len(event_data):
raise IndexError("More files to write than contained events!")
bjorkenx = event_data.Bx bjorkenx = event_data.Bx
bjorkeny = event_data.By bjorkeny = event_data.By
...@@ -696,6 +697,7 @@ def write_detector_file(gibuu_output, ...@@ -696,6 +697,7 @@ def write_detector_file(gibuu_output,
w2 = gibuu_output.w2weights(geometry.volume, targets_per_volume, 4 * np.pi) w2 = gibuu_output.w2weights(geometry.volume, targets_per_volume, 4 * np.pi)
global_generation_weight = gibuu_output.global_generation_weight(4 * np.pi) global_generation_weight = gibuu_output.global_generation_weight(4 * np.pi)
mean_xsec_func = gibuu_output.mean_xsec
head = ROOT.Head() head = ROOT.Head()
header_dct = EMPTY_KM3NET_HEADER_DICT.copy() header_dct = EMPTY_KM3NET_HEADER_DICT.copy()
...@@ -714,95 +716,113 @@ def write_detector_file(gibuu_output, ...@@ -714,95 +716,113 @@ def write_detector_file(gibuu_output,
version, timestamp.strftime("%Y%m%d %H%M%S")) version, timestamp.strftime("%Y%m%d %H%M%S"))
header_dct["primary"] = "{:d}".format(nu_type) header_dct["primary"] = "{:d}".format(nu_type)
for k, v in header_dct.items(): for i in range(no_files):
head.set_line(k, v) start_id = 0
head.Write("Head") stop_id = len(event_data)
tmp_filename = ofile
mean_xsec_func = gibuu_output.mean_xsec if no_files > 1:
tmp_filename = ofile.replace(".root", ".{}.root".format(i + 1))
for mc_event_id, event in enumerate(event_data): bunch_size = stop_id // no_files
evt.clear() start_id = i * bunch_size
evt.id = mc_event_id stop_id = (i + 1) * bunch_size if i < (no_files - 1) else stop_id
evt.mc_run_id = mc_event_id
# Weights evt = ROOT.Evt()
evt.w.push_back(geometry.volume) #w1 (can volume) outfile = ROOT.TFile.Open(tmp_filename, "RECREATE")
evt.w.push_back(w2[mc_event_id]) #w2 tree = ROOT.TTree("E", "KM3NeT Evt Tree")
evt.w.push_back(-1.0) #w3 (= w2*flux) tree.Branch("Evt", evt, 32000, 4)
# Event Information (w2list) mc_trk_id = 0
evt.w2list.resize(W2LIST_LENGTH)
evt.w2list[W2LIST_LOOKUP["PS"]] = global_generation_weight for k, v in header_dct.items():
evt.w2list[W2LIST_LOOKUP["EG"]] = gibuu_output.flux_index head.set_line(k, v)
evt.w2list[W2LIST_LOOKUP["XSEC_MEAN"]] = mean_xsec_func(event.lepIn_E) head.Write("Head")
evt.w2list[W2LIST_LOOKUP["XSEC"]] = event.xsec
evt.w2list[W2LIST_LOOKUP["TARGETA"]] = gibuu_output.A for mc_event_id, event in enumerate(event_data[start_id:stop_id]):
evt.w2list[W2LIST_LOOKUP["TARGETZ"]] = gibuu_output.Z evt.clear()
evt.w2list[W2LIST_LOOKUP["BX"]] = bjorkenx[mc_event_id] evt.id = mc_event_id
evt.w2list[W2LIST_LOOKUP["BY"]] = bjorkeny[mc_event_id] evt.mc_run_id = mc_event_id
evt.w2list[W2LIST_LOOKUP["CC"]] = ichan # Weights
evt.w2list[W2LIST_LOOKUP["ICHAN"]] = SCATTERING_TYPE_TO_GENIE[ evt.w.push_back(geometry.volume) #w1 (can volume)
event.evType] evt.w.push_back(w2[mc_event_id]) #w2
evt.w2list[W2LIST_LOOKUP["VERINCAN"]] = 1 evt.w.push_back(-1.0) #w3 (= w2*flux)
evt.w2list[W2LIST_LOOKUP["LEPINCAN"]] = 1 # Event Information (w2list)
evt.w2list[W2LIST_LOOKUP["GIBUU_WEIGHT"]] = event.weight evt.w2list.resize(W2LIST_LENGTH)
evt.w2list[W2LIST_LOOKUP["GIBUU_SCAT_TYPE"]] = event.evType evt.w2list[W2LIST_LOOKUP["PS"]] = global_generation_weight
#TODO evt.w2list[W2LIST_LOOKUP["EG"]] = gibuu_output.flux_index
evt.w2list[W2LIST_LOOKUP["DXSEC"]] = np.nan evt.w2list[W2LIST_LOOKUP["XSEC_MEAN"]] = mean_xsec_func(
evt.w2list[W2LIST_LOOKUP["COLUMN_DEPTH"]] = np.nan event.lepIn_E)
evt.w2list[W2LIST_LOOKUP["P_EARTH"]] = np.nan evt.w2list[W2LIST_LOOKUP["XSEC"]] = event.xsec
evt.w2list[W2LIST_LOOKUP["WATER_INT_LEN"]] = np.nan evt.w2list[W2LIST_LOOKUP["TARGETA"]] = gibuu_output.A
evt.w2list[W2LIST_LOOKUP["TARGETZ"]] = gibuu_output.Z
# Vertex Position evt.w2list[W2LIST_LOOKUP["BX"]] = bjorkenx[mc_event_id]
vtx_pos = np.array(geometry.random_pos()) evt.w2list[W2LIST_LOOKUP["BY"]] = bjorkeny[mc_event_id]
# Direction evt.w2list[W2LIST_LOOKUP["CC"]] = ichan
phi = np.random.uniform(0, 2 * np.pi) evt.w2list[W2LIST_LOOKUP["ICHAN"]] = SCATTERING_TYPE_TO_GENIE[
cos_theta = np.random.uniform(-1, 1) event.evType]
sin_theta = np.sqrt(1 - cos_theta**2) evt.w2list[W2LIST_LOOKUP["VERINCAN"]] = 1
evt.w2list[W2LIST_LOOKUP["LEPINCAN"]] = 1
dir_x = np.cos(phi) * sin_theta evt.w2list[W2LIST_LOOKUP["GIBUU_WEIGHT"]] = event.weight
dir_y = np.sin(phi) * sin_theta evt.w2list[W2LIST_LOOKUP["GIBUU_SCAT_TYPE"]] = event.evType
dir_z = cos_theta #TODO
evt.w2list[W2LIST_LOOKUP["DXSEC"]] = np.nan
direction = np.array([dir_x, dir_y, dir_z]) evt.w2list[W2LIST_LOOKUP["COLUMN_DEPTH"]] = np.nan
theta = np.arccos(cos_theta) evt.w2list[W2LIST_LOOKUP["P_EARTH"]] = np.nan
R = Rotation.from_euler("yz", [theta, phi]) evt.w2list[W2LIST_LOOKUP["WATER_INT_LEN"]] = np.nan
timestamp = np.random.uniform(0, livetime) # Vertex Position
vtx_pos = np.array(geometry.random_pos())
nu_in_trk = ROOT.Trk() # Direction
nu_in_trk.id = mc_trk_id phi = np.random.uniform(0, 2 * np.pi)
mc_trk_id += 1 cos_theta = np.random.uniform(-1, 1)
nu_in_trk.mother_id = -1 sin_theta = np.sqrt(1 - cos_theta**2)
nu_in_trk.type = nu_type
nu_in_trk.pos.set(*vtx_pos) dir_x = np.cos(phi) * sin_theta
nu_in_trk.dir.set(*direction) dir_y = np.sin(phi) * sin_theta
nu_in_trk.E = event.lepIn_E dir_z = cos_theta
nu_in_trk.t = timestamp
nu_in_trk.status = PARTICLE_MC_STATUS["TRK_ST_PRIMARYNEUTRINO"] direction = np.array([dir_x, dir_y, dir_z])
evt.mc_trks.push_back(nu_in_trk) theta = np.arccos(cos_theta)
R = Rotation.from_euler("yz", [theta, phi])
if tau_secondaries is not None:
event_tau_sec = tau_secondaries[mc_event_id] timestamp = np.random.uniform(0, livetime)
add_particles(event_tau_sec, vtx_pos, R, mc_trk_id, timestamp,
PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"]) nu_in_trk = ROOT.Trk()
mc_trk_id += len(event_tau_sec.E) nu_in_trk.id = mc_trk_id
else:
lep_out_trk = ROOT.Trk()
lep_out_trk.id = mc_trk_id
mc_trk_id += 1 mc_trk_id += 1
lep_out_trk.mother_id = 0 nu_in_trk.mother_id = -1
lep_out_trk.type = sec_lep_type nu_in_trk.type = nu_type
lep_out_trk.pos.set(*vtx_pos) nu_in_trk.pos.set(*vtx_pos)
mom = np.array([event.lepOut_Px, event.lepOut_Py, event.lepOut_Pz]) nu_in_trk.dir.set(*direction)
p_dir = R.apply(mom / np.linalg.norm(mom)) nu_in_trk.E = event.lepIn_E
lep_out_trk.dir.set(*p_dir) nu_in_trk.t = timestamp
lep_out_trk.E = event.lepOut_E nu_in_trk.status = PARTICLE_MC_STATUS["TRK_ST_PRIMARYNEUTRINO"]
lep_out_trk.t = timestamp evt.mc_trks.push_back(nu_in_trk)
lep_out_trk.status = PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"]
evt.mc_trks.push_back(lep_out_trk) if tau_secondaries is not None:
event_tau_sec = tau_secondaries[mc_event_id]
add_particles(event, vtx_pos, R, mc_trk_id, timestamp, add_particles(event_tau_sec, vtx_pos, R, mc_trk_id, timestamp,
PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"]) PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"])
tree.Fill() mc_trk_id += len(event_tau_sec.E)
outfile.Write() else:
outfile.Close() lep_out_trk = ROOT.Trk()
del outfile lep_out_trk.id = mc_trk_id
mc_trk_id += 1
lep_out_trk.mother_id = 0
lep_out_trk.type = sec_lep_type
lep_out_trk.pos.set(*vtx_pos)
mom = np.array(
[event.lepOut_Px, event.lepOut_Py, event.lepOut_Pz])
p_dir = R.apply(mom / np.linalg.norm(mom))
lep_out_trk.dir.set(*p_dir)
lep_out_trk.E = event.lepOut_E
lep_out_trk.t = timestamp
lep_out_trk.status = PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"]
evt.mc_trks.push_back(lep_out_trk)
add_particles(event, vtx_pos, R, mc_trk_id, timestamp,
PARTICLE_MC_STATUS["TRK_ST_FINALSTATE"])
tree.Fill()
outfile.Write()
outfile.Close()
del outfile
del evt
del tree
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment