Streamlined, save new data as root file

ELoss/PCEnergyAnalysis.py

@@ -11,9 +11,6 @@ from scipy.interpolate import interp1d
 import argparse
 import uproot
 import pycatima as catima
-import numpy as np
-import pandas as pd
-from scipy.interpolate import interp1d
 from scipy.integrate import cumulative_trapezoid
 import matplotlib
 matplotlib.use("TkAgg")
@@ -46,6 +43,9 @@ alpha_data = [2, 4.0026, 40, "alpha"]
 proton_data = [1, 1.0078, 20, "proton"]
 deuteron_data = [1, 2.014102, 30, "deuteron"]
 
+dA = 3.2 #cm
+dC = 4.2 #cm
+dsx3 = 8.8 #cm
 
 particles = {
     "alpha": alpha_data,
@@ -53,6 +53,8 @@ particles = {
     "deuteron": deuteron_data
 }
 
+interp_cache = {}
+
 def make_E_vs_x(z, mass_u, emax_mev, label, npoints, P_TORR, TEMP_K):
     # GAS SETUP
     R = 8.3144
@@ -149,29 +151,25 @@ def load_table(filename):
 
     data = pd.read_csv(
         filename,
-        sep='\s+',
+        sep=r'\s+',
         comment="#",
         header=None,
         skiprows=1
     )
 
-    x = data.iloc[:, 0].values
-    E = data.iloc[:, 1].values
+    x = data.iloc[:, 0].to_numpy()
+    E = data.iloc[:, 1].to_numpy()
 
     return x, E
 
-def energy_loss(particle, Ei, dl):
-    filename = f"E_vs_x_{particle}.dat"
-    data = pd.read_csv(
-        filename,
-        sep='\s+',
-        comment="#",
-        header=None,
-        skiprows=1
-    )
+def get_interpolators(particle):
 
-    x = data.iloc[:, 0].values
-    E = data.iloc[:, 1].values
+    if particle in interp_cache:
+        return interp_cache[particle]
+
+    filename = f"E_vs_x_{particle}.dat"
+
+    x, E = load_table(filename)
 
     E_of_x = interp1d(
         x,
@@ -181,81 +179,51 @@ def energy_loss(particle, Ei, dl):
     )
 
     x_of_E = interp1d(
-        E[::-1],      # reverse so energy increases
+        E[::-1],
         x[::-1],
         bounds_error=False,
         fill_value="extrapolate"
     )
 
-    xi = float(x_of_E(Ei))
+    interp_cache[particle] = (E_of_x, x_of_E)
 
-    xf = xi + float(dl)
+    return E_of_x, x_of_E
 
-    Ef = float(E_of_x(xf))
+def energy_loss(particle, Ei, dl):
 
-    return max(Ef, 0.0)
+    E_of_x, x_of_E = get_interpolators(particle)
+
+    xi = x_of_E(Ei)
+
+    xf = xi + dl
+
+    Ef = E_of_x(xf)
+
+    return np.maximum(Ef, 0.0)
 
 #def energy_loss_angular(particle, Ei, theta):
     
 def energy_reconstruction(particle, Ef, dl):
-    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
+    E_of_x, x_of_E = get_interpolators(particle)
 
-    E_of_x = interp1d(
-        x,
-        E,
-        bounds_error=False,
-        fill_value=0.0
-    )
-    
-    x_of_E = interp1d(
-        E[::-1],      # reverse so energy increases
-        x[::-1],
-        bounds_error=False,
-        fill_value="extrapolate"
-    )
-    
-    xf = float(x_of_E(Ef))
+    xf = x_of_E(Ef)
 
     xi = xf - dl
 
-    Ei = float(E_of_x(xi))
+    Ei = E_of_x(xi)
 
-    return max(Ei, 0.0)
+    return np.maximum(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 = get_interpolators(particle)
 
-    x_of_E = interp1d(
-        E[::-1],      # reverse so energy increases
-        x[::-1],
-        bounds_error=False,
-        fill_value="extrapolate"
-    )
+    xi = x_of_E(Ei)
 
-    xi = float(x_of_E(Ei))
-    xf = float(x_of_E(Ef))
+    xf = x_of_E(Ef)
 
-    return abs(xf - xi)
+    return np.abs(xf - xi)
 
 
 def resolve_particle(name):
@@ -307,11 +275,15 @@ class MyInteractiveApp(cmd.Cmd):
         
     def print_params(self):
         """Helper method to display current state"""
-        print(f"Current Parameters: T={self.T}, P={self.P}")
+        print(f"Current Parameters: T={self.T} K, P={self.P} Torr")
         
     def initialize_file(self):
         """Load in default root file for anasen"""
-        self.file = uproot.open(f"../Armory/{self.rootFile}")
+        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."
@@ -510,6 +482,7 @@ class MyInteractiveApp(cmd.Cmd):
             
     def do_uproot_file(self, arg):
         """Open a specific root file for inspection"""
+        global rootfile
         args = shlex.split(arg)
         if len(args) > 0:
             filename = args[0]
@@ -517,9 +490,12 @@ class MyInteractiveApp(cmd.Cmd):
             filename = self.rootFile
         try:
             print(f"Opening {filename}")
-            with uproot.open(f"../Armory/{filename}") as tempfile:
-                self.file = tempfile
+            with uproot.open(f"../Armory/{filename}") as rootfile:
+                self.file = rootfile
         except FileNotFoundError:
+            with uproot.open(f"{filename}") as rootfile:
+                self.file = rootfile
+        except:
             print("Error: file not found")
     
     def do_print_file(self, arg):
@@ -535,11 +511,12 @@ class MyInteractiveApp(cmd.Cmd):
         """Set a specific tree from the file (default to 'tree')"""
         file = self.file
         if len(arg) > 0:
-            treeName = f"tree{arg}"
+            treeName = f"{arg}"
         else:
             treeName = "tree"
         try:
             self.tree = file[treeName]
+            global tree
             tree = self.tree
             print(f"Tree: {tree}")
             print("Branches: ", tree.keys())
@@ -583,14 +560,10 @@ class MyInteractiveApp(cmd.Cmd):
         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)
     
@@ -613,5 +586,130 @@ class MyInteractiveApp(cmd.Cmd):
         look up 'uproot' for more details"""
         self.run_command_line()
         
+    def do_energy_analysis(self, arg):
+    
+        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}")
+    
+        branches = [
+            "Tb",
+            "thetab"
+        ]
+    
+        if max_events:
+            n_events = max_events
+        else:
+            n_events = self.tree.num_entries
+    
+        print(f"Loading {n_events} events...")
+    
+        data = self.tree.arrays(
+            branches,
+            library="np",
+            entry_stop=max_events
+        )
+    
+        Ei = data["Tb"]
+    
+        theta = np.radians(data["thetab"])
+    
+        # Remove theta = 0 events
+        mask = np.sin(theta) != 0
+    
+        Ei = Ei[mask]
+        theta = theta[mask]
+    
+        sin_theta = np.sin(theta)
+    
+        radii = np.array([3.2, 4.2, 6.6])
+    
+        dA = radii[0] / sin_theta
+        dC = radii[1] / sin_theta
+        dsx3 = radii[2] / sin_theta
+    
+        print("Calculating energy losses...")
+    
+        EA = energy_loss(particle, Ei, dA)
+    
+        EC = energy_loss(particle, Ei, dC)
+    
+        Esx3 = energy_loss(particle, Ei, dsx3)
+    
+        Eprop = EA - EC
+    
+        print("Analysis complete")
+    
+        print(f"Processed events: {len(Ei)}")
+    
+        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"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)
+        
+        # Put values back into valid entries
+        EA_full[mask] = EA
+        EC_full[mask] = EC
+        Esx3_full[mask] = Esx3
+        Eprop_full[mask] = Eprop
+        
+        # Add new branches
+        all_data["EA"] = EA_full
+        all_data["EC"] = EC_full
+        all_data["Esx3"] = Esx3_full
+        all_data["Eprop"] = Eprop_full
+        
+        # Write new ROOT file
+        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")
+        
+                
+            
+
 if __name__ == "__main__":
     MyInteractiveApp().cmdloop()