#!/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 # 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 = ["Tb", "thetab", "sx3Z", "vX", "vY", "vZ", "sx3ID", "aID", "cID", "aDist", "cDist"] if treename == 'tree1': branches.extend(["sx3X", "sx3Y"]) 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 = data["thetab"] >= 0 #data = {key: value[mask] for key, value in data.items()} Ei = data["Tb"] 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) 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) Ei = Ei[mask] theta = theta[mask] vX = vX[mask] vY = vY[mask] vZ = vZ[mask] if treename == 'tree1': sx3X = sx3X[mask] sx3Y = sx3Y[mask] sx3Z = sx3Z[mask] dsx3 = calculate_distance_tree1(vX, vY, vZ, sx3X, sx3Y, sx3Z) else: dsx3 = calculate_distance_tree2(vZ, 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((vX/10)**2 + (vY/10)**2))/ sin_theta dC = (radii[1] - np.sqrt((vX/10)**2 + (vY/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 distance_mask = (dA > 0.1) & (dC > 0.1) & (sx3ID >= 0) & aID_valid & cID_valid Ei = Ei[distance_mask] theta = theta[distance_mask] dA = dA[distance_mask] dC = dC[distance_mask] dsx3 = dsx3[distance_mask] vX = vX[distance_mask] vY = vY[distance_mask] vZ = vZ[distance_mask] 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, Ei, dA) EC = energy_loss(particle, Ei, dC) Esx3 = energy_loss(particle, Ei, dsx3) Elost = Ei - Esx3 Eprop = EA - EC return { "particle": particle, "Ei": Ei, "sx3Z": sx3Z, "EA": EA, "EC": EC, "Esx3": Esx3, "Elost": Elost, "Eprop": Eprop, "dA": dA, "dC": dC, "thetab": np.degrees(theta) } 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"] sx3Z = data["sx3Z"] EA = data["EA"] EC = data["EC"] Esx3 = data["Esx3"] Elost = data["Elost"] Eprop = data["Eprop"] dA = data["dA"] dC = data["dC"] 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})") 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() try: plt.figure(figsize=(7,5)) plt.hist(sx3Z, bins=100) plt.xlabel("SX3 Z") plt.ylabel("Counts") plt.title(f"{particle} ({treename}) SX3 Position Distribution") plt.grid(True) plt.tight_layout() plt.savefig(f"{base}/sx3Z_hist.png", dpi=300) plt.show() except ValueError: print("Value error") 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: plt.figure(figsize=(7,6)) plt.hist2d(EA, Esx3, 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], 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: mask1 = ~np.isnan(Ei) & (Esx3 > 0) plt.figure(figsize=(7,6)) plt.hist2d(data["thetab"], Ei, 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() 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 #thetab_mask = all_branches["thetab"] >= 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 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) if len(args) < 2: try: args = ['SimAnasenProton.root', 'SimAnasenAlpha.root'] except: print("Usage: make dual plots proton_data.root alpha_data.root") return file1 = args[0] file2 = args[1] #tree1_name = args[2] if len(args) > 2 else 'tree2' #tree2_name = args[3] if len(args) > 3 else 'tree1' outdir = "dual_plots" # load both trees for file1 and combine their arrays into a single dataset data1_tree1 = process_file(os.path.join("..", "Armory", file1), "tree1") data1_tree2 = process_file(os.path.join("..", "Armory", file1), "tree2") # concatenate matching array entries data1 = {"particle": f"{data1_tree1['particle']}_combined"} for key in data1_tree1: if key == "particle": continue try: data1[key] = np.concatenate([data1_tree1[key], data1_tree2[key]]) except Exception: # fallback: prefer tree1 value if concat fails data1[key] = data1_tree1[key] data2 = process_file(os.path.join("..", "Armory", file2), "tree1") #print(f"File one {file1} ({tree1_name}) length {len(data1['Ei'])} \nFile two {file2} ({tree2_name}) length {len(data2['Ei'])}") os.makedirs(outdir, exist_ok=True) print(f"Saving plots to: {outdir}") #Overlay histogram: Elost plt.figure(figsize=(8,6)) plt.hist(data1["Elost"],bins=200,histtype='step',linewidth=2,density=True,label=data1["particle"]) plt.hist(data2["Elost"],bins=200,histtype='step',linewidth=2,density=True,label=data2["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() #Overlay histogram: sx3Z plt.figure(figsize=(8,6)) plt.hist(data1["sx3Z"],bins=150,histtype='step',linewidth=2,density=True,label=data1["particle"]) plt.hist(data2["sx3Z"],bins=150,histtype='step',linewidth=2,density=True,label=data2["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() try: #Side-by-side 2D plots fig, axes = plt.subplots(1, 2, figsize=(14,6)) h1 = axes[0].hist2d(data1["sx3Z"],data1["Elost"],bins=200) axes[0].set_title(f'{data1["particle"]} Elost vs SX3') axes[0].set_xlabel("SX3 Z") axes[0].set_ylabel("Energy Loss (MeV)") h2 = axes[1].hist2d(data2["sx3Z"],data2["Elost"],bins=200) axes[1].set_title(f'{data2["particle"]} Elost vs SX3') axes[1].set_xlabel("SX3 Z") axes[1].set_ylabel("Energy Loss (MeV)") fig.colorbar(h1[3], ax=axes[0], label="Counts") fig.colorbar(h2[3], ax=axes[1], label="Counts") plt.tight_layout() plt.savefig(f"{outdir}/Elost_vs_sx3_comparison.png", dpi=300) plt.show() except: print("Error with side-by-side plots") #EA vs Esx3 overlay scatter plt.figure(figsize=(8,6)) plt.scatter(data1["EA"],data1["Esx3"],s=1,alpha=0.3,label=data1["particle"]) plt.scatter(data2["EA"],data2["Esx3"],s=1,alpha=0.3,label=data2["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() #PCE vs SiE mask1 = data1["Esx3"] > 1 mask2 = data2["Esx3"] > 1 thetab1 = np.deg2rad(data1["thetab"][mask1]) thetab2 = np.deg2rad(data2["thetab"][mask2]) #theta smear (spatial uncertainty) if True: sigma = 0.2 * np.sqrt(np.maximum(thetab1, 0)) thetab1 += np.random.normal(0,sigma,len(thetab1)) thetab1 += np.random.normal(0,0.02*np.sqrt(thetab1),len(thetab1)) sigma = 0.2 * np.sqrt(np.maximum(thetab2, 0)) thetab2 += np.random.normal(0,sigma,len(thetab2)) thetab2 += np.random.normal(0,0.02*np.sqrt(thetab2),len(thetab2)) combined_Esx3 = np.concatenate([ (data1["Esx3"][mask1]), data2["Esx3"][mask2] ]) combined_Eprop = np.concatenate([ data1["Eprop"][mask1] * np.sin(thetab1) * 3, data2["Eprop"][mask2] * np.sin(thetab2) * 3 ]) combined_Esx3 = ( combined_Esx3 + np.random.normal(0,0.08,len(combined_Esx3)) ) #Artifical smear if False: sigma = 0.02 * np.sqrt(np.maximum(combined_Eprop, 0)) combined_Eprop += np.random.normal( 0, sigma, len(combined_Eprop) ) combined_Eprop += np.random.normal( 0, 0.02*np.sqrt(combined_Eprop), len(combined_Eprop) ) mask = (np.isfinite(combined_Esx3)& np.isfinite(combined_Eprop)) combined_Esx3 = combined_Esx3[mask] combined_Eprop = combined_Eprop[mask] #combined_Eprop = combined_Eprop * .686 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.xlim(0,30) #plt.ylim(0,0.45) 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], #range=[[0,35],[0,0.6]], cmap=yellow_jet, cmin=1) plt.margins(0) plt.xlabel("SX3 Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.xlim(0,35) plt.ylim(0,0.6) 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() if False: plt.figure(figsize=(7,6)) plt.hist2d( combined_Esx3, combined_Eprop, bins=100, range=[[0,30],[0,0.45]], cmap='viridis' ) mask = data1["Esx3"] > 0 power_fit_and_plot( data1["Esx3"][mask], data1["Eprop"][mask], label=data1["particle"], color='red' ) mask2 = data2["Esx3"] > 0 power_fit_and_plot( data2["Esx3"][mask2], data2["Eprop"][mask2], label=data2["particle"], color='cyan' ) plt.ylabel("PCEnergy") plt.xlabel("SX3 Energy (MeV)") plt.title( f'{data1["particle"]} + {data2["particle"]} ' 'Energy Propagation Difference vs SX3 Energy' ) plt.colorbar(label="Counts") #plt.xlim(0,30) #plt.ylim(0,0.45) plt.legend() plt.tight_layout() plt.savefig(f"{outdir}/Combined_Eprop_vs_Esx3_fit.png",dpi=300) plt.show() print("Dual plotting complete.") #exec(open("PCEnergyAnalysis.py").read()) if __name__ == "__main__": MyInteractiveApp().cmdloop()