from minicalo import GeometryDescriptor from minicalo import G4System as _G4System import plotly.graph_objects as go import pandas as pd import numpy as np import os import glob import subprocess import uproot import time from IPython.display import Image, display class __G4System(_G4System): def run_batch( self, nEvents: int, particleSpec: str, minEnergy_GeV: float, maxEnergy_GeV: float = -1.0, filename: str = "", ): if maxEnergy_GeV < 0: maxEnergy_GeV = minEnergy_GeV save_file = len(filename) > 0 # filename without file ending(!) filename = "_" + str(time.perf_counter_ns()) + ".root" _G4System.run_batch(self, nEvents, particleSpec, minEnergy_GeV, maxEnergy_GeV, filename=filename) # TO FIX: Geant4 adds "t" to the filename, circumvent this for one thread, but this is not a good solution file = glob.glob(filename.replace(".root", "*.root")) print(file) 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 if save_file: df.to_pickle(filename.replace(".root", ".pkl")) os.system("rm "+filename) return df def displayEvent(self, logE = False, renderer=None): # for loop over all layers to_plot = [] material_dict = {} z0=0 # sum up total deposited energy total_dep_energy = 0 for layer in self.getGeometryDescriptor().getLayers(): for sensor in layer.sensors: total_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 # loop over materials for layer in self.getGeometryDescriptor().getLayers(): layer_width = layer.nx * layer.sens_xwidth * 10. # in mm # # plot layers # layer_hx = layer_width / 2. layer_hy = layer_width / 2. layer_z = layer.thickness * 10. # in mm layer_material = layer.material # add material to materials if it is not already in there if layer_material not in material_dict.keys(): material_dict[layer_material] = {'name': layer_material, 'color': col_dict[layer_material], 'showlegend': False, 'flatshading': True, 'opacity': 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'], showlegend=True, name=layer_material)) to_plot.append(go.Mesh3d( # 8 vertices of a cube x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * layer_hx, z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * layer_hy, y = np.array([0, 0, 0, 0, layer_z, layer_z, layer_z, layer_z]) + z0, **ijk_cube, **material_dict[layer_material] )) # # sensors # # if there are sensors in the current layer, add them if layer.sensors != []: z = layer.sensors[0].getdz() corr = 0. #layer.sensors[0].getX() - layer.sensors[0].getdx()/2. + layer_width/2. # loop over all sensors in current layer and add them to plot for sensor in layer.sensors: x_center = sensor.getX() - corr y_center = sensor.getY() - corr hwidth = sensor.getdx() /2. energy = sensor.getEnergy() use_energy = float(energy / total_dep_energy) if logE: raise NotImplementedError use_energy = np.log(use_energy+1.) # - np.log(total_dep_energy) to_plot.append(go.Mesh3d( # 8 vertices of a cube x = np.array([-1, -1, 1, 1, -1, -1, 1, 1]) * hwidth + x_center, z = np.array([-1, 1, 1, -1, -1, 1, 1, -1]) * hwidth + y_center, y = np.array([0, 0, 0, 0, z, z, z, z]) + z0, **ijk_cube, flatshading=True, color='black', name='Sensor', opacity= max(0.03, use_energy), showlegend=False, )) z0 += layer_z # add legend entry for sensors to_plot.append(go.Mesh3d(x=[None], y=[None], z=[None], i=[0], j=[0], k=[0], color='black', showlegend=True, name='Sensors')) # add black-white colorbar for sensor hits to_plot.append(go.Surface( z=[[0, 0], [0, 0]], x=[[0, 0], [0, 0]], y=[[0, 0], [0, 0]], colorscale=[[0, 'white'], [1, 'black']], showscale=True, cmin=0, cmax=1, colorbar=dict( title='Fraction of total deposited Energy', tickvals=[0, 1], ticktext=['0', '1'], ticks='outside', ticklen=10, ), )) # # add red arrow for incoming particle # # Define the start and end points of the line start_point = [0, 0, - z0*0.1] end_point = [0, 0, - z0*0.25] # Create the line trace line_trace = go.Scatter3d( x=[start_point[0], end_point[0]], z=[start_point[1], end_point[1]], y=[start_point[2], end_point[2]], mode='lines', line=dict(color='red', width=5), 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=[start_point[0]], z=[start_point[1]], y=[start_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 # fig = go.Figure(data=[ *to_plot ]) fig.update_layout(legend=dict(x=0)) # add legend if renderer is not None: fig.show(renderer=renderer) else: fig.show() # # some helpers # def calculate_sensor_centers(X, square_size): sensor_size = square_size / np.sqrt(X) hsensor_size = sensor_size / 2. centers = [] for j in range(int(np.sqrt(X))): for i in range(int(np.sqrt(X))): x_center = -square_size/2 + (i + 0.5) * sensor_size y_center = -square_size/2 + (j + 0.5) * sensor_size centers.append((x_center[0], y_center[0])) return centers, hsensor_size[0] ijk_cube = { "i": [7, 0, 0, 0, 4, 4, 6, 6, 4, 0, 3, 2], "j": [3, 4, 1, 2, 5, 6, 5, 2, 0, 1, 6, 3], "k": [0, 7, 2, 3, 6, 7, 1, 1, 5, 5, 7, 6], } col_dict = { # active materials "G4_POLYSTYRENE": 'red', "G4_PLASTIC_SC_VINYLTOLUENE": 'blue', "G4_BGO": 'green', "G4_LSO": 'yellow', "G4_LYSO": 'orange', "G4_CESIUM_IODIDE": 'purple', "G4_PbWO4": 'pink', "G4_Si": 'brown', # passive materials "G4_Pb": 'darkgrey', "G4_Fe": 'lightgrey', "G4_W": 'grey', "G4_Cu": 'dimgray', "G4_BRASS": 'slategrey', } 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) G4System = __G4System()#singleton instance