Merge a9df8a2 into 67689d5

openmdao/core/group.py

@@ -2177,7 +2177,7 @@ def compute_sys_graph(self, comps_only=False):
         if comps_only:
             systems = [s.pathname for s in self.system_iter(recurse=True, typ=Component)]
         else:
-            systems = [s.pathname for s in self._subsystems_myproc]
+            systems = [s.name for s in self._subsystems_myproc]
 
         if MPI:
             sysbyproc = self.comm.allgather(systems)
@@ -2205,6 +2205,7 @@ def compute_sys_graph(self, comps_only=False):
 
         for key in edge_data:
             src_sys, tgt_sys = key
-            graph.add_edge(src_sys, tgt_sys, conns=edge_data[key])
+            if comps_only or src_sys != tgt_sys:
+                graph.add_edge(src_sys, tgt_sys, conns=edge_data[key])
 
         return graph

openmdao/core/problem.py

@@ -662,6 +662,60 @@ def run_driver(self, case_prefix=None, reset_iter_counts=True):
         self.model._clear_iprint()
         return self.driver.run()
 
+    def compute_jacvec_product(self, of, wrt, mode, seed):
+        """
+        Given a seed and 'of' and 'wrt' variables, compute the total jacobian vector product.
+
+        Parameters
+        ----------
+        of : list of str
+            Variables whose derivatives will be computed.
+        wrt : list of str
+            Derivatives will be computed with respect to these variables.
+        mode : str
+            Derivative direction ('fwd' or 'rev').
+        seed : dict or list
+            Either a dict keyed by 'wrt' varnames (fwd) or 'of' varnames (rev), containing
+            dresidual (fwd) or doutput (rev) values, OR a list of dresidual or doutput
+            values that matches the corresponding 'wrt' (fwd) or 'of' (rev) varname list.
+
+        Returns
+        -------
+        dict
+            The total jacobian vector product, keyed by variable name.
+        """
+        if mode == 'fwd':
+            if len(wrt) != len(seed):
+                raise RuntimeError("seed and 'wrt' list must be the same length in fwd mode.")
+            lnames, rnames = of, wrt
+            lkind, rkind = 'output', 'residual'
+        else:  # rev
+            if len(of) != len(seed):
+                raise RuntimeError("seed and 'of' list must be the same length in rev mode.")
+            lnames, rnames = wrt, of
+            lkind, rkind = 'residual', 'output'
+
+        rvec = self.model._vectors[rkind]['linear']
+        lvec = self.model._vectors[lkind]['linear']
+
+        # set seed values into dresids (fwd) or doutputs (rev)
+        try:
+            seed[rnames[0]]
+        except (IndexError, TypeError):
+            for i, name in enumerate(rnames):
+                rvec[name] = seed[i]
+        else:
+            for name in rnames:
+                rvec[name] = seed[name]
+
+        # We apply a -1 here because the derivative of the output is minus the derivative of
+        # the residual in openmdao.
+        rvec._data *= -1.
+
+        self.model.run_solve_linear(['linear'], mode)
+
+        return {n: lvec[n].copy() for n in lnames}
+
     def run_once(self):
         """
         Backward compatible call for run_model.

openmdao/core/tests/test_compute_jacvec_prod.py

@@ -0,0 +1,222 @@
+
+import sys
+import unittest
+
+from six import assertRaisesRegex, StringIO, assertRegex, iteritems
+
+import numpy as np
+
+import openmdao.api as om
+from openmdao.core.system import get_relevant_vars
+from openmdao.core.driver import Driver
+from openmdao.utils.assert_utils import assert_rel_error, assert_warning
+from openmdao.test_suite.components.paraboloid import Paraboloid
+from openmdao.test_suite.components.sellar import SellarDerivatives
+
+
+def get_comp(size):
+    return om.ExecComp('out = sum(3*x**2) + inp', x=np.ones(size - 1), inp=0.0, out=0.0)
+
+
+class SubProbComp(om.ExplicitComponent):
+    """
+    This component contains a sub-Problem with a component that will be solved over num_nodes
+    points instead of creating num_nodes instances of that same component and connecting them
+    together.
+    """
+    def __init__(self, input_size, num_nodes, mode, **kwargs):
+        super(SubProbComp, self).__init__(**kwargs)
+        self.prob = None
+        self.size = input_size
+        self.num_nodes = num_nodes
+        self.mode = mode
+
+    def _setup_subprob(self):
+        self.prob = p = om.Problem(comm=self.comm)
+        model = self.prob.model
+
+        indep = model.add_subsystem('indep', om.IndepVarComp('x', val=np.zeros(self.size - 1)))
+        indep.add_output('inp', val=0.0)
+
+        model.add_subsystem('comp', get_comp(self.size))
+
+        model.connect('indep.x', 'comp.x')
+        model.connect('indep.inp', 'comp.inp')
+
+        p.setup()
+        p.final_setup()
+
+    def setup(self):
+        self._setup_subprob()
+
+        self.add_input('x', np.zeros(self.size - 1))
+        self.add_input('inp', val=0.0)
+        self.add_output('out', val=0.0)
+        self.declare_partials('*', '*')
+
+    def compute(self, inputs, outputs):
+        p = self.prob
+        p['indep.x'] = inputs['x']
+        p['indep.inp'] = inputs['inp']
+        inp = inputs['inp']
+        for i in range(self.num_nodes):
+            p['indep.inp'] = inp
+            p.run_model()
+            inp = p['comp.out']
+
+        outputs['out'] = p['comp.out']
+
+    def _compute_partials_fwd(self, inputs, partials):
+        p = self.prob
+        x = inputs['x']
+        p['indep.x'] = x
+        p['indep.inp'] = inputs['inp']
+
+        seed = {'indep.x':np.zeros(x.size), 'indep.inp': np.zeros(1)}
+        p.run_model()
+        p.model._linearize(None)
+        for rhsname in seed:
+            for rhs_i in range(seed[rhsname].size):
+                seed['indep.x'][:] = 0.0
+                seed['indep.inp'][:] = 0.0
+                seed[rhsname][rhs_i] = 1.0
+                for i in range(self.num_nodes):
+                    p.model._vectors['output']['linear'].set_const(0.0)
+                    p.model._vectors['residual']['linear'].set_const(0.0)
+                    jvp = p.compute_jacvec_product(of=['comp.out'], wrt=['indep.x','indep.inp'], mode='fwd', seed=seed)
+                    seed['indep.inp'][:] = jvp['comp.out']
+
+                if rhsname == 'indep.x':
+                    partials[self.pathname + '.out', self.pathname +'.x'][0, rhs_i] = jvp[self.pathname + '.out']
+                else:
+                    partials[self.pathname + '.out', self.pathname + '.inp'][0, 0] = jvp[self.pathname + '.out']
+
+    def _compute_partials_rev(self, inputs, partials):
+        p = self.prob
+        p['indep.x'] = inputs['x']
+        p['indep.inp'] = inputs['inp']
+        seed = {'comp.out': np.ones(1)}
+
+        stack = []
+        comp = p.model.comp
+        comp._inputs['inp'] = inputs['inp']
+        # store the inputs to each comp (the comp at each node point) by doing nonlinear solves
+        # and storing what the inputs are for each node point.  We'll set these inputs back
+        # later when we linearize about each node point.
+        for i in range(self.num_nodes):
+            stack.append(comp._inputs['inp'][0])
+            comp._inputs['x'] = inputs['x']
+            comp._solve_nonlinear()
+            comp._inputs['inp'] = comp._outputs['out']
+
+        for i in range(self.num_nodes):
+            p.model._vectors['output']['linear'].set_const(0.0)
+            p.model._vectors['residual']['linear'].set_const(0.0)
+            comp._inputs['inp'] = stack.pop()
+            comp._inputs['x'] = inputs['x']
+            p.model._linearize(None)
+            jvp = p.compute_jacvec_product(of=['comp.out'], wrt=['indep.x','indep.inp'], mode='rev', seed=seed)
+            seed['comp.out'][:] = jvp['indep.inp']
+
+            # all of the comp.x's are connected to indep.x, so we have to accumulate their
+            # contributions together
+            partials[self.pathname + '.out', self.pathname + '.x'] += jvp['indep.x']
+
+            # this one doesn't get accumulated because each comp.inp contributes to the
+            # previous comp's .out (or to indep.inp in the case of the first comp) only.
+            # Note that we have to handle this explicitly here because normally in OpenMDAO
+            # we accumulate derivatives when we do reverse transfers.  We can't do that
+            # here because we only have one instance of our component, so instead of
+            # accumulating into separate 'comp.out' variables for each comp instance,
+            # we would be accumulating into a single comp.out variable, which would make
+            # our derivative too big.
+            partials[self.pathname + '.out', self.pathname + '.inp'] = jvp['indep.inp']
+
+    def compute_partials(self, inputs, partials):
+        # note that typically you would only have to define partials for one direction,
+        # either fwd OR rev, not both.
+        if self.mode == 'fwd':
+            self._compute_partials_fwd(inputs, partials)
+        else:
+            self._compute_partials_rev(inputs, partials)
+
+
+class TestPComputeJacvecProd(unittest.TestCase):
+
+    def _build_om_model(self, size):
+        p = om.Problem()
+        model = p.model
+        indep = model.add_subsystem('indep', om.IndepVarComp('x', val=np.zeros(size - 1)))
+        indep.add_output('inp', val=0.0)
+
+        C1 = model.add_subsystem('C1', get_comp(size))
+        C2 = model.add_subsystem('C2', get_comp(size))
+        C3 = model.add_subsystem('C3', get_comp(size))
+
+        model.connect('indep.x', ['C1.x', 'C2.x', 'C3.x'])
+        model.connect('indep.inp', 'C1.inp')
+        model.connect('C1.out', 'C2.inp')
+        model.connect('C2.out', 'C3.inp')
+
+        return p
+
+    def _build_cjv_model(self, size, mode):
+        p = om.Problem()
+        indep = p.model.add_subsystem('indep', om.IndepVarComp('x', val=np.zeros(size - 1)))
+        indep.add_output('inp', val=0.0)
+
+        comp = p.model.add_subsystem('comp', SubProbComp(input_size=size, num_nodes=3, mode=mode))
+        p.model.connect('indep.x', 'comp.x')
+        p.model.connect('indep.inp', 'comp.inp')
+
+        #p.model.linear_solver = om.DirectSolver()
+        p.setup(mode=mode)
+
+        return p
+
+    def test_fwd(self):
+        size = 5
+        p = self._build_om_model(size)
+        p.setup(mode='fwd')
+        p['indep.x'] = np.arange(size-1, dtype=float) + 1. #np.random.random(size - 1)
+        p['indep.inp'] = np.array([7.])  #np.random.random(1)[0]
+        p.final_setup()
+
+        p2 = self._build_cjv_model(size, 'fwd')
+
+        p2['indep.x'] = p['indep.x']
+        p2['indep.inp'] = p['indep.inp']
+
+        p2.run_model()
+        J2 = p2.compute_totals(of=['comp.out'], wrt=['indep.x', 'indep.inp'], return_format='array')
+
+        p.run_model()
+        J = p.compute_totals(of=['C3.out'], wrt=['indep.x', 'indep.inp'], return_format='array')
+
+        self.assertEqual(p['C3.out'], p2['comp.out'])
+        np.testing.assert_allclose(J2, J)
+
+
+    def test_rev(self):
+        size = 5
+        p = self._build_om_model(size)
+        p.setup(mode='rev')
+        p['indep.x'] = np.arange(size-1, dtype=float) + 1. #np.random.random(size - 1)
+        p['indep.inp'] = np.array([7.])  #np.random.random(1)[0]
+        p.final_setup()
+
+        p2 = self._build_cjv_model(size, 'rev')
+
+        p2['indep.x'] = p['indep.x']
+        p2['indep.inp'] = p['indep.inp']
+
+        p.run_model()
+        J = p.compute_totals(of=['C3.out'], wrt=['indep.x', 'indep.inp'], return_format='array')
+
+        p2.run_model()
+        self.assertEqual(p['C3.out'], p2['comp.out'])
+
+        J2 = p2.compute_totals(of=['comp.out'], wrt=['indep.x', 'indep.inp'], return_format='array')
+
+        np.testing.assert_allclose(J2, J)
+

openmdao/core/tests/test_problem.py

@@ -2,8 +2,9 @@
 
 import sys
 import unittest
+import itertools
 
-from six import assertRaisesRegex, StringIO, assertRegex, iteritems
+from six import assertRaisesRegex, StringIO, assertRegex
 
 import numpy as np
 
@@ -14,6 +15,11 @@
 from openmdao.test_suite.components.paraboloid import Paraboloid
 from openmdao.test_suite.components.sellar import SellarDerivatives
 
+try:
+    from parameterized import parameterized
+except ImportError:
+    from openmdao.utils.assert_utils import SkipParameterized as parameterized
+
 
 class SellarOneComp(om.ImplicitComponent):
 
@@ -470,6 +476,53 @@ def test_compute_totals_no_args_promoted(self):
 
         assert_rel_error(self, derivs['calc.y', 'des_vars.x'], [[2.0]], 1e-6)
 
+    @parameterized.expand(itertools.product(['fwd', 'rev']))
+    def test_compute_jacvec_product(self, mode):
+
+        prob = om.Problem()
+        prob.model = SellarDerivatives()
+        prob.model.nonlinear_solver = om.NonlinearBlockGS()
+
+        prob.setup(mode=mode)
+        prob.run_model()
+
+        of = ['obj', 'con1']
+        wrt = ['x', 'z']
+
+        if mode == 'fwd':
+            seed_names = wrt
+            result_names = of
+            rvec = prob.model._vectors['output']['linear']
+            lvec = prob.model._vectors['residual']['linear']
+        else:
+            seed_names = of
+            result_names = wrt
+            rvec = prob.model._vectors['residual']['linear']
+            lvec = prob.model._vectors['output']['linear']
+
+        J = prob.compute_totals(of, wrt, return_format='array')
+
+        seed = []
+        rvec._data[:] = 0.
+        lvec._data[:] = 0.
+        for name in seed_names:
+            seed.append(np.random.random(rvec[name].size))
+
+        resdict = prob.compute_jacvec_product(of, wrt, mode, seed)
+        result = []
+        for name in result_names:
+            result.append(resdict[name].flat)
+        result = np.hstack(result)
+
+        testvec = np.hstack(seed)
+
+        if mode == 'fwd':
+            checkvec = J.dot(testvec)
+        else:
+            checkvec = J.T.dot(testvec)
+
+        np.testing.assert_allclose(checkvec, result)
+
     def test_feature_set_indeps(self):
         import openmdao.api as om
         from openmdao.test_suite.components.paraboloid import Paraboloid