In-Memory Assembly of Jacobians#
When you have groups, or entire models, that are small enough to fit the entire Jacobian into memory, you can have OpenMDAO actually assemble the partial-derivative Jacobian in memory. In many cases this can yield a substantial speed up over the default, matrix-free implementation in OpenMDAO.
Note
Assembled Jacobians are especially effective when you have a deeply-nested hierarchy with a large number of components and/or variables. See the matrix-free for more details on how to best select which type of Jacobian to use.
To use an assembled Jacobian, you set the assemble_jac
option of the linear solver that will use it to True. The type of the assembled Jacobian will be determined by the value of options['assembled_jac_type']
in the solver’s containing system. There are two options of assembled_jac_type
to choose from, dense
and csc
.
Note
csc
is an abbreviation for compressed sparse column. csc
is one of many sparse storage schemes that allocate contiguous storage in memory for the nonzero elements of the matrix, and perhaps a limited number of zeros. For more information, see Compressed sparse column.
For example:
model.options['assembled_jac_type'] = 'dense'
model.linear_solver = DirectSolver(assemble_jac=True)
‘csc’ is the default, and you should try that first if you’re not sure of which one to use. Most problems, even if they have dense sub-Jacobians from each component, are fairly sparse at the model level and the DirectSolver will usually be much faster with a sparse factorization.
Note
You are allowed to use multiple assembled Jacobians at multiple different levels of your model hierarchy. This may be useful if you have nested nonlinear solvers to converge very difficult problems.
Usage Example#
In the following example, borrowed from the newton solver tutorial, we assemble the Jacobian at the same level of the model hierarchy as the NewtonSolver and DirectSolver. In general, you will always do the assembly at the same level as the linear solver that will make use of the Jacobian matrix.
Resistor
class definition
class Resistor(om.ExplicitComponent):
"""Computes current across a resistor using Ohm's law."""
def initialize(self):
self.options.declare('R', default=1., desc='Resistance in Ohms')
def setup(self):
self.add_input('V_in', units='V')
self.add_input('V_out', units='V')
self.add_output('I', units='A')
def setup_partials(self):
self.declare_partials('I', 'V_in', method='fd')
self.declare_partials('I', 'V_out', method='fd')
def compute(self, inputs, outputs):
deltaV = inputs['V_in'] - inputs['V_out']
outputs['I'] = deltaV / self.options['R']
Diode
class definition
class Diode(om.ExplicitComponent):
"""Computes current across a diode using the Shockley diode equation."""
def initialize(self):
self.options.declare('Is', default=1e-15, desc='Saturation current in Amps')
self.options.declare('Vt', default=.025875, desc='Thermal voltage in Volts')
def setup(self):
self.add_input('V_in', units='V')
self.add_input('V_out', units='V')
self.add_output('I', units='A')
def setup_partials(self):
self.declare_partials('I', 'V_in', method='fd')
self.declare_partials('I', 'V_out', method='fd')
def compute(self, inputs, outputs):
deltaV = inputs['V_in'] - inputs['V_out']
Is = self.options['Is']
Vt = self.options['Vt']
outputs['I'] = Is * (np.exp(deltaV / Vt) - 1)
Node
class definition
class Node(om.ImplicitComponent):
"""Computes voltage residual across a node based on incoming and outgoing current."""
def initialize(self):
self.options.declare('n_in', default=1, types=int, desc='number of connections with + assumed in')
self.options.declare('n_out', default=1, types=int, desc='number of current connections + assumed out')
def setup(self):
self.add_output('V', val=5., units='V')
for i in range(self.options['n_in']):
i_name = 'I_in:{}'.format(i)
self.add_input(i_name, units='A')
for i in range(self.options['n_out']):
i_name = 'I_out:{}'.format(i)
self.add_input(i_name, units='A')
def setup_partials(self):
#note: we don't declare any partials wrt `V` here,
# because the residual doesn't directly depend on it
self.declare_partials('V', 'I*', method='fd')
def apply_nonlinear(self, inputs, outputs, residuals):
residuals['V'] = 0.
for i_conn in range(self.options['n_in']):
residuals['V'] += inputs['I_in:{}'.format(i_conn)]
for i_conn in range(self.options['n_out']):
residuals['V'] -= inputs['I_out:{}'.format(i_conn)]
import openmdao.api as om
from openmdao.test_suite.scripts.circuit_analysis import Resistor, Diode, Node
class Circuit(om.Group):
def setup(self):
self.add_subsystem('n1', Node(n_in=1, n_out=2), promotes_inputs=[('I_in:0', 'I_in')])
self.add_subsystem('n2', Node()) # leaving defaults
self.add_subsystem('R1', Resistor(R=100.), promotes_inputs=[('V_out', 'Vg')])
self.add_subsystem('R2', Resistor(R=10000.))
self.add_subsystem('D1', Diode(), promotes_inputs=[('V_out', 'Vg')])
self.connect('n1.V', ['R1.V_in', 'R2.V_in'])
self.connect('R1.I', 'n1.I_out:0')
self.connect('R2.I', 'n1.I_out:1')
self.connect('n2.V', ['R2.V_out', 'D1.V_in'])
self.connect('R2.I', 'n2.I_in:0')
self.connect('D1.I', 'n2.I_out:0')
self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
self.nonlinear_solver.options['iprint'] = 2
self.nonlinear_solver.options['maxiter'] = 20
##################################################################
# Assemble at the group level. Default assembled jac type is 'csc'
##################################################################
self.options['assembled_jac_type'] = 'csc'
self.linear_solver = om.DirectSolver(assemble_jac=True)
p = om.Problem()
model = p.model
model.add_subsystem('circuit', Circuit())
p.setup()
p.set_val('circuit.I_in', 0.1)
p.set_val('circuit.Vg', 0.)
# set some initial guesses
p.set_val('circuit.n1.V', 10.)
p.set_val('circuit.n2.V', 1.)
p.run_model()
=======
circuit
=======
NL: Newton 0 ; 59.9046598 1
NL: Newton 1 ; 22.3887968 0.373740489
NL: Newton 2 ; 8.23629948 0.13749013
NL: Newton 3 ; 3.02978539 0.0505767899
NL: Newton 4 ; 1.11437824 0.0186025301
NL: Newton 5 ; 0.409725119 0.00683962016
NL: Newton 6 ; 0.150492273 0.00251219644
NL: Newton 7 ; 0.0551239641 0.000920194928
NL: Newton 8 ; 0.0200403261 0.000334537016
NL: Newton 9 ; 0.00713760492 0.000119149411
NL: Newton 10 ; 0.00240281896 4.01107188e-05
NL: Newton 11 ; 0.000692025597 1.15521163e-05
NL: Newton 12 ; 0.000129140479 2.15576684e-06
NL: Newton 13 ; 7.60315072e-06 1.26920856e-07
NL: Newton 14 ; 3.10658299e-08 5.1858787e-10
NL: Newton 15 ; 1.12045605e-12 1.87039882e-14
NL: Newton Converged
print(p['circuit.n1.V'])
print(p['circuit.n2.V'])
print(p['circuit.R1.I'])
print(p['circuit.R2.I'])
print(p['circuit.D1.I'])
# sanity check: should sum to .1 Amps
print(p['circuit.R1.I'] + p['circuit.D1.I'])
[9.90804735]
[0.71278185]
[0.09908047]
[0.00091953]
[0.00091953]
[0.1]