180 lines
3.3 KiB
Python
180 lines
3.3 KiB
Python
#!/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") |