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
------------------

1. **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

2. **Data cleaning:**
   Removes any rows where the second QOI column is ``-9999`` (failed runs).

3. **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.

4. **QOI filtering:**
   QOIs with zero variance (``min == max``) are excluded from training and stored in ``qoi_good.txt``.

5. **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

6. **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:

1. **Generate ensemble samples** using ``model_surrogate.py`` with a parameter distribution file (``parm_list`` format):

   .. code-block:: bash

      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/``.

2. **Train the surrogate**:

   .. code-block:: bash

      python surrogate_NN.py --case mycase

   This reads the datasets from step 1, performs hyperparameter search, and writes ``NNmodel.pkl`` and diagnostics.

3. **Use the surrogate** for forward prediction or sensitivity analysis:

   .. code-block:: python

      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.