#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Wed May 20 13:32:14 2026 @author: jamesszalkie """ import gc import numpy as np import pandas as pd from scipy.interpolate import interp1d import uproot import pycatima as catima from scipy.integrate import cumulative_trapezoid #matplotlib.use("Agg") import matplotlib.pyplot as plt import cmd import shlex import textwrap import os import periodictable as pt import re from matplotlib.colors import LinearSegmentedColormap from mpl_toolkits.mplot3d import Axes3D import nbformat as nbf # ROOT-like styling plt.rcParams.update({ # Figure "figure.facecolor": "white", "axes.facecolor": "white", # ROOT-like axes "axes.linewidth": 1.2, # Ticks on all sides "xtick.top": True, "ytick.right": True, # Tick direction inward "xtick.direction": "in", "ytick.direction": "in", # Major/minor ticks "xtick.major.size": 10, "ytick.major.size": 10, "xtick.minor.size": 5, "ytick.minor.size": 5, # Smaller ROOT-like fonts "font.size": 12, }) base_cmap = plt.cm.jet # Stop before the red region colors = base_cmap(np.linspace(0.2, 0.65, 256)) yellow_jet = LinearSegmentedColormap.from_list( 'yellow_jet', colors ) #data = [z, mass_u, maximum MeV, name] alpha_data = [2, 4.0015, 40, "alpha"] proton_data = [1, 1.0073, 20, "proton"] deuteron_data = [1, 2.014102, 30, "deuteron"] particles = { "alpha": alpha_data, "proton": proton_data, "deuteron": deuteron_data } interp_cache = {} SX3_SI_THICKNESS_CM = 0.1 QQQ_SI_THICKNESS_CM = 0.1 SX3_THETA_MIN_DEG = 0.0 def get_loss_table_path(medium, particle_label): script_dir = os.path.dirname(os.path.abspath(__file__)) return os.path.join(script_dir, f"{medium}Loss", f"E_vs_x_{particle_label}.dat") def clear_interpolator_cache(): """Drop cached energy-loss interpolators and force Python to release memory.""" global interp_cache interp_cache.clear() gc.collect() def make_E_vs_x( z, mass_u, emax_mev, medium, label, npoints, P_TORR, TEMP_K ): """ Builds energy vs depth table using Catima stopping powers. Output: x [cm], E [MeV] """ R = 8.3144 # J/mol/K p_pa = P_TORR * 133.322 molar_density = p_pa / (R * TEMP_K) # mol/m^3 # --------------------------- # Medium definition # --------------------------- if medium == "He": m_he = 4.0026 m_c = 12.0000 m_o = 15.9949 # FIX: do NOT double count oxygen material_def = [ (m_he, 2, 0.96), ((m_c + 2 * m_o), 18, 0.04) # CO2 lumped correctly ] #m_mix_avg = 0.96 * m_he + 0.04 * (m_c + 2 * m_o) m_mix_avg = m_he #rho_g_cm3 = (molar_density * m_mix_avg) / 1e6 if z == 2: rho_g_cm3 = 0.00017 #rho_g_cm3 = 8.2928e-5 else: rho_g_cm3 = 8.2928e-5 gas = catima.Material(material_def) gas.density(rho_g_cm3) elif medium == "Si": m_si = 28.084 rho_g_cm3 = 2.33 gas = catima.Material([(m_si, 14, 1.0)]) gas.density(rho_g_cm3) else: raise ValueError("Unsupported medium") print(f"[INFO] density = {rho_g_cm3:.3e} g/cm^3") projectile = catima.Projectile(mass_u, z) E = np.linspace(0.1, emax_mev, npoints) S_mass = np.zeros_like(E) for i, energy in enumerate(E): projectile.T(energy) S_mass[i] = catima.dedx(projectile, gas) # MeV/(g/cm^2) # convert to MeV/cm S_linear = S_mass * rho_g_cm3 invS = 1.0 / np.clip(S_linear, 1e-30, None) x = cumulative_trapezoid(invS[::-1], E[::-1], initial=0) x = x[::-1] x = -x df = pd.DataFrame({ "Distance_cm": x, "Energy_MeV": E }) outfile = get_loss_table_path(medium, label) os.makedirs(os.path.dirname(outfile), exist_ok=True) df.to_csv(outfile, sep="\t", index=False) print(f"[INFO] saved: {outfile}") interp_cache.pop((label.lower(), medium.lower()), None) gc.collect() return x, E #Generate energy loss tables from file def load_table(filename): """ Load table with columns: x(cm) E(MeV) Returns: x_array, E_array """ data = pd.read_csv( filename, sep=r'\s+', comment="#", header=None, skiprows=1 ) x = data.iloc[:, 0].to_numpy() E = data.iloc[:, 1].to_numpy() return x, E def get_interpolators(particle, medium): cache_key = (particle.lower(), medium.lower()) if cache_key in interp_cache: return interp_cache[cache_key] filename = get_loss_table_path(medium, particle) x, E = load_table(filename) E_of_x = interp1d( x, E, bounds_error=False, fill_value="extrapolate" ) x_of_E = interp1d( E[::-1], x[::-1], bounds_error=False, fill_value="extrapolate" ) interp_cache[cache_key] = (E_of_x, x_of_E) return E_of_x, x_of_E def energy_loss(particle, medium, Ei, dl): E_of_x, x_of_E = get_interpolators(particle, medium) xi = x_of_E(Ei) xf = xi + dl xmax = E_of_x.x.max() # maximum tabulated range Ef = E_of_x(xf) Ef = np.where(xf >= xmax, 0.0, Ef) return np.maximum(Ef, 0.0) def energy_reconstruction(particle, medium, Ef, dl): E_of_x, x_of_E = get_interpolators(particle, medium) xf = x_of_E(Ef) xi = xf - dl Ei = E_of_x(xi) return np.maximum(Ei, 0.0) def energy_distance(particle, medium, Ei, Ef): _, x_of_E = get_interpolators(particle, medium) xi = x_of_E(Ei) xf = x_of_E(Ef) return np.abs(xf - xi) def resolve_particle(name): name = name.lower().strip().rstrip("s") if name in particles: return particles[name] match = re.match(r"([a-zA-Z]+)[-\s]?(\d+)", name) if match: element_symbol = match.group(1).capitalize() A = int(match.group(2)) try: element = pt.elements.symbol(element_symbol) isotope = element[A] return ( isotope.number, # Z isotope.mass, # mass in u 30.0, f"{element_symbol}-{A}" ) except Exception: raise ValueError(f"Unknown isotope: {name}") try: elem = pt.elements.symbol(name.capitalize()) return elem.number, elem.mass, 30.0, name except Exception: raise ValueError(f"Unknown particle/isotope: {name}") def update_plot_data(name, values): for i, (existing_name, _) in enumerate(plot_data): if existing_name == name: plot_data[i] = (name, values) return plot_data.append((name, values)) def calculate_distance_tree1(vx, vy, vz, sx3x, sx3y, sx3z): """Calculate 3D distance from vertex to SX3 hit position for tree1""" dx = (sx3x - vx) / 10 dy = (sx3y - vy) / 10 dz = (sx3z - vz) / 10 return np.sqrt(dx*dx + dy*dy + dz*dz) def calculate_distance_tree2(vz, theta, z_max=34.86): """Calculate distance along trajectory from vertex to z_max for tree2 theta is polar angle from z-axis in radians""" cos_theta = np.cos(theta) cos_theta = np.where(np.abs(cos_theta) < 1e-10, 1e-10, cos_theta) return np.abs((z_max - vz) / cos_theta) def calculate_hit_theta(vx, vy, vz, hitx, hity, hitz): """Calculate polar angle from vertex to reconstructed hit position.""" dx = hitx - vx dy = hity - vy dz = hitz - vz distance = np.sqrt(dx * dx + dy * dy + dz * dz) distance = np.where(distance > 1e-12, distance, 1e-12) cos_theta = np.clip(dz / distance, -1.0, 1.0) return np.arccos(cos_theta) def calculate_sx3_thickness_path1(theta, phi, sx3x, sx3y, thickness_cm=0.1): rxy = np.sqrt(sx3x * sx3x + sx3y * sx3y) nx = np.where(rxy > 1e-12, sx3x / rxy, 0.0) ny = np.where(rxy > 1e-12, sx3y / rxy, 0.0) ux = np.sin(theta) * np.cos(phi) uy = np.sin(theta) * np.sin(phi) cos_incidence = np.abs(ux * nx + uy * ny) cos_incidence = np.clip(cos_incidence, 1e-4, 1.0) return thickness_cm / cos_incidence """ def calculate_sx3_thickness_path2(vx, vy, vz, sx3X, sx3Y, sx3Z, thickness_cm=0.1): dx = (sx3X - vx) / 100 dy = (sx3Y - vy) / 100 dz = (sx3Z - vz) / 100 distance = np.sqrt(dx * dx + dy * dy + dz * dz) distance = np.where(distance > 1e-12, distance, 1e-12) uz = dz / distance sin_incidence = np.abs(uz) sin_incidence = np.clip(sin_incidence, 1e-4, 1.0) return thickness_cm / sin_incidence""" def calculate_sx3_thickness_path(vx, vy, vz, sx3X, sx3Y, sx3Z, thickness_cm=0.1): dx = (sx3X - vx) dy = (sx3Y - vy) dz = (sx3Z - vz) distance = np.sqrt(dx * dx + dy * dy + dz * dz) radius = np.sqrt(dx * dx + dy * dy) #return thickness_cm return (distance * thickness_cm) / radius def calculate_qqq_thickness_path(vx, vy, vz, qqqx, qqqy, qqqz, thickness_cm=0.1): """Path length through planar QQQ silicon using hit direction and z-normal.""" dx = qqqx - vx dy = qqqy - vy dz = qqqz - vz distance = np.sqrt(dx * dx + dy * dy + dz * dz) distance = np.where(distance > 1e-12, distance, 1e-12) uz = dz / distance cos_incidence = np.abs(uz) cos_incidence = np.clip(cos_incidence, 1e-4, 1.0) return thickness_cm / cos_incidence def load_tree_arrays(tree, treename, max_events=None): if tree is None: raise ValueError( f"Tree '{treename}' is not loaded. Run 'set tree {treename}' " "after loading a ROOT file with 'uproot file '." ) branches = tree.keys(recursive=False) return tree.arrays(branches, library="np", entry_stop=max_events) def prepare_tree_data(tree, treename, particle, max_events=None, z_max=34.86): data = load_tree_arrays(tree, treename, max_events) initial_energy = data["Tb"] initial_energy_nan_count = np.isnan(initial_energy).sum() print(f"sx3 NaN entries: {initial_energy_nan_count}") vertex_x = data["vX"] vertex_y = data["vY"] vertex_z = data["vZ"] sx3_theta_rad = np.radians(data["thetab"]) sx3_phi_rad = np.radians(data["phib"]) sx3_hit_x = data["sx3X"] sx3_hit_y = data["sx3Y"] sx3_hit_z = data["sx3Z"] sx3_angle_mask = data["thetab"] > SX3_THETA_MIN_DEG sx3_valid_mask = ~np.isnan(sx3_hit_x) & ~np.isnan(sx3_hit_y) & ~np.isnan(sx3_hit_z) sx3_event_mask = sx3_angle_mask & sx3_valid_mask qqq_hit_x = data["qqqX"] qqq_hit_y = data["qqqY"] qqq_hit_z = data["qqqZ"] qqq_valid_mask = ( ~np.isnan(qqq_hit_x) & ~np.isnan(qqq_hit_y) & ~np.isnan(qqq_hit_z) & ~np.isnan(initial_energy) ) qqq_nan_count = np.isnan(qqq_hit_z).sum() print(f"qqq NaN entries: {qqq_nan_count}") print(f"Total: {initial_energy_nan_count + qqq_nan_count}") qqq_angle_mask = data["thetab"] > SX3_THETA_MIN_DEG qqq_event_mask = qqq_valid_mask & qqq_angle_mask #else: # sx3Z = np.full_like(vZ, z_max) # Assume sx3Z is at z_max for tree2 # mask = ~np.isnan(Ei) & ~np.isnan(theta) & ~np.isnan(vZ) sx3_initial_energy = initial_energy[sx3_event_mask] sx3_theta_rad = sx3_theta_rad[sx3_event_mask] sx3_phi_rad = sx3_phi_rad[sx3_event_mask] sx3_vertex_x = vertex_x[sx3_event_mask] sx3_vertex_y = vertex_y[sx3_event_mask] sx3_vertex_z = vertex_z[sx3_event_mask] qqq_initial_energy = initial_energy[qqq_event_mask] qqq_vertex_x = vertex_x[qqq_event_mask] qqq_vertex_y = vertex_y[qqq_event_mask] qqq_vertex_z = vertex_z[qqq_event_mask] #theta = calculate_hit_theta(vX, vY, vZ, sx3X, sx3Y, sx3Z)[mask] #if treename == 'tree1': sx3_hit_x = sx3_hit_x[sx3_event_mask] sx3_hit_y = sx3_hit_y[sx3_event_mask] sx3_hit_z = sx3_hit_z[sx3_event_mask] sx3_track_distance_cm = calculate_distance_tree1( sx3_vertex_x, sx3_vertex_y, sx3_vertex_z, sx3_hit_x, sx3_hit_y, sx3_hit_z, ) qqq_hit_x = qqq_hit_x[qqq_event_mask] qqq_hit_y = qqq_hit_y[qqq_event_mask] qqq_hit_z = qqq_hit_z[qqq_event_mask] qqq_track_distance_cm = calculate_distance_tree1( qqq_vertex_x, qqq_vertex_y, qqq_vertex_z, qqq_hit_x, qqq_hit_y, qqq_hit_z, ) qqq_theta_rad = calculate_hit_theta( qqq_vertex_x, qqq_vertex_y, qqq_vertex_z, qqq_hit_x, qqq_hit_y, qqq_hit_z, ) #else: # dsx3 = calculate_distance_tree2(vZsx3, theta, z_max=z_max) sx3_sin_theta = np.sin(sx3_theta_rad) sx3_sin_theta = np.where(sx3_sin_theta != 0, sx3_sin_theta, 1e-10) qqq_sin_theta = np.sin(qqq_theta_rad) qqq_sin_theta = np.where(qqq_sin_theta != 0, qqq_sin_theta, 1e-10) radii = np.array([3.7, 4.3]) sx3_anode_path_cm = ( radii[0] - np.sqrt((sx3_vertex_x / 10) ** 2 + (sx3_vertex_y / 10) ** 2) ) / sx3_sin_theta sx3_cathode_path_cm = ( radii[1] - np.sqrt((sx3_vertex_x / 10) ** 2 + (sx3_vertex_y / 10) ** 2) ) / sx3_sin_theta qqq_anode_path_cm = ( radii[0] - np.sqrt((qqq_vertex_x / 10) ** 2 + (qqq_vertex_y / 10) ** 2) ) / qqq_sin_theta qqq_cathode_path_cm = ( radii[1] - np.sqrt((qqq_vertex_x / 10) ** 2 + (qqq_vertex_y / 10) ** 2) ) / qqq_sin_theta sx3_silicon_path_cm = calculate_sx3_thickness_path( sx3_vertex_x, sx3_vertex_y, sx3_vertex_z, sx3_hit_x, sx3_hit_y, sx3_hit_z, SX3_SI_THICKNESS_CM,) * sx3_sin_theta qqq_silicon_path_cm = calculate_qqq_thickness_path( qqq_vertex_x, qqq_vertex_y, qqq_vertex_z, qqq_hit_x, qqq_hit_y, qqq_hit_z, QQQ_SI_THICKNESS_CM,) dA = sx3_anode_path_cm dC = sx3_cathode_path_cm # Filter out unphysical distances (negative or unreasonably small) # These typically occur at large angles where trajectory doesn't properly intersect proportional counters sx3ID = data["sx3ID"] aID = data["aID"] cID = data["cID"] # Keep only events with valid PW wire assignments and positive distances # sx3ID >= 0 means SX3 was hit, aID/cID valid means tracks found # Handle both 1D and 2D array structures from ROOT if aID.ndim == 2: aID_valid = aID[:, 0] >= 0 cID_valid = cID[:, 0] >= 0 else: aID_valid = aID >= 0 cID_valid = cID >= 0 print(f"Computing energies for {particle} ({treename})...") #print(f" Retained {np.sum(distance_mask)} / {len(distance_mask)} events after distance filter") clear_interpolator_cache() EA = energy_loss(particle, "He", sx3_initial_energy, sx3_anode_path_cm) EC = energy_loss(particle, "He", sx3_initial_energy, sx3_cathode_path_cm) EAq = energy_loss(particle, "He", qqq_initial_energy, qqq_anode_path_cm) ECq = energy_loss(particle, "He", qqq_initial_energy, qqq_cathode_path_cm) Esx3 = energy_loss(particle, "He", sx3_initial_energy, sx3_track_distance_cm) Eqqq = energy_loss(particle, "He", qqq_initial_energy, qqq_track_distance_cm) clear_interpolator_cache() Efinal = energy_loss(particle, "Si", Esx3, sx3_silicon_path_cm) EfinalQ = energy_loss(particle, "Si", Eqqq, qqq_silicon_path_cm) Edet = (Esx3 - Efinal) EdetQ = Eqqq - EfinalQ #Edet = np.where(Prange > lsx3, Efinal + Esx3, Esx3) Elost = sx3_initial_energy - Esx3 Elostqqq = qqq_initial_energy - Eqqq Eprop = EA - EC EpropQ = EAq - ECq return { "particle": particle, "Ei": sx3_initial_energy, "thetabi": data["thetab"], "thetabqqq": np.degrees(qqq_theta_rad), "sx3Z": sx3_hit_z, "EA": EA, "EC": EC, "Esx3": Esx3, "Elost": Elost, "Elostqqq": Elostqqq, "Eprop": Eprop, "EpropQ": EpropQ, "dA": dA, "dC": dC, "thetab": np.degrees(sx3_theta_rad), "sx3X": sx3_hit_x, "sx3Y": sx3_hit_y, "qqqX": qqq_hit_x, "qqqY": qqq_hit_y, "qqqZ": qqq_hit_z, "qqqE": Eqqq, "Eqqq": Eqqq, "Edet": Edet, "Efinal": Efinal, "lsx3": sx3_silicon_path_cm, "lqqq": qqq_silicon_path_cm, "EdetQ": EdetQ } def infer_particle_from_filename(filename): lower = os.path.basename(filename).lower() if "proton" in lower: return "proton" if "alpha" in lower: return "alpha" if "deuteron" in lower: return "deuteron" return None def process_file(filename, treename, particle=None, max_events=None): tree = uproot.open(filename)[f"{treename}"] if particle is None: particle = infer_particle_from_filename(filename) if particle is None: particle = "proton" print(f"File {filename} particle {particle}, tree {treename}") return prepare_tree_data(tree, treename, particle, max_events=max_events) def power_fit_and_plot(x, y, label, color=None): x = np.array(x) y = np.array(y) mask = ( np.isfinite(x) & np.isfinite(y) & (x > 0) & (y > 0) ) x = x[mask] y = y[mask] logx = np.log(x) logy = np.log(y) b, loga = np.polyfit(logx, logy, 1) a = np.exp(loga) y_pred = a * x**b ss_res = np.sum((y - y_pred)**2) ss_tot = np.sum((y - np.mean(y))**2) r2 = 1 - (ss_res / ss_tot) x_fit = np.linspace(np.min(x), np.max(x), 1000) y_fit = a * x_fit**b plt.plot( x_fit, y_fit, linewidth=3, color=color, label=f'{label} fit ($R^2$={r2:.4f})' ) print(f"\n{label} power-law fit:") print(f"y = {a:.6f} * x^{b:.6f}") print(f"R^2 = {r2:.6f}") return a, b, r2 class MyInteractiveApp(cmd.Cmd): def __init__(self): super().__init__() # Initial value set when the script starts self.buckets = 10000 self.T = 293.15 self.P = 379 self.temp_particle = [0, 0.0, 0.0, ""] self.rootFile = "SimAnasen1.root" self.file = None self.initialize_file() self.tree = None print("-" * 30) print("INTERACTIVE SHELL STARTED") self.print_params() print("Type 'help' for commands.") print("Type 'exit' to end program") print("-" * 30) def print_params(self): """Helper method to display current state""" print(f"Current Parameters: T={self.T} K, P={self.P} Torr") def initialize_file(self): """Load in default root file for anasen""" try: self.file = uproot.open(f"../Armory/{self.rootFile}") except: self.file = None print("\nATTENTION: Root file not found, continue without uproot functions or uproot file manually") #intro = "Interactive Shell Started. Type 'help' to see commands." prompt = ">> " def default(self, line): # Check if the command starts with our multi-word phrase if line.startswith("make table "): # Extract everything after "make table " args = line[len("make table "):].strip() self.do_make_table(args) elif line.startswith("set t") or line.startswith("Set T") or line.startswith("set T") or line.startswith("Set t"): args = line[len("set t "):].strip() self.do_set_T(args) elif line.startswith("set p") or line.startswith("Set P") or line.startswith("set P") or line.startswith("Set p"): args = line[len("set p "):].strip() self.do_set_P(args) elif line.startswith("set buckets") or line.startswith("Set Buckets") or line.startswith("set Buckets") or line.startswith("Set buckets"): args = line[len("set buckets "):].strip() self.do_set_Buckets(args) elif line.startswith("energy loss") or line.startswith("Energy Loss") or line.startswith("Energy loss"): args = line[len("energy loss "):].strip() self.do_energy_loss(args) elif line.startswith("energy reconstruction") or line.startswith("Energy Reconstruction") or line.startswith("Energy reconstruction"): args = line[len("energy reconstruction "):].strip() self.do_energy_reconstruction(args) elif line.startswith("energy distance") or line.startswith("Energy Distance") or line.startswith("Energy distance"): args = line[len("energy distance "):].strip() self.do_energy_distance(args) else: print(f"*** Unknown syntax: {line}") def do_T(self, arg): """Print value of T""" print(self.T) def do_Buckets(self, arg): """print number of histogram buckets and points""" print(self.buckets) def do_P(self, arg): """Print value of P (pressure)""" print(self.P) def do_make_particle(self, arg): """Enter parameters for new particle to be given to catyma system Ex: >> make_particle """ args = shlex.split(arg) z = int(args[0]) mass_u = float(args[1]) emax_mev = float(args[2]) particle = args[3] self.temp_particle = [z, mass_u, emax_mev, particle] print(f"New particle '{particle}': z = {z}, mass_u = {mass_u}, emax_mev = {emax_mev}") def do_set_T(self, arg): """Changes the value of T. Usage: set_t 300""" try: self.T = float(arg) print(f"T has been updated to {self.T}") except ValueError: print("Please enter a valid number for T.") def do_set_P(self, arg): """Changes the value of P in Torr. Usage: set_ 400""" try: self.P = float(arg) print(f"P has been updated to {self.P}") except ValueError: print("Please enter a valid number for P.") def do_set_Buckets(self, arg): """Changes the value of buckets""" try: self.buckets = int(arg) print(f"Number of buckets has been updated to {self.buckets}") except ValueError: print("Please enter a valid number of buckets.") def do_exit(self, arg): """Exits the application.""" print("Closing application...") return True # Returning True stops the cmdloop() def do_make_table(self, arg): """Create E vs X tables for particle, or isotopes Ex: >> make table proton Ex: >> make table Co60 Ex: >> make table N17 """ try: args = shlex.split(arg) if not args: print("Please enter desired reaction particle") return name = args[0] if len(args) > 1: emax_mev = float(args[1]) else: emax_mev = None if len(args) > 2: medium = args[2] else: medium = "He" z, mass_u, default_emax, label = resolve_particle(name) if emax_mev is None: emax_mev = default_emax x, E = make_E_vs_x( z, mass_u, emax_mev, medium, label, self.buckets, self.P, self.T ) plt.figure(figsize=(8,6)) plt.plot(x, E) plt.xlabel("Distance (cm)") plt.ylabel("Energy (MeV)") plt.title(f"Energy Loss Curve {label.capitalize()} {medium}") plt.grid(True) textstr = f"T = {self.T:.2f} K\nP = {self.P:.2f} Torr" plt.gca().text( 0.02, 0.02, textstr, transform=plt.gca().transAxes, fontsize=10, verticalalignment='bottom', bbox=dict(boxstyle="round", facecolor="white", alpha=0.7) ) plt.tight_layout() filename = f"{medium}Loss/Energy_Loss_Curve_{label}.png" plt.savefig(filename, dpi=300, bbox_inches="tight") plt.show() print(f"Saved plot: {filename}") except Exception as e: print(f"Error in make_table: {e}") def load_table(self, arg): args = shlex.split(arg) particle = args[0] filename = f"E_vs_x_{particle}.dat" load_table(filename) print(f"Loaded table {filename}") def do_energy_loss(self, arg): """Find a final energy given an initial energy and distance travelled Ex: >> energy loss """ args = shlex.split(arg) try: particle = args[0] medium = args[1] Ei = float(args[2]) dl = float(args[3]) #Ei_offset = Ei * 1.1 #table_specs = f"{particle} {Ei_offset}" try: #self.do_make_table(table_specs) Ef = energy_loss(particle, medium, Ei, dl) print(f"\nFinal energy: {Ef:.6f} MeV\n") except: return except IndexError: print("Please input particle, initial energy, and distance travelled") def do_energy_reconstruction(self, arg): """Find a vertex energy given an final energy and distance travelled Ex: >> energy reconstruction """ args = shlex.split(arg) try: particle = args[0] medium = args[1] Ef = float(args[2]) dl = float(args[3]) try: Ei = energy_reconstruction(particle, medium, Ef, dl) if Ei > 0: print(f"\nInitial energy: {Ei:.6f} MeV\n") else: print("Error: remake table with larger value, fallen off map") except: print("Particle energy table not made yet, please do so using 'make table'") except IndexError: print("Please input particle, final energy from detector, and distance travelled") def do_energy_distance(self, arg): """Find a distance travelled given an initial and final energy Ex: >> energy distance """ args = shlex.split(arg) try: particle = args[0] medium = args[1] Ei = float(args[2]) Ef = float(args[3]) dE = Ei - Ef try: dl = energy_distance(particle, medium, Ei, Ef) if dl > 0: print(f"\nChange in energy: {dE:.6f} MeV") print(f"Distance travelled: {dl:.6f} cm\n") print(f"{(dl * 10000):.6f} microns") else: print("Error: remake table with larger value, fallen off map") except: print("Particle energy table not made yet, please do so using 'make table'") except IndexError: print("Please input particle, final energy from detector, and distance travelled") def do_uproot_file(self, arg): """Open a specific root file for inspection""" args = shlex.split(arg) if len(args) > 0: filename = args[0] else: filename = self.rootFile try: print(f"Opening {filename}") # Try Armory path first try: self.file = uproot.open(f"../Armory/{filename}") except FileNotFoundError: self.file = uproot.open(filename) print("File loaded successfully.") print("Keys:", self.file.keys()) except Exception as e: print("Error opening file:", e) def do_print_file(self, arg): """Print contents of ROOT file""" args = shlex.split(arg) file = self.file if "keys" in args or len(args) == 0: print("Keys in file: ", file.keys()) if "class_names" in args or len(args) == 0: print("Class names: ", file.classnames()) def do_set_tree(self, arg): if self.file is None: print("No ROOT file loaded.") return keys = self.file.keys() print("Available trees:", keys) treeName = arg if len(arg) > 0 else "tree1" try: # uproot automatically resolves ";1" self.tree = self.file[treeName] print(f"Tree loaded: {self.tree}") print("Branches:", self.tree.keys()) except Exception as e: print("Error loading tree:", treeName) print("Exception:", e) def run_command_line(self): import readline import atexit import os histfile = os.path.expanduser("~/.uproot_shell_history") try: readline.read_history_file(histfile) except FileNotFoundError: pass readline.set_history_length(1000) readline.parse_and_bind("tab: complete") atexit.register(readline.write_history_file, histfile) print("Custom Python CMD (type 'exit' to stop)") local_vars = { "file": self.file, "uproot": uproot, "np": np, "plt": plt, "self": self } startup_script = textwrap.dedent(""" tree = file["tree"] """) exec(startup_script, globals(), local_vars) while True: try: entry = input(">>> ").strip() if entry.lower() in ["exit", "quit"]: break if not entry: continue try: result = eval(entry, globals(), local_vars) if result is not None: print(result) except SyntaxError: exec(entry, globals(), local_vars) if "tree" in local_vars: self.tree = local_vars["tree"] except OverflowError(): print("Arrays too large, causing crash") except Exception as e: print(f"Error: {e}") self.shell_vars = local_vars def do_uproot(self, arg): """Start up an in-program command line to use root tools with python, look up 'uproot' for more details""" self.run_command_line() def do_energy_analysis(self, arg): #geometry needs update args = shlex.split(arg) try: particle = args[0] except IndexError: print("Please indicate reactant for analysis") return try: max_events = int(args[1]) except IndexError: max_events = None if self.tree is None: self.do_set_tree("") print(f"Using TTree: {self.tree}") if self.tree is None: print("No tree is loaded. Use 'uproot file ' then 'set tree tree1' or 'set tree tree2'.") return treename = self.tree.name if hasattr(self.tree, 'name') else "tree1" branches = [ "Tb", "thetab", "sx3Z", "vX", "vY", "vZ", "sx3X", "sx3Y", "sx3XExit", "sx3YExit", "sx3ZExit", ] initial_energy_key = "Tb" polar_angle_key = "thetab" if max_events: n_events = max_events else: n_events = self.tree.num_entries print(f"Loading {n_events} events from {treename}...") data = self.tree.arrays( branches, library="np", entry_stop=max_events ) initial_energy = data[initial_energy_key] global sx3Z sx3Z = data["sx3Z"] sx3XExit = data["sx3XExit"] if "sx3XExit" in data else None sx3YExit = data["sx3YExit"] if "sx3YExit" in data else None sx3ZExit = data["sx3ZExit"] if "sx3ZExit" in data else None polar_angle_rad = np.radians(data[polar_angle_key]) vertex_x = data["vX"] vertex_y = data["vY"] vertex_z = data["vZ"] if treename == 'tree1': sx3_hit_x = data["sx3X"] sx3_hit_y = data["sx3Y"] event_mask = ~np.isnan(sx3_hit_x) & ~np.isnan(sx3_hit_y) & ~np.isnan(sx3Z) else: event_mask = ~np.isnan(initial_energy) & ~np.isnan(polar_angle_rad) initial_energy = initial_energy[event_mask] polar_angle_rad = polar_angle_rad[event_mask] sx3Z = sx3Z[event_mask] vertex_z_masked = vertex_z[event_mask] sin_theta = np.sin(polar_angle_rad) sin_theta = np.where(sin_theta != 0, sin_theta, 1e-10) if treename == 'tree1': vertex_x_masked = vertex_x[event_mask] vertex_y_masked = vertex_y[event_mask] sx3_hit_x_masked = sx3_hit_x[event_mask] sx3_hit_y_masked = sx3_hit_y[event_mask] sx3_track_distance_cm = calculate_distance_tree1( vertex_x_masked, vertex_y_masked, vertex_z_masked, sx3_hit_x_masked, sx3_hit_y_masked, sx3Z, ) sx3_silicon_path_cm = calculate_sx3_thickness_path( vertex_x_masked, vertex_y_masked, vertex_z_masked, sx3_hit_x_masked, sx3_hit_y_masked, sx3Z, SX3_SI_THICKNESS_CM,) else: sx3_track_distance_cm = calculate_distance_tree2(vertex_z_masked, polar_angle_rad, z_max=34.86) sx3_silicon_path_cm = np.full_like(sx3_track_distance_cm, np.nan, dtype=float) radii = np.array([3.7, 4.2]) anode_path_cm = radii[0] / sin_theta cathode_path_cm = radii[1] / sin_theta print("Calculating energy losses...") global EA EA = energy_loss(particle, "He", initial_energy, anode_path_cm) global EC EC = energy_loss(particle, "He", initial_energy, cathode_path_cm) global Esx3 Esx3 = energy_loss(particle, "He", initial_energy, sx3_track_distance_cm) global Eprop Eprop = EA - EC global Elost Elost = initial_energy - Esx3 print("Analysis complete") print(f"Processed events: {len(initial_energy)}") print(f"Anode average energy: {np.mean(EA):.3f} MeV") print(f"Cathode average energy: {np.mean(EC):.3f} MeV") print(f"sx3 average energy: {np.mean(Esx3):.3f} MeV") print(f"Average total energy loClassSX3ss to sx3: {np.mean(Elost):.3f} MeV") print(f"Maximum total energy loss to sx3: {np.max(Elost):.3f} MeV") print(f"Minimum total energy loss to sx3: {np.min(Elost):.3f} MeV") print(f"Proportion counter average energy difference: {np.mean(Eprop):.3f} MeV") print(f"Maximum proportion counter energy difference: {np.max(Eprop):.3f} MeV") print(f"Minimum proportion counter energy difference: {np.min(Eprop):.3f} MeV") output_filename = "energy_analysis.root" print(f"Writing new tree to {output_filename}") # Load ALL original branches all_data = self.tree.arrays(library="np",entry_stop=max_events) # Create full-length arrays initialized to NaN n_total = len(data["Tb"]) EA_full = np.full(n_total, np.nan) EC_full = np.full(n_total, np.nan) Esx3_full = np.full(n_total, np.nan) Eprop_full = np.full(n_total, np.nan) Elost_full = np.full(n_total, np.nan) # Put values back into valid entries EA_full[event_mask] = EA EC_full[event_mask] = EC Esx3_full[event_mask] = Esx3 Eprop_full[event_mask] = Eprop Elost_full[event_mask] = Elost # Add new branches all_data["EA"] = EA_full all_data["EC"] = EC_full all_data["Esx3"] = Esx3_full all_data["Eprop"] = Eprop_full all_data["Elost"] = Elost_full if treename == 'tree1': lsx3_full = np.full(n_total, np.nan) lsx3_full[event_mask] = sx3_silicon_path_cm all_data["lsx3"] = lsx3_full # Write new ROOT file as a classic TTree with uproot.recreate(output_filename) as fout: branch_types = { name: array.dtype for name, array in all_data.items() } fout.mktree("tree", branch_types) fout["tree"].extend(all_data) print("Finished writing augmented ROOT file") def do_make_plots(self, arg): import os global plot_data plot_data = [] args = shlex.split(arg) particle = args[0] if len(args) > 0 else "proton" max_events = int(args[1]) if len(args) > 1 else None if self.tree is None: self.do_set_tree("") print(f"Using TTree: {self.tree}") if self.tree is None: print("No tree is loaded. Use 'uproot file ' then 'set tree tree1' or 'set tree tree2'.") return treename = self.tree.name if hasattr(self.tree, 'name') else "tree1" data = prepare_tree_data(self.tree, treename, particle, max_events=max_events) Ei = data["Ei"] thetab = data["thetab"] thetabqqq = data["thetabqqq"] sx3X = data["sx3X"] sx3Y = data["sx3Y"] sx3Z = data["sx3Z"] EA = data["EA"] EC = data["EC"] Esx3 = data["Esx3"] Efinal = data["Efinal"] Elost = data["Elost"] Eprop = data["Eprop"] EpropQ = data["EpropQ"] dA = data["dA"] dC = data["dC"] qqqX = data["qqqX"] qqqY = data["qqqY"] qqqZ = data["qqqZ"] qqqE = data["Eqqq"] Elostqqq = data["Elostqqq"] Edet = data["Edet"] EdetQ = data["EdetQ"] update_plot_data(f"{particle}_{treename}_Ei", Ei) update_plot_data(f"{particle}_{treename}_sx3Z", sx3Z) update_plot_data(f"{particle}_{treename}_EA", EA) update_plot_data(f"{particle}_{treename}_EC", EC) update_plot_data(f"{particle}_{treename}_Esx3", Esx3) update_plot_data(f"{particle}_{treename}_Elost", Elost) update_plot_data(f"{particle}_{treename}_Eprop", Eprop) update_plot_data(f"{particle}_treename_Edert", Edet) base = f"{particle}_{treename}_plots" os.makedirs(base, exist_ok=True) print(f"Saving plots to folder: {base} ({treename})") # --- clean data --- x = np.asarray(sx3X, dtype=float) y = np.asarray(sx3Y, dtype=float) z = np.asarray(sx3Z, dtype=float) c = np.asarray(Esx3, dtype=float) qx = np.asarray(qqqX, dtype=float) qy = np.asarray(qqqY, dtype=float) qz = np.asarray(qqqZ, dtype=float) qc = np.asarray(qqqE, dtype=float) sx3_theta_plot_mask = np.isfinite(thetab) & (thetab > SX3_THETA_MIN_DEG) qqq_theta_plot_mask = np.isfinite(thetabqqq) & (thetabqqq > SX3_THETA_MIN_DEG) if x.size > 0 and False: fig = plt.figure(figsize=(8,6)) ax = fig.add_subplot(111, projection='3d') sc = ax.scatter( x, y, z, c=c, cmap='viridis', s=3, alpha=0.8 ) ax.set_xlabel("SX3 X") ax.set_ylabel("SX3 Y") ax.set_zlabel("SX3 Z") ax.set_title(f"{particle} ({treename}) SX3 3D Hit Map colored by Esx3") cb = plt.colorbar(sc, ax=ax, pad=0.1) cb.set_label("Esx3 (MeV)") plt.tight_layout() out_file = f"{base}/sx3_3D_energy.png" plt.savefig(out_file, dpi=300) plt.show() fig = plt.figure(figsize=(8,6)) ax = fig.add_subplot(111, projection='3d') sc = ax.scatter( qx, qy, qz, c=qc, cmap='viridis', s=3, alpha=0.8 ) ax.set_xlabel("SX3 X") ax.set_ylabel("SX3 Y") ax.set_zlabel("SX3 Z") ax.set_title(f"{particle} ({treename}) QQQ 3D Hit Map colored by qqqTb") cb = plt.colorbar(sc, ax=ax, pad=0.1) cb.set_label("qqqTb (MeV)") plt.tight_layout() out_file = f"{base}/QQQ_3D_energy.png" plt.savefig(out_file, dpi=300) plt.show() mask1 = ~np.isnan(Ei) & ~np.isnan(thetab) & sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(np.radians(thetab[mask1]), Ei[mask1], bins=200) plt.xlabel("thetab") plt.ylabel("Event Energy (MeV)") plt.title(f"{particle} ({treename}) Energy vs Theta") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/E_vs_theta.png", dpi=300) plt.show() plt.figure(figsize=(7,6)) mask1 = ~np.isnan(Esx3) & ~np.isnan(Edet) & sx3_theta_plot_mask plt.hist2d(Esx3[mask1], Edet[mask1], bins=300) plt.xlabel("Esx3") plt.ylabel("Edet") plt.title("Esx3 vs Edet") #plt.yscale("log") #plt.xscale("log") plt.colorbar(label="counts") plt.tight_layout() plt.savefig(f"{base}/Esx3_vs_Edet.png") plt.show() mask1 = ~np.isnan(qqqE) & ~np.isnan(EdetQ) & qqq_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(qqqE[mask1], EdetQ[mask1], bins=300) plt.xlabel("Eqqq") plt.ylabel("Edet") plt.title("Eqqq vs Edet") plt.yscale("log") plt.xscale("log") plt.colorbar(label="counts") plt.tight_layout() plt.savefig(f"{base}/Eqqq_vs_Edet.png") plt.show() mask1 = ~np.isnan(qqqE) & ~np.isnan(thetabqqq) & qqq_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(qqqE[mask1], (np.sin(np.deg2rad(thetabqqq)) * Elostqqq)[mask1], bins=200) plt.ylabel("Elost x sin(theta)") plt.xlabel("Eqqq (MeV)") plt.title(f"{particle} ({treename}) Energy QQQ vs Elost * Theta") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/Eqqq_vs_Elostxsintheta.png", dpi=300) plt.show() mask1 = (Esx3 > 0) & ~np.isnan(thetab) & sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(thetab[mask1], Esx3[mask1], bins=200) plt.xlabel("thetab") plt.ylabel("Esx3") plt.title(f"{particle} ({treename}) sx3 Energy vs Theta") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/sx3E_vs_theta.png", dpi=300) plt.show() mask1 = (qqqE > 0) & ~np.isnan(thetabqqq) & qqq_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(thetabqqq[mask1], qqqE[mask1], bins=200) plt.xlabel("thetab") plt.ylabel("E-QQQ") plt.title(f"{particle} ({treename}) QQQ Energy vs Theta") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/QQQE_vs_theta.png", dpi=300) plt.show() plt.figure(figsize=(7,5)) plt.hist(Elost[sx3_theta_plot_mask], bins=200) plt.xlabel("Energy Loss (MeV)") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) Total Energy Loss Distribution") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/Elost_hist.png", dpi=300) plt.show() plt.figure(figsize=(7,5)) plt.hist(Edet[sx3_theta_plot_mask], bins=200) plt.xlabel("Energy Deposition (MeV)") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) TEnergy Deposition in SX3 (MeV)") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/Edet_hist.png", dpi=300) plt.show() plt.figure(figsize=(7,5)) plt.hist(Efinal[sx3_theta_plot_mask], bins=200) plt.xlabel("Final Energy (MeV)") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) Final Energy after sx3") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/Efinal_hist.png", dpi=300) plt.show() try: plt.figure(figsize=(7,5)) plt.hist(data["lsx3"][sx3_theta_plot_mask], bins=200) plt.xlabel("Distance travelled in sx3") plt.ylabel("Counts") #plt.xlim(0, 2) plt.yscale("log") plt.title(f"{particle} ({treename}) Distance Travelled in sx3") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/lsx3.png", dpi=300) plt.show() except: print("lsx3 not able to plot") plt.figure(figsize=(7,5)) plt.hist(data["lqqq"], bins=200) plt.xlabel("Distance travelled in QQQ") plt.ylabel("Counts") #plt.xlim(0, 2) plt.yscale("log") plt.title(f"{particle} ({treename}) Distance Travelled in QQQ") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/lqqq.png", dpi=300) plt.show() plt.figure(figsize=(7,5)) plt.hist(Ei[sx3_theta_plot_mask], bins=100, histtype='step', label='Vertex Energy') plt.hist(Esx3[(Esx3 != 0) & sx3_theta_plot_mask], bins=100, histtype='step', label='SX3 Energy') plt.xlabel("Energies (MeV)") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) Vertex and SX3 Energy") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/Vertex_vs_SX3_Energy.png", dpi=300) plt.legend(loc='upper right') plt.show() plt.figure(figsize=(7,5)) plt.hist(dA[sx3_theta_plot_mask], bins=100, label='dA', histtype='step') plt.hist(dC[sx3_theta_plot_mask], bins=100, label='dC', histtype='step') plt.xlabel("Distance") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) Anode/ Cathode distance") plt.legend(loc='upper right') plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/d_hist.png", dpi=300) plt.show() try: mask1 = sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(sx3Z[mask1], Elost[mask1], bins=200) plt.ylabel("Energy Loss (MeV)") plt.xlabel("SX3 Z") plt.title(f"{particle} ({treename}) Energy Loss vs SX3 Position") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig(f"{base}/Elost_vs_sx3Z.png", dpi=300) plt.show() except ValueError: print("Value error") try: mask = (Esx3 > 0) & sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(EA[mask], Esx3[mask], bins=200) plt.xlabel("EA (MeV)") plt.ylabel("Esx3 (MeV)") plt.title(f"{particle} ({treename}) Anode vs SX3 Energy") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig(f"{base}/EA_vs_Esx3.png", dpi=300) plt.show() except ValueError: print("Value error") try: mask1 = sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(sx3Z[mask1], Eprop[mask1], bins=200) plt.ylabel("EA - EC (MeV)") plt.xlabel("SX3 Z") plt.title(f"{particle} ({treename}) Energy Propagation Difference vs Position") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig(f"{base}/Eprop_vs_sx3Z.png", dpi=300) plt.show() except ValueError: print("Value error") try: mask1 = ~np.isnan(Esx3) & (Esx3 > 0) & sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(Esx3[mask1], np.sin(np.deg2rad(data["thetab"][mask1])) * Eprop[mask1], bins=200) plt.ylabel("PCEnergy x sin(theta)") plt.xlabel("SX3 Energy (MeV)") plt.title(f"{particle} ({treename}) Energy Propagation Difference vs SX3 Energy") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/Eprop_vs_Esx3.png", dpi=300) plt.show() except ValueError: print("Value error") try: mask1 = ~np.isnan(Edet) & (Edet > 0) & sx3_theta_plot_mask plt.figure(figsize=(7,6)) plt.hist2d(Edet[mask1], np.sin(np.deg2rad(data["thetab"][mask1])) * Eprop[mask1], bins=200) plt.ylabel("PCEnergy x sin(theta)") plt.xlabel("SX3 Energy Detected (MeV)") plt.title(f"{particle} ({treename}) Energy Propagation Difference vs SX3 Energy Detected") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,.5) plt.tight_layout() plt.savefig(f"{base}/Eprop_vs_Edet.png", dpi=300) plt.show() except ValueError: print("Value error") mask1 = ( ~np.isnan(EdetQ) & (EdetQ > 0) & ~np.isnan(EpropQ) & ~np.isnan(thetabqqq) & qqq_theta_plot_mask ) plt.figure(figsize=(7,6)) plt.hist2d(EdetQ[mask1], (np.sin(np.deg2rad(thetabqqq[mask1])) * EpropQ[mask1]), bins=200) plt.ylabel("PCEnergy x sin(theta)") plt.xlabel("Eqqq detected (MeV)") plt.title(f"{particle} ({treename}) EDet QQQ vs PCEnergy * sin(theta)") plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.tight_layout() plt.savefig(f"{base}/Eprop_vs_EdetQ.png", dpi=300) plt.show() branch_names = [] for key in self.tree.keys(): if isinstance(key, bytes): branch_names.append(key.decode("ascii")) elif hasattr(key, "name"): branch_names.append(key.name) else: branch_names.append(str(key)) branch_names = list(dict.fromkeys(branch_names)) if branch_names and True: print(f"Creating histograms for {len(branch_names)} branches...") all_branches = self.tree.arrays(branch_names, library="np", entry_stop=max_events) for branch in branch_names: values = all_branches[branch] try: values = np.asarray(values, dtype=float) except Exception: print(f"Skipping non-numeric branch: {branch}") continue if values.ndim != 1: print(f"Skipping non-scalar branch: {branch}") continue # Apply the thetab cut #values = values[mask] values = values[~np.isnan(values)] if values.size == 0: print(f"Skipping empty branch: {branch}") continue plt.figure(figsize=(7,5)) try: plt.hist(values, bins=100) except: plt.hist(values, bins=10) plt.xlabel(branch) plt.ylabel("Counts") plt.title(f"{particle} ({treename}) {branch} distribution") plt.grid(True) plt.tight_layout() safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", branch) plt.savefig(f"{base}/{safe_name}_hist.png", dpi=300) plt.close() else: print("No branches found to histogram.") print(f"Plotting complete ({treename}).") def do_hist_comp(self, arg): particle = arg base = f"{particle}_comp_plots" os.makedirs(base, exist_ok=True) branch_names = [] for key in self.tree.keys(): if isinstance(key, bytes): branch_names.append(key.decode("ascii")) elif hasattr(key, "name"): branch_names.append(key.name) else: branch_names.append(str(key)) branch_names = list(dict.fromkeys(branch_names)) print(branch_names) if branch_names: print(f"Creating histograms for {len(branch_names)} branches...") self.do_set_tree("tree1") all_branches1 = self.tree.arrays(branch_names, library="np", entry_stop=None) self.do_set_tree("tree2") all_branches2 = self.tree.arrays(branch_names, library="np", entry_stop=None) for branch in branch_names: values1 = all_branches1[branch] values2 = all_branches2[branch] try: values1 = np.asarray(values1, dtype=float) values2 = np.asarray(values2, dtype=float) except Exception: print(f"Skipping non-numeric branch: {branch}") continue if values1.ndim != 1 or values2.ndim != 1: print(f"Skipping non-scalar branch: {branch}") continue values1 = values1[~np.isnan(values1)] values2 = values2[~np.isnan(values2)] if values1.size == 0 or values2.size == 0: print(f"Skipping empty branch: {branch}") continue plt.figure(figsize=(7,5)) plt.hist(values1, bins=100, histtype='step') plt.hist(values2, bins=100, histtype='step') plt.xlabel(branch) plt.ylabel("Counts") plt.title(f"{particle} {branch} distribution") plt.grid(True) plt.tight_layout() safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", branch) plt.savefig(f"{base}/{safe_name}_hist.png", dpi=300) plt.show() else: print("No branches found to histogram.") def do_dual_plotter(self, arg): args = shlex.split(arg) # Default files if none provided if len(args) == 0: files = [ "SimAnasenProton.root", "SimAnasenAlpha.root" ] else: files = args outdir = "dual_plots" os.makedirs(outdir, exist_ok=True) print(f"Saving plots to: {outdir}") datasets = [] script_dir = os.path.dirname(os.path.abspath(__file__)) project_root = os.path.abspath(os.path.join(script_dir, "..")) armory_dir = os.path.join(project_root, "Armory") for file in files: try: if os.path.isabs(file): resolved_path = file elif os.path.exists(file): resolved_path = os.path.abspath(file) else: resolved_path = os.path.join(armory_dir, file) resolved_particle = infer_particle_from_filename(resolved_path) tree1 = process_file( resolved_path, "tree1", particle=resolved_particle ) try: tree2 = process_file( resolved_path, "tree2", particle=resolved_particle ) data = { "particle": f"{resolved_particle or tree1['particle']}_combined" } for key in tree1: if key == "particle": continue try: data[key] = np.concatenate( [tree1[key], tree2[key]] ) except Exception: data[key] = tree1[key] except Exception: data = tree1 datasets.append(data) print( f"Loaded {file} " f"({len(data['Ei'])} events)" ) except Exception as e: print(f"Failed to load {file}: {e}") if len(datasets) == 0: print("No valid datasets loaded.") return def finite_values(values): array = np.asarray(values, dtype=float) return array[np.isfinite(array)] def finite_pairs(x_values, y_values): x_array = np.asarray(x_values, dtype=float) y_array = np.asarray(y_values, dtype=float) valid_mask = np.isfinite(x_array) & np.isfinite(y_array) return x_array[valid_mask], y_array[valid_mask] def save_overlay_histogram(value_getter, xlabel, ylabel, title, output_name): plt.figure(figsize=(8, 6)) for data in datasets: values = finite_values(value_getter(data)) if values.size == 0: continue plt.hist( values, bins=200, histtype="step", linewidth=2, density=True, label=data["particle"], ) plt.xlabel(xlabel) plt.ylabel(ylabel) plt.title(title) plt.legend() plt.grid(True) plt.tight_layout() plt.savefig(f"{outdir}/{output_name}", dpi=300) plt.show() def save_overlay_hist2d(x_getter, y_getter, xlabel, ylabel, title, output_name, bins=200): n = len(datasets) fig, axes = plt.subplots(1, n, figsize=(7 * n, 6)) if n == 1: axes = [axes] for ax, data in zip(axes, datasets): x_values, y_values = finite_pairs(x_getter(data), y_getter(data)) if x_values.size == 0 or y_values.size == 0: continue h = ax.hist2d(x_values, y_values, bins=bins) ax.set_title(f'{data["particle"]}\n{title}') ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) fig.colorbar(h[3], ax=ax, label="Counts") plt.tight_layout() plt.savefig(f"{outdir}/{output_name}", dpi=300) plt.show() def stack_by_key(key): stacked = [] for data in datasets: values = finite_values(data.get(key, [])) if values.size > 0: stacked.append(values) return stacked plt.figure(figsize=(8, 6)) for data in datasets: dep_values = data.get("Edet", data.get("Elost", [])) plt.hist( dep_values, bins=200, histtype="step", linewidth=2, density=True, label=data["particle"] ) plt.xlabel("Energy Deposition (MeV)") plt.ylabel("Normalized Counts") plt.title("Energy Deposition Comparison") plt.legend() plt.grid(True) plt.tight_layout() plt.savefig( f"{outdir}/Edet_overlay.png", dpi=300 ) plt.show() save_overlay_histogram( lambda data: data["sx3Z"], "SX3 Z", "Normalized Counts", "SX3 Position Comparison", "sx3Z_overlay.png", ) save_overlay_histogram( lambda data: data["EdetQ"], "QQQ Energy Deposition (MeV)", "Normalized Counts", "QQQ Energy Deposition Comparison", "EdetQ_overlay.png", ) save_overlay_histogram( lambda data: data["thetabqqq"], "QQQ Theta (deg)", "Normalized Counts", "QQQ Angle Comparison", "thetabqqq_overlay.png", ) n = len(datasets) fig, axes = plt.subplots( 1, n, figsize=(7 * n, 6) ) if n == 1: axes = [axes] for ax, data in zip(axes, datasets): dep_values = data.get("Edet", data.get("Elost", [])) h = ax.hist2d( data["sx3Z"], dep_values, bins=200 ) ax.set_title( f'{data["particle"]}\nEnergy Deposition vs SX3' ) ax.set_xlabel("SX3 Z") ax.set_ylabel("Energy Deposition (MeV)") fig.colorbar( h[3], ax=ax, label="Counts" ) plt.tight_layout() plt.savefig( f"{outdir}/Edet_vs_sx3_comparison.png", dpi=300 ) plt.show() save_overlay_hist2d( lambda data: data.get("qqqE", data.get("Eqqq", [])), lambda data: data["EdetQ"], "QQQ Initial Energy (MeV)", "QQQ Energy Deposition (MeV)", "QQQ Energy Deposition vs Initial Energy", "EdetQ_vs_qqqE_comparison.png", bins=200, ) save_overlay_hist2d( lambda data: data["thetabqqq"], lambda data: data.get("qqqE", data.get("Eqqq", [])), "QQQ Theta (deg)", "QQQ Initial Energy (MeV)", "QQQ Energy vs Theta", "QQQE_vs_theta_comparison.png", bins=200, ) save_overlay_hist2d( lambda data: data["qqqZ"], lambda data: data["EdetQ"], "QQQ Z", "QQQ Energy Deposition (MeV)", "QQQ Energy Deposition vs QQQ Position", "EdetQ_vs_qqqZ_comparison.png", bins=200, ) for data in datasets: mask = ( ~np.isnan(data.get("Edet", [])) & (data.get("Edet", []) > 0) & ~np.isnan(data.get("Eprop", [])) & ~np.isnan(data.get("thetab", [])) ) if not np.any(mask): print(f"No valid events for {data['particle']} when plotting Eprop vs Edet.") continue x = np.asarray(data["Edet"])[mask] y = np.sin(np.deg2rad(np.asarray(data["thetab"])[mask])) * np.asarray(data["Eprop"])[mask] plt.figure(figsize=(7, 6)) plt.hist2d(x, y, bins=200) plt.ylabel("PCEnergy x Sin(theta)") plt.xlabel("SX3 Energy Detected (MeV)") plt.title(f"{data['particle']} Energy Propagation Difference vs SX3 Energy Detected") plt.colorbar(label="Counts") plt.tight_layout() safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", data["particle"]) plt.savefig( f"{outdir}/{safe_name}_Eprop_vs_Edet.png", dpi=300 ) plt.show() plt.figure(figsize=(8, 6)) for data in datasets: plt.scatter( data["EA"], data["Esx3"], s=1, alpha=0.3, label=data["particle"] ) plt.xlabel("EA (MeV)") plt.ylabel("Esx3 (MeV)") plt.title("Anode vs SX3 Energy") plt.legend() plt.grid(True) plt.tight_layout() plt.savefig( f"{outdir}/EA_vs_Esx3_overlay.png", dpi=300 ) plt.show() all_Esx3 = [] all_Eprop = [] all_Edet = [] for data in datasets: mask = ( ~np.isnan(data["Edet"]) & (data["Edet"] > 0) & ~np.isnan(data["Eprop"]) & ~np.isnan(data["thetab"]) ) thetab = data["thetab"][mask] all_Esx3.append( data["Esx3"][mask] ) all_Eprop.append( np.sin(np.deg2rad(thetab)) * data["Eprop"][mask] ) all_Edet.append( data["Edet"][mask] ) all_EQQQ = [] all_EpropQQQ = [] all_EdetQQQ = [] all_thetabQQQ = [] for data in datasets: qqq_mask = ( ~np.isnan(data["EdetQ"]) & (data["EdetQ"] > 0) & ~np.isnan(data["EpropQ"]) & ~np.isnan(data["thetabqqq"]) ) qqq_thetab = data["thetabqqq"][qqq_mask] all_EQQQ.append( np.asarray(data.get("qqqE", data.get("Eqqq", [])))[qqq_mask] ) all_EpropQQQ.append( np.sin(np.deg2rad(qqq_thetab)) * data["EpropQ"][qqq_mask] ) all_EdetQQQ.append( data["EdetQ"][qqq_mask] ) all_thetabQQQ.append( qqq_thetab ) combined_Esx3 = np.concatenate(all_Esx3) combined_Eprop = np.concatenate(all_Eprop) combined_Edet = np.concatenate(all_Edet) combined_Esx3 += np.random.normal( 0, 0.08, len(combined_Esx3) ) mask = ( np.isfinite(combined_Esx3) & np.isfinite(combined_Eprop) & np.isfinite(combined_Edet) ) combined_Esx3 = combined_Esx3[mask] combined_Eprop = combined_Eprop[mask] combined_Edet = combined_Edet[mask] combined_EQQQ = np.concatenate(all_EQQQ) combined_EpropQQQ = np.concatenate(all_EpropQQQ) combined_EdetQQQ = np.concatenate(all_EdetQQQ) combined_thetabQQQ = np.concatenate(all_thetabQQQ) qqq_mask = ( np.isfinite(combined_EQQQ) & np.isfinite(combined_EpropQQQ) & np.isfinite(combined_EdetQQQ) & np.isfinite(combined_thetabQQQ) ) combined_EQQQ = combined_EQQQ[qqq_mask] combined_EpropQQQ = combined_EpropQQQ[qqq_mask] combined_EdetQQQ = combined_EdetQQQ[qqq_mask] combined_thetabQQQ = combined_thetabQQQ[qqq_mask] plt.figure(figsize=(8, 6)) plt.hist2d( combined_Esx3, combined_Eprop, bins=200 ) plt.xlabel("SX3 Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/Eprop_vs_Esx3.png", dpi=300 ) plt.show() plt.figure(figsize=(8, 6)) plt.hist2d( combined_Edet, combined_Eprop, bins=200) plt.xlabel("SX3 Energy Deposition (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.title("PCEnergy vs Energy Deposition in sx3") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/Eprop_vs_Edet.png", dpi=300) plt.show() plt.figure(figsize=(8, 6)) plt.hist2d( combined_EdetQQQ, combined_EpropQQQ, bins=200) plt.xlabel("QQQ Energy Deposition (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.title("PCEnergy vs Energy Deposition in QQQ") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/Eprop_vs_EdetQ.png", dpi=300) plt.show() plt.figure(figsize=(8, 6)) plt.hist2d( combined_thetabQQQ, combined_EQQQ, bins=200, ) plt.xlabel("QQQ Theta (deg)") plt.ylabel("QQQ Initial Energy (MeV)") plt.title("QQQ Energy vs Theta") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/QQQE_vs_theta.png", dpi=300, ) plt.show() plt.figure(figsize=(8, 6)) plt.hist2d( combined_EQQQ, combined_EdetQQQ, bins=200, ) plt.xlabel("QQQ Initial Energy (MeV)") plt.ylabel("QQQ Energy Deposition (MeV)") plt.title("QQQ Energy Deposition vs Initial Energy") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/EdetQ_vs_qqqE.png", dpi=300, ) plt.show() plt.figure(figsize=(8, 6)) plt.hist2d( combined_EQQQ, combined_EpropQQQ, bins=200, ) plt.xlabel("QQQ Initial Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.title("PCEnergy x Sin(theta) vs QQQ Initial Energy") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{outdir}/Eprop_vs_QQQE.png", dpi=300, ) plt.show() plt.figure( figsize=(14, 6), facecolor="white" ) plt.hist2d( combined_Esx3, combined_Eprop, bins=[500, 500], cmap=yellow_jet, cmin=1 ) plt.xlabel("SX3 Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") #plt.yscale("log") cbar = plt.colorbar() cbar.set_label("Counts") plt.minorticks_on() plt.grid(False) plt.tight_layout() plt.savefig( f"{outdir}/ROOT_style_plot.png", dpi=300, facecolor="white" ) plt.show() plt.figure( figsize=(14, 6), facecolor="white" ) plt.hist2d( combined_EQQQ, combined_EpropQQQ, bins=[500, 500], cmap=yellow_jet, cmin=1 ) plt.xlabel("QQQ Initial Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") cbar = plt.colorbar() cbar.set_label("Counts") plt.minorticks_on() plt.grid(False) plt.tight_layout() plt.savefig( f"{outdir}/ROOT_style_plot_QQQ.png", dpi=300, facecolor="white" ) plt.show() print( f"Completed plotting " f"{len(datasets)} datasets." ) if __name__ == "__main__": MyInteractiveApp().cmdloop()