diff --git a/ELoss/PCEnergyAnalysis.py b/ELoss/PCEnergyAnalysis.py index ad86d04..cba0b38 100644 --- a/ELoss/PCEnergyAnalysis.py +++ b/ELoss/PCEnergyAnalysis.py @@ -13,7 +13,7 @@ import uproot import pycatima as catima from scipy.integrate import cumulative_trapezoid import matplotlib -matplotlib.use("TkAgg") +matplotlib.use("Agg") import matplotlib.pyplot as plt import threading import time @@ -482,21 +482,28 @@ 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] else: filename = self.rootFile + try: print(f"Opening {filename}") - 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") + + # Try Armory path first + try: + self.file = uproot.open(f"../Armory/{filename}") + except FileNotFoundError: + self.file = uproot.open(filename) + + print("File loaded successfully.") + print("Keys:", self.file.keys()) + + except Exception as e: + print("Error opening file:", e) def do_print_file(self, arg): """Print contents of ROOT file""" @@ -508,20 +515,27 @@ class MyInteractiveApp(cmd.Cmd): print("Class names: ", file.classnames()) def do_set_tree(self, arg): - """Set a specific tree from the file (default to 'tree')""" - file = self.file - if len(arg) > 0: - treeName = f"{arg}" - else: - treeName = "tree" + + if self.file is None: + print("No ROOT file loaded.") + return + + keys = self.file.keys() + + print("Available trees:", keys) + + treeName = arg if len(arg) > 0 else "tree" + try: - self.tree = file[treeName] - global tree - tree = self.tree - print(f"Tree: {tree}") - print("Branches: ", tree.keys()) - except: - print("\nError, trees include ", file.keys()) + # uproot automatically resolves ";1" + self.tree = self.file[treeName] + + print(f"Tree loaded: {self.tree}") + print("Branches:", self.tree.keys()) + + except Exception as e: + print("Error loading tree:", treeName) + print("Exception:", e) def run_command_line(self): import readline @@ -607,7 +621,8 @@ class MyInteractiveApp(cmd.Cmd): branches = [ "Tb", - "thetab" + "thetab", + "sx3Z" ] if max_events: @@ -622,8 +637,10 @@ class MyInteractiveApp(cmd.Cmd): library="np", entry_stop=max_events ) - + global Ei Ei = data["Tb"] + global sx3Z + sx3Z = data["sx3Z"] theta = np.radians(data["thetab"]) @@ -642,74 +659,210 @@ class MyInteractiveApp(cmd.Cmd): dsx3 = radii[2] / sin_theta print("Calculating energy losses...") - + global EA EA = energy_loss(particle, Ei, dA) - + global EC EC = energy_loss(particle, Ei, dC) - + global Esx3 Esx3 = energy_loss(particle, Ei, dsx3) - + global Eprop Eprop = EA - EC - + global Elost + Elost = Ei - Esx3 + 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"Average total energy loss to sx3: {np.mean(Elost):.3f} MeV") + + print(f"Maximum total energy loss to sx3: {np.max(Elost):.3f} MeV") + + print(f"Minimum total energy loss to sx3: {np.min(Elost):.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) - + 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) - + + Elost_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 - + + Elost_full[mask] = Elost + # 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 + + all_data["Elost"] = Elost_full + + # Write new ROOT file as a classic TTree 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") - print("Finished writing augmented ROOT file") + def do_make_plots(self, arg): + + import os + + args = shlex.split(arg) + + particle = args[0] if len(args) > 0 else "proton" + max_events = int(args[1]) if len(args) > 1 else None + + if self.tree is None: + self.do_set_tree("") + print(f"Using TTree: {self.tree}") + + branches = ["Tb", "thetab", "sx3Z"] + + data = self.tree.arrays( + branches, + library="np", + entry_stop=max_events + ) + + 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 + + print("Computing energies...") + + EA = energy_loss(particle, Ei, dA) + EC = energy_loss(particle, Ei, dC) + Esx3 = energy_loss(particle, Ei, dsx3) + + Elost = Ei - Esx3 + Eprop = EA - EC + + base = f"{particle}_plots" + os.makedirs(base, exist_ok=True) + + print(f"Saving plots to folder: {base}") + + # 1. Histogram: energy loss + plt.figure(figsize=(7,5)) + plt.hist(Elost, bins=200) + plt.xlabel("Energy Loss (MeV)") + plt.ylabel("Counts") + plt.title("Total Energy Loss Distribution") + plt.grid(True) + plt.tight_layout() + plt.savefig(f"{base}/Elost_hist.png", dpi=300) + plt.close() + + # 2. 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.grid(True) + plt.tight_layout() + plt.savefig(f"{base}/sx3Z_hist.png", dpi=300) + plt.close() + + # 3. 2D: Elost vs sx3Z + plt.figure(figsize=(7,6)) + plt.hist2d(Elost, sx3Z, bins=200) + plt.xlabel("Energy Loss (MeV)") + plt.ylabel("SX3 Z") + plt.title("Energy Loss vs SX3 Position") + plt.colorbar(label="Counts") + plt.tight_layout() + plt.savefig(f"{base}/Elost_vs_sx3Z.png", dpi=300) + plt.close() + + #Anode energy vs sx3 energy + plt.figure(figsize=(7,6)) + plt.hist2d(EA, Esx3, bins=200) + plt.xlabel("EA (MeV)") + plt.ylabel("Esx3 (MeV)") + plt.title("Anode vs SX3 Energy") + plt.colorbar(label="Counts") + plt.tight_layout() + plt.savefig(f"{base}/EA_vs_Esx3.png", dpi=300) + plt.close() + #Prop counter energy loss + plt.figure(figsize=(7,6)) + plt.hist2d(Eprop, sx3Z, bins=200) + plt.xlabel("EA - EC (MeV)") + plt.ylabel("SX3 Z") + plt.title("Energy Propagation Difference vs Position") + plt.colorbar(label="Counts") + plt.tight_layout() + plt.savefig(f"{base}/Eprop_vs_sx3Z.png", dpi=300) + plt.close() + + print("Plotting complete.") - +#exec(open("PCEnergyAnalysis.py").read()) if __name__ == "__main__": MyInteractiveApp().cmdloop()