Restarting Linear Solutions for Expensive Linear Solves#
When using iterative linear solvers, it is often desirable to use the converged solution from a previous linear solve as the initial guess for the current one. There is some memory cost associated with this feature, because the solution for each quantity of interest will be saved separately. However, the benefit is reduced computational cost for the subsequent linear solves.
Note
This feature should not be used when using the DirectSolver at the top level of your model. It won’t offer any computational savings in that situation.
To use this feature, provide cache_linear_solution=True
as an argument to add_design_var(), add_objective() , or add_constraint().
If you are using one of the OpenMDAO iterative solvers (ScipyKrylov, PETScKrylov, LinearBlockGS, or LinearBlockJac, then simply adding that argument is enough to activate this feature.
If you have implemented the solve_linear
() method for an ImplicitComponent, then you will need to make sure to use the provided guess solution in your implementation. The cached solution will be put into the solution vector for you to use as an initial guess. Note that you will override those values with the final solution. In fwd
mode, the guess will be in the d_outputs
vector. In rev
mode, the guess will be in the d_residuals
vector.
Below is a toy example problem that illustrates how the restart vectors should work. The restart is passed in via the x0
argument to gmres.
import numpy as np
from scipy.sparse.linalg import gmres
import openmdao.api as om
class QuadraticComp(om.ImplicitComponent):
"""
A Simple Implicit Component representing a Quadratic Equation.
R(a, b, c, x) = ax^2 + bx + c
Solution via Quadratic Formula:
x = (-b + sqrt(b^2 - 4ac)) / 2a
"""
def setup(self):
self.add_input('a', val=1.)
self.add_input('b', val=1.)
self.add_input('c', val=1.)
self.add_output('states', val=[0,0])
def setup_partials(self):
self.declare_partials(of='*', wrt='*')
def apply_nonlinear(self, inputs, outputs, residuals):
a = inputs['a']
b = inputs['b']
c = inputs['c']
x = outputs['states'][0]
y = outputs['states'][1]
residuals['states'][0] = a * x ** 2 + b * x + c
residuals['states'][1] = a * y + b
def solve_nonlinear(self, inputs, outputs):
a = inputs['a']
b = inputs['b']
c = inputs['c']
outputs['states'][0] = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)
outputs['states'][1] = -b/a
def linearize(self, inputs, outputs, partials):
a = inputs['a'][0]
b = inputs['b'][0]
# c = inputs['c'][0] # c is not needed
x = outputs['states'][0]
y = outputs['states'][1]
partials['states', 'a'] = [[x**2],[y]]
partials['states', 'b'] = [[x],[1]]
partials['states', 'c'] = [[1.0],[0]]
partials['states', 'states'] = [[2*a*x+b, 0],[0, a]]
self.state_jac = np.array([[2*a*x+b, 0],[0, a]])
def solve_linear(self, d_outputs, d_residuals, mode):
if mode == 'fwd':
print("incoming initial guess", d_outputs['states'])
d_outputs['states'] = gmres(self.state_jac, d_residuals['states'], x0=d_outputs['states'])[0]
elif mode == 'rev':
d_residuals['states'] = gmres(self.state_jac, d_outputs['states'], x0=d_residuals['states'])[0]
p = om.Problem()
p.driver = om.ScipyOptimizeDriver()
p.driver.options['optimizer'] = 'SLSQP'
p.model.set_input_defaults('b', val=4.)
p.model.add_subsystem('quad', QuadraticComp(), promotes_inputs=['a', 'b', 'c'], promotes_outputs=['states'])
p.model.add_subsystem('obj', om.ExecComp('y = (x[1]-x[0])**2', x=np.ones(2)))
p.model.connect('states', 'obj.x')
p.model.add_design_var('a', upper=4., cache_linear_solution=True)
p.model.add_constraint('states', upper=10)
p.model.add_objective('obj.y')
p.setup(mode='fwd')
p.run_driver()
print(p['a'], p['b'], p['c'])
print(p['states'])
print(p['obj.y'])
incoming initial guess [0. 0.]
incoming initial guess [-0.02072594 4. ]
Optimization terminated successfully (Exit mode 0)
Current function value: 0.2500005362356651
Iterations: 6
Function evaluations: 2
Gradient evaluations: 2
Optimization Complete
-----------------------------------
[4.] [4.] [1.]
[-0.49999946 -1. ]
[0.25000054]
/tmp/ipykernel_20244/3428580991.py:40: 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.)
outputs['states'][0] = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)
/tmp/ipykernel_20244/3428580991.py:41: 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.)
outputs['states'][1] = -b/a
Note
In the example shown above, the model is solving for derivatives in ‘fwd’ mode, specified by p.setup(mode=’fwd’)
. In that case the cache_linear_solution
arg may be passed when adding design variables as shown above with p.model.add_design_var(‘a’, cache_linear_solution=True)
. However, if the model were to solve for derivatives in ‘rev’ mode instead, then the cache_linear_solution
arg would be applied to the objective and/or the constraint variables. For example, p.model.add_constraint(‘states’, upper=10, cache_linear_solution=True)
.