Sparse Partial Derivatives

When a partial derivative is sparse (few nonzero entries compared to the total size of the matrix), it may be advantageous to utilize a format that stores only the nonzero entries. To use sparse partial derivatives, they must first be declared with the sparsity pattern in setup using the declare_partials method.

Usage

1. To specify the sparsity pattern in the AIJ format (alternatively known as COO format), use the rows and cols arguments to declare_partials. For example, to declare a sparsity pattern of nonzero entries in the (0, 0), (1, 1), (1, 2), and (1,3) positions, one would use rows=[0, 1, 1, 1], cols=[0, 1, 2, 3]. When using compute_partials, you do not need to pass the sparsity pattern again. Instead, you simply give the values for the entries in the same order as given in declare_partials.

import numpy as np

from openmdao.api import Problem, Group, IndepVarComp, ExplicitComponent

class SparsePartialComp(ExplicitComponent):
    def setup(self):
        self.add_input('x', shape=(4,))
        self.add_output('f', shape=(2,))

        self.declare_partials(of='f', wrt='x', rows=[0, 1, 1, 1], cols=[0, 1, 2, 3])

    def compute_partials(self, inputs, partials):
        # Corresponds to the [(0,0), (1,1), (1,2), (1,3)] entries.
        partials['f', 'x'] = [1., 2., 3., 4.]

model = Group()
comp = IndepVarComp()
comp.add_output('x', np.ones(4))

model.add_subsystem('input', comp)
model.add_subsystem('example', SparsePartialComp())

model.connect('input.x', 'example.x')

problem = Problem(model=model)
problem.setup(check=False)
problem.run_model()
totals = problem.compute_totals(['example.f'], ['input.x'])

print(totals['example.f', 'input.x'])
[[ 1. -0. -0. -0.]
 [-0.  2.  3.  4.]]
  1. If only some of your Jacobian entries change across iterations, or if you wish to avoid creating intermediate arrays, you may update the entries in-place.
import numpy as np

from openmdao.api import Problem, Group, IndepVarComp, ExplicitComponent

class SparsePartialComp(ExplicitComponent):
    def setup(self):
        self.add_input('x', shape=(4,))
        self.add_output('f', shape=(2,))

        self.declare_partials(of='f', wrt='x', rows=[0,1,1,1], cols=[0,1,2,3])

    def compute_partials(self, inputs, partials):
        pd = partials['f', 'x']

        # Corresponds to the (0, 0) entry
        pd[0] = 1.

        # (1,1) entry
        pd[1] = 2.

        # (1, 2) entry
        pd[2] = 3.

        # (1, 3) entry
        pd[3] = 4


model = Group()
comp = IndepVarComp()
comp.add_output('x', np.ones(4))

model.add_subsystem('input', comp)
model.add_subsystem('example', SparsePartialComp())

model.connect('input.x', 'example.x')

problem = Problem(model=model)
problem.setup(check=False)
problem.run_model()
totals = problem.compute_totals(['example.f'], ['input.x'])

print(totals['example.f', 'input.x'])
[[ 1. -0. -0. -0.]
 [-0.  2.  3.  4.]]
  1. If your partial derivative is constant and sparse, or if you simply wish to provide an initial value for the derivative, you can pass in the values using the val argument. If you are using the AIJ format, val should receive the nonzero entries in the same order as given for rows and cols. Alternatively, you may provide a Scipy sparse matrix, from which the sparsity pattern is deduced.
import numpy as np
import scipy as sp

from openmdao.api import Problem, Group, IndepVarComp, ExplicitComponent

class SparsePartialComp(ExplicitComponent):
    def setup(self):
        self.add_input('x', shape=(4,))
        self.add_input('y', shape=(2,))
        self.add_output('f', shape=(2,))

        self.declare_partials(of='f', wrt='x', rows=[0,1,1,1], cols=[0,1,2,3],
                              val=[1. , 2., 3., 4.])
        self.declare_partials(of='f', wrt='y', val=sp.sparse.eye(2, format='csc'))

    def compute_partials(self, inputs, partials):
        pass

model = Group()
comp = IndepVarComp()
comp.add_output('x', np.ones(4))
comp.add_output('y', np.ones(2))

model.add_subsystem('input', comp)
model.add_subsystem('example', SparsePartialComp())

model.connect('input.x', 'example.x')
model.connect('input.y', 'example.y')

problem = Problem(model=model)
problem.setup(check=False)
problem.run_model()
totals = problem.compute_totals(['example.f'], ['input.x', 'input.y'])

print(totals['example.f', 'input.x'])
[[ 1. -0. -0. -0.]
 [-0.  2.  3.  4.]]
print(totals['example.f', 'input.y'])
[[ 1. -0.]
 [-0.  1.]]