LinearUserDefined

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.

The function needs to have the following signature:

    def my_solve_function(d_outputs, d_residuals, mode):
        r"""
        Apply inverse jac product. The model is assumed to be in an unscaled state.

        Parameters
        ----------
        d_outputs: Vector
            unscaled, dimensional quantities read via d_outputs[key]
        d_residuals: Vector
            unscaled, dimensional quantities read via d_residuals[key]
        mode: str
            either 'fwd' or 'rev'

        Returns
        -------
        None or bool or (bool, float, float)
            The bool is the failure flag; and the two floats are absolute and relative error.
        """

Here is a rather contrived example where an identity preconditioner is used by giving the component’s “mysolve” method to a LinearUserDefined solver.

import numpy as np
import openmdao.api as om

from openmdao.utils.array_utils import evenly_distrib_idxs
from openmdao.utils.mpi import MPI


class CustomSolveImplicit(om.ImplicitComponent):

    def setup(self):

        self.add_input('a', val=10., units='m')

        rank = self.comm.rank
        GLOBAL_SIZE = 15
        sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)

        self.add_output('states', shape=int(sizes[rank]))

        self.add_output('out_var', shape=1)
        self.local_size = sizes[rank]

        self.linear_solver = om.PETScKrylov()
        self.linear_solver.precon = om.LinearUserDefined(solve_function=self.mysolve)

    def solve_nonlinear(self, i, o):
        o['states'] = i['a']

        local_sum = np.zeros(1)
        local_sum[0] = np.sum(o['states'])
        tmp = np.zeros(1)

        o['out_var'] = tmp[0]

    def apply_nonlinear(self, i, o, r):
        r['states'] = o['states'] - i['a']

        local_sum = np.zeros(1)
        local_sum[0] = np.sum(o['states'])
        tmp = np.zeros(1)

        r['out_var'] = o['out_var'] - tmp[0]

    def apply_linear(self, i, o, d_i, d_o, d_r, mode):
        if mode == 'fwd':
            if 'states' in d_o:
                d_r['states'] += d_o['states']

                local_sum = np.array([np.sum(d_o['states'])])
                global_sum = np.zeros(1)
                self.comm.Allreduce(local_sum, global_sum, op=MPI.SUM)
                d_r['out_var'] -= global_sum

            if 'out_var' in d_o:
                    d_r['out_var'] += d_o['out_var']

            if 'a' in d_i:
                    d_r['states'] -= d_i['a']

        elif mode == 'rev':
            if 'states' in d_o:
                d_o['states'] += d_r['states']

                tmp = np.zeros(1)
                if self.comm.rank == 0:
                    tmp[0] = d_r['out_var'].copy()
                self.comm.Bcast(tmp, root=0)

                d_o['states'] -= tmp

            if 'out_var' in d_o:
                d_o['out_var'] += d_r['out_var']

            if 'a' in d_i:
                    d_i['a'] -= np.sum(d_r['states'])

    def mysolve(self, d_outputs, d_residuals, mode):
        r"""
        Apply inverse jac product. The model is assumed to be in an unscaled state.

        If mode is:
            'fwd': d_residuals \|-> d_outputs

            'rev': d_outputs \|-> d_residuals

        Parameters
        ----------
        d_outputs : Vector
            unscaled, dimensional quantities read via d_outputs[key]
        d_residuals : Vector
            unscaled, dimensional quantities read via d_residuals[key]
        mode: str
            either 'fwd' or 'rev'
        """
        # Note: we are just preconditioning with Identity as a proof of concept.
        if mode == 'fwd':
            d_outputs.set_vec(d_residuals)
        elif mode == 'rev':
            d_residuals.set_vec(d_outputs)

prob = om.Problem()


prob.model.add_subsystem('icomp', CustomSolveImplicit(), promotes=['*'])
prob.model.set_input_defaults('a', 10., units='m')

model = prob.model

model.linear_solver = om.PETScKrylov()
model.linear_solver.precon = om.LinearRunOnce()

prob.setup(mode='rev', check=False)
prob.run_model()
jac = prob.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')

print(jac['out_var']['a'][0][0])

15.000000000000002

/tmp/ipykernel_23935/3256407146.py:66: DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)
  tmp[0] = d_r['out_var'].copy()

LinearUserDefined Options

Option	Default	Acceptable Values	Acceptable Types	Description
assemble_jac	False	[True, False]	['bool']	Activates use of assembled jacobian by this solver.
atol	1e-10	N/A	N/A	absolute error tolerance
err_on_non_converge	False	[True, False]	['bool']	When True, AnalysisError will be raised if we don't converge.
iprint	1	N/A	['int']	whether to print output
maxiter	10	N/A	['int']	maximum number of iterations
rtol	1e-10	N/A	N/A	relative error tolerance

LinearUserDefined Constructor

The call signature for the LinearUserDefined constructor is:

LinearUserDefined.__init__(solve_function=None, **kwargs)

Initialize all attributes.