{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb9b77d8",
   "metadata": {},
   "source": [
    "# Surrogate Replacement Workflow\n",
    "\n",
    "This notebook demonstrates the full workflow for replacing a computationally expensive component in a multi-component OpenMDAO assembly with a surrogate model. The workflow uses `standard_evaluator` to:\n",
    "\n",
    "1. **Build** a multi-component assembly with aerodynamics, structures, and performance groups\n",
    "2. **Capture** the assembly interface and extract the expensive sub-group\n",
    "3. **Train** a polynomial surrogate on the isolated sub-group\n",
    "4. **Replace** the expensive component with the surrogate in the assembly description\n",
    "5. **Compare** three evaluation approaches: original assembly, locally-replaced assembly, and a global surrogate\n",
    "\n",
    "All components use `ExecComp` (no external dependencies), making this notebook fully reproducible. The aerodynamics group uses internal `connect()` calls to demonstrate proper handling of connection metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a12561",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import openmdao.api as om\n",
    "import standard_evaluator as se\n",
    "from standard_evaluator.evaluators import OpenMDAOEvaluator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9693b9a",
   "metadata": {},
   "source": [
    "## 1. Creating the Multi-Component Assembly\n",
    "\n",
    "We build an assembly with three sub-groups:\n",
    "\n",
    "- **aero** — Aerodynamics group computing dynamic pressure, lift, and drag. Uses internal `connect()` to pass the pressure output to lift and drag calculations. Mark the variables and responses that are internal with the 'internal' tag.\n",
    "- **structures** — Structural analysis computing weight and stress from material/geometry parameters.\n",
    "- **performance** — Performance metric combining aero and fuel parameters to compute range.\n",
    "\n",
    "The top-level assembly uses promotions (`promotes=['*']`) to connect the groups together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d364047",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the assembly\n",
    "assembly = om.Group()\n",
    "\n",
    "# --- Aero group (with internal connect) ---\n",
    "aero = om.Group()\n",
    "aero.add_subsystem(\n",
    "    'pressure_calc',\n",
    "    om.ExecComp('pressure = 0.5 * rho * v**2', pressure={'tags': 'internal'}),\n",
    "    promotes_inputs=['rho', 'v']\n",
    ")\n",
    "aero.add_subsystem(\n",
    "    'lift_calc',\n",
    "    om.ExecComp('lift = cl * q * area', q={'tags': 'internal'}),\n",
    "    promotes_inputs=['cl', 'area'],\n",
    "    promotes_outputs=['lift']\n",
    ")\n",
    "aero.add_subsystem(\n",
    "    'drag_calc',\n",
    "    om.ExecComp('drag = cd * q * area', q={'tags': 'internal'}),\n",
    "    promotes_inputs=['cd', 'area'],\n",
    "    promotes_outputs=['drag']\n",
    ")\n",
    "# Internal connections: pressure output feeds lift and drag calculations\n",
    "aero.connect('pressure_calc.pressure', 'lift_calc.q')\n",
    "aero.connect('pressure_calc.pressure', 'drag_calc.q')\n",
    "\n",
    "# --- Structures group ---\n",
    "structures = om.Group()\n",
    "structures.add_subsystem(\n",
    "    'weight_calc',\n",
    "    om.ExecComp('weight = density * volume * 9.81'),\n",
    "    promotes_inputs=['density', 'volume'],\n",
    "    promotes_outputs=['weight']\n",
    ")\n",
    "structures.add_subsystem(\n",
    "    'stress_calc',\n",
    "    om.ExecComp('stress = force / area_struct'),\n",
    "    promotes_inputs=['force', 'area_struct'],\n",
    "    promotes_outputs=['stress']\n",
    ")\n",
    "\n",
    "# --- Performance group ---\n",
    "performance = om.Group()\n",
    "performance.add_subsystem(\n",
    "    'range_calc',\n",
    "    om.ExecComp('range_out = (lift / drag) * (fuel / sfc)'),\n",
    "    promotes_inputs=['lift', 'drag', 'fuel', 'sfc'],\n",
    "    promotes_outputs=['range_out']\n",
    ")\n",
    "\n",
    "# --- Top-level assembly with promotions ---\n",
    "assembly.add_subsystem('aero', aero, promotes=['*'])\n",
    "assembly.add_subsystem('structures', structures, promotes=['*'])\n",
    "assembly.add_subsystem('performance', performance, promotes=['*'])\n",
    "\n",
    "# Create and setup the problem\n",
    "prob = om.Problem(model=assembly)\n",
    "prob.setup()\n",
    "prob.final_setup()\n",
    "\n",
    "# Set input values\n",
    "prob.set_val('rho', 1.225)       # air density [kg/m^3]\n",
    "prob.set_val('v', 50.0)          # velocity [m/s]\n",
    "prob.set_val('cl', 0.5)          # lift coefficient\n",
    "prob.set_val('cd', 0.03)         # drag coefficient\n",
    "prob.set_val('area', 20.0)       # wing area [m^2]\n",
    "prob.set_val('density', 2700.0)  # material density [kg/m^3]\n",
    "prob.set_val('volume', 0.01)     # structural volume [m^3]\n",
    "prob.set_val('force', 5000.0)    # applied force [N]\n",
    "prob.set_val('area_struct', 0.005)  # structural cross-section area [m^2]\n",
    "prob.set_val('fuel', 500.0)      # fuel mass [kg]\n",
    "prob.set_val('sfc', 0.5)         # specific fuel consumption [1/hr]\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print('Assembly outputs:')\n",
    "print(f\"  pressure = {prob.get_val('aero.pressure_calc.pressure')[0]:.4f} Pa\")\n",
    "print(f\"  lift     = {prob.get_val('lift')[0]:.4f} N\")\n",
    "print(f\"  drag     = {prob.get_val('drag')[0]:.4f} N\")\n",
    "print(f\"  weight   = {prob.get_val('weight')[0]:.4f} N\")\n",
    "print(f\"  stress   = {prob.get_val('stress')[0]:.4f} Pa\")\n",
    "print(f\"  range    = {prob.get_val('range_out')[0]:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1c3f8d2",
   "metadata": {},
   "source": [
    "## 2. Capturing the Interface and Extracting the Aero Sub-Group\n",
    "\n",
    "With the assembly built and verified, we now use `standard_evaluator` to:\n",
    "\n",
    "1. **Capture** the full assembly interface as a serializable `GroupInfo` description using `se.get_interface()`\n",
    "2. **Display** the hierarchical structure with `se.show_structure()` to see all components, inputs, outputs, promotions, and linkages\n",
    "3. **Extract** the aerodynamics sub-group from the assembly description and reconstruct it as a standalone OpenMDAO Problem using `se.create_problem()`\n",
    "4. **Wrap** the extracted problem as an `OpenMDAOEvaluator` so it can be sampled for surrogate training\n",
    "\n",
    "This demonstrates the key capability of `standard_evaluator`: converting between live OpenMDAO assemblies and portable interface descriptions that can be inspected, modified, and reconstructed.\n",
    "\n",
    "Note: If there are connections that are made inside a group that should stay internal it is important to add the `internal` tag to the inputs or outputs.\n",
    "\n",
    "We also set bounds for the variables and create an optimization problem that will be used later when exposing the OpenMDAO problem as an evaluator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2d4e7f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Capture the full assembly interface as a GroupInfo description\n",
    "assembly_info = se.get_interface(assembly)\n",
    "\n",
    "se.set_variable_bounds(assembly_info, {\n",
    "    'rho': (0.5, 2.0), 'v': (10, 100),\n",
    "    'cl': (0.1, 1.5), 'cd': (0.01, 0.3), 'area': (5, 50),\n",
    "    'density': (2000, 4000), 'volume': (0.001, 0.5),\n",
    "    'force': (4000.0, 6000.0), 'area_struct': (0.002, 0.008),\n",
    "    'fuel': (300.0, 600.0), 'sfc': (0.3, 0.8) })\n",
    "# Build the se problem with those bounds\n",
    "full_opt_problem = se.build_opt_problem(assembly_info, prob)\n",
    "\n",
    "# Display the hierarchical assembly structure\n",
    "print(\"Assembly structure:\")\n",
    "print(\"=\" * 60)\n",
    "se.show_structure(assembly_info)\n",
    "print(full_opt_problem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3e5a8g2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "# Extract the aero sub-group info from the assembly description\n",
    "aero_info = copy.deepcopy(assembly_info.components['aero'])\n",
    "\n",
    "# Set bounds on the serializable interface description\n",
    "se.set_variable_bounds(aero_info, {\n",
    "    'rho': (0.5, 2.0), 'v': (10, 100),\n",
    "    'cl': (0.1, 1.5), 'cd': (0.01, 0.3), 'area': (5, 50),\n",
    "})\n",
    "\n",
    "# Reconstruct the aero sub-group as a standalone Problem\n",
    "aero_prob = se.create_problem(aero_info)\n",
    "\n",
    "# Build the se problem and evaluator with those bounds\n",
    "opt_problem = se.build_opt_problem(aero_info)\n",
    "aero_evaluator = OpenMDAOEvaluator(aero_prob, opt_problem=opt_problem)\n",
    "\n",
    "print(\"Aero evaluator inputs:\", aero_evaluator.inputs)\n",
    "print(\"Aero evaluator outputs:\", aero_evaluator.outputs)\n",
    "print(\"Bounds on the inputs:\")\n",
    "for var in aero_evaluator.opt_problem.variables:\n",
    "    print(f\"{var.name}: {var.bounds}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4f1a2b3",
   "metadata": {},
   "source": [
    "## 3. Training a Polynomial Surrogate\n",
    "\n",
    "With the aero sub-group isolated as an evaluator, we can now train a polynomial surrogate to approximate it. The steps are:\n",
    "\n",
    "1. **Generate training sites** — Sample input combinations uniformly within the variable bounds of the aero evaluator\n",
    "2. **Evaluate training sites** — Run the aero evaluator on all training sites to get true response values\n",
    "3. **Build the surrogate** — Fit a degree-4 polynomial model to the training data\n",
    "4. **Wrap as OpenMDAO component** — Use `EvaluatorOpenMdaoComponent` to make the surrogate usable inside an OpenMDAO assembly\n",
    "5. **Create surrogate evaluator** — Wrap the surrogate component in a Problem and then in an `OpenMDAOEvaluator`\n",
    "\n",
    "A degree-2 polynomial is sufficient for these simple ExecComp equations and provides a fast approximation.\n",
    "\n",
    "For this comparison we are using the same number of training points for both the local and global model. In many cases additional savings might be able to be had by the fact that the local model has fewer inputs so we might be able to build a good model with fewer runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5g2h3i4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from standard_evaluator.surrogate_models import PolynomialModel, PolynomialModelOptions\n",
    "from standard_evaluator.components import EvaluatorOpenMdaoComponent\n",
    "from scipy.stats.qmc import LatinHypercube, scale\n",
    "\n",
    "# --- Step 1: Generate training sites within the aero evaluator bounds ---\n",
    "num_independent = len(aero_evaluator.inputs)\n",
    "n_train = 60\n",
    "\n",
    "# Use an optimized Latin Hypercube for better space-filling\n",
    "sampler = LatinHypercube(d=num_independent, optimization='random-cd', seed=42)\n",
    "sample_unit = sampler.random(n=n_train)\n",
    "\n",
    "# Scale from [0,1]^d to the variable bounds\n",
    "lower_bounds = [var.bounds[0] for var in aero_evaluator.opt_problem.variables]\n",
    "upper_bounds = [var.bounds[1] for var in aero_evaluator.opt_problem.variables]\n",
    "sample_scaled = scale(sample_unit, lower_bounds, upper_bounds)\n",
    "\n",
    "# Build training DataFrame\n",
    "train_data = {}\n",
    "for i, var in enumerate(aero_evaluator.opt_problem.variables):\n",
    "    train_data[var.name] = sample_scaled[:, i]\n",
    "    print(f\"Bounds on {var.name}: [{var.bounds[0]}, {var.bounds[1]}]\")\n",
    "\n",
    "train_sites = pd.DataFrame(train_data)\n",
    "\n",
    "# --- Step 2: Evaluate training sites using the aero evaluator ---\n",
    "aero_evaluator(train_sites)\n",
    "\n",
    "print(f\"Training data shape: {train_sites.shape}, using {n_train} training points\")\n",
    "print(f\"Input columns: {aero_evaluator.inputs}\")\n",
    "print(f\"Output columns: {aero_evaluator.outputs}\")\n",
    "train_sites.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6h3i4j5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Step 3: Build a degree-2 polynomial surrogate from training data ---\n",
    "# Create the options for the Polynomial model to have degree 2.\n",
    "poly_options = PolynomialModelOptions(degree=2)\n",
    "\n",
    "surrogate = PolynomialModel(\n",
    "    sites=train_sites,\n",
    "    options=poly_options,\n",
    "    opt_problem=aero_evaluator.opt_problem,\n",
    ")\n",
    "\n",
    "print(f\"Surrogate inputs: {surrogate.inputs}\")\n",
    "print(f\"Surrogate outputs: {surrogate.outputs}\")\n",
    "print(f\"Number of polynomial terms: {surrogate.nterms}\")\n",
    "print(f\"Polynomial degree: {surrogate.degree}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8333cee4",
   "metadata": {},
   "source": [
    "### 3.1 Checking the quality of the local model\n",
    "\n",
    "We explore how well the local model does in approximating the aero component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d59292",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Generate test sites (distinct from training data) ---\n",
    "n_test = 30\n",
    "\n",
    "# Use an optimized Latin Hypercube with a different seed for test points\n",
    "test_sampler = LatinHypercube(d=num_independent, optimization='random-cd', seed=999)\n",
    "test_unit = test_sampler.random(n=n_test)\n",
    "test_scaled = scale(test_unit, lower_bounds, upper_bounds)\n",
    "\n",
    "test_data = {}\n",
    "for i, var in enumerate(aero_evaluator.opt_problem.variables):\n",
    "    test_data[var.name] = test_scaled[:, i]\n",
    "\n",
    "# Create separate DataFrames for each evaluator (they modify in place)\n",
    "test_sites_original = pd.DataFrame(test_data).copy()\n",
    "test_sites_modified = pd.DataFrame(test_data).copy()\n",
    "\n",
    "# --- Evaluate all three ---\n",
    "# 1. Original assembly (ground truth)\n",
    "aero_evaluator(test_sites_original)\n",
    "\n",
    "# 2. Locally-replaced assembly\n",
    "surrogate(test_sites_modified)\n",
    "\n",
    "print(f\"Test sites: {n_test} points\")\n",
    "print(f\"Responses being compared: {aero_evaluator.outputs}\")\n",
    "# --- Compute error metrics ---\n",
    "# Using original as baseline truth\n",
    "error_records = []\n",
    "\n",
    "for response in aero_evaluator.outputs:\n",
    "    actual = test_sites_original[response].values\n",
    "\n",
    "    # Local surrogate-replaced assembly errors\n",
    "    if response in test_sites_modified.columns:\n",
    "        pred_local = test_sites_modified[response].values\n",
    "        abs_err_local = np.abs(pred_local - actual)\n",
    "        rel_err_local = np.abs((pred_local - actual)/actual)\n",
    "        error_records.append({\n",
    "            'Response': response,\n",
    "            'Evaluator': 'Local Surrogate',\n",
    "            'Mean Absolute Error': abs_err_local.mean(),\n",
    "            'Max Absolute Error': abs_err_local.max(),\n",
    "            'Mean Relative Error': rel_err_local.mean(),\n",
    "            'Max Relative Error': rel_err_local.max(),\n",
    "        })\n",
    "\n",
    "error_df = pd.DataFrame(error_records)\n",
    "\n",
    "# Display summary table\n",
    "print(\"Error Summary (relative to original assembly):\")\n",
    "print(\"=\" * 70)\n",
    "error_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "g7i4j5k6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Step 4: Wrap the surrogate as an OpenMDAO component ---\n",
    "surrogate_component = EvaluatorOpenMdaoComponent(surrogate)\n",
    "\n",
    "# --- Step 5: Create an OpenMDAO Problem wrapping the surrogate component ---\n",
    "surrogate_prob = om.Problem()\n",
    "surrogate_prob.model.add_subsystem(\n",
    "    'surrogate', surrogate_component, promotes=['*']\n",
    ")\n",
    "surrogate_prob.setup()\n",
    "surrogate_prob.final_setup()\n",
    "\n",
    "# Wrap the surrogate Problem as an OpenMDAOEvaluator\n",
    "surrogate_evaluator = OpenMDAOEvaluator(\n",
    "    surrogate_prob, scan_model=True, use_defined_problem=False\n",
    ")\n",
    "\n",
    "print(f\"Surrogate evaluator inputs: {surrogate_evaluator.inputs}\")\n",
    "print(f\"Surrogate evaluator outputs: {surrogate_evaluator.outputs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "h8j5k6l7",
   "metadata": {},
   "source": [
    "## 4. Replacing the Aero Component with the Surrogate\n",
    "\n",
    "Now we perform the key step: replacing the expensive aero sub-group in the assembly's interface description with the trained surrogate. The process is:\n",
    "\n",
    "1. **Capture** the surrogate evaluator's interface using `se.get_interface()` to obtain its `EvaluatorInfo`\n",
    "2. **Replace** the `'aero'` entry in the original assembly's `GroupInfo.components` instance with the surrogate's `EvaluatorInfo`\n",
    "3. **Reconstruct** the modified assembly using `se.create_problem()` to demonstrate the full round-trip\n",
    "4. **Wrap** the modified assembly as an `OpenMDAOEvaluator` for evaluation\n",
    "\n",
    "This demonstrates that `standard_evaluator`'s interface descriptions are fully editable — you can swap components at the description level and rebuild a working assembly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "i9k6l7m8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Step 1: Capture the surrogate's interface as an EvaluatorInfo ---\n",
    "surrogate_group_info = se.get_interface(surrogate_prob.model)\n",
    "surrogate_eval_info = surrogate_group_info.components['surrogate']\n",
    "\n",
    "print(f\"Surrogate interface type: {surrogate_eval_info.class_type}\")\n",
    "print(f\"Surrogate inputs: {[v.name for v in surrogate_eval_info.inputs]}\")\n",
    "print(f\"Surrogate outputs: {[v.name for v in surrogate_eval_info.outputs]}\")\n",
    "\n",
    "# --- Step 2: Replace aero in the assembly description ---\n",
    "modified_info = copy.deepcopy(assembly_info)\n",
    "modified_info.components['aero'] = surrogate_eval_info\n",
    "\n",
    "# --- Step 3: Instantiate the modified assembly ---\n",
    "modified_prob = se.create_problem(modified_info)\n",
    "\n",
    "# --- Step 4: Wrap as OpenMDAOEvaluator ---\n",
    "modified_evaluator = OpenMDAOEvaluator(\n",
    "    modified_prob, opt_problem=full_opt_problem\n",
    ")\n",
    "\n",
    "print(f\"\\nModified assembly evaluator inputs: {modified_evaluator.inputs}\")\n",
    "print(f\"Modified assembly evaluator outputs: {modified_evaluator.outputs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "j0l7m8n9",
   "metadata": {},
   "source": [
    "## 5. Building a Global Surrogate\n",
    "\n",
    "As an alternative to the local replacement strategy, we can build a **global surrogate** that maps the full assembly's inputs directly to all outputs, bypassing the component structure entirely.\n",
    "\n",
    "The process is:\n",
    "\n",
    "1. **Wrap the original assembly** as an `OpenMDAOEvaluator` so we can sample it\n",
    "2. **Generate training data** — Evaluate the original assembly on Latin Hypercube samples\n",
    "3. **Build a PolynomialModel** (degree=2) that approximates the full input→output mapping\n",
    "\n",
    "This approach trades structural fidelity for simplicity: a single polynomial replaces the entire assembly. It works well when the assembly's overall behavior is smooth and low-dimensional, but loses the ability to swap individual components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "k1m8n9o0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# --- Step 1: Wrap the original assembly as an OpenMDAOEvaluator ---\n",
    "original_evaluator = OpenMDAOEvaluator(prob, opt_problem=full_opt_problem)\n",
    "\n",
    "\n",
    "\n",
    "print(f\"Original assembly evaluator inputs: {original_evaluator.inputs}\")\n",
    "print(f\"Original assembly evaluator outputs: {original_evaluator.outputs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "l2n9o0p1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Step 2: Generate training data for the global surrogate ---\n",
    "global_num_independent = len(original_evaluator.inputs)\n",
    "n_train_global = n_train\n",
    "\n",
    "# Use an optimized Latin Hypercube for the global training data\n",
    "global_num_vars = len(original_evaluator.opt_problem.variables)\n",
    "global_sampler = LatinHypercube(d=global_num_vars, optimization='random-cd', seed=123)\n",
    "global_sample_unit = global_sampler.random(n=n_train_global)\n",
    "\n",
    "global_lower = [var.bounds[0] for var in original_evaluator.opt_problem.variables]\n",
    "global_upper = [var.bounds[1] for var in original_evaluator.opt_problem.variables]\n",
    "global_sample_scaled = scale(global_sample_unit, global_lower, global_upper)\n",
    "\n",
    "# Build training DataFrame\n",
    "global_train_data = {}\n",
    "for i, var in enumerate(original_evaluator.opt_problem.variables):\n",
    "    print(f\"{var.name}: {var.bounds}\")\n",
    "    global_train_data[var.name] = global_sample_scaled[:, i]\n",
    "\n",
    "global_train_sites = pd.DataFrame(global_train_data)\n",
    "\n",
    "# Evaluate the original assembly on all training sites\n",
    "original_evaluator(global_train_sites)\n",
    "\n",
    "print(f\"Global training data shape: {global_train_sites.shape}, using {n_train_global} training points (compared to {n_train} for local model)\")\n",
    "print(f\"Input columns: {original_evaluator.inputs}\")\n",
    "print(f\"Output columns: {original_evaluator.outputs}\")\n",
    "global_train_sites.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "m3o0p1q2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Step 3: Build a degree-2 polynomial mapping all assembly inputs to all outputs ---\n",
    "\n",
    "global_poly_options = PolynomialModelOptions(degree=2)\n",
    "\n",
    "global_surrogate = PolynomialModel(\n",
    "    sites=global_train_sites,\n",
    "    options=global_poly_options,\n",
    "    opt_problem=full_opt_problem,\n",
    ")\n",
    "\n",
    "print(f\"Global surrogate inputs: {global_surrogate.inputs}\")\n",
    "print(f\"Global surrogate outputs: {global_surrogate.outputs}\")\n",
    "print(f\"Number of polynomial terms: {global_surrogate.nterms}\")\n",
    "print(f\"Polynomial degree: {global_surrogate.degree}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "n4p1q2r3",
   "metadata": {},
   "source": [
    "## 6. Comparing All Three Evaluators\n",
    "\n",
    "We now compare the three evaluation approaches on a fresh set of **test sites** (distinct from training data) to assess accuracy:\n",
    "\n",
    "1. **Original assembly** — The ground truth reference (full OpenMDAO evaluation)\n",
    "2. **Locally-replaced assembly** — The assembly with the aero sub-group replaced by a local polynomial surrogate, while structures and performance groups remain as original components\n",
    "3. **Global surrogate** — A single polynomial mapping all assembly inputs directly to outputs\n",
    "\n",
    "**Methodology:**\n",
    "- Generate 30 random test sites using a distinct random seed (`seed=999`)\n",
    "- Evaluate all three approaches on the same test sites\n",
    "- Use the original assembly outputs as ground truth\n",
    "- Compute Mean Absolute Error (MAE) and Max Absolute Error for each response per evaluator\n",
    "- Visualize with scatter plots (predicted vs actual) and percent error box plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "o5q2r3s4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Generate test sites (distinct from training data) ---\n",
    "n_test = 30\n",
    "\n",
    "# Use an optimized Latin Hypercube with a different seed\n",
    "global_test_sampler = LatinHypercube(d=global_num_vars, optimization='random-cd', seed=999)\n",
    "global_test_unit = global_test_sampler.random(n=n_test)\n",
    "global_test_scaled = scale(global_test_unit, global_lower, global_upper)\n",
    "\n",
    "test_data = {}\n",
    "for i, var in enumerate(original_evaluator.opt_problem.variables):\n",
    "    test_data[var.name] = global_test_scaled[:, i]\n",
    "\n",
    "# Create separate DataFrames for each evaluator (they modify in place)\n",
    "test_sites_original = pd.DataFrame(test_data).copy()\n",
    "test_sites_modified = pd.DataFrame(test_data).copy()\n",
    "test_sites_global = pd.DataFrame(test_data).copy()\n",
    "\n",
    "# --- Evaluate all three ---\n",
    "# 1. Original assembly (ground truth)\n",
    "original_evaluator(test_sites_original)\n",
    "\n",
    "# 2. Locally-replaced assembly\n",
    "modified_evaluator(test_sites_modified)\n",
    "\n",
    "# 3. Global surrogate (PolynomialModel, call directly with DataFrame)\n",
    "global_surrogate(test_sites_global)\n",
    "\n",
    "print(f\"Test sites: {n_test} points (seed=999)\")\n",
    "print(f\"Responses being compared: {original_evaluator.outputs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "p6r3s4t5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Compute error metrics ---\n",
    "# Using original as baseline truth\n",
    "error_records = []\n",
    "\n",
    "for response in original_evaluator.outputs:\n",
    "    actual = test_sites_original[response].values\n",
    "\n",
    "    # Local surrogate-replaced assembly errors\n",
    "    if response in test_sites_modified.columns:\n",
    "        pred_local = test_sites_modified[response].values\n",
    "        abs_err_local = np.abs(pred_local - actual)\n",
    "        rel_err_local = np.abs((pred_local - actual)/actual)\n",
    "        error_records.append({\n",
    "            'Response': response,\n",
    "            'Evaluator': 'Local Surrogate',\n",
    "            'Mean Absolute Error': abs_err_local.mean(),\n",
    "            'Max Absolute Error': abs_err_local.max(),\n",
    "            'Mean Relative Error': rel_err_local.mean(),\n",
    "            'Max Relative Error': rel_err_local.max(),\n",
    "        })\n",
    "\n",
    "    # Global surrogate errors\n",
    "    if response in test_sites_global.columns:\n",
    "        pred_global = test_sites_global[response].values\n",
    "        abs_err_global = np.abs(pred_global - actual)\n",
    "        rel_err_global = np.abs((pred_global - actual)/actual)\n",
    "        error_records.append({\n",
    "            'Response': response,\n",
    "            'Evaluator': 'Global Surrogate',\n",
    "            'Mean Absolute Error': abs_err_global.mean(),\n",
    "            'Max Absolute Error': abs_err_global.max(),\n",
    "            'Mean Relative Error': rel_err_global.mean(),\n",
    "            'Max Relative Error': rel_err_global.max(),\n",
    "        })\n",
    "\n",
    "error_df = pd.DataFrame(error_records)\n",
    "\n",
    "# Display summary table\n",
    "print(\"Error Summary (relative to original assembly):\")\n",
    "print(\"=\" * 70)\n",
    "error_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "q7s4t5u6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Scatter plots: predicted vs actual for each response ---\n",
    "n_responses = len(original_evaluator.outputs)\n",
    "fig, axes = plt.subplots(1, n_responses, figsize=(4 * n_responses, 4))\n",
    "\n",
    "if n_responses == 1:\n",
    "    axes = [axes]\n",
    "\n",
    "for idx, response in enumerate(original_evaluator.outputs):\n",
    "    ax = axes[idx]\n",
    "    actual = test_sites_original[response].values\n",
    "\n",
    "    # Plot local surrogate predictions\n",
    "    if response in test_sites_modified.columns:\n",
    "        ax.scatter(actual, test_sites_modified[response].values,\n",
    "                   alpha=0.7, label='Local Surrogate', marker='o', s=30)\n",
    "\n",
    "    # Plot global surrogate predictions\n",
    "    if response in test_sites_global.columns:\n",
    "        ax.scatter(actual, test_sites_global[response].values,\n",
    "                   alpha=0.7, label='Global Surrogate', marker='^', s=30)\n",
    "\n",
    "    # Perfect prediction line\n",
    "    lims = [min(actual.min(), actual.min()), max(actual.max(), actual.max())]\n",
    "    ax.plot(lims, lims, 'k--', alpha=0.5, label='Perfect')\n",
    "\n",
    "    ax.set_xlabel('Actual (Original)')\n",
    "    ax.set_ylabel('Predicted')\n",
    "    ax.set_title(response)\n",
    "    ax.legend(fontsize=8)\n",
    "\n",
    "plt.suptitle('Predicted vs Actual for Each Response', fontsize=13, y=1.02)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "r8t5u6v7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Percent error box plots grouped by evaluator ---\n",
    "fig, axes = plt.subplots(1, n_responses, figsize=(4 * n_responses, 4))\n",
    "\n",
    "if n_responses == 1:\n",
    "    axes = [axes]\n",
    "\n",
    "for idx, response in enumerate(original_evaluator.outputs):\n",
    "    ax = axes[idx]\n",
    "    actual = test_sites_original[response].values\n",
    "\n",
    "    # Compute percent errors (avoid division by zero)\n",
    "    pct_errors = {}\n",
    "    mask = np.abs(actual) > 1e-10  # Only compute where actual is non-zero\n",
    "\n",
    "    if response in test_sites_modified.columns:\n",
    "        pred_local = test_sites_modified[response].values\n",
    "        pct_err_local = np.where(mask, (pred_local - actual) / actual * 100, 0.0)\n",
    "        pct_errors['Local\\nSurrogate'] = pct_err_local\n",
    "\n",
    "    if response in test_sites_global.columns:\n",
    "        pred_global = test_sites_global[response].values\n",
    "        pct_err_global = np.where(mask, (pred_global - actual) / actual * 100, 0.0)\n",
    "        pct_errors['Global\\nSurrogate'] = pct_err_global\n",
    "\n",
    "    if pct_errors:\n",
    "        ax.boxplot(pct_errors.values(), tick_labels=pct_errors.keys())\n",
    "\n",
    "    ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)\n",
    "    ax.set_ylabel('Percent Error (%)')\n",
    "    ax.set_title(response)\n",
    "\n",
    "plt.suptitle('Percent Error by Evaluator for Each Response', fontsize=13, y=1.02)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "435bef57",
   "metadata": {},
   "source": [
    "Looking at the results we see that the local model of course is perfect for `stress` and `weight` since those are not approximated by a surrogate model. The results for the responses from the aero component are not as clear."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}