text:
import itertools
import unittest
import math
import os
import shutil
import tempfile

import numpy as np
from numpy.testing import assert_almost_equal
import scipy
from io import StringIO

from distutils.version import LooseVersion

try:
    from parameterized import parameterized
except ImportError:
    from openmdao.utils.assert_utils import SkipParameterized as parameterized

import openmdao.api as om
from openmdao.components.exec_comp import _expr_dict, _temporary_expr_dict
from openmdao.utils.assert_utils import assert_near_equal, assert_check_partials, assert_warning
from openmdao.utils.om_warnings import OMDeprecationWarning

_ufunc_test_data = {
    'min': {
        'str': 'f=min(x)',
        'check_func': np.min,
        'args': { 'f': {'val': np.zeros(6)},
                  'x': {'val': np.random.random(6)}}},
    'max': {
        'str': 'f=max(x)',
        'check_func': np.max,
        'args': { 'f': {'val': np.zeros(6)},
                  'x': {'val': np.random.random(6)}}},
    'diff': {
        'str': 'f=diff(x)',
        'check_func': np.diff,
        'args': { 'f': {'val': np.zeros(5)},
                  'x': {'val': np.random.random(6)}}},
    'abs': {
        'str': 'f=abs(x)',
        'check_func': np.abs,
        'args': { 'f': {'val': np.zeros(6)},
                  'x': {'val': np.random.random(6)}}},
    'acos': {
        'str': 'f=acos(x)',
        'check_func': np.arccos,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) - 0.5}}},
    'arccos': {
        'str': 'f=arccos(x)',
        'check_func': np.arccos,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) - 0.5}}},
    'arccosh': {
        'str': 'f=arccosh(x)',
        'check_func': np.arccosh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': 1.1 + np.random.random(6)}}},
    'acosh': {
        'str': 'f=acosh(x)',
        'check_func': np.arccosh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': 1.1 + np.random.random(6)}}},
    'arange': {
        'str': 'f=arange(0,10,2)',
        'check_val': np.arange(0, 10, 2),
        'args': {'f': {'val': np.zeros(5)}}},
    'arcsin': {
        'str': 'f=arcsin(x)',
        'check_func': np.arcsin,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) - .5}}},
    'arcsinh': {
        'str': 'f=arcsinh(x)',
        'check_func': np.arcsinh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'asinh': {
        'str': 'f=asinh(x)',
        'check_func': np.arcsinh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'asin': {
        'str': 'f=asin(x)',
        'check_func': np.arcsin,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) - .5}}},
    'arctan': {
        'str': 'f=arctan(x)',
        'check_func': np.arctan,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'arctan2': {
        'str': 'f=arctan2(y, x)',
        'check_val': np.array([-2.35619449, -0.78539816,  0.78539816,  2.35619449]),
        'args': {'f': {'val': np.zeros(4)},
                 'x': {'val': np.array([-1, +1, +1, -1])},
                 'y': {'val': np.array([-1, -1, +1, +1])}}},
    'atan': {
        'str': 'f=atan(x)',
        'check_func': np.arctan,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'cos': {
        'str': 'f=cos(x)',
        'check_func': np.cos,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'cosh': {
        'str': 'f=cosh(x)',
        'check_func': np.cosh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'dot': {
        'str': 'f=dot(x, y)',
        'check_func': np.dot,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'e': {
        'str': 'f=e',
        'check_val': np.e,
        'args': {'f': {'val': 0.0}}},
    'erf': {
        'str': 'f=erf(x)',
        'check_func': scipy.special.erf,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'erfc': {
        'str': 'f=erfc(x)',
        'check_func': scipy.special.erfc,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'exp': {
        'str': 'f=exp(x)',
        'check_func': np.exp,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'expm1': {
        'str': 'f=expm1(x)',
        'check_func': np.expm1,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'fmax': {
        'str': 'f=fmax(x, y)',
        'check_func': np.fmax,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'fmin': {
        'str': 'f=fmin(x, y)',
        'check_func': np.fmin,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'inner': {
        'str': 'f=inner(x, y)',
        'check_func': np.inner,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'isinf': {
        'str': 'f=isinf(x)',
        'check_func': np.isinf,
        'args': {'f': {'val': np.zeros(3)},
                 'x': {'val': [0, np.inf, 5.0]}}},
    'isnan': {
        'str': 'f=isnan(x)',
        'check_func': np.isnan,
        'args': {'f': {'val': np.zeros(3)},
                 'x': {'val': [0, np.nan, np.nan]}}},
    'kron': {
        'str': 'f=kron(x, y)',
        'check_func': np.kron,
        'args': {'f': {'val': np.zeros(36)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'linspace': {
        'str': 'f=linspace(0,10,50)',
        'check_val': np.linspace(0, 10, 50),
        'args': {'f': {'val': np.zeros(50)}}},
    'log': {
        'str': 'f=log(x)',
        'check_func': np.log,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) + 0.1}}},
    'log10': {
        'str': 'f=log10(x)',
        'check_func': np.log10,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6) + 0.1}}},
    'log1p': {
        'str': 'f=log1p(x)',
        'check_func': np.log1p,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'matmul': {
        'str': 'f=matmul(x, y)',
        'check_func': np.matmul,
        'args': {'f': {'val': np.zeros((3, 1))},
                 'x': {'val': np.random.random((3, 3))},
                 'y': {'val': np.random.random((3, 1))}}},
    'maximum': {
        'str': 'f=maximum(x, y)',
        'check_func': np.maximum,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'minimum': {
        'str': 'f=minimum(x, y)',
        'check_func': np.minimum,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'ones': {
        'str': 'f=ones(21)',
        'check_val': np.ones(21),
         'args': {'f': {'val': np.zeros(21)}}},
    'outer': {
        'str': 'f=outer(x, y)',
        'check_func': np.outer,
        'args': {'f': {'val': np.zeros((6, 6))},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6)}}},
    'pi': {
        'str': 'f=pi',
        'check_val': np.pi,
        'args': {'f': {'val': 0.0}}},
    'power': {
        'str': 'f=power(x, y)',
        'check_func': np.power,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)},
                 'y': {'val': np.random.random(6) + 1.0}}},
    'prod': {
        'str': 'f=prod(x)',
        'check_func': np.prod,
        'args': {'f': {'val': 0.0},
                 'x': {'val': np.random.random(6)}}},
    'sin': {
        'str': 'f=sin(x)',
        'check_func': np.sin,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'sinh': {
        'str': 'f=sinh(x)',
        'check_func': np.sinh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'sum': {
        'str': 'f=sum(x)',
        'check_func': np.sum,
        'args': {'f': {'val': 0.0},
                 'x': {'val': np.random.random(6)}}},
    'tan': {
        'str': 'f=tan(x)',
        'check_func': np.tan,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'tanh': {
        'str': 'f=tanh(x)',
        'check_func': np.tanh,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}}},
    'tensordot': {
        'str': 'f=tensordot(x, y)',
        'check_func': np.tensordot,
        'args': {'f': {'val': 0.0},
                 'x': {'val': np.random.random((6, 6))},
                 'y': {'val': np.random.random((6, 6))}}},
    'zeros': {
        'str': 'f=zeros(21)',
        'check_val': np.zeros(21),
        'args': {'f': {'val': np.zeros(21)}}},
}

# 'factorial' will raise a RuntimeError or a deprecation warning depending on scipy version
if LooseVersion(scipy.__version__) >= LooseVersion("1.5.0"):
    _ufunc_test_data['factorial'] = {
        'str': 'f=factorial(x)',
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}},
        'error': (RuntimeError,
                  "The 'factorial' function is not supported for SciPy "
                  f"versions >= 1.5, current version: {scipy.__version__}")
    }
else:
    _ufunc_test_data['factorial'] = {
        'str': 'f=factorial(x)',
        'check_func': scipy.special.factorial,
        'args': {'f': {'val': np.zeros(6)},
                 'x': {'val': np.random.random(6)}},
        'warning': (OMDeprecationWarning,
                    "The 'factorial' function is deprecated. "
                    "It is no longer supported for SciPy versions >= 1.5.")
    }

class TestExecComp(unittest.TestCase):

    def test_missing_partial_warn(self):
        p = om.Problem()
        model = p.model
        comp = om.ExecComp('z=3.0*x + 2.5*y')
        model.add_subsystem('comp', comp)
        comp.declare_partials('z', 'x', method='fd')
        p.setup()
        with assert_warning(UserWarning, "'comp' <class ExecComp>: The following partial derivatives have not been declared so they are assumed to be zero: ['z' wrt 'y']."):
            p.final_setup()

    def test_colon_vars(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=foo:bar+1.'))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: failed to compile expression 'y=foo:bar+1.'.")

    def test_bad_kwargs(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=x+1.', xx=2.0))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: arg 'xx' in call to ExecComp() does not refer to any variable "
                         "in the expressions ['y=x+1.']")

    def test_bad_kwargs_meta(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=x+1.',
                                                   x={'val': 2, 'low': 0, 'high': 10, 'units': 'ft'}))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: the following metadata names were not recognized for "
                         "variable 'x': ['high', 'low']")

    def test_name_collision_const(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('e=x+1.'))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: cannot assign to variable 'e' because it's already defined "
                         "as an internal function or constant.")

    def test_name_collision_func(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('sin=x+1.'))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: cannot assign to variable 'sin' because it's already defined "
                         "as an internal function or constant.")

    def test_func_as_rhs_var(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=sin+1.'))
        with self.assertRaises(Exception) as context:
            prob.setup()
        self.assertEqual(str(context.exception),
                         "'C1' <class ExecComp>: cannot use 'sin' as a variable because it's already defined "
                         "as an internal function or constant.")

    def test_mixed_type(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=sum(x)',
                                                        x=np.arange(10, dtype=float)))
        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 45.0, 0.00001)

    def test_simple(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x+1.', x=2.0))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 3.0, 0.00001)

    def test_for_spaces(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y = pi * x', x=2.0))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)
        self.assertTrue('pi' not in C1._inputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 2 * math.pi, 0.00001)

    def test_units(self):
        prob = om.Problem()
        prob.model.add_subsystem('indep', om.IndepVarComp('x', 100.0, units='cm'))
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x+z+1.',
                                                        x={'val': 2.0, 'units': 'm'},
                                                        y={'units': 'm'},
                                                        z=2.0))
        prob.model.connect('indep.x', 'C1.x')

        prob.setup()

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 4.0, 0.00001)

    def test_units_varname(self):
        prob = om.Problem()

        with self.assertRaises(TypeError) as cm:
            prob.model.add_subsystem('C1', om.ExecComp('y=x+units+1.',
                                                       x={'val': 2.0, 'units': 'm'},
                                                       y={'units': 'm'},
                                                       units=2.0))

        self.assertEqual(str(cm.exception),
                         "ExecComp: Value (2.0) of option 'units' has type 'float', "
                         "but type 'str' was expected.")

    def test_units_varname_str(self):
        prob = om.Problem()

        with self.assertRaises(ValueError) as cm:
            prob.model.add_subsystem('C1', om.ExecComp('y=x+units+1.',
                                                       x={'val': 2.0, 'units': 'm'},
                                                       y={'units': 'm'},
                                                       units='two'))

        self.assertEqual(str(cm.exception), "The units 'two' are invalid.")

    def test_units_varname_novalue(self):
        prob = om.Problem()
        prob.model.add_subsystem('indep', om.IndepVarComp('x', 100.0, units='cm'))
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x+units+1.',
                                                        x={'val': 2.0, 'units': 'm'},
                                                        y={'units': 'm'}))
        prob.model.connect('indep.x', 'C1.x')

        with self.assertRaises(NameError) as cm:
            prob.setup()

        self.assertEqual(str(cm.exception),
                         "'C1' <class ExecComp>: cannot use variable name 'units' because it's a reserved keyword.")

    def test_common_units(self):
        # all variables in the ExecComp have the same units
        prob = om.Problem()

        prob.model.add_subsystem('indep', om.IndepVarComp('x', 100.0, units='cm'))
        prob.model.add_subsystem('comp', om.ExecComp('y=x+z+1.', units='m',
                                                     x={'val': 2.0},
                                                     z=2.0))
        prob.model.connect('indep.x', 'comp.x')

        prob.setup()
        prob.run_model()

        assert_near_equal(prob['comp.y'], 4.0, 0.00001)

    def test_common_units_no_meta(self):
        # make sure common units are assigned when no metadata is provided
        prob = om.Problem()

        prob.model.add_subsystem('indep', om.IndepVarComp('x', 2.0, units='km'))
        prob.model.add_subsystem('comp', om.ExecComp('y = x+1', units='m'))

        prob.model.connect('indep.x', 'comp.x')

        prob.setup()
        prob.run_model()

        assert_near_equal(prob['comp.y'], 2001., 0.00001)

    def test_conflicting_units(self):
        prob = om.Problem()
        prob.model.add_subsystem('indep', om.IndepVarComp('x', 100.0, units='cm'))
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x+z+1.', units='m',
                                                        x={'val': 2.0, 'units': 'km'},
                                                        z=2.0))
        prob.model.connect('indep.x', 'C1.x')

        with self.assertRaises(RuntimeError) as cm:
            prob.setup()

        self.assertEqual(str(cm.exception),
                         "'C1' <class ExecComp>: units of 'km' have been specified for variable 'x', but "
                         "units of 'm' have been specified for the entire component.")

    def test_shape_and_value(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5',
                                                x={'shape': (5,), 'val': np.zeros(5)},
                                                y={'shape': (5,), 'val': np.zeros(5)}))

        p.setup()
        p.run_model()

        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')

        assert_almost_equal(J, np.eye(5)*3., decimal=6)

    def test_conflicting_shape(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5',
                                                x={'shape': (5,), 'val': 5},
                                                y={'shape': (5,)}))

        with self.assertRaises(Exception) as context:
            p.setup()

        self.assertEqual(str(context.exception).replace('L,', ','),  # L on Windows
                         "'comp' <class ExecComp>: shape of (5,) has been specified for variable 'x', "
                         "but a value of shape (1,) has been provided.")

    def test_common_shape(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5', shape=(5,)))

        p.setup()
        p.run_model()

        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')

        assert_almost_equal(J, np.eye(5)*3., decimal=6)

    def test_common_shape_with_values(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5', shape=(5,),
                                                x={'val': np.zeros(5)},
                                                y={'val': np.zeros(5)}))

        p.setup()
        p.run_model()

        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')

        assert_almost_equal(J, np.eye(5)*3., decimal=6)

    def test_common_shape_conflicting_shape(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5', shape=(5,),
                                                y={'shape': (10,)}))

        with self.assertRaises(Exception) as context:
            p.setup()

        self.assertEqual(str(context.exception).replace('L,', ','),  # L on Windows
                         "'comp' <class ExecComp>: shape of (10,) has been specified for variable 'y', "
                         "but shape of (5,) has been specified for the entire component.")

    def test_common_shape_conflicting_value(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5', shape=(5,),
                                                x={'val': 5}))

        with self.assertRaises(Exception) as context:
            p.setup()

        self.assertEqual(str(context.exception).replace('1L,', '1,'),  # 1L on Windows
                         "'comp' <class ExecComp>: value of shape (1,) has been specified for variable 'x', "
                         "but shape of (5,) has been specified for the entire component.")

    def test_math(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=sin(x)', x=2.0))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], math.sin(2.0), 0.00001)

    def test_np(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp(['y1=sin(x)', 'y2=np.cos(x)'], x=2.0))
        prob.setup()

        with self.assertRaises(Exception) as context:
          prob.run_model()

        self.assertEqual(str(context.exception),
            "'C1' <class ExecComp>: Error occurred evaluating 'y2=np.cos(x)':\n"
            "    ExecComp supports a subset of numpy functions directly, without the 'np' prefix.\n"
            "    'cos' is supported (See the documentation).")

    def test_numpy(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=numpy.sin(x)', x=2.0))
        prob.setup()

        with self.assertRaises(Exception) as context:
          prob.run_model()

        self.assertEqual(str(context.exception),
            "'C1' <class ExecComp>: Error occurred evaluating 'y=numpy.sin(x)':\n"
            "    ExecComp supports a subset of numpy functions directly, without the 'numpy' prefix.\n"
            "    'sin' is supported (See the documentation).")

    def test_numpy_fft(self):
        prob = om.Problem()
        prob.model.add_subsystem('C1', om.ExecComp('y=numpy.fft(x)', x=2.0))
        prob.setup()

        with self.assertRaises(Exception) as context:
          prob.run_model()

        self.assertEqual(str(context.exception),
            "'C1' <class ExecComp>: Error occurred evaluating 'y=numpy.fft(x)':\n"
            "    ExecComp supports a subset of numpy functions directly, without the 'numpy' prefix.\n"
            "    'fft' is not supported (See the documentation).")

    def test_array(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x[1]',
                                                        x=np.array([1., 2., 3.]),
                                                        y=0.0))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 2.0, 0.00001)

    def test_array_lhs(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp(['y[0]=x[1]', 'y[1]=x[0]'],
                                                        x=np.array([1., 2., 3.]),
                                                        y=np.array([0., 0.])))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], np.array([2., 1.]), 0.00001)

    def test_simple_array_model(self):
        prob = om.Problem()
        prob.model.add_subsystem('comp', om.ExecComp(['y[0]=2.0*x[0]+7.0*x[1]',
                                                      'y[1]=5.0*x[0]-3.0*x[1]'],
                                                     x=np.zeros([2]), y=np.zeros([2])))

        prob.setup()
        prob.set_solver_print(level=0)
        prob.run_model()

        data = prob.check_partials(out_stream=None)

        assert_check_partials(data, atol=1e-5, rtol=1e-5)

    def test_simple_array_model2(self):
        prob = om.Problem()
        prob.model.add_subsystem('comp', om.ExecComp('y = mat.dot(x)',
                                                     x=np.zeros((2,)), y=np.zeros((2,)),
                                                     mat=np.array([[2., 7.], [5., -3.]])))

        prob.setup()
        prob.set_solver_print(level=0)
        prob.run_model()

        data = prob.check_partials(out_stream=None)

        assert_check_partials(data, atol=1e-5, rtol=1e-5)

    def test_complex_step(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp(['y=2.0*x+1.'], x=2.0))

        prob.setup()

        # Conclude setup but don't run model.
        prob.final_setup()

        self.assertTrue('x' in C1._inputs)
        self.assertTrue('y' in C1._outputs)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 5.0, 0.00001)

        C1._linearize()

        assert_near_equal(C1._jacobian[('y', 'x')], [[2.0]], 0.00001)

    def test_complex_step2(self):
        prob = om.Problem(om.Group())
        prob.model.add_subsystem('comp', om.ExecComp('y=x*x + x*2.0', x=2.0))
        prob.set_solver_print(level=0)

        prob.setup(check=False, mode='fwd')
        prob.run_model()

        J = prob.compute_totals(['comp.y'], ['comp.x'], return_format='flat_dict')
        assert_near_equal(J['comp.y', 'comp.x'], np.array([[6.0]]), 0.00001)

        prob.setup(check=False, mode='rev')
        prob.run_model()

        J = prob.compute_totals(['comp.y'], ['comp.x'], return_format='flat_dict')
        assert_near_equal(J['comp.y', 'comp.x'], np.array([[6.0]]), 0.00001)

    def test_abs_complex_step(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=2.0*abs(x)', x=-2.0))

        prob.setup()
        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], 4.0, 0.00001)

        # any positive C1.x should give a 2.0 derivative for dy/dx
        C1._inputs['x'] = 1.0e-10
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], [[2.0]], 0.00001)

        C1._inputs['x'] = -3.0
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], [[-2.0]], 0.00001)

        C1._inputs['x'] = 0.0
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], [[2.0]], 0.00001)

    def test_abs_complex_step(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=2.0*arctan2(y, x)', x=np.array([1+2j]), y=1))

        prob.setup()
        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], np.array([1.57079633]), 1e-8)

    def test_abs_array_complex_step(self):
        prob = om.Problem()
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=2.0*abs(x)',
                                                        x=np.ones(3)*-2.0, y=np.zeros(3)))

        prob.setup()
        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(C1._outputs['y'], np.ones(3)*4.0, 0.00001)

        # any positive C1.x should give a 2.0 derivative for dy/dx
        C1._inputs['x'] = np.ones(3)*1.0e-10
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], np.eye(3)*2.0, 0.00001)

        C1._inputs['x'] = np.ones(3)*-3.0
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], np.eye(3)*-2.0, 0.00001)

        C1._inputs['x'] = np.zeros(3)
        C1._linearize()
        assert_near_equal(C1._jacobian['y', 'x'], np.eye(3)*2.0, 0.00001)

        C1._inputs['x'] = np.array([1.5, -0.6, 2.4])
        C1._linearize()
        expect = np.zeros((3, 3))
        expect[0, 0] = 2.0
        expect[1, 1] = -2.0
        expect[2, 2] = 2.0

        assert_near_equal(C1._jacobian['y', 'x'], expect, 0.00001)

    def test_has_diag_partials_error(self):
        p = om.Problem()
        model = p.model

        mat = np.arange(15).reshape((3,5))
        model.add_subsystem('comp', om.ExecComp('y=A.dot(x)', has_diag_partials=True, A=mat,
                                                x=np.ones(5), y=np.ones(3)))

        p.setup()

        with self.assertRaises(Exception) as context:
            p.final_setup()

        self.assertEqual(str(context.exception),
                         "'comp' <class ExecComp>: has_diag_partials is True but partial(y, A) is not square (shape=(3, 15)).")

    def test_has_diag_partials(self):
        # Really check to see that the has_diag_partials argument had its intended effect

        # run with has_diag_partials=False
        p = om.Problem()
        model = p.model
        comp = om.ExecComp('y=3.0*x + 2.5', has_diag_partials=False, x=np.ones(5), y=np.ones(5))
        model.add_subsystem('comp', comp)
        p.setup()

        declared_partials = comp._declared_partials[('y','x')]
        self.assertTrue('rows' not in declared_partials )
        self.assertTrue('cols' not in declared_partials )

        # run with has_diag_partials=True
        p = om.Problem()
        model = p.model
        comp = om.ExecComp('y=3.0*x + 2.5', has_diag_partials=True, x=np.ones(5), y=np.ones(5))
        model.add_subsystem('comp', comp)
        p.setup()
        p.final_setup()

        declared_partials = comp._declared_partials[('y','x')]
        self.assertTrue('rows' in declared_partials )
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y','x')]['rows']))
        self.assertTrue('cols' in declared_partials )
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y','x')]['cols']))

    def test_exec_comp_deriv_sparsity(self):
        # Check to make sure that when an ExecComp has more than one
        # expression that only the partials that are needed are declared and computed

        # with has_diag_partials set to the default of False and just scalars
        p = om.Problem()
        model = p.model
        comp = om.ExecComp(['y1=2.0*x1+1.', 'y2=3.0*x2-1.'],x1=1.0, x2=2.0)
        model.add_subsystem('comp', comp)
        p.setup()
        p.final_setup()

        # make sure only the partials that are needed are declared
        declared_partials = comp._declared_partials
        self.assertListEqual( sorted([('y1', 'x1'), ('y2', 'x2') ]),
                              sorted(declared_partials.keys()))

        p.run_model()

        # make sure only what is needed was computed
        subjacs_info = comp._jacobian._subjacs_info
        self.assertListEqual(sorted([('comp.y1', 'comp.x1'), ('comp.y2', 'comp.x2'),
                                     ('comp.y1', 'comp.y1'),('comp.y2', 'comp.y2')]),
                             sorted(subjacs_info.keys()))

        # make sure the result of compute_partials is correct
        J = p.compute_totals(of=['comp.y1'], wrt=['comp.x1'], return_format='array')
        self.assertEqual(2.0, J)
        J = p.compute_totals(of=['comp.y2'], wrt=['comp.x2'], return_format='array')
        self.assertEqual(3.0, J)

        # make sure this works with arrays and when has_diag_partials is the default of False
        p = om.Problem()
        model = p.model
        comp = om.ExecComp(['y1=2.0*x1+1.', 'y2=3.0*x2-1.'],
                           x1=np.ones(5), y1=np.ones(5), x2=np.ones(5), y2=np.ones(5))
        model.add_subsystem('comp', comp)
        p.setup()
        p.final_setup()

        declared_partials = comp._declared_partials
        self.assertListEqual( sorted([('y1', 'x1'), ('y2', 'x2') ]),
                              sorted(declared_partials.keys()))

        p.run_model()
        J = p.compute_totals(of=['comp.y1'], wrt=['comp.x1'], return_format='array')
        self.assertTrue(np.all(2.0*np.identity(5) == J))
        J = p.compute_totals(of=['comp.y2'], wrt=['comp.x2'], return_format='array')
        self.assertTrue(np.all(3.0*np.identity(5) == J))

        # with has_diag_partials True to make sure that still works with arrays
        p = om.Problem()
        model = p.model
        comp = om.ExecComp(['y1=2.0*x1+1.', 'y2=3.0*x2-1.'], has_diag_partials=True,
                           x1=np.ones(5), y1=np.ones(5), x2=np.ones(5), y2=np.ones(5) )
        model.add_subsystem('comp', comp)
        p.setup()
        p.final_setup()

        declared_partials = comp._declared_partials
        self.assertListEqual( sorted([('y1', 'x1'), ('y2', 'x2') ]),
                              sorted(declared_partials.keys()))
        self.assertTrue('cols' in declared_partials[('y1', 'x1')] )
        self.assertTrue('rows' in declared_partials[('y1', 'x1')] )
        self.assertTrue('cols' in declared_partials[('y2', 'x2')] )
        self.assertTrue('rows' in declared_partials[('y2', 'x2')] )
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y1','x1')]['rows']))
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y1','x1')]['cols']))
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y2','x2')]['rows']))
        self.assertListEqual([0,1,2,3,4], list( comp._declared_partials[('y2','x2')]['cols']))

        p.run_model()

        J = p.compute_totals(of=['comp.y1'], wrt=['comp.x1'], return_format='array')
        self.assertTrue(np.all(2.0*np.identity(5) == J))
        J = p.compute_totals(of=['comp.y2'], wrt=['comp.x2'], return_format='array')
        self.assertTrue(np.all(3.0*np.identity(5) == J))

    def test_has_diag_partials_shape_only(self):
        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5', has_diag_partials=True,
                                                x={'shape': (5,)}, y={'shape': (5,)}))

        p.setup()
        p.run_model()

        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')

        assert_almost_equal(J, np.eye(5)*3., decimal=6)

    def test_tags(self):
        prob = om.Problem(model=om.Group())
        C1 = prob.model.add_subsystem('C1', om.ExecComp('y=x+z+1.',
                                                     x={'val': 1.0, 'units': 'm', 'tags': 'tagx'},
                                                     y={'units': 'm', 'tags': ['tagy','tagq']},
                                                     z={'val': 2.0, 'tags': 'tagz'},
                                                     ))

        prob.setup(check=False)

        prob.set_solver_print(level=0)
        prob.run_model()

        # Inputs no tags
        inputs = prob.model.list_inputs(values=False, out_stream=None)
        self.assertEqual(sorted(inputs), [
            ('C1.x', {}),
            ('C1.z', {}),
        ])

        # Inputs with tags
        inputs = prob.model.list_inputs(values=False, out_stream=None, tags="tagx")
        self.assertEqual(sorted(inputs), [
            ('C1.x', {}),
        ])

        # Inputs with multiple tags
        inputs = prob.model.list_inputs(values=False, out_stream=None, tags=["tagx", "tagz"])
        self.assertEqual(sorted(inputs), [
            ('C1.x', {}),
            ('C1.z', {}),
        ])

        # Inputs with tag that does not match
        inputs = prob.model.list_inputs(values=False, out_stream=None, tags="tag_wrong")
        self.assertEqual(sorted(inputs), [])

        # Outputs no tags
        outputs = prob.model.list_outputs(values=False, out_stream=None)
        self.assertEqual(sorted(outputs), [
            ('C1.y', {}),
        ])

        # Outputs with tags
        outputs = prob.model.list_outputs(values=False, out_stream=None, tags="tagy")
        self.assertEqual(sorted(outputs), [
            ('C1.y', {}),
        ])

        # Outputs with multiple tags
        outputs = prob.model.list_outputs(values=False, out_stream=None, tags=["tagy", "tagx"])
        self.assertEqual(sorted(outputs), [
            ('C1.y', {}),
        ])

        # Outputs with tag that does not match
        outputs = prob.model.list_outputs(values=False, out_stream=None, tags="tag_wrong")
        self.assertEqual(sorted(outputs), [])

    def test_feature_has_diag_partials(self):

        p = om.Problem()
        model = p.model

        model.add_subsystem('comp', om.ExecComp('y=3.0*x + 2.5',
                                                has_diag_partials=True,
                                                x=np.ones(5), y=np.ones(5)))

        p.setup()

        p.set_val('comp.x', np.ones(5))

        p.run_model()

        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')

        assert_almost_equal(J, np.eye(5)*3., decimal=6)

    def test_feature_multi_output(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp(['y1=x+1.', 'y2=x-1.']), promotes=['x'])

        prob.setup()

        prob.set_val('x', 2.0)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob.get_val('comp.y1'), 3.0, 0.00001)
        assert_near_equal(prob.get_val('comp.y2'), 1.0, 0.00001)

    def test_feature_multi_output2(self):
        # verify that expressions can have multiple LHS variables.

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp(['y1, y2 = x+1., x-1.']), promotes=['x'])

        prob.setup()

        prob.set_val('x', 2.0)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob.get_val('comp.y1'), 3.0, 0.00001)
        assert_near_equal(prob.get_val('comp.y2'), 1.0, 0.00001)

    def test_feature_array(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp('y=x[1]',
                                                x=np.array([1., 2., 3.]),
                                                y=0.0))

        prob.setup()

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob.get_val('comp.y'), 2.0, 0.00001)

    def test_feature_math(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp('z = sin(x)**2 + cos(y)**2'))

        prob.setup()

        prob.set_val('comp.x', np.pi/2.0)
        prob.set_val('comp.y', np.pi/2.0)

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob.get_val('comp.z'), 1.0, 0.00001)

    def test_feature_numpy(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp('y=sum(x)', x=np.array([1., 2., 3.])))

        prob.setup()

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob['comp.y'], 6.0, 0.00001)

    def test_feature_metadata(self):

        prob = om.Problem()
        model = prob.model

        model.add_subsystem('comp', om.ExecComp('z=x+y',
                                                x={'val': 0.0, 'units': 'inch'},
                                                y={'val': 0.0, 'units': 'inch'},
                                                z={'val': 0.0, 'units': 'inch'}))

        prob.setup()

        prob.set_val('comp.x', 12.0, units='inch')
        prob.set_val('comp.y', 1.0, units='ft')

        prob.set_solver_print(level=0)
        prob.run_model()

        assert_near_equal(prob.get_val('comp.z'), 24.0, 0.00001)

    def test_feature_options(self):

        model = om.Group()

        xcomp = model.add_subsystem('comp', om.ExecComp('y=2*x', shape=(2,)))

        xcomp.options['units'] = 'm'

        prob = om.Problem(model)
        prob.setup()

        prob.set_val('comp.x', [100., 200.], units='cm')

        prob.run_model()

        assert_near_equal(prob.get_val('comp.y'), [2., 4.], 0.00001)

    def test_list_outputs_resids_tol(self):
        prob = om.Problem()
        model = prob.model

        model.add_subsystem(
            "quad_1",
            om.ExecComp(
                "y = a * x ** 2 + b * x + c",
                a={"value": 2.0},
                b={"value": 5.0},
                c={"value": 3.0},
                x={"shape": (2,)},
                y={"shape": (2,)},
            ),
        )

        balance = model.add_subsystem("balance", om.BalanceComp())
        balance.add_balance("x_1", val=np.array([1, -1]), rhs_val=np.array([0., 0.]))
        model.connect("balance.x_1", "quad_1.x")
        model.connect("quad_1.y", "balance.lhs:x_1")

        prob.model.linear_solver = om.ScipyKrylov()
        prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=2)

        prob.setup()
        prob.model.nonlinear_solver.options["maxiter"] = 0
        prob.run_model()

        stream = StringIO()
        outputs = prob.model.list_outputs(residuals=True, residuals_tol=1e-5, out_stream=stream)

        text = stream.getvalue()
        self.assertTrue("balance" in text)
        self.assertTrue("x_1" in text)

    def test_add_expr(self):
        p = om.Problem()

        excomp = om.ExecComp('y=x',
                             x={'val' : 3.0, 'units' : 'mm'},
                             y={'shape' : (1, ), 'units' : 'cm'})

        excomp.add_expr('z = 2.9*x',
                        z={'shape' : (1, ), 'units' : 's'})

        p.model.add_subsystem('comp', excomp, promotes=['*'])
        p.setup()
        p.run_model()

        assert_almost_equal(p.get_val('z'), 8.7, 1e-8)
        assert_almost_equal(p.get_val('y'), 3.0, 1e-8)

    def test_add_expr_bare(self):
        p = om.Problem()

        excomp = om.ExecComp()

        excomp.add_expr('z = 2.9*x',
                        x={'val' : 3.0, 'units' : 'mm'},
                        z={'shape' : (1, ), 'units' : 's'})

        p.model.add_subsystem('comp', excomp, promotes=['*'])
        p.setup()
        p.run_model()

        assert_almost_equal(p.get_val('z'), 8.7, 1e-8)

    def test_add_expr_configure(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp('y=x',
                                     x={'val' : 3.0, 'units' : 'mm'},
                                     y={'shape' : (1, ), 'units' : 'cm'})

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                self.excomp.add_expr('z = 2.9*x',
                                     z={'shape' : (1, ), 'units' : 's'})

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.run_model()

        assert_almost_equal(p.get_val('z'), 8.7, 1e-8)
        assert_almost_equal(p.get_val('y'), 3.0, 1e-8)

    def test_add_expr_configure_no_shape_or_val(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp('x_constraint=x',
                                     x_constraint={'units': None},
                                     x={'units': 's'})

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                self.excomp.add_expr('y_constraint = y',
                                     y_constraint={'units': None},
                                     y={'units': 's'})

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.set_val('x', 3.0)
        p.set_val('y', 4.0)
        p.run_model()

        assert_almost_equal(p.get_val('x_constraint'), 3.0)
        assert_almost_equal(p.get_val('y_constraint'), 4.0)

    def test_no_expressions_until_configure(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp()

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                self.excomp.add_expr('y_constraint = y**2',
                                     y_constraint={'units': None},
                                     y={'units': 's'})

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.set_val('y', 4.0)
        p.run_model()

        assert_almost_equal(p.get_val('y_constraint'), 16.0)

    def test_no_expressions(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp()

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                pass

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.run_model()

    def test_add_expr_configure_delay_defaults(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp('y=x',
                                     y={'shape' : (1, ), 'units' : 'cm'})

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                self.excomp.add_expr('z = 2.9*x',
                                     x={'val' : 3.0, 'units' : 'mm'},
                                     z={'shape' : (1, ), 'units' : 's'})

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.run_model()

        assert_almost_equal(p.get_val('z'), 8.7, 1e-8)
        assert_almost_equal(p.get_val('y'), 3.0, 1e-8)

    def test_add_expr_errors(self):
        p = om.Problem()

        excomp = om.ExecComp('y=x',
                             x={'val' : 3.0, 'units' : 'mm'},
                             y={'shape' : (1, ), 'units' : 'cm'})

        with self.assertRaises(NameError) as cm:
            excomp.add_expr('z = 2.9*x',
                            x={'val' : 3.0, 'units' : 'cm'},
                            z={'shape' : (1, ), 'units' : 's'})

        self.assertEquals(cm.exception.args[0],
                          "Defaults for 'x' have already been defined in a previous "
                          "expression.")

        with self.assertRaises(TypeError) as cm:
            excomp.add_expr(p)

        self.assertEquals(cm.exception.args[0],
                          "Argument 'expr' must be of type 'str', but type 'Problem' was found.")

        excomp.add_expr('y = 2.9*x')
        p.model.add_subsystem('zzz', excomp)
        with self.assertRaises(RuntimeError) as cm:
            p.setup()

        self.assertEquals(cm.exception.args[0],
                          "'zzz' <class ExecComp>: The output 'y' has already been defined by an expression.")

    def test_value_deprecation(self):
        p = om.Problem()

        msg = ("'zzz' <class ExecComp>: 'value' will be deprecated in 4.0. Please use 'val' in the future.")

        excomp = om.ExecComp('y=x**2', x={'value': np.ones(10)}, y={'val': np.ones(10)})

        p.model.add_subsystem('zzz', excomp)
        with assert_warning(OMDeprecationWarning, msg):
            p.setup()

    def test_val_value_error(self):
        p = om.Problem()

        msg = ("Cannot use 'val' and 'value' at the same time, use 'val'.")

        excomp = om.ExecComp('y=x**2', x={'value': np.ones(10), 'val': np.ones(10)},
                                       y={'val': np.ones(10)})

        p.model.add_subsystem('zzz', excomp)

        with self.assertRaises(RuntimeError) as cm:
            p.setup()
            self.assertEquals(cm.exception.args[0], msg)

    def test_feature_add_expr(self):

        class ConfigGroup(om.Group):
            def setup(self):
                excomp = om.ExecComp('y=x',
                                     x={'val' : 3.0, 'units' : 'mm'},
                                     y={'shape' : (1, ), 'units' : 'cm'})

                self.add_subsystem('excomp', excomp, promotes=['*'])

            def configure(self):
                self.excomp.add_expr('z = 2.9*x',
                                     z={'shape' : (1, ), 'units' : 's'})

        p = om.Problem()
        p.model.add_subsystem('sub', ConfigGroup(), promotes=['*'])
        p.setup()
        p.run_model()

        assert_almost_equal(p.get_val('z'), 8.7, 1e-8)
        assert_almost_equal(p.get_val('y'), 3.0, 1e-8)

class TestFunctionRegistration(unittest.TestCase):

    # These 2 tests don't run normally unless you run testflo with a -m "featuretest_*"
    # I don't want to include them in our normal tests because they don't contain
    # the _temporary_expr_dict contextmanager (to avoid user confusion)
    # which keeps the contents of the ExecComp _expr_dict clean
    def featuretest_register_simple(self):

        om.ExecComp.register("myfunc", lambda x: x * x, complex_safe=True)
        p = om.Problem()
        comp = p.model.add_subsystem("comp", om.ExecComp("y = 2 * myfunc(x)"))

        p.setup()
        p.run_model()
        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'])
        assert_near_equal(J['comp.y', 'comp.x'][0][0], 4., 1e-10)

    def featuretest_register_simple_unsafe(self):

        # the following function isn't really complex unsafe, but we'll call it unsafe anyway
        # for demonstration purposes
        om.ExecComp.register("unsafe", lambda x: x * x, complex_safe=False)
        p = om.Problem()
        comp = p.model.add_subsystem("comp", om.ExecComp("y = 2 * unsafe(x)"))

        # because our function is complex unsafe, we must declare that the partials
        # with respect to 'x' use 'fd' instead of 'cs'
        comp.declare_partials('*', 'x', method='fd')

        p.setup()
        p.run_model()
        J = p.compute_totals(of=['comp.y'], wrt=['comp.x'])
        assert_near_equal(J['comp.y', 'comp.x'][0][0], 4., 1e-5)

    def test_register_simple(self):
        with _temporary_expr_dict():
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)'))
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            self.assertEqual(p['comp.area_square'], 9.)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            self.assertEqual(J['comp.area_square', 'comp.x'], 6.)

    def test_register_simple_arr(self):
        with _temporary_expr_dict():
            size = 5
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            p.setup()
            p['comp.x'] = x = np.arange(1, size+1, dtype=float)
            p.run_model()
            assert_near_equal(p['comp.area_square'], x * x, 1e-11)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * x * 2., 1e-11)

    def test_register_check_partials_not_safe(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=False)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            comp.declare_partials('*', '*', method='fd')
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-6)

            data = p.check_partials(out_stream=None, step=1e-7)
            self.assertEqual(list(data), ['comp'])

    def test_register_check_partials_not_safe_err(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=False)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            p.setup()
            p['comp.x'] = 3.
            # calling run_model should NOT raise an exception
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-6)

            with self.assertRaises(Exception) as cm:
                data = p.check_partials(out_stream=None)
            self.assertEquals(cm.exception.args[0],
                              "'comp' <class ExecComp>: expression contains functions ['area'] that are not complex safe. To fix this, call declare_partials('*', ['x'], method='fd') on this component prior to setup.")

    def test_register_check_partials_not_safe_mult_expr(self):
        with _temporary_expr_dict():
            size = 5
            om.ExecComp.register('unsafe', lambda x: x**2, complex_safe=False)
            om.ExecComp.register('safe', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp(['out1 = unsafe(x) * z',
                                                              'out2 = safe(y) + z'], shape=size))
            rows = cols = np.arange(size)
            comp.declare_partials('out1', ['x', 'z'], rows=rows, cols=cols, method='fd')
            comp.declare_partials('out2', ['y'], rows=rows, cols=cols, method='cs')
            comp.declare_partials('out2', ['z'], rows=rows, cols=cols, method='fd')
            p.setup()
            xx = np.arange(1, size + 1, dtype=float)
            p['comp.x'] = x = xx * 3.
            p['comp.y'] = y = xx * 4.
            p['comp.z'] = z = xx * 5.
            p.run_model()
            assert_near_equal(p['comp.out1'], x * x * z, 1e-10)
            assert_near_equal(p['comp.out2'], y * y + z, 1e-10)

            J = p.compute_totals(of=['comp.out1', 'comp.out2'], wrt=['comp.x', 'comp.y', 'comp.z'])
            assert_near_equal(J['comp.out1', 'comp.x'], np.eye(size) * 2. * x * z, 1e-6)
            assert_near_equal(J['comp.out1', 'comp.y'], np.zeros((size, size)), 1e-6)
            assert_near_equal(J['comp.out1', 'comp.z'], np.eye(size) * x**2, 1e-6)
            assert_near_equal(J['comp.out2', 'comp.x'], np.zeros((size, size)), 1e-11)
            assert_near_equal(J['comp.out2', 'comp.y'], np.eye(size) * 2. * y, 1e-11)
            assert_near_equal(J['comp.out2', 'comp.z'], np.eye(size), 1e-6)

            data = p.check_partials(out_stream=None, step=1e-7)
            self.assertEqual(list(data), ['comp'])

    def test_register_check_partials_not_safe_mult_expr_err(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('unsafe', lambda x: x**2, complex_safe=False)
            om.ExecComp.register('safe', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp(['out1 = unsafe(x) * z',
                                                              'out2 = safe(y) + z'], shape=size))
            comp.declare_partials('*', ['x'], method='fd')
            p.setup()
            p['comp.x'] = 3.
            p['comp.y'] = 4.
            p['comp.z'] = 5.
            p.run_model()
            assert_near_equal(p['comp.out1'], np.ones(size) * 45., 1e-6)
            assert_near_equal(p['comp.out2'], np.ones(size) * 21., 1e-6)

            with self.assertRaises(Exception) as cm:
                data = p.check_partials(out_stream=None, step=1e-7)
            self.assertEquals(cm.exception.args[0],
                              "'comp' <class ExecComp>: expression contains functions ['unsafe'] that are not complex safe. To fix this, call declare_partials('*', ['z'], method='fd') on this component prior to setup.")

    def test_register_check_partials_safe(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-11)

            data = p.check_partials(out_stream=None)
            self.assertEqual(list(data), [])

    def test_register_simple_arr_manual_partials_cs(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            comp.declare_partials('area_square', 'x', method='cs')
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-11)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * 6., 1e-11)

    def test_register_simple_arr_manual_partials_fd(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=False)
            p = om.Problem()
            comp = p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size))
            comp.declare_partials('area_square', 'x', method='fd')
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-6)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * 6., 1e-6)

    def test_register_simple_arr_diag(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size, has_diag_partials=True))
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-11)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * 6., 1e-11)

            # verify diagonal subjac
            self.assertTrue(np.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['rows'] == np.arange(size)))
            self.assertTrue(np.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['cols'] == np.arange(size)))

    def test_register_shape_by_conn(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('part', lambda x: x[2:], complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('indeps', om.IndepVarComp('x', np.ones(size)))
            p.model.add_subsystem('comp', om.ExecComp('y = part(x) * 3.', shape_by_conn=True))
            p.model.add_subsystem('sink', om.ExecComp('y=x', shape=size-2))
            p.model.connect("indeps.x", "comp.x")
            p.model.connect("comp.y", "sink.x")
            p.setup()
            p['indeps.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.x'], np.ones(size) * 3., 1e-11)
            assert_near_equal(p['comp.y'], np.ones(size-2) * 9., 1e-11)
            assert_near_equal(p['sink.y'], np.ones(size-2) * 9., 1e-11)
            J = p.compute_totals(of=['sink.y'], wrt=['indeps.x'])
            assert_near_equal(J['sink.y', 'indeps.x'], np.eye(size)[2:, :] * 3., 1e-11)

    def test_register_shape_by_conn_err(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('double', lambda x: x * 2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('indeps', om.IndepVarComp('x', np.ones(size)))
            p.model.add_subsystem('comp', om.ExecComp('y = double(x) * 3.', shape_by_conn=True))
            p.model.add_subsystem('sink', om.ExecComp('y=x', shape=size-2))
            p.model.connect("indeps.x", "comp.x")
            p.model.connect("comp.y", "sink.x")
            p.setup()
            p['indeps.x'] = 3.

            # if shape_by_conn results in inputs and outputs being different shapes, we can't detect if that's
            # an error or not because we don't know how the expressions evaluate, so we have to wait until
            # compute runs and just report as an error during expression evaluation.

            # have to use regex to handle differences in numpy print formats for shape
            msg = "'comp' <class ExecComp>: Error occurred evaluating 'y = double\(x\) \* 3\.':\n" \
                  "'comp' <class ExecComp>: Failed to set value of 'y': could not broadcast " \
                  "input array from shape \(10.*\) into shape \(8.*\)."
            with self.assertRaisesRegex(Exception, msg) as cm:
                p.run_model()

    def test_shape_by_conn_bug_has_diag_partials_bug(self):
        # this is for a bug where has_diag_partials was being ignored when shape_by_conn
        # and/or copy_shape was used.

        prob = om.Problem()

        size = 100000
        t = np.linspace(0, 1, size)

        prob.model.add_subsystem("ivc",
                                om.IndepVarComp("t", val=t),
                                promotes_outputs=["*"])

        comp1 = prob.model.add_subsystem("comp1", om.ExecComp(["x = t + 1"],
                t={"shape_by_conn":True},
                x={"shape_by_conn":True, "copy_shape":"t"},
                has_diag_partials=True), promotes=["*"])

        comp2 = prob.model.add_subsystem("comp2", om.ExecComp(["y = t + 2"],
                t={"shape":size},
                y={"copy_shape":"t"},
                has_diag_partials=True), promotes=["*"])

        comp3 = prob.model.add_subsystem("comp3", om.ExecComp(["z = t + 3"],
                t={"shape_by_conn":True},
                z={"shape":size},
                has_diag_partials=True), promotes=["*"])

        comp4 = prob.model.add_subsystem("comp4", om.ExecComp(["w = t + 4"],
                t={"shape":size},
                w={"shape":size},
                has_diag_partials=True), promotes=["*"])

        prob.setup()
        prob.final_setup()

        # all subjac values should be size == size from above instead of (size, size)
        self.assertEqual(comp1._subjacs_info[('comp1.x', 'comp1.t')]['val'].size, size)
        self.assertEqual(comp2._subjacs_info[('comp2.y', 'comp2.t')]['val'].size, size)
        self.assertEqual(comp3._subjacs_info[('comp3.z', 'comp3.t')]['val'].size, size)
        self.assertEqual(comp4._subjacs_info[('comp4.w', 'comp4.t')]['val'].size, size)

    def test_register_err_keyword(self):
        with _temporary_expr_dict():
            with self.assertRaises(Exception) as cm:
                om.ExecComp.register('shape', lambda x: x, complex_safe=True)
            self.assertEquals(cm.exception.args[0], "ExecComp: cannot register name 'shape' because "
                              "it's a reserved keyword.")

    def test_register_err_not_callable(self):
        with _temporary_expr_dict():
            with self.assertRaises(Exception) as cm:
                om.ExecComp.register('foo', 99, complex_safe=True)
            self.assertEquals(cm.exception.args[0], "ExecComp: 'foo' passed to register() of type 'int' is not callable.")

    def test_register_err_dup(self):
        with _temporary_expr_dict():
            with self.assertRaises(Exception) as cm:
                om.ExecComp.register('exp', lambda x: x, complex_safe=True)
            self.assertEquals(cm.exception.args[0], "ExecComp: 'exp' has already been registered.")

_MASK = np.array(
    [[1, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1],
    [0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1],
    [0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1],
    [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0],
    [0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0],
    [0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1],
    [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1],
    [0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0],
    [0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0],
    [1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0],
    [0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1]]
)

def setup_sparsity(mask):
    sparsity = np.random.random(mask.shape) + 1e-5
    return sparsity * mask

class TestFunctionRegistrationColoring(unittest.TestCase):
    def setUp(self):
        np.random.seed(11)
        self.startdir = os.getcwd()
        self.tempdir = tempfile.mkdtemp(prefix=self.__class__.__name__ + '_')
        os.chdir(self.tempdir)

    def tearDown(self):
        os.chdir(self.startdir)
        try:
            shutil.rmtree(self.tempdir)
        except OSError:
            pass

    def test_coloring(self):
        with _temporary_expr_dict():

            prob = om.Problem(coloring_dir=self.tempdir)
            model = prob.model

            sparsity = setup_sparsity(_MASK)

            def mydot(x):
                return sparsity.dot(x)

            om.ExecComp.register('mydot', mydot, complex_safe=True)

            comp = model.add_subsystem('comp', om.ExecComp('y=mydot(x)',
                                                            x=np.ones(sparsity.shape[1]),
                                                            y=np.ones(sparsity.shape[0])))
            comp.declare_coloring('x', method='cs')

            prob.setup(mode='fwd')
            prob.set_solver_print(level=0)
            prob.run_model()

            J = prob.compute_totals('comp.y', 'comp.x')

            assert_near_equal(J['comp.y', 'comp.x'], sparsity)

            self.assertTrue(np.all(comp._coloring_info['coloring'].get_dense_sparsity() == _MASK))

class TestExecCompParameterized(unittest.TestCase):

    @parameterized.expand(itertools.product([
      func_name for func_name in _expr_dict if func_name not in ('np', 'numpy', '__builtins__')
    ]), name_func=lambda f, n, p: 'test_exec_comp_value_' + '_'.join(a for a in p.args))
    def test_exec_comp_value(self, f):
        test_data = _ufunc_test_data[f]

        prob = om.Problem()
        model = prob.model

        if len(test_data['args']) > 1:
            ivc = model.add_subsystem(name='ivc', subsys=om.IndepVarComp())
            for arg_name, arg_value in test_data['args'].items():
                if arg_name == 'f':
                    continue
                ivc.add_output(name=arg_name, val=arg_value['val'])
                model.connect('ivc.{0}'.format(arg_name), 'comp.{0}'.format(arg_name))

        model.add_subsystem('comp', om.ExecComp(test_data['str'], **test_data['args']),
                            promotes_outputs=['f'])
        prob.setup()

        if 'error' in test_data:
            err, msg = test_data['error']
            with self.assertRaises(err) as cm:
                prob.run_model()
            self.assertTrue(msg in str(cm.exception))
            return
        elif 'warning' in test_data:
            with assert_warning(*test_data['warning']):
                prob.run_model()
        else:
            prob.run_model()

        if 'check_func' in test_data:
            check_args = []
            try:
                check_args.append(test_data['args']['x']['val'])
            except Exception:
                pass
            try:
                check_args.append(test_data['args']['y']['val'])
            except Exception:
                pass
            check_args = tuple(check_args)

            expected = test_data['check_func'](*check_args)
        else:
            expected = test_data['check_val']
        np.testing.assert_almost_equal(prob['f'], expected)

        if 'check_val' not in test_data:
            try:
                prob.check_partials(out_stream=None)
            except TypeError as e:
                print(f, 'does not support complex-step differentiation')

    @parameterized.expand(itertools.product([
      func_name for func_name in _expr_dict if func_name not in ('np', 'numpy', '__builtins__')
    ]), name_func=lambda f, n, p: 'test_exec_comp_jac_' + '_'.join(a for a in p.args))
    def test_exec_comp_jac(self, f):
        test_data = _ufunc_test_data[f]

        prob = om.Problem()
        model = prob.model

        if len(test_data['args']) > 1:
            ivc = model.add_subsystem(name='ivc', subsys=om.IndepVarComp())
            for arg_name, arg_value in test_data['args'].items():
                if arg_name == 'f':
                    continue
                ivc.add_output(name=arg_name, val=arg_value['val'])
                model.connect('ivc.{0}'.format(arg_name),
                              '{0}_comp.{1}'.format(f, arg_name))

        model.add_subsystem('{0}_comp'.format(f),
                            om.ExecComp(test_data['str'], **test_data['args']),
                            promotes_outputs=['f'])
        prob.setup()

        if 'error' in test_data:
            err, msg = test_data['error']
            with self.assertRaises(err) as cm:
                prob.run_model()
            self.assertTrue(msg in str(cm.exception))
            return
        elif 'warning' in test_data:
            with assert_warning(*test_data['warning']):
                prob.run_model()
        else:
            prob.run_model()

        if 'check_val' not in test_data:
            cpd = prob.check_partials(out_stream=None)

            for comp in cpd:
                for (var, wrt) in cpd[comp]:
                    np.testing.assert_almost_equal(cpd[comp][var, wrt]['abs error'][0], 0, decimal=4)

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: eytanadler/OpenMDAO

max_issues_repo_name: eytanadler/OpenMDAO

max_forks_repo_name: eytanadler/OpenMDAO

reason: import numpy,from numpy,import scipy


text:
# Copyright 2019 NREL

# Licensed under the Apache License, Version 2.0 (the "License"); you may not use
# this file except in compliance with the License. You may obtain a copy of the
# License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

# See read the https://floris.readthedocs.io for documentation

import matplotlib.pyplot as plt
import floris.tools as wfct
import floris.tools.visualization as vis
import floris.tools.cut_plane as cp
from floris.tools.optimization import optimize_yaw
from floris.utilities import Vec3
import numpy as np

print('Running FLORIS with no yaw...')
fi = wfct.floris_utilities.FlorisInterface("example_input.json")

# set turbine locations to 3 turbines in a row
D = fi.floris.farm.turbines[0].rotor_diameter
layout_x = [0, 7*D, 14*D]
layout_y = [0, 0, 0]
fi.reinitialize_flow_field(layout_array=(layout_x, layout_y))
fi.calculate_wake()

# initial power output
power_initial = np.sum(fi.get_turbine_power())

# ================================================================================
print('Plotting the FLORIS flowfield...')
# ================================================================================

# Initialize the horizontal cut
hor_plane = wfct.cut_plane.HorPlane(
    fi.get_hub_height_flow_data(),
    fi.floris.farm.turbines[0].hub_height
)

# Plot and show
fig, ax = plt.subplots()
wfct.visualization.visualize_cut_plane(hor_plane, ax=ax)

# ================================================================================
print('Finding optimal yaw angles in FLORIS...')
# ================================================================================
# set bounds for allowable wake steering
min_yaw = 0.0
max_yaw = 25.0
yaw_angles = optimize_yaw(fi, min_yaw, max_yaw)

print('yaw angles = ')
for i in range(len(yaw_angles)):
    print('Turbine ', i, '=', yaw_angles[i], ' deg')

# assign yaw angles to turbines and calculate wake
fi.calculate_wake(yaw_angles=yaw_angles)
power_opt = np.sum(fi.get_turbine_power())

print('==========================================')
print('Total Power Gain = %.1f%%' %
      (100.*(power_opt - power_initial)/power_initial))
print('==========================================')
# ================================================================================
print('Plotting the FLORIS flowfield with yaw...')
# ================================================================================

# Initialize the horizontal cut
hor_plane = wfct.cut_plane.HorPlane(
    fi.get_hub_height_flow_data(),
    fi.floris.farm.turbines[0].hub_height
)

# Plot and show
fig, ax = plt.subplots()
wfct.visualization.visualize_cut_plane(hor_plane, ax=ax)
plt.show()

ext: py

reason: import numpy


text:
import os
import sys
import numpy as np
import pandas as pd

import logging

if '../../' not in sys.path:
    sys.path.append('../../')
import src.optimization as optimization

model = 'lmg'
model_parameters = dict(num_spins=50)
protocol = 'crab'
num_frequencies = 2
optimization_method = 'Powell'

# ------ build and check name for output file
additional_file_name_qualifiers = None
protocol_name = protocol
if protocol == 'crab':
    protocol_name += str(num_frequencies) + 'freq'

output_file_name = (model + '_' + protocol_name + '_' +
                    optimization_method.replace('-', '').lower())
if additional_file_name_qualifiers is not None:
    output_file_name += '_' + additional_file_name_qualifiers
filenum = 1
_output_file_name = output_file_name
while os.path.isfile(_output_file_name + '.csv'):
    _output_file_name = output_file_name + '({:02})'.format(filenum)
    filenum += 1
output_file_name = _output_file_name + '.csv'

# ------ set up logger
logFormatter = logging.Formatter("%(asctime)s [%(threadName)-12.12s]"
                                 "[%(levelname)-5.5s]  %(message)s")
rootLogger = logging.getLogger()
rootLogger.setLevel(logging.DEBUG)
# consoleHandler = logging.StreamHandler()
# consoleHandler.setFormatter(logFormatter)
# rootLogger.addHandler(consoleHandler)
fileHandler = logging.FileHandler(output_file_name[:-4] + '.log')
fileHandler.setFormatter(logFormatter)
fileHandler.setLevel(logging.DEBUG)
rootLogger.addHandler(fileHandler)

logging.info('Output file name will be "{}"'.format(output_file_name))

# ------ start optimization
results = optimization.find_best_protocol(
    problem_specification=dict(
        model=model,
        model_parameters=model_parameters,
        task='critical point state generation'
    ),
    optimization_specs=dict(
        protocol=protocol,
        optimization_method=optimization_method
    ),
    other_options=dict(
        scan_times=np.linspace(0.1, 2, 100)
    )
)

# ------ save results to file
results.to_csv(output_file_name)


text:
import os
from os.path import join as pjoin
import h5py
import cPickle as cp

import numpy as np
import tensorflow as tf
from sampling_utils import *
from acinVAE import  ACIN_VAE
try:
    import PIL.Image as Image
except ImportError:
    import Image

# command line arguments
flags = tf.flags
flags.DEFINE_integer("batchSize", 100, "batch size.")
flags.DEFINE_integer("nEpochs", 1000, "number of epochs to train.")
flags.DEFINE_float("adamLr", 3e-4, "AdaM learning rate.")
flags.DEFINE_integer("hidden_size", 200, "number of hidden units in en/decoder.")
flags.DEFINE_integer("latent_size", 50, "dimensionality of latent variables.")
flags.DEFINE_integer("K", 10, "Maximun number of loops in RNN.")
flags.DEFINE_string("experimentDir", "MNIST/", "directory to save training artifacts.")
inArgs = flags.FLAGS

def get_file_name(expDir, vaeParams, trainParams):
    # concat hyperparameters into file name
    output_file_base_name = '_'+''.join('{}_{}_'.format(key, val) for key, val in sorted(vaeParams.items()) if key not in ['prior', 'input_d'])
    # output_file_base_name += ''.join('{}_{}_'.format(key, "-".join([str(x) for x in vaeParams['prior'][key]])) for key in sorted(['mu', 'sigma']))
    output_file_base_name += '_adamLR_'+str(trainParams['adamLr'])

    # check if results file already exists, if so, append a number
    results_file_name = pjoin(expDir, "train_logs/gaussMM_vae_trainResults"+output_file_base_name+".txt")
    file_exists_counter = 0
    while os.path.isfile(results_file_name):
        file_exists_counter += 1
        results_file_name = pjoin(expDir, "train_logs/gaussMM_vae_trainResults"+output_file_base_name+"_"+str(file_exists_counter)+".txt")
    if file_exists_counter > 0:
        output_file_base_name += "_" + str(file_exists_counter)

    return output_file_base_name

#def trainVAE(data, vae_hyperParams, hyperParams):
def trainVAE(data, vae_hyperParams, hyperParams, param_save_path, logFile=None):
    N_train, d = data['train'].shape
    N_valid, d = data['valid'].shape
    nTrainBatches = N_train / hyperParams['batchSize']
    nValidBatches = N_valid / hyperParams['batchSize']
    vae_hyperParams['batchSize'] = hyperParams['batchSize']

    # init Mix Density VAE
    model = ACIN_VAE(vae_hyperParams)

    # get training op
    optimizer = tf.train.AdamOptimizer(hyperParams['adamLr']).minimize(-model.elbo_obj)

    persister = tf.train.Saver()
    reg_weights = np.linspace(0., 1., int(hyperParams['nEpochs'] * .75)).astype('float32').tolist()
    with tf.Session() as s:
        s.run(tf.initialize_all_variables())

        # for early stopping
        best_elbo = -10000000.
        best_epoch = 0

        for epoch_idx in xrange(hyperParams['nEpochs']):
            # perform update
            if epoch_idx < len(reg_weights):
                w = reg_weights[epoch_idx]
            else:
                w = reg_weights[-1]

            # training
            train_elbo = 0.
            for batch_idx in xrange(nTrainBatches):
                x = data['train'][batch_idx * hyperParams['batchSize']:(batch_idx + 1) * hyperParams['batchSize'],:]
                _, elbo_train = s.run([optimizer, model.elbo_obj], {model.X: x, model.regularization_weight: w})
                train_elbo += elbo_train

            # validation
            valid_elbo = 0.
            for batch_idx in xrange(nValidBatches):
                x = data['valid'][batch_idx * hyperParams['batchSize']:(batch_idx + 1) * hyperParams['batchSize'], :]
                valid_elbo += s.run(model.elbo_obj, {model.X: x,model.regularization_weight: w})

            # check for ELBO improvement
            star_printer = ""
            train_elbo /= nTrainBatches
            valid_elbo /= nValidBatches
            if valid_elbo > best_elbo:
                best_elbo = valid_elbo
                best_epoch = epoch_idx
                star_printer = "***"
                # save the parameters
                persister.save(s, param_save_path)

            # log training progress
            logging_str = "Epoch %d.  Train ELBO: %.3f,  Validation ELBO: %.3f %s" % (epoch_idx + 1, train_elbo, valid_elbo, star_printer)
            print logging_str
            if logFile:
                logFile.write(logging_str + "\n")
                logFile.flush()

            # check for convergence
            if epoch_idx - best_epoch > hyperParams['lookahead_epochs'] or np.isnan(train_elbo): break

    return model

### Marginal Likelihood Calculation
def calc_margLikelihood(data, model, param_file_path, vae_hyperParams, nSamples=50):
    N, d = data.shape

    # get op to load the model
    persister = tf.train.Saver()

    with tf.Session() as s:
        persister.restore(s, param_file_path)

        sample_collector = []
        for s_idx in xrange(nSamples):
            samples = s.run(model.get_log_margLL(), {model.X: data})
            if not np.isnan(samples.mean()) and not np.isinf(samples.mean()):
                sample_collector.append(samples)

    if len(sample_collector) < 1:
        print "\tMARG LIKELIHOOD CALC: No valid samples were collected!"
        return np.nan

    all_samples = np.hstack(sample_collector)
    m = np.amax(all_samples, axis=1)
    mLL = m + np.log(np.mean(np.exp(all_samples - m[np.newaxis].T), axis=1))
    return mLL.mean()

### Sample Images
def sample_from_model(model, param_file_path, vae_hyperParams, image_file_path, nImages=100):
    # get op to load the model
    persister = tf.train.Saver()

    with tf.Session() as s:
        persister.restore(s, param_file_path)
        sample_list = s.run(model.get_samples(nImages))

    for i, samples in enumerate(sample_list):
        image = Image.fromarray(tile_raster_images(X=samples, img_shape=(28, 28), tile_shape=(int(np.sqrt(nImages)), int(np.sqrt(nImages))), tile_spacing=(1, 1)))
        image.save(image_file_path+"_component"+str(i)+".png")

if __name__ == "__main__":
    # mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

    # load MNIST
    f = h5py.File('./binarized_mnist.h5')
    mnist = {'train': np.copy(f['train']), 'valid': np.copy(f['valid']), 'test': np.copy(f['test'])}
    np.random.shuffle(mnist['train'])

    # set architecture params

    vae_hyperParams = {'input_d': mnist['train'].shape[1], 'rnn_hidden_d': inArgs.hidden_size,
                       'hidden_d': inArgs.hidden_size,
                       'z_space_d': inArgs.latent_size, 'K': inArgs.K,
                       'prior': {'mu': [0.] * inArgs.K, 'sigma': [1.] * inArgs.K}}

    # set training hyperparameters
    train_hyperParams = {'adamLr': inArgs.adamLr, 'nEpochs': inArgs.nEpochs, 'batchSize': inArgs.batchSize,
                         'lookahead_epochs': 25}

    # setup files to write results and save parameters
    outfile_base_name = get_file_name(inArgs.experimentDir, vae_hyperParams, train_hyperParams)
    logging_file = open(inArgs.experimentDir+"train_logs/gaussMM_vae_trainResults"+outfile_base_name+".txt", 'w')
    param_file_name = inArgs.experimentDir+"params/gaussMM_vae_params"+outfile_base_name+".ckpt"

    # train
    print "Training model..."
    model = trainVAE(mnist, vae_hyperParams, train_hyperParams, param_file_name, logging_file)

    # evaluate marginal likelihood
    print "Calculating the marginal likelihood..."
    margll_valid = calc_margLikelihood(mnist['valid'], model, param_file_name, vae_hyperParams)
    margll_test = calc_margLikelihood(mnist['test'], model, param_file_name, vae_hyperParams)
    logging_str = "\n\nValidation Marginal Likelihood: %.3f,  Test Marginal Likelihood: %.3f" % (margll_valid, margll_test)
    print logging_str
    logging_file.write(logging_str+"\n")
    logging_file.close()

    # draw some samples
    print "Drawing samples..."
    sample_from_model(model, param_file_name, vae_hyperParams, inArgs.experimentDir+'samples/gaussMM_vae_samples'+outfile_base_name)