ANASEN_analysis/ELoss/PCEnergyAnalysis.py
2026-06-03 10:35:26 -04:00

1551 lines
45 KiB
Python

#!/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
import matplotlib as mpl
# 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 = 400
TEMP_K = 293.15
R = 8.3144
# 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_loss_angular(particle, Ei, theta):
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)
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.2])
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": data["thetab"]
}
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 <z> <mass u> <maximum E value> <particle name>"""
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 <max energy (optional) >
Ex: >> make table Co60 <max energy (optional) >
Ex: >> make table N17 <max energy (optional) >"""
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 <particle> <initial energy MeV> <distance travelled cm>"""
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 <particle> <final energy MeV> <distance travelled cm>"""
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 <particle> <initial energy> <final energy MeV>"""
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
import awkward as ak
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,
#"ak": ak,
"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):
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, ClassSX3dC)
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)
mask = data["thetab"] >= 0
Ei = data["Ei"][mask]
sx3Z = data["sx3Z"][mask]
EA = data["EA"][mask]
EC = data["EC"][mask]
Esx3 = data["Esx3"][mask]
Elost = data["Elost"][mask]
Eprop = data["Eprop"][mask]
dA = data["dA"][mask]
dC = data["dC"][mask]
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
thetab_mask = mask
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[thetab_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
self.do_set_tree("tree1")
data1 = prepare_tree_data(self.tree, "tree1", particle, max_events=None)
self.do_set_tree("tree2")
data2 = prepare_tree_data(self.tree, "tree2", particle, max_events=None)
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"
data1 = process_file(os.path.join("..", "Armory", file1), "tree1")
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()
#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()
#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"] > 0
mask2 = data2["Esx3"] > 0
combined_Esx3 = np.concatenate([
(data1["Esx3"][mask1]),
data2["Esx3"][mask2]
])
combined_Eprop = np.concatenate([
data1["Eprop"][mask1] * 3,
data2["Eprop"][mask2] * 2.7
])
combined_Esx3 = (
combined_Esx3
+ np.random.normal(0,0.08,len(combined_Esx3))
)
if True:
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]# * 3
#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")
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=(11,6), facecolor='white')
plt.hist2d(
combined_Esx3,
combined_Eprop,
bins=[500,500],
#range=[[0,35],[0,0.6]],
cmap=yellow_jet,
cmin=1
)
plt.xlabel("SX3 Energy (MeV)")
plt.ylabel("PCEnergy")
plt.xlim(0,30)
plt.ylim(0,0.45)
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] * .686,
label=data1["particle"],
color='red'
)
mask2 = data2["Esx3"] > 0
power_fit_and_plot(
data2["Esx3"][mask2],
data2["Eprop"][mask2] * .686,
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()