From 9a1ad9bd3962b3b739cac173af9921d6a4f9897b Mon Sep 17 00:00:00 2001 From: james szalkie Date: Thu, 6 Aug 2026 12:02:51 -0400 Subject: [PATCH] AnasenML --- Armory/ANASEN_ML.py | 178 +++++++++++++++++++++++++++++++++++++ Armory/TrainingBuilder.py | 180 ++++++++++++++++++++++++++++++++++++++ Armory/anasenMS.cpp | 49 +++++++++-- 3 files changed, 398 insertions(+), 9 deletions(-) create mode 100644 Armory/ANASEN_ML.py create mode 100644 Armory/TrainingBuilder.py diff --git a/Armory/ANASEN_ML.py b/Armory/ANASEN_ML.py new file mode 100644 index 0000000..3881eac --- /dev/null +++ b/Armory/ANASEN_ML.py @@ -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") \ No newline at end of file diff --git a/Armory/TrainingBuilder.py b/Armory/TrainingBuilder.py new file mode 100644 index 0000000..8ccb613 --- /dev/null +++ b/Armory/TrainingBuilder.py @@ -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") \ No newline at end of file diff --git a/Armory/anasenMS.cpp b/Armory/anasenMS.cpp index 458418f..e3c4a60 100644 --- a/Armory/anasenMS.cpp +++ b/Armory/anasenMS.cpp @@ -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 ExAList = {0}; // 18Ne projectile excitations in MeV - std::vector ExList = {0}; // 21Na* excitation in MeV + std::vector ExAList = {0}; // Beam excited energy + std::vector 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);