Merge pull request #325 from robfalck/balance_update

user can set default RHS (target) value for BalanceComp

openmdao/components/balance_comp.py

@@ -16,7 +16,7 @@ class BalanceComp(ImplicitComponent):
     """
 
     def __init__(self, name=None, eq_units=None, lhs_name=None,
-                 rhs_name=None, mult_name=None, mult_val=1.0, **kwargs):
+                 rhs_name=None, rhs_val=0.0, mult_name=None, mult_val=1.0, **kwargs):
         r"""
         Initialize a BalanceComp, optionally creating a new implicit state variable.
 
@@ -108,6 +108,9 @@ def __init__(self, name=None, eq_units=None, lhs_name=None,
         rhs_name : str or None
             Optional name for the RHS variable associated with the implicit state variable.  If
             None, the default will be used:  'rhs:{name}'.
+        rhs_val : int, float, or np.array
+            Default value for the RHS of the given state.  Must be compatible
+            with the shape (optionally) given by the val option in kwargs.
         mult_name : str or None
             Optional name for the LHS multiplier variable associated with the implicit state
             variable. If None, the default will be used: 'mult:{name}'.
@@ -120,7 +123,8 @@ def __init__(self, name=None, eq_units=None, lhs_name=None,
         super(BalanceComp, self).__init__()
         self._state_vars = {}
         if name is not None:
-            self.add_balance(name, eq_units, lhs_name, rhs_name, mult_name, mult_val, **kwargs)
+            self.add_balance(name, eq_units, lhs_name, rhs_name, rhs_val,
+                             mult_name, mult_val, **kwargs)
 
     def setup(self):
         """
@@ -142,9 +146,19 @@ def setup(self):
             self._state_vars[name]['size'] = n
 
             self.add_output(name, **options['kwargs'])
-            self.add_input(options['lhs_name'], val=np.ones(n), units=options['eq_units'])
-            self.add_input(options['rhs_name'], val=np.ones(n), units=options['eq_units'])
-            self.add_input(options['mult_name'], val=options['mult_val'] * np.ones(n), units=None)
+
+            self.add_input(options['lhs_name'],
+                           val=np.ones(n),
+                           units=options['eq_units'])
+
+            self.add_input(options['rhs_name'],
+                           val=options['rhs_val'] * np.ones(n),
+                           units=options['eq_units'])
+
+            self.add_input(options['mult_name'],
+                           val=options['mult_val'] * np.ones(n),
+                           units=None)
+
             self._scale_factor = np.ones(n)
             self._dscale_drhs = np.ones(n)
 
@@ -211,7 +225,7 @@ def linearize(self, inputs, outputs, jacobian):
             jacobian[name, mult_name] = lhs * self._scale_factor
 
     def add_balance(self, name, eq_units=None, lhs_name=None,
-                    rhs_name=None, mult_name=None, mult_val=1.0, **kwargs):
+                    rhs_name=None, rhs_val=0.0, mult_name=None, mult_val=1.0, **kwargs):
         """
         Add a new state variable and associated equation to be balanced.
 
@@ -231,6 +245,9 @@ def add_balance(self, name, eq_units=None, lhs_name=None,
         rhs_name : str or None
             Optional name for the RHS variable associated with the implicit state variable.  If
             None, the default will be used:  'rhs:{name}'.
+        rhs_val : int, float, or np.array
+            Default value for the RHS.  Must be compatible with the shape (optionally)
+            given by the val option in kwargs.
         mult_name : str or None
             Optional name for the LHS multiplier variable associated with the implicit state
             variable. If None, the default will be used: 'mult:{name}'.
@@ -244,5 +261,6 @@ def add_balance(self, name, eq_units=None, lhs_name=None,
                                   'eq_units': eq_units,
                                   'lhs_name': lhs_name,
                                   'rhs_name': rhs_name,
+                                  'rhs_val': rhs_val,
                                   'mult_name': mult_name,
                                   'mult_val': mult_val}

openmdao/components/tests/test_balance_comp.py

@@ -280,6 +280,47 @@ def test_scalar(self):
         for (of, wrt) in cpd['balance']:
             assert_almost_equal(cpd['balance'][of, wrt]['abs error'], 0.0, decimal=5)
 
+    def test_rhs_val(self):
+        """ Test solution with a default RHS value and no connected RHS variable. """
+
+        n = 1
+
+        prob = Problem(model=Group())
+
+        bal = BalanceComp('x', rhs_val=4.0)
+
+        exec_comp = ExecComp('y=x**2', x={'value': 1}, y={'value': 1})
+
+        prob.model.add_subsystem(name='exec', subsys=exec_comp)
+
+        prob.model.add_subsystem(name='balance', subsys=bal)
+
+        prob.model.connect('balance.x', 'exec.x')
+        prob.model.connect('exec.y', 'balance.lhs:x')
+
+        prob.model.linear_solver = DirectSolver()
+
+        prob.model.nonlinear_solver = NewtonSolver()
+        prob.model.nonlinear_solver.options['maxiter'] = 100
+        prob.model.nonlinear_solver.options['iprint'] = 0
+
+        prob.model.jacobian = DenseJacobian()
+
+        prob.setup()
+
+        prob['balance.x'] = np.random.rand(n)
+
+        prob.run_model()
+
+        assert_almost_equal(prob['balance.x'], 2.0, decimal=7)
+
+        np.set_printoptions(linewidth=1024)
+
+        cpd = prob.check_partials()
+
+        for (of, wrt) in cpd['balance']:
+            assert_almost_equal(cpd['balance'][of, wrt]['abs error'], 0.0, decimal=5)
+
     def test_scalar_with_mult(self):
 
         n = 1