cmd line feature for eloss app

This commit is contained in:
James Szalkie 2026-05-21 11:59:47 -04:00
parent ca33ba17b3
commit 6ef916ba1c
2 changed files with 678 additions and 505 deletions

File diff suppressed because it is too large Load Diff

View File

@ -15,15 +15,29 @@ import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
from scipy.integrate import cumulative_trapezoid
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import threading
import time
import sys
import cmd
import shlex
import textwrap
import readline
import atexit
import os
#Run program from terminal or IDE, and prompts will provide user steps
histfile = os.path.expanduser("~/.uproot_shell_history")
try:
readline.read_history_file(histfile)
except FileNotFoundError:
pass
atexit.register(readline.write_history_file, histfile)
#data = [z, mass_u, maximum MeV, name]
alpha_data = [2, 4.0026, 40, "alpha"]
@ -216,6 +230,31 @@ def energy_reconstruction(particle, Ef, dl):
return max(Ei, 0.0)
def energy_distance(particle, Ei, Ef):
filename = f"E_vs_x_{particle}.dat"
data = pd.read_csv(
filename,
sep='\s+',
comment="#",
header=None,
skiprows=1
)
x = data.iloc[:, 0].values
E = data.iloc[:, 1].values
x_of_E = interp1d(
E[::-1], # reverse so energy increases
x[::-1],
bounds_error=False,
fill_value="extrapolate"
)
xi = float(x_of_E(Ei))
xf = float(x_of_E(Ef))
return abs(xf - xi)
class MyInteractiveApp(cmd.Cmd):
@ -226,6 +265,10 @@ class MyInteractiveApp(cmd.Cmd):
self.T = 293.15
self.P = 400
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()
@ -236,9 +279,16 @@ class MyInteractiveApp(cmd.Cmd):
def print_params(self):
"""Helper method to display current state"""
print(f"Current Parameters: T={self.T}, P={self.P}")
def initialize_file(self):
"""Load in default root file for anasen"""
self.file = uproot.open(f"../Armory/{self.rootFile}")
#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 "):
@ -260,6 +310,9 @@ class MyInteractiveApp(cmd.Cmd):
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}")
@ -395,7 +448,7 @@ class MyInteractiveApp(cmd.Cmd):
try:
self.do_make_table(table_specs)
Ef = energy_loss(particle, Ei, dl)
print(f"\nFinal energy: {Ef:.6f} MeV")
print(f"\nFinal energy: {Ef:.6f} MeV\n")
except:
return
except IndexError:
@ -412,7 +465,7 @@ class MyInteractiveApp(cmd.Cmd):
try:
Ei = energy_reconstruction(particle, Ef, dl)
if Ei > 0:
print(f"\nInitial energy: {Ei:.6f} MeV")
print(f"\nInitial energy: {Ei:.6f} MeV\n")
else:
print("Error: remake table with larger value, fallen off map")
except:
@ -420,7 +473,127 @@ class MyInteractiveApp(cmd.Cmd):
except IndexError:
print("Please input particle, final energy from detector, and distance travelled")
def do_energy_distance(self, arg):
"""Find a distance travelled given an initial and final energy
Ex: >> energy distance <particle> <initial energy> <final energy MeV>"""
args = shlex.split(arg)
try:
particle = args[0]
Ei = float(args[1])
Ef = float(args[2])
dE = Ei - Ef
try:
dl = energy_distance(particle, Ei, Ef)
if dl > 0:
print(f"\nChange in energy: {dE:.6f} MeV")
print(f"Distance travelled: {dl:.6f} cm\n")
else:
print("Error: remake table with larger value, fallen off map")
except:
print("Particle energy table not made yet, please do so using 'make table'")
except IndexError:
print("Please input particle, final energy from detector, and distance travelled")
def do_uproot_file(self, arg):
args = shlex.split(arg)
if len(args) > 0:
filename = args[0]
else:
filename = self.rootFile
try:
print(f"Opening {filename}")
with uproot.open(f"../Armory/{filename}") as tempfile:
self.file = tempfile
except FileNotFoundError:
print("Error: file not found")
def do_print_file(self, arg):
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):
file = self.file
if len(arg) > 0:
treeName = f"tree{arg}"
else:
treeName = "tree"
try:
self.tree = file[treeName]
tree = self.tree
print(f"Tree: {tree}")
print("Branches: ", tree.keys())
except:
print("\nError, trees include ", file.keys())
def run_command_line(self):
import readline
import atexit
import os
import awkward as ak
histfile = os.path.expanduser("~/.uproot_shell_history")
try:
readline.read_history_file(histfile)
except FileNotFoundError:
pass
readline.set_history_length(1000)
readline.parse_and_bind("tab: complete")
atexit.register(readline.write_history_file, histfile)
print("Custom Python CMD (type 'exit' to stop)")
local_vars = {
"file": self.file,
"uproot": uproot,
"np": np,
#"ak": ak,
"plt": plt,
"self": self
}
startup_script = textwrap.dedent("""
tree = file["tree"]
""")
exec(startup_script, globals(), local_vars)
while True:
try:
entry = input(">>> ").strip()
if entry.lower() in ["exit", "quit"]:
break
if not entry:
continue
# Try eval first so expressions print naturally
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):
self.run_command_line()
if __name__ == "__main__":
MyInteractiveApp().cmdloop()