text:
#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES), and is copyright (c) 2018-2021
# by the software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia University
# Research Corporation, et al.  All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and
# license information.
#################################################################################
"""
Tests for FPhx state formulation.
Authors: Andrew Lee
"""

import pytest
from sys import modules

from pyomo.environ import ConcreteModel, Constraint, Var, units as pyunits
from pyomo.util.check_units import (
    check_units_equivalent, assert_units_consistent)

# Need define_default_scaling_factors, even though it is not used directly
from idaes.generic_models.properties.core.state_definitions.FPhx import \
    FPhx, define_state, set_metadata, define_default_scaling_factors
from idaes.core import (MaterialFlowBasis,
                        MaterialBalanceType,
                        EnergyBalanceType,
                        declare_process_block_class)
from idaes.generic_models.properties.core.generic.generic_property import (
        GenericParameterData)
from idaes.generic_models.properties.core.generic.tests.dummy_eos import DummyEoS
from idaes.core.util.testing import PhysicalParameterTestBlock
from idaes.core.util.exceptions import ConfigurationError
import idaes.logger as idaeslog

# Mark module as all unit tests
pytestmark = pytest.mark.unit

@declare_process_block_class("DummyParameterBlock")
class DummyParameterData(GenericParameterData):
    pass

@pytest.mark.unit
def test_set_metadata():
    m = ConcreteModel()
    m.props = PhysicalParameterTestBlock()

    # Set metadata to make sure it is overwritten
    m.props.get_metadata().properties["enth_mol"] = {'method': "test"}

    set_metadata(m.props)

    assert m.props.get_metadata().properties["enth_mol"] == {'method': None}

class TestInvalidBounds(object):

    @pytest.mark.unit
    def test_bad_name(self):
        m = ConcreteModel()

        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K},
                "state_bounds": {"foo": (None, None, None)}})

        with pytest.raises(
                ConfigurationError,
                match="props\[1\] - found unexpected state_bounds key foo. "
                "Please ensure bounds are provided only for expected state "
                "variables and that you have typed the variable names "
                "correctly."):
            # Build state block
            m.props = m.params.build_state_block(
                [1], default={"defined_state": False})

    @pytest.mark.unit
    def test_mole_frac(self, caplog):
        m = ConcreteModel()

        caplog.set_level(
            idaeslog.WARNING,
            logger=("idaes.generic_models.properties.core."))

        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K},
                "state_bounds": {"mole_frac_comp": (None, None, None)}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": False})
        m.props[1].enth_mol_phase["p1"] = 1

        define_state(m.props[1])

        assert ("props[1] - found state_bounds argument for mole_frac_comp."
                " Mole fraction bounds are set automatically and "
                "this argument will be ignored." in caplog.text)

class Test1PhaseDefinedStateFalseNoBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": False})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value is None
        assert check_units_equivalent(frame.props[1].flow_mol,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3
            assert check_units_equivalent(frame.props[1].mole_frac_comp, None)

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value is None
        assert check_units_equivalent(frame.props[1].pressure,
                                      pyunits.Pa)

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value is None
        assert check_units_equivalent(frame.props[1].enth_mol,
                                      pyunits.J/pyunits.mol)

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value is None
        assert check_units_equivalent(frame.props[1].temperature,
                                      pyunits.K)

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 1
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value is None
        assert check_units_equivalent(frame.props[1].flow_mol_phase,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 1
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 1
        assert check_units_equivalent(frame.props[1].phase_frac,
                                      None)

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 3
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3
        assert check_units_equivalent(frame.props[1].mole_frac_phase_comp,
                                      None)

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].total_flow_balance, Constraint)
        assert len(frame.props[1].total_flow_balance) == 1
        assert str(frame.props[1].total_flow_balance.body) == str(
            frame.props[1].flow_mol -
            frame.props[1].flow_mol_phase[frame.params.phase_list[1]])

        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for i in frame.props[1].component_flow_balances:
            assert i in frame.props[1].params.component_list
            assert str(frame.props[1].component_flow_balances[i].body) == str(
                frame.props[1].mole_frac_comp[i] -
                frame.props[1].mole_frac_phase_comp[
                    frame.params.phase_list[1], i])

        assert isinstance(frame.props[1].sum_mole_frac_out, Constraint)
        assert len(frame.props[1].sum_mole_frac_out) == 1
        assert str(frame.props[1].sum_mole_frac_out.body) == str(
                sum(frame.props[1].mole_frac_comp[i]
                    for i in frame.props[1].params.component_list))

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 1
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i])

        assert_units_consistent(frame.props[1])

class Test1PhaseDefinedStateTrueWithBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        # Create a dummy parameter block
        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "state_bounds": {"flow_mol": (0, 100, 200),
                                 "temperature": (290, 345, 400),
                                 "pressure": (1e5, 3e5, 5e5),
                                 "enth_mol": (0, 500, 1000)},
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": True})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value == 100
        assert frame.props[1].flow_mol.lb == 0
        assert frame.props[1].flow_mol.ub == 200

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value == 3e5
        assert frame.props[1].pressure.lb == 1e5
        assert frame.props[1].pressure.ub == 5e5

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value == 500
        assert frame.props[1].enth_mol.lb == 0
        assert frame.props[1].enth_mol.ub == 1000

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value == 345
        assert frame.props[1].temperature.lb == 290
        assert frame.props[1].temperature.ub == 400

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 1
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value == 100
            assert frame.props[1].flow_mol_phase[i].lb == 0
            assert frame.props[1].flow_mol_phase[i].ub == 200

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 1
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 1

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 3
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].total_flow_balance, Constraint)
        assert len(frame.props[1].total_flow_balance) == 1
        assert str(frame.props[1].total_flow_balance.body) == str(
            frame.props[1].flow_mol -
            frame.props[1].flow_mol_phase[frame.params.phase_list[1]])

        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for i in frame.props[1].component_flow_balances:
            assert i in frame.props[1].params.component_list
            assert str(frame.props[1].component_flow_balances[i].body) == str(
                frame.props[1].mole_frac_comp[i] -
                frame.props[1].mole_frac_phase_comp[
                    frame.params.phase_list[1], i])

        assert not hasattr(frame.props[1], "sum_mole_frac_out")

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 1
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i])

        assert_units_consistent(frame.props[1])

class Test2PhaseDefinedStateFalseNoBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        # Create a dummy parameter block
        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS},
                    "p2": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": False})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value is None
        assert check_units_equivalent(frame.props[1].flow_mol,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3
            assert check_units_equivalent(frame.props[1].mole_frac_comp, None)

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value is None
        assert check_units_equivalent(frame.props[1].pressure,
                                      pyunits.Pa)

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value is None
        assert check_units_equivalent(frame.props[1].enth_mol,
                                      pyunits.J/pyunits.mol)

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value is None
        assert check_units_equivalent(frame.props[1].temperature,
                                      pyunits.K)

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 2
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value is None
        assert check_units_equivalent(frame.props[1].flow_mol_phase,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 2
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 1/2
        assert check_units_equivalent(frame.props[1].phase_frac,
                                      None)

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 6
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3"),
                         ("p2", "c1"), ("p2", "c2"), ("p2", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3
        assert check_units_equivalent(frame.props[1].mole_frac_phase_comp,
                                      None)

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].total_flow_balance, Constraint)
        assert len(frame.props[1].total_flow_balance) == 1
        assert str(frame.props[1].total_flow_balance.body) == str(
                sum(frame.props[1].flow_mol_phase[p]
                    for p in frame.props[1].params.phase_list) -
                frame.props[1].flow_mol)

        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for i in frame.props[1].component_flow_balances:
            assert i in frame.props[1].params.component_list
            assert str(frame.props[1].component_flow_balances[i].body) == str(
                frame.props[1].flow_mol * frame.props[1].mole_frac_comp[i] -
                sum(frame.props[1].flow_mol_phase[p] *
                    frame.props[1].mole_frac_phase_comp[p, i]
                    for p in frame.props[1].params.phase_list))

        assert isinstance(frame.props[1].sum_mole_frac, Constraint)
        assert len(frame.props[1].sum_mole_frac) == 1
        assert str(frame.props[1].sum_mole_frac.body) == str(
                sum(frame.props[1].mole_frac_phase_comp[
                    frame.props[1].params.phase_list[1], i]
                    for i in frame.props[1].params.component_list) -
                sum(frame.props[1].mole_frac_phase_comp[
                    frame.props[1].params.phase_list[2], i]
                    for i in frame.props[1].params.component_list))

        assert isinstance(frame.props[1].sum_mole_frac_out, Constraint)
        assert len(frame.props[1].sum_mole_frac_out) == 1
        assert str(frame.props[1].sum_mole_frac_out.body) == str(
                sum(frame.props[1].mole_frac_comp[i]
                    for i in frame.props[1].params.component_list))

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 2
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i]*frame.props[1].flow_mol -
                    frame.props[1].flow_mol_phase[i])

        assert_units_consistent(frame.props[1])

class Test2PhaseDefinedStateTrueWithBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        # Create a dummy parameter block
        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS},
                    "p2": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "state_bounds": {"flow_mol": (0, 100, 200),
                                 "temperature": (290, 345, 400),
                                 "pressure": (1e5, 3e5, 5e5),
                                 "enth_mol": (0, 500, 1000)},
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": True})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value == 100
        assert frame.props[1].flow_mol.lb == 0
        assert frame.props[1].flow_mol.ub == 200

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value == 3e5
        assert frame.props[1].pressure.lb == 1e5
        assert frame.props[1].pressure.ub == 5e5

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value == 500
        assert frame.props[1].enth_mol.lb == 0
        assert frame.props[1].enth_mol.ub == 1000

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value == 345
        assert frame.props[1].temperature.lb == 290
        assert frame.props[1].temperature.ub == 400

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 2
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value == 50
            assert frame.props[1].flow_mol_phase[i].lb == 0
            assert frame.props[1].flow_mol_phase[i].ub == 200

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 2
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 0.5

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 6
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3"),
                         ("p2", "c1"), ("p2", "c2"), ("p2", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].total_flow_balance, Constraint)
        assert len(frame.props[1].total_flow_balance) == 1
        assert str(frame.props[1].total_flow_balance.body) == str(
                sum(frame.props[1].flow_mol_phase[p]
                    for p in frame.props[1].params.phase_list) -
                frame.props[1].flow_mol)

        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for i in frame.props[1].component_flow_balances:
            assert i in frame.props[1].params.component_list
            assert str(frame.props[1].component_flow_balances[i].body) == str(
                frame.props[1].flow_mol * frame.props[1].mole_frac_comp[i] -
                sum(frame.props[1].flow_mol_phase[p] *
                    frame.props[1].mole_frac_phase_comp[p, i]
                    for p in frame.props[1].params.phase_list))

        assert isinstance(frame.props[1].sum_mole_frac, Constraint)
        assert len(frame.props[1].sum_mole_frac) == 1
        assert str(frame.props[1].sum_mole_frac.body) == str(
                sum(frame.props[1].mole_frac_phase_comp[
                    frame.props[1].params.phase_list[1], i]
                    for i in frame.props[1].params.component_list) -
                sum(frame.props[1].mole_frac_phase_comp[
                    frame.props[1].params.phase_list[2], i]
                    for i in frame.props[1].params.component_list))

        assert not hasattr(frame.props[1], "sum_mole_frac_out")

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 2
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i]*frame.props[1].flow_mol -
                    frame.props[1].flow_mol_phase[i])

        assert_units_consistent(frame.props[1])

class Test3PhaseDefinedStateFalseNoBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        # Create a dummy parameter block
        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS},
                    "p2": {"equation_of_state": DummyEoS},
                    "p3": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": False})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value is None
        assert check_units_equivalent(frame.props[1].flow_mol,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3
            assert check_units_equivalent(frame.props[1].mole_frac_comp, None)

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value is None
        assert check_units_equivalent(frame.props[1].pressure,
                                      pyunits.Pa)

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value is None
        assert check_units_equivalent(frame.props[1].enth_mol,
                                      pyunits.J/pyunits.mol)

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value is None
        assert check_units_equivalent(frame.props[1].temperature,
                                      pyunits.K)

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 3
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value is None
        assert check_units_equivalent(frame.props[1].flow_mol_phase,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 3
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 1/3
        assert check_units_equivalent(frame.props[1].phase_frac,
                                      None)

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 9
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3"),
                         ("p2", "c1"), ("p2", "c2"), ("p2", "c3"),
                         ("p3", "c1"), ("p3", "c2"), ("p3", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3
        assert check_units_equivalent(frame.props[1].mole_frac_phase_comp,
                                      None)

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for j in frame.props[1].component_flow_balances:
            assert j in frame.params.component_list
            assert str(frame.props[1].component_flow_balances[j].body) == str(
                frame.props[1].flow_mol*frame.props[1].mole_frac_comp[j] -
                sum(frame.props[1].flow_mol_phase[p] *
                    frame.props[1].mole_frac_phase_comp[p, j]
                    for p in frame.props[1].params.phase_list))

        assert isinstance(frame.props[1].sum_mole_frac, Constraint)
        assert len(frame.props[1].sum_mole_frac) == 3
        for p in frame.props[1].sum_mole_frac:
            assert p in frame.params.phase_list
            assert str(frame.props[1].sum_mole_frac[p].body) == str(
                    sum(frame.props[1].mole_frac_phase_comp[p, i]
                        for i in frame.props[1].params.component_list))

        assert isinstance(frame.props[1].sum_mole_frac_out, Constraint)
        assert len(frame.props[1].sum_mole_frac_out) == 1
        assert str(frame.props[1].sum_mole_frac_out.body) == str(
                sum(frame.props[1].mole_frac_comp[i]
                    for i in frame.props[1].params.component_list))

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 3
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i]*frame.props[1].flow_mol -
                    frame.props[1].flow_mol_phase[i])

        assert_units_consistent(frame.props[1])

class Test3PhaseDefinedStateTrueWithBounds(object):
    # Test define_state method with no bounds and defined_State = False
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        # Create a dummy parameter block
        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "p1": {"equation_of_state": DummyEoS},
                    "p2": {"equation_of_state": DummyEoS},
                    "p3": {"equation_of_state": DummyEoS}},
                "state_definition": modules[__name__],
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "state_bounds": {"flow_mol": (0, 100, 200),
                                 "temperature": (290, 345, 400),
                                 "pressure": (1e5, 3e5, 5e5),
                                 "enth_mol": (0, 500, 1000)},
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": True})

        # Add necessary variables that would be built by other methods
        m.props[1].enth_mol_phase = Var(m.params.phase_list,
                                        initialize=1,
                                        units=pyunits.J/pyunits.mol)

        return m

    @pytest.mark.unit
    def test_always_flash(self, frame):
        define_state(frame.props[1])

        assert frame.props[1].always_flash

    @pytest.mark.unit
    def test_vars(self, frame):
        # Check that all necessary variables have been constructed and have
        # the correct values
        assert isinstance(frame.props[1].flow_mol, Var)
        assert frame.props[1].flow_mol.value == 100
        assert frame.props[1].flow_mol.lb == 0
        assert frame.props[1].flow_mol.ub == 200
        assert check_units_equivalent(frame.props[1].flow_mol,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].mole_frac_comp, Var)
        assert len(frame.props[1].mole_frac_comp) == 3
        for i in frame.props[1].mole_frac_comp:
            assert i in frame.props[1].params.component_list
            assert frame.props[1].mole_frac_comp[i].value == 1/3
            assert check_units_equivalent(frame.props[1].mole_frac_comp, None)

        assert isinstance(frame.props[1].pressure, Var)
        assert frame.props[1].pressure.value == 3e5
        assert frame.props[1].pressure.lb == 1e5
        assert frame.props[1].pressure.ub == 5e5
        assert check_units_equivalent(frame.props[1].pressure,
                                      pyunits.Pa)

        assert isinstance(frame.props[1].enth_mol, Var)
        assert frame.props[1].enth_mol.value == 500
        assert frame.props[1].enth_mol.lb == 0
        assert frame.props[1].enth_mol.ub == 1000
        assert check_units_equivalent(frame.props[1].enth_mol,
                                      pyunits.J/pyunits.mol)

        assert isinstance(frame.props[1].temperature, Var)
        assert frame.props[1].temperature.value == 345
        assert frame.props[1].temperature.lb == 290
        assert frame.props[1].temperature.ub == 400
        assert check_units_equivalent(frame.props[1].temperature,
                                      pyunits.K)

        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 3
        for i in frame.props[1].flow_mol_phase:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].flow_mol_phase[i].value == 100/3
            assert frame.props[1].flow_mol_phase[i].lb == 0
            assert frame.props[1].flow_mol_phase[i].ub == 200
        assert check_units_equivalent(frame.props[1].flow_mol_phase,
                                      pyunits.mol/pyunits.s)

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 3
        for i in frame.props[1].phase_frac:
            assert i in frame.props[1].params.phase_list
            assert frame.props[1].phase_frac[i].value == 1/3
        assert check_units_equivalent(frame.props[1].phase_frac,
                                      None)

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 9
        for i in frame.props[1].mole_frac_phase_comp:
            assert i in [("p1", "c1"), ("p1", "c2"), ("p1", "c3"),
                         ("p2", "c1"), ("p2", "c2"), ("p2", "c3"),
                         ("p3", "c1"), ("p3", "c2"), ("p3", "c3")]
            assert frame.props[1].mole_frac_phase_comp[i].value == 1/3
        assert check_units_equivalent(frame.props[1].mole_frac_phase_comp,
                                      None)

    @pytest.mark.unit
    def test_constraints(self, frame):
        # Check that the correct constraints are present
        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3
        for j in frame.props[1].component_flow_balances:
            assert j in frame.params.component_list
            assert str(frame.props[1].component_flow_balances[j].body) == str(
                frame.props[1].flow_mol*frame.props[1].mole_frac_comp[j] -
                sum(frame.props[1].flow_mol_phase[p] *
                    frame.props[1].mole_frac_phase_comp[p, j]
                    for p in frame.props[1].params.phase_list))

        assert isinstance(frame.props[1].sum_mole_frac, Constraint)
        assert len(frame.props[1].sum_mole_frac) == 3
        for p in frame.props[1].sum_mole_frac:
            assert p in frame.params.phase_list
            assert str(frame.props[1].sum_mole_frac[p].body) == str(
                    sum(frame.props[1].mole_frac_phase_comp[p, i]
                        for i in frame.props[1].params.component_list))

        assert not hasattr(frame.props[1], "sum_mole_frac_out")

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 3
        for i in frame.props[1].phase_fraction_constraint:
            assert i in frame.props[1].params.phase_list
            assert str(frame.props[1].phase_fraction_constraint[i].body) == \
                str(frame.props[1].phase_frac[i]*frame.props[1].flow_mol -
                    frame.props[1].flow_mol_phase[i])

        assert_units_consistent(frame.props[1])

class TestCommon(object):
    @pytest.fixture(scope="class")
    def frame(self):
        m = ConcreteModel()

        m.params = DummyParameterBlock(default={
                "components": {"c1": {}, "c2": {}, "c3": {}},
                "phases": {
                    "a": {"equation_of_state": DummyEoS},
                    "b": {"equation_of_state": DummyEoS}},
                "state_definition": FPhx,
                "pressure_ref": 1e5,
                "temperature_ref": 300,
                "state_bounds": {
                    "flow_mol": (0, 0.1, 0.2, pyunits.kmol/pyunits.s),
                    "temperature": (522, 621, 720, pyunits.degR),
                    "pressure": (1, 3, 5, pyunits.bar),
                    "enth_mol": (0, 0.5, 1, pyunits.kJ/pyunits.mol)},
                "base_units": {"time": pyunits.s,
                               "length": pyunits.m,
                               "mass": pyunits.kg,
                               "amount": pyunits.mol,
                               "temperature": pyunits.K}})

        # Build state block
        m.props = m.params.build_state_block(
            [1], default={"defined_state": False})

        # Add necessary variables that would be built by other methods
        m.props[1].dens_mol_phase = Var(m.params.phase_list, initialize=1)
        m.props[1].enth_mol_phase = Var(m.params.phase_list, initialize=1)

        return m

    @pytest.mark.unit
    def test_convert_vars(self, frame):
        # Check that all state var values and bounds were converted correctly
        assert frame.props[1].flow_mol.value == 100
        assert frame.props[1].flow_mol.lb == 0
        assert frame.props[1].flow_mol.ub == 200
        assert check_units_equivalent(frame.props[1].flow_mol,
                                      pyunits.mol/pyunits.s)

        assert frame.props[1].pressure.value == 3e5
        assert frame.props[1].pressure.lb == 1e5
        assert frame.props[1].pressure.ub == 5e5
        assert check_units_equivalent(frame.props[1].pressure,
                                      pyunits.Pa)

        assert frame.props[1].enth_mol.value == 500
        assert frame.props[1].enth_mol.lb == 0
        assert frame.props[1].enth_mol.ub == 1000
        assert check_units_equivalent(frame.props[1].enth_mol,
                                      pyunits.J/pyunits.mol)

        assert frame.props[1].temperature.value == 345
        assert frame.props[1].temperature.lb == 290
        assert frame.props[1].temperature.ub == 400
        assert check_units_equivalent(frame.props[1].temperature,
                                      pyunits.K)

        # Check supporting variables
        assert isinstance(frame.props[1].flow_mol_phase, Var)
        assert len(frame.props[1].flow_mol_phase) == 2

        assert isinstance(frame.props[1].mole_frac_phase_comp, Var)
        assert len(frame.props[1].mole_frac_phase_comp) == 6

        assert isinstance(frame.props[1].phase_frac, Var)
        assert len(frame.props[1].phase_frac) == 2

        assert isinstance(frame.props[1].total_flow_balance, Constraint)
        assert len(frame.props[1].total_flow_balance) == 1

        assert isinstance(frame.props[1].component_flow_balances, Constraint)
        assert len(frame.props[1].component_flow_balances) == 3

        assert isinstance(frame.props[1].sum_mole_frac, Constraint)
        assert len(frame.props[1].sum_mole_frac) == 1

        assert isinstance(frame.props[1].sum_mole_frac_out, Constraint)

        assert isinstance(frame.props[1].phase_fraction_constraint, Constraint)
        assert len(frame.props[1].phase_fraction_constraint) == 2

    @pytest.mark.unit
    def test_calculate_scaling_factors(self, frame):
        frame.props[1].calculate_scaling_factors()

        assert len(frame.props[1].scaling_factor) == 23

        assert frame.props[1].scaling_factor[
            frame.props[1].dens_mol_phase["a"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].dens_mol_phase["b"]] == 1e-2
        assert frame.props[1].scaling_factor[frame.props[1].enth_mol] == 1e-4
        assert frame.props[1].scaling_factor[frame.props[1].flow_mol] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase["a"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase["b"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["a", "c1"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["a", "c2"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["a", "c3"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["b", "c1"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["b", "c2"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].flow_mol_phase_comp["b", "c3"]] == 1e-2
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_comp["c1"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_comp["c2"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_comp["c3"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["a", "c1"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["a", "c2"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["a", "c3"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["b", "c1"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["b", "c2"]] == 1e3
        assert frame.props[1].scaling_factor[
            frame.props[1].mole_frac_phase_comp["b", "c3"]] == 1e3
        assert frame.props[1].scaling_factor[frame.props[1].pressure] == 1e-5
        assert frame.props[1].scaling_factor[
            frame.props[1].temperature] == 1e-2

    # Test General Methods
    @pytest.mark.unit
    def test_get_material_flow_terms(self, frame):
        for (p, j) in frame.params._phase_component_set:
            assert str(frame.props[1].get_material_flow_terms(p, j)) == str(
                frame.props[1].flow_mol_phase_comp[p, j])

    @pytest.mark.unit
    def test_get_enthalpy_flow_terms(self, frame):
        for p in frame.params.phase_list:
            assert str(frame.props[1].get_enthalpy_flow_terms(p)) == str(
                frame.props[1]._enthalpy_flow_term[p])
            assert str(frame.props[1]._enthalpy_flow_term[p].expr) == str(
                frame.props[1].flow_mol_phase[p] *
                frame.props[1].enth_mol_phase[p])

    @pytest.mark.unit
    def test_get_material_density_terms(self, frame):
        for (p, j) in frame.params._phase_component_set:
            assert str(frame.props[1].get_material_density_terms(p, j)) == str(
                frame.props[1]._material_density_term[p, j])
            assert str(
                frame.props[1]._material_density_term[p, j].expr) == str(
                frame.props[1].dens_mol_phase[p] *
                frame.props[1].mole_frac_phase_comp[p, j])

    @pytest.mark.unit
    def test_get_energy_density_terms(self, frame):
        for p in frame.params.phase_list:
            assert str(frame.props[1].get_energy_density_terms(p)) == str(
                frame.props[1]._energy_density_term[p])
            assert str(
                frame.props[1]._energy_density_term[p].expr) == str(
                frame.props[1].dens_mol_phase[p] *
                frame.props[1].energy_internal_mol_phase[p])

    @pytest.mark.unit
    def test_default_material_balance_type(self, frame):
        assert frame.props[1].default_material_balance_type() == \
            MaterialBalanceType.componentTotal

    @pytest.mark.unit
    def test_default_energy_balance_type(self, frame):
        assert frame.props[1].default_energy_balance_type() == \
            EnergyBalanceType.enthalpyTotal

    @pytest.mark.unit
    def test_get_material_flow_basis(self, frame):
        assert frame.props[1].get_material_flow_basis() == \
            MaterialFlowBasis.molar

    @pytest.mark.unit
    def test_define_state_vars(self, frame):
        assert frame.props[1].define_state_vars() == \
            {"flow_mol": frame.props[1].flow_mol,
             "mole_frac_comp": frame.props[1].mole_frac_comp,
             "enth_mol": frame.props[1].enth_mol,
             "pressure": frame.props[1].pressure}

    @pytest.mark.unit
    def test_define_display_vars(self, frame):
        assert frame.props[1].define_display_vars() == \
            {"Total Molar Flowrate": frame.props[1].flow_mol,
             "Total Mole Fraction": frame.props[1].mole_frac_comp,
             "Molar Enthalpy": frame.props[1].enth_mol,
             "Pressure": frame.props[1].pressure}

max_stars_repo_name: carldlaird/idaes-pse

max_issues_repo_name: carldlaird/idaes-pse

max_forks_repo_name: carldlaird/idaes-pse

reason: from pyomo


text:
# Copyright 2017 reinforce.io. All Rights Reserved.
#
# 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.
# ==============================================================================
"""
Standard Gaussian policy.
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import tensorflow as tf

from tensorforce import util
from tensorforce.core.networks import layers
from tensorforce.core.distributions import Distribution

class Gaussian(Distribution):

    def __init__(self, distribution=None):
        assert distribution is None or len(distribution) == 2
        super(Gaussian, self).__init__(distribution)
        if self.distribution is not None:
            self.mean, self.std_dev = self.distribution

    def create_tf_operations(self, x, sample=True):
        self.mean = tf.squeeze(input=layers['linear'](x=x, size=1), axis=1)
        # self.std_dev = tf.exp(tf.squeeze(input=layers['linear'](x=x, size=1), axis=1))
        # self.std_dev = tf.ones_like(self.mean) * tf.exp(tf.Variable(initial_value=tf.random_normal(shape=(), stddev=0.01)))
        self.std_dev = tf.ones_like(self.mean)
        self.distribution = (self.mean, self.std_dev)
        if sample:
            self.value = self.mean + tf.multiply(x=self.std_dev, y=tf.random_normal(shape=tf.shape(self.mean)))
        else:
            self.value = self.mean

    def log_probability(self, action):
        l2_dist = tf.square(action - self.mean)
        sqr_std_dev = tf.square(self.std_dev)
        log_std_dev = tf.log(self.std_dev + util.epsilon)
        log_prob = -l2_dist / (2 * sqr_std_dev + util.epsilon) - 0.5 * tf.log(tf.constant(2 * np.pi)) - log_std_dev
        return log_prob

    def entropy(self):
        log_std_dev = tf.log(self.std_dev + util.epsilon)
        entropy = tf.reduce_mean(log_std_dev + tf.constant(0.5 * np.log(2 * np.pi * np.e), tf.float32), axis=0)

        return entropy

    def kl_divergence(self, other):
        assert isinstance(other, Gaussian)
        log_std_dev1 = tf.log(self.std_dev + util.epsilon)
        log_std_dev2 = tf.log(other.std_dev + util.epsilon)
        l2_dist = tf.square(self.mean - other.mean)
        sqr_std_dev = tf.square(other.std_dev)
        kl_div = tf.reduce_mean(log_std_dev1 - log_std_dev2 + (self.std_dev + l2_dist) / (2 * sqr_std_dev + util.epsilon) - 0.5, axis=0)
        return kl_div

ext: py

max_stars_repo_name: matthewwilfred/tensorforce

max_issues_repo_name: matthewwilfred/tensorforce

max_forks_repo_name: matthewwilfred/tensorforce

reason: import numpy


text:
import os,cv2
import numpy as np
from torch.utils.data import Dataset
from skimage.transform import warp, AffineTransform
from nltk.corpus import wordnet as wn
import pickle

def random_transform(img):
    if np.random.random() < 0.5:
        img = img[:,::-1,:]

    if np.random.random() < 0.5:
        sx = np.random.uniform(0.7, 1.3)
        sy = np.random.uniform(0.7, 1.3)
    else:
        sx = 1.0
        sy = 1.0

    if np.random.random() < 0.5:
        rx = np.random.uniform(-30.0*2.0*np.pi/360.0,+30.0*2.0*np.pi/360.0)
    else:
        rx = 0.0

    if np.random.random() < 0.5:
        tx = np.random.uniform(-10,10)
        ty = np.random.uniform(-10,10)
    else:
        tx = 0.0
        ty = 0.0

    aftrans = AffineTransform(scale=(sx, sy), rotation=rx, translation=(tx,ty))
    img_aug = warp(img,aftrans.inverse,preserve_range=True).astype('uint8')

    return img_aug

class SketchyDataset(Dataset):
    def __init__(self, split='train',
                 root_dir='../dataset/Sketchy/',
                 version='sketch_tx_000000000000_ready', zero_version='zeroshot1',
                 cid_mask = False, transform=None, aug=False, shuffle=False,
                 first_n_debug=9999999, contrastive_transform=None):

        self.root_dir = root_dir
        self.version = version
        self.split = split

        self.img_dir = self.root_dir

        if self.split == 'train':
            file_ls_file = os.path.join(self.root_dir, zero_version, self.version+'_filelist_train.txt')
        elif self.split == 'val':
            file_ls_file = os.path.join(self.root_dir, zero_version, self.version+'_filelist_test.txt')
        elif self.split == 'zero':
            file_ls_file = os.path.join(self.root_dir, zero_version, self.version+'_filelist_zero.txt')
        else:
            print('unknown split for dataset initialization: ' + self.split)
            return

        with open(file_ls_file, 'r') as fh:
            file_content = fh.readlines()

        self.file_ls = np.array([' '.join(ff.strip().split()[:-1]) for ff in file_content])
        self.labels = np.array([int(ff.strip().split()[-1]) for ff in file_content])
        if shuffle:
            self.shuffle()

        self.file_ls = self.file_ls[:first_n_debug]
        self.labels = self.labels[:first_n_debug]

        self.transform = transform
        self.contrastive_transform = contrastive_transform
        self.aug = aug

        self.cid_mask = cid_mask
        if cid_mask:
            cid_mask_file = os.path.join(self.root_dir, zero_version, 'cid_mask.pickle')
            with open(cid_mask_file, 'rb') as fh:
                self.cid_matrix = pickle.load(fh)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        img = cv2.imread(os.path.join(self.img_dir, self.file_ls[idx]))[:,:,::-1]
        if self.aug and np.random.random()<0.7:
            img = random_transform(img)

        if self.contrastive_transform is not None:
            img1 = self.contrastive_transform(img)
            img2 = self.contrastive_transform(img)

        if self.transform is not None:
            img = self.transform(img)

        label = self.labels[idx]

        if self.cid_mask:
            mask = self.cid_matrix[label]
            if self.contrastive_transform is not None:
                return img, img1, img2, label, mask
            else:
                return img, label, mask

        else:
            if self.contrastive_transform is not None:
                return img, img1, img2, label
            else:
                return img, label

    def shuffle(self):
        s_idx = np.random.shuffle(np.arange(len(self.labels)))
        self.file_ls = self.file_ls[s_idx]
        self.labels = self.labels[s_idx]

class SketchImagePairedDataset(Dataset):
    def __init__(self, dataset='sketchy', zero_version='zeroshot1',
                 cid_mask=False, transform=None, contrastive_transform=None,
                 aug=False, shuffle=False, first_n_debug=9999999):
        if dataset == 'sketchy':
            self.root_dir = '../dataset/Sketchy/'
            self.sketch_version = 'sketch_tx_000000000000_ready'
            self.image_version = 'all_photo'
        else:
            self.root_dir = '../dataset/TUBerlin/'
            self.sketch_version = 'png_ready'
            self.image_version = 'ImageResized_ready'

        self.zero_version = zero_version
        self.transform = transform
        self.contrastive_transform = contrastive_transform
        self.aug = aug

        file_ls_sketch = os.path.join(self.root_dir, self.zero_version, self.sketch_version + '_filelist_test.txt')
        file_ls_image = os.path.join(self.root_dir, self.zero_version, self.image_version + '_filelist_train.txt')
        if cid_mask:
            file_dict_cid = os.path.join(self.root_dir, self.zero_version, 'cid_mask.pickle')
            with open(file_dict_semantic, 'rb') as fh:
                self.cid_dict = pickle.load(fh)

        with open(file_ls_sketch, 'r') as fh:
            file_content_sketch = fh.readlines()
        with open(file_ls_image, 'r') as fh:
            file_content_image = fh.readlines()

        self.file_ls_sketch = np.array([' '.join(ff.strip().split()[:-1]) for ff in file_content_sketch])
        self.labels_sketch = np.array([int(ff.strip().split()[-1]) for ff in file_content_sketch])
        self.names_sketch = np.array([' '.join(ff.strip().split()[:-1]).split('/')[-2] for ff in file_content_sketch])
        self.file_ls_image = np.array([' '.join(ff.strip().split()[:-1]) for ff in file_content_image])
        self.labels_image = np.array([int(ff.strip().split()[-1]) for ff in file_content_image])
        self.names_image = np.array([' '.join(ff.strip().split()[:-1]).split('/')[-2] for ff in file_content_image])

        if shuffle:
            self.shuffle()

        self.file_ls_sketch = self.file_ls_sketch[:first_n_debug]
        self.labels_sketch = self.labels_sketch[:first_n_debug]
        self.names_sketch = self.names_sketch[:first_n_debug]
        self.file_ls_image = self.file_ls_image[:first_n_debug]
        self.labels_image = self.labels_image[:first_n_debug]
        self.names_image = self.names_image[:first_n_debug]

    def __getitem__(self, idx):
        label = self.labels_image[idx]

        select_idx = np.random.choice(np.argwhere(self.labels_sketch == label).reshape(-1), 1)
        sketch = cv2.imread(os.path.join(self.root_dir, self.file_ls_sketch[select_idx][0]))[:, :, ::-1]
        if self.transform is not None:
            sketch = self.transform(sketch)

        image = cv2.imread(os.path.join(self.root_dir, self.file_ls_image[idx]))[:, :, ::-1]
        if self.transform is not None:
            image = self.transform(image)

        return sketch, image, label

    def __len__(self):
        # if self.opt.dataset_name == 'QuickDraw':
        #     return len(self.labels_sketch)
        # else:
        #     return len(self.labels_image)
        return len(self.labels_image)

    def shuffle(self):
        s_idx = np.arange(len(self.labels_sketch))
        np.random.shuffle(s_idx)
        self.file_ls_sketch = self.file_ls_sketch[s_idx]
        self.labels_sketch = self.labels_sketch[s_idx]
        self.names_sketch = self.names_sketch[s_idx]
        s_idx = np.arange(len(self.labels_image))
        np.random.shuffle(s_idx)
        self.file_ls_image = self.file_ls_image[s_idx]
        self.labels_image = self.labels_image[s_idx]
        self.names_image = self.names_image[s_idx]

def wnid_to_synset(wnid):
    return wn.synset_from_pos_and_offset(wnid[0], int(wnid[1:]))


text:
from __future__ import print_function
import fcl
import os
import matplotlib
import matplotlib.pyplot as plt
import time
from os.path import abspath, join, dirname
from fcl import kmeans
from fcl.datasets import load_sector_dataset, load_usps_dataset
from fcl.matrix.csr_matrix import get_csr_matrix_from_object, csr_matrix_to_libsvm_string
from sklearn.decomposition import TruncatedSVD, PCA
from scipy.sparse import csr_matrix
from sklearn.datasets import dump_svmlight_file
import numpy as np

def do_evaluations(dataset_path, dataset_name):

  # Load dataset directly into csr matrix format this way it only needs to be converted once
  data_as_csrmatrix = get_csr_matrix_from_object(dataset_path)

  print("Doing evaluations for dataset %s"%dataset_path)

  algorithm_results = {'bv_kmeans': {}
                       , 'yinyang': {}
                       , 'bv_yinyang': {}
                       , 'elkan': {}
                       , 'pca_kmeans': {}
                       , 'pca_elkan': {}
                       , 'pca_yinyang': {}
                       }

  additional_algo_data =  {}

  calculate_svd = False
  for algo in algorithm_results:
    if "pca" in algo:
      calculate_svd = True
      break

  if calculate_svd:
    p = TruncatedSVD(n_components = int(data_as_csrmatrix.annz * 0.1))
    start = time.time()
    p.fit(data_as_csrmatrix.to_numpy())
    # convert to millis
    fin = (time.time() - start) * 1000
    print(fin)
    pca_projection_csrmatrix = get_csr_matrix_from_object(p.components_)
    for algorithm in algorithm_results:
      if algorithm.startswith("pca_"):
        additional_algo_data[algorithm] = {'data': pca_projection_csrmatrix, 'duration': fin}

  clusters = [100, 250, 1000]

  # Do the evaluations for every algorithm and every no_clusters
  for algorithm in sorted(algorithm_results.keys()):
    for no_clusters in clusters:
      print("Executing k-means with algorithm: %s and k=%d"%(algorithm, no_clusters))
      km = kmeans.KMeans(n_jobs=1, no_clusters=no_clusters, algorithm=algorithm, init='random', seed = 0, verbose = True)
      if algorithm not in additional_algo_data:
        km.fit(data_as_csrmatrix)
      else:
        km.fit(data_as_csrmatrix, external_vectors = additional_algo_data[algorithm]['data'])

      algorithm_results[algorithm][no_clusters] = km.get_tracked_params()
      if algorithm in additional_algo_data:
        algorithm_results[algorithm][no_clusters]['duration_kmeans'] += additional_algo_data[algorithm]['duration']

  # plot the results
  plot_overall_duration(algorithm_results, dataset_name)

def plot_overall_duration(algorithm_results, dataset_name):
  # create subplot showing the overall duration of the algorithm

  fig = plt.figure()
  ax = plt.subplot(111)

  i = 0
  tick_pos = []
  sorted_algorithms = sorted(algorithm_results.keys())
  no_clusters_list = sorted(algorithm_results[sorted_algorithms[0]])
  width = 1 / float(len(algorithm_results) + 1)
  base_tick_pos = (float(width) * len(algorithm_results)) / 2
  tick_pos = []

  algorithm_xvalues = {}
  algorithm_yvalues = {}
  algorithm_legend_data = {}
  for i in range(len(sorted_algorithms)):
    algo_color = next(ax._get_lines.prop_cycler)['color']
    algo_name = sorted_algorithms[i]
    algorithm_legend_data[i] = (algo_color, algo_name)
    algorithm_xvalues[i] = []
    algorithm_yvalues[i] = []

  j = 0
  for no_clusters in no_clusters_list:
    i = 0
    for algorithm in sorted_algorithms:
      algorithm_xvalues[i].append(j + (i * width))
      algorithm_yvalues[i].append(algorithm_results[algorithm][no_clusters]['duration_kmeans'] / 1000)
      i += 1

    tick_pos.append(j + base_tick_pos)
    j += 1

  legend_rects = []
  for i in range(len(sorted_algorithms)):
    (algo_color, algo_name) = algorithm_legend_data[i]
    r = ax.bar(algorithm_xvalues[i], algorithm_yvalues[i], width, color=algo_color)
    legend_rects.append(r[0])

  # Shrink current axis by 20%
  box = ax.get_position()
  ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])

  ax.set_xticks(tick_pos, minor=False)
  ax.set_xticklabels(no_clusters_list, fontdict=None, minor=False)
  ax.set_ylabel('time / s')
  ax.set_xlabel('number of clusters')
  ax.set_title("Overall duration of various algorithms for dataset %s"%dataset_name)

  lgd = ax.legend(legend_rects, sorted_algorithms, loc='center left', bbox_to_anchor=(1, 0.5))

  fig.tight_layout()
  destination_filename = join(dirname( __file__ ), "algorithm_speed_comparison.png")
  plt.savefig(destination_filename, bbox_extra_artists=(lgd,), bbox_inches='tight')
  print("plot was saved in the current folder to: %s"%destination_filename)

if __name__ == "__main__":
  ds_folder = abspath(join(dirname( __file__ ), os.pardir, os.pardir, os.pardir, 'datasets'))
  do_evaluations(load_sector_dataset(ds_folder), 'sector')

max_stars_repo_name: thomas-bottesch/fcl

max_issues_repo_name: thomas-bottesch/fcl

max_forks_repo_name: thomas-bottesch/fcl

reason: import numpy,from scipy


text:
# This file is part of GridCal.
#
# GridCal is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GridCal is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GridCal.  If not, see <http://www.gnu.org/licenses/>.

from numpy import linalg
from scipy.sparse import issparse, csr_matrix as sparse, hstack as hstack_sp, vstack as vstack_sp, diags
from scipy.sparse.linalg import spsolve, splu
import scipy
scipy.ALLOW_THREADS = True
import time
import numpy as np

np.set_printoptions(precision=8, suppress=True, linewidth=320000)

def convert_to_reduced(pq, pv, vd, n):
    """
    Convert the PQ and PV lists of indices to PQ and PV lists of a reduced scheme
    :param pq: Array of PQ indices
    :param pv: Array of PV indices
    :param vd: Array of slack indices
    :param n: number of buses
    :return: Pq, PV reduced arrays
    """
    nsl_counted = np.zeros(n, dtype=int)
    compt = 0
    for i in range(n):
        if i in vd:
            compt += 1
        nsl_counted[i] = compt
    pq_ = pq - nsl_counted[pq]
    pv_ = pv - nsl_counted[pv]
    return pq_, pv_

def ASD1(Ybus, S0, V0, I0, pv, pq, vd, tol, max_it=15, gamma=0.5):
    """
    Alternate Search Directions Power Flow

    As defined in the paper:

        Unified formulation of a family of iterative solvers for power system analysis
        by Domenico Borzacchiello et. Al.

    :param Ybus: Admittance matrix
    :param S0: Power injections
    :param V0: Initial Voltage Vector
    :param I0: Current injections @V=1.0 p.u.
    :param pv: Array of PV bus indices
    :param pq: Array of PQ bus indices
    :param vd: Array of Slack bus indices
    :param tol: Tolerance
    :param max_it: Maximum number of iterations
    :param gamma: Reactive power acceleration factor
    :return: V, converged, normF, Scalc, iterations, elapsed
    """
    start = time.time()

    # initialize
    V = V0.copy()
    Vset_pv = np.abs(V0[pv])
    Scalc = S0.copy()
    Scalc[pv] = Scalc[pv].real + 0j
    normF = 1e20
    iterations = 0
    converged = False
    n = len(V0)
    npq = len(pq)
    npv = len(pv)
    pqpv = np.r_[pq, pv]

    npqpv = len(pqpv)

    # kron
    if npv:
        Y11 = Ybus[pq, :][:, pq]
        Y12 = Ybus[pq, :][:, pv]
        Y21 = Ybus[pv, :][:, pq]
        Y22 = Ybus[pv, :][:, pv]
        Ykron = Y22 + Y21 * (spsolve(Y11, Y12))

        Yred = vstack_sp((hstack_sp((Y11, Y12)), hstack_sp((Y21, Y22))))

    else:
        Yred = Ybus[pq, :][:, pq]

    # reduced system
    Sred = S0[pqpv]
    I0_red = I0[pqpv]

    Yslack = - Ybus[np.ix_(pqpv, vd)]  # yes, it is the negative of this
    Vslack = V0[vd]

    # compute alpha
    Vabs = np.ones(npqpv)
    alpha = sp.diags(np.conj(Sred) / (Vabs * Vabs))

    # compute Y-alpha
    Y_alpha = (Yred - alpha)

    # compute beta as a vector
    B = Y_alpha.diagonal()
    beta = diags(B)
    Y_beta = Yred - beta

    # get the first voltage approximation V0, or simply V(l
    Ivd = Yslack * Vslack  # slack currents
    V_l = spsolve(Y_alpha, Ivd)  # slack voltages influence
    V_l = V0[pqpv]

    while not converged and iterations < max_it:

        iterations += 1

        # Global step
        rhs_global = np.conj(Sred) / np.conj(V_l) - alpha * V_l + I0_red + Ivd
        V_l12 = spsolve(Y_alpha, rhs_global)

        # PV correction
        if npv:
            V_pv = V_l12[npq:]
            V_pv_abs = np.abs(V_pv)
            dV_l12 = (Vset_pv - V_pv_abs) * V_pv / V_pv_abs
            dI_l12 = Ykron * dV_l12
            dQ_l12 = (V_pv * np.conj(dI_l12)).imag  # correct the reactive power
            Sred[npq:] = Sred[npq:].real + 1j * (Sred[npq:].imag + gamma * dQ_l12)  # assign the reactive power
            V_l12[npq:] += dV_l12  # correct the voltage

        # local step
        A = (Y_beta * V_l12 - I0_red - Ivd) / B
        Sigma = - np.conj(Sred) / (A * np.conj(A) * B)
        U = (-1 - np.sqrt(1 - 4 * (Sigma.imag * Sigma.imag + Sigma.real))) / 2.0 + 1j * Sigma.imag
        V_l = U * A

        # Assign the reduced solution
        V[pq] = V_l[:npq]
        V[pv] = V_l[npq:]

        # compute the calculated power injection and the error of the voltage solution
        Scalc = V * np.conj(Ybus * V - I0)
        mis = Scalc - S0  # complex power mismatch
        mismatch = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]  # concatenate again
        normF = np.linalg.norm(mismatch, np.Inf)

        converged = normF < tol

        print(normF)

    end = time.time()
    elapsed = end - start

    return V, converged, normF, Scalc, iterations, elapsed

def ASD2(Ybus, S0, V0, I0, pv, pq, vd, tol, max_it=15, gamma=0.1):
    """
    Alternate Search Directions Power Flow

    As defined in the paper:

        Unified formulation of a family of iterative solvers for power system analysis
        by Domenico Borzacchiello et. Al.

    Modifications:
    - Better reactive power computation

    :param Ybus: Admittance matrix
    :param S0: Power injections
    :param V0: Initial Voltage Vector
    :param I0: Current injections @V=1.0 p.u.
    :param pv: Array of PV bus indices
    :param pq: Array of PQ bus indices
    :param vd: Array of Slack bus indices
    :param tol: Tolerance
    :param max_it: Maximum number of iterations
    :return: V, converged, normF, Scalc, iterations, elapsed
    """
    start = time.time()

    # initialize
    V = V0.copy()
    Vset_pv = np.abs(V0[pv])
    Scalc = S0.copy()
    Scalc[pv] = Scalc[pv].real + 0j
    normF = 1e20
    iterations = 0
    converged = False
    n = len(V0)
    npq = len(pq)
    npv = len(pv)
    pqpv = np.r_[pq, pv]

    npqpv = len(pqpv)

    # kron
    Y11 = Ybus[pq, :][:, pq]
    Y12 = Ybus[pq, :][:, pv]
    Y21 = Ybus[pv, :][:, pq]
    Y22 = Ybus[pv, :][:, pv]

    # reduced system
    Yred = vstack_sp((hstack_sp((Y11, Y12)), hstack_sp((Y21, Y22))))
    Sred = S0[pqpv]
    I0_red = I0[pqpv]

    Yslack = - Ybus[np.ix_(pqpv, vd)]  # yes, it is the negative of this
    Vslack = V0[vd]

    # compute alpha
    Vabs = np.ones(npqpv)
    alpha = sp.diags(np.conj(Sred) / (Vabs * Vabs))

    # compute Y-alpha
    Y_alpha = (Yred - alpha)

    # compute beta as a vector
    B = Y_alpha.diagonal()
    beta = diags(B)
    Y_beta = Yred - beta

    # get the first voltage approximation V0, or simply V(l
    Ivd = Yslack * Vslack  # slack currents
    V_l = spsolve(Y_alpha, Ivd)  # slack voltages influence
    # V_l = V0[pqpv]

    while not converged and iterations < max_it:

        iterations += 1

        # Global step
        rhs_global = np.conj(Sred) / np.conj(V_l) - alpha * V_l + I0_red + Ivd
        V_l12 = spsolve(Y_alpha, rhs_global)

        # PV correction (using Lynn's formulation)
        V_pv = V_l12[npq:]
        V_l12[npq:] = V_pv * Vset_pv / np.abs(V_pv)
        V[pq] = V_l12[:npq]
        V[pv] = V_l12[npq:]
        Qpv = (V[pv] * np.conj((Ybus[pv, :] * V))).imag
        Sred[npq:] = Sred[npq:].real + 1j * gamma * Qpv

        # local step
        A = (Y_beta * V_l12 - I0_red - Ivd) / B
        Sigma = - np.conj(Sred) / (A * np.conj(A) * B)
        U = (-1 - np.sqrt(1 - 4 * (Sigma.imag * Sigma.imag + Sigma.real))) / 2.0 + 1j * Sigma.imag
        V_l = U * A

        # Assign the reduced solution
        V[pq] = V_l[:npq]
        V[pv] = V_l[npq:]

        # compute the calculated power injection and the error of the voltage solution
        Scalc = V * np.conj(Ybus * V - I0)
        mis = Scalc - S0  # complex power mismatch
        mismatch = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]  # concatenate again
        normF = np.linalg.norm(mismatch, np.Inf)

        converged = normF < tol

        print(normF)

    end = time.time()
    elapsed = end - start

    return V, converged, normF, Scalc, iterations, elapsed

def ASD3(Ybus, S0, V0, I0, pv, pq, vd, tol, max_it=15, gamma=0.1):
    """
    Alternate Search Directions Power Flow

    As defined in the paper:

        Unified formulation of a family of iterative solvers for power system analysis
        by Domenico Borzacchiello et. Al.

    Modifications:
    - Better reactive power computation
    - No need to partition the buses in PQ|PV
    - No need to perform the partitioning of Y, nor the Kron reduction for the PV buses correction

    :param Ybus: Admittance matrix
    :param S0: Power injections
    :param V0: Initial Voltage Vector
    :param I0: Current injections @V=1.0 p.u.
    :param pv: Array of PV bus indices
    :param pq: Array of PQ bus indices
    :param vd: Array of Slack bus indices
    :param tol: Tolerance
    :param max_it: Maximum number of iterations
    :return: V, converged, normF, Scalc, iterations, elapsed
    """
    start = time.time()

    # initialize
    V = V0.copy()
    Vset_pv = np.abs(V0[pv])
    Scalc = S0.copy()
    normF = 1e20
    iterations = 0
    converged = False
    n = len(V0)
    pqpv = np.r_[pq, pv]
    pqpv.sort()

    # Create the zero-based indices in the internal reduced scheme
    pq_, pv_ = convert_to_reduced(pq, pv, vd, n)

    # reduced system
    Yred = Ybus[np.ix_(pqpv, pqpv)]
    Sred = S0[pqpv]
    I0_red = I0[pqpv]

    Yslack = - Ybus[np.ix_(pqpv, vd)]  # yes, it is the negative of this
    Vslack = V0[vd]

    # compute alpha
    # Vabs = np.abs(Vslack)
    # alpha = sp.diags(np.conj(Sred) / (Vabs * Vabs))
    alpha = -sp.linalg.inv(sp.diags(Yred.diagonal()).tocsc())

    # compute Y-alpha
    Y_alpha = (Yred - alpha).tocsc()

    # compute beta as a vector
    B = Yred.diagonal()
    beta = diags(B)
    Y_beta = Yred - beta

    # get the first voltage approximation V0, or simply V(l
    Ivd = Yslack * Vslack  # slack currents
    V_l = spsolve(Y_alpha, Ivd)  # slack voltages influence

    while not converged and iterations < max_it:

        # Global step
        rhs_global = np.conj(Sred) / np.conj(V_l) - alpha * V_l + I0_red + Ivd
        V_l12 = spsolve(Y_alpha, rhs_global)

        # Better PV correction than the paper, using Lynn's formulation
        V_pv = V_l12[pv_]
        V_l12[pv_] = V_pv * Vset_pv / np.abs(V_pv)  # compute the corrected bus voltage
        V[pqpv] = V_l12  # Assign the reduced solution

        # compute the reactive power at the PV nodes (actual scheme)
        Qpv = (V[pv] * np.conj((Ybus[pv, :] * V - I0[pv]))).imag
        Sred[pv_] = Sred[pv_].real + 1j * gamma * Qpv  # assign the PV reactive power in the reduced scheme

        # local step
        A = (Y_beta * V_l12 - I0_red - Ivd) / B
        Sigma = - np.conj(Sred) / (A * np.conj(A) * B)
        U = (-1.0 - np.sqrt(1.0 - 4.0 * (Sigma.imag * Sigma.imag + Sigma.real))) / 2.0 + 1.0j * Sigma.imag
        V_l = U * A

        # Assign the reduced solution
        V[pqpv] = V_l

        # compute the calculated power injection and the error of the voltage solution
        Scalc = V * np.conj(Ybus * V - I0)
        mis = Scalc - S0  # complex power mismatch
        mismatch = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]  # concatenate again
        normF = 0.5 * np.dot(mismatch, mismatch)  #np.linalg.norm(mismatch, np.Inf)

        converged = normF < tol

        print(normF)
        iterations += 1

    end = time.time()
    elapsed = end - start

    return V, converged, normF, Scalc, iterations, elapsed

def func(Sred, pv_, x, Qpv, Y_beta, V_l12, I0_red, Ivd, B, V, pqpv, Ybus, I0, S0, pv, pq):
    Sred[pv_] = Sred[pv_].real + 1j * x * Qpv
    # local step
    A = (Y_beta * V_l12 - I0_red - Ivd) / B
    Sigma = - np.conj(Sred) / (A * np.conj(A) * B)
    U = (-1.0 - np.sqrt(1.0 - 4.0 * (Sigma.imag * Sigma.imag + Sigma.real))) / 2.0 + 1.0j * Sigma.imag
    V_l = U * A
    V[pqpv] = V_l
    Scalc = V * np.conj(Ybus * V - I0)
    mis = Scalc - S0  # complex power mismatch
    mismatch = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]  # concatenate again
    f = 0.5 * np.dot(mismatch, mismatch)
    return f

def ASD4(Ybus, S0, V0, I0, pv, pq, vd, tol, max_it=15, gamma=0.1):
    """
    Alternate Search Directions Power Flow

    As defined in the paper:

        Unified formulation of a family of iterative solvers for power system analysis
        by Domenico Borzacchiello et. Al.

    Modifications:
    - Better reactive power computation
    - No need to partition the buses in PQ|PV
    - No need to perform the partitioning of Y, nor the Kron reduction for the PV buses correction
    - Automatic finding of gamma

    :param Ybus: Admittance matrix
    :param S0: Power injections
    :param V0: Initial Voltage Vector
    :param I0: Current injections @V=1.0 p.u.
    :param pv: Array of PV bus indices
    :param pq: Array of PQ bus indices
    :param vd: Array of Slack bus indices
    :param tol: Tolerance
    :param max_it: Maximum number of iterations
    :return: V, converged, normF, Scalc, iterations, elapsed
    """
    start = time.time()

    # initialize
    V = V0.copy()
    Vset_pv = np.abs(V0[pv])
    Scalc = S0.copy()
    normF = 1e20
    iterations = 0
    converged = False
    n = len(V0)
    pqpv = np.r_[pq, pv]
    pqpv.sort()

    # Create the zero-based indices in the internal reduced scheme
    pq_, pv_ = convert_to_reduced(pq, pv, vd, n)

    # reduced system
    Yred = Ybus[np.ix_(pqpv, pqpv)]
    Sred = S0[pqpv]
    I0_red = I0[pqpv]

    Yslack = - Ybus[np.ix_(pqpv, vd)]  # yes, it is the negative of this
    Vslack = V0[vd]

    # compute alpha
    # Vabs = np.abs(Vslack)
    # alpha = sp.diags(np.conj(Sred) / (Vabs * Vabs))
    alpha = -sp.linalg.inv(sp.diags(Yred.diagonal()).tocsc())

    # compute Y-alpha
    Y_alpha = (Yred - alpha).tocsc()

    # compute beta as a vector
    B = Yred.diagonal()
    beta = diags(B)
    Y_beta = Yred - beta

    # get the first voltage approximation V0, or simply V(l
    Ivd = Yslack * Vslack  # slack currents
    V_l = spsolve(Y_alpha, Ivd)  # slack voltages influence

    while not converged and iterations < max_it:

        # Global step
        rhs_global = np.conj(Sred) / np.conj(V_l) - alpha * V_l + I0_red + Ivd
        V_l12 = spsolve(Y_alpha, rhs_global)

        # Better PV correction than the paper, using Lynn's formulation
        V_pv = V_l12[pv_]
        V_l12[pv_] = V_pv * Vset_pv / np.abs(V_pv)  # compute the corrected bus voltage
        V[pqpv] = V_l12  # Assign the reduced solution

        # compute the reactive power at the PV nodes (actual scheme)
        Qpv = (V[pv] * np.conj((Ybus[pv, :] * V - I0[pv]))).imag

        # Find gamma the first time
        # if iterations == 0:
        N = 1
        a = 1.0
        b = 0
        found = False
        c = (a + b) / 2
        while not found or N <= max_it:  # limit iterations to prevent infinite loop
            c = (a + b) / 2  # new midpoint

            f_c = func(Sred.copy(), pv_, c, Qpv, Y_beta, V_l12, I0_red, Ivd, B, V.copy(), pqpv, Ybus, I0, S0, pv, pq)
            f_a = func(Sred.copy(), pv_, a, Qpv, Y_beta, V_l12, I0_red, Ivd, B, V.copy(), pqpv, Ybus, I0, S0, pv, pq)

            sign_fa = f_a < 0
            sign_fc = f_c < 0

            if f_c <= tol or (a - b) / 2 < tol:  # solution found
                found = True

            N += 1  # increment step  counter
            if f_a > f_c:
                b = c
            else:
                a = c  # new interval

        gamma = c

        # compute the reactive power at the PV nodes (actual scheme)
        Sred[pv_] = Sred[pv_].real + 1j * gamma * Qpv  # assign the PV reactive power in the reduced scheme

        # local step
        A = (Y_beta * V_l12 - I0_red - Ivd) / B
        Sigma = - np.conj(Sred) / (A * np.conj(A) * B)
        U = (-1.0 - np.sqrt(1.0 - 4.0 * (Sigma.imag * Sigma.imag + Sigma.real))) / 2.0 + 1.0j * Sigma.imag
        V_l = U * A

        # Assign the reduced solution
        V[pqpv] = V_l

        # compute the calculated power injection and the error of the voltage solution
        Scalc = V * np.conj(Ybus * V - I0)
        mis = Scalc - S0  # complex power mismatch
        mismatch = np.r_[mis[pv].real, mis[pq].real, mis[pq].imag]  # concatenate again
        normF = 0.5 * np.dot(mismatch, mismatch)  # np.linalg.norm(mismatch, np.Inf)

        converged = normF < tol

        print(normF)
        iterations += 1

    end = time.time()
    elapsed = end - start

    return V, converged, normF, Scalc, iterations, elapsed

########################################################################################################################
#  MAIN
########################################################################################################################
if __name__ == "__main__":
    from GridCal.Engine import *

    fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE14_from_raw.gridcal'
    # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE 30 Bus.gridcal'
    # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/Lynn 5 Bus pv.gridcal'
    # fname ='/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/Brazil11_loading05.gridcal'
    # fname = '/home/santi/Documentos/GitHub/GridCal/Grids_and_profiles/grids/IEEE 5 Bus.xlsx'

    grid = FileOpen(fname).open()

    numeric_circuit = compile_snapshot_circuit(grid)
    islands = split_into_islands(numeric_circuit)
    circuit = islands[0]

    print('\nYbus:\n', circuit.Ybus.todense())
    print('\nYseries:\n', circuit.Yseries.todense())
    print('\nYshunt:\n', circuit.Yshunt)
    print('\nSbus:\n', circuit.Sbus)
    print('\nIbus:\n', circuit.Ibus)
    print('\nVbus:\n', circuit.Vbus)
    print('\ntypes:\n', circuit.bus_types)
    print('\npq:\n', circuit.pq)
    print('\npv:\n', circuit.pv)
    print('\nvd:\n', circuit.vd)

    import time
    print('ASD')
    start_time = time.time()

    V1, converged_, err, S, iter_, elapsed_ = ASD1(Ybus=circuit.Ybus,
                                                   S0=circuit.Sbus,
                                                   V0=circuit.Vbus,
                                                   I0=circuit.Ibus,
                                                   pv=circuit.pv,
                                                   pq=circuit.pq,
                                                   vd=circuit.vd,
                                                   tol=1e-9,
                                                   max_it=20,
                                                   gamma=0.1)

    print("--- %s seconds ---" % (time.time() - start_time))
    print('V:\n', np.abs(V1))
    print('error: \t', err)

reason: import numpy,from numpy,import scipy,from scipy


text:
from ..imports import *

def split_punct(tok):
    toks=tokenize_agnostic(tok)
    wordil=[]

    try:
        for i,x in enumerate(toks):
            if any(y.isalpha() for y in x): wordil+=[i]
        pref=''.join([
            x
            for i,x in enumerate(toks)
            if i<wordil[-1]
        ])
        suf=''.join([
            x
            for i,x in enumerate(toks)
            if i>wordil[-1]
        ])
        words=''.join(
            toks[wordil[0] : wordil[-1]+1]
        ) if wordil else ''
        return (pref,words,suf)
    except IndexError:
        return ('',tok,'')

def zero_punc(token):
    return ''.join(x for x in token if x.isalpha())

def tokenize_agnostic(txt):
    return re.findall(r"[\w']+|[.,!?; -—–\n]", txt)

def tokenize_fast(line,lower=False):
	line = line.lower() if lower else line
	import re
	# tokenize using reg ex (fast)
	tokens = re.findall("[A-Z]{2,}(?![a-z])|[A-Z][a-z]+(?=[A-Z])|[\'\w\-]+",line)
	# remove punctuation on either end
	from string import punctuation
	tokens = [tok.strip(punctuation) for tok in tokens]
	# make sure each thing in list isn't empty
	tokens = [tok for tok in tokens if tok]
	return tokens

# tokenize
def tokenize_nltk(txt,lower=False):
	# lowercase
	txt_l = txt.lower() if lower else txt
	# use nltk
	tokens = nltk.word_tokenize(txt_l)
	# weed out punctuation
	tokens = [
		tok
		for tok in tokens
		if tok[0].isalpha()
	]
	# return
	return tokens

def tokenize_nice(xstr,starters=set("([‘“"),**kwargs):
    l=tokenize_agnostic(xstr)
    l2=[]
    addback=False
    for x in l:
        if not x:#.strip():
            pass
        elif x[0].isalpha():
            if addback and len(l2):
                l2[-1]+=x
                addback=False
            else:
                l2+=[x]
        else:
            if x in starters:
                l2+=[x]
                addback=True
            elif len(l2):
                l2[-1]+=x
            else:
                l2+=[x]
                addback=True
    return l2

# def tokenize(txt,*x,**y):
# 	return tokenize_fast(txt,*x,**y)

def subfinder(mylist, pattern):
    matches = []
    for i in range(len(mylist)):
        if mylist[i] == pattern[0] and mylist[i:i+len(pattern)] == pattern:
            matches.append(pattern)
    return matches

DTOK_TREEBANK=None
def detokenize_treebank(x):
    global DTOK_TREEBANK
    if DTOK_TREEBANK is None:
        from nltk.tokenize.treebank import TreebankWordDetokenizer
        DTOK_TREEBANK = TreebankWordDetokenizer()
    return DTOK_TREEBANK.detokenize(x)

def detokenize(x):
    x=detokenize_treebank(x)
    return x

def setindex(df,key=LINEKEY,badcols={'index','level_0'},sort=True):
    df = df[set(df.columns) - set(df.index.names)]
    df = resetindex(df)

    cols=[]
    for x in key:
        if not x in set(cols) and x in set(df.columns):
            cols.append(x)

    odf=df.set_index(cols)
    if sort: odf=odf.sort_index()
    odf=odf[[col for col in odf.columns if not col in badcols and col not in key]]

    constraints = [c for c in odf.columns if c.startswith('*')]
    resortcols = ['*total'] + [c for c in constraints if c!='*total']
    resortcols+= list(sorted([c for c in odf.columns if c not in set(constraints)]))
    return odf[[c for c in resortcols if c in odf.columns]]

def resetindex(df,badcols={'level_0','index'},**y):
    cols=set(df.columns)
    inds=set(df.index.names)
    both=cols&inds
    # print(cols)
    # print(inds)
    # print(both)
    newdf=df[cols - both - badcols].reset_index(**y)
    # print(newdf.columns)
    # print()
    return newdf

def occurrences(string, sub):
    count = start = 0
    while True:
        start = string.find(sub, start) + 1
        if start > 0:
            count+=1
        else:
            return count

import numpy as np
def rolling_slices(df,window_len=3,incl_empty=True,keep_last_keys=[]):
    nrad=(window_len - 1)//2
    for i in range(len(df)):
        mini = i-nrad
        maxi = i+nrad
        empty_row=pd.Series(dict((k,np.nan) for k in df.columns))
        for xx in keep_last_keys: empty_row[xx]=df.iloc[0][xx]
        irows=[]
        inames=[]
        for ii in range(mini,maxi+1):
            inames+=[ii]
            if ii<0 or ii>=len(df):
                if incl_empty:

                    irows+=[empty_row]
            else:
                irows+=[df.iloc[ii]]
        yield pd.DataFrame(irows,index=inames)
def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i:i + n]

# def product(*args):
# 	if not args:
# 		return iter(((),)) # yield tuple()
# 	return (items + (item,)
# 		for items in product(*args[:-1]) for item in args[-1])

def hashstr(*x):
    import hashlib
    return hashlib.sha224(str(x).encode('utf-8')).hexdigest()

def pmap_do(inp):
    func,obj,args,kwargs = inp
    return func(obj,*args,**kwargs)

def pmap_iter(func, objs, args=[], kwargs={}, num_proc=DEFAULT_NUM_PROC, use_threads=False, progress=True, desc=None, position=0, leave=True, **y):
    """
    Yields results of func(obj) for each obj in objs
    Uses multiprocessing.Pool(num_proc) for parallelism.
    If use_threads, use ThreadPool instead of Pool.
    Results in any order.
    """

    # imports
    import multiprocessing as mp
    from tqdm import tqdm

    # check num proc
    num_cpu = mp.cpu_count()
    if num_proc>num_cpu: num_proc=num_cpu

    # if parallel
    if not desc: desc=f'Mapping {func.__name__}()'
    if desc: desc=f'{desc} [x{num_proc}]'
    if num_proc>1 and len(objs)>1:

        # real objects
        objects = [(func,obj,args,kwargs) for obj in objs]

        # create pool
        pool=mp.Pool(num_proc) if not use_threads else mp.pool.ThreadPool(num_proc)

        # yield iter
        iterr = pool.imap(pmap_do, objects)
        iterr=tqdm(
            iterr,
            total=len(objects),
            desc=desc,
            disable=not progress,
            leave=leave,
            position=position
        )
        yield from iterr

        # Close the pool?
        pool.close()
        pool.join()
    else:
        # yield
        iterr=tqdm(
            objs,
            desc=desc,
            disable=not progress,
            leave=leave,
            position=position
        )
        for obj in iterr:
            yield func(obj,*args,**kwargs)

def pmap(*x,**y):
    """
    Non iterator version of pmap_iter
    """
    # return as list
    return list(pmap_iter(*x,**y))

def index_by_truth(x):
    i=0
    for y in x:
        if y:
            yield i
            i+=1
        else:
            yield np.nan

def do_pmap_group(obj,*args,**kwargs):
    import pandas as pd
    import types
    func,group_df,group_key,group_name = obj

    if type(group_name) not in {list,tuple}:group_name=[group_name]
    if type(group_df)==str: group_df=pd.read_pickle(group_df)
    out=func(group_df,*args,**kwargs)
    if isinstance(out, types.GeneratorType):
        out=[x for x in out]
    if type(out)==list:
        try:
            out=[x for x in out if type(x) in {pd.DataFrame, pd.Series}]
            out=pd.concat(out)
        except ValueError:
            return out
    if type(out)==pd.DataFrame:
        for x,y in zip(group_key,group_name): out[x]=y
    return out

def joindfs(a,b):
    bcols = set(b.columns) - set(a.columns)
    return a.join(b[bcols])

def hashstr(x):
	import hashlib
	return hashlib.sha224(str(x).encode('utf-8')).hexdigest()

# Use sqlite dictionary
def get_db(
        prefix,
        mode='c',
        folders=[],
        autocommit=False,
        ):
    if mode=='w': autocommit=True
    ofnfn=os.path.join(PATH_DATA, f'db.{prefix}.sqlite')
    odir=os.path.dirname(ofnfn)
    if not os.path.exists(odir): os.makedirs(odir)
    if not os.path.exists(ofnfn): mode='c'
    return SqliteDict(
        ofnfn,
        tablename='data',
        autocommit=autocommit,
        flag='r' if mode=='r' else 'c',
        timeout=30
    )

def slices(l,n,strict=True):
    o=[]
    for x in l:
        o.append(x)
        if len(o)>n: o.pop(0)
        if not strict or len(o)==n:
            yield list(o)
def apply_combos(df,group1,group2,combo_key='combo_i'):
    # combo of indices?
    combo_opts = [
        [x for ii,x in grp.groupby(group2)]
        for i,grp in df.groupby(group1)
    ]

    # poss
    for combo in product(*combo_opts):
        if not len(combo): continue
        yield pd.concat(combo)# if len(combo) else pd.DataFrame()
def pmap_iter_groups(func,df_grouped,use_cache=False,num_proc=DEFAULT_NUM_PROC,iter=False,**attrs):
    import os,tempfile,pandas as pd
    from tqdm import tqdm

    # get index/groupby col name(s)
    group_key=df_grouped.grouper.names
    # if not using cache
    # if not use_cache or attrs.get('num_proc',1)<2:
    if not use_cache or len(df_grouped)<2 or num_proc<2:
        objs=[
            (func,group_df,group_key,group_name)
            for group_name,group_df in df_grouped
        ]
    else:
        objs=[]
        tmpdir=tempfile.mkdtemp()
        # for i,(group_name,group_df) in enumerate(tqdm(list(df_grouped),desc='Preparing input')):
        for i,(group_name,group_df) in enumerate(df_grouped):
            tmp_path = os.path.join(tmpdir, str(i)+'.pkl')
            # print([i,group_name,tmp_path,group_df])
            group_df.to_pickle(tmp_path)
            objs+=[(func,tmp_path,group_key,group_name)]

    # desc?
    if not attrs.get('desc'): attrs['desc']=f'Mapping {func.__name__}'

    #iterfunc = pmap if not iter else pmap_iter
    #return pd.concat(iterfunc) if not iter else
    return pmap_iter(
        do_pmap_group,
        objs,
        num_proc=num_proc,
        **attrs
    )
def pmap_groups(*x,**y):
    res = list(pmap_iter_groups(*x,**y))
    # print([type(x) for x in res])
    # for y in res[-1]:
        # print(type(y), y)
    resl = []
    for x in res:
        if type(x)==list:
            for y in x:
                resl.append(y)
        else:
            resl.append(x)
    return pd.concat(resl) if len(resl) else pd.DataFrame()

def check_basic_config():
    # check basic config
    try:
        if not os.path.exists(PATH_HOME): os.makedirs(PATH_HOME)
    except Exception:
        pass
    if not os.path.exists(PATH_DATA): os.makedirs(PATH_DATA)
    if not os.path.exists(os.path.join(PATH_DATA,'en')):
        zipfn=os.path.join(PATH_HOME,'data_cadence.zip')
        download(DATA_URL, zipfn)
        unzip(zipfn, PATH_HOME)
    try:
        nltk.word_tokenize('testing')
    except LookupError:
        nltk.download('punkt')
    try:
        nltk.corpus.stopwords.words('english')
    except LookupError:
        nltk.download('stopwords')

### utils
def printm(x):
    try:
        from IPython.display import Markdown
        return display(Markdown(x))
    except ImportError:
        print(x)

def download_wget(url, save_to, **attrs):
    import wget
    save_to_dir,save_to_fn=os.path.split(save_to)
    if save_to_dir:
        if not os.path.exists(save_to_dir): os.makedirs(save_to_dir)
        os.chdir(save_to_dir)
    fn=wget.download(url,bar=wget.bar_adaptive)
    os.rename(fn,save_to_fn)
    # print('\n>> saved:',save_to)

def download(url,save_to,force=False,desc=''):
    here=os.getcwd()
    download_wget(url,save_to,desc=desc)
    os.chdir(here)

def unzip(zipfn, dest='.', flatten=False, overwrite=False, replace_in_filenames={},desc='',progress=True):
    from zipfile import ZipFile
    from tqdm import tqdm

    # Open your .zip file
    if not desc: desc=f'Extracting {os.path.basename(zipfn)}'
    with ZipFile(zipfn) as zip_file:
        namelist=zip_file.namelist()

        # Loop over each file
        iterr=tqdm(iterable=namelist, total=len(namelist),desc=desc) if progress else namelist
        for member in iterr:
            # Extract each file to another directory
            # If you want to extract to current working directory, don't specify path
            filename = os.path.basename(member)
            if not filename: continue
            target_fnfn = os.path.join(dest,member) if not flatten else os.path.join(dest,filename)
            for k,v in replace_in_filenames.items(): target_fnfn = target_fnfn.replace(k,v)
            if not overwrite and os.path.exists(target_fnfn): continue
            target_dir = os.path.dirname(target_fnfn)
            try:
                if not os.path.exists(target_dir): os.makedirs(target_dir)
            except FileExistsError:
                pass
            except FileNotFoundError:
                continue

            with zip_file.open(member) as source, open(target_fnfn,'wb') as target:
                shutil.copyfileobj(source, target)

def get_num_lines(filename):
    from smart_open import open

    def blocks(files, size=65536):
        while True:
            b = files.read(size)
            if not b: break
            yield b

    with open(filename, 'r', errors='ignore') as f:
        numlines=sum(bl.count("\n") for bl in blocks(f))

    return numlines

max_stars_repo_name: quadrismegistus/cadence

max_issues_repo_name: quadrismegistus/cadence

max_forks_repo_name: quadrismegistus/cadence