178 lines
3.5 KiB
Python
178 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Wed Aug 5 15:52:39 2026
|
|
|
|
@author: jamesszalkie
|
|
"""
|
|
|
|
import uproot
|
|
import joblib
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import tensorflow as tf
|
|
from sklearn.metrics import mean_absolute_error
|
|
|
|
# User settings
|
|
|
|
ROOT_FILE = "/Users/jamesszalkie/ANASEN_analysis/Armory/SimAnasen1.root"
|
|
TREE_NAME = "tree1"
|
|
|
|
MODEL = "beam_predictor.keras"
|
|
INPUT_SCALER = "input_scaler.pkl"
|
|
OUTPUT_SCALER = "output_scaler.pkl"
|
|
|
|
INPUT_BRANCHES = [
|
|
"Tb",
|
|
"thetab",
|
|
"vZ",
|
|
"MBeam",
|
|
"MTarget",
|
|
"MLight",
|
|
"MHeavy",
|
|
]
|
|
|
|
# Load model
|
|
|
|
model = tf.keras.models.load_model(MODEL)
|
|
|
|
input_scaler = joblib.load(INPUT_SCALER)
|
|
output_scaler = joblib.load(OUTPUT_SCALER)
|
|
|
|
# Read ROOT file
|
|
|
|
with uproot.open(ROOT_FILE) as f:
|
|
tree = f[TREE_NAME]
|
|
arrays = tree.arrays(INPUT_BRANCHES, library="np")
|
|
|
|
X = np.column_stack([arrays[b] for b in INPUT_BRANCHES])
|
|
|
|
# Remove NaN events
|
|
mask = np.all(np.isfinite(X), axis=1)
|
|
X = X[mask]
|
|
|
|
print(f"Predicting {len(X)} events")
|
|
|
|
# Predict
|
|
|
|
X_scaled = input_scaler.transform(X)
|
|
|
|
pred_scaled = model.predict(X_scaled, verbose=0)
|
|
|
|
pred = output_scaler.inverse_transform(pred_scaled)
|
|
|
|
beam = pred[:,0]
|
|
Ex = pred[:,1]
|
|
|
|
# See whether truth branches exist
|
|
|
|
truth_beam = None
|
|
truth_Ex = None
|
|
|
|
with uproot.open(ROOT_FILE) as f:
|
|
|
|
tree = f[TREE_NAME]
|
|
branches = tree.keys()
|
|
|
|
if "beamEnergy" in branches:
|
|
truth_beam = tree["beamEnergy"].array(library="np")[mask]
|
|
|
|
if "Ex" in branches:
|
|
truth_Ex = tree["Ex"].array(library="np")[mask]
|
|
|
|
print("\n============================")
|
|
print("Model Performance")
|
|
print("============================")
|
|
|
|
if truth_beam is not None:
|
|
beam_mae = mean_absolute_error(truth_beam, beam)
|
|
beam_rmse = np.sqrt(np.mean((truth_beam - beam)**2))
|
|
|
|
print(f"Beam Energy MAE : {beam_mae:.4f} MeV")
|
|
print(f"Beam Energy RMSE: {beam_rmse:.4f} MeV")
|
|
|
|
if truth_Ex is not None:
|
|
Ex_mae = mean_absolute_error(truth_Ex, Ex)
|
|
Ex_rmse = np.sqrt(np.mean((truth_Ex - Ex)**2))
|
|
|
|
print(f"Excitation MAE : {Ex_mae:.4f} MeV")
|
|
print(f"Excitation RMSE : {Ex_rmse:.4f} MeV")
|
|
|
|
# Plot Beam Energy
|
|
plt.figure(figsize=(8,6))
|
|
|
|
plt.hist(
|
|
beam,
|
|
bins=250,
|
|
histtype="step",
|
|
linewidth=2,
|
|
label="Predicted",
|
|
)
|
|
|
|
if truth_beam is not None:
|
|
plt.hist(
|
|
truth_beam,
|
|
bins=250,
|
|
histtype="step",
|
|
linewidth=2,
|
|
label="Truth",
|
|
)
|
|
|
|
plt.xlabel("Beam Energy (MeV)")
|
|
plt.ylabel("Counts")
|
|
plt.title("Beam Energy")
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
|
|
# Plot Excitation Energy
|
|
|
|
plt.figure(figsize=(8,6))
|
|
|
|
plt.hist(
|
|
Ex,
|
|
bins=250,
|
|
histtype="step",
|
|
linewidth=2,
|
|
label="Predicted",
|
|
)
|
|
|
|
if truth_Ex is not None:
|
|
plt.hist(
|
|
truth_Ex,
|
|
bins=250,
|
|
histtype="step",
|
|
linewidth=2,
|
|
label="Truth",
|
|
)
|
|
|
|
plt.xlabel("Excitation Energy (MeV)")
|
|
plt.ylabel("Counts")
|
|
plt.title("Excitation Energy")
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
|
|
plt.figure(figsize=(6,6))
|
|
|
|
plt.scatter(truth_beam, beam, s=2)
|
|
|
|
mn = min(truth_beam.min(), beam.min())
|
|
mx = max(truth_beam.max(), beam.max())
|
|
|
|
plt.plot([mn, mx], [mn, mx], 'k--')
|
|
|
|
plt.xlabel("True Beam Energy (MeV)")
|
|
plt.ylabel("Predicted Beam Energy (MeV)")
|
|
plt.title("Beam Energy Reconstruction")
|
|
|
|
plt.figure(figsize=(6,6))
|
|
|
|
plt.scatter(truth_Ex, Ex, s=2)
|
|
|
|
mn = min(truth_Ex.min(), Ex.min())
|
|
mx = max(truth_Ex.max(), Ex.max())
|
|
|
|
plt.plot([mn, mx], [mn, mx], 'k--')
|
|
|
|
plt.xlabel("True Excitation Energy (MeV)")
|
|
plt.ylabel("Predicted Excitation Energy (MeV)")
|
|
plt.title("Excitation Energy Reconstruction") |