{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Surrogate Models\n",
    "\n",
    "This notebook demonstrates how to train, evaluate, and serialize surrogate models using the Standard Evaluator library. Surrogate models are fast mathematical approximations of expensive evaluators, enabling rapid prediction of response values without re-running the underlying analysis. You will learn how to use the `PolynomialModel`, SMT-based models (such as RBF), serialize and reconstruct trained models, and configure model options."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "The SMT-based surrogate models require the optional `smt` dependency. Install it with:\n",
    "\n",
    "```bash\n",
    "pip install standard-evaluator[smt]\n",
    "```\n",
    "\n",
    "The `PolynomialModel` is available with the base installation and does not require any extra dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import standard_evaluator as se"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Hierarchy\n",
    "\n",
    "The surrogate model classes follow this inheritance hierarchy:\n",
    "\n",
    "- **`SurrogateModel`** (base class, inherits from `NumpyEvaluator`)\n",
    "  - **`PolynomialModel`** — Polynomial response surface model (base dependencies only)\n",
    "  - **`AbstractSmtModel`** — Base class for all SMT-based models (requires `smt`)\n",
    "    - **`RadialBasisFunctionModel`** (RBF) — Radial basis function interpolation\n",
    "    - **`InverseDistanceWeightingModel`** (IDW) — Inverse distance weighting interpolation\n",
    "    - **`GradientEnhancedNeuralNetworksModel`** (GENN) — Gradient-enhanced neural network\n",
    "    - **`LeastSquaresApproximationModel`** (LS) — Least-squares polynomial approximation\n",
    "    - **`RegularizedMinimalEnergyTensorProductBSplines`** (RMTS) — Regularized tensor-product B-splines\n",
    "    - **`SecondOrderPolynomialApproximationModel`** (SOPA) — Second-order polynomial approximation\n",
    "\n",
    "All surrogate models share a common API:\n",
    "- **Training**: Pass a `pd.DataFrame` of sites (inputs + true responses) to the constructor\n",
    "- **Prediction**: Call the model with a DataFrame of new input sites (modifies in-place)\n",
    "- **Serialization**: Use `to_dict()` and `from_dict()` for round-trip persistence\n",
    "- **Options**: Configure model behavior via a Pydantic options model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a PolynomialModel\n",
    "\n",
    "We start by training a `PolynomialModel` on sample data from the Sphere test evaluator. First we generate training sites by sampling within the variable bounds, then we evaluate the true responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a test evaluator (Sphere function: f = x0^2 + x1^2)\n",
    "evaluator = se.evaluators.test.Sphere(num_independent=2, num_dependent=1)\n",
    "\n",
    "# Get variable bounds from the opt_problem\n",
    "opt_problem = evaluator.opt_problem\n",
    "variables = opt_problem.variables\n",
    "\n",
    "# Generate 20 random training sites within bounds\n",
    "np.random.seed(42)\n",
    "n_train = 20\n",
    "train_data = {}\n",
    "for var in variables:\n",
    "    lb, ub = var.bounds\n",
    "    train_data[var.name] = np.random.uniform(lb, ub, n_train)\n",
    "\n",
    "train_sites = pd.DataFrame(train_data)\n",
    "\n",
    "# Evaluate true responses at training sites\n",
    "evaluator(train_sites)\n",
    "print(f\"Training sites shape: {train_sites.shape}\")\n",
    "train_sites.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we train a degree-2 polynomial model on these sites and predict at new test points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from standard_evaluator.surrogate_models import PolynomialModel, PolynomialModelOptions\n",
    "\n",
    "# Train the polynomial model (degree 2 fits a quadratic — perfect for Sphere)\n",
    "poly_model = PolynomialModel(\n",
    "    sites=train_sites,\n",
    "    options=PolynomialModelOptions(degree=2),\n",
    "    opt_problem=opt_problem,\n",
    ")\n",
    "\n",
    "# Generate 5 new test sites for prediction\n",
    "n_test = 5\n",
    "test_data = {}\n",
    "for var in variables:\n",
    "    lb, ub = var.bounds\n",
    "    test_data[var.name] = np.random.uniform(lb, ub, n_test)\n",
    "\n",
    "test_sites = pd.DataFrame(test_data)\n",
    "\n",
    "# Predict using the polynomial model (modifies DataFrame in place)\n",
    "poly_model(test_sites)\n",
    "predicted = test_sites[\"f\"].values.copy()\n",
    "\n",
    "# Compute true values for comparison\n",
    "true_sites = test_sites[evaluator.inputs].copy()\n",
    "evaluator(true_sites)\n",
    "true_values = true_sites[\"f\"].values\n",
    "\n",
    "# Display predicted vs true\n",
    "comparison = pd.DataFrame({\n",
    "    \"x0\": test_sites[\"x0\"],\n",
    "    \"x1\": test_sites[\"x1\"],\n",
    "    \"Predicted\": predicted,\n",
    "    \"True\": true_values,\n",
    "    \"Error\": np.abs(predicted - true_values),\n",
    "})\n",
    "print(\"Polynomial Model Predictions vs True Values:\")\n",
    "comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The degree-2 polynomial model fits the Sphere function exactly (since the Sphere function is itself a degree-2 polynomial), resulting in near-zero prediction errors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serialization and Reconstruction\n",
    "\n",
    "Surrogate models can be serialized to a dictionary using `to_dict()` and reconstructed using `SurrogateModel.from_dict()`. This enables saving trained models to disk and loading them later without retraining."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from standard_evaluator.surrogate_models import SurrogateModel\n",
    "\n",
    "# Serialize the trained model\n",
    "model_dict = poly_model.to_dict()\n",
    "print(f\"Serialized model type: {model_dict['type']}\")\n",
    "print(f\"Dictionary keys: {list(model_dict.keys())}\")\n",
    "\n",
    "# Reconstruct from dictionary\n",
    "reconstructed_model = SurrogateModel.from_dict(model_dict)\n",
    "\n",
    "# Verify predictions match\n",
    "verify_sites = pd.DataFrame(test_data)  # fresh copy of test inputs\n",
    "reconstructed_model(verify_sites)\n",
    "reconstructed_predictions = verify_sites[\"f\"].values\n",
    "\n",
    "# Compare original and reconstructed predictions\n",
    "max_diff = np.max(np.abs(predicted - reconstructed_predictions))\n",
    "print(f\"\\nMax difference between original and reconstructed predictions: {max_diff:.2e}\")\n",
    "print(\"Predictions match:\", np.allclose(predicted, reconstructed_predictions))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reconstructed model produces identical predictions to the original, confirming that the serialization round-trip preserves the model state completely."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training an RBF Model (SMT)\n",
    "\n",
    "The `RadialBasisFunctionModel` uses radial basis function interpolation from the SMT library. Let's train it on a more complex function and assess its accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from standard_evaluator.surrogate_models import RadialBasisFunctionModel\n",
    "\n",
    "# Use the Rosenbrock test evaluator (a harder, non-polynomial function)\n",
    "rosenbrock = se.evaluators.test.Rosenbrock(num_independent=2, num_dependent=1)\n",
    "rb_opt_problem = rosenbrock.opt_problem\n",
    "rb_variables = rb_opt_problem.variables\n",
    "\n",
    "# Generate 30 training sites (Latin hypercube-like random sampling)\n",
    "np.random.seed(123)\n",
    "n_rb_train = 30\n",
    "rb_train_data = {}\n",
    "for var in rb_variables:\n",
    "    lb, ub = var.bounds\n",
    "    rb_train_data[var.name] = np.random.uniform(lb, ub, n_rb_train)\n",
    "\n",
    "rb_train_sites = pd.DataFrame(rb_train_data)\n",
    "rosenbrock(rb_train_sites)\n",
    "\n",
    "# Train the RBF model\n",
    "rbf_model = RadialBasisFunctionModel(\n",
    "    sites=rb_train_sites,\n",
    "    opt_problem=rb_opt_problem,\n",
    ")\n",
    "\n",
    "# Generate 15 test sites\n",
    "n_rb_test = 15\n",
    "rb_test_data = {}\n",
    "for var in rb_variables:\n",
    "    lb, ub = var.bounds\n",
    "    rb_test_data[var.name] = np.random.uniform(lb, ub, n_rb_test)\n",
    "\n",
    "rb_test_sites = pd.DataFrame(rb_test_data)\n",
    "\n",
    "# Predict with RBF model\n",
    "rbf_model(rb_test_sites)\n",
    "rbf_predicted = rb_test_sites[\"f\"].values.copy()\n",
    "\n",
    "# Get true values\n",
    "rb_true_sites = rb_test_sites[rosenbrock.inputs].copy()\n",
    "rosenbrock(rb_true_sites)\n",
    "rbf_true = rb_true_sites[\"f\"].values\n",
    "\n",
    "# Compute RMSE\n",
    "rmse = np.sqrt(np.mean((rbf_predicted - rbf_true) ** 2))\n",
    "# Compute R² score\n",
    "ss_res = np.sum((rbf_true - rbf_predicted) ** 2)\n",
    "ss_tot = np.sum((rbf_true - np.mean(rbf_true)) ** 2)\n",
    "r_squared = 1 - ss_res / ss_tot\n",
    "\n",
    "print(f\"RBF Model Accuracy on Rosenbrock function:\")\n",
    "print(f\"  RMSE: {rmse:.4f}\")\n",
    "print(f\"  R²:   {r_squared:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The RBF model provides good interpolation accuracy. Since RBF is an interpolation method, it passes exactly through training points and approximates well between them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Options\n",
    "\n",
    "Each surrogate model class defines its configurable options via a Pydantic model. Options control model behavior such as polynomial degree, regularization parameters, or basis function settings.\n",
    "\n",
    "- **`PolynomialModelOptions`**: `degree` (int, default=1), `coefficient_ordering` (enum)\n",
    "- **`RadialBasisFunctionModelOptions`**: `d0` (float, default=1.0), `poly_degree` (int, default=-1), `reg` (float, default=1e-10)\n",
    "\n",
    "You can pass an options instance to the model constructor to override defaults. Let's see the effect of changing the polynomial degree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train a degree-1 (linear) polynomial model on the same Sphere data\n",
    "poly_linear = PolynomialModel(\n",
    "    sites=train_sites,\n",
    "    options=PolynomialModelOptions(degree=1),  # Linear fit (non-default)\n",
    "    opt_problem=opt_problem,\n",
    ")\n",
    "\n",
    "# Predict at test sites\n",
    "linear_test = pd.DataFrame(test_data)\n",
    "poly_linear(linear_test)\n",
    "linear_predicted = linear_test[\"f\"].values\n",
    "\n",
    "# Compare linear vs quadratic fit\n",
    "options_comparison = pd.DataFrame({\n",
    "    \"x0\": test_sites[\"x0\"],\n",
    "    \"x1\": test_sites[\"x1\"],\n",
    "    \"True\": true_values,\n",
    "    \"Degree 1 (linear)\": linear_predicted,\n",
    "    \"Degree 2 (quadratic)\": predicted,\n",
    "})\n",
    "print(\"Effect of polynomial degree on predictions:\")\n",
    "options_comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The degree-1 model cannot capture the quadratic Sphere function accurately, while the degree-2 model fits it exactly. This demonstrates how the `degree` option directly affects model capability and accuracy."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}