cmd line feature for eloss app

ELoss/PCEnergyAnalysis.py

@@ -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()
-