model_surrogate.py

Purpose

model_surrogate.py provides a Python class (MyModel) that loads and evaluates a trained neural-network surrogate model for UQ/calibration workflows. It wraps the pickled scikit-learn MLPRegressor created by surrogate_NN.py, applies parameter and output normalization, and exposes a simple run(parms) interface for fast ensemble prediction without rerunning the full ELM simulation.

The surrogate model replaces computationally expensive ensemble runs in MCMC sampling (MCMC.py) and global sensitivity analysis (run_GSA.py).

Architecture

MyModel is not a command-line script — it is a class imported by other UQ scripts:

import model_surrogate as models
model = models.MyModel(case="my_case_name")
model.run(parameter_samples)  # shape (n_samples, n_params)
predictions = model.output     # shape (n_samples, n_obs)

Initialization

MyModel(case="")

Loads pre-trained surrogate data from ./UQ_output/<case>/:

data/ptrain.dat — parameter training samples (n_train × n_params)
data/ytrain.dat — output training samples (n_train × n_obs)
data/pnames.txt — parameter names (one per line)
data/obs.dat — observed values and uncertainties (two columns)
data/outnames.txt — output variable names (one per line)
NN_surrogate/NNmodel.pkl — pickled scikit-learn MLPRegressor
NN_surrogate/qoi_good.txt — indices of outputs with non-zero variance

The constructor computes parameter bounds (pmin, pmax) and output ranges (yrange) from the training data and initializes a default parameter vector (pdef) at the midpoint of each range.

Class Attributes

After initialization, the following attributes are available:

nparms — number of parameters
nobs — number of output quantities of interest (QOIs)
ntrain — number of training samples
parm_names — list of parameter names (length nparms)
obs_name — list of output variable names (length nobs)
pmin, pmax — parameter bounds (shape nparms)
pdef — default parameter values (midpoint of bounds, shape nparms)
obs — observed values (shape nobs)
obs_err — observation uncertainties (shape nobs)
yrange — min/max range of each output in training data (shape [2, nobs])
qoi_good — indices of outputs with non-zero training variance
nnmodel — unpickled scikit-learn MLPRegressor object
output — last prediction result (populated by run())

Methods

run(parms)

Evaluate the surrogate model on one or more parameter vectors.

Arguments:

parms — NumPy array, shape (nparms,) for a single sample or (nsamples, nparms) for a batch.

Behavior:

Normalize input parameters to [0, 1] using training min/max.
Call self.nnmodel.predict() on normalized inputs.
Denormalize outputs to original scale using yrange.
For outputs with zero training variance (excluded during training), assign the training maximum.
Store result in self.output (shape (nsamples, nobs)).

Returns:

None (result is stored in self.output).

Training Data Requirements

The surrogate expects a fixed directory structure under ./UQ_output/<case>/:

UQ_output/
  └── <case>/
      ├── data/
      │   ├── ptrain.dat         # n_train × n_params, whitespace-separated
      │   ├── ytrain.dat         # n_train × n_obs, whitespace-separated
      │   ├── pval.dat           # optional validation parameters (used by surrogate_NN.py)
      │   ├── yval.dat           # optional validation outputs
      │   ├── pnames.txt         # parameter names, one per line
      │   ├── outnames.txt       # output variable names, one per line
      │   └── obs.dat            # observed values (col 1) and uncertainties (col 2)
      └── NN_surrogate/
          ├── NNmodel.pkl        # pickled MLPRegressor (created by surrogate_NN.py)
          ├── qoi_good.txt       # indices of non-constant outputs (one per line)
          └── fitstats.txt       # training/validation statistics (text summary)

Data generation:

Training/validation data are typically generated by:

Running an ensemble of ELM cases with manage_ensemble.py (see ensemble_run.py, ensemble_copy.py).
Post-processing output into the required .dat format.

The exact procedure is project-specific and not automated by OLMT.

Relationship to surrogate_NN.py

surrogate_NN.py is the training script that creates NNmodel.pkl.

Workflow:

User runs surrogate_NN.py --case <name> on pre-generated training/validation data.
surrogate_NN.py performs a randomized search over neural-network architectures (2–3 hidden layers, adaptive layer sizes), fits 100 candidate models, and saves the one with the best validation R².
The trained model is pickled to UQ_output/<case>/NN_surrogate/NNmodel.pkl.
model_surrogate.py loads the pickled model at runtime.

Key training parameters (hardcoded in surrogate_NN.py ):

Solver: Adam optimizer with early stopping
Tolerance: 1e-7
Max iterations: 200
Validation fraction: 0.2 (split from training data)
Hidden layer sizes: randomly sampled from \([10\sqrt{n+1}, 20\sqrt{n+1}]\) for layers 1 and 2

Outputs from surrogate_NN.py:

NNmodel.pkl — pickled MLPRegressor
qoi_good.txt — list of non-constant output indices
fitstats.txt — validation R² and RMSE per output
nnfit_qoi<i>.pdf — scatter plots of model vs. surrogate predictions

Command-line Arguments

None. model_surrogate.py is a library module, not a standalone script. It has no if __name__ == "__main__" block.

To train a surrogate, use:

python surrogate_NN.py --case <case_name>

To evaluate a trained surrogate in a UQ workflow, import MyModel in your script:

import model_surrogate as models
model = models.MyModel(case="my_case")
model.run(samples)

Surrogate Types Supported

Multi-layer perceptron (MLP) regressor only.

The current implementation is hard-coded to use sklearn.neural_network.MLPRegressor with:

2 or 3 fully connected hidden layers
Adam solver
Early stopping on a 20% validation holdout

No other surrogate types (Gaussian process, polynomial chaos expansion, random forest, etc.) are currently supported. To add a new surrogate type, you would need to:

Modify surrogate_NN.py to train the new model type.
Modify MyModel.run() to invoke the new model’s prediction interface.
Ensure the new model can be pickled/unpickled via pickle.

Example Usage

Typical usage in MCMC or GSA:

import numpy as np
import model_surrogate as models

# Load the trained surrogate
model = models.MyModel(case="my_ensemble_case")

# Generate parameter samples (Latin hypercube, Sobol, etc.)
n_samples = 10000
samples = np.zeros((n_samples, model.nparms))
for p in range(model.nparms):
    samples[:, p] = np.random.uniform(
        model.pmin[p], model.pmax[p], n_samples
    )

# Evaluate surrogate
model.run(samples)

# Access predictions
predictions = model.output  # shape (10000, nobs)

# Compare to observations
for i in range(model.nobs):
    print(f"{model.obs_name[i]}: obs={model.obs[i]}, "
          f"mean_pred={predictions[:, i].mean()}")