added multi particles

This commit is contained in:
Jan Kieseler
2024-10-09 13:05:18 +02:00
parent b7946e2e37
commit 7410049a72
7 changed files with 184 additions and 125 deletions
+61 -35
View File
@@ -32,7 +32,7 @@ class __G4System(_G4System):
def run_batch(
self,
nEvents: int,
particleSpec: str,
particleSpec ,
minEnergy_GeV: float,
maxEnergy_GeV: float = -1.0,
filename: str = "",
@@ -41,6 +41,17 @@ class __G4System(_G4System):
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]
save_file = len(filename) > 0
# filename without file ending(!)
filename = "._" + str(time.perf_counter_ns()) + ".root"
@@ -141,7 +152,7 @@ def index_out_of_bounds_workaround(tbranch):
def _run_mini_batch(
cw : GeometryDescriptor,
nEvents: int,
particleSpec: str,
particleSpec ,
minEnergy_GeV: float,
maxEnergy_GeV: float = -1.0,
batch_seed : int = 0,
@@ -173,7 +184,7 @@ def init_g4calo_threading_lock():
def run_batch(
gd : GeometryDescriptor,
nEvents: int,
particleSpec: str,
particleSpec ,
minEnergy_GeV: float,
maxEnergy_GeV: float = -1.0,
filename: str = "",
@@ -184,6 +195,11 @@ def run_batch(
assert nEvents > 0
assert minEnergy_GeV > 0
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"
nCores = multiprocessing.cpu_count()
#make sure to adjust cores such that at least 80 events are run per core
nCores = min(nCores, nEvents // 80 + 1)
@@ -230,7 +246,7 @@ def run_batch(
return alldf
def _fill_event(gd : GeometryDescriptor,
particleSpec: str,
particleSpec,
energy: float,
seed: int = -1):
@@ -241,27 +257,34 @@ def _fill_event(gd : GeometryDescriptor,
return gd
def display_event(gd : GeometryDescriptor,
particleSpec: str,
particleSpec,
energy: float,
logE = False, renderer=None, seed = -1):
logE = False, renderer=None, seed = -1, outfile : str = ""):
#run _fill_event in forked mode using 1-core multiprocessing to avoid G4 singletons to interfere
with multiprocessing.Pool(1) as pool:
gd = pool.apply(_fill_event, (gd, particleSpec, energy))
print('creating plot')
#use gd to plot
# for loop over all layers
to_plot = []
material_dict = {}
z0=0
# sum up total deposited energy
# sum up total deposited energy, this is just a normalization factor
total_dep_energy = 0
max_dep_energy = 0
for layer in gd.getLayers():
for sensor in layer.sensors:
total_dep_energy += sensor.getEnergy()
max_dep_energy = max(max_dep_energy, sensor.getEnergy())
if total_dep_energy == 0:
print("No energy deposited in calorimeter!")
total_dep_energy = 10**-8 # to avoid division by zero
if max_dep_energy == 0:
print("No energy deposited in any sensor!")
max_dep_energy = 10**-8
# loop over layers, invert order
for layer in gd.getLayers()[::-1]:
@@ -282,7 +305,7 @@ def display_event(gd : GeometryDescriptor,
'color': col_dict[layer_material],
'showlegend': False,
'flatshading': True,
'opacity': 0.2 * 20. / len(gd.getLayers())}
'opacity': min(0.2 * 20. / len(gd.getLayers()) + 1e-3, 0.2)}
# add legend entry
to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0],
color=material_dict[layer_material]['color'],
@@ -310,7 +333,7 @@ def display_event(gd : GeometryDescriptor,
y_center = sensor.getY() - corr
hwidth = sensor.getdx() /2.
energy = sensor.getEnergy()
use_energy = float(energy / total_dep_energy)
use_energy = float(energy / max_dep_energy)
if logE:
raise NotImplementedError
use_energy = np.log(use_energy+1.) # - np.log(total_dep_energy)
@@ -343,7 +366,7 @@ def display_event(gd : GeometryDescriptor,
cmin=0,
cmax=1,
colorbar=dict(
title='Fraction of total deposited Energy',
title='Energy/max(Energy)',
tickvals=[0, 1],
ticktext=['0', '1'],
ticks='outside',
@@ -366,26 +389,27 @@ def display_event(gd : GeometryDescriptor,
name='Incoming particle',
showlegend=True,
)
# Calculate the direction vector for the arrow
direction_vector = [(end_point[0] - start_point[0]), (end_point[1] - start_point[1]), (end_point[2] - start_point[2])]
# Create the arrowhead at the start point with the opposite direction
arrowhead_trace = go.Cone(
x=[end_point[0]],
z=[end_point[1]],
y=[end_point[2]],
u=[direction_vector[0]],
w=[direction_vector[1]],
v=[direction_vector[2]],
sizemode='scaled',
sizeref=0.8,
showscale=False,
colorscale='Reds',
opacity=1.0,
anchor='tail',
)
# Create the 3D scatter plot with both traces
to_plot.append(line_trace)
to_plot.append(arrowhead_trace)
if (not isinstance(particleSpec, list)) or len(particleSpec) == 1:
# Calculate the direction vector for the arrow
direction_vector = [(end_point[0] - start_point[0]), (end_point[1] - start_point[1]), (end_point[2] - start_point[2])]
# Create the arrowhead at the start point with the opposite direction
arrowhead_trace = go.Cone(
x=[end_point[0]],
z=[end_point[1]],
y=[end_point[2]],
u=[direction_vector[0]],
w=[direction_vector[1]],
v=[direction_vector[2]],
sizemode='scaled',
sizeref=0.8,
showscale=False,
colorscale='Reds',
opacity=1.0,
anchor='tail',
)
# Create the 3D scatter plot with both traces
to_plot.append(line_trace)
to_plot.append(arrowhead_trace)
#
# finally show plot
#
@@ -403,12 +427,14 @@ def display_event(gd : GeometryDescriptor,
#name the axes in the HEP way, so y and z switch names
fig.update_layout(scene=dict(xaxis_title='x [mm]', yaxis_title='z [mm]', zaxis_title='y [mm]'))
#fig.write_html("temp.html")
#exit()
if renderer is not None:
fig.show(renderer=renderer)
if len(outfile) > 0:
fig.write_html(outfile)
else:
fig.show()
if renderer is not None:
fig.show(renderer=renderer)
else:
fig.show()