{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e2955a17",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:20.006204Z",
     "iopub.status.busy": "2026-10-02T14:40:20.005895Z",
     "iopub.status.idle": "2026-10-02T14:40:20.011657Z",
     "shell.execute_reply": "2026-10-02T14:40:20.010670Z"
    },
    "papermill": {
     "duration": 0.009389,
     "end_time": "2026-10-02T14:40:20.012287+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:20.002898+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": "99831a74",
   "metadata": {
    "papermill": {
     "duration": 0.001258,
     "end_time": "2026-10-02T14:40:20.015056+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:20.013798+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# A Guide to Using Complex Step to Compute Derivatives\n",
    "\n",
    "The intent of this guide is to summarize in detail how to use complex step to compute derivatives.\n",
    "It is assumed that you're already familiar with OpenMDAO usage in general. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d7bdc9",
   "metadata": {
    "papermill": {
     "duration": 0.001104,
     "end_time": "2026-10-02T14:40:20.017305+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:20.016201+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Setting Up Complex Step\n",
    "\n",
    "When complex step is used somewhere in an OpenMDAO model, OpenMDAO must allocate sufficient memory to\n",
    "contain a complex version of the nonlinear vector, and in many cases also a complex version of the\n",
    "linear vector.  OpenMDAO can figure this out automatically most of the time.  For example, if any \n",
    "component in your model calls either `declare_partials` or `declare_coloring` with `method='cs'`,\n",
    "complex vectors will be allocated automatically.  \n",
    "\n",
    "The main situation where complex vectors are \n",
    "needed but are *not* allocated automatically is when you call either `check_totals` or `check_partials`\n",
    "with `method='cs'` and nothing in your model natively uses complex step.  In that case, you must\n",
    "tell OpenMDAO that complex vectors are required by passing a `force_alloc_complex=True` argument\n",
    "when calling `setup` on your `Problem`.  The `force_alloc_complex` flag will force OpenMDAO to \n",
    "allocate complex nonlinear vectors regardless of what it detects in the model.\n",
    "\n",
    "Note that while `ExecComp` components use complex step to compute derivatives by default, they do not \n",
    "require that the OpenMDAO nonlinear vectors are complex because they perform their own internal\n",
    "complex step operation.  However, if you declare your own partials on an `ExecComp` using \n",
    "`declare_partials` or `declare_coloring` with `method='cs'`, then that component will use the\n",
    "framework level complex step routines and will be treated as any other component with partials\n",
    "declared in that manner."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07be1e22",
   "metadata": {
    "papermill": {
     "duration": 0.00106,
     "end_time": "2026-10-02T14:40:20.019483+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:20.018423+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# How to Tell if a Component is Running Under Complex Step\n",
    "\n",
    "A component can tell when it's running under complex step by checking the value of its `under_complex_step`\n",
    "attribute.  A similar flag, `under_finite_difference` can be used to tell if a component is running\n",
    "under finite difference.  Here's an example of a component that checks its complex step status in \n",
    "its `compute` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "068b51e9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:20.022836Z",
     "iopub.status.busy": "2026-10-02T14:40:20.022619Z",
     "iopub.status.idle": "2026-10-02T14:40:22.841861Z",
     "shell.execute_reply": "2026-10-02T14:40:22.840720Z"
    },
    "papermill": {
     "duration": 2.821831,
     "end_time": "2026-10-02T14:40:22.842479+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:20.020648+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1790952022.706992] [runnervm8df0l:5807 :0]        ib_iface.c:1269 UCX  ERROR mana_0: iface 0x55c86bc44410 failed to create UD QP TX wr:256 sge:6 inl:64 resp:0 RX wr:4096 sge:1 resp:0 failed: Operation not supported\n",
      "[1790952022.707278] [runnervm8df0l:5807 :0]      ucp_worker.c:1412 UCX  ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n",
      "not under complex step\n",
      "under complex step\n",
      "[[2.]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[runnervm8df0l:05807] pml_ucx.c:313  Error: Failed to create UCP worker\n"
     ]
    }
   ],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "class MyCheckComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('a', 1.0)\n",
    "        self.add_output('x', 0.0)\n",
    "        # because we set method='cs' here, OpenMDAO automatically knows to allocate\n",
    "        # complex nonlinear vectors\n",
    "        self.declare_partials(of='*', wrt='*', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        a = inputs['a']\n",
    "        if self.under_complex_step:\n",
    "            print('under complex step')\n",
    "        else:\n",
    "            print('not under complex step')\n",
    "        outputs['x'] = a * 2.\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', MyCheckComp())\n",
    "# don't need to set force_alloc_complex=True here since we call declare_partials with\n",
    "# method='cs' in our model.\n",
    "p.setup()\n",
    "\n",
    "# during run_model, our component's compute will *not* be running under complex step\n",
    "p.run_model()\n",
    "\n",
    "# during compute_partials, our component's compute *will* be running under complex step\n",
    "J = p.compute_totals(of=['comp.x'], wrt=['comp.a'])\n",
    "print(J['comp.x', 'comp.a'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4cc5eb5",
   "metadata": {
    "papermill": {
     "duration": 0.039988,
     "end_time": "2026-10-02T14:40:22.926860+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:22.886872+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Complex Stepping through Solvers\n",
    "\n",
    "See [Complex Step Guidelines](complex-step-guidelines) for important issues to consider when your model has nonlinear solvers \n",
    "under a group that is performing complex step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a73c351",
   "metadata": {
    "papermill": {
     "duration": 0.035456,
     "end_time": "2026-10-02T14:40:22.983895+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:22.948439+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Using Complex Step to Compute a Partial Jacobian Coloring\n",
    "\n",
    "OpenMDAO can compute a coloring of the partial jacobian matrix for a component that uses complex step\n",
    "or finite difference to compute its derivatives.  For a detailed explanation of this, see \n",
    "[Simultaneous Coloring of Approximated Derivatives](../features/experimental/approx_coloring.ipynb).\n",
    "Assuming your component is complex safe, generally using complex step is more accurate than finite difference\n",
    "and should be preferred when computing a coloring of the partial jacobian."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca44074",
   "metadata": {
    "papermill": {
     "duration": 0.041912,
     "end_time": "2026-10-02T14:40:23.072661+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:23.030749+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# The cs_safe Module\n",
    "\n",
    "The `openmdao.utils.cs_safe` module contains complex-safe versions of a few common functions, namely,\n",
    "`abs`, `norm`, and `arctan2`.  The `numpy` versions of these functions are not complex-safe and so\n",
    "must be replaced with the safe versions if you intend to use such functions in your component under\n",
    "complex step.  Note that the `ExecComp` component, which uses complex step by default, automatically\n",
    "uses the complex safe versions of these functions if they are referenced in one of its expressions.\n",
    "\n",
    "The following example shows how to make a component that uses `abs` safe to use under complex step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ef367198",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:23.136035Z",
     "iopub.status.busy": "2026-10-02T14:40:23.135636Z",
     "iopub.status.idle": "2026-10-02T14:40:23.152665Z",
     "shell.execute_reply": "2026-10-02T14:40:23.151880Z"
    },
    "papermill": {
     "duration": 0.042218,
     "end_time": "2026-10-02T14:40:23.153394+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:23.111176+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[4.]]\n",
      "[[-2.]]\n"
     ]
    }
   ],
   "source": [
    "import openmdao.api as om\n",
    "from openmdao.utils.cs_safe import abs as cs_abs\n",
    "\n",
    "class MyComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('a', 1.0)\n",
    "        self.add_input('b', -2.0)\n",
    "        self.add_output('x', 0.0)\n",
    "        # because we set method='cs' here, OpenMDAO automatically knows to allocate\n",
    "        # complex nonlinear vectors\n",
    "        self.declare_partials(of='*', wrt='*', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        a, b = inputs.values()\n",
    "        # normal abs isn't complex safe, so use cs_abs\n",
    "        outputs['x'] = a * cs_abs(b) * 2.\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', MyComp())\n",
    "# don't need to set force_alloc_complex=True here since we call declare_partials with\n",
    "# method='cs' in our model.\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "J = p.compute_totals(of=['comp.x'], wrt=['comp.a', 'comp.b'])\n",
    "print(J['comp.x', 'comp.a'])\n",
    "print(J['comp.x', 'comp.b'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14799aeb",
   "metadata": {
    "papermill": {
     "duration": 0.079217,
     "end_time": "2026-10-02T14:40:23.342907+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:23.263690+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model Level Complex Step Mode for Debugging\n",
    "\n",
    "Sometimes when debugging it may be useful to run a model or part of a model using complex inputs. This\n",
    "can be done by calling `set_complex_step_mode(True)` on the `Problem` instance. You can then call\n",
    "\n",
    "```python\n",
    "prob['some_var'] = a_complex_val\n",
    "```\n",
    "\n",
    "or\n",
    "\n",
    "```python\n",
    "prob.set_val('some_var', a_complex_val)\n",
    "```\n",
    "\n",
    "then run the model by calling\n",
    "\n",
    "```python\n",
    "prob.run_model()\n",
    "```\n",
    "\n",
    "The complex values will carry through the model as it runs.  Note that this only works if all of\n",
    "the components in the model are complex-safe.  After the model run has completed, the outputs\n",
    "can then be inspected using one of the following:\n",
    "\n",
    "```python\n",
    "x = prob['some_output']\n",
    "```\n",
    "\n",
    "or\n",
    "\n",
    "```python\n",
    "x = prob.get_val('some_output')\n",
    "```\n",
    "\n",
    "or \n",
    "\n",
    "```python\n",
    "prob.model.list_outputs()\n",
    "```\n",
    "\n",
    "Here's a short example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9a108fbd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:23.551964Z",
     "iopub.status.busy": "2026-10-02T14:40:23.551726Z",
     "iopub.status.idle": "2026-10-02T14:40:23.557990Z",
     "shell.execute_reply": "2026-10-02T14:40:23.557324Z"
    },
    "papermill": {
     "duration": 0.11164,
     "end_time": "2026-10-02T14:40:23.558436+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:23.446796+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float [9.]\n",
      "complex [-7.48446667+12.99933685j]\n"
     ]
    }
   ],
   "source": [
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', MyComp())\n",
    "p.setup()\n",
    "\n",
    "p['comp.a'] = 1.5\n",
    "p['comp.b'] = -3.\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "# output x should be a float here\n",
    "print('float', p['comp.x'])\n",
    "\n",
    "# now we're setting the problem to use complex step\n",
    "p.set_complex_step_mode(True)\n",
    "\n",
    "p['comp.a'] = 1.5+2j\n",
    "p['comp.b'] = -3.-7j\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "# output x should be complex here\n",
    "print('complex', p['comp.x'])"
   ]
  }
 ],
 "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.14"
  },
  "orphan": true,
  "papermill": {
   "default_parameters": {},
   "duration": 4.751136,
   "end_time": "2026-10-02T14:40:24.029483+00:00",
   "environment_variables": {},
   "exception": null,
   "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/advanced_user_guide/complex_step.ipynb",
   "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/advanced_user_guide/complex_step.ipynb",
   "parameters": {},
   "start_time": "2026-10-02T14:40:19.278347+00:00",
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}