In-Memory Assembly of Jacobians

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='cs')
        self.declare_partials('I', 'V_out', method='cs')

    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='cs')
        self.declare_partials('I', 'V_out', method='cs')

    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='cs')

    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.3883642 0.373733267
NL: Newton 2 ; 8.23598116 0.137484817
NL: Newton 3 ; 3.02960973 0.0505738576
NL: Newton 4 ; 1.11429207 0.0186010917
NL: Newton 5 ; 0.409685491 0.00683895866
NL: Newton 6 ; 0.150474776 0.00251190436
NL: Newton 7 ; 0.055116451 0.000920069511
NL: Newton 8 ; 0.0200371642 0.000334484233
NL: Newton 9 ; 0.00713629394 0.000119127526
NL: Newton 10 ; 0.00240228312 4.01017739e-05
NL: Newton 11 ; 0.000691815028 1.15486012e-05
NL: Newton 12 ; 0.000129072995 2.15464031e-06
NL: Newton 13 ; 7.59363152e-06 1.26761951e-07
NL: Newton 14 ; 3.08435461e-08 5.14877244e-10
NL: Newton 15 ; 5.1439842e-13 8.5869517e-15
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]

In-Memory Assembly of Jacobians

Contents

In-Memory Assembly of Jacobians#

Usage Example#