fixed memory issues and dual plotting

ELoss/PCEnergyAnalysis.py

@@ -8,16 +8,12 @@ Created on Wed May 20 13:32:14 2026
 import numpy as np
 import pandas as pd
 from scipy.interpolate import interp1d
-import argparse
 import uproot
 import pycatima as catima
 from scipy.integrate import cumulative_trapezoid
 import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
-import threading
-import time
-import sys
 import cmd
 import shlex
 import textwrap
@@ -26,17 +22,18 @@ import atexit
 import os
 import periodictable as pt
 import re
+from matplotlib.colors import PowerNorm
 
 #Run program from terminal or IDE, and prompts will provide user steps
 
-histfile = os.path.expanduser("~/.uproot_shell_history")
+#histfile = os.path.expanduser("~/.uproot_shell_history")
 
-try:
+#try:
-    readline.read_history_file(histfile)
+#    readline.read_history_file(histfile)
-except FileNotFoundError:
+#except FileNotFoundError:
-    pass
+#    pass
 
-atexit.register(readline.write_history_file, histfile)
+#atexit.register(readline.write_history_file, histfile)
 
 #data = [z, mass_u, maximum MeV, name]
 alpha_data = [2, 4.0026, 40, "alpha"]
@@ -254,6 +251,69 @@ def resolve_particle(name):
         raise ValueError(f"Unknown particle/isotope: {name}")
     
     
+def update_plot_data(name, values):
+    for i, (existing_name, _) in enumerate(plot_data):
+        if existing_name == name:
+            plot_data[i] = (name, values)
+            return
+    plot_data.append((name, values))
+    
+def process_file(filename):
+
+    tree = uproot.open(filename)["tree"]
+
+    branches = ["Tb", "thetab", "sx3Z"]
+
+    data = tree.arrays(branches, library="np")
+
+    Ei = data["Tb"]
+    theta = np.radians(data["thetab"])
+    sx3Z = data["sx3Z"]
+
+    mask = np.sin(theta) != 0
+
+    Ei = Ei[mask]
+    theta = theta[mask]
+    sx3Z = sx3Z[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
+
+    # Determine particle from filename
+    lower = filename.lower()
+
+    if "proton" in lower:
+        particle = "proton"
+    elif "alpha" in lower:
+        particle = "alpha"
+    else:
+        particle = "proton"
+
+    print(f"Computing energies for {particle}...")
+
+    EA = energy_loss(particle, Ei, dA)
+    EC = energy_loss(particle, Ei, dC)
+    Esx3 = energy_loss(particle, Ei, dsx3)
+
+    Elost = Ei - Esx3
+    Eprop = EA - EC
+
+    return {
+        "particle": particle,
+        "Ei": Ei,
+        "sx3Z": sx3Z,
+        "EA": EA,
+        "EC": EC,
+        "Esx3": Esx3,
+        "Elost": Elost,
+        "Eprop": Eprop
+    }
+
 class MyInteractiveApp(cmd.Cmd):
     def __init__(self):
         super().__init__()
@@ -405,11 +465,28 @@ class MyInteractiveApp(cmd.Cmd):
     
             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()}")
             plt.grid(True)
-            plt.show()
+            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"Energy_Loss_Curve_{label}.png"
+            plt.savefig(filename, dpi=300, bbox_inches="tight")
+            plt.close()
+            
+            print(f"Saved plot: {filename}")
     
         except Exception as e:
             print(f"Error in make_table: {e}")
@@ -756,6 +833,10 @@ class MyInteractiveApp(cmd.Cmd):
     
         import os
         
+        global plot_data
+        
+        plot_data = []
+    
         args = shlex.split(arg)
     
         particle = args[0] if len(args) > 0 else "proton"
@@ -783,6 +864,9 @@ class MyInteractiveApp(cmd.Cmd):
         theta = theta[mask]
         sx3Z = sx3Z[mask]
         
+        update_plot_data(f"{particle}_Ei", Ei)
+        update_plot_data(f"{particle}_sx3Z", sx3Z)
+    
         sin_theta = np.sin(theta)
     
         radii = np.array([3.2, 4.2, 6.6])
@@ -800,6 +884,12 @@ class MyInteractiveApp(cmd.Cmd):
         Elost = Ei - Esx3
         Eprop = EA - EC
         
+        update_plot_data(f"{particle}_EA", EA)
+        update_plot_data(f"{particle}_EC", EC)
+        update_plot_data(f"{particle}_Esx3", Esx3)
+        update_plot_data(f"{particle}_Elost", Elost)
+        update_plot_data(f"{particle}_Eprop", Eprop)
+    
         base = f"{particle}_plots"
         os.makedirs(base, exist_ok=True)
     
@@ -810,29 +900,29 @@ class MyInteractiveApp(cmd.Cmd):
         plt.hist(Elost, bins=200)
         plt.xlabel("Energy Loss (MeV)")
         plt.ylabel("Counts")
-        plt.title("Total Energy Loss Distribution")
+        plt.title(f"{particle} Total Energy Loss Distribution")
         plt.grid(True)
         plt.tight_layout()
         plt.savefig(f"{base}/Elost_hist.png", dpi=300)
         plt.close()
     
-        # 2. Histogram: sx3Z
+        #Histogram: sx3Z
         plt.figure(figsize=(7,5))
         plt.hist(sx3Z, bins=100)
         plt.xlabel("SX3 Z")
         plt.ylabel("Counts")
-        plt.title("SX3 Position Distribution")
+        plt.title(f"{particle} SX3 Position Distribution")
         plt.grid(True)
         plt.tight_layout()
         plt.savefig(f"{base}/sx3Z_hist.png", dpi=300)
         plt.close()
     
-        # 3. 2D: Elost vs sx3Z
+        #2D: Elost vs sx3Z
         plt.figure(figsize=(7,6))
-        plt.hist2d(Elost, sx3Z, bins=200)
+        plt.hist2d(sx3Z, Elost, bins=200)
-        plt.xlabel("Energy Loss (MeV)")
+        plt.ylabel("Energy Loss (MeV)")
-        plt.ylabel("SX3 Z")
+        plt.xlabel("SX3 Z")
-        plt.title("Energy Loss vs SX3 Position")
+        plt.title(f"{particle} Energy Loss vs SX3 Position")
         plt.colorbar(label="Counts")
         plt.tight_layout()
         plt.savefig(f"{base}/Elost_vs_sx3Z.png", dpi=300)
@@ -843,7 +933,7 @@ class MyInteractiveApp(cmd.Cmd):
         plt.hist2d(EA, Esx3, bins=200)
         plt.xlabel("EA (MeV)")
         plt.ylabel("Esx3 (MeV)")
-        plt.title("Anode vs SX3 Energy")
+        plt.title(f"{particle} Anode vs SX3 Energy")
         plt.colorbar(label="Counts")
         plt.tight_layout()
         plt.savefig(f"{base}/EA_vs_Esx3.png", dpi=300)
@@ -854,14 +944,214 @@ class MyInteractiveApp(cmd.Cmd):
         plt.hist2d(Eprop, sx3Z, bins=200)
         plt.xlabel("EA - EC (MeV)")
         plt.ylabel("SX3 Z")
-        plt.title("Energy Propagation Difference vs Position")
+        plt.title(f"{particle} Energy Propagation Difference vs Position")
         plt.colorbar(label="Counts")
         plt.tight_layout()
         plt.savefig(f"{base}/Eprop_vs_sx3Z.png", dpi=300)
         plt.close()
         
+        plt.figure(figsize=(7,6))
+        plt.hist2d(Esx3, Eprop, bins=200)
+        plt.ylabel("PCEnergy")
+        plt.xlabel("SX3 Energy (MeV)")
+        plt.title(f"{particle} Energy Propagation Difference vs sx3 Energy")
+        plt.colorbar(label="Counts")
+        plt.xlim(0,30)
+        plt.ylim(0,.45)
+        plt.tight_layout()
+        plt.savefig(f"{base}/Eprop_vs_Esx3.png", dpi=300)
+        plt.close()
+    
         print("Plotting complete.")
                 
+    def do_dual_plotter(self, arg):
+        args = shlex.split(arg)
+    
+        if len(args) < 2:
+            print("Usage: make dual plots proton_data.root alpha_data.root")
+            return
+    
+        file1 = args[0]
+        file2 = args[1]
+    
+        #Helper function
+
+    
+        #Process both files
+        data1 = process_file(f"../Armory/{file1}")
+        data2 = process_file(f"../Armory/{file2}")
+    
+        outdir = "dual_plots"
+        os.makedirs(outdir, exist_ok=True)
+    
+        print(f"Saving plots to: {outdir}")
+    
+        #Overlay histogram: Elost
+        plt.figure(figsize=(8,6))
+    
+        plt.hist(
+            data1["Elost"],
+            bins=200,
+            histtype='step',
+            linewidth=2,
+            density=True,
+            label=data1["particle"]
+        )
+    
+        plt.hist(
+            data2["Elost"],
+            bins=200,
+            histtype='step',
+            linewidth=2,
+            density=True,
+            label=data2["particle"]
+        )
+    
+        plt.xlabel("Energy Loss (MeV)")
+        plt.ylabel("Normalized Counts")
+        plt.title("Energy Loss Comparison")
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+    
+        plt.savefig(f"{outdir}/Elost_overlay.png", dpi=300)
+        plt.close()
+    
+        #Overlay histogram: sx3Z
+        plt.figure(figsize=(8,6))
+    
+        plt.hist(
+            data1["sx3Z"],
+            bins=150,
+            histtype='step',
+            linewidth=2,
+            density=True,
+            label=data1["particle"]
+        )
+    
+        plt.hist(
+            data2["sx3Z"],
+            bins=150,
+            histtype='step',
+            linewidth=2,
+            density=True,
+            label=data2["particle"]
+        )
+    
+        plt.xlabel("SX3 Z")
+        plt.ylabel("Normalized Counts")
+        plt.title("SX3 Position Comparison")
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+    
+        plt.savefig(f"{outdir}/sx3Z_overlay.png", dpi=300)
+        plt.close()
+    
+        #Side-by-side 2D plots
+        fig, axes = plt.subplots(1, 2, figsize=(14,6))
+    
+        h1 = axes[0].hist2d(
+            data1["sx3Z"],
+            data1["Elost"],
+            bins=200
+        )
+    
+        axes[0].set_title(f'{data1["particle"]} Elost vs SX3')
+        axes[0].set_xlabel("SX3 Z")
+        axes[0].set_ylabel("Energy Loss (MeV)")
+    
+        h2 = axes[1].hist2d(
+            data2["sx3Z"],
+            data2["Elost"],
+            bins=200
+        )
+    
+        axes[1].set_title(f'{data2["particle"]} Elost vs SX3')
+        axes[1].set_xlabel("SX3 Z")
+        axes[1].set_ylabel("Energy Loss (MeV)")
+    
+        fig.colorbar(h1[3], ax=axes[0], label="Counts")
+        fig.colorbar(h2[3], ax=axes[1], label="Counts")
+    
+        plt.tight_layout()
+    
+        plt.savefig(f"{outdir}/Elost_vs_sx3_comparison.png", dpi=300)
+        plt.close()
+        
+        #EA vs Esx3 overlay scatter
+        plt.figure(figsize=(8,6))
+    
+        plt.scatter(
+            data1["EA"],
+            data1["Esx3"],
+            s=1,
+            alpha=0.3,
+            label=data1["particle"]
+        )
+    
+        plt.scatter(
+            data2["EA"],
+            data2["Esx3"],
+            s=1,
+            alpha=0.3,
+            label=data2["particle"]
+        )
+    
+        plt.xlabel("EA (MeV)")
+        plt.ylabel("Esx3 (MeV)")
+        plt.title("Anode vs SX3 Energy")
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+    
+        plt.savefig(f"{outdir}/EA_vs_Esx3_overlay.png", dpi=300)
+        plt.close()
+        
+        #PCE vs SiE
+        
+        combined_Esx3 = np.concatenate([
+            data1["Esx3"],
+            data2["Esx3"]
+        ])
+        
+        combined_Eprop = np.concatenate([
+            data1["Eprop"],
+            data2["Eprop"]
+        ])
+        
+        plt.figure(figsize=(7,6))
+        
+        plt.hist2d(
+            combined_Esx3,
+            combined_Eprop,
+            bins=100,
+            range=[[0,30],[0,0.45]],
+            cmap='viridis'   # same default matplotlib style
+        )
+        
+        plt.ylabel("PCEnergy")
+        plt.xlabel("SX3 Energy (MeV)")
+        
+        plt.title(
+            f'{data1["particle"]} + {data2["particle"]} '
+            'Energy Propagation Difference vs SX3 Energy'
+        )
+        
+        plt.colorbar(label="Counts")
+        
+        plt.xlim(0,30)
+        plt.ylim(0,0.45)
+        
+        plt.tight_layout()
+        
+        plt.savefig(f"{outdir}/Combined_Eprop_vs_Esx3.png", dpi=300)
+        
+        plt.close()
+    
+        print("Dual plotting complete.")
+        
+        
 #exec(open("PCEnergyAnalysis.py").read())
 
 if __name__ == "__main__":