ANASEN_analysis/ELoss/EnergyLoss.py
James Szalkie 55e04f228d straggling
2026-08-26 13:45:34 -04:00

594 lines
17 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 20 12:46:47 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
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
# Energy and range straggling at each cumulative depth, starting from emax_mev
e_u_init = emax_mev / mass_u
sigma_E = np.zeros_like(E)
sigma_x = np.zeros_like(E)
for i, depth_cm in enumerate(x):
gas.density(rho_g_cm3).thickness(depth_cm * rho_g_cm3)
projectile.T(e_u_init)
result = catima.calculate(projectile, gas)
sigma_E[i] = result.sigma_E
sigma_x[i] = result.sigma_x
df = pd.DataFrame({
"Distance_cm": x,
"Energy_MeV": E,
"Sigma_E_MeV": sigma_E,
"Sigma_x_cm": sigma_x
})
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) Sigma_E(MeV) [optional] Sigma_x(cm) [optional]
Returns:
x_array, E_array, sigma_E_array, sigma_x_array (None if column absent)
"""
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()
sigma_E = data.iloc[:, 2].to_numpy() if data.shape[1] > 2 else None
sigma_x = data.iloc[:, 3].to_numpy() if data.shape[1] > 3 else None
return x, E, sigma_E, sigma_x
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, sigma_E, sigma_x = 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"
)
if sigma_E is not None:
sigma_of_x = interp1d(
x,
sigma_E,
bounds_error=False,
fill_value="extrapolate"
)
else:
sigma_of_x = None
if sigma_x is not None:
sigma_x_of_x = interp1d(
x,
sigma_x,
bounds_error=False,
fill_value="extrapolate"
)
else:
sigma_x_of_x = None
interp_cache[cache_key] = (E_of_x, x_of_E, sigma_of_x, sigma_x_of_x)
return E_of_x, x_of_E, sigma_of_x, sigma_x_of_x
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}")
class MyInteractiveApp(cmd.Cmd):
def __init__(self):
super().__init__()
# Initial value set when the script starts
self.T = 293.15
self.P = 250#379
self.temp_particle = [0, 0.0, 0.0, ""]
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")
#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,
100000,
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")
if __name__ == "__main__":
MyInteractiveApp().cmdloop()