Allow BalanceComp to be sized from rhs_val

openmdao/components/balance_comp.py

@@ -314,7 +314,13 @@ def add_balance(self, name, eq_units=None, lhs_name=None, rhs_name=None, rhs_val
 
         self._state_vars[name] = options
 
-        if val is not None:
+        if val is None:
+            # If user doesn't specify initial guess for val, we cann size problem from initial
+            # rhs_val.
+            if 'shape' not in kwargs and not np.isscalar(rhs_val):
+                kwargs['shape'] = rhs_val.shape
+
+        else:
             options['kwargs']['val'] = val
 
         meta = self.add_output(name, **options['kwargs'])

openmdao/components/tests/test_balance_comp.py

@@ -1,7 +1,6 @@
 """
 Unit tests for the BalanceComp.
 """
-import os
 import unittest
 import warnings
 
@@ -897,6 +896,27 @@ def test_specified_shape(self):
 
         assert_check_partials(cpd, atol=1e-5, rtol=1e-5)
 
+    def test_shape_from_rhs_val(self):
+        n = 5
+        p = om.Problem()
+
+        p.model.add_subsystem('comp_A',
+                              om.ExecComp('y = x**2', y={'shape': (n, )}, x={'shape': (n, )}))
+        p.model.add_subsystem('comp_B',
+                              om.ExecComp('y = x + 7', y={'shape': (n, )}, x={'shape': (n, )}))
+
+        p.model.add_subsystem('bal', om.BalanceComp('x', rhs_val=np.ones((n, ))))
+        p.model.connect('bal.x', 'comp_A.x')
+        p.model.connect('bal.x', 'comp_B.x')
+        p.model.connect('comp_A.y', 'bal.lhs:x')
+        p.model.connect('comp_B.y', 'bal.rhs:x')
+
+        p.model.linear_solver = om.DirectSolver()
+        p.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)
+
+        # Bug was a size mismatch exception raised during setup.
+        p.setup()
+
 
 if __name__ == '__main__':  # pragma: no cover
     unittest.main()