#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Wed May 20 13:32:14 2026 @author: jamesszalkie """ 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.0026, 40, "alpha"] proton_data = [1, 1.0078, 20, "proton"] deuteron_data = [1, 2.014102, 30, "deuteron"] dA = 3.2 #cm dC = 4.2 #cm dsx3 = 8.8 #cm particles = { "alpha": alpha_data, "proton": proton_data, "deuteron": deuteron_data } interp_cache = {} def make_E_vs_x(z, mass_u, emax_mev, label, npoints, P_TORR, TEMP_K): # GAS SETUP R = 8.3144 P_TORR = P_TORR TEMP_K = TEMP_K # Gas density p_pa = P_TORR * 133.322 molar_density = p_pa / (R * TEMP_K) m_he = 4.0026 m_c = 12.0000 m_o = 15.9949 m_mix_avg = (0.96 * m_he) + (0.04 * (m_c + 2*m_o)) rho_g_cm3 = (molar_density * m_mix_avg) / 1e6 print(f"Gas density = {rho_g_cm3:.6e} g/cm^3") # MATERIAL material_def = [ (m_he, 2, 0.96), (m_c, 6, 0.04), (m_o, 8, 0.08) ] gas_mix = catima.Material(material_def) gas_mix.density(rho_g_cm3) projectile = catima.Projectile(mass_u, z) # Energy grid E = np.linspace(0.01, emax_mev, npoints) # Stopping power array S_mass = np.zeros_like(E) for i, energy in enumerate(E): projectile.T(energy / mass_u) # MeV / (g/cm^2) S_mass[i] = catima.dedx(projectile, gas_mix) # Convert to MeV/cm S_linear = S_mass * rho_g_cm3 # Sort descending energy sort_idx = np.argsort(E)[::-1] E = E[sort_idx] S_linear = S_linear[sort_idx] # Integrate dx/dE = 1/S(E) invS = 1.0 / S_linear x = cumulative_trapezoid( invS, E, initial=0 ) x = -x # Output table output = pd.DataFrame({ "Distance_cm": x, "Energy_MeV": E }) outfile = f"E_vs_x_{label}.dat" output.to_csv( outfile, sep='\t', index=False ) print(f"Saved: {outfile}") 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): if particle in interp_cache: return interp_cache[particle] filename = f"E_vs_x_{particle}.dat" 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[particle] = (E_of_x, x_of_E) return E_of_x, x_of_E def energy_loss(particle, Ei, dl): E_of_x, x_of_E = get_interpolators(particle) xi = x_of_E(Ei) xf = xi + dl Ef = E_of_x(xf) return np.maximum(Ef, 0.0) def energy_reconstruction(particle, Ef, dl): E_of_x, x_of_E = get_interpolators(particle) xf = x_of_E(Ef) xi = xf - dl Ei = E_of_x(xi) return np.maximum(Ei, 0.0) def energy_distance(particle, Ei, Ef): _, x_of_E = get_interpolators(particle) xi = x_of_E(Ei) xf = x_of_E(Ef) * .686 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] # <-- THIS is the correct way 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 load_tree_arrays(tree, treename, max_events=None): 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) #mask = (~np.isnan(data["sx3X"])) & (~np.isnan(data["Tb"])) #data = {key: value[mask] for key, value in data.items()} Ei = data["Tb"] Einan = np.isnan(Ei).sum() print(f"sx3 NaN entries: {Einan}") theta = np.radians(data["thetab"]) vX = data["vX"] vY = data["vY"] vZ = data["vZ"] #if treename == 'tree1': sx3X = data["sx3X"] sx3Y = data["sx3Y"] sx3Z = data["sx3Z"] mask = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z) qqqX = data["qqqX"] qqqY = data["qqqY"] qqqZ = data["qqqZ"] qqqnan = np.isnan(qqqZ).sum() print(f"qqq NaN entries: {qqqnan}") print(f"Total: {Einan + qqqnan}") mask2 = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z) #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) Eisx3 = Ei[mask] theta = theta[mask] vXsx3 = vX[mask] vYsx3 = vY[mask] vZsx3 = vZ[mask] Eiqqq = Ei[mask2] vXqqq = vX[mask2] vYqqq = vY[mask2] vZqqq = vZ[mask2] #if treename == 'tree1': sx3X = sx3X[mask] sx3Y = sx3Y[mask] sx3Z = sx3Z[mask] dsx3 = calculate_distance_tree1(vXsx3, vYsx3, vZsx3, sx3X, sx3Y, sx3Z) qqqX = qqqX[mask2] qqqY = qqqY[mask2] qqqZ = qqqZ[mask2] dqqq = calculate_distance_tree1(vXqqq, vYqqq, vZqqq, qqqX, qqqY, 128) #else: # dsx3 = calculate_distance_tree2(vZsx3, theta, z_max=z_max) sin_theta = np.sin(theta) sin_theta = np.where(sin_theta != 0, sin_theta, 1e-10) radii = np.array([3.7, 4.3]) dA = (radii[0] - np.sqrt((vXsx3/10)**2 + (vYsx3/10)**2))/ sin_theta dC = (radii[1] - np.sqrt((vXsx3/10)**2 + (vYsx3/10)**2))/ sin_theta # 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") EA = energy_loss(particle, Eisx3, dA) EC = energy_loss(particle, Eisx3, dC) Esx3 = energy_loss(particle, Eisx3, dsx3) Eqqq = energy_loss(particle, Eiqqq, dqqq) Elost = Eisx3 - Esx3 Elostqqq = Eiqqq - Eqqq Eprop = EA - EC return { "particle": particle, "Ei": Eisx3, "thetabi": data["thetab"], "sx3Z": sx3Z, "EA": EA, "EC": EC, "Esx3": Esx3, "Elost": Elost, "Elostqqq": Elostqqq, "Eprop": Eprop, "dA": dA, "dC": dC, "thetab": np.degrees(theta), "sx3X": sx3X, "sx3Y": sx3Y, "qqqX": qqqX, "qqqY": qqqY, "qqqZ": qqqZ, "Eqqq": Eqqq } def process_file(filename, treename, particle=None, max_events=None): tree = uproot.open(filename)[f"{treename}"] if particle is None: lower = filename.lower() if "proton" in lower: particle = "proton" elif "alpha" in lower: particle = "alpha" else: 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 = 500 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 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, 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()}") 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"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] Ei = float(args[1]) dl = float(args[2]) Ei_offset = Ei * 1.1 table_specs = f"{particle} {Ei_offset}" try: self.do_make_table(table_specs) Ef = energy_loss(particle, 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] Ef = float(args[1]) dl = float(args[2]) try: Ei = energy_reconstruction(particle, 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] Ei = float(args[1]) Ef = float(args[2]) dE = Ei - Ef try: dl = energy_distance(particle, Ei, Ef) if dl > 0: print(f"\nChange in energy: {dE:.6f} MeV") print(f"Distance travelled: {dl:.6f} cm\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_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}") treename = self.tree.name if hasattr(self.tree, 'name') else "tree1" branches = ["Tb", "thetab", "sx3Z", "vX", "vY", "vZ", "sx3X", "sx3Y"] Tb_key = "Tb" theta_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 ) global Ei Ei = data[Tb_key] global sx3Z sx3Z = data["sx3Z"] theta = np.radians(data[theta_key]) vX = data["vX"] vY = data["vY"] vZ = data["vZ"] if treename == 'tree1': sx3X = data["sx3X"] sx3Y = data["sx3Y"] mask = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z) else: mask = ~np.isnan(Ei) & ~np.isnan(theta) Ei = Ei[mask] theta = theta[mask] sx3Z = sx3Z[mask] vZ_masked = vZ[mask] sin_theta = np.sin(theta) sin_theta = np.where(sin_theta != 0, sin_theta, 1e-10) if treename == 'tree1': vX_masked = vX[mask] vY_masked = vY[mask] sx3X_masked = sx3X[mask] sx3Y_masked = sx3Y[mask] dsx3 = calculate_distance_tree1(vX_masked, vY_masked, vZ_masked, sx3X_masked, sx3Y_masked, sx3Z) else: dsx3 = calculate_distance_tree2(vZ_masked, theta, z_max=34.86) radii = np.array([3.7, 4.2]) dA = radii[0] / sin_theta dC = radii[1] / sin_theta print("Calculating energy losses...") global EA EA = energy_loss(particle, Ei, dA) global EC EC = energy_loss(particle, Ei, dC) global Esx3 Esx3 = energy_loss(particle, Ei, dsx3) global Eprop Eprop = EA - EC global Elost Elost = Ei - Esx3 print("Analysis complete") print(f"Processed events: {len(Ei)}") 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[mask] = EA EC_full[mask] = EC Esx3_full[mask] = Esx3 Eprop_full[mask] = Eprop Elost_full[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 # 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}") 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"] sx3X = data["sx3X"] sx3Y = data["sx3Y"] sx3Z = data["sx3Z"] EA = data["EA"] EC = data["EC"] Esx3 = data["Esx3"] Elost = data["Elost"] Eprop = data["Eprop"] dA = data["dA"] dC = data["dC"] qqqX = data["qqqX"] qqqY = data["qqqY"] qqqZ = data["qqqZ"] qqqE = data["Eqqq"] Elostqqq = data["Elostqqq"] 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) 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) 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() else: print("No valid SX3 points for 3D plot.") mask1 = ~np.isnan(Ei) & ~np.isnan(thetab) plt.figure(figsize=(7,6)) plt.hist2d(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() mask1 = ~np.isnan(qqqE) & ~np.isnan(thetab) plt.figure(figsize=(7,6)) plt.hist2d(qqqE[mask1], (np.sin(np.deg2rad(thetab)) * 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) 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(thetab) plt.figure(figsize=(7,6)) plt.hist2d(thetab[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, 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(Ei, bins=100, histtype='step', label='Vertex Energy') plt.hist(Esx3[Esx3 != 0], 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, bins=100, label='dA', histtype='step') plt.hist(dC, 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: plt.figure(figsize=(7,6)) plt.hist2d(sx3Z, Elost, 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 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: plt.figure(figsize=(7,6)) plt.hist2d(sx3Z, Eprop, 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) plt.figure(figsize=(7,6)) plt.hist2d(Esx3[mask1], Eprop[mask1] * data["thetab"][mask1], bins=200) plt.ylabel("PCEnergy") 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: #except ValueError: # print("Value error") 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: print(f"Creating histograms for {len(branch_names)} branches...") all_branches = self.tree.arrays(branch_names, library="np", entry_stop=max_events) # Keep only events with thetab >= 45 #mask = (all_branches["Tb"] > 0) & (all_branches["qqqTb"] > 0) 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)) plt.hist(values, bins=100) 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 = [] for file in files: try: # If you want to combine tree1 + tree2: tree1 = process_file( os.path.join("..", "Armory", file), "tree1" ) try: tree2 = process_file( os.path.join("..", "Armory", file), "tree2" ) data = { "particle": f"{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 # -------------------------------------------------- # Elost overlay # -------------------------------------------------- plt.figure(figsize=(8, 6)) for data in datasets: plt.hist( data["Elost"], bins=200, histtype="step", linewidth=2, density=True, label=data["particle"] ) plt.xlabel("Energy Loss (MeV)") plt.ylabel("Normalized Counts") plt.title("Energy Loss Comparison") plt.legend() plt.grid(True) plt.tight_layout() plt.savefig( f"{outdir}/Elost_overlay.png", dpi=300 ) plt.show() # -------------------------------------------------- # SX3 position overlay # -------------------------------------------------- plt.figure(figsize=(8, 6)) for data in datasets: plt.hist( data["sx3Z"], bins=150, histtype="step", linewidth=2, density=True, label=data["particle"] ) plt.xlabel("SX3 Z") plt.ylabel("Normalized Counts") plt.title("SX3 Position Comparison") plt.legend() plt.grid(True) plt.tight_layout() plt.savefig( f"{outdir}/sx3Z_overlay.png", dpi=300 ) plt.show() # -------------------------------------------------- # Individual Elost vs SX3 plots # -------------------------------------------------- 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): h = ax.hist2d( data["sx3Z"], data["Elost"], bins=200 ) ax.set_title( f'{data["particle"]}\nElost vs SX3' ) ax.set_xlabel("SX3 Z") ax.set_ylabel("Energy Loss (MeV)") fig.colorbar( h[3], ax=ax, label="Counts" ) plt.tight_layout() plt.savefig( f"{outdir}/Elost_vs_sx3_comparison.png", dpi=300 ) plt.show() # -------------------------------------------------- # EA vs Esx3 overlay scatter # -------------------------------------------------- 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() # -------------------------------------------------- # Combined PCE vs SX3 # -------------------------------------------------- all_Esx3 = [] all_Eprop = [] for data in datasets: mask = data["Esx3"] > 1 thetab = np.deg2rad( data["thetab"][mask] ) all_Esx3.append( data["Esx3"][mask] ) all_Eprop.append( data["Eprop"][mask] * np.sin(thetab) * 3 ) combined_Esx3 = np.concatenate(all_Esx3) combined_Eprop = np.concatenate(all_Eprop) combined_Esx3 += np.random.normal( 0, 0.08, len(combined_Esx3) ) mask = ( np.isfinite(combined_Esx3) & np.isfinite(combined_Eprop) ) combined_Esx3 = combined_Esx3[mask] combined_Eprop = combined_Eprop[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=(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)") 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() print( f"Completed plotting " f"{len(datasets)} datasets." ) #exec(open("PCEnergyAnalysis.py").read()) if __name__ == "__main__": MyInteractiveApp().cmdloop()