diff --git a/Armory/ANASEN_ML.py b/Armory/ANASEN_ML.py
index 3881eac..9e7f28a 100644
--- a/Armory/ANASEN_ML.py
+++ b/Armory/ANASEN_ML.py
@@ -15,13 +15,16 @@ from sklearn.metrics import mean_absolute_error
 
 # User settings
 
-ROOT_FILE = "/Users/jamesszalkie/ANASEN_analysis/Armory/SimAnasen1.root"
+ROOT_FILE = "/Users/jamesszalkie/ANASEN_analysis/Armory/Ne18Protons.root"
 TREE_NAME = "tree1"
 
 MODEL = "beam_predictor.keras"
 INPUT_SCALER = "input_scaler.pkl"
 OUTPUT_SCALER = "output_scaler.pkl"
 
+# Candidate excitation energies (MeV) for nearest-state snapping
+EX_CANDIDATES = np.array([0, 0.3, 1.7, 2.4, 2.8], dtype=np.float32)
+
 INPUT_BRANCHES = [
     "Tb",
     "thetab",
@@ -30,6 +33,10 @@ INPUT_BRANCHES = [
     "MTarget",
     "MLight",
     "MHeavy",
+    "ZBeam",
+    "ZHeavy",
+    "ZLight",
+    "ZTarget"
 ]
 
 # Load model
@@ -64,6 +71,10 @@ pred = output_scaler.inverse_transform(pred_scaled)
 beam = pred[:,0]
 Ex = pred[:,1]
 
+# Snap each predicted excitation to the nearest candidate energy.
+nearest_idx = np.argmin(np.abs(Ex[:, None] - EX_CANDIDATES[None, :]), axis=1)
+Ex_snapped = EX_CANDIDATES[nearest_idx]
+
 # See whether truth branches exist
 
 truth_beam = None
@@ -95,8 +106,15 @@ if truth_Ex is not None:
     Ex_mae = mean_absolute_error(truth_Ex, Ex)
     Ex_rmse = np.sqrt(np.mean((truth_Ex - Ex)**2))
 
+    Ex_snap_mae = mean_absolute_error(truth_Ex, Ex_snapped)
+    Ex_snap_rmse = np.sqrt(np.mean((truth_Ex - Ex_snapped)**2))
+
     print(f"Excitation MAE  : {Ex_mae:.4f} MeV")
     print(f"Excitation RMSE : {Ex_rmse:.4f} MeV")
+    print(f"Ex Snapped MAE  : {Ex_snap_mae:.4f} MeV")
+    print(f"Ex Snapped RMSE : {Ex_snap_rmse:.4f} MeV")
+
+print("\nEx candidate states (MeV):", ", ".join(f"{x:.3f}" for x in EX_CANDIDATES))
 
 # Plot Beam Energy
 plt.figure(figsize=(8,6))
@@ -133,7 +151,15 @@ plt.hist(
     bins=250,
     histtype="step",
     linewidth=2,
-    label="Predicted",
+    label="Predicted (raw)",
+)
+
+plt.hist(
+    Ex_snapped,
+    bins=250,
+    histtype="step",
+    linewidth=2,
+    label="Predicted (snapped)",
 )
 
 if truth_Ex is not None:
@@ -175,4 +201,18 @@ plt.plot([mn, mx], [mn, mx], 'k--')
 
 plt.xlabel("True Excitation Energy (MeV)")
 plt.ylabel("Predicted Excitation Energy (MeV)")
-plt.title("Excitation Energy Reconstruction")
\ No newline at end of file
+plt.title("Excitation Energy Reconstruction")
+
+if truth_Ex is not None:
+    plt.figure(figsize=(6,6))
+
+    plt.scatter(truth_Ex, Ex_snapped, s=2)
+
+    mn = min(truth_Ex.min(), Ex_snapped.min())
+    mx = max(truth_Ex.max(), Ex_snapped.max())
+
+    plt.plot([mn, mx], [mn, mx], 'k--')
+
+    plt.xlabel("True Excitation Energy (MeV)")
+    plt.ylabel("Snapped Excitation Energy (MeV)")
+    plt.title("Excitation Energy Reconstruction (Snapped)")
\ No newline at end of file