{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4",
   "metadata": {},
   "source": [
    "# OpenMDAO Component Integration\n",
    "\n",
    "This notebook demonstrates how to wrap any Standard Evaluator as an OpenMDAO `ExplicitComponent` using `EvaluatorOpenMdaoComponent`. You will learn how OptProblem variables map to OpenMDAO inputs, how responses map to outputs (with bounds and scaling), and how to serialize an evaluator for later reconstruction inside an OpenMDAO model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c3d4e5",
   "metadata": {},
   "source": [
    "## How EvaluatorOpenMdaoComponent Works\n",
    "\n",
    "The `EvaluatorOpenMdaoComponent` class bridges the Standard Evaluator API and OpenMDAO:\n",
    "\n",
    "- **Variables → Inputs**: Each variable in the evaluator's `OptProblem` becomes an OpenMDAO input. The variable's `default` value sets the input default, and `bounds` define the feasible range.\n",
    "- **Responses → Outputs**: Each response becomes an OpenMDAO output. Response `bounds` map to `lower`/`upper`, and `shift`/`scale` map to OpenMDAO's `ref0`/`ref` scaling parameters.\n",
    "- **Compute**: When OpenMDAO calls `compute()`, the component builds a pandas DataFrame from the current inputs, calls the evaluator, and extracts the response values back into the OpenMDAO output vector.\n",
    "\n",
    "This means any evaluator — whether backed by a NumPy function, a surrogate model, or an external executable — can participate in an OpenMDAO model with a single line of wrapping code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d4e5f6",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We start by importing the Standard Evaluator library and OpenMDAO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e5f6a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import standard_evaluator as se\n",
    "from standard_evaluator.evaluators.numpy_evaluator import NumpyEvaluator\n",
    "from standard_evaluator.components import EvaluatorOpenMdaoComponent\n",
    "import openmdao.api as om\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f6a7b8",
   "metadata": {},
   "source": [
    "## Wrapping an Evaluator as an OpenMDAO Component\n",
    "\n",
    "We define a simple evaluator that computes $f(x, y) = x^2 + y^2$. First we create an `OptProblem` with two variables and one response, then implement the `NumpyEvaluator` subclass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6a7b8c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the optimization problem interface\n",
    "opt_problem = se.utilities.create_opt_problem(num_independent=2, num_dependent=1)\n",
    "opt_problem.variables[0].name = \"x\"\n",
    "opt_problem.variables[0].bounds = (-5.0, 5.0)\n",
    "opt_problem.variables[0].default = 1.0\n",
    "opt_problem.variables[1].name = \"y\"\n",
    "opt_problem.variables[1].bounds = (-5.0, 5.0)\n",
    "opt_problem.variables[1].default = 2.0\n",
    "opt_problem.responses[0].name = \"f\"\n",
    "opt_problem.responses[0].bounds = (0.0, 50.0)\n",
    "opt_problem.responses[0].scale = 0.02\n",
    "opt_problem.responses[0].shift = -25.0\n",
    "opt_problem.objectives = [\"f\"]\n",
    "opt_problem.constraints = []\n",
    "\n",
    "\n",
    "# Implement the evaluator\n",
    "class SumOfSquares(NumpyEvaluator):\n",
    "    \"\"\"Computes f = x^2 + y^2.\"\"\"\n",
    "\n",
    "    def eval_np(self, sites, **kwargs):\n",
    "        x = sites[:, 0]\n",
    "        y = sites[:, 1]\n",
    "        return (x**2 + y**2).reshape(-1, 1)\n",
    "\n",
    "\n",
    "evaluator = SumOfSquares(opt_problem=opt_problem)\n",
    "print(\"Inputs:\", evaluator.inputs)\n",
    "print(\"Outputs:\", evaluator.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7b8c9d0",
   "metadata": {},
   "source": [
    "Now we wrap the evaluator in an `EvaluatorOpenMdaoComponent`, add it to an OpenMDAO Problem, and run the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c9d0e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an OpenMDAO Problem with the wrapped evaluator\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem(\n",
    "    \"comp\", EvaluatorOpenMdaoComponent(evaluator), promotes=[\"*\"]\n",
    ")\n",
    "prob.setup()\n",
    "\n",
    "# Set input values and run\n",
    "prob.set_val(\"x\", 3.0)\n",
    "prob.set_val(\"y\", 4.0)\n",
    "prob.run_model()\n",
    "\n",
    "print(f\"x = {prob.get_val('x')[0]}\")\n",
    "print(f\"y = {prob.get_val('y')[0]}\")\n",
    "print(f\"f = x^2 + y^2 = {prob.get_val('f')[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d0e1f2",
   "metadata": {},
   "source": [
    "The component correctly computes $f = 3^2 + 4^2 = 25$. The evaluator's `OptProblem` defined the variable names and bounds, and the component automatically registered them as OpenMDAO inputs and outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e1f2a3",
   "metadata": {},
   "source": [
    "## Inspecting OpenMDAO Input and Output Metadata\n",
    "\n",
    "The component maps `OptProblem` variable/response properties to OpenMDAO metadata. We can retrieve this metadata to verify how bounds, defaults, and scaling are configured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f2a3b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "comp = prob.model.comp\n",
    "\n",
    "# Input metadata: default values and shape\n",
    "input_meta = comp.get_io_metadata(\n",
    "    iotypes=\"input\", metadata_keys=[\"val\", \"shape\", \"units\"]\n",
    ")\n",
    "print(\"Input metadata:\")\n",
    "print(\"-\" * 50)\n",
    "for abs_name, meta in input_meta.items():\n",
    "    print(f\"  {meta['prom_name']:>4s}: default = {meta['val']}, shape = {meta['shape']}\")\n",
    "\n",
    "print()\n",
    "\n",
    "# Output metadata: bounds and scaling\n",
    "output_meta = comp.get_io_metadata(\n",
    "    iotypes=\"output\", metadata_keys=[\"val\", \"shape\", \"lower\", \"upper\", \"ref0\", \"ref\"]\n",
    ")\n",
    "print(\"Output metadata:\")\n",
    "print(\"-\" * 50)\n",
    "for abs_name, meta in output_meta.items():\n",
    "    print(f\"  {meta['prom_name']:>4s}: lower = {meta['lower']}, upper = {meta['upper']}\")\n",
    "    print(f\"        ref0 = {meta['ref0']}, ref = {meta['ref']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a3b4c5",
   "metadata": {},
   "source": [
    "The metadata shows:\n",
    "\n",
    "- **Inputs** receive their `default` values from the OptProblem variables (x defaults to 1.0, y to 2.0).\n",
    "- **Outputs** receive `lower`/`upper` bounds from the response bounds (0.0 to 50.0), and `ref0`/`ref` are computed from the response's `shift` and `scale` parameters. OpenMDAO uses `ref0` and `ref` to normalize the output for optimizers: `normalized = (physical - ref0) / (ref - ref0)`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b4c5d6",
   "metadata": {},
   "source": [
    "## Serialization and Reconstruction via evaluator_options\n",
    "\n",
    "Surrogate models (which inherit from `NumpyEvaluator`) support serialization via `to_dict()`. The resulting dictionary can be passed as the `evaluator_options` keyword argument to construct a new `EvaluatorOpenMdaoComponent` — without needing the original Python class definition. This is useful for saving trained models and loading them into OpenMDAO workflows later.\n",
    "\n",
    "We demonstrate this by training a `PolynomialModel` on data from our evaluator, wrapping it in a component, then reconstructing an equivalent component from the serialized dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c5d6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from standard_evaluator.surrogate_models import PolynomialModel, PolynomialModelOptions\n",
    "\n",
    "# Generate training data from the evaluator\n",
    "np.random.seed(42)\n",
    "train_df = pd.DataFrame({\n",
    "    \"x\": np.random.uniform(-5, 5, 20),\n",
    "    \"y\": np.random.uniform(-5, 5, 20),\n",
    "})\n",
    "evaluator(train_df)\n",
    "\n",
    "# Train a degree-2 polynomial model (exact fit for x^2 + y^2)\n",
    "poly_model = PolynomialModel(\n",
    "    sites=train_df,\n",
    "    options=PolynomialModelOptions(degree=2),\n",
    "    opt_problem=opt_problem,\n",
    ")\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())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d6e7f8",
   "metadata": {},
   "source": [
    "Now we use the original polynomial model in an OpenMDAO component and compare its output to a reconstructed component built from the serialized dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6e7f8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original component using the trained polynomial model\n",
    "prob_original = om.Problem()\n",
    "prob_original.model.add_subsystem(\n",
    "    \"comp\", EvaluatorOpenMdaoComponent(poly_model), promotes=[\"*\"]\n",
    ")\n",
    "prob_original.setup()\n",
    "prob_original.set_val(\"x\", 3.0)\n",
    "prob_original.set_val(\"y\", 4.0)\n",
    "prob_original.run_model()\n",
    "original_f = prob_original.get_val(\"f\")[0]\n",
    "print(f\"Original component output:      f = {original_f:.6f}\")\n",
    "\n",
    "# Reconstructed component from evaluator_options\n",
    "prob_reconstructed = om.Problem()\n",
    "prob_reconstructed.model.add_subsystem(\n",
    "    \"comp\", EvaluatorOpenMdaoComponent(evaluator_options=model_dict), promotes=[\"*\"]\n",
    ")\n",
    "prob_reconstructed.setup()\n",
    "prob_reconstructed.set_val(\"x\", 3.0)\n",
    "prob_reconstructed.set_val(\"y\", 4.0)\n",
    "prob_reconstructed.run_model()\n",
    "reconstructed_f = prob_reconstructed.get_val(\"f\")[0]\n",
    "print(f\"Reconstructed component output: f = {reconstructed_f:.6f}\")\n",
    "\n",
    "print(f\"\\nOutputs match: {np.isclose(original_f, reconstructed_f)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7f8a9b0",
   "metadata": {},
   "source": [
    "The reconstructed component produces identical output. The `evaluator_options` keyword internally calls `SurrogateModel.from_dict()` to rebuild the trained model, so the full prediction capability is preserved without needing the original training data or class definition in scope."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}