# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   orphan: true
# ---

# + tags=["remove-input", "active-ipynb", "remove-output"]
# try:
#     import openmdao.api as om
# except ImportError:
#     !python -m pip install openmdao[notebooks]
#     import openmdao.api as om
# -

# # Vectorizing Linear Solves for Feed-Forward Models
#
# If you have an optimization constraint composed of a large array, or similarly a large array design variable, then there will be one linear solve for each entry of that array. It is possible to speed up the derivative computation by vectorizing the linear solve associated with the design variable or constraint, though the speed up comes at the cost of some additional memory allocation within OpenMDAO.
#
# ```{Note}
# Vectorizing derivatives is only viable for variables/constraints that have a purely feed-forward data path through the model. If there are any solvers in the path between your variable and the objective/constraint of your model then you should not use this feature! See the [Theory Manual on vectorized derivatives](../../../theory_manual/advanced_linear_solvers_special_cases/vectorized) for a detailed explanation of how this feature works.
# ```
#
# You can vectorize derivatives in either `fwd` or `rev` modes. Below is an example of how to do it for `rev` mode, where you specify an argument to `add_constraint`(). See [add_design_var()](../adding_desvars_cons_objs/adding_design_variables) and [add_constraint()](../adding_desvars_cons_objs/adding_constraint) for the full call signature of the relevant methods.

# ## Usage Example

# +
import numpy as np
import openmdao.api as om

SIZE = 5

class ExpensiveAnalysis(om.ExplicitComponent):

    def setup(self):

        self.add_input('x', val=np.ones(SIZE))
        self.add_input('y', val=np.ones(SIZE))

        self.add_output('f', shape=1)

        self.declare_partials('f', 'x')
        self.declare_partials('f', 'y')

    def compute(self, inputs, outputs):

        outputs['f'] = np.sum(inputs['x']**inputs['y'])

    def compute_partials(self, inputs, J):

        J['f', 'x'] = inputs['y']*inputs['x']**(inputs['y']-1)
        J['f', 'y'] = (inputs['x']**inputs['y'])*np.log(inputs['x'])

class CheapConstraint(om.ExplicitComponent):

    def setup(self):

        self.add_input('y', val=np.ones(SIZE))
        self.add_output('g', shape=SIZE)

        row_col = np.arange(SIZE, dtype=int)
        self.declare_partials('g', 'y', rows=row_col, cols=row_col)

        self.limit = 2*np.arange(SIZE)

    def compute(self, inputs, outputs):

        outputs['g'] = inputs['y']**2 - self.limit

    def compute_partials(self, inputs, J):

        J['g', 'y'] = 2*inputs['y']

p = om.Problem()


p.model.set_input_defaults('x', val=2*np.ones(SIZE))
p.model.set_input_defaults('y', val=2*np.ones(SIZE))
p.model.add_subsystem('obj', ExpensiveAnalysis(), promotes=['x', 'y', 'f'])
p.model.add_subsystem('constraint', CheapConstraint(), promotes=['y', 'g'])

p.model.add_design_var('x', lower=.1, upper=10000)
p.model.add_design_var('y', lower=-1000, upper=10000)
p.model.add_constraint('g', upper=0, vectorize_derivs=True)
p.model.add_objective('f')

p.setup(mode='rev')

p.run_model()

p.driver = om.ScipyOptimizeDriver()
p.run_driver()

print(p['x'])
print(p['y'])

# + tags=["remove-input", "remove-output"]
from openmdao.utils.assert_utils import assert_near_equal
import numpy as np

assert_near_equal(p['x'], np.array([0.10000691, 0.1, 0.1, 0.1, 0.1]), tolerance=1e-7)
assert_near_equal(p['y'], np.array([8.19470198e-07, 1.41421356e+00, 2.00000000e+00, 2.44948974e+00, 2.82842712e+00]), tolerance=1e-7)

StarcoderdataJupyterScriptsDedupFiltered