{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "26c79ba8", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:42:39.022344Z", "iopub.status.busy": "2026-10-02T14:42:39.022099Z", "iopub.status.idle": "2026-10-02T14:42:39.026376Z", "shell.execute_reply": "2026-10-02T14:42:39.025858Z" }, "papermill": { "duration": 0.006869, "end_time": "2026-10-02T14:42:39.026968+00:00", "exception": false, "start_time": "2026-10-02T14:42:39.020099+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": "0e9c223d", "metadata": { "papermill": { "duration": 0.000914, "end_time": "2026-10-02T14:42:39.033029+00:00", "exception": false, "start_time": "2026-10-02T14:42:39.032115+00:00", "status": "completed" }, "tags": [] }, "source": [ "# LinearUserDefined\n", "\n", "LinearUserDefined is a solver that lets you define a custom method for performing a linear solve on a component. The default method is named \"solve_linear\", but you can give it any name by passing in the function or method handle to the \"solve_function\" attribute.\n", "\n", "The function needs to have the following signature:\n", "\n", "```python\n", " def my_solve_function(d_outputs, d_residuals, mode):\n", " r\"\"\"\n", " Apply inverse jac product. The model is assumed to be in an unscaled state.\n", "\n", " Parameters\n", " ----------\n", " d_outputs: Vector\n", " unscaled, dimensional quantities read via d_outputs[key]\n", " d_residuals: Vector\n", " unscaled, dimensional quantities read via d_residuals[key]\n", " mode: str\n", " either 'fwd' or 'rev'\n", "\n", " Returns\n", " -------\n", " None or bool or (bool, float, float)\n", " The bool is the failure flag; and the two floats are absolute and relative error.\n", " \"\"\"\n", "```\n", "\n", "Here is a rather contrived example where an identity preconditioner is used by giving the component's \"mysolve\"\n", "method to a LinearUserDefined solver." ] }, { "cell_type": "code", "execution_count": 2, "id": "aeb5c163", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:42:39.035303Z", "iopub.status.busy": "2026-10-02T14:42:39.035162Z", "iopub.status.idle": "2026-10-02T14:42:41.654980Z", "shell.execute_reply": "2026-10-02T14:42:41.654340Z" }, "papermill": { "duration": 2.622017, "end_time": "2026-10-02T14:42:41.655830+00:00", "exception": false, "start_time": "2026-10-02T14:42:39.033813+00:00", "status": "completed" }, "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[1790952161.590947] [runnervm8df0l:7232 :0] ib_iface.c:1269 UCX ERROR mana_0: iface 0x56499c7115e0 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", "[1790952161.591216] [runnervm8df0l:7232 :0] ucp_worker.c:1412 UCX ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n", "15.000000000000002\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[runnervm8df0l:07232] pml_ucx.c:313 Error: Failed to create UCP worker\n" ] } ], "source": [ "import numpy as np\n", "import openmdao.api as om\n", "\n", "from openmdao.utils.array_utils import evenly_distrib_idxs\n", "from openmdao.utils.mpi import MPI\n", "\n", "\n", "class CustomSolveImplicit(om.ImplicitComponent):\n", "\n", " def setup(self):\n", "\n", " self.add_input('a', val=10., units='m')\n", "\n", " rank = self.comm.rank\n", " GLOBAL_SIZE = 15\n", " sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)\n", "\n", " self.add_output('states', shape=int(sizes[rank]))\n", "\n", " self.add_output('out_var', shape=1)\n", " self.local_size = sizes[rank]\n", "\n", " self.linear_solver = om.PETScKrylov()\n", " self.linear_solver.precon = om.LinearUserDefined(solve_function=self.mysolve)\n", "\n", " def solve_nonlinear(self, i, o):\n", " o['states'] = i['a']\n", "\n", " local_sum = np.zeros(1)\n", " local_sum[0] = np.sum(o['states'])\n", " tmp = np.zeros(1)\n", "\n", " o['out_var'] = tmp[0]\n", "\n", " def apply_nonlinear(self, i, o, r):\n", " r['states'] = o['states'] - i['a']\n", "\n", " local_sum = np.zeros(1)\n", " local_sum[0] = np.sum(o['states'])\n", " tmp = np.zeros(1)\n", "\n", " r['out_var'] = o['out_var'] - tmp[0]\n", "\n", " def apply_linear(self, i, o, d_i, d_o, d_r, mode):\n", " if mode == 'fwd':\n", " if 'states' in d_o:\n", " d_r['states'] += d_o['states']\n", "\n", " local_sum = np.array([np.sum(d_o['states'])])\n", " global_sum = np.zeros(1)\n", " self.comm.Allreduce(local_sum, global_sum, op=MPI.SUM)\n", " d_r['out_var'] -= global_sum\n", "\n", " if 'out_var' in d_o:\n", " d_r['out_var'] += d_o['out_var']\n", "\n", " if 'a' in d_i:\n", " d_r['states'] -= d_i['a']\n", "\n", " elif mode == 'rev':\n", " if 'states' in d_o:\n", " d_o['states'] += d_r['states']\n", "\n", " tmp = np.zeros(1)\n", " if self.comm.rank == 0:\n", " tmp[:] = d_r['out_var']\n", " self.comm.Bcast(tmp, root=0)\n", "\n", " d_o['states'] -= tmp\n", "\n", " if 'out_var' in d_o:\n", " d_o['out_var'] += d_r['out_var']\n", "\n", " if 'a' in d_i:\n", " d_i['a'] -= np.sum(d_r['states'])\n", "\n", " def mysolve(self, d_outputs, d_residuals, mode):\n", " r\"\"\"\n", " Apply inverse jac product. The model is assumed to be in an unscaled state.\n", "\n", " If mode is:\n", " 'fwd': d_residuals \\|-> d_outputs\n", "\n", " 'rev': d_outputs \\|-> d_residuals\n", "\n", " Parameters\n", " ----------\n", " d_outputs : Vector\n", " unscaled, dimensional quantities read via d_outputs[key]\n", " d_residuals : Vector\n", " unscaled, dimensional quantities read via d_residuals[key]\n", " mode: str\n", " either 'fwd' or 'rev'\n", " \"\"\"\n", " # Note: we are just preconditioning with Identity as a proof of concept.\n", " if mode == 'fwd':\n", " d_outputs.set_vec(d_residuals)\n", " elif mode == 'rev':\n", " d_residuals.set_vec(d_outputs)\n", "\n", "prob = om.Problem()\n", "\n", "\n", "prob.model.add_subsystem('icomp', CustomSolveImplicit(), promotes=['*'])\n", "prob.model.set_input_defaults('a', 10., units='m')\n", "\n", "model = prob.model\n", "\n", "model.linear_solver = om.PETScKrylov()\n", "model.linear_solver.precon = om.LinearRunOnce()\n", "\n", "prob.setup(mode='rev', check=False)\n", "prob.run_model()\n", "jac = prob.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')\n", "\n", "print(jac['out_var']['a'][0][0])" ] }, { "cell_type": "code", "execution_count": 3, "id": "a905ac54", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:42:41.659720Z", "iopub.status.busy": "2026-10-02T14:42:41.659223Z", "iopub.status.idle": "2026-10-02T14:42:41.666423Z", "shell.execute_reply": "2026-10-02T14:42:41.665585Z" }, "papermill": { "duration": 0.009788, "end_time": "2026-10-02T14:42:41.666972+00:00", "exception": false, "start_time": "2026-10-02T14:42:41.657184+00:00", "status": "completed" }, "tags": [ "remove-input", "remove-output" ] }, "outputs": [ { "data": { "text/plain": [ "np.float64(1.1842378929335e-16)" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from openmdao.utils.assert_utils import assert_near_equal\n", "\n", "assert_near_equal(15.0, jac['out_var']['a'][0][0])" ] }, { "cell_type": "markdown", "id": "4e7837cb", "metadata": { "papermill": { "duration": 0.001069, "end_time": "2026-10-02T14:42:41.669258+00:00", "exception": false, "start_time": "2026-10-02T14:42:41.668189+00:00", "status": "completed" }, "tags": [] }, "source": [ "## LinearUserDefined Options" ] }, { "cell_type": "code", "execution_count": 4, "id": "55a9c211", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:42:41.672587Z", "iopub.status.busy": "2026-10-02T14:42:41.672399Z", "iopub.status.idle": "2026-10-02T14:42:41.677318Z", "shell.execute_reply": "2026-10-02T14:42:41.676718Z" }, "papermill": { "duration": 0.007742, "end_time": "2026-10-02T14:42:41.678170+00:00", "exception": false, "start_time": "2026-10-02T14:42:41.670428+00:00", "status": "completed" }, "tags": [ "remove-input" ] }, "outputs": [ { "data": { "text/html": [ "\n", "\n", "\n", "\n", " \n", "\n", "\n", "

\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
OptionDefaultAcceptable ValuesAcceptable TypesDescription
assemble_jacFalse[True, False]['bool']Activates use of assembled jacobian by this solver.
atol1e-10N/AN/Aabsolute error tolerance
err_on_non_convergeFalse[True, False]['bool']When True, AnalysisError will be raised if we don't converge.
iprint1N/A['int']whether to print output
maxiter10N/A['int']maximum number of iterations
rtol1e-10N/AN/Arelative error tolerance
\n", "\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "om.show_options_table(\"openmdao.solvers.linear.user_defined.LinearUserDefined\")" ] }, { "cell_type": "markdown", "id": "22d4b7e0", "metadata": { "papermill": { "duration": 0.001185, "end_time": "2026-10-02T14:42:41.680562+00:00", "exception": false, "start_time": "2026-10-02T14:42:41.679377+00:00", "status": "completed" }, "tags": [] }, "source": [ "## LinearUserDefined Constructor\n", "\n", "The call signature for the `LinearUserDefined` constructor is:\n", "\n", "```{eval-rst}\n", " .. automethod:: openmdao.solvers.linear.user_defined.LinearUserDefined.__init__\n", " :noindex:\n", "```" ] } ], "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" }, "orphan": true, "papermill": { "default_parameters": {}, "duration": 4.134893, "end_time": "2026-10-02T14:42:42.498071+00:00", "environment_variables": {}, "exception": null, "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/building_blocks/solvers/linear_user_defined.ipynb", "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/building_blocks/solvers/linear_user_defined.ipynb", "parameters": {}, "start_time": "2026-10-02T14:42:38.363178+00:00", "version": "2.7.0" } }, "nbformat": 4, "nbformat_minor": 5 }