BalanceComp#

BalanceComp is a specialized implementation of ImplicitComponent that is intended to provide a simple way to implement most implicit equations without the need to define your own residuals.

BalanceComp Options#

OptionDefaultAcceptable ValuesAcceptable TypesDescription
always_optFalse[True, False]['bool']If True, force nonlinear operations on this component to be included in the optimization loop even if this component is not relevant to the design variables and responses.
assembled_jac_typecsc['csc', 'dense']N/ALinear solver(s) in this group or implicit component, if using an assembled jacobian, will use this type.
derivs_methodN/A['jax', 'cs', 'fd', None]N/AThe method to use for computing derivatives
distributedFalse[True, False]['bool']If True, set all variables in this component as distributed across multiple processes
guess_funcN/AN/A['function']A callable function in the form f(inputs, outputs, residuals) that can provide an initial "guess" value of the state variable(s) based on the inputs, outputs and residuals.
run_root_onlyFalse[True, False]['bool']If True, call compute, compute_partials, linearize, apply_linear, apply_nonlinear, and compute_jacvec_product only on rank 0 and broadcast the results to the other ranks.
use_jitTrue[True, False]['bool']If True, attempt to use jit on compute_primal, assuming jax or some other AD package is active.

BalanceComp Constructor#

The call signature for the BalanceComp constructor is:

BalanceComp.__init__()[source]

Initialize a BalanceComp, optionally creating a new implicit state variable.

The BalanceComp allows for the creation of one or more implicit state variables, and computes the residuals for those variables based on the following equation.

\[\mathcal{R}_{name} = \frac{f_{mult}(x,...) \times f_{lhs}(x,...) - f_{rhs}(x,...)}{f_{norm}(f_{rhs}(x,...))}\]

Where \(f_{lhs}\) represents the left-hand-side of the equation, \(f_{rhs}\) represents the right-hand-side, and \(f_{mult}\) is an optional multiplier on the left hand side. At least one of these quantities should be a function of the associated state variable. If use_mult is True the default value of the multiplier is 1. The optional normalization function \(f_{norm}(f_{rhs}(x,...))\) is computed as:

\[\begin{split}f_{norm}(f_{rhs}(x,...)) = \begin{cases} \left| f_{rhs} \right|, & \text{if normalize and } \left| f_{rhs} \right| \geq 2 \\ 0.25 f_{rhs}^2 + 1, & \text{if normalize and } \left| f_{rhs} \right| < 2 \\ 1, & \text{if not normalize} \end{cases}\end{split}\]

New state variables, and their associated residuals are created by calling add_balance. As an example, solving the equation \(x**2 = 2\) implicitly can be be accomplished as follows:

prob = Problem()
bal = BalanceComp()
bal.add_balance('x', val=1.0)
exec_comp = ExecComp('y=x**2')
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(solve_subsystems=False)
prob.setup()
prob.set_val('exec.x', 2)
prob.run_model()

The arguments to add_balance can be provided on initialization to provide a balance with a one state/residual without the need to call add_balance:

prob = Problem()
bal = BalanceComp('x', val=1.0)
exec_comp = ExecComp('y=x**2')
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(solve_subsystems=False)
prob.setup()
prob.set_val('exec.x', 2)
prob.run_model()

Using the BalanceComp#

BalanceComp allows you to add one or more state variables and its associated implicit equations. For each balance added to the component it solves the following equation:

\[ \begin{align} \mathcal{R}_{name} = \frac{f_{mult}(x,...) \times f_{lhs}(x,...) - f_{rhs}(x,...)}{f_{norm}(f_{rhs}(x,...))} \end{align} \]

The optional normalization function \(f_{norm}(f_{rhs})\) is computed as:

\[\begin{split} \begin{align} f_{norm}(f_{rhs}(x,...)) = \begin{cases} \left| f_{rhs} \right|, & \text{if normalize and } \left| f_{rhs} \right| \geq 2 \\ 0.25 f_{rhs}^2 + 1, & \text{if normalize and } \left| f_{rhs} \right| < 2 \\ 1, & \text{if not normalize} \end{cases} \end{align} \end{split}\]

The following inputs and outputs are associated with each implicit state.

Name

I/O

Description

{name}

output

implicit state variable

lhs:{name}

input

left-hand side of equation to be balanced

rhs:{name}

input

right-hand side of equation to be balanced

mult:{name}

input

left-hand side multiplier of equation to be balanced

The default value for the rhs:{name} input can be set to via the rhs_val argument (see arguments below). If the rhs value is fixed (e.g. 0), then the input can be left unconnected. The lhs:{name} input must always have something connected to it and should be dependent upon the value of the implicit state variable.

The multiplier is optional and will default to 1.0 if not connected.

BalanceComp supports vectorized implicit states. Simply provide a default value or shape when adding the balance that reflects the correct shape.

You can provide the arguments to create a balance when instantiating a BalanceComp or you can use the add_balance method to create one or more state variables after instantiation. The constructor accepts all the same arguments as the add_balance method:

BalanceComp.add_balance(name, eq_units=None, lhs_name=None, rhs_name=None, rhs_val=0.0, use_mult=False, mult_name=None, mult_val=1.0, normalize=True, val=None, **kwargs)[source]

Add a new state variable and associated equation to be balanced.

This will create new inputs lhs:name, rhs:name, and mult:name that will define the left and right sides of the equation to be balanced, and a multiplier for the left-hand-side.

Parameters:
namestr

The name of the state variable to be created.

eq_unitsstr or None

Units for the left-hand-side and right-hand-side of the equation to be balanced.

lhs_namestr or None

Optional name for the LHS variable associated with the implicit state variable. If None, the default will be used: ‘lhs:{name}’.

rhs_namestr or None

Optional name for the RHS variable associated with the implicit state variable. If None, the default will be used: ‘rhs:{name}’.

rhs_valint, float, or np.array

Default value for the RHS. Must be compatible with the shape (optionally) given by the val or shape option in kwargs.

use_multbool

Specifies whether the LHS multiplier is to be used. If True, then an additional input mult_name is created, with the default value given by mult_val, that multiplies lhs. Default is False.

mult_namestr or None

Optional name for the LHS multiplier variable associated with the implicit state variable. If None, the default will be used: ‘mult:{name}’.

mult_valint, float, or np.array

Default value for the LHS multiplier. Must be compatible with the shape (optionally) given by the val or shape option in kwargs.

normalizebool

Specifies whether or not the resulting residual should be normalized by a quadratic function of the RHS.

valfloat, int, or np.ndarray

Set initial value for the state.

**kwargsdict

Additional arguments to be passed for the creation of the implicit state variable. (see add_output method).

Note that the kwargs arguments can include any of the keyword arguments normally available when creating an output variable with the add_output method of a Component.

Example: Scalar Root Finding#

The following example uses a BalanceComp to implicitly solve the equation:

\[ 2 \cdot x^2 = 4 \]

Here, our LHS is connected to a computed value for \(x^2\), the multiplier is 2, and the RHS is 4. The expected solution is \(x=\sqrt{2}\). We initialize \(x\) with a value of 1 so that it finds the positive root.

import openmdao.api as om

prob = om.Problem()

bal = om.BalanceComp()
bal.add_balance('x', use_mult=True)

exec_comp = om.ExecComp('y=x**2', x={'val': 1}, y={'val': 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 = om.DirectSolver(assemble_jac=True)
prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)

prob.setup()

prob.set_val('balance.rhs:x', 4)
prob.set_val('balance.mult:x', 2.)

# A reasonable initial guess to find the positive root.
prob['balance.x'] = 1.0

prob.run_model()

print(prob.get_val('balance.x'))

[1.41421356]

Alternatively, we could simplify the code by using the mult_val argument.

prob = om.Problem()

bal = om.BalanceComp()
bal.add_balance('x', use_mult=True, mult_val=2.0)

exec_comp = om.ExecComp('y=x**2', x={'val': 1}, y={'val': 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 = om.DirectSolver(assemble_jac=True)
prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)

prob.setup()

prob.set_val('balance.rhs:x', 4)

# A reasonable initial guess to find the positive root.
prob.set_val('balance.x', 1.0)

prob.run_model()
print(prob.get_val('balance.x'))

[1.41421356]

Example: Vectorized Root Finding#

The following example uses a BalanceComp to implicitly solve the equation:

\[ b \cdot x + c = 0 \]

for various values of \(b\), and \(c\). Here, our LHS is connected to a computed value of the linear equation. The multiplier is one and the RHS is zero (the defaults), and thus they need not be connected.

n = 100

prob = om.Problem()

exec_comp = om.ExecComp('y=b*x+c',
                        b={'val': np.random.uniform(0.01, 100, size=n)},
                        c={'val': np.random.rand(n)},
                        x={'val': np.zeros(n)},
                        y={'val': np.ones(n)})

prob.model.add_subsystem(name='exec', subsys=exec_comp)
prob.model.add_subsystem(name='balance', subsys=om.BalanceComp('x', val=np.ones(n)))

prob.model.connect('balance.x', 'exec.x')
prob.model.connect('exec.y', 'balance.lhs:x')

prob.model.linear_solver = om.DirectSolver(assemble_jac=True)
prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)

prob.setup()

prob.set_val('balance.x', np.random.rand(n))

prob.run_model()

b = prob.get_val('exec.b')
c = prob.get_val('exec.c')

print(prob.get_val('balance.x'))

[-1.23904308e-01 -4.76907643e-03 -1.10526767e-02 -1.34274473e-02
 -2.35552093e-02 -1.20340853e-02 -2.40475208e-01 -6.85344686e-03
 -3.81620578e-02 -4.73320484e-02 -2.03003631e-02 -4.26421107e-03
 -1.18110315e-01 -1.05246108e-02 -1.07622380e-02 -6.28236223e-03
 -3.86940127e-02 -6.64769337e-03 -1.06325243e-02 -8.72811414e-04
 -4.28650315e-04 -5.46215582e-03 -7.99678801e-02 -3.45883432e-02
 -2.12332732e-02 -1.18683376e-02 -3.88292088e-02 -4.05974560e-01
 -2.10409732e-01 -4.13397304e-02 -2.12380705e-02 -3.17737934e-02
 -6.41528603e-03 -8.02328547e-03 -2.48223391e-01 -9.12733391e-03
 -4.96661061e-03 -1.30242459e-02 -1.12538249e-02 -2.14681622e-04
 -1.17863351e-02 -4.95967685e-02 -1.06457879e-02 -6.31924450e-01
 -9.86666199e-02 -1.78584745e-02 -8.81560085e-03 -8.80471207e-03
 -4.49769862e-02 -1.94237867e-02 -1.16780938e-02 -1.67750825e-02
 -4.39032508e-03 -1.28942018e-02 -1.05610830e-03 -1.71408675e-02
 -5.65895423e-02 -1.12548938e-02 -7.20617325e-03 -9.91231120e-03
 -4.91428478e-03 -1.03253444e-02 -3.16667837e-02 -1.71700603e-02
 -2.19343908e-03 -5.94791972e-03 -4.62433824e-02 -7.25013281e-03
 -3.09109763e-04 -5.00221374e-03 -2.00848686e-02 -2.13602548e-02
 -4.99566193e-03 -8.92124070e-02 -2.82730304e-01 -2.72574560e-03
 -7.75026361e-04 -2.53598527e-02 -5.42035493e-03 -1.84991281e-02
 -4.49718365e-01 -1.98444555e-03 -8.94594162e-03 -2.00825054e-03
 -1.54284350e-02 -1.30318286e-02 -9.78276630e-03 -4.58061253e-03
 -1.82606298e-02 -4.51666421e-03 -1.63270519e-01 -1.00997507e-01
 -1.61865790e-02 -1.38341889e-04 -2.76501809e-02 -1.01368503e-02
 -8.35901706e-03 -1.77528926e-02 -1.31717155e-02 -1.95986433e-02]

print(-c/b)  # expected

[-1.23904308e-01 -4.76907643e-03 -1.10526767e-02 -1.34274473e-02
 -2.35552093e-02 -1.20340853e-02 -2.40475208e-01 -6.85344686e-03
 -3.81620578e-02 -4.73320484e-02 -2.03003631e-02 -4.26421107e-03
 -1.18110315e-01 -1.05246108e-02 -1.07622380e-02 -6.28236223e-03
 -3.86940127e-02 -6.64769337e-03 -1.06325243e-02 -8.72811414e-04
 -4.28650315e-04 -5.46215582e-03 -7.99678801e-02 -3.45883432e-02
 -2.12332732e-02 -1.18683376e-02 -3.88292088e-02 -4.05974560e-01
 -2.10409732e-01 -4.13397304e-02 -2.12380705e-02 -3.17737934e-02
 -6.41528603e-03 -8.02328547e-03 -2.48223391e-01 -9.12733391e-03
 -4.96661061e-03 -1.30242459e-02 -1.12538249e-02 -2.14681622e-04
 -1.17863351e-02 -4.95967685e-02 -1.06457879e-02 -6.31924450e-01
 -9.86666199e-02 -1.78584745e-02 -8.81560085e-03 -8.80471207e-03
 -4.49769862e-02 -1.94237867e-02 -1.16780938e-02 -1.67750825e-02
 -4.39032508e-03 -1.28942018e-02 -1.05610830e-03 -1.71408675e-02
 -5.65895423e-02 -1.12548938e-02 -7.20617325e-03 -9.91231120e-03
 -4.91428478e-03 -1.03253444e-02 -3.16667837e-02 -1.71700603e-02
 -2.19343908e-03 -5.94791972e-03 -4.62433824e-02 -7.25013281e-03
 -3.09109763e-04 -5.00221374e-03 -2.00848686e-02 -2.13602548e-02
 -4.99566193e-03 -8.92124070e-02 -2.82730304e-01 -2.72574560e-03
 -7.75026361e-04 -2.53598527e-02 -5.42035493e-03 -1.84991281e-02
 -4.49718365e-01 -1.98444555e-03 -8.94594162e-03 -2.00825054e-03
 -1.54284350e-02 -1.30318286e-02 -9.78276630e-03 -4.58061253e-03
 -1.82606298e-02 -4.51666421e-03 -1.63270519e-01 -1.00997507e-01
 -1.61865790e-02 -1.38341889e-04 -2.76501809e-02 -1.01368503e-02
 -8.35901706e-03 -1.77528926e-02 -1.31717155e-02 -1.95986433e-02]

Example: Providing an Initial Guess for a State Variable#

BalanceComp has a guess_func option that can be used to supply an initial guess value for the state variables. This option provides the same functionality as the guess_nonlinear method of ImplicitComponent.

The Kepler example script shows how guess_func can be used.

prob = om.Problem()

bal = om.BalanceComp()

bal.add_balance(name='E', val=0.0, units='rad', eq_units='rad', rhs_name='M')

# Use M (mean anomaly) as the initial guess for E (eccentric anomaly)
def guess_function(inputs, outputs, residuals):
    if np.abs(residuals['E']) > 1.0E-2:
        outputs['E'] = inputs['M']

bal.options['guess_func'] = guess_function

# ExecComp used to compute the LHS of Kepler's equation.
lhs_comp = om.ExecComp('lhs=E - ecc * sin(E)',
                       lhs={'val': 0.0, 'units': 'rad'},
                       E={'val': 0.0, 'units': 'rad'},
                       ecc={'val': 0.0})

prob.model.add_subsystem(name='balance', subsys=bal,
                         promotes_inputs=['M'],
                         promotes_outputs=['E'])

prob.model.set_input_defaults('M', 85.0, units='deg')

prob.model.add_subsystem(name='lhs_comp', subsys=lhs_comp,
                         promotes_inputs=['E', 'ecc'])

# Explicit connections
prob.model.connect('lhs_comp.lhs', 'balance.lhs:E')

# Set up solvers
prob.model.linear_solver = om.DirectSolver()
prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=2)

prob.setup()

prob.set_val('ecc', 0.6)

prob.run_model()

print(np.degrees(prob.get_val('E')))

NL: Newton 0 ; 1.30402513 1
NL: Newton 1 ; 0.117134648 0.0898254536
NL: Newton 2 ; 0.00208780933 0.00160104992
NL: Newton 3 ; 7.36738647e-07 5.64972738e-07
NL: Newton 4 ; 9.19264664e-14 7.04943981e-14
NL: Newton Converged
[115.91942563]