{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a1b2c3d4-0001-0001-0001-000000000001",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:07.091396Z",
     "iopub.status.busy": "2026-06-25T10:29:07.091196Z",
     "iopub.status.idle": "2026-06-25T10:29:07.096506Z",
     "shell.execute_reply": "2026-06-25T10:29:07.095259Z"
    },
    "papermill": {
     "duration": 0.010807,
     "end_time": "2026-06-25T10:29:07.097363+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:07.086556+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c423cbb6-af1f-4bb2-a7bf-3bc0f7344462",
   "metadata": {
    "papermill": {
     "duration": 0.003719,
     "end_time": "2026-06-25T10:29:07.105787+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:07.102068+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# OpenMDAO `Problem` Functional Interface\n",
    "\n",
    "The functional interface lets you wrap an assembled OpenMDAO `Problem` in a\n",
    "Python callable so that external tools — optimizers, surrogate trainers,\n",
    "uncertainty-quantification frameworks, etc. — can drive it using ordinary\n",
    "NumPy arrays instead of the OpenMDAO variable-name API.\n",
    "\n",
    "The entry point is `Problem.get_callback()`, which returns a\n",
    "`_FunctionalCallback` object.  Calling that object\n",
    "\n",
    "1. takes a flat input vector as an argument and writes it into the problem,\n",
    "2. calls `Problem.run_model()`, and\n",
    "3. returns outputs and/or total derivatives as NumPy arrays, depending on\n",
    "   which *form* was requested.\n",
    "\n",
    "## The three forms\n",
    "\n",
    "| `form` | Signature | Returns |\n",
    "|--------|-----------|----------|\n",
    "| `'f'` | `cb(x[, y])` | flat output vector `y` |\n",
    "| `'dfdx'` | `cb(x[, J])` | Jacobian `J` of shape `(n_outputs, n_inputs)` |\n",
    "| `'fdfdx'` | `cb(x[, y[, J]])` | tuple `(y, J)` |\n",
    "\n",
    "All three forms accept pre-allocated output arrays as optional arguments to\n",
    "avoid repeated memory allocation.\n",
    "\n",
    "## Specifying variables\n",
    "\n",
    "The `input_vars` and `output_vars` arguments to `get_callback()` control which\n",
    "model variables are included as inputs and outputs of the callback function.\n",
    "\n",
    "Each entry in either list may be:\n",
    "\n",
    "* a **plain string** — the variable name; or\n",
    "* a **dict** mapping an alias to a metadata dict with optional keys:\n",
    "  * `'name'` — actual variable name when the alias differs;\n",
    "  * `'indices'` — a subset of elements, in any format accepted by\n",
    "    `Problem.get_val()`;\n",
    "  * `'units'` — the units in which the variable should be expressed.  For\n",
    "    `form='f'` this controls the units used when reading back output values.\n",
    "    For the derivative forms (`'dfdx'` and `'fdfdx'`) this scales the\n",
    "    corresponding rows (output variable) or columns (input variable) of the\n",
    "    Jacobian so that derivatives are expressed in the requested units.\n",
    "\n",
    "For the derivative forms (`'dfdx'` and `'fdfdx'`), omitting `input_vars`\n",
    "causes the callback to use the driver's registered design variables, and\n",
    "omitting `output_vars` causes it to use the driver's registered responses\n",
    "(objectives + constraints).\n",
    "\n",
    "For `form='f'`, both arguments are required.\n",
    "\n",
    "Variables are packed into the flat vector in the order they appear in the\n",
    "list, with multi-element variables occupying contiguous slices.\n",
    "\n",
    "## Helper methods on the callback object\n",
    "\n",
    "| Method | Description |\n",
    "|--------|-------------|\n",
    "| `create_input_vector()` | Allocate a flat input array pre-filled with the current problem values. |\n",
    "| `create_output_vector()` | Allocate a flat output array pre-filled with the current problem values. |\n",
    "| `create_jacobian_matrix()` | Allocate a zero Jacobian of shape `(n_outputs, n_inputs)`. |\n",
    "| `get_input_val(name)` | Return the current flat value of one registered input variable. |\n",
    "| `get_output_val(name)` | Return the current flat value of one registered output variable. |\n",
    "| `input_var_names` | List of registered input variable names (or aliases). |\n",
    "| `output_var_names` | List of registered output variable names (or aliases). |\n",
    "| `form` | The form string this callback was created with. |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0002-0002-0002-000000000002",
   "metadata": {
    "papermill": {
     "duration": 0.003576,
     "end_time": "2026-06-25T10:29:07.112411+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:07.108835+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 1: evaluating outputs (`form='f'`)\n",
    "\n",
    "This example uses everyone's favorite `Paraboloid` component,\n",
    "$f(x, y) = (x-3)^2 + xy + (y+4)^2 - 3$, to show how to evaluate model\n",
    "outputs through the functional interface. In this example the OpenMDAO `Model` will consist of just this one `Component`, but the `Model` could be arbitrarily complex.\n",
    "\n",
    "Both `input_vars` and `output_vars` must be supplied when using `form='f'`.\n",
    "The callback accepts a flat vector whose elements correspond to the variables\n",
    "listed in `input_vars` in order, and returns a flat vector whose elements\n",
    "correspond to `output_vars` in order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a1b2c3d4-0003-0003-0003-000000000003",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:07.119899Z",
     "iopub.status.busy": "2026-06-25T10:29:07.119715Z",
     "iopub.status.idle": "2026-06-25T10:29:09.935675Z",
     "shell.execute_reply": "2026-06-25T10:29:09.934896Z"
    },
    "papermill": {
     "duration": 2.828752,
     "end_time": "2026-06-25T10:29:09.944416+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:07.115664+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input variable names : ['x', 'y']\n",
      "output variable names: ['f_xy']\n",
      "f(3, -4) = -15.0\n",
      "f(6.6667, -7.3333) = -27.333333330000006\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem('comp', Paraboloid(),\n",
    "                         promotes_inputs=['x', 'y'],\n",
    "                         promotes_outputs=['f_xy'])\n",
    "prob.setup()\n",
    "prob.final_setup()\n",
    "\n",
    "# Create a callback that maps [x, y] -> [f_xy].\n",
    "f = prob.get_callback('f', input_vars=['x', 'y'], output_vars=['f_xy'])\n",
    "\n",
    "print('input variable names :', f.input_var_names)\n",
    "print('output variable names:', f.output_var_names)\n",
    "\n",
    "# create_input_vector() returns a 1-D array pre-filled with the current\n",
    "# problem values.  Modify it in-place before each call.\n",
    "x = f.create_input_vector()\n",
    "x[0] = 3.0   # x\n",
    "x[1] = -4.0  # y\n",
    "\n",
    "y = f(x)\n",
    "print('f(3, -4) =', y[0])\n",
    "\n",
    "x[0] = 6.6667\n",
    "x[1] = -7.3333\n",
    "print('f(6.6667, -7.3333) =', f(x)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1b2c3d4-0004-0004-0004-000000000004",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.050879Z",
     "iopub.status.busy": "2026-06-25T10:29:10.050577Z",
     "iopub.status.idle": "2026-06-25T10:29:10.058879Z",
     "shell.execute_reply": "2026-06-25T10:29:10.058069Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.058235,
     "end_time": "2026-06-25T10:29:10.059606+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.001371+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "x[0], x[1] = 8.0, 9.0\n",
    "assert_near_equal(f(x)[0], (x[0]-3.0)**2 + x[0]*x[1] + (x[1]+4.0)**2 - 3.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0005-0005-0005-000000000005",
   "metadata": {
    "papermill": {
     "duration": 0.005548,
     "end_time": "2026-06-25T10:29:10.071095+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.065547+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Passing a pre-allocated output array\n",
    "\n",
    "Pass a pre-allocated array as `y` to avoid a memory allocation on every call.\n",
    "Use `create_output_vector()` to get a correctly-sized array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a1b2c3d4-0006-0006-0006-000000000006",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.083993Z",
     "iopub.status.busy": "2026-06-25T10:29:10.083731Z",
     "iopub.status.idle": "2026-06-25T10:29:10.088189Z",
     "shell.execute_reply": "2026-06-25T10:29:10.087542Z"
    },
    "papermill": {
     "duration": 0.011973,
     "end_time": "2026-06-25T10:29:10.088839+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.076866+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(9, 10) = 319.0\n"
     ]
    }
   ],
   "source": [
    "y_buf = f.create_output_vector()  # allocate once\n",
    "\n",
    "x[0], x[1] = 9.0, 10.0\n",
    "f(x, y=y_buf)  # result written into y_buf; return value is also y_buf\n",
    "print('f(9, 10) =', y_buf[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0007-0007-0007-000000000007",
   "metadata": {
    "papermill": {
     "duration": 0.006082,
     "end_time": "2026-06-25T10:29:10.100900+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.094818+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Holding some inputs fixed\n",
    "\n",
    "List only the variables you want to vary in `input_vars`.  Any variable\n",
    "omitted from the list keeps whatever value is currently stored in the problem.\n",
    "Set that fixed value with `Problem.set_val()` before creating the callback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a1b2c3d4-0008-0008-0008-000000000008",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.113429Z",
     "iopub.status.busy": "2026-06-25T10:29:10.113174Z",
     "iopub.status.idle": "2026-06-25T10:29:10.117925Z",
     "shell.execute_reply": "2026-06-25T10:29:10.117334Z"
    },
    "papermill": {
     "duration": 0.012154,
     "end_time": "2026-06-25T10:29:10.118528+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.106374+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(x=3, y=-4) = -15.0\n"
     ]
    }
   ],
   "source": [
    "prob.set_val('y', -4.0)  # fix y; only x will be varied\n",
    "\n",
    "f_x_only = prob.get_callback('f', input_vars=['x'], output_vars=['f_xy'])\n",
    "\n",
    "x1 = f_x_only.create_input_vector()  # length 1: only x\n",
    "x1[0] = 3.0\n",
    "print('f(x=3, y=-4) =', f_x_only(x1)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0009-0009-0009-000000000009",
   "metadata": {
    "papermill": {
     "duration": 0.005795,
     "end_time": "2026-06-25T10:29:10.130226+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.124431+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 2: total derivatives (`form='dfdx'`)\n",
    "\n",
    "Use `form='dfdx'` to obtain the Jacobian $\\partial f / \\partial x$.  The\n",
    "callback returns a 2-D array of shape `(n_outputs, n_inputs)`.\n",
    "\n",
    "When `input_vars` and `output_vars` are both omitted the callback\n",
    "automatically uses the driver's registered design variables as inputs and its\n",
    "registered responses (objectives + constraints) as outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a1b2c3d4-0010-0010-0010-000000000010",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.143660Z",
     "iopub.status.busy": "2026-06-25T10:29:10.143370Z",
     "iopub.status.idle": "2026-06-25T10:29:10.337626Z",
     "shell.execute_reply": "2026-06-25T10:29:10.336412Z"
    },
    "papermill": {
     "duration": 0.202206,
     "end_time": "2026-06-25T10:29:10.338412+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.136206+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input  vars: ['x', 'y']\n",
      "output vars: ['f_xy']\n",
      "J = [[-0.5 14.5]]\n",
      "shape: (1, 2)\n"
     ]
    }
   ],
   "source": [
    "prob2 = om.Problem()\n",
    "prob2.model.add_subsystem('comp', Paraboloid(),\n",
    "                          promotes_inputs=['x', 'y'],\n",
    "                          promotes_outputs=['f_xy'])\n",
    "prob2.model.add_design_var('x', lower=-50, upper=50)\n",
    "prob2.model.add_design_var('y', lower=-50, upper=50)\n",
    "prob2.model.add_objective('f_xy')\n",
    "prob2.driver = om.ScipyOptimizeDriver()\n",
    "prob2.setup()\n",
    "prob2.final_setup()\n",
    "\n",
    "# No input_vars / output_vars: falls back to driver design vars and responses.\n",
    "dfdx = prob2.get_callback('dfdx')\n",
    "\n",
    "print('input  vars:', dfdx.input_var_names)\n",
    "print('output vars:', dfdx.output_var_names)\n",
    "\n",
    "x = dfdx.create_input_vector()\n",
    "x[0] = 1.5  # x\n",
    "x[1] = 2.5  # y\n",
    "\n",
    "J = dfdx(x)\n",
    "print('J =', J)\n",
    "print('shape:', J.shape)  # (1, 2): one output, two inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a1b2c3d4-0011-0011-0011-000000000011",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.352332Z",
     "iopub.status.busy": "2026-06-25T10:29:10.351861Z",
     "iopub.status.idle": "2026-06-25T10:29:10.359843Z",
     "shell.execute_reply": "2026-06-25T10:29:10.357603Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.015681,
     "end_time": "2026-06-25T10:29:10.360521+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.344840+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J_expected = np.array([[2*(x[0]-3.0) + x[1], x[0] + 2*(x[1]+4.0)]])\n",
    "assert_near_equal(J, J_expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0012-0012-0012-000000000012",
   "metadata": {
    "papermill": {
     "duration": 0.006324,
     "end_time": "2026-06-25T10:29:10.373345+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.367021+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Passing a pre-allocated Jacobian\n",
    "\n",
    "Use `create_jacobian_matrix()` to allocate a zero matrix of the right shape,\n",
    "then pass it as the `J` keyword argument to avoid re-allocation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a1b2c3d4-0013-0013-0013-000000000013",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.387089Z",
     "iopub.status.busy": "2026-06-25T10:29:10.386837Z",
     "iopub.status.idle": "2026-06-25T10:29:10.391967Z",
     "shell.execute_reply": "2026-06-25T10:29:10.391238Z"
    },
    "papermill": {
     "duration": 0.01301,
     "end_time": "2026-06-25T10:29:10.392740+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.379730+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "J = [[-0.2 14.8]]\n"
     ]
    }
   ],
   "source": [
    "J_buf = dfdx.create_jacobian_matrix()  # shape (1, 2), all zeros\n",
    "\n",
    "x[0], x[1] = 1.6, 2.6\n",
    "dfdx(x, J=J_buf)\n",
    "print('J =', J_buf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6d50afc7-b67c-4028-b261-7f5466205b30",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.406199Z",
     "iopub.status.busy": "2026-06-25T10:29:10.405929Z",
     "iopub.status.idle": "2026-06-25T10:29:10.411862Z",
     "shell.execute_reply": "2026-06-25T10:29:10.411177Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.013423,
     "end_time": "2026-06-25T10:29:10.412477+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.399054+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J_expected = np.array([[2*(x[0]-3.0) + x[1], x[0] + 2*(x[1]+4.0)]])\n",
    "assert_near_equal(J_buf, J_expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0014-0014-0014-000000000014",
   "metadata": {
    "papermill": {
     "duration": 0.006453,
     "end_time": "2026-06-25T10:29:10.425439+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.418986+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 3: outputs and derivatives together (`form='fdfdx'`)\n",
    "\n",
    "Use `form='fdfdx'` when you need both the function value and the Jacobian from\n",
    "the same model evaluation.  The callback returns the tuple `(y, J)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a1b2c3d4-0015-0015-0015-000000000015",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.439561Z",
     "iopub.status.busy": "2026-06-25T10:29:10.439299Z",
     "iopub.status.idle": "2026-06-25T10:29:10.445774Z",
     "shell.execute_reply": "2026-06-25T10:29:10.444923Z"
    },
    "papermill": {
     "duration": 0.0149,
     "end_time": "2026-06-25T10:29:10.446491+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.431591+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(3, -4)  = -15.0\n",
      "df/dx     = -4.0\n",
      "df/dy     = 3.0\n"
     ]
    }
   ],
   "source": [
    "fdfdx = prob2.get_callback('fdfdx', input_vars=['x', 'y'],\n",
    "                            output_vars=['f_xy'])\n",
    "\n",
    "x = fdfdx.create_input_vector()\n",
    "x[0] = 3.0\n",
    "x[1] = -4.0\n",
    "\n",
    "y, J = fdfdx(x)\n",
    "print('f(3, -4)  =', y[0])\n",
    "print('df/dx     =', J[0, 0])\n",
    "print('df/dy     =', J[0, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a1b2c3d4-0016-0016-0016-000000000016",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.461225Z",
     "iopub.status.busy": "2026-06-25T10:29:10.460973Z",
     "iopub.status.idle": "2026-06-25T10:29:10.468059Z",
     "shell.execute_reply": "2026-06-25T10:29:10.467251Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.015754,
     "end_time": "2026-06-25T10:29:10.468784+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.453030+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert_near_equal(y[0], (x[0]-3.0)**2 + x[0]*x[1] + (x[1]+4.0)**2 - 3.0)\n",
    "assert_near_equal(J[0, 0], 2*(x[0]-3.0) + x[1])\n",
    "assert_near_equal(J[0, 1], x[0] + 2*(x[1]+4.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0017-0017-0017-000000000017",
   "metadata": {
    "papermill": {
     "duration": 0.006654,
     "end_time": "2026-06-25T10:29:10.482153+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.475499+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Pre-allocating both output arrays\n",
    "\n",
    "Pass pre-allocated `y` and `J` buffers together to eliminate all allocation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a1b2c3d4-0018-0018-0018-000000000018",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.495925Z",
     "iopub.status.busy": "2026-06-25T10:29:10.495653Z",
     "iopub.status.idle": "2026-06-25T10:29:10.501515Z",
     "shell.execute_reply": "2026-06-25T10:29:10.500834Z"
    },
    "papermill": {
     "duration": 0.013668,
     "end_time": "2026-06-25T10:29:10.502082+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.488414+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f(6.6667, -7.3333) = -27.333333330000006\n",
      "J = [[1.e-04 1.e-04]]\n"
     ]
    }
   ],
   "source": [
    "y_buf = fdfdx.create_output_vector()\n",
    "J_buf = fdfdx.create_jacobian_matrix()\n",
    "\n",
    "x[0], x[1] = 6.6667, -7.3333\n",
    "fdfdx(x, y=y_buf, J=J_buf)\n",
    "print('f(6.6667, -7.3333) =', y_buf[0])\n",
    "print('J =', J_buf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0019-0019-0019-000000000019",
   "metadata": {
    "papermill": {
     "duration": 0.00618,
     "end_time": "2026-06-25T10:29:10.514712+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.508532+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 4: selecting variable subsets with `indices`\n",
    "\n",
    "When a model variable is an array, you can select only a subset of its\n",
    "elements to be included as either an input or output to the functional interface by supplying an `'indices'` key in the variable\n",
    "metadata dict.  The indices may be an integer list (flat indexing) or a tuple\n",
    "of arrays (multi-dimensional indexing).\n",
    "\n",
    "The example below uses a component with two scalar inputs `x` and `y` of\n",
    "shape `(2,)` and one output `f` of shape `(2,)`.  The callback is configured\n",
    "to expose only `x[0]` and `f[0]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a1b2c3d4-0020-0020-0020-000000000020",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.528359Z",
     "iopub.status.busy": "2026-06-25T10:29:10.528098Z",
     "iopub.status.idle": "2026-06-25T10:29:10.543212Z",
     "shell.execute_reply": "2026-06-25T10:29:10.542515Z"
    },
    "papermill": {
     "duration": 0.023106,
     "end_time": "2026-06-25T10:29:10.543833+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.520727+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x3 = [-2.]\n",
      "f[0]         = -4.0\n",
      "df[0]/dx[0]  = -0.0\n"
     ]
    }
   ],
   "source": [
    "class QuadComp(om.ExplicitComponent):\n",
    "    \"\"\"f[0] = (x[0]+2)^2 + (x[1]+3)^2 - 4,  f[1] = 2*x[0] + 3*x[1].\"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=0.0, shape=(2,))\n",
    "        self.add_output('f', val=0.0, shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('*', '*', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['x']\n",
    "        outputs['f'][0] = (x[0]+2)**2 + (x[1]+3)**2 - 4\n",
    "        outputs['f'][1] = 2*x[0] + 3*x[1]\n",
    "\n",
    "\n",
    "prob3 = om.Problem()\n",
    "prob3.model.add_subsystem('comp', QuadComp(),\n",
    "                           promotes_inputs=['x'], promotes_outputs=['f'])\n",
    "prob3.model.add_design_var('x', lower=-50, upper=50)\n",
    "prob3.model.add_objective('f', index=0)\n",
    "prob3.driver = om.ScipyOptimizeDriver()\n",
    "prob3.setup(force_alloc_complex=True)\n",
    "prob3.final_setup()\n",
    "prob3.set_val('x', [-2.0, -3.0])\n",
    "prob3.run_model()\n",
    "\n",
    "# Expose only x[0] as input and f[0] as output.\n",
    "fdfdx3 = prob3.get_callback(\n",
    "    'fdfdx',\n",
    "    input_vars=[{'x': {'indices': [0]}}],\n",
    "    output_vars=[{'f': {'indices': [0]}}],\n",
    ")\n",
    "\n",
    "x3 = fdfdx3.create_input_vector()  # length 1\n",
    "print('x3 =', x3)                  # reflects current prob value of x[0]\n",
    "\n",
    "f3, J3 = fdfdx3(x3)\n",
    "print('f[0]         =', f3[0])\n",
    "print('df[0]/dx[0]  =', J3[0, 0])   # shape (1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a1b2c3d4-0021-0021-0021-000000000021",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.558083Z",
     "iopub.status.busy": "2026-06-25T10:29:10.557842Z",
     "iopub.status.idle": "2026-06-25T10:29:10.562318Z",
     "shell.execute_reply": "2026-06-25T10:29:10.561648Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.01293,
     "end_time": "2026-06-25T10:29:10.563315+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.550385+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert_near_equal(J3[0, 0], 2*(x3[0]+2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0022-0022-0022-000000000022",
   "metadata": {
    "papermill": {
     "duration": 0.006275,
     "end_time": "2026-06-25T10:29:10.576207+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.569932+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 5: driver variables used automatically\n",
    "\n",
    "When a `Problem` has design variables and responses registered on the driver,\n",
    "calling `get_callback('dfdx')` (or `'fdfdx'`) without `input_vars` /\n",
    "`output_vars` automatically selects those driver variables.  This is\n",
    "convenient after running an optimization: the callback immediately reflects\n",
    "the driver's variable set including any index subsets specified on the\n",
    "constraint or objective.\n",
    "\n",
    "The `input_var_names` and `output_var_names` properties contain which variables were selected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a1b2c3d4-0023-0023-0023-000000000023",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.591802Z",
     "iopub.status.busy": "2026-06-25T10:29:10.588639Z",
     "iopub.status.idle": "2026-06-25T10:29:10.609755Z",
     "shell.execute_reply": "2026-06-25T10:29:10.609229Z"
    },
    "papermill": {
     "duration": 0.027964,
     "end_time": "2026-06-25T10:29:10.610666+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.582702+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization terminated successfully    (Exit mode 0)\n",
      "            Current function value: 47.99999999999999\n",
      "            Iterations: 3\n",
      "            Function evaluations: 4\n",
      "            Gradient evaluations: 3\n",
      "Optimization Complete\n",
      "-----------------------------------\n",
      "inputs : ['x']\n",
      "outputs: ['f_obj', 'f_con']\n",
      "x at optimum: [2. 3.]\n",
      "f4 = [48. 13.]\n",
      "J4 =\n",
      " [[ 8. 12.]\n",
      " [ 2.  3.]]\n"
     ]
    }
   ],
   "source": [
    "prob4 = om.Problem()\n",
    "prob4.model.add_subsystem('comp', QuadComp(),\n",
    "                           promotes_inputs=['x'], promotes_outputs=['f'])\n",
    "prob4.model.add_design_var('x', lower=-50, upper=50)\n",
    "prob4.model.add_objective('f', index=0, alias='f_obj')\n",
    "prob4.model.add_constraint('f', indices=[1], equals=13.0, alias='f_con')\n",
    "prob4.driver = om.ScipyOptimizeDriver()\n",
    "prob4.setup(force_alloc_complex=True)\n",
    "prob4.set_val('x', [0.0, 0.0])\n",
    "prob4.run_driver()\n",
    "\n",
    "# Both input_vars and output_vars omitted: uses driver design vars + responses.\n",
    "fdfdx4 = prob4.get_callback('fdfdx')\n",
    "print('inputs :', fdfdx4.input_var_names)\n",
    "print('outputs:', fdfdx4.output_var_names)\n",
    "\n",
    "x4 = fdfdx4.create_input_vector()\n",
    "print('x at optimum:', x4)   # pre-filled with the post-optimization values\n",
    "\n",
    "f4, J4 = fdfdx4(x4)\n",
    "print('f4 =', f4)\n",
    "print('J4 =\\n', J4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a1b2c3d4-0024-0024-0024-000000000024",
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.624105Z",
     "iopub.status.busy": "2026-06-25T10:29:10.623887Z",
     "iopub.status.idle": "2026-06-25T10:29:10.629417Z",
     "shell.execute_reply": "2026-06-25T10:29:10.629023Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.013242,
     "end_time": "2026-06-25T10:29:10.630380+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.617138+00:00",
     "status": "completed"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# At the optimum x = [2, 3].\n",
    "assert_near_equal(x4, [2.0, 3.0], tolerance=1e-5)\n",
    "\n",
    "J4_expected = np.array([\n",
    "    [2*(x4[0]+2), 2*(x4[1]+3)],  # df_obj/dx\n",
    "    [2,           3           ],  # df_con/dx\n",
    "])\n",
    "assert_near_equal(J4, J4_expected, tolerance=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0025-0025-0025-000000000025",
   "metadata": {
    "papermill": {
     "duration": 0.005794,
     "end_time": "2026-06-25T10:29:10.642229+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.636435+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Inspecting current values after a call\n",
    "\n",
    "After calling the callback the problem is left in the state corresponding to\n",
    "the last `x` that was passed in.  The helper methods `get_input_val()` and\n",
    "`get_output_val()` retrieve the current flat values of any registered\n",
    "variable without another model evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a1b2c3d4-0026-0026-0026-000000000026",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.655409Z",
     "iopub.status.busy": "2026-06-25T10:29:10.655018Z",
     "iopub.status.idle": "2026-06-25T10:29:10.661941Z",
     "shell.execute_reply": "2026-06-25T10:29:10.661067Z"
    },
    "papermill": {
     "duration": 0.014349,
     "end_time": "2026-06-25T10:29:10.662489+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.648140+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x from callback  : [3.]\n",
      "y from callback  : [-15.]\n"
     ]
    }
   ],
   "source": [
    "# Continue from Example 3.\n",
    "x_in = fdfdx.create_input_vector()\n",
    "x_in[0] = 3.0\n",
    "x_in[1] = -4.0\n",
    "fdfdx(x_in)\n",
    "\n",
    "print('x from callback  :', fdfdx.get_input_val('x'))\n",
    "print('y from callback  :', fdfdx.get_output_val('f_xy'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2vld9hd5yhb",
   "metadata": {
    "papermill": {
     "duration": 0.005883,
     "end_time": "2026-06-25T10:29:10.674773+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.668890+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 6: unit conversion in Jacobians\n",
    "\n",
    "When model variables have physical units the Jacobian entries are expressed in those units.  The functional interface lets you request different\n",
    "units for any input or output variable by adding a `'units'` key to the\n",
    "variable metadata dict.  The Jacobian is then automatically scaled so that\n",
    "derivatives are expressed in the requested units.\n",
    "\n",
    "Consider a component that computes `f = 2 * x` where `x` is in metres and\n",
    "`f` is in Newtons.  The native Jacobian is $\\partial f / \\partial x = 2\\ \\text{N/m}$.\n",
    "\n",
    "* Requesting `f` in **kN** scales the output rows by $10^{-3}$, giving\n",
    "  $0.002\\ \\text{kN/m}$.\n",
    "* Requesting `x` in **km** scales the input columns by $10^{3}$, giving\n",
    "  $2000\\ \\text{N/km}$.\n",
    "* Requesting both gives $2\\ \\text{kN/km}$ (the two scalings cancel)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4eb5dcaf-24ee-4902-b2e0-d72fbabf1015",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.688027Z",
     "iopub.status.busy": "2026-06-25T10:29:10.687691Z",
     "iopub.status.idle": "2026-06-25T10:29:10.706165Z",
     "shell.execute_reply": "2026-06-25T10:29:10.705630Z"
    },
    "papermill": {
     "duration": 0.026279,
     "end_time": "2026-06-25T10:29:10.706957+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.680678+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "df[N]/dx[m]   = 2.0000  (expected 2.0)"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "df[kN]/dx[m]  = 0.002000  (expected 0.002)\n",
      "df[N]/dx[km]  = 2000.0  (expected 2000.0)\n",
      "df[kN]/dx[km] = 2.0000  (expected 2.0)\n"
     ]
    }
   ],
   "source": [
    "class ForceComp(om.ExplicitComponent):\n",
    "    \"\"\"f = 2 * x,  x in metres,  f in Newtons.\"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=1.0, units='m')\n",
    "        self.add_output('f', val=0.0, units='N')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('f', 'x')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['f'] = 2.0 * inputs['x']\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        partials['f', 'x'] = 2.0\n",
    "\n",
    "\n",
    "prob5 = om.Problem()\n",
    "prob5.model.add_subsystem('comp', ForceComp(), promotes=['*'])\n",
    "prob5.model.add_design_var('x', lower=0.0, upper=10.0)\n",
    "prob5.model.add_objective('f')\n",
    "prob5.driver = om.ScipyOptimizeDriver()\n",
    "prob5.setup()\n",
    "prob5.final_setup()\n",
    "prob5.set_val('x', 3.0, units='m')\n",
    "prob5.run_model()\n",
    "\n",
    "# --- native units: df[N] / dx[m] ---\n",
    "dfdx_native = prob5.get_callback('dfdx', input_vars=['x'], output_vars=['f'])\n",
    "x5 = dfdx_native.create_input_vector()\n",
    "J_native = dfdx_native(x5)\n",
    "print(f'df[N]/dx[m]   = {J_native[0, 0]:.4f}  (expected 2.0)')\n",
    "\n",
    "# --- output in kN: df[kN] / dx[m] ---\n",
    "dfdx_kN = prob5.get_callback('dfdx',\n",
    "                              input_vars=['x'],\n",
    "                              output_vars=[{'f': {'units': 'kN'}}])\n",
    "J_kN = dfdx_kN(x5)\n",
    "print(f'df[kN]/dx[m]  = {J_kN[0, 0]:.6f}  (expected 0.002)')\n",
    "\n",
    "# --- input in km: df[N] / dx[km] ---\n",
    "dfdx_km = prob5.get_callback('dfdx',\n",
    "                              input_vars=[{'x': {'units': 'km'}}],\n",
    "                              output_vars=['f'])\n",
    "J_km = dfdx_km(x5)\n",
    "print(f'df[N]/dx[km]  = {J_km[0, 0]:.1f}  (expected 2000.0)')\n",
    "\n",
    "# --- both: df[kN] / dx[km] ---\n",
    "dfdx_both = prob5.get_callback('dfdx',\n",
    "                                input_vars=[{'x': {'units': 'km'}}],\n",
    "                                output_vars=[{'f': {'units': 'kN'}}])\n",
    "J_both = dfdx_both(x5)\n",
    "print(f'df[kN]/dx[km] = {J_both[0, 0]:.4f}  (expected 2.0)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "inmzmf0slj",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.720463Z",
     "iopub.status.busy": "2026-06-25T10:29:10.720250Z",
     "iopub.status.idle": "2026-06-25T10:29:10.726378Z",
     "shell.execute_reply": "2026-06-25T10:29:10.725486Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.013782,
     "end_time": "2026-06-25T10:29:10.726880+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.713098+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert_near_equal(J_native[0, 0], 2.0)\n",
    "assert_near_equal(J_kN[0, 0],     2.0 / 1000.0)\n",
    "assert_near_equal(J_km[0, 0],     2.0 * 1000.0)\n",
    "assert_near_equal(J_both[0, 0],   2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ywzi2efjn",
   "metadata": {
    "papermill": {
     "duration": 0.006038,
     "end_time": "2026-06-25T10:29:10.739297+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.733259+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 7: using `return_index_map` to locate variables in flat vectors\n",
    "\n",
    "The `create_input_vector()`, `create_output_vector()`, and\n",
    "`create_jacobian_matrix()` methods all accept a `return_index_map` keyword\n",
    "argument.  When set to `True`, each method returns a two-element tuple; the\n",
    "second element is a `dict` that maps variable names to `slice` objects\n",
    "selecting each variable's elements within the flat array.\n",
    "\n",
    "This is useful when a callback covers several variables and you want to read\n",
    "or write a specific variable by name rather than by manually tracking index\n",
    "offsets.\n",
    "\n",
    "For `create_jacobian_matrix(return_index_map=True)`, the dict keys are\n",
    "``(output_name, input_name)`` tuples and the values are\n",
    "``(row_slice, col_slice)`` tuples that identify the corresponding sub-block\n",
    "of the Jacobian matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "yr3mbz6u65",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.752459Z",
     "iopub.status.busy": "2026-06-25T10:29:10.752239Z",
     "iopub.status.idle": "2026-06-25T10:29:10.761672Z",
     "shell.execute_reply": "2026-06-25T10:29:10.760750Z"
    },
    "papermill": {
     "duration": 0.01684,
     "end_time": "2026-06-25T10:29:10.762237+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.745397+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input  index map: {'x': slice(0, 1, None), 'y': slice(1, 2, None)}\n",
      "output index map: {'f_xy': slice(0, 1, None)}\n",
      "f(3, -4) = [-15.]\n",
      "\n",
      "Jacobian index map:\n",
      "  ('f_xy', 'x'): (slice(0, 1, None), slice(0, 1, None))\n",
      "  ('f_xy', 'y'): (slice(0, 1, None), slice(1, 2, None))\n",
      "df/dx = [[-0.5]]\n",
      "df/dy = [[14.5]]\n"
     ]
    }
   ],
   "source": [
    "# --- create_input_vector and create_output_vector with return_index_map ---\n",
    "# Reuse the 'f' callback from Example 1 (inputs: x, y; output: f_xy).\n",
    "x_rim, x_map = f.create_input_vector(return_index_map=True)\n",
    "y_rim, y_map = f.create_output_vector(return_index_map=True)\n",
    "\n",
    "print('input  index map:', x_map)\n",
    "print('output index map:', y_map)\n",
    "\n",
    "# Use the maps to address variables by name without tracking offsets manually.\n",
    "x_rim[x_map['x']] = 3.0\n",
    "x_rim[x_map['y']] = -4.0\n",
    "y_rim = f(x_rim, y=y_rim)\n",
    "print('f(3, -4) =', y_rim[y_map['f_xy']])\n",
    "\n",
    "# --- create_jacobian_matrix with return_index_map ---\n",
    "# Reuse the 'dfdx' callback from Example 2 (inputs: x, y; output: f_xy).\n",
    "J_rim, J_map = dfdx.create_jacobian_matrix(return_index_map=True)\n",
    "\n",
    "print('\\nJacobian index map:')\n",
    "for key, slices in J_map.items():\n",
    "    print(f'  {key}: {slices}')\n",
    "\n",
    "x2_rim = dfdx.create_input_vector()\n",
    "x2_rim[0], x2_rim[1] = 1.5, 2.5\n",
    "dfdx(x2_rim, J=J_rim)\n",
    "\n",
    "print('df/dx =', J_rim[J_map['f_xy', 'x']])\n",
    "print('df/dy =', J_rim[J_map['f_xy', 'y']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "l7k0hv8g5z",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.775796Z",
     "iopub.status.busy": "2026-06-25T10:29:10.775571Z",
     "iopub.status.idle": "2026-06-25T10:29:10.782253Z",
     "shell.execute_reply": "2026-06-25T10:29:10.781330Z"
    },
    "hide_input": true,
    "papermill": {
     "duration": 0.014023,
     "end_time": "2026-06-25T10:29:10.782816+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.768793+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.0)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assert x_map == {'x': slice(0, 1), 'y': slice(1, 2)}\n",
    "assert y_map == {'f_xy': slice(0, 1)}\n",
    "assert_near_equal(y_rim[y_map['f_xy']], [-15.0])\n",
    "\n",
    "assert J_map == {('f_xy', 'x'): (slice(0, 1), slice(0, 1)),\n",
    "                 ('f_xy', 'y'): (slice(0, 1), slice(1, 2))}\n",
    "assert_near_equal(J_rim[J_map['f_xy', 'x']], [[2*(x2_rim[0]-3.0) + x2_rim[1]]])\n",
    "assert_near_equal(J_rim[J_map['f_xy', 'y']], [[x2_rim[0] + 2*(x2_rim[1]+4.0)]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08aa7ec-024f-43d2-8bec-089a9f6c6402",
   "metadata": {
    "papermill": {
     "duration": 0.006246,
     "end_time": "2026-06-25T10:29:10.795340+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.789094+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 8: `AnalysisDriver`-like functionality\n",
    "\n",
    "OpenMDAO's [`AnalysisDriver`](../building_blocks/drivers/analysis_driver) is an OpenMDAO `Driver` that runs an OpenMDAO `Model` for a range of user-specified inputs.\n",
    "The `Problem` functional interface described here provides a similar capability that might be preferable.\n",
    "We'll recreate the `Paraboloid` example from the [`AnalysisDriver` docs](../building_blocks/drivers/analysis_driver) here with the functional interface.\n",
    "\n",
    "The first step is to create the OpenMDAO `Problem`, just like usual:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d0e4339e-4eb2-4a9d-8649-4a0c023c93ee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.807520Z",
     "iopub.status.busy": "2026-06-25T10:29:10.807328Z",
     "iopub.status.idle": "2026-06-25T10:29:10.813452Z",
     "shell.execute_reply": "2026-06-25T10:29:10.812755Z"
    },
    "papermill": {
     "duration": 0.012619,
     "end_time": "2026-06-25T10:29:10.814040+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.801421+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<openmdao.core.problem.Problem at 0x7ff855634950>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create and setup the OpenMDAO Problem, just like usual.\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "prob6 = om.Problem()\n",
    "prob6.model.add_subsystem('comp', Paraboloid(), promotes=['*'])\n",
    "prob6.setup()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af3e25cd-b28e-48cc-abf2-4e3c3b1a7c56",
   "metadata": {
    "papermill": {
     "duration": 0.004211,
     "end_time": "2026-06-25T10:29:10.822996+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.818785+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "With an `AnalysisDriver` we're usually only interested in the outputs, not derivatives, so we'll use the `f` version of the functional interface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "9e3d3c5b-1cad-4f96-905a-187f4a20ad9b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.832146Z",
     "iopub.status.busy": "2026-06-25T10:29:10.831951Z",
     "iopub.status.idle": "2026-06-25T10:29:10.834632Z",
     "shell.execute_reply": "2026-06-25T10:29:10.834057Z"
    },
    "papermill": {
     "duration": 0.008169,
     "end_time": "2026-06-25T10:29:10.835232+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.827063+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Create the callback function with the inputs and outputs we're interested in.\n",
    "f = prob6.get_callback('f', input_vars=['x', 'y'], output_vars=['f_xy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e8b4c8b-5379-4de7-83b6-02e6e7b94775",
   "metadata": {
    "papermill": {
     "duration": 0.004071,
     "end_time": "2026-06-25T10:29:10.843469+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.839398+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can use the fancy [case generating functionality from the `AnalysisDriver` docs](../building_blocks/drivers/analysis_driver.ipynb#generating-cases) with the functional interface.\n",
    "Here we'll use the `ProductGenerator`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "831075ea-7457-410f-81bc-dd6e99799a5d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.852915Z",
     "iopub.status.busy": "2026-06-25T10:29:10.852567Z",
     "iopub.status.idle": "2026-06-25T10:29:10.855394Z",
     "shell.execute_reply": "2026-06-25T10:29:10.854880Z"
    },
    "papermill": {
     "duration": 0.00839,
     "end_time": "2026-06-25T10:29:10.855943+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.847553+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Can still use the fancy `AnalysisGenerator` thingies:\n",
    "gen = om.ProductGenerator({'x': {'val': [0.0, 0.5, 1.0], 'units': None},\n",
    "                           'y': {'val': [0.0, 0.5, 1.0], 'units': None}})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8182382b-348a-4405-a6d5-76ccd1a5328a",
   "metadata": {
    "papermill": {
     "duration": 0.00404,
     "end_time": "2026-06-25T10:29:10.864219+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.860179+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now we can create an input and output vector from the functional interface, and iterate over the cases like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "41d2b730-ffb4-4f6b-85a3-37359d31490b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.874288Z",
     "iopub.status.busy": "2026-06-25T10:29:10.874106Z",
     "iopub.status.idle": "2026-06-25T10:29:10.881455Z",
     "shell.execute_reply": "2026-06-25T10:29:10.880893Z"
    },
    "papermill": {
     "duration": 0.013633,
     "end_time": "2026-06-25T10:29:10.882187+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.868554+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0.000      0.000     22.000\n",
      " 0.000      0.500     26.250\n",
      " 0.000      1.000     31.000\n",
      " 0.500      0.000     19.250\n",
      " 0.500      0.500     23.750\n",
      " 0.500      1.000     28.750\n",
      " 1.000      0.000     17.000\n",
      " 1.000      0.500     21.750\n",
      " 1.000      1.000     27.000\n"
     ]
    }
   ],
   "source": [
    "# Create the input and output vectors that we can use for each call.\n",
    "X, X_idxmap = f.create_input_vector(return_index_map=True)\n",
    "Y, Y_idxmap = f.create_output_vector(return_index_map=True)\n",
    "\n",
    "# Iterate over the samples.\n",
    "for sample in gen:\n",
    "    # Set inputs.\n",
    "    X[X_idxmap['x']] = sample['x']['val']\n",
    "    X[X_idxmap['y']] = sample['y']['val']\n",
    "    # Call the function, using the pre-allocated output array we created outside the loop.\n",
    "    f(X, y=Y)\n",
    "    # Print the results for this sample.\n",
    "    x = X[X_idxmap['x']][0]\n",
    "    y = X[X_idxmap['y']][0]\n",
    "    f_xy = Y[Y_idxmap['f_xy']][0]\n",
    "    print(f'{x:6.3f}     {y:6.3f}     {f_xy:6.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ttus9yao2n",
   "metadata": {
    "papermill": {
     "duration": 0.004166,
     "end_time": "2026-06-25T10:29:10.890847+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.886681+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Example 9: Optimization with `scipy.optimize`\n",
    "\n",
    "The functional interface can be used with external optimizers such as\n",
    "`scipy.optimize.minimize`.  Using `form='fdfdx'` the callback returns both the\n",
    "objective value and its gradient in a single model evaluation, which is\n",
    "the signature expected by `scipy.optimize.minimize` when `jac=True`.\n",
    "\n",
    "The example below solves the same unconstrained paraboloid minimization as the\n",
    "[Simple Optimization](../../examples/paraboloid) example — finding the minimum of\n",
    "$f(x,y) = (x-3)^2 + xy + (y+4)^2 - 3$ — but drives it with\n",
    "`scipy.optimize.minimize` rather than OpenMDAO's built-in driver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d7wx89k15ie",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:10.900775Z",
     "iopub.status.busy": "2026-06-25T10:29:10.900567Z",
     "iopub.status.idle": "2026-06-25T10:29:11.013995Z",
     "shell.execute_reply": "2026-06-25T10:29:11.013551Z"
    },
    "papermill": {
     "duration": 0.121698,
     "end_time": "2026-06-25T10:29:11.016811+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:10.895113+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x* = 6.666666666666667\n",
      "y* = -7.333333333333334\n",
      "f* = -27.33333333333333\n"
     ]
    }
   ],
   "source": [
    "import openmdao.api as om\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "prob9 = om.Problem()\n",
    "prob9.model.add_subsystem('paraboloid', om.ExecComp('f = (x-3)**2 + x*y + (y+4)**2 - 3'))\n",
    "prob9.setup()\n",
    "\n",
    "# Set initial values.\n",
    "prob9.set_val('paraboloid.x', 3.0)\n",
    "prob9.set_val('paraboloid.y', -4.0)\n",
    "prob9.final_setup()\n",
    "\n",
    "# Create a callback that returns both f and the gradient in one model evaluation.\n",
    "fdfdx_opt = prob9.get_callback('fdfdx',\n",
    "                               input_vars=['paraboloid.x', 'paraboloid.y'],\n",
    "                               output_vars=['paraboloid.f'])\n",
    "\n",
    "# scipy.optimize.minimize with jac=True expects the function to return (f, grad).\n",
    "def objective(x):\n",
    "    y, J = fdfdx_opt(x)\n",
    "    return float(y[0]), J[0, :]\n",
    "\n",
    "x0 = fdfdx_opt.create_input_vector()  # pre-filled with the current problem values\n",
    "bounds = [(-50, 50), (-50, 50)]\n",
    "\n",
    "result = minimize(objective, x0, method='SLSQP', jac=True, bounds=bounds)\n",
    "print('x* =', result.x[0])\n",
    "print('y* =', result.x[1])\n",
    "print('f* =', result.fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "gceazsuezjl",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-25T10:29:11.045272Z",
     "iopub.status.busy": "2026-06-25T10:29:11.044372Z",
     "iopub.status.idle": "2026-06-25T10:29:11.053174Z",
     "shell.execute_reply": "2026-06-25T10:29:11.052147Z"
    },
    "papermill": {
     "duration": 0.02586,
     "end_time": "2026-06-25T10:29:11.053730+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:11.027870+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(1.2195109964002665e-10)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(result.x[0], 6.66666666, tolerance=1.0E-5)\n",
    "assert_near_equal(result.x[1], -7.33333333, tolerance=1.0E-5)\n",
    "assert_near_equal(result.fun, -27.33333333, tolerance=1.0E-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4-0027-0027-0027-000000000027",
   "metadata": {
    "papermill": {
     "duration": 0.015977,
     "end_time": "2026-06-25T10:29:11.080621+00:00",
     "exception": false,
     "start_time": "2026-06-25T10:29:11.064644+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Limitations\n",
    "\n",
    "* **`form='f'` requires explicit variable lists**: unlike the derivative\n",
    "  forms, `form='f'` cannot fall back to driver variables because no\n",
    "  `_TotalJacInfo` object is created.  Both `input_vars` and `output_vars`\n",
    "  must be provided.\n",
    "\n",
    "* **`run_model()` on every call**: each invocation of the callback calls\n",
    "  `Problem.run_model()`.  There is no caching; if you need to evaluate\n",
    "  outputs and derivatives at the same point prefer `form='fdfdx'` over\n",
    "  separate `'f'` and `'dfdx'` calls."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.13"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 5.65084,
   "end_time": "2026-06-25T10:29:11.814419+00:00",
   "environment_variables": {},
   "exception": null,
   "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/experimental/functional_interface.ipynb",
   "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/experimental/functional_interface.ipynb",
   "parameters": {},
   "start_time": "2026-06-25T10:29:06.163579+00:00",
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}