892 lines
27 KiB
Python
892 lines
27 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Tue Aug 18 09:57:50 2026
|
|
|
|
@author: jamesszalkie
|
|
"""
|
|
|
|
import gc
|
|
import numpy as np
|
|
import pandas as pd
|
|
from scipy.interpolate import interp1d
|
|
import uproot
|
|
import pycatima as catima
|
|
from scipy.integrate import cumulative_trapezoid
|
|
#matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import cmd
|
|
import shlex
|
|
import textwrap
|
|
import os
|
|
import periodictable as pt
|
|
import re
|
|
from matplotlib.colors import LinearSegmentedColormap
|
|
from mpl_toolkits.mplot3d import Axes3D
|
|
import nbformat as nbf
|
|
import shutil
|
|
|
|
alpha_data = [2, 4.0015, 40, "alpha"]
|
|
proton_data = [1, 1.0073, 20, "proton"]
|
|
deuteron_data = [1, 2.014102, 30, "deuteron"]
|
|
|
|
interp_cache = {}
|
|
|
|
particles = {
|
|
"alpha": alpha_data,
|
|
"proton": proton_data,
|
|
"deuteron": deuteron_data
|
|
}
|
|
|
|
|
|
def get_loss_table_path(medium, particle_label):
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
return os.path.join(script_dir, f"{medium}Loss", f"E_vs_x_{particle_label}.dat")
|
|
|
|
|
|
def normalize_particle_label(particle):
|
|
"""Return the canonical table label for a particle or isotope name."""
|
|
if isinstance(particle, str):
|
|
try:
|
|
_, _, _, label = resolve_particle(particle)
|
|
return label
|
|
except Exception:
|
|
return particle.strip()
|
|
return str(particle)
|
|
|
|
|
|
def clear_interpolator_cache():
|
|
"""Drop cached energy-loss interpolators and force Python to release memory."""
|
|
global interp_cache
|
|
interp_cache.clear()
|
|
gc.collect()
|
|
|
|
|
|
def make_E_vs_x(
|
|
z,
|
|
mass_u,
|
|
emax_mev,
|
|
medium,
|
|
label,
|
|
npoints,
|
|
P_TORR,
|
|
TEMP_K
|
|
):
|
|
"""
|
|
Builds energy vs depth table using Catima stopping powers.
|
|
Output:
|
|
x [cm], E [MeV]
|
|
"""
|
|
R = 8.3144 # J/mol/K
|
|
|
|
p_pa = P_TORR * 133.322
|
|
molar_density = p_pa / (R * TEMP_K) # mol/m^3
|
|
|
|
# Medium definition
|
|
if medium == "He":
|
|
|
|
m_he = 4.0026
|
|
m_c = 12.0000
|
|
m_o = 15.9949
|
|
"""
|
|
material_def = [
|
|
(m_he, 2, 0.96),
|
|
(m_c, 6, 0.04),
|
|
(m_o, 8, 0.08)
|
|
] """
|
|
material_def = [(m_he, 2, 1.0)]
|
|
|
|
#m_mix_avg = 0.96 * m_he + 0.04 * (m_c + 2 * m_o)
|
|
m_mix_avg = 1.0 * m_he
|
|
rho_g_cm3 = (molar_density * m_mix_avg) / 1e6
|
|
|
|
gas = catima.Material(material_def)
|
|
gas.density(rho_g_cm3)
|
|
|
|
elif medium == "Si":
|
|
|
|
m_si = 28.084
|
|
rho_g_cm3 = 2.33
|
|
|
|
gas = catima.Material([(m_si, 14, 1.0)])
|
|
gas.density(rho_g_cm3)
|
|
|
|
elif medium == "kapton":
|
|
# Kapton (C22H10N2O5)
|
|
# Density: 1.42 g/cm3
|
|
m_h = 1.0078
|
|
m_n = 14.0067
|
|
m_c = 12.0000
|
|
m_o = 15.9949
|
|
kapton_molar_mass = 22*m_c + 10*m_h + 2*m_n + 5*m_o
|
|
|
|
material_def = [
|
|
(m_c, 6, 22/kapton_molar_mass * kapton_molar_mass / m_c),
|
|
(m_h, 1, 10/kapton_molar_mass * kapton_molar_mass / m_h),
|
|
(m_n, 7, 2/kapton_molar_mass * kapton_molar_mass / m_n),
|
|
(m_o, 8, 5/kapton_molar_mass * kapton_molar_mass / m_o)
|
|
]
|
|
|
|
rho_g_cm3 = 1.42
|
|
gas = catima.Material(material_def)
|
|
gas.density(rho_g_cm3)
|
|
|
|
elif medium == "mylar":
|
|
# Mylar (C10H8O4, polyethylene terephthalate)
|
|
# Density: 1.39 g/cm3
|
|
m_h = 1.0078
|
|
m_c = 12.0000
|
|
m_o = 15.9949
|
|
mylar_molar_mass = 10*m_c + 8*m_h + 4*m_o
|
|
|
|
material_def = [
|
|
(m_c, 6, 10/mylar_molar_mass * mylar_molar_mass / m_c),
|
|
(m_h, 1, 8/mylar_molar_mass * mylar_molar_mass / m_h),
|
|
(m_o, 8, 4/mylar_molar_mass * mylar_molar_mass / m_o)
|
|
]
|
|
|
|
rho_g_cm3 = 1.39
|
|
gas = catima.Material(material_def)
|
|
gas.density(rho_g_cm3)
|
|
|
|
else:
|
|
raise ValueError("Unsupported medium")
|
|
|
|
print(f"[INFO] density = {rho_g_cm3:.3e} g/cm^3")
|
|
|
|
projectile = catima.Projectile(mass_u, z)
|
|
E = np.linspace(0.1, emax_mev, npoints)
|
|
S_mass = np.zeros_like(E)
|
|
for i, energy in enumerate(E):
|
|
projectile.T(energy / mass_u)
|
|
S_mass[i] = catima.dedx(projectile, gas) # MeV/(g/cm^2)
|
|
|
|
S_linear = S_mass * rho_g_cm3
|
|
invS = 1.0 / np.clip(S_linear, 1e-30, None)
|
|
x = cumulative_trapezoid(invS[::-1], E[::-1], initial=0)
|
|
x = x[::-1]
|
|
x = -x
|
|
df = pd.DataFrame({
|
|
"Distance_cm": x,
|
|
"Energy_MeV": E
|
|
})
|
|
|
|
outfile = get_loss_table_path(medium, label)
|
|
os.makedirs(os.path.dirname(outfile), exist_ok=True)
|
|
df.to_csv(outfile, sep="\t", index=False)
|
|
|
|
print(f"[INFO] saved: {outfile}")
|
|
|
|
interp_cache.pop((label.lower(), medium.lower()), None)
|
|
gc.collect()
|
|
|
|
return x, E
|
|
|
|
#Generate energy loss tables from file
|
|
def load_table(filename):
|
|
"""
|
|
Load table with columns:
|
|
x(cm) E(MeV)
|
|
|
|
Returns:
|
|
x_array, E_array
|
|
"""
|
|
|
|
data = pd.read_csv(
|
|
filename,
|
|
sep=r'\s+',
|
|
comment="#",
|
|
header=None,
|
|
skiprows=1
|
|
)
|
|
|
|
x = data.iloc[:, 0].to_numpy()
|
|
E = data.iloc[:, 1].to_numpy()
|
|
|
|
return x, E
|
|
|
|
def get_interpolators(particle, medium):
|
|
|
|
canonical_particle = normalize_particle_label(particle)
|
|
cache_key = (canonical_particle.lower(), medium.lower())
|
|
|
|
if cache_key in interp_cache:
|
|
return interp_cache[cache_key]
|
|
|
|
candidate_labels = [canonical_particle]
|
|
if isinstance(particle, str):
|
|
raw_particle = particle.strip()
|
|
candidate_labels.extend([
|
|
raw_particle,
|
|
canonical_particle.replace("-", "")
|
|
])
|
|
|
|
filename = None
|
|
seen_candidates = set()
|
|
for label in candidate_labels:
|
|
normalized_label = label.strip()
|
|
if not normalized_label:
|
|
continue
|
|
key = normalized_label.lower()
|
|
if key in seen_candidates:
|
|
continue
|
|
seen_candidates.add(key)
|
|
|
|
candidate_filename = get_loss_table_path(medium, normalized_label)
|
|
if os.path.exists(candidate_filename):
|
|
filename = candidate_filename
|
|
break
|
|
|
|
if filename is None:
|
|
filename = get_loss_table_path(medium, canonical_particle)
|
|
|
|
x, E = load_table(filename)
|
|
|
|
E_of_x = interp1d(
|
|
x,
|
|
E,
|
|
bounds_error=False,
|
|
fill_value="extrapolate"
|
|
)
|
|
|
|
x_of_E = interp1d(
|
|
E[::-1],
|
|
x[::-1],
|
|
bounds_error=False,
|
|
fill_value="extrapolate"
|
|
)
|
|
|
|
interp_cache[cache_key] = (E_of_x, x_of_E)
|
|
|
|
return E_of_x, x_of_E
|
|
|
|
def energy_loss(particle, medium, Ei, dl):
|
|
|
|
E_of_x, x_of_E = get_interpolators(particle, medium)
|
|
|
|
xi = x_of_E(Ei)
|
|
|
|
xf = xi + dl
|
|
|
|
xmax = E_of_x.x.max() # maximum tabulated range
|
|
|
|
Ef = E_of_x(xf)
|
|
|
|
Ef = np.where(xf >= xmax, 0.0, Ef)
|
|
|
|
return np.maximum(Ef, 0.0)
|
|
|
|
def energy_reconstruction(particle, medium, Ef, dl):
|
|
|
|
E_of_x, x_of_E = get_interpolators(particle, medium)
|
|
|
|
xf = x_of_E(Ef)
|
|
|
|
xi = xf - dl
|
|
|
|
Ei = E_of_x(xi)
|
|
|
|
return np.maximum(Ei, 0.0)
|
|
|
|
def energy_distance(particle, medium, Ei, Ef):
|
|
|
|
_, x_of_E = get_interpolators(particle, medium)
|
|
|
|
xi = x_of_E(Ei)
|
|
|
|
xf = x_of_E(Ef)
|
|
|
|
return np.abs(xf - xi)
|
|
|
|
def resolve_particle(name):
|
|
name = name.lower().strip().rstrip("s")
|
|
if name in particles:
|
|
return particles[name]
|
|
match = re.match(r"([a-zA-Z]+)[-\s]?(\d+)$", name)
|
|
if match:
|
|
element_symbol = match.group(1).capitalize()
|
|
A = int(match.group(2))
|
|
|
|
try:
|
|
element = pt.elements.symbol(element_symbol)
|
|
isotope = element[A]
|
|
|
|
return (
|
|
isotope.number, # Z
|
|
isotope.mass, # mass in u
|
|
30.0,
|
|
f"{element_symbol}-{A}"
|
|
)
|
|
except Exception:
|
|
raise ValueError(f"Unknown isotope: {name}")
|
|
|
|
match = re.match(r"(\d+)[-\s]?([a-zA-Z]+)$", name)
|
|
if match:
|
|
A = int(match.group(1))
|
|
element_symbol = match.group(2).capitalize()
|
|
|
|
try:
|
|
element = pt.elements.symbol(element_symbol)
|
|
isotope = element[A]
|
|
|
|
return (
|
|
isotope.number, # Z
|
|
isotope.mass, # mass in u
|
|
30.0,
|
|
f"{element_symbol}-{A}"
|
|
)
|
|
except Exception:
|
|
raise ValueError(f"Unknown isotope: {name}")
|
|
try:
|
|
elem = pt.elements.symbol(name.capitalize())
|
|
return elem.number, elem.mass, 30.0, name
|
|
except Exception:
|
|
raise ValueError(f"Unknown particle/isotope: {name}")
|
|
|
|
def load_tree_arrays(tree, treename, max_events=None):
|
|
if tree is None:
|
|
raise ValueError(
|
|
f"Tree '{treename}' is not loaded. Run 'set tree {treename}' "
|
|
"after loading a ROOT file with 'uproot file <filename>'."
|
|
)
|
|
branches = tree.keys(recursive=False)
|
|
return tree.arrays(branches, library="np", entry_stop=max_events)
|
|
|
|
def infer_particle_from_filename(filename):
|
|
lower = os.path.basename(filename).lower()
|
|
if "proton" in lower:
|
|
return "proton"
|
|
if "alpha" in lower:
|
|
return "alpha"
|
|
if "deuteron" in lower:
|
|
return "deuteron"
|
|
return None
|
|
|
|
class MyInteractiveApp(cmd.Cmd):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# Initial value set when the script starts
|
|
self.T = 293.15
|
|
self.P = 379
|
|
self.temp_particle = [0, 0.0, 0.0, ""]
|
|
self.rootFile = "SimAnasen1.root"
|
|
self.file = None
|
|
self.initialize_file()
|
|
self.tree = None
|
|
print("-" * 30)
|
|
print("INTERACTIVE SHELL STARTED")
|
|
self.print_params()
|
|
print("Type 'help' for commands.")
|
|
print("Type 'exit' to end program")
|
|
print("-" * 30)
|
|
|
|
def print_params(self):
|
|
"""Helper method to display current state"""
|
|
print(f"Current Parameters: T={self.T} K, P={self.P} Torr")
|
|
|
|
def initialize_file(self):
|
|
"""Load in default root file for anasen"""
|
|
try:
|
|
self.file = uproot.open(f"../Armory/{self.rootFile}")
|
|
except:
|
|
self.file = None
|
|
print("\nATTENTION: Root file not found, continue without uproot functions or uproot file manually")
|
|
|
|
#intro = "Interactive Shell Started. Type 'help' to see commands."
|
|
prompt = ">> "
|
|
|
|
def default(self, line):
|
|
# Check if the command starts with our multi-word phrase
|
|
if line.startswith("make table "):
|
|
# Extract everything after "make table "
|
|
args = line[len("make table "):].strip()
|
|
self.do_make_table(args)
|
|
elif line.startswith("set t") or line.startswith("Set T") or line.startswith("set T") or line.startswith("Set t"):
|
|
args = line[len("set t "):].strip()
|
|
self.do_set_T(args)
|
|
elif line.startswith("set p") or line.startswith("Set P") or line.startswith("set P") or line.startswith("Set p"):
|
|
args = line[len("set p "):].strip()
|
|
self.do_set_P(args)
|
|
elif line.startswith("energy loss") or line.startswith("Energy Loss") or line.startswith("Energy loss"):
|
|
args = line[len("energy loss "):].strip()
|
|
self.do_energy_loss(args)
|
|
elif line.startswith("energy reconstruction") or line.startswith("Energy Reconstruction") or line.startswith("Energy reconstruction"):
|
|
args = line[len("energy reconstruction "):].strip()
|
|
self.do_energy_reconstruction(args)
|
|
elif line.startswith("energy distance") or line.startswith("Energy Distance") or line.startswith("Energy distance"):
|
|
args = line[len("energy distance "):].strip()
|
|
self.do_energy_distance(args)
|
|
else:
|
|
print(f"*** Unknown syntax: {line}")
|
|
|
|
def do_exit(self, arg):
|
|
"""Exits the application."""
|
|
print("Closing application...")
|
|
return True # Returning True stops the cmdloop()
|
|
|
|
def do_T(self, arg):
|
|
"""Print value of T"""
|
|
print(self.T)
|
|
|
|
def do_P(self, arg):
|
|
"""Print value of P (pressure)"""
|
|
print(self.P)
|
|
|
|
def do_set_T(self, arg):
|
|
"""Changes the value of T. Usage: set_t 300"""
|
|
try:
|
|
self.T = float(arg)
|
|
print(f"T has been updated to {self.T}")
|
|
except ValueError:
|
|
print("Please enter a valid number for T.")
|
|
|
|
def do_set_P(self, arg):
|
|
"""Changes the value of P in Torr. Usage: set_ 400"""
|
|
try:
|
|
self.P = float(arg)
|
|
print(f"P has been updated to {self.P}")
|
|
except ValueError:
|
|
print("Please enter a valid number for P.")
|
|
|
|
|
|
def do_make_table(self, arg):
|
|
"""Create E vs X tables for particle, or isotopes
|
|
Ex: >> make table proton <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,
|
|
10000,
|
|
self.P,
|
|
self.T
|
|
)
|
|
|
|
plt.figure(figsize=(8,6))
|
|
plt.plot(x, E)
|
|
|
|
plt.xlabel("Distance (cm)")
|
|
plt.ylabel("Energy (MeV)")
|
|
plt.title(f"Energy Loss Curve {label.capitalize()} {medium}")
|
|
plt.grid(True)
|
|
textstr = f"T = {self.T:.2f} K\nP = {self.P:.2f} Torr"
|
|
|
|
plt.gca().text(
|
|
0.02, 0.02,
|
|
textstr,
|
|
transform=plt.gca().transAxes,
|
|
fontsize=10,
|
|
verticalalignment='bottom',
|
|
bbox=dict(boxstyle="round", facecolor="white", alpha=0.7)
|
|
)
|
|
|
|
plt.tight_layout()
|
|
filename = f"{medium}Loss/Energy_Loss_Curve_{label}.png"
|
|
plt.savefig(filename, dpi=300, bbox_inches="tight")
|
|
plt.show()
|
|
|
|
print(f"Saved plot: {filename}")
|
|
|
|
except Exception as e:
|
|
print(f"Error in make_table: {e}")
|
|
|
|
def do_energy_loss(self, arg):
|
|
"""Find a final energy given an initial energy and distance travelled
|
|
Ex: >> energy loss <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_set_tree(self, arg):
|
|
|
|
if self.file is None:
|
|
print("No ROOT file loaded.")
|
|
return
|
|
|
|
keys = self.file.keys()
|
|
|
|
print("Available trees:", keys)
|
|
|
|
treeName = arg if len(arg) > 0 else "tree1"
|
|
|
|
try:
|
|
# uproot automatically resolves ";1"
|
|
self.tree = self.file[treeName]
|
|
|
|
print(f"Tree loaded: {self.tree}")
|
|
print("Branches:", self.tree.keys())
|
|
|
|
except Exception as e:
|
|
print("Error loading tree:", treeName)
|
|
print("Exception:", e)
|
|
|
|
def do_print_file(self, arg):
|
|
"""Print contents of ROOT file"""
|
|
args = shlex.split(arg)
|
|
file = self.file
|
|
if "keys" in args or len(args) == 0:
|
|
print("Keys in file: ", file.keys())
|
|
if "class_names" in args or len(args) == 0:
|
|
print("Class names: ", file.classnames())
|
|
|
|
def do_root_histograms(self, arg):
|
|
self.do_set_tree("")
|
|
|
|
if arg:
|
|
args = shlex.split(arg)
|
|
particle = args[0]
|
|
else:
|
|
particle = "proton"
|
|
|
|
base = f"{particle}_roothist_plots"
|
|
|
|
if os.path.exists(base):
|
|
shutil.rmtree(base)
|
|
|
|
os.makedirs(base, exist_ok=True)
|
|
|
|
branch_names = []
|
|
for key in self.tree.keys():
|
|
if isinstance(key, bytes):
|
|
branch_names.append(key.decode("ascii"))
|
|
elif hasattr(key, "name"):
|
|
branch_names.append(key.name)
|
|
else:
|
|
branch_names.append(str(key))
|
|
branch_names = list(dict.fromkeys(branch_names))
|
|
|
|
if branch_names and True:
|
|
print(f"Creating histograms for {len(branch_names)} branches...")
|
|
#all_branches = self.tree.arrays(branch_names, library="np", entry_stop=max_events)
|
|
|
|
for branch in self.tree.keys():
|
|
|
|
branch = branch.decode() if isinstance(branch, bytes) else str(branch)
|
|
|
|
try:
|
|
values = self.tree[branch].array(
|
|
library="np",
|
|
)
|
|
|
|
values = np.asarray(values, dtype=float)
|
|
|
|
except Exception as e:
|
|
print(f"Skipping branch {branch}: {e}")
|
|
continue
|
|
|
|
if values.ndim != 1:
|
|
print(f"Skipping non-scalar branch: {branch}")
|
|
continue
|
|
if branch == "qqqZ":
|
|
continue
|
|
|
|
values = values[np.isfinite(values)]
|
|
|
|
if values.size == 0:
|
|
print(f"Skipping empty branch: {branch}")
|
|
continue
|
|
|
|
plt.figure(figsize=(7, 5))
|
|
|
|
plt.hist(
|
|
values,
|
|
bins=100,
|
|
histtype="stepfilled"
|
|
)
|
|
|
|
plt.xlabel(branch)
|
|
plt.ylabel("Counts")
|
|
plt.title(f"{particle} {branch} distribution")
|
|
plt.grid(True)
|
|
plt.tight_layout()
|
|
|
|
safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", branch)
|
|
|
|
filename = f"{base}/{safe_name}_hist.png"
|
|
|
|
counter = 1
|
|
while os.path.exists(filename):
|
|
filename = f"{base}/{safe_name}_{counter}_hist.png"
|
|
counter += 1
|
|
|
|
plt.savefig(filename, dpi=300)
|
|
plt.close()
|
|
|
|
def do_make_plots(self, arg):
|
|
plot_data = []
|
|
|
|
args = shlex.split(arg)
|
|
particle = args[0] if len(args) > 0 else "proton"
|
|
max_events = int(args[1]) if len(args) > 1 else None
|
|
base = f"{particle}_plots"
|
|
|
|
if os.path.exists(base):
|
|
shutil.rmtree(base)
|
|
|
|
os.makedirs(base, exist_ok=True)
|
|
|
|
if self.tree is None:
|
|
self.do_set_tree("")
|
|
print(f"Using TTree: {self.tree}")
|
|
if self.tree is None:
|
|
print("No tree is loaded. Use 'uproot file <filename>' then 'set tree tree1' or 'set tree tree2'.")
|
|
return
|
|
|
|
data = self.tree.arrays(library="pd")
|
|
|
|
#Combine sx3 and QQQ data
|
|
data['X'] = data['sx3X'].fillna(data['qqqX'])
|
|
data['Y'] = data['sx3Y'].fillna(data['qqqY'])
|
|
data['Z'] = data['sx3Z'].fillna(data['qqqZ'])
|
|
data['Ei'] = data['Tb'].fillna(data['qqqTb'])
|
|
|
|
data['dl'] = np.sqrt(((data['X'] - data['vX']) ** 2) + ((data['Y'] - data['vY']) ** 2) + ((data['Z'] - data['vZ']) ** 2)) * .1
|
|
sx3dl = np.sqrt(((data['sx3X'] - data['vX']) ** 2) + ((data['sx3Y'] - data['vY']) ** 2) + ((data['sx3Z'] - data['vZ']) ** 2)) * .1
|
|
qqqdl = np.sqrt(((data['qqqX'] - data['vX']) ** 2) + ((data['qqqY'] - data['vY']) ** 2) + ((data['qqqZ'] - data['vZ']) ** 2)) * .1
|
|
Adl = np.sqrt(((data['aX'] - data['vX']) ** 2) + ((data['aY'] - data['vY']) ** 2) + ((data['aZ'] - data['vZ']) ** 2)) * .1
|
|
Cdl = np.sqrt(((data['cX'] - data['vX']) ** 2) + ((data['cY'] - data['vY']) ** 2) + ((data['cZ'] - data['vZ']) ** 2)) * .1
|
|
|
|
Ei = np.asarray(data['Ei'], dtype=float)
|
|
dl = np.asarray(data['dl'], dtype=float)
|
|
data['EatDet'] = energy_loss(particle, "He", Ei, dl)
|
|
data['Esx3'] = energy_loss(particle, "He", data['Tb'], sx3dl)
|
|
data['Eqqq'] = energy_loss(particle, "He", data['qqqTb'], qqqdl)
|
|
data['EA'] = energy_loss(particle, "He", data['Ei'], Adl)
|
|
data['EC'] = energy_loss(particle, "He", data['Ei'], Cdl)
|
|
|
|
|
|
#Filter out events with 0 final energy, does not reach detector
|
|
data = data[data['EA'] > .2].copy()
|
|
|
|
x = np.asarray(data['X'], dtype=float)
|
|
y = np.asarray(data['Y'], dtype=float)
|
|
z = np.asarray(data['Z'], dtype=float)
|
|
|
|
#print("\nHit detector counts:")
|
|
#print(data['hitDetector'].value_counts())
|
|
|
|
if x.size > 0 and False:
|
|
|
|
fig = plt.figure(figsize=(8,6))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
sc = ax.scatter(
|
|
x, y, z,
|
|
c=np.asarray(data['EatDet'], dtype=float),
|
|
cmap='viridis',
|
|
s=3,
|
|
alpha=0.8
|
|
)
|
|
|
|
ax.set_xlabel("X")
|
|
ax.set_ylabel("Y")
|
|
ax.set_zlabel("Z")
|
|
|
|
ax.set_title(f"{particle} SX3 3D Hit Map")
|
|
|
|
cb = plt.colorbar(sc, ax=ax, pad=0.1)
|
|
cb.set_label("Energy (MeV)")
|
|
|
|
plt.tight_layout()
|
|
out_file = f"{base}/3D_energy.png"
|
|
plt.savefig(out_file, dpi=300)
|
|
plt.show()
|
|
|
|
Ax = np.asarray(data['aX'], dtype=float)
|
|
Ay = np.asarray(data['aY'], dtype=float)
|
|
Az = np.asarray(data['aZ'], dtype=float)
|
|
Cx = np.asarray(data['cX'], dtype=float)
|
|
Cy = np.asarray(data['cY'], dtype=float)
|
|
Cz = np.asarray(data['cZ'], dtype=float)
|
|
|
|
|
|
if x.size > 0 and True:
|
|
|
|
fig = plt.figure(figsize=(8,6))
|
|
ax = fig.add_subplot(111, projection='3d')
|
|
|
|
sc = ax.scatter(
|
|
Ax, Ay, Az,
|
|
c=np.asarray(data['EA'], dtype=float),
|
|
cmap='viridis',
|
|
s=3,
|
|
alpha=0.8
|
|
)
|
|
|
|
ax.scatter(
|
|
Cx, Cy, Cz,
|
|
c=np.asarray(data['EC'], dtype=float),
|
|
s=3,
|
|
alpha=0.8,
|
|
marker='^',
|
|
label='Cathode'
|
|
)
|
|
|
|
ax.set_xlabel("X")
|
|
ax.set_ylabel("Y")
|
|
ax.set_zlabel("Z")
|
|
|
|
ax.set_title(f"{particle} SX3 3D Hit Map")
|
|
|
|
cb = plt.colorbar(sc, ax=ax, pad=0.1)
|
|
cb.set_label("Energy (MeV)")
|
|
|
|
plt.tight_layout()
|
|
out_file = f"{base}/Anode_Cathode_Hits.png"
|
|
plt.savefig(out_file, dpi=300)
|
|
plt.show()
|
|
|
|
EPC = abs(data['EA'] - data['EC'])
|
|
|
|
plt.figure(figsize=(8, 6))
|
|
plt.hist(data['Esx3'], bins=200)
|
|
|
|
|
|
|
|
mask = ~np.isnan(data['Esx3'])
|
|
plt.figure(figsize=(8, 6))
|
|
plt.hist2d(
|
|
data['Esx3'][mask] * np.sin(np.deg2rad(data['thetab'][mask])),
|
|
EPC[mask],
|
|
bins=200
|
|
)
|
|
plt.xlabel("SX3 Energy (MeV)")
|
|
plt.ylabel("PCEnergy x Sin(theta)")
|
|
plt.colorbar(label="Counts")
|
|
plt.tight_layout()
|
|
plt.savefig(
|
|
f"{base}/Eprop_vs_Esx3.png",
|
|
dpi=300
|
|
)
|
|
plt.show()
|
|
|
|
plt.figure(figsize=(8,6))
|
|
plt.hist(data['Esx3'], bins=200)
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
mask = ~np.isnan(data['Esx3'])
|
|
plt.figure(figsize=(8,6))
|
|
plt.hist2d(data['EPC'][mask], data['Esx3'][mask], bins=200)
|
|
|
|
plt.show()
|
|
|
|
data.to_csv('data.csv', index=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
MyInteractiveApp().cmdloop()
|
|
|
|
|