added multi particles
This commit is contained in:
+61
-35
@@ -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()
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user