This commit is contained in:
James Szalkie 2026-08-06 12:02:51 -04:00
parent 6ef362fca3
commit 9a1ad9bd39
3 changed files with 398 additions and 9 deletions

178
Armory/ANASEN_ML.py Normal file
View File

@ -0,0 +1,178 @@
#!/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")

180
Armory/TrainingBuilder.py Normal file
View File

@ -0,0 +1,180 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Train a TensorFlow model from a directory of ROOT files.
Inputs:
Tb
thetab
vZ
MBeam
MTarget
MLight
MHeavy
Outputs:
beamEnergy
Ex
"""
import numpy as np
from pathlib import Path
import joblib
import pandas as pd
import uproot
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.callbacks import EarlyStopping
ROOT_FOLDER = Path("/Users/jamesszalkie/ANASEN_analysis/Armory/Training_Data_Ne/")
TREE_NAME = "tree1"
INPUT_BRANCHES = [
"Tb",
"thetab",
"vZ",
"MBeam",
"MTarget",
"MLight",
"MHeavy",
]
OUTPUT_BRANCHES = [
"beamEnergy",
"Ex",
]
MODEL_NAME = "beam_predictor.keras"
print("Reading ROOT files...")
dfs = []
for root_file in sorted(ROOT_FOLDER.glob("*.root")):
print(f" {root_file.name}")
with uproot.open(root_file) as f:
tree = f[TREE_NAME]
arrays = tree.arrays(
INPUT_BRANCHES + OUTPUT_BRANCHES,
library="np"
)
dfs.append(pd.DataFrame(arrays))
dataset = pd.concat(dfs, ignore_index=True)
# Remove events with missing or invalid values
dataset = dataset.replace([np.inf, -np.inf], np.nan)
dataset = dataset.dropna()
print(f"Training on {len(dataset)} complete events.")
print(dataset.describe())
print()
print(dataset.isna().sum())
print()
print(np.isinf(dataset).sum())
print()
print("Total events:", len(dataset))
X = dataset[INPUT_BRANCHES].values
Y = dataset[OUTPUT_BRANCHES].values
X_train, X_test, Y_train, Y_test = train_test_split(
X,
Y,
test_size=0.20,
random_state=42,
)
input_scaler = StandardScaler()
output_scaler = StandardScaler()
X_train = input_scaler.fit_transform(X_train)
X_test = input_scaler.transform(X_test)
Y_train = output_scaler.fit_transform(Y_train)
Y_test = output_scaler.transform(Y_test)
model = Sequential([
Dense(128, activation="relu"),
Dense(128, activation="relu"),
Dense(64, activation="relu"),
Dense(32, activation="relu"),
Dense(2)
])
model.build((None, len(INPUT_BRANCHES)))
model.compile(
optimizer="adam",
loss="mse",
metrics=["mae"]
)
model.summary()
early_stop = EarlyStopping(
monitor="val_loss",
patience=20,
restore_best_weights=True
)
history = model.fit(
X_train,
Y_train,
epochs=25,
batch_size=512,
validation_split=0.20,
callbacks=[early_stop],
verbose=1,
)
loss, mae = model.evaluate(
X_test,
Y_test,
verbose=0,
)
pred_scaled = model.predict(X_test, verbose=0)
pred = output_scaler.inverse_transform(pred_scaled)
truth = output_scaler.inverse_transform(Y_test)
print()
print(f"Test Loss : {loss:.6f}")
print(f"Test MAE : {mae:.6f}")
beam_error = np.mean(np.abs(pred[:,0] - truth[:,0]))
Ex_error = np.mean(np.abs(pred[:,1] - truth[:,1]))
print(f"Beam Energy MAE: {beam_error:.6f} MeV")
print(f"Ex MAE: {Ex_error:.6f} MeV")
model.save(MODEL_NAME)
joblib.dump(input_scaler, "input_scaler.pkl")
joblib.dump(output_scaler, "output_scaler.pkl")
print()
print("Training complete.")
print("Saved:")
print(" ", MODEL_NAME)
print(" input_scaler.pkl")
print(" output_scaler.pkl")

View File

@ -98,14 +98,19 @@ int main(int argc, char **argv){
TransferReaction transfer;
//To set beam energy loss, use energy loss app, and create table with target isotope, set Initial beam energy as max energy
transfer.SetA(27, 13, 0); // 18Ne projectile
TGraph* elossBeam = LoadELoss("../ELoss/HeLoss/E_vs_x_Al-27.dat");
transfer.SetA(18, 9, 0); // 18Ne projectile
TGraph* elossBeam = LoadELoss("../ELoss/HeLoss/E_vs_x_Na-21.dat");
transfer.Seta(4, 2); // 4He target
transfer.Setb(1, 1); // outgoing proton from the primary transfer
transfer.SetB(30, 14); // 21Na* heavy product
const double beamA = 27; // mass number of 27Al beam
const double beamE = 72.0 / 27.0; // beam energy in MeV
transfer.SetB(21, 11); // 21Na* heavy product
const ReactionConfig reactionConfig = transfer.GetRectionConfig();
const double beamA = reactionConfig.beamA; // mass number of 14N beam
const double beamE = 42.82 / beamA; // beam energy in MeV
const int kMBeam = reactionConfig.beamA; // mass number of beam
const int kMTarget = reactionConfig.targetA; // mass number of target
const int kMLight = reactionConfig.recoilLightA; // mass number of light ejectile
const int kMHeavy = reactionConfig.recoilHeavyA; // mass number of heavy product
bool enableSequentialDecay = false; // turning to false to disable sequential decay for now, can be set to true to enable
const int decayDaughterA = 20;
const int decayDaughterZ = 10;
@ -113,13 +118,13 @@ int main(int argc, char **argv){
const int decayEjectZ = 1;
// Excited state lists (projectile and heavy-product excitation states)
std::vector<float> ExAList = {0}; // 18Ne projectile excitations in MeV
std::vector<float> ExList = {0}; // 21Na* excitation in MeV
std::vector<float> ExAList = {0}; // Beam excited energy
std::vector<float> ExList = {0, .3, 1.7, 2.4, 2.8, 3.5, 3.9, 4, 4.3, 4.5}; // Heavy product excited energy
// define vertex position uniform distribution ranges (mm)
double vertexXRange[2] = { -5, 5}; // mm - 5, 5
double vertexYRange[2] = { -5, 5}; // -5, 5
double vertexZRange[2] = { -174.3, 34.86}; // -174.3, 174.3 (full length of gas volume, centered at 0)
double vertexZRange[2] = { -174.3, 174.3}; // -174.3, 174.3 (full length of gas volume, centered at 0)
const double beamEntranceZ = -280; //vertexZRange[0]; // mm, assumed beam entrance into the gas
@ -168,9 +173,10 @@ int main(int argc, char **argv){
TString saveFileName = "SimAnasen1.root";
printf("\e[32m#################################### building Tree in %s\e[0m\n", saveFileName.Data());
TFile * saveFile = new TFile(saveFileName, "recreate");
//TFile * saveFile2 = new TFile("SimAnasen2.root", "recreate");
TTree * tree1 = new TTree("tree1", "tree1");
TTree * tree2 = new TTree("tree2", "tree2");
//TTree * tree3 = new TTree("tree3", "tree3"); // only includes Tb and thetab
// beam and CM variables saved in tree
double KEA;
@ -178,6 +184,10 @@ int main(int argc, char **argv){
double beamPath_cm;
double beamEnergy;
double beamEnergyLoss;
int MBeamOut;
int MTargetOut;
int MLightOut;
int MHeavyOut;
tree1->Branch("beamKEA", &KEA, "beamKEA/D");
tree2->Branch("beamKEA", &KEA2, "beamKEA/D");
tree1->Branch("beamPath_cm", &beamPath_cm, "beamPath_cm/D");
@ -186,6 +196,20 @@ int main(int argc, char **argv){
tree2->Branch("beamEnergy", &beamEnergy, "beamEnergy/D");
tree1->Branch("beamEnergyLoss", &beamEnergyLoss, "beamEnergyLoss/D");
tree2->Branch("beamEnergyLoss", &beamEnergyLoss, "beamEnergyLoss/D");
tree1->Branch("MBeam", &MBeamOut, "MBeam/I");
tree1->Branch("MTarget", &MTargetOut, "MTarget/I");
tree1->Branch("MLight", &MLightOut, "MLight/I");
tree1->Branch("MHeavy", &MHeavyOut, "MHeavy/I");
tree2->Branch("MBeam", &MBeamOut, "MBeam/I");
tree2->Branch("MTarget", &MTargetOut, "MTarget/I");
tree2->Branch("MLight", &MLightOut, "MLight/I");
tree2->Branch("MHeavy", &MHeavyOut, "MHeavy/I");
// constant reaction mass numbers stored in every event entry
MBeamOut = kMBeam;
MTargetOut = kMTarget;
MLightOut = kMLight;
MHeavyOut = kMHeavy;
double thetaCM, phiCM;
double thetaCM2, phiCM2;
@ -221,6 +245,9 @@ int main(int argc, char **argv){
tree2->Branch("qqqTb", &qqqTb2, "qqqTb/D");
tree2->Branch("qqqTB", &qqqTB2, "qqqTB/D");
//tree3->Branch("Tb", &Tb, "Tb/D");
//tree3->Branch("thetab", &thetab, "thetab/D");
// excitation state identifiers
int ExAID;
double ExA;
@ -479,6 +506,7 @@ int main(int argc, char **argv){
z0 = pw->GetZ0();
tree1->Fill();
//tree3->Fill();
// fill tree2 using the secondary proton from 21Na* decay
TVector3 dir2(1, 0, 0);
@ -592,6 +620,7 @@ int main(int argc, char **argv){
}
tree1->Fill();
//tree3->Fill();
TVector3 dir2(1, 0, 0);
dir2.SetTheta(thetab2 * TMath::DegToRad());
@ -676,8 +705,10 @@ int main(int argc, char **argv){
// write results to ROOT file and close
tree1->Write("", TObject::kOverwrite);
tree2->Write("", TObject::kOverwrite);
//tree3->Write("", TObject::kOverwrite);
int count1 = tree1->GetEntries();
int count2 = tree2->GetEntries();
//int count3 = tree3->GetEntries();
saveFile->Close();
printf("=============== done. saved as %s. tree1 entries: %d, tree2 entries: %d\n", saveFileName.Data(), count1, count2);