#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Tue Aug 18 09:57:50 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 import shutil alpha_data = [2, 4.0015, 40, "alpha"] proton_data = [1, 1.0073, 20, "proton"] deuteron_data = [1, 2.014102, 30, "deuteron"] interp_cache = {} particles = { "alpha": alpha_data, "proton": proton_data, "deuteron": deuteron_data } 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 normalize_particle_label(particle): """Return the canonical table label for a particle or isotope name.""" if isinstance(particle, str): try: _, _, _, label = resolve_particle(particle) return label except Exception: return particle.strip() return str(particle) 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 """ material_def = [ (m_he, 2, 0.96), (m_c, 6, 0.04), (m_o, 8, 0.08) ] """ material_def = [(m_he, 2, 1.0)] #m_mix_avg = 0.96 * m_he + 0.04 * (m_c + 2 * m_o) m_mix_avg = 1.0 * m_he rho_g_cm3 = (molar_density * m_mix_avg) / 1e6 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) elif medium == "kapton": # Kapton (C22H10N2O5) # Density: 1.42 g/cm3 m_h = 1.0078 m_n = 14.0067 m_c = 12.0000 m_o = 15.9949 kapton_molar_mass = 22*m_c + 10*m_h + 2*m_n + 5*m_o material_def = [ (m_c, 6, 22/kapton_molar_mass * kapton_molar_mass / m_c), (m_h, 1, 10/kapton_molar_mass * kapton_molar_mass / m_h), (m_n, 7, 2/kapton_molar_mass * kapton_molar_mass / m_n), (m_o, 8, 5/kapton_molar_mass * kapton_molar_mass / m_o) ] rho_g_cm3 = 1.42 gas = catima.Material(material_def) gas.density(rho_g_cm3) elif medium == "mylar": # Mylar (C10H8O4, polyethylene terephthalate) # Density: 1.39 g/cm3 m_h = 1.0078 m_c = 12.0000 m_o = 15.9949 mylar_molar_mass = 10*m_c + 8*m_h + 4*m_o material_def = [ (m_c, 6, 10/mylar_molar_mass * mylar_molar_mass / m_c), (m_h, 1, 8/mylar_molar_mass * mylar_molar_mass / m_h), (m_o, 8, 4/mylar_molar_mass * mylar_molar_mass / m_o) ] rho_g_cm3 = 1.39 gas = catima.Material(material_def) 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 / mass_u) S_mass[i] = catima.dedx(projectile, gas) # MeV/(g/cm^2) 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): canonical_particle = normalize_particle_label(particle) cache_key = (canonical_particle.lower(), medium.lower()) if cache_key in interp_cache: return interp_cache[cache_key] candidate_labels = [canonical_particle] if isinstance(particle, str): raw_particle = particle.strip() candidate_labels.extend([ raw_particle, canonical_particle.replace("-", "") ]) filename = None seen_candidates = set() for label in candidate_labels: normalized_label = label.strip() if not normalized_label: continue key = normalized_label.lower() if key in seen_candidates: continue seen_candidates.add(key) candidate_filename = get_loss_table_path(medium, normalized_label) if os.path.exists(candidate_filename): filename = candidate_filename break if filename is None: filename = get_loss_table_path(medium, canonical_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}") match = re.match(r"(\d+)[-\s]?([a-zA-Z]+)$", name) if match: A = int(match.group(1)) element_symbol = match.group(2).capitalize() 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 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 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 class MyInteractiveApp(cmd.Cmd): def __init__(self): super().__init__() # Initial value set when the script starts 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("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_exit(self, arg): """Exits the application.""" print("Closing application...") return True # Returning True stops the cmdloop() def do_T(self, arg): """Print value of T""" print(self.T) def do_P(self, arg): """Print value of P (pressure)""" print(self.P) 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_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, 10000, 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 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 * 10):.6f} mm") 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_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 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_root_histograms(self, arg): self.do_set_tree("") if arg: args = shlex.split(arg) particle = args[0] else: particle = "proton" base = f"{particle}_roothist_plots" if os.path.exists(base): shutil.rmtree(base) 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)) 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 self.tree.keys(): branch = branch.decode() if isinstance(branch, bytes) else str(branch) try: values = self.tree[branch].array( library="np", ) values = np.asarray(values, dtype=float) except Exception as e: print(f"Skipping branch {branch}: {e}") continue if values.ndim != 1: print(f"Skipping non-scalar branch: {branch}") continue if branch == "qqqZ": continue values = values[np.isfinite(values)] if values.size == 0: print(f"Skipping empty branch: {branch}") continue plt.figure(figsize=(7, 5)) plt.hist( values, bins=100, histtype="stepfilled" ) 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) filename = f"{base}/{safe_name}_hist.png" counter = 1 while os.path.exists(filename): filename = f"{base}/{safe_name}_{counter}_hist.png" counter += 1 plt.savefig(filename, dpi=300) plt.close() def do_make_plots(self, arg): 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 base = f"{particle}_plots" if os.path.exists(base): shutil.rmtree(base) os.makedirs(base, exist_ok=True) 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 data = self.tree.arrays(library="pd") #Combine sx3 and QQQ data data['X'] = data['sx3X'].fillna(data['qqqX']) data['Y'] = data['sx3Y'].fillna(data['qqqY']) data['Z'] = data['sx3Z'].fillna(data['qqqZ']) data['Ei'] = data['Tb'].fillna(data['qqqTb']) data['dl'] = np.sqrt(((data['X'] - data['vX']) ** 2) + ((data['Y'] - data['vY']) ** 2) + ((data['Z'] - data['vZ']) ** 2)) * .1 sx3dl = np.sqrt(((data['sx3X'] - data['vX']) ** 2) + ((data['sx3Y'] - data['vY']) ** 2) + ((data['sx3Z'] - data['vZ']) ** 2)) * .1 qqqdl = np.sqrt(((data['qqqX'] - data['vX']) ** 2) + ((data['qqqY'] - data['vY']) ** 2) + ((data['qqqZ'] - data['vZ']) ** 2)) * .1 Adl = np.sqrt(((data['aX'] - data['vX']) ** 2) + ((data['aY'] - data['vY']) ** 2) + ((data['aZ'] - data['vZ']) ** 2)) * .1 Cdl = np.sqrt(((data['cX'] - data['vX']) ** 2) + ((data['cY'] - data['vY']) ** 2) + ((data['cZ'] - data['vZ']) ** 2)) * .1 Ei = np.asarray(data['Ei'], dtype=float) dl = np.asarray(data['dl'], dtype=float) data['EatDet'] = energy_loss(particle, "He", Ei, dl) data['Esx3'] = energy_loss(particle, "He", data['Tb'], sx3dl) data['Eqqq'] = energy_loss(particle, "He", data['qqqTb'], qqqdl) data['EA'] = energy_loss(particle, "He", data['Ei'], Adl) data['EC'] = energy_loss(particle, "He", data['Ei'], Cdl) #Filter out events with 0 final energy, does not reach detector data = data[data['EA'] > .2].copy() x = np.asarray(data['X'], dtype=float) y = np.asarray(data['Y'], dtype=float) z = np.asarray(data['Z'], dtype=float) #print("\nHit detector counts:") #print(data['hitDetector'].value_counts()) 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=np.asarray(data['EatDet'], dtype=float), cmap='viridis', s=3, alpha=0.8 ) ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") ax.set_title(f"{particle} SX3 3D Hit Map") cb = plt.colorbar(sc, ax=ax, pad=0.1) cb.set_label("Energy (MeV)") plt.tight_layout() out_file = f"{base}/3D_energy.png" plt.savefig(out_file, dpi=300) plt.show() Ax = np.asarray(data['aX'], dtype=float) Ay = np.asarray(data['aY'], dtype=float) Az = np.asarray(data['aZ'], dtype=float) Cx = np.asarray(data['cX'], dtype=float) Cy = np.asarray(data['cY'], dtype=float) Cz = np.asarray(data['cZ'], dtype=float) if x.size > 0 and True: fig = plt.figure(figsize=(8,6)) ax = fig.add_subplot(111, projection='3d') sc = ax.scatter( Ax, Ay, Az, c=np.asarray(data['EA'], dtype=float), cmap='viridis', s=3, alpha=0.8 ) ax.scatter( Cx, Cy, Cz, c=np.asarray(data['EC'], dtype=float), s=3, alpha=0.8, marker='^', label='Cathode' ) ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") ax.set_title(f"{particle} SX3 3D Hit Map") cb = plt.colorbar(sc, ax=ax, pad=0.1) cb.set_label("Energy (MeV)") plt.tight_layout() out_file = f"{base}/Anode_Cathode_Hits.png" plt.savefig(out_file, dpi=300) plt.show() EPC = abs(data['EA'] - data['EC']) plt.figure(figsize=(8, 6)) plt.hist(data['Esx3'], bins=200) mask = ~np.isnan(data['Esx3']) plt.figure(figsize=(8, 6)) plt.hist2d( data['Esx3'][mask] * np.sin(np.deg2rad(data['thetab'][mask])), EPC[mask], bins=200 ) plt.xlabel("SX3 Energy (MeV)") plt.ylabel("PCEnergy x Sin(theta)") plt.colorbar(label="Counts") plt.tight_layout() plt.savefig( f"{base}/Eprop_vs_Esx3.png", dpi=300 ) plt.show() plt.figure(figsize=(8,6)) plt.hist(data['Esx3'], bins=200) plt.tight_layout() plt.show() mask = ~np.isnan(data['Esx3']) plt.figure(figsize=(8,6)) plt.hist2d(data['EPC'][mask], data['Esx3'][mask], bins=200) plt.show() data.to_csv('data.csv', index=False) if __name__ == "__main__": MyInteractiveApp().cmdloop()