{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3d699335",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:59.630668Z",
     "iopub.status.busy": "2026-10-02T14:40:59.630458Z",
     "iopub.status.idle": "2026-10-02T14:40:59.635039Z",
     "shell.execute_reply": "2026-10-02T14:40:59.634525Z"
    },
    "hide_input": false,
    "papermill": {
     "duration": 0.007114,
     "end_time": "2026-10-02T14:40:59.635611+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:59.628497+00:00",
     "status": "completed"
    },
    "tags": [
     "active-ipynb",
     "remove-input",
     "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": "4b442fa4",
   "metadata": {
    "papermill": {
     "duration": 0.062532,
     "end_time": "2026-10-02T14:40:59.713787+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:59.651255+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Kepler’s Equation\n",
    "\n",
    "This example will demonstrate the use of OpenMDAO for solving an implicit equation commonly found in astrodynamics, Kepler’s Equation:\n",
    "\n",
    "\\begin{align}\n",
    "E - e \\sin{E} = M\n",
    "\\end{align}\n",
    "\n",
    "Here $M$ is the mean anomaly, $E$ is the eccentric anomaly, and $e$ is the eccentricity of the orbit.\n",
    "\n",
    "If we know the eccentric anomaly, computing the mean anomaly is trivial. However, solving for the eccentric anomaly when given the mean anomaly must be done numerically. We’ll do so using a nonlinear solver. In OpenMDAO, solvers converge all implicit state variables in a Group by driving their residuals to zero."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e45b74d",
   "metadata": {
    "papermill": {
     "duration": 0.001221,
     "end_time": "2026-10-02T14:40:59.716630+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:59.715409+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "\n",
    "## Using a BalanceComp and NewtonSolver\n",
    "\n",
    "In an effort to simplify things for users, OpenMDAO features a Balance component. For each implicit state variable we assign to the balance, it solves the following equation:\n",
    "\n",
    "\\begin{align}\n",
    "lhs(var) \\cdot mult(var) = rhs(var)\n",
    "\\end{align}\n",
    "\n",
    "The _mult_ term is an optional multiplier than can be applied to the left-hand side (LHS) of the equation. For our example, we will assign the right-hand side (RHS) to the mean anomaly ($M$), and the left-hand side to $E - e \\sin{E}$.\n",
    "\n",
    "In this implementation, we rely on an ExecComp to compute the value of the LHS.\n",
    "\n",
    "BalanceComp also provides a way to supply the starting value for the implicit state variable ($E$ in this case), via the guess_func argument. The supplied function should have a similar signature to the guess_nonlinear function of ImplicitComponent. When solving Kepler’s equation, using $M$ as the initial guess for $E$ is a good starting point.\n",
    "\n",
    "In summary, the recipe for solving Kepler’s equation with a NewtonSolver is as follows:\n",
    "\n",
    "1. Define a Group to contain the implicit system.\n",
    "2. To that Group, add components which provide, $M$, $e$, and the left-hand side of Kepler’s equation.\n",
    "3. Add a linear and nonlinear solver to the Group, since the default solvers do not iterate.\n",
    "4. Setup the problem, set values for the inputs, and run the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "49063098",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:40:59.724251Z",
     "iopub.status.busy": "2026-10-02T14:40:59.724019Z",
     "iopub.status.idle": "2026-10-02T14:41:02.431727Z",
     "shell.execute_reply": "2026-10-02T14:41:02.430859Z"
    },
    "papermill": {
     "duration": 2.71472,
     "end_time": "2026-10-02T14:41:02.432722+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:40:59.718002+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1790952062.410867] [runnervm8df0l:6219 :0]        ib_iface.c:1269 UCX  ERROR mana_0: iface 0x5623102180b0 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",
      "[1790952062.411136] [runnervm8df0l:6219 :0]      ucp_worker.c:1412 UCX  ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n",
      "NL: Newton 0 ; 0.3321557 1\n",
      "NL: Newton 1 ; 0.117134648 0.352649822\n",
      "NL: Newton 2 ; 0.00208780933 0.00628563452\n",
      "NL: Newton 3 ; 7.36738647e-07 2.2180521e-06\n",
      "NL: Newton 4 ; 9.19264664e-14 2.76757155e-13\n",
      "NL: Newton Converged\n",
      "M = [85.]\n",
      "E = [2.02317564]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[runnervm8df0l:06219] pml_ucx.c:313  Error: Failed to create UCP worker\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "bal = om.BalanceComp()\n",
    "\n",
    "bal.add_balance(name='E', val=0.0, units='rad', eq_units='rad', rhs_name='M')\n",
    "\n",
    "# Use M (mean anomaly) as the initial guess for E (eccentric anomaly)\n",
    "def guess_function(inputs, outputs, residuals):\n",
    "    if np.abs(residuals['E']) > 1.0E-2:\n",
    "        outputs['E'] = inputs['M']\n",
    "\n",
    "bal.options['guess_func'] = guess_function\n",
    "\n",
    "# ExecComp used to compute the LHS of Kepler's equation.\n",
    "lhs_comp = om.ExecComp('lhs=E - ecc * sin(E)',\n",
    "                       lhs={'units': 'rad'},\n",
    "                       E={'units': 'rad'},\n",
    "                       ecc={'units': None})\n",
    "\n",
    "prob.model.set_input_defaults('M', 85.0, units='deg')\n",
    "\n",
    "prob.model.add_subsystem(name='lhs_comp', subsys=lhs_comp,\n",
    "                         promotes_inputs=['E', 'ecc'])\n",
    "\n",
    "prob.model.add_subsystem(name='balance', subsys=bal,\n",
    "                         promotes_inputs=['M'],\n",
    "                         promotes_outputs=['E'])\n",
    "\n",
    "# Explicit connections\n",
    "prob.model.connect('lhs_comp.lhs', 'balance.lhs:E')\n",
    "\n",
    "# Set up solvers\n",
    "prob.model.linear_solver = om.DirectSolver()\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=2)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('M', 85.0, units='deg')\n",
    "prob.set_val('E', 85.0, units='deg')\n",
    "prob.set_val('ecc', 0.6)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "M = prob.get_val('M')\n",
    "E = prob.get_val('E')\n",
    "print(f'M = {M}')\n",
    "print(f'E = {E}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1def23e4",
   "metadata": {
    "papermill": {
     "duration": 0.003068,
     "end_time": "2026-10-02T14:41:02.440459+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.437391+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "```{note}\n",
    "BalanceComp should be placed after those components which compute its inputs.  Otherwise errors can arise due to timing issues.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88370de0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:02.445010Z",
     "iopub.status.busy": "2026-10-02T14:41:02.444623Z",
     "iopub.status.idle": "2026-10-02T14:41:02.451608Z",
     "shell.execute_reply": "2026-10-02T14:41:02.451012Z"
    },
    "papermill": {
     "duration": 0.010021,
     "end_time": "2026-10-02T14:41:02.452201+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.442180+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(2.1087335717807284e-09)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check that results are correct.\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('E'), 2.02317564, tolerance=1.0E-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a7c624c",
   "metadata": {
    "papermill": {
     "duration": 0.001812,
     "end_time": "2026-10-02T14:41:02.455575+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.453763+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Using Fixed-Point Iteration with an ImplicitComponent\n",
    "\n",
    "One approach to solving Kepler's Equation is to use fixed-point iteration:\n",
    "\n",
    "\\begin{align}\n",
    "E_{n+1} = M + e \\sin{E_{n}}\n",
    "\\end{align}\n",
    "\n",
    "We can achieve this behavior by defining a simple ImplicitComponent that \"solves itself\" by overriding its `solve_nonlinear` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a7392825",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:02.521367Z",
     "iopub.status.busy": "2026-10-02T14:41:02.521166Z",
     "iopub.status.idle": "2026-10-02T14:41:02.526244Z",
     "shell.execute_reply": "2026-10-02T14:41:02.525487Z"
    },
    "papermill": {
     "duration": 0.069665,
     "end_time": "2026-10-02T14:41:02.526748+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.457083+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class KeplerComp(om.ImplicitComponent):\n",
    "    \n",
    "    def initialize(self):\n",
    "        self.options.declare('tol', types=float, default=1.0E-12, desc='convergence tolerance')\n",
    "        self.options.declare('max_its', types=int, default=50, desc='maximum number of iterations')\n",
    "    \n",
    "    def setup(self):\n",
    "        self.add_input('M', val=1.0, units='rad', desc='Mean anomaly')\n",
    "        self.add_input('ecc', val=0.0, units=None, desc='eccentricity')\n",
    "        self.add_output('E', val=1.0, units='rad', desc='Eccentric anomaly')\n",
    "        \n",
    "        self.declare_partials(of='E', wrt='*', method='cs')\n",
    "        \n",
    "    def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "        \"\"\" Compute the residual of output E. \"\"\"\n",
    "        M = inputs['M']\n",
    "        e = inputs['ecc']\n",
    "        E = outputs['E']\n",
    "        residuals['E'] = E - M - e * np.sin(E)\n",
    "        \n",
    "    def solve_nonlinear(self, inputs, outputs):\n",
    "        \"\"\" Determine the values of the implicit outputs that eliminate the residual. \"\"\"\n",
    "        M = inputs['M']\n",
    "        e = inputs['ecc']\n",
    "        E = outputs['E']\n",
    "        \n",
    "        # Compute the initial guess for iteration if the initial redisual is large.\n",
    "        E = inputs['M'] if np.abs(E - M - e * np.sin(E)) > 1.0E-2 else outputs['E']\n",
    "            \n",
    "        for i in range(self.options['max_its']):\n",
    "            E_old = E\n",
    "            E = M + e * np.sin(E)\n",
    "            if i == 0:\n",
    "                print('iter ' + '   ' + 'max abs. error')\n",
    "                print(5*'-' + '   '+ 14*'-')\n",
    "            print(f'{i:>5} {np.max(np.abs(E - E_old)):16.12f}')\n",
    "            if np.abs(E - E_old) < self.options['tol']:\n",
    "                outputs['E'] = E\n",
    "                break\n",
    "        else:\n",
    "            raise om.AnalysisError('Iteration Limit Exceeded')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50704280",
   "metadata": {
    "papermill": {
     "duration": 0.001186,
     "end_time": "2026-10-02T14:41:02.529321+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.528135+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Here the KeplerComp uses the `solve_nonlinear` method to drive it residual to zero through fixed point iteration.\n",
    "\n",
    "```{important}\n",
    "In OpenMDAO it doesn't matter how the residuals of a system of eliminated.  Inside solve_nonlinear you're free to use a variety of techniques.  As long as the resulting residual is _reasonably close to zero_, the derivatives across the implicit system will be accurate.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d675e33",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:02.532653Z",
     "iopub.status.busy": "2026-10-02T14:41:02.532475Z",
     "iopub.status.idle": "2026-10-02T14:41:02.539931Z",
     "shell.execute_reply": "2026-10-02T14:41:02.539181Z"
    },
    "papermill": {
     "duration": 0.009863,
     "end_time": "2026-10-02T14:41:02.540327+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.530464+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter    max abs. error\n",
      "-----   --------------\n",
      "    0   0.597716818855\n",
      "    1   0.074202079572\n",
      "    2   0.020291240188\n",
      "    3   0.005255936885\n",
      "    4   0.001382786018\n",
      "    5   0.000362353275\n",
      "    6   0.000095052939\n",
      "    7   0.000024927544\n",
      "    8   0.000006537697\n",
      "    9   0.000001714596\n",
      "   10   0.000000449677\n",
      "   11   0.000000117934\n",
      "   12   0.000000030930\n",
      "   13   0.000000008112\n",
      "   14   0.000000002127\n",
      "   15   0.000000000558\n",
      "   16   0.000000000146\n",
      "   17   0.000000000038\n",
      "   18   0.000000000010\n",
      "   19   0.000000000003\n",
      "   20   0.000000000001\n",
      "M = [1.48352986]\n",
      "E = [2.02317564]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem(name='kep_comp', subsys=KeplerComp(),\n",
    "                         promotes_inputs=['M', 'ecc'],\n",
    "                         promotes_outputs=['E'])\n",
    "\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('M', 85.0, units='deg')\n",
    "prob.set_val('E', 1.0, units='deg')\n",
    "prob.set_val('ecc', 0.6)\n",
    "\n",
    "prob.run_model()\n",
    "M = prob.get_val('M')\n",
    "E = prob.get_val('E')\n",
    "print(f'M = {M}')\n",
    "print(f'E = {E}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2e61ba60",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:02.568835Z",
     "iopub.status.busy": "2026-10-02T14:41:02.568615Z",
     "iopub.status.idle": "2026-10-02T14:41:02.572465Z",
     "shell.execute_reply": "2026-10-02T14:41:02.571846Z"
    },
    "papermill": {
     "duration": 0.031222,
     "end_time": "2026-10-02T14:41:02.572928+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.541706+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(2.1087686919513014e-09)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check that results are correct.\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('E'), 2.02317564, tolerance=1.0E-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1747605",
   "metadata": {
    "papermill": {
     "duration": 0.001273,
     "end_time": "2026-10-02T14:41:02.575346+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.574073+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In contrast to the NewtonSolver approach, this approach roughly needs the following steps:\n",
    "\n",
    "1. Define an ImplicitComponent where the `apply_nonlinear` method computes the residuals for Kepler's Equation and `solve_nonlinear` find the value of the implicit output that eliminates the residuals.\n",
    "2. Add the ImplicitComponent to the problem model.\n",
    "3. Setup the problem, set values for the inputs, and run the model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972d7fea",
   "metadata": {
    "papermill": {
     "duration": 0.001438,
     "end_time": "2026-10-02T14:41:02.578196+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:02.576758+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Which approach should you use?\n",
    "\n",
    "It's really up to you.  In this case, the Newton approach involves slightly more complicated `System` consisting of two components (the `ExecComp` and the `BalanceComp`) but it converges in just 4 iteratoins.\n",
    "\n",
    "The ImplicitComponent involves only a single system, but writing the `solve_nonlinear` method with some error handling, initial guess generation, and print capability made for more coding on the part of the user."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 4.421135,
   "end_time": "2026-10-02T14:41:03.397425+00:00",
   "environment_variables": {},
   "exception": null,
   "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/examples/keplers_equation.ipynb",
   "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/examples/keplers_equation.ipynb",
   "parameters": {},
   "start_time": "2026-10-02T14:40:58.976290+00:00",
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}