surrogate_NN.py

Purpose

surrogate_NN.py trains a neural network surrogate model to approximate the relationship between ELM parameter samples and model outputs (quantities of interest, QOIs). The surrogate is trained on data generated by model_surrogate.py and can be used for efficient sensitivity analysis and uncertainty quantification workflows that require many thousands of model evaluations.

The script implements a random-search hyperparameter optimization procedure that trains up to 100 candidate neural networks and selects the best-performing architecture based on validation R² scores.

Architecture

The neural network is implemented using scikit-learn’s MLPRegressor with the following characteristics:

Layer structure:

2 or 3 hidden layers (randomly selected during hyperparameter search)
Layer sizes scale with search iteration n:
- Layer 1: between 10√(n+1) and 20√(n+1) neurons
- Layer 2: between 20√(n+1) and 40√(n+1) neurons
- Layer 3 (if used): between 10√(n+1) and 20√(n+1) neurons
Input layer: number of parameters (from ptrain.dat)
Output layer: number of QOIs (from ytrain.dat)

Solver and training:

Solver: adam (adaptive moment estimation optimizer)
Early stopping: enabled with 20% internal validation split
Convergence tolerance: 1e-7
Maximum iterations: 200
Activation function: ReLU (default for MLPRegressor)

Training Procedure

Data loading: Reads training and validation datasets from UQ_output/<casename>/data/:
- ptrain.dat — training parameter samples
- ytrain.dat — training model outputs
- pval.dat — validation parameter samples
- yval.dat — validation model outputs
- outnames.txt — variable names for QOIs
Data cleaning: Removes any rows where the second QOI column is -9999 (failed runs).
Normalization: All inputs and outputs are normalized to [0, 1]:
- Parameters: (x - min(x_train)) / (max(x_train) - min(x_train))
- QOIs: (y - min(y_train)) / (max(y_train) - min(y_train))
- Validation data is clamped to [0, 1] to avoid extrapolation artifacts.
QOI filtering: QOIs with zero variance (min == max) are excluded from training and stored in qoi_good.txt.
Hyperparameter search: For up to 100 iterations (or until all QOIs achieve R² > 0.99):
- Randomly sample layer sizes as described above
- Randomly decide whether to use 2 or 3 hidden layers
- Train the network on normalized training data
- Evaluate on validation set
- Track best model by sum of R² across all QOIs
Model selection: The best-performing model is serialized to NNmodel.pkl using pickle.

Outputs

All outputs are written to UQ_output/<casename>/NN_surrogate/:

NNmodel.pkl

Serialized scikit-learn model, ready for use in prediction workflows

fitstats.txt

Training summary including:

Number of parameters, outputs, training/validation samples
Best network architecture (layer sizes)
Per-QOI R² and RMSE on validation set

qoi_good.txt

Indices of QOIs with nonzero variance (used during prediction)

nnfit_qoi<i>.pdf

Scatter plots of model vs. surrogate predictions for each QOI (validation set)

Hyperparameters

Fixed hyperparameters (not tuned):

Solver: adam
Convergence tolerance: 1e-7
Maximum iterations: 200
Early stopping: enabled (20% internal validation fraction)

Tuned hyperparameters (random search):

Number of hidden layers: 2 or 3
Layer sizes: scale with iteration number n as 10√(n+1) to 20√(n+1) (layer 1 and 3) or 20√(n+1) to 40√(n+1) (layer 2)
Search iterations: up to 100 (terminates early if all R² > 0.99)

Selection criterion:

Model with maximum sum of R² across all QOIs on the validation set.

Usage with model_surrogate.py

surrogate_NN.py is typically invoked after model_surrogate.py has generated training and validation datasets. The standard workflow is:

Generate ensemble samples using model_surrogate.py with a parameter distribution file (parm_list format):
```
python model_surrogate.py --case mycase --ns 500 --parm_list myparms.txt \
                          --postproc_file mypostproc_vars --model_root /path/to/E3SM
```
This runs 500 ensemble members and writes ptrain.dat, ytrain.dat, pval.dat, yval.dat to UQ_output/mycase/data/.
Train the surrogate:
```
python surrogate_NN.py --case mycase
```
This reads the datasets from step 1, performs hyperparameter search, and writes NNmodel.pkl and diagnostics.

Use the surrogate for forward prediction or sensitivity analysis:

import pickle
import numpy as np

# Load trained model
with open('UQ_output/mycase/NN_surrogate/NNmodel.pkl', 'rb') as f:
    model = pickle.load(f)

# Load normalization info (stored in fitstats.txt or recompute from ptrain/ytrain)
ptest_norm = (ptest - pmin) / (pmax - pmin)  # normalize new parameters

# Predict
ypredict_norm = model.predict(ptest_norm)
ypredict = ypredict_norm * (ymax - ymin) + ymin  # denormalize

The surrogate can then replace expensive ELM runs in workflows like MCMC calibration (MCMC.py) or global sensitivity analysis (run_GSA.py).

Notes

MPI is disabled: The script imports mpi4py but the relevant lines are commented out. All training happens in serial.
No cross-validation: The script uses a fixed training/validation split provided by model_surrogate.py. It does not perform k-fold cross-validation.
Overfitting risk: With up to 100 random architectures and no separate test set, there is a risk of overfitting to the validation set. Consider holding out a true test set if the surrogate will be used for extrapolation.
Normalization state: The normalization ranges (pmin, pmax, ymin, ymax) are not saved alongside the model. Users must recompute them from the original training data or parse ptrain.dat/ytrain.dat when making predictions.
Performance: Training 100 networks can take several minutes to hours depending on dataset size and layer sizes. Early stopping typically triggers well before the 200-iteration limit.