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