text:
"""Classes used to construct investment planning subproblem"""

import time
import pickle
import logging
from collections import OrderedDict

from pyomo.environ import *
from pyomo.util.infeasible import log_infeasible_constraints

from data import ModelData
from components import CommonComponents

class Subproblem:
    # Pre-processed model data
    data = ModelData()

    def __init__(self):
        # Solver options
        self.keepfiles = False
        self.solver_options = {'Method': 1}  # 'MIPGap': 0.0005
        self.opt = SolverFactory('gurobi', solver_io='lp')

        # Setup logger
        logging.basicConfig(filename='subproblem.log', filemode='a',
                            format='%(asctime)s %(name)s %(levelname)s %(message)s',
                            datefmt='%Y-%m-%d %H:%M:%S',
                            level=logging.DEBUG)

        logging.info("Running subproblem")
        self.logger = logging.getLogger('Subproblem')

    def construct_model(self):
        """Construct subproblem model components"""

        # Used to define sets and parameters common to both master and subproblem
        common_components = CommonComponents()

        # Initialise base model object
        m = ConcreteModel()

        # Define sets - common to both master and subproblem
        m = common_components.define_sets(m)

        # Define parameters - common to both master and subproblem
        m = common_components.define_parameters(m)

        # Define variables - common to both master and subproblem
        m = common_components.define_variables(m)

        # Define expressions - common to both master and subproblem
        m = common_components.define_expressions(m)

        # Define expressions
        m = self.define_expressions(m)

        # Define blocks
        m = self.define_blocks(m)

        # Define constraints
        m = self.define_constraints(m)

        # Define objective
        m = self.define_objective(m)

        return m

    @staticmethod
    def define_expressions(m):
        """Define subproblem expressions"""

        def min_power_output_rule(_m, g, i):
            """
            Minimum generator output in MW

            Note: Min power output can be a function of max capacity which is only known ex-post for candidate
            generators.
            """

            return m.P_MIN_PROP[g] * m.P_MAX[g, i]

        # Expression for minimum power output
        m.P_MIN = Expression(m.G, m.I, rule=min_power_output_rule)

        return m

    def define_blocks(self, m):
        """Define blocks for each operating scenario"""

        def operating_scenario_block_rule(s, i, o):
            """Operating scenario block"""

            def define_block_parameters(s):
                """Define all parameters within a given block"""

                # Fixed shutdown indicator binary variable value
                s.FIXED_W = Param(m.G_E_THERM.union(m.G_C_THERM), m.T, initialize=0, within=Binary, mutable=True)

                # Fixed startup indicator binary variable value
                s.FIXED_V = Param(m.G_E_THERM.union(m.G_C_THERM), m.T, initialize=0, within=Binary, mutable=True)

                # Fixed on-state binary variable value
                s.FIXED_U = Param(m.G_E_THERM.union(m.G_C_THERM), m.T, initialize=0, within=Binary, mutable=True)

                # Indicates if unit is available (assume available for all periods for now)
                s.AVAILABILITY_INDICATOR = Param(m.G, initialize=1, within=Binary, mutable=True)

                # Power output in interval prior to model start (assume = 0 for now)
                s.P0 = Param(m.G, mutable=True, within=NonNegativeReals, initialize=0)

                # Energy in battering in interval prior to model start (assume battery initially completely discharged)
                s.Y0 = Param(m.G_C_STORAGE, initialize=0)

                def wind_capacity_factor_rule(_s, b, t):
                    """Wind capacity factors for each operating scenario"""

                    # Capacity factor for given wind plant and operating scenario
                    # capacity_factor = float(self.data.input_traces.loc[(i, o), ('WIND', b, t)])
                    capacity_factor = float(self.data.input_traces_dict[('WIND', b, t)][(i, o)])

                    # Handle small values - prevent numerical instability when solving
                    if capacity_factor < 0.01:
                        return float(0)
                    else:
                        return capacity_factor

                # Wind capacity factors
                s.Q_WIND = Param(m.B, m.T, rule=wind_capacity_factor_rule)

                def solar_capacity_factor_rule(_s, z, g, t):
                    """Solar capacity factors in each NEM zone"""

                    # Replace FFP with FFP2 - so technology type names are consistent
                    if g == 'FFP':
                        g = 'FFP2'

                    # Column name is a composite of the zone and technology type (g)
                    col = f'{z}|{g}'

                    # capacity_factor = float(self.data.input_traces.loc[(i, o), ('SOLAR', col, t)])
                    capacity_factor = float(self.data.input_traces_dict[('SOLAR', col, t)][(i, o)])

                    if capacity_factor < 0.01:
                        return float(0)
                    else:
                        return capacity_factor

                # Solar capacity factors
                s.Q_SOLAR = Param(m.Z, m.G_C_SOLAR_TECHNOLOGIES, m.T, rule=solar_capacity_factor_rule)

                def demand_rule(_s, z, t):
                    """NEM demand in each zone"""

                    # Demand in given NEM zone
                    # demand = float(self.data.input_traces.loc[(i, o), ('DEMAND', z, t)])
                    demand = float(self.data.input_traces_dict[('DEMAND', z, t)][(i, o)])

                    return demand

                # Demand in each NEM zone
                s.DEMAND = Param(m.Z, m.T, rule=demand_rule)

                def scenario_duration_rule(_s):
                    """Normalised duration for each operation scenario"""

                    return float(self.data.input_traces.loc[(i, o), ('K_MEANS', 'METRIC', 'NORMALISED_DURATION')])

                # Normalised duration for each operating scenario
                s.NORMALISED_DURATION = Param(rule=scenario_duration_rule)

                def historic_hydro_output_rule(_s, g, t):
                    """Historic output for existing hydro generators"""
                    # output = float(self.data.input_traces.loc[(i, o), ('HYDRO', g, t)])
                    output = float(self.data.input_traces_dict[('HYDRO', g, t)][(i, o)])

                    # Remove small values - prevent numerical instability when solving
                    if output < 1:
                        return float(0)
                    else:
                        return output

                # Assumed output for hydro generators (based on historic values)
                s.P_HYDRO_HISTORIC = Param(m.G_E_HYDRO, m.T, rule=historic_hydro_output_rule)

                return s

            def define_block_variables(s):
                """Define variables for each block"""

                # Startup state variable
                s.v = Var(m.G_E_THERM.union(m.G_C_THERM), m.T, within=NonNegativeReals, initialize=0)

                # Shutdown state variable
                s.w = Var(m.G_E_THERM.union(m.G_C_THERM), m.T, within=NonNegativeReals, initialize=0)

                # On-state variable
                s.u = Var(m.G_E_THERM.union(m.G_C_THERM), m.T, within=NonNegativeReals, initialize=1)

                # Power output above minimum dispatchable level of output [MW]
                s.p = Var(m.G, m.T, within=NonNegativeReals, initialize=0)

                # Upward reserve allocation [MW]
                s.r_up = Var(m.G_E_THERM.union(m.G_C_THERM).union(m.G_C_STORAGE), m.T, within=NonNegativeReals,
                             initialize=0)

                # Storage unit charging (power in) [MW]
                s.p_in = Var(m.G_C_STORAGE, m.T, within=NonNegativeReals, initialize=0)

                # Storage unit discharging (power out) [MW]
                s.p_out = Var(m.G_C_STORAGE, m.T, within=NonNegativeReals, initialize=0)

                # Energy in storage unit [MWh]
                s.y = Var(m.G_C_STORAGE, m.T, within=NonNegativeReals, initialize=0)

                # Powerflow between NEM zones [MW]
                s.p_flow = Var(m.L, m.T, initialize=0)

                # Lost load - up [MW]
                s.p_lost_up = Var(m.Z, m.T, within=NonNegativeReals, initialize=0)

                # Lost load - down [MW]
                s.p_lost_down = Var(m.Z, m.T, within=NonNegativeReals, initialize=0)

                # Amount by which upward reserve is violated [MW]
                s.upward_reserve_violation = Var(m.R, m.T, within=NonNegativeReals, initialize=0)

                return s

            def define_block_expressions(s):
                """Define expressions used when constructing each operating scenario"""

                def total_power_output_rule(_s, g, t):
                    """Total power output [MW]"""

                    if g in m.G_THERM_QUICK:
                        if t < m.T.last():
                            return ((m.P_MIN[g, i] * (s.u[g, t] + s.v[g, t + 1])) + s.p[g, t]) * \
                                   s.AVAILABILITY_INDICATOR[
                                       g]
                        else:
                            return ((m.P_MIN[g, i] * s.u[g, t]) + s.p[g, t]) * s.AVAILABILITY_INDICATOR[g]

                    # Define startup and shutdown trajectories for 'slow' generators
                    elif g in m.G_THERM_SLOW:
                        # Startup duration
                        SU_D = ceil(m.P_MIN[g, i].expr() / m.RR_SU[g])

                        # Startup power output trajectory increment
                        ramp_up_increment = m.P_MIN[g, i].expr() / SU_D

                        # Startup power output trajectory
                        P_SU = OrderedDict({k + 1: ramp_up_increment * k for k in range(0, SU_D + 1)})

                        # Shutdown duration
                        SD_D = ceil(m.P_MIN[g, i].expr() / m.RR_SD[g])

                        # Shutdown power output trajectory increment
                        ramp_down_increment = m.P_MIN[g, i].expr() / SD_D

                        # Shutdown power output trajectory
                        P_SD = OrderedDict(
                            {k + 1: m.P_MIN[g, i].expr() - (ramp_down_increment * k) for k in range(0, SD_D + 1)})

                        if t < m.T.last():
                            return (((m.P_MIN[g, i] * (s.u[g, t] + s.v[g, t + 1])) + s.p[g, t]
                                     + sum(P_SU[k] * s.v[g, t - k + SU_D + 2] if t - k + SU_D + 2 in m.T else 0 for k in
                                           range(1, SU_D + 1))
                                     + sum(P_SD[k] * s.w[g, t - k + 2] if t - k + 2 in m.T else 0 for k in
                                           range(2, SD_D + 2))
                                     ) * s.AVAILABILITY_INDICATOR[g])
                        else:
                            return (((m.P_MIN[g, i] * s.u[g, t]) + s.p[g, t]
                                     + sum(P_SU[k] * s.v[g, t - k + SU_D + 2] if t - k + SU_D + 2 in m.T else 0 for k in
                                           range(1, SU_D + 1))
                                     + sum(P_SD[k] * s.w[g, t - k + 2] if t - k + 2 in m.T else 0 for k in
                                           range(2, SD_D + 2))
                                     ) * s.AVAILABILITY_INDICATOR[g])

                    # Remaining generators with no startup / shutdown directory defined
                    else:
                        return (m.P_MIN[g, i] + s.p[g, t]) * s.AVAILABILITY_INDICATOR[g]

                # Total power output [MW]
                s.P_TOTAL = Expression(m.G, m.T, rule=total_power_output_rule)

                def generator_energy_output_rule(_s, g, t):
                    """
                    Total generator energy output [MWh]

                    No information regarding generator dispatch level in period prior to
                    t=0. Assume that all generators are at their minimum dispatch level.
                    """

                    if t != m.T.first():
                        return (s.P_TOTAL[g, t - 1] + s.P_TOTAL[g, t]) / 2

                    # Else, first interval (t-1 will be out of range)
                    else:
                        return (s.P0[g] + s.P_TOTAL[g, t]) / 2

                # Energy output for a given generator
                s.ENERGY = Expression(m.G, m.T, rule=generator_energy_output_rule)

                def storage_unit_energy_output_rule(_s, g, t):
                    """Energy output from storage units"""

                    if t != m.T.first():
                        return (s.p_out[g, t] + s.p_out[g, t - 1]) / 2

                    else:
                        return s.p_out[g, t] / 2

                # Energy output for a given storage unit
                s.ENERGY_OUT = Expression(m.G_C_STORAGE, m.T, rule=storage_unit_energy_output_rule)

                def storage_unit_energy_input_rule(_s, g, t):
                    """Energy input (charging) from storage units"""

                    if t != m.T.first():
                        return (s.p_in[g, t] + s.p_in[g, t - 1]) / 2

                    else:
                        return s.p_in[g, t] / 2

                # Energy input for a given storage unit
                s.ENERGY_IN = Expression(m.G_C_STORAGE, m.T, rule=storage_unit_energy_input_rule)

                def thermal_operating_costs_rule(_s):
                    """Cost to operate existing and candidate thermal units"""

                    return (
                        sum((m.C_MC[g, i] + (m.EMISSIONS_RATE[g] - m.baseline[i]) * m.permit_price[i]) * s.ENERGY[g, t]
                            + (m.C_SU[g, i] * s.v[g, t]) + (m.C_SD[g, i] * s.w[g, t])
                            for g in m.G_E_THERM.union(m.G_C_THERM) for t in m.T))

                # Existing and candidate thermal unit operating costs for given scenario
                s.C_OP_THERM = Expression(rule=thermal_operating_costs_rule)

                def hydro_operating_costs_rule(_s):
                    """Cost to operate existing hydro generators"""

                    return sum(m.C_MC[g, i] * s.ENERGY[g, t] for g in m.G_E_HYDRO for t in m.T)

                # Existing hydro unit operating costs (no candidate hydro generators)
                s.C_OP_HYDRO = Expression(rule=hydro_operating_costs_rule)

                def solar_operating_costs_rule(_s):
                    """Cost to operate existing and candidate solar units"""

                    return (sum(m.C_MC[g, i] * s.ENERGY[g, t] for g in m.G_E_SOLAR for t in m.T)
                            + sum(
                                (m.C_MC[g, i] - m.baseline[i] * m.permit_price[i]) * s.ENERGY[g, t] for g in m.G_C_SOLAR
                                for t in
                                m.T))

                # Existing and candidate solar unit operating costs (only candidate solar eligible for credits)
                s.C_OP_SOLAR = Expression(rule=solar_operating_costs_rule)

                def wind_operating_costs_rule(_s):
                    """Cost to operate existing and candidate wind generators"""

                    return (sum(m.C_MC[g, i] * s.ENERGY[g, t] for g in m.G_E_WIND for t in m.T)
                            + sum(
                                (m.C_MC[g, i] - m.baseline[i] * m.permit_price[i]) * s.ENERGY[g, t] for g in m.G_C_WIND
                                for t in
                                m.T))

                # Existing and candidate solar unit operating costs (only candidate solar eligible for credits)
                s.C_OP_WIND = Expression(rule=wind_operating_costs_rule)

                def storage_unit_charging_cost_rule(_s):
                    """Cost to charge storage unit"""

                    return sum(m.C_MC[g, i] * s.ENERGY_IN[g, t] for g in m.G_C_STORAGE for t in m.T)

                # Charging cost rule - no subsidy received when purchasing energy
                s.C_OP_STORAGE_CHARGING = Expression(rule=storage_unit_charging_cost_rule)

                def storage_unit_discharging_cost_rule(_s):
                    """
                    Cost to charge storage unit

                    Note: If storage units are included in the scheme this could create an undesirable outcome. Units
                    would be subsidised for each MWh they generate. Therefore they could be incentivised to continually charge
                    and then immediately discharge in order to receive the subsidy. For now assume the storage units are not
                    eligible to receive a subsidy for each MWh under the policy.
                    """
                    # - (m.baseline * m.permit_price))
                    return sum(m.C_MC[g, i] * s.ENERGY_OUT[g, t] for g in m.G_C_STORAGE for t in m.T)

                # Discharging cost rule - assumes storage units are eligible under REP scheme
                s.C_OP_STORAGE_DISCHARGING = Expression(rule=storage_unit_discharging_cost_rule)

                # Candidate storage unit operating costs
                s.C_OP_STORAGE = Expression(expr=s.C_OP_STORAGE_CHARGING + s.C_OP_STORAGE_DISCHARGING)

                def lost_load_cost_rule(_s):
                    """Value of lost-load"""

                    return sum((s.p_lost_up[z, t] + s.p_lost_down[z, t]) * m.C_LOST_LOAD for z in m.Z for t in m.T)

                # Total cost of lost-load
                s.C_OP_LOST_LOAD = Expression(rule=lost_load_cost_rule)

                def total_operating_cost_rule(_s):
                    """Total operating cost"""

                    return s.C_OP_THERM + s.C_OP_HYDRO + s.C_OP_SOLAR + s.C_OP_WIND + s.C_OP_STORAGE + s.C_OP_LOST_LOAD

                # Total operating cost
                s.C_OP_TOTAL = Expression(rule=total_operating_cost_rule)

                def storage_unit_energy_capacity_rule(_s, g):
                    """Energy capacity depends on installed capacity (variable in master problem)"""

                    return sum(m.x_c[g, y] for y in m.I if y <= i)

                # Capacity of storage unit [MWh]
                s.STORAGE_UNIT_ENERGY_CAPACITY = Expression(m.G_C_STORAGE, rule=storage_unit_energy_capacity_rule)

                def storage_unit_max_energy_interval_end_rule(_s, g):
                    """Maximum energy at end of storage interval. Assume equal to unit capacity."""

                    return s.STORAGE_UNIT_ENERGY_CAPACITY[g]

                # Max state of charge for storage unit at end of operating scenario (assume = unit capacity)
                s.STORAGE_INTERVAL_END_MAX_ENERGY = Expression(m.G_C_STORAGE,
                                                               rule=storage_unit_max_energy_interval_end_rule)

                def storage_unit_max_power_out_rule(_s, g):
                    """
                    Maximum discharging power of storage unit - set equal to energy capacity. Assumes
                    storage unit can completely discharge in 1 hour
                    """

                    return s.STORAGE_UNIT_ENERGY_CAPACITY[g]

                # Max MW out of storage device - discharging
                s.P_STORAGE_MAX_OUT = Expression(m.G_C_STORAGE, rule=storage_unit_max_power_out_rule)

                def storage_unit_max_power_in_rule(_s, g):
                    """
                    Maximum charging power of storage unit - set equal to energy capacity. Assumes
                    storage unit can completely charge in 1 hour
                    """

                    return s.STORAGE_UNIT_ENERGY_CAPACITY[g]

                # Max MW into storage device - charging
                s.P_STORAGE_MAX_IN = Expression(m.G_C_STORAGE, rule=storage_unit_max_power_in_rule)

                def scenario_emissions_rule(_s):
                    """
                    Total emissions for a given scenario

                    Note: Only thermal generators assumed to emit
                    """
                    return sum(
                        s.ENERGY[g, t] * m.EMISSIONS_RATE[g] for g in m.G_E_THERM.union(m.G_C_THERM) for t in m.T)

                # Total emissions in a given scenario
                # TODO: Figure out how to handle scenario duration here
                s.TOTAL_EMISSIONS = Expression(rule=scenario_emissions_rule)

                def scenario_revenue_rule(_s):
                    """
                    Net policy revenue for a given scenario

                    Note: Only existing and candidate thermal, and candidate wind and solar eligible under scheme
                    (payments to existing renewables wouldn't lead to any emissions abatement)
                    """
                    gens = m.G_E_THERM.union(m.G_C_THERM).union(m.G_C_WIND).union(m.G_C_SOLAR)

                    return sum((m.EMISSIONS_RATE[g] - m.baseline[i]) * s.ENERGY[g, t] for g in gens for t in m.T)

                # Net scheme revenue for a given scenario
                # TODO: Figure out how to handle scenario duration here
                s.TOTAL_REVENUE = Expression(rule=scenario_revenue_rule)

                # Penalty imposed on violating upward reserve requirement
                s.UPWARD_RESERVE_VIOLATION_PENALTY = Expression(
                    expr=sum(m.C_LOST_LOAD * s.upward_reserve_violation[r, t] for r in m.R for t in m.T))

                return s

            def define_block_constraints(s):
                """Define model constraints"""

                def upward_power_reserve_rule(_s, r, t):
                    """Upward reserve constraint"""

                    # NEM region-zone map
                    # region_zone_map = self._get_nem_region_zone_map()
                    region_zone_map = self.data.nem_region_zone_map_dict

                    # Mapping describing the zone to which each generator is assigned
                    # gen_zone_map = self._get_generator_zone_map()
                    gen_zone_map = self.data.generator_zone_map

                    # Existing and candidate thermal gens + candidate storage units
                    gens = m.G_E_THERM.union(m.G_C_THERM).union(m.G_C_STORAGE)

                    return (sum(s.r_up[g, t] for g in gens if gen_zone_map[g] in region_zone_map[r])
                            + s.upward_reserve_violation[r, t] == m.RESERVE_UP[r])

                # Upward power reserve rule for each NEM region
                s.UPWARD_POWER_RESERVE = Constraint(m.R, m.T, rule=upward_power_reserve_rule)

                def power_production_rule(_s, g, t):
                    """Power production rule"""

                    if t != m.T.last():
                        return (s.p[g, t] + s.r_up[g, t]
                                <= ((m.P_MAX[g, i] - m.P_MIN[g, i]) * s.u[g, t])
                                - ((m.P_MAX[g, i] - m.RR_SD[g]) * s.w[g, t + 1])
                                + ((m.RR_SU[g] - m.P_MIN[g, i]) * s.v[g, t + 1]))
                    else:
                        return s.p[g, t] + s.r_up[g, t] <= (m.P_MAX[g, i] - m.P_MIN[g, i]) * s.u[g, t]

                # Power production
                s.POWER_PRODUCTION = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=power_production_rule)

                def ramp_rate_up_rule(_s, g, t):
                    """Ramp-rate up constraint"""

                    if t > m.T.first():
                        return (s.p[g, t] + s.r_up[g, t]) - s.p[g, t - 1] <= m.RR_UP[g]

                    else:
                        # Ramp-rate for first interval
                        return s.p[g, t] - s.P0[g] <= m.RR_UP[g]

                # Ramp-rate up limit
                s.RAMP_RATE_UP = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=ramp_rate_up_rule)

                def ramp_rate_down_rule(_s, g, t):
                    """Ramp-rate down constraint"""

                    if t > m.T.first():
                        return - s.p[g, t] + s.p[g, t - 1] <= m.RR_DOWN[g]

                    else:
                        # Ramp-rate for first interval
                        return - s.p[g, t] + s.P0[g] <= m.RR_DOWN[g]

                # Ramp-rate up limit
                s.RAMP_RATE_DOWN = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=ramp_rate_down_rule)

                def existing_wind_output_min_rule(_s, g, t):
                    """Constrain minimum output for existing wind generators"""

                    return s.P_TOTAL[g, t] >= 0

                # Minimum wind output for existing generators
                s.EXISTING_WIND_MIN_OUTPUT = Constraint(m.G_E_WIND, m.T, rule=existing_wind_output_min_rule)

                def wind_output_max_rule(_s, g, t):
                    """
                    Constrain maximum output for wind generators

                    Note: Candidate unit output depends on investment decisions
                    """

                    # Get wind bubble to which generator belongs
                    if g in m.G_E_WIND:

                        # If an existing generator
                        # bubble = self.existing_wind_bubble_map.loc[g, 'BUBBLE']
                        bubble = self.data.existing_wind_bubble_map_dict['BUBBLE'][g]

                    elif g in m.G_C_WIND:

                        # If a candidate generator
                        # bubble = self.candidate_units.loc[g, ('PARAMETERS', 'WIND_BUBBLE')]
                        bubble = self.data.candidate_units_dict[('PARAMETERS', 'WIND_BUBBLE')][g]

                    else:
                        raise Exception(f'Unexpected generator: {g}')

                    return s.P_TOTAL[g, t] <= s.Q_WIND[bubble, t] * m.P_MAX[g, i]

                # Max output from existing wind generators
                s.EXISTING_WIND_MAX_OUTPUT = Constraint(m.G_E_WIND.union(m.G_C_WIND), m.T, rule=wind_output_max_rule)

                def solar_output_max_rule(_s, g, t):
                    """
                    Constrain maximum output for solar generators

                    Note: Candidate unit output depends on investment decisions
                    """

                    # Get NEM zone
                    if g in m.G_E_SOLAR:
                        # If an existing generator
                        # zone = self.existing_units.loc[g, ('PARAMETERS', 'NEM_ZONE')]
                        zone = self.data.existing_units_dict[('PARAMETERS', 'NEM_ZONE')][g]

                        # Assume existing arrays are single-axis tracking arrays
                        technology = 'SAT'

                    elif g in m.G_C_SOLAR:
                        # If a candidate generator
                        # zone = self.candidate_units.loc[g, ('PARAMETERS', 'ZONE')]
                        zone = self.data.candidate_units_dict[('PARAMETERS', 'ZONE')][g]

                        # Extract technology type from unit ID
                        technology = g.split('-')[-1]

                    else:
                        raise Exception(f'Unexpected generator: {g}')

                    return s.P_TOTAL[g, t] <= s.Q_SOLAR[zone, technology, t] * m.P_MAX[g, i]

                # Max output from existing wind generators
                s.EXISTING_SOLAR_MAX_OUTPUT = Constraint(m.G_E_SOLAR.union(m.G_C_SOLAR), m.T,
                                                         rule=solar_output_max_rule)

                def hydro_output_max_rule(_s, g, t):
                    """
                    Constrain hydro output to registered capacity of existing plant

                    Note: Assume no investment in hydro over model horizon (only consider existing units)
                    """

                    return s.P_TOTAL[g, t] <= s.P_HYDRO_HISTORIC[g, t]

                # Max output from existing hydro generators
                s.EXISTING_HYDRO_MAX_OUTPUT = Constraint(m.G_E_HYDRO, m.T, rule=hydro_output_max_rule)

                # TODO: May want to add energy constraint for hydro units

                def thermal_generator_max_output_rule(_s, g, t):
                    """Max MW output for thermal generators"""

                    return s.P_TOTAL[g, t] <= m.P_MAX[g, i]

                # Max output for existing and candidate thermal generators
                s.EXISTING_THERMAL_MAX_OUTPUT = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T,
                                                           rule=thermal_generator_max_output_rule)

                def storage_max_charge_rate_rule(_s, g, t):
                    """Maximum charging power for storage units [MW]"""

                    return s.p_in[g, t] <= s.P_STORAGE_MAX_IN[g]

                # Storage unit max charging power
                s.STORAGE_MAX_CHARGE_RATE = Constraint(m.G_C_STORAGE, m.T, rule=storage_max_charge_rate_rule)

                def storage_max_discharge_rate_rule(_s, g, t):
                    """Maximum discharging power for storage units [MW]"""

                    return s.p_out[g, t] + s.r_up[g, t] <= s.P_STORAGE_MAX_OUT[g]

                # Storage unit max charging power
                s.STORAGE_MAX_DISCHARGE_RATE = Constraint(m.G_C_STORAGE, m.T, rule=storage_max_discharge_rate_rule)

                def state_of_charge_rule(_s, g, t):
                    """Energy within a given storage unit [MWh]"""

                    return s.y[g, t] <= s.STORAGE_UNIT_ENERGY_CAPACITY[g]

                # Energy within storage unit [MWh]
                s.STATE_OF_CHARGE = Constraint(m.G_C_STORAGE, m.T, rule=state_of_charge_rule)

                def storage_energy_transition_rule(_s, g, t):
                    """Constraint that couples energy + power between periods for storage units"""

                    if t != m.T.first():
                        return s.y[g, t] == s.y[g, t - 1] + (m.BATTERY_EFFICIENCY[g] * s.p_in[g, t]) - (
                                s.p_out[g, t] / m.BATTERY_EFFICIENCY[g])
                    else:
                        # Assume battery completely discharged in first period
                        return s.y[g, t] == s.Y0[g] + (m.BATTERY_EFFICIENCY[g] * s.p_in[g, t]) - (
                                s.p_out[g, t] / m.BATTERY_EFFICIENCY[g])

                # Account for inter-temporal energy transition within storage units
                s.STORAGE_ENERGY_TRANSITION = Constraint(m.G_C_STORAGE, m.T, rule=storage_energy_transition_rule)

                def storage_interval_end_min_energy_rule(_s, g):
                    """Lower bound on permissible amount of energy in storage unit at end of operating scenario"""

                    return m.STORAGE_INTERVAL_END_MIN_ENERGY[g] <= s.y[g, m.T.last()]

                # Minimum amount of energy that must be in storage unit at end of operating scenario
                s.STORAGE_INTERVAL_END_MIN_ENERGY_CONS = Constraint(m.G_C_STORAGE,
                                                                    rule=storage_interval_end_min_energy_rule)

                def storage_interval_end_max_energy_rule(_s, g):
                    """Upper bound on permissible amount of energy in storage unit at end of operating scenario"""

                    return s.y[g, m.T.last()] <= s.STORAGE_INTERVAL_END_MAX_ENERGY[g]

                # Maximum amount of energy that must be in storage unit at end of operating scenario
                s.STORAGE_INTERVAL_END_MAX_ENERGY_CONS = Constraint(m.G_C_STORAGE,
                                                                    rule=storage_interval_end_max_energy_rule)

                def power_balance_rule(_s, z, t):
                    """Power balance for each NEM zone"""

                    # Existing units within zone
                    # existing_units = self.existing_units[
                    #     self.existing_units[('PARAMETERS', 'NEM_ZONE')] == z].index.tolist()
                    existing_units = [gen for gen, zone in
                                      self.data.existing_units_dict[('PARAMETERS', 'NEM_ZONE')].items()
                                      if zone == z]

                    # Candidate units within zone
                    # candidate_units = self.candidate_units[
                    #     self.candidate_units[('PARAMETERS', 'ZONE')] == z].index.tolist()
                    candidate_units = [gen for gen, zone in
                                       self.data.candidate_units_dict[('PARAMETERS', 'ZONE')].items()
                                       if zone == z]

                    # All generators within a given zone
                    generators = existing_units + candidate_units

                    # Storage units within a given zone
                    # storage_units = (self.battery_properties.loc[self.battery_properties['NEM_ZONE'] == z, 'NEM_ZONE']
                    #                  .index.tolist())
                    storage_units = [gen for gen, zone in self.data.battery_properties_dict['NEM_ZONE'].items() if
                                     zone == z]

                    return (sum(s.P_TOTAL[g, t] for g in generators) - s.DEMAND[z, t]
                            - sum(m.INCIDENCE_MATRIX[l, z] * s.p_flow[l, t] for l in m.L)
                            + sum(s.p_out[g, t] - s.p_in[g, t] for g in storage_units)
                            + s.p_lost_up[z, t] - s.p_lost_down[z, t] == 0)

                # Power balance constraint for each zone and time period
                s.POWER_BALANCE = Constraint(m.Z, m.T, rule=power_balance_rule)

                def powerflow_min_constraint_rule(_s, l, t):
                    """Minimum powerflow over a link connecting adjacent NEM zones"""

                    return s.p_flow[l, t] >= m.POWERFLOW_MIN[l]

                # Constrain max power flow over given network link
                s.POWERFLOW_MIN_CONS = Constraint(m.L_I, m.T, rule=powerflow_min_constraint_rule)

                def powerflow_max_constraint_rule(_s, l, t):
                    """Maximum powerflow over a link connecting adjacent NEM zones"""

                    return s.p_flow[l, t] <= m.POWERFLOW_MAX[l]

                # Constrain max power flow over given network link
                s.POWERFLOW_MAX_CONS = Constraint(m.L_I, m.T, rule=powerflow_max_constraint_rule)

                def fixed_on_state_rule(_s, g, t):
                    """Used to obtain sensitivities for fixed on-state binary variables in master problem"""

                    return s.u[g, t] == s.FIXED_U[g, t]

                # Set fixed on-state variables to given value - used to recover sensitivities
                s.FIXED_U_CONS = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=fixed_on_state_rule)

                def fixed_startup_state_rule(_s, g, t):
                    """Used to obtain sensitivities for fixed startup state binary variables in master problem"""

                    return s.v[g, t] == s.FIXED_V[g, t]

                # Set fixed on-state variables to given value - used to recover sensitivities
                s.FIXED_V_CONS = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=fixed_startup_state_rule)

                def fixed_shutdown_state_rule(_s, g, t):
                    """Used to obtain sensitivities for fixed shutdown state binary variables in master problem"""

                    return s.w[g, t] == s.FIXED_W[g, t]

                # Set fixed on-state variables to given value - used to recover sensitivities
                s.FIXED_W_CONS = Constraint(m.G_E_THERM.union(m.G_C_THERM), m.T, rule=fixed_shutdown_state_rule)

            def construct_block(s):
                """Construct block for each operating scenario"""
                print(f'Constructing block: Year - {i}, Scenario - {o}')
                print('###############################################')

                # Define block parameters for given operating scenario
                start = time.time()
                s = define_block_parameters(s)
                print(f'Defined block parameters in: {time.time() - start}s')

                # Define block variables
                start = time.time()
                s = define_block_variables(s)
                print(f'Defined block variables in: {time.time() - start}s')

                # Define block expressions
                start = time.time()
                s = define_block_expressions(s)
                print(f'Defined block expressions in: {time.time() - start}s')

                # Define block constraints
                start = time.time()
                s = define_block_constraints(s)
                print(f'Defined block constraints in: {time.time() - start}s')

                return s

            # Construct block
            construct_block(s)

        # Construct block
        m.SCENARIO = Block(m.I, m.O, rule=operating_scenario_block_rule)

        return m

    @staticmethod
    def define_constraints(m):
        """Define subproblem constraints"""

        def fixed_baseline_rule(_m, i):
            """Fix emissions intensity baseline to given value in sub-problems"""

            return m.baseline[i] == m.FIXED_BASELINE[i]

        # Fix emissions intensity baseline for each year in model horizon
        m.FIXED_BASELINE_CONS = Constraint(m.I, rule=fixed_baseline_rule)

        def fixed_permit_price_rule(_m, i):
            """Fixed permit price rule"""

            return m.permit_price[i] == m.FIXED_PERMIT_PRICE[i]

        # Fix permit price to given value in sub-problems
        m.FIXED_PERMIT_PRICE_CONS = Constraint(m.I, rule=fixed_permit_price_rule)

        def fixed_capacity_continuous_rule(m, g, i):
            """Fixed installed capacities"""

            return m.FIXED_X_C[g, i] == m.x_c[g, i]

        # Fixed installed capacity for candidate units - continuous sizing
        m.FIXED_CAPACITY_CONT = Constraint(m.G_C_SOLAR.union(m.G_C_WIND).union(m.G_C_STORAGE), m.I,
                                           rule=fixed_capacity_continuous_rule)

        def fixed_capacity_discrete_rule(m, g, i, n):
            """Fix installed capacity for discrete unit sizing"""
            # TODO: Need to fix how binary variables for discrete sizing are updated in subproblem
            return m.FIXED_D[g, i, n] == m.d[g, i, n]

        # Fixed installed capacity for candidate units - discrete sizing
        m.FIXED_CAPACITY_DISC = Constraint(m.G_C_THERM, m.I, m.G_C_THERM_SIZE_OPTIONS,
                                           rule=fixed_capacity_discrete_rule)

        def scheme_revenue_constraint_rule(m):
            """Scheme must be revenue neutral over model horizon"""

            return (sum(m.SCENARIO[i, o].TOTAL_REVENUE for i in m.I for o in m.O) + m.revenue_shortfall
                    >= m.REVENUE_TARGET)

        # Revenue constraint - must break-even over model horizon
        m.REVENUE_CONSTRAINT = Constraint(rule=scheme_revenue_constraint_rule)

        def scheme_emissions_constraint_rule(m):
            """Emissions limit over model horizon"""

            return (sum(m.SCENARIO[i, o].TOTAL_EMISSIONS for i in m.I for o in m.O)
                    <= m.EMISSIONS_TARGET + m.emissions_target_exceeded)

        # Emissions constraint - must be less than some target, else penalty imposed for each unit above target
        m.EMISSIONS_CONSTRAINT = Constraint(rule=scheme_emissions_constraint_rule)

        return m

    @staticmethod
    def define_objective(m):
        """Define subproblem objective"""

        # Minimise total operating cost - include penalty for revenue / emissions constraint violations
        m.OBJECTIVE = Objective(expr=sum(m.SCENARIO[i, o].C_OP_TOTAL for i in m.I for o in m.O)
                                     + sum(m.SCENARIO[i, o].UPWARD_RESERVE_VIOLATION_PENALTY for i in m.I for o in m.O)
                                     + m.C_FOM_TOTAL
                                     + m.C_INV_TOTAL
                                     + m.C_EMISSIONS_VIOLATION + m.C_REVENUE_VIOLATION, sense=minimize)

        return m

    @staticmethod
    def update_fixed_baseline(m, fixed_baselines):
        """Update fixed baseline - obtained from master problem"""

        for i in m.I:
            m.FIXED_BASELINE[i] = fixed_baselines[i]
            # m.FIXED_BASELINE[i] = 0

        return m

    @staticmethod
    def update_fixed_permit_price(m, fixed_permit_prices):
        """Update fixed permit price - obtained from master problem"""

        for i in m.I:
            m.FIXED_PERMIT_PRICE[i] = fixed_permit_prices[i]
            # m.FIXED_PERMIT_PRICE[i] = 0
        return m

    @staticmethod
    def update_fixed_capacity_continuous(m, continuous_capacities):
        """Update installed capacity for units with continuous capacity sizing option"""

        for i in m.I:
            for g in m.G_C_WIND.union(m.G_C_SOLAR).union(m.G_C_STORAGE):
                m.FIXED_X_C[g, i] = continuous_capacities[(g, i)]
                # m.FIXED_X_C[g, i] = 0

        return m

    @staticmethod
    def update_fixed_capacity_discrete(m, discrete_capacities):
        """Update installed capacity for units with discrete capacity sizing option"""

        for i in m.I:
            for g in m.G_C_THERM:
                for n in m.G_C_THERM_SIZE_OPTIONS:
                    m.FIXED_D[g, i, n] = discrete_capacities[(g, i, n)]
                    # if n == 0:
                    #     m.FIXED_D[g, i, n] = 1
                    # else:
                    #     m.FIXED_D[g, i, n] = 0

        return m

    @staticmethod
    def update_fixed_shutdown_state(m, shutdown_states):
        """Update fixed shutdown state variables"""
        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].FIXED_W[g, t] = shutdown_states[i][o][(g, t)]
                        # m.SCENARIO[i, o].FIXED_W[g, t] = 0

        return m

    @staticmethod
    def update_fixed_startup_state(m, startup_states):
        """Update fixed startup state variables"""
        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].FIXED_V[g, t] = startup_states[i][o][(g, t)]
                        # m.SCENARIO[i, o].FIXED_V[g, t] = 0

        return m

    @staticmethod
    def update_fixed_on_state(m, on_states):
        """Update fixed startup state variables"""
        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].FIXED_U[g, t] = on_states[i][o][(g, t)]
                        # m.SCENARIO[i, o].FIXED_U[g, t] = 1

        return m

    @staticmethod
    def fix_baseline(m):
        """Fix emissions intensity baseline"""
        for i in m.I:
            m.baseline[i].fix(m.FIXED_BASELINE[i].value)

        return m

    @staticmethod
    def unfix_baseline(m):
        """Fix emissions intensity baseline"""
        for i in m.I:
            m.baseline[i].unfix()

        return m

    @staticmethod
    def fix_permit_price(m):
        """Fix permit price"""
        for i in m.I:
            m.permit_price[i].fix(m.FIXED_PERMIT_PRICE[i].value)

        return m

    @staticmethod
    def unfix_permit_price(m):
        """Unfix permit price"""
        for i in m.I:
            m.permit_price[i].unfix()

        return m

    @staticmethod
    def fix_discrete_capacity_options(m):
        """Fix binary variables used in discrete capacity sizing"""
        for g in m.G_C_THERM:
            for i in m.I:
                for n in m.G_C_THERM_SIZE_OPTIONS:
                    m.d[g, i, n].fix(m.FIXED_D[g, i, n].value)

        return m

    @staticmethod
    def unfix_discrete_capacity_options(m):
        """Unfix binary variables used in discrete capacity sizing"""
        for g in m.G_C_THERM:
            for i in m.I:
                for n in m.G_C_THERM_SIZE_OPTIONS:
                    m.d[g, i, n].unfix()
        return m

    @staticmethod
    def fix_continuous_capacity_options(m):
        """Fix variables used in continuous capacity sizing options"""
        for g in m.G_C_WIND.union(m.G_C_SOLAR).union(m.G_C_STORAGE):
            for i in m.I:
                m.x_c[g, i].fix(m.FIXED_X_C[g, i].value)

        return m

    @staticmethod
    def unfix_continuous_capacity_options(m):
        """Unfix variables used in continuous capacity sizing options"""
        for g in m.G_C_WIND.union(m.G_C_SOLAR).union(m.G_C_STORAGE):
            for i in m.I:
                # Unfix capacity size to zero for now
                m.x_c[g, i].unfix()

        return m

    @staticmethod
    def fix_uc_integer_variables(m, last_value=True):
        """
        Fix unit commitment integer variables - allows duals to be obtained when re-solving

        Parameters
        ----------
        m : pyomo model object
            Object defining optimisation model

        last_value : bool (default=True)
            Fix variable to value obtained in last solution (e.g. from previous model solution). Else fix to value
            to be same as corresponding parameter.
        """

        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        if last_value:
                            m.SCENARIO[i, o].v[g, t].fix()
                            m.SCENARIO[i, o].w[g, t].fix()
                            m.SCENARIO[i, o].u[g, t].fix()
                        else:
                            m.SCENARIO[i, o].v[g, t].fix(m.SCENARIO[i, o].FIXED_V[g, t].value)
                            m.SCENARIO[i, o].w[g, t].fix(m.SCENARIO[i, o].FIXED_W[g, t].value)
                            m.SCENARIO[i, o].u[g, t].fix(m.SCENARIO[i, o].FIXED_U[g, t].value)
        return m

    @staticmethod
    def unfix_uc_integer_variables(m):
        """Unfix unit commitment integer variables"""

        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].v[g, t].unfix()
                        m.SCENARIO[i, o].w[g, t].unfix()
                        m.SCENARIO[i, o].u[g, t].unfix()
        return m

    @staticmethod
    def fix_uc_integer_variable_on_state(m):
        """Fix on-state integer variable"""

        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].u[g, t].fix(m.SCENARIO[i, o].FIXED_U[g, t].value)

        return m

    @staticmethod
    def fix_uc_integer_variable_startup_state(m):
        """Fix shutdown state integer variable"""

        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].w[g, t].fix(m.SCENARIO[i, o].FIXED_V[g, t].value)

        return m

    @staticmethod
    def fix_uc_integer_variable_shutdown_state(m):
        """Fix shutdown state integer variable"""

        for g in m.G_E_THERM.union(m.G_C_THERM):
            for i in m.I:
                for o in m.O:
                    for t in m.T:
                        m.SCENARIO[i, o].w[g, t].fix(m.SCENARIO[i, o].FIXED_W[g, t].value)

        return m

    @staticmethod
    def fix_uc_continuous_variables(m):
        """Fix unit commitment model continuous  variables"""
        for i in m.I:
            for o in m.O:
                for t in m.T:
                    for g in m.G:

                        # Power output above minimum dispatchable level of output [MW]
                        m.SCENARIO[i, o].p[g, t].fix()

                        if g in m.G_E_THERM.union(m.G_C_THERM).union(m.G_C_STORAGE):
                            # Upward reserve allocation [MW]
                            m.SCENARIO[i, o].r_up[g, t].fix()

                        if g in m.G_C_STORAGE:
                            # Storage unit charging (power in) [MW]
                            m.SCENARIO[i, o].p_in[g, t].fix()

                            # Storage unit discharging (power out) [MW]
                            m.SCENARIO[i, o].p_out[g, t].fix()

                            # Energy in storage unit [MWh]
                            m.SCENARIO[i, o].y[g, t].fix()

                    for l in m.L:
                        # Powerflow between NEM zones [MW]
                        m.SCENARIO[i, o].p_flow[l, t].fix()

                    for z in m.Z:
                        # Lost load - up [MW]
                        m.SCENARIO[i, o].p_lost_up[z, t].fix()

                        # Lost load - down [MW]
                        m.SCENARIO[i, o].p_lost_down[z, t].fix()

        return m

    def _unfix_variables(self, m):
        """Unfix subproblem variables - return model variable fixed state to what it was prior to first solve"""

        # Unfix policy variables
        m = self.unfix_permit_price(m)
        m = self.unfix_baseline(m)

        # Unfix capacity sizing option variables
        m = self.unfix_discrete_capacity_options(m)
        m = self.unfix_continuous_capacity_options(m)

        # Unfix UC integer variables
        m = self.unfix_uc_integer_variables(m)

        return m

    def update_subproblem_parameters(self, m, master_results):
        """Update parameters - values obtained from master problem"""

        # Policy variables
        m = self.update_fixed_baseline(m, master_results['baseline'])
        m = self.update_fixed_permit_price(m, master_results['permit_price'])

        # Capacity sizing variables
        m = self.update_fixed_capacity_discrete(m, master_results['d'])
        m = self.update_fixed_capacity_continuous(m, master_results['x_c'])

        # UC binary variables
        m = self.update_fixed_on_state(m, master_results['u'])
        m = self.update_fixed_startup_state(m, master_results['v'])
        m = self.update_fixed_shutdown_state(m, master_results['w'])

        return m

    @staticmethod
    def save_results(results):
        """
        Save (potentially intermediary) results

        Parameters
        ----------
        results : dict
            Sub-problem results stored in a dictionary
        """

        # Save results
        with open('subproblem_results.pickle', 'wb') as f:
            pickle.dump(results, f)

    def solve_model(self, m):
        """Solve model"""

        # Solve model
        self.opt.solve(m, tee=True, options=self.solver_options, keepfiles=self.keepfiles)

        # Log infeasible constraints if they exist
        log_infeasible_constraints(m)

        return m

    def solve_subproblem(self, m, master_results):
        """Solve subproblem"""

        # Update model parameters (values obtained from master problem solution)
        # ----------------------------------------------------------------------
        self.logger.info('Updating subproblem parameters')
        m = self.update_subproblem_parameters(m, master_results)
        self.logger.info('Finished updating subproblem parameters')

        # Solve model for primal variables (energy output etc.)
        # -----------------------------------------------------
        # Fix master problem variables
        self.logger.info('Fixing master problem variables')
        m = self.fix_permit_price(m)
        m = self.fix_baseline(m)
        m = self.fix_discrete_capacity_options(m)
        m = self.fix_continuous_capacity_options(m)
        m = self.fix_uc_integer_variables(m, last_value=False)
        self.logger.info('Finished fixing master problem variables')

        # Solve model
        self.logger.info('Solving model for primal variables')
        start = time.time()
        m = self.solve_model(m)
        primal_variable_results = self._extract_primal_varialble_results(m)
        print(f'Solved model for primal variables in: {time.time() - start}s')
        self.logger.info(f'Solved model for primal variables in: {time.time() - start}s')

        # Get sensitivities for fixed capacity sizing options
        # ---------------------------------------------------
        self.logger.info(f'Fixing variables to obtain capacity sizing option sensitivities')

        # Fix UC integer variables
        m = self.fix_uc_integer_variables(m, last_value=True)

        # Unfix capacity variables
        m = self.unfix_continuous_capacity_options(m)
        m = self.unfix_discrete_capacity_options(m)

        # Solve model
        self.logger.info(f'Solving model to obtain capacity sizing option sensitivities')
        start = time.time()
        m = self.solve_model(m)
        print(f'Solved model for capacity sizing sensitivities in: {time.time() - start}s')
        self.logger.info(f'Solved model to obtain capacity sizing option sensitivities in: {time.time() - start}s')

        capacity_dual_variable_results = self._extract_capacity_dual_variable_results(m)

        # Get sensitivities for permit price and baseline
        # -----------------------------------------------
        self.logger.info('Get sensitivities for permit price and baseline')

        # Fix capacity variables
        m = self.fix_continuous_capacity_options(m)
        m = self.fix_discrete_capacity_options(m)

        # Fix all unit commitment primal variables
        m = self.fix_uc_continuous_variables(m)

        # Unfix permit price and baseline
        m = self.unfix_permit_price(m)
        m = self.unfix_baseline(m)

        # Solve model
        self.logger.info('Solving model for permit price and baseline sensitivities')
        start = time.time()
        m = self.solve_model(m)
        print(f'Solved model for policy variable sensitivities in: {time.time() - start}s')
        self.logger.info(f'Solved model for permit price and baseline sensitivities in: {time.time() - start}s')

        policy_variable_results = self._extract_policy_dual_variable_results(m)

        # Get sensitivities for UC binary variables
        # -----------------------------------------
        self.logger.info('Obtaining sensitivities for UC binary variables')

        # Fixing policy variables
        m = self.fix_permit_price(m)
        m = self.fix_baseline(m)

        # Unfix UC integer variables
        m = self.unfix_uc_integer_variables(m)

        # Solve model
        self.logger.info('Solving model to obtain sensitivities for UC binary variables')
        start = time.time()
        m = self.solve_model(m)
        print(f'Solved model for UC integer variable sensitivities in: {time.time() - start}')
        self.logger.info(f'Solved model to obtain sensitivities for UC binary variables in: {time.time() - start}s')

        uc_integer_variable_results = self._extract_uc_integer_dual_variable_results(m)

        # Combine all results
        results = {'primal': primal_variable_results,
                   'dual': {**capacity_dual_variable_results, **policy_variable_results, **uc_integer_variable_results}}

        return m, results

    @staticmethod
    def _extract_primal_varialble_results(m):
        """Extract results"""

        results = {
            'ENERGY': {
                i: {o: {key: val.expr() for key, val in m.SCENARIO[i, o].ENERGY.items()} for o in m.O} for i in m.I},
            'ENERGY_IN':
                {i: {o: {key: val.expr() for key, val in m.SCENARIO[i, o].ENERGY_IN.items()} for o in m.O} for i in
                 m.I},
            'ENERGY_OUT':
                {i: {o: {key: val.expr() for key, val in m.SCENARIO[i, o].ENERGY_OUT.items()} for o in m.O} for i in
                 m.I},
            'p_lost_up': {i: {o: m.SCENARIO[i, o].p_lost_up.get_values() for o in m.O} for i in m.I},
            'p_lost_down': {i: {o: m.SCENARIO[i, o].p_lost_down.get_values() for o in m.O} for i in m.I},
            'upward_reserve_violation':
                {i: {o: m.SCENARIO[i, o].upward_reserve_violation.get_values() for o in m.O} for i in m.I},
        }

        return results

    @staticmethod
    def _extract_capacity_dual_variable_results(m):
        """Extract sensitivities associated with capacity sizing decisions"""

        results = {'x_c': {i: {g: m.dual[m.FIXED_CAPACITY_CONT[(g, i)]] for g in
                               m.G_C_STORAGE.union(m.G_C_WIND).union(m.G_C_SOLAR)} for i in m.I},
                   'd': {key: m.dual[m.FIXED_CAPACITY_DISC[key]] for key in
                         m.G_C_THERM.cross(m.I).cross(m.G_C_THERM_SIZE_OPTIONS)}
                   }

        return results

    @staticmethod
    def _extract_policy_dual_variable_results(m):
        """Extract sensitivities associated with the emissions intensity baseline and the permit price"""

        results = {'permit_price': {i: m.dual[m.FIXED_PERMIT_PRICE_CONS[i]] for i in m.I},
                   'baseline': {i: m.dual[m.FIXED_BASELINE_CONS[i]] for i in m.I},
                   }

        return results

    @staticmethod
    def _extract_uc_integer_dual_variable_results(m):
        """Extract sensitivities associated with integer variables in the UC problem"""

        results = {
            'u': {i: {
                o: {key: m.dual[m.SCENARIO[i, o].FIXED_U_CONS[key]] for key in
                    m.G_E_THERM.union(m.G_C_THERM).cross(m.T)}
                for o in m.O} for i in m.I},
            'v': {i: {
                o: {key: m.dual[m.SCENARIO[i, o].FIXED_V_CONS[key]] for key in
                    m.G_E_THERM.union(m.G_C_THERM).cross(m.T)}
                for o in m.O} for i in m.I},
            'w': {i: {
                o: {key: m.dual[m.SCENARIO[i, o].FIXED_W_CONS[key]] for key in
                    m.G_E_THERM.union(m.G_C_THERM).cross(m.T)}
                for o in m.O} for i in m.I},
        }

        return results

if __name__ == '__main__':
    # Start timer
    start_timer = time.time()

    # Define object used to construct subproblem model
    subproblem = Subproblem()

    # Construct subproblem model
    subproblem_model = subproblem.construct_model()

    # Prepare to read suffix values (dual values)
    subproblem_model.dual = Suffix(direction=Suffix.IMPORT)
    print(f'Constructed model in: {time.time() - start_timer}s')

    # Data from master problem
    with open('master_results.pickle', 'rb') as f:
        master_results = pickle.load(f)

    # Solve subproblem
    subproblem.solve_subproblem(subproblem_model, master_results)

max_stars_repo_name: akxen/rep-gep

max_issues_repo_name: akxen/rep-gep

max_forks_repo_name: akxen/rep-gep

reason: from pyomo


text:
"""
mpi_compare_phases_improvements.py: generate the phases comparision pdf files.
"""
from glob import glob
from os.path import basename, join

import click
import numpy as np
from mpi4py import MPI

from ...setting import SURFACE_THRESHOLD
from ...utils.load_files import load_pickle
from .compare_phases_improvements import kernel

comm = MPI.COMM_WORLD  # pylint: disable=c-extension-no-member
size = comm.Get_size()
rank = comm.Get_rank()

def get_all_gcmtids_with_filtered(misfit_windows_dir, info_dir, plot_surface):
    all_misfit_windows_path = sorted(glob(join(misfit_windows_dir, "*pkl")))
    all_gcmtids = [basename(item).split(".")[0]
                   for item in all_misfit_windows_path]
    evdp_dict = load_pickle(join(info_dir, "extra.evdp.pkl"))
    surface_result = {}
    for each_gcmtid in all_gcmtids:
        rep_net_sta = list(evdp_dict[each_gcmtid].keys())[0]
        depth = evdp_dict[each_gcmtid][rep_net_sta]
        if (depth <= SURFACE_THRESHOLD):
            surface_result[each_gcmtid] = True
        else:
            surface_result[each_gcmtid] = False
    if (plot_surface):
        final_gcmtids = []
        for each_gcmtid in all_gcmtids:
            if (surface_result[each_gcmtid]):
                final_gcmtids.append(each_gcmtid)
    else:
        final_gcmtids = all_gcmtids
    return final_gcmtids

def get_gcmtids_each_rank(all_gcmtids):
    return np.array_split(all_gcmtids, size)[rank]

@click.command()
@click.option('--obs_dir', required=True, type=str)
@click.option('--syn_dir1', required=True, type=str)
@click.option('--syn_dir2', required=True, type=str)
@click.option('--azimuth_width', required=True, type=int)
@click.option('--output_dir', required=True, type=str)
@click.option('--waves_perpage', required=True, type=int)
@click.option('--info_dir', required=True, type=str)
@click.option('--misfit_windows_dir1', required=True, type=str)
@click.option('--misfit_windows_dir2', required=True, type=str)
@click.option('--snr', required=True, type=float)
@click.option('--cc', required=True, type=float)
@click.option('--deltat', required=True, type=float)
@click.option('--band', required=True, type=str)
@click.option('--plot_surface', is_flag=True, default=False)
def main(obs_dir, syn_dir1, syn_dir2, azimuth_width, output_dir, waves_perpage, info_dir, misfit_windows_dir1,
         misfit_windows_dir2, snr, cc, deltat, band, plot_surface):
    all_gcmtids = get_all_gcmtids_with_filtered(
        misfit_windows_dir1, info_dir, plot_surface)
    gcmtids_this_rank = get_gcmtids_each_rank(all_gcmtids)
    min_time, max_time = map(int, band.split(","))
    for each_gcmtid in gcmtids_this_rank:
        # * build up parameters
        obs_asdf = join(
            obs_dir, f"{each_gcmtid}.preprocessed_{min_time}s_to_{max_time}s.h5")
        syn_asdf_1 = join(
            syn_dir1, f"{each_gcmtid}.preprocessed_{min_time}s_to_{max_time}s.h5")
        syn_asdf_2 = join(
            syn_dir2, f"{each_gcmtid}.preprocessed_{min_time}s_to_{max_time}s.h5")
        misfit_windows_path_1 = join(
            misfit_windows_dir1, f"{each_gcmtid}.pkl")
        misfit_windows_path_2 = join(
            misfit_windows_dir2, f"{each_gcmtid}.pkl")
        # * now we call the plotting script to plot
        kernel(obs_asdf, syn_asdf_1, syn_asdf_2, azimuth_width, output_dir, waves_perpage,
               info_dir, misfit_windows_path_1, misfit_windows_path_2, snr, cc, deltat, False, plot_surface)

if __name__ == "__main__":
    main()  # pylint: disable=no-value-for-parameter

ext: py

max_stars_repo_name: ziyixi/seisflow

max_issues_repo_name: ziyixi/seisflow

max_forks_repo_name: ziyixi/seisflow

reason: import numpy


text:
#!/usr/bin/env python

import util
import socket
import numpy as np
from sbf import sbf
from random import shuffle

class SbfUser(object):
    """User Class of the location privacy protcol.
       The protocol is realized in the loop() function of this class.

       1. The user receives the following sbf info:
            - the encrypted and precomputed SBF Enc(b#)
            - the set of k hash functions H
            - the value m
            - the conventional grid E

       2. The user determines his geographic position (long. lat.) at regular
           time intervals or whenever a specific application requires it, then he has to
           selects the corresponding grid region e_u in E. (In this implementation
           we use, as convetional grid, longitude and latitute trimmed at the
           fourth decimal)
           He creates the sbf b_u#, using the received sbf info, over e_u and
           counts the numbers of non zero values therein.

       3. The user computes e# = Enc(b#)*b_u# using the homomorphic properties
           of the Pailler cryptosystem. Then applies a random permutation over the
           e#. After that he send:
            - non zero values
            - e#
    """
    def __init__(self, HOST="localhost", PORT=2324):
        """This method creates the socket stream."""
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((HOST, PORT))

    def close_stream(self):
        """This method closes the socket stream."""
        self.sock.close()

    def generate_sbf(self, position, sbf_data):
        """Generates client sbf over the user's position.

        Args:
            position (str): longitude latitude data of the user's postion
            sbf_data (dict): data for the sbf creation

        Returns:
            tuple: the user's sbf (list) and the number of non zero values
            therein (int)
        """
        user_sbf = sbf(sbf_data['bit_mapping'], sbf_data['hash_family'], sbf_data['hash_runs'], 1, sbf_data['salt'])
        user_sbf.insert(position, 1)
        non_zero = np.count_nonzero(user_sbf.filter == 1)
        return user_sbf, non_zero

    def mask_message(self, user_sbf_vector, provider_sbf_vector_enc):
        """Computes element-wise multiplication:
            mask = provider_sbf_vector_enc * user_sbf_vector.

        Args:
            user_sbf_vector (list) the user's sbf
            provider_sbf_vector_enc (list) the encrypted provider's sbf

        Returns:
            mask (list)
            """
        mask = []
        for i, enc in enumerate(provider_sbf_vector_enc):
            mask.append(user_sbf_vector[i] * enc)
        shuffle(mask)
        return mask

    def loop(self):
        """Main user's method. It realize the comunication"""
        util.packet_send(self.sock, 'S')
        sbf_data, _ = util.packet_recv(self.sock)
        print('Received sbf data')

        while True:
            position = input("Enter your position (longitude, latitude) as [xx.xxxx#yy.yyyy]\n").rstrip()
            user_sbf, non_zero = self.generate_sbf(position, sbf_data)
            mask = self.mask_message(user_sbf.filter, sbf_data['sbf_vector_enc'])
            data = {
                    "non_zero": non_zero,
                    "mask": mask
            }
            util.packet_send(self.sock, 'C', data)
        self.close_stream()

def main():
    try:
        u = SbfUser()
        u.loop()
    except socket.error:
        print('socket error')
    finally:
        u.close_stream()

if __name__ == '__main__':
    main()

max_stars_repo_name: GloriaSam/Location-Privacy-Protocol

max_issues_repo_name: GloriaSam/Location-Privacy-Protocol

max_forks_repo_name: GloriaSam/Location-Privacy-Protocol


text:
from sklearn.model_selection import train_test_split, ShuffleSplit
from sklearn.utils import shuffle
import numpy as np
from models_SS import *
from datasets import *
from time import time
from Laplacian import laplacian

# set random seed
np.random.seed(42)

# hyperparameters
lams = [10**i for i in range(-6, 7)] # regularization parameter, lambda
lap_lams = [[10**i, 10**j] for i in range(-6, 7) for j in range(-6, 7)] # permutation of all possible lams

def partition_data(X_train, y_train, partition):
    L, U, V, T = partition[0], partition[1], partition[2], partition[3]
    _labeled = []
    _validation = []
    used = set()
    # fill up labeled set with one element from each class
    present = []
    for i in range(len(X_train)):
        if len(_labeled) == L:
            break
        if (y_train[i] not in present) and (i not in used):
            _labeled.append(i)
            present.append(y_train[i])
            used.add(i)
    # fill up validation set with one element from each class
    present = []
    for i in range(len(X_train)):
        if len(_validation) == V:
            break
        if (y_train[i] not in present) and (i not in used):
            _validation.append(i)
            present.append(y_train[i])
            used.add(i)
    # remaining training samples
    remaining = [i for i in range(len(X_train)) if i not in used]
    # print(len(remaining))
    lab_missing = L - len(_labeled)
    val_missing = V - len(_validation)
    # labeled
    if lab_missing <= 0:
        X_lab, y_lab = X_train[_labeled], y_train[_labeled]
        X_remain, y_remain = X_train[remaining], y_train[remaining]
    else:
        X_lab, X_remain, y_lab, y_remain = train_test_split(X_train[remaining], y_train[remaining], train_size=lab_missing)
        X_lab, y_lab = np.append(X_lab, X_train[_labeled], 0), np.append(y_lab, y_train[_labeled], 0)
    # validation, unlabelled
    X_val, X_unlab, y_val, y_unlab = train_test_split(X_remain, y_remain, train_size=val_missing)
    X_unlab, y_unlab = X_unlab[:U], y_unlab[:U]
    X_val, y_val = np.append(X_val, X_train[_validation], 0), np.append(y_val, y_train[_validation], 0)

    assert len(X_lab) == L
    assert len(y_lab) == L
    assert len(X_unlab) == U
    assert len(y_unlab) == U
    assert len(X_val) == V
    assert len(y_val) == V

    return X_lab, y_lab, X_unlab, y_unlab, X_val, y_val

def run_SS(dataset, model_class, lam=1, n_layer=1, activation='sigmoid'):
    n_reps = 1
    n_node = 2000 # num of nodes in hidden layer
    w_range = [-1, 1] # range of random weights
    b_range = [-1, 1] # range of random biases

    if dataset == g50c:
        NN = 50
        # sigma = 17.5
        partition = [50, 314, 50, 136]
    elif dataset in [coil20b, coil20]:
        NN = 2
        # sigma = 0.6
        partition = [40, 1000, 40, 360]
    elif dataset in [uspstb, uspst]:
        NN = 15
        # sigma = 9.4
        partition = [50, 1409, 50, 498]
    else:
        return 0

    data, label, n_class = dataset()

    unlab_accs = []
    val_accs = []
    test_accs = []
    train_time = []

    data, label = shuffle(data, label)
    # Cross-validation split
    for i in range(n_reps):
        ss = ShuffleSplit(n_splits=4, test_size=partition[3], random_state=42)
        for i, ss_values in enumerate(ss.split(data, label)):
            train_index, test_index = ss_values
            X_train, X_test = data[train_index], data[test_index]
            y_train, y_test = label[train_index], label[test_index]

            assert len(X_test) == partition[3]
            assert len(y_test) == partition[3]

            # partition data to labelled, unlabelled, validation
            X_lab, y_lab, X_unlab, y_unlab, X_val, y_val = partition_data(X_train, y_train, partition)
            # predict unlabelled dataset
            X_train, y_train = X_lab, y_lab
            t = time()
            if model_class in [ELM, RVFL, DeepRVFL, EnsembleDeepRVFL]:
                model = model_class(n_node, lam, w_range, b_range, n_layer, activation=activation)
            elif model_class in [LapELM, LapRVFL, LapDeepRVFL, LapEnsembleDeepRVFL]:
                X_train, y_train = np.vstack([X_lab, X_unlab]), np.append(y_lab, [False]*len(y_unlab), 0)
                L = laplacian(X_lab, X_unlab, NN, sigma=2)
                model = model_class(n_node, lam, w_range, b_range, NN, L, n_layer, activation=activation)
            # elif model_class in [BRVFL, BDeepRVFL, BEnsembleDeepRVFL]:
            #     model = model_class(n_node, lam, w_range, b_range, n_layer, tol=10**(-3), activation=activation)
            model.train(X_train, y_train, n_class)
            train_time.append(time() - t)
            unlab_acc, unlab_pred = model.eval(X_unlab, y_unlab)
            # print('Unlabelled Accuracy: ', unlab_acc)
            unlab_accs.append(unlab_acc)
            # retrain on y_lab and unlab prediction, eval on val set
            new_X_train, new_y_train = np.vstack([X_lab, X_unlab]), np.append(y_lab, unlab_pred, 0)
            if model_class in [RVFL, DeepRVFL, EnsembleDeepRVFL]:
                model = model_class(n_node, lam, w_range, b_range, n_layer, activation=activation)
            elif model_class in [LapRVFL, LapDeepRVFL, LapEnsembleDeepRVFL]:
                model = model_class(n_node, lam, w_range, b_range, NN, L, n_layer, activation=activation)
            # elif model_class in [BRVFL, BDeepRVFL, BEnsembleDeepRVFL]:
            #     model = model_class(n_node, lam, w_range, b_range, n_layer, tol=10**(-5), activation=activation)
            t = time()
            model.train(new_X_train, new_y_train, n_class)
            train_time.append(time() - t)
            val_acc, _ = model.eval(X_val, y_val)
            val_accs.append(val_acc)

            test_acc, _ = model.eval(X_test, y_test)
            # print('Test accuracy: ', test_acc)val_acc
            test_accs.append(test_acc)
        break

    # find average acc
    mean_unlab_acc = 1 - np.mean(unlab_accs), np.std(unlab_accs)
    mean_val_acc = 1 - np.mean(val_accs), np.std(val_accs)
    mean_test_acc = 1 - np.mean(test_accs), np.std(test_accs)
    mean_train_time = np.mean(train_time)
    # print('Validation Accuracy', mean_val_acc)
    # print('Test accuracy: ', mean_test_acc)

    return mean_unlab_acc*100, mean_val_acc*100, mean_test_acc*100, mean_train_time

def run_all(dataset):
    run_ELM(dataset)
    run_RVFL(dataset)
    run_dRVFL(dataset)
    run_edRVFL(dataset)
    run_LapELM(dataset)
    run_LapRVFL(dataset)
    run_LapdRVFL(dataset)
    run_LapedRVFL(dataset)

def run_ELM(dataset):
    print('running ELM...')
    model_class = ELM
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=1, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_RVFL(dataset):
    print('running RVFL...')
    model_class = RVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=1, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_dRVFL(dataset):
    print('running Deep RVFL...')
    model_class = DeepRVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=10, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_edRVFL(dataset):
    print('running Ensemble Deep RVFL...')
    model_class = EnsembleDeepRVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=10, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_LapELM(dataset):
    print('running Laplacian ELM...')
    model_class = LapELM
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lap_lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=1, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lap_lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_LapRVFL(dataset):
    print('running Laplacian RVFL...')
    model_class = LapRVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lap_lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=1, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lap_lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_LapdRVFL(dataset):
    print('running Laplacian Deep RVFL...')
    model_class = LapDeepRVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lap_lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=10, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lap_lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

def run_LapedRVFL(dataset):
    print('running Laplacian Ensemble Deep RVFL...')
    model_class = LapEnsembleDeepRVFL
    unlab = []
    val = []
    test = []
    t = []
    for i, lam in enumerate(lap_lams):
        unlab_acc, val_acc, test_acc, train_time = run_SS(dataset, model_class, lam=lam, n_layer=10, activation='sigmoid')
        unlab.append(unlab_acc)
        val.append(val_acc)
        test.append(test_acc)
        t.append(train_time)

    max_index = np.argmin(val, axis=0)[0]
    opt_lam = lap_lams[max_index]
    print('Unlab: ', unlab[max_index][0], u"\u00B1", unlab[max_index][1])
    print('Val: ', val[max_index][0], u"\u00B1", val[max_index][1])
    print('Test: ', test[max_index][0], u"\u00B1", test[max_index][1])
    print('Lambda: ', opt_lam)
    print('Train time: ', t[max_index])

# def run_BRVFL(data, label, n_class, partition):
#     print('running Bayesian RVFL...')
#     model_class = BRVFL
#     unlab_acc, val_acc, test_acc, train_time = run_SS(data, label, n_class, model_class, partition)

#     print('Unlab: ', unlab_acc[0], u"\u00B1", unlab_acc[1])
#     print('Val: ', val_acc[0], u"\u00B1", val_acc[1])
#     print('Test: ', test_acc[0], u"\u00B1", test_acc[1])
#     print('Train time: ', train_time)

# def run_BdRVFL(data, label, n_class, partition):
#     print('running Bayesian Deep RVFL...')
#     model_class = BDeepRVFL
#     unlab_acc, val_acc, test_acc, train_time = run_SS(data, label, n_class, model_class, partition, n_layer=10)

#     print('Unlab: ', unlab_acc[0], u"\u00B1", unlab_acc[1])
#     print('Val: ', val_acc[0], u"\u00B1", val_acc[1])
#     print('Test: ', test_acc[0], u"\u00B1", test_acc[1])
#     print('Train time: ', train_time)

# def run_BedRVFL(data, label, n_class, partition):
#     print('running Bayesian Ensemble Deep RVFL...')
#     model_class = BEnsembleDeepRVFL
#     unlab_acc, val_acc, test_acc, train_time = run_SS(data, label, n_class, model_class, partition, n_layer=10)

#     print('Unlab: ', unlab_acc[0], u"\u00B1", unlab_acc[1])
#     print('Val: ', val_acc[0], u"\u00B1", val_acc[1])
#     print('Test: ', test_acc[0], u"\u00B1", test_acc[1])
#     print('Train time: ', train_time)


text:
import numpy as np
import torch

from environments import Env

class DreamerMemory:
    def __init__(self, capacity, sequence_length, action_size, obs_size, n_agents, device, env_type):
        self.capacity = capacity
        self.sequence_length = sequence_length
        self.action_size = action_size
        self.obs_size = obs_size
        self.device = device
        self.env_type = env_type
        self.init_buffer(n_agents, env_type)

    def init_buffer(self, n_agents, env_type):
        self.observations = np.empty((self.capacity, n_agents, self.obs_size), dtype=np.float32)
        self.actions = np.empty((self.capacity, n_agents, self.action_size), dtype=np.float32)
        self.av_actions = np.empty((self.capacity, n_agents, self.action_size),
                                   dtype=np.float32) if env_type == Env.STARCRAFT else None
        self.rewards = np.empty((self.capacity, n_agents, 1), dtype=np.float32)
        self.dones = np.empty((self.capacity, n_agents, 1), dtype=np.float32)
        self.fake = np.empty((self.capacity, n_agents, 1), dtype=np.float32)
        self.last = np.empty((self.capacity, n_agents, 1), dtype=np.float32)
        self.next_idx = 0
        self.n_agents = n_agents
        self.full = False

    def append(self, obs, action, reward, done, fake, last, av_action):
        if self.actions.shape[-2] != action.shape[-2]:
            self.init_buffer(action.shape[-2], self.env_type)
        for i in range(len(obs)):
            self.observations[self.next_idx] = obs[i]
            self.actions[self.next_idx] = action[i]
            if av_action is not None:
                self.av_actions[self.next_idx] = av_action[i]
            self.rewards[self.next_idx] = reward[i]
            self.dones[self.next_idx] = done[i]
            self.fake[self.next_idx] = fake[i]
            self.last[self.next_idx] = last[i]
            self.next_idx = (self.next_idx + 1) % self.capacity
            self.full = self.full or self.next_idx == 0

    def tenzorify(self, nparray):
        return torch.from_numpy(nparray).float()

    def sample(self, batch_size):
        return self.get_transitions(self.sample_positions(batch_size))

    def process_batch(self, val, idxs, batch_size):
        return torch.as_tensor(val[idxs].reshape(self.sequence_length, batch_size, self.n_agents, -1)).to(self.device)

    def get_transitions(self, idxs):
        batch_size = len(idxs)
        vec_idxs = idxs.transpose().reshape(-1)
        observation = self.process_batch(self.observations, vec_idxs, batch_size)[1:]
        reward = self.process_batch(self.rewards, vec_idxs, batch_size)[:-1]
        action = self.process_batch(self.actions, vec_idxs, batch_size)[:-1]
        av_action = self.process_batch(self.av_actions, vec_idxs, batch_size)[1:] if self.env_type == Env.STARCRAFT else None
        done = self.process_batch(self.dones, vec_idxs, batch_size)[:-1]
        fake = self.process_batch(self.fake, vec_idxs, batch_size)[1:]
        last = self.process_batch(self.last, vec_idxs, batch_size)[1:]

        return {'observation': observation, 'reward': reward, 'action': action, 'done': done,
                'fake': fake, 'last': last, 'av_action': av_action}

    def sample_position(self):
        valid_idx = False
        while not valid_idx:
            idx = np.random.randint(0, self.capacity if self.full else self.next_idx - self.sequence_length)
            idxs = np.arange(idx, idx + self.sequence_length) % self.capacity
            valid_idx = self.next_idx not in idxs[1:]  # Make sure data does not cross the memory index
        return idxs

    def sample_positions(self, batch_size):
        return np.asarray([self.sample_position() for _ in range(batch_size)])

    def __len__(self):
        return self.capacity if self.full else self.next_idx

    def clean(self):
        self.memory = list()
        self.position = 0

max_stars_repo_name: jbr-ai-labs/mamba

max_issues_repo_name: jbr-ai-labs/mamba

max_forks_repo_name: jbr-ai-labs/mamba


text:
#	 BASILISK: A generative, probabilistic model of side chains in proteins
#
#	 Copyright (C) 2010		Tim Harder and Jes Frellsen
#
#	 BASILISK 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.
#
#	 BASILISK 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 BASILISK.	 If not, see <http://www.gnu.org/licenses/>.
#
############################################################################

from numpy import zeros, exp, log, transpose, sum, alltrue

one = 0.99

# Likelihood calculator for IOHMM using the Forward algorithm
class LikelihoodIOHMM_Fw:
    """
    Likelihood calculator for IOHMM using the Forward algorithm
    """
    def __init__(self, dbn, hidden_node):
        self.dbn=dbn
        # List of node objects
        nodes_0=self.nodes_0=dbn.nodes_0
        nodes_1=self.nodes_1=dbn.nodes_1

        self.hidden_node = hidden_node
        self.hidden_node_index_cpd = 0
        self.hidden_node_size = nodes_1[hidden_node].node_size

        self.forward = None
        self.parentmap_list = None

        # Assert that the hidden node has same size in both slices
        assert(nodes_0[hidden_node].node_size==nodes_1[hidden_node].node_size)

    def set_data_in_dbn(self, seq):
        dbn = self.dbn
        parentmap_list = ([],[])

        # Put the sequence into the dbn
        for node in dbn.unique_nodes:
            parentmap=node.get_parentmap(seq, weight=1)
            node.set_parentmap(parentmap)
            parentmap_list[0].append(parentmap)

    def calc_ll(self, seq, mismask, start=None, end=None, ignore_child_mism=True, set_data_in_dbn=True):
        # Make some local copies for speed.
        dbn = self.dbn
        nodes_0 = self.nodes_0
        nodes_1 = self.nodes_1
        hidden_node = self.hidden_node
        hnic = self.hidden_node_index_cpd
        hidden_node_size = self.hidden_node_size

        # Assert that seq and mismask has same shape
        assert(seq.shape == mismask.shape)

        # Put the data into the dbn
        if set_data_in_dbn:
            self.set_data_in_dbn(seq)

        # Set the start- and end-variable correct
        if start == None:
            start = 0

        if end == None:
            end = seq.shape[0]

        # Set the rest of the variables
        slice_count = (end-start)
        scales = zeros(slice_count)

        # Calculate the forward array
        forward = zeros((slice_count, hidden_node_size))
        for i,l in enumerate(xrange(start, end)):
            # For the first slice
            if l==0 or i==0:
                if l==0:
                    node=nodes_0[hidden_node]
                    nodes = nodes_0
                else:
                    node=nodes_1[hidden_node]
                    nodes = nodes_1

                # Get the values the parents
                pv = tuple(node.parentmap[l][:-1])

                # Get the transistion probability
                trans_prob = node.cpd[pv]
                forward[i] = trans_prob

            # For all following slices
            else:
                node=nodes_1[hidden_node]
                nodes = nodes_1

                # Get the values the parents
                pv = list(node.parentmap[l][:-1])

                # Slice out the previous value of the hidden node
                pv[hnic] = slice(None)
                pv = tuple(pv)

                # Get the transistion probability
                trans_prob = node.cpd[pv]

                # Multiply transistionsvalues g->h by the
                # forwardvalues [i-1, g] and sum over nodevalues of g
                forward[i] = sum(transpose(trans_prob)*forward[i-1], axis=1)

            # Multiply the forward values by the probability of the observed children
            # Note that the value of the hidden node is preserved in prev_node_value
            prev_node_value = node.parentmap[l][-1]
            for child in node.children_1:
                if ignore_child_mism or mismask[l,child.node_index] < one:
                    for j in xrange(hidden_node_size):
                        node.parentmap[l] = j
                        forward[i,j] *= exp( child.get_slice_log_likelihood(l)[0] )
            node.parentmap[l] = prev_node_value

            # Scale forward to be a probability distribution (this is
            # allowed according to BSA sec. 3.6 (p. 78)
            s=sum(forward[i])
            forward[i]/=s
            scales[i] = s

        self.forward = forward

        return sum(log(scales)), slice_count

max_stars_repo_name: berquist/fragbuilder

max_issues_repo_name: berquist/fragbuilder

max_forks_repo_name: berquist/fragbuilder

reason: from numpy


text:
"""

Path tracking simulation with pure pursuit steering control and PID speed control.

author: Atsushi Sakai (@Atsushi_twi)
        Guillaume Jacquenot (@Gjacquenot)

"""
import numpy as np
import math
import matplotlib.pyplot as plt

k = 0.1  # look forward gain
Lfc = 2.0  # look-ahead distance
Kp = 1.0  # speed proportional gain
dt = 0.1  # [s]
L = 2.9  # [m] wheel base of vehicle

show_animation = True

class State:

    def __init__(self, x=0.0, y=0.0, yaw=0.0, v=0.0):
        self.x = x
        self.y = y
        self.yaw = yaw
        self.v = v
        self.rear_x = self.x - ((L / 2) * math.cos(self.yaw))
        self.rear_y = self.y - ((L / 2) * math.sin(self.yaw))

    def update(self, a, delta):

        self.x += self.v * math.cos(self.yaw) * dt
        self.y += self.v * math.sin(self.yaw) * dt
        self.yaw += self.v / L * math.tan(delta) * dt
        self.v += a * dt
        self.rear_x = self.x - ((L / 2) * math.cos(self.yaw))
        self.rear_y = self.y - ((L / 2) * math.sin(self.yaw))

    def calc_distance(self, point_x, point_y):

        dx = self.rear_x - point_x
        dy = self.rear_y - point_y
        return math.hypot(dx, dy)

class States:

    def __init__(self):
        self.x = []
        self.y = []
        self.yaw = []
        self.v = []
        self.t = []

    def append(self, t , state):
        self.x.append(state.x)
        self.y.append(state.y)
        self.yaw.append(state.yaw)
        self.v.append(state.v)
        self.t.append(t)

def PIDControl(target, current):
    a = Kp * (target - current)

    return a

class Trajectory:
    def __init__(self, cx, cy):
        self.cx = cx
        self.cy = cy
        self.old_nearest_point_index = None

    def search_target_index(self, state):
        if self.old_nearest_point_index is None:
            # search nearest point index
            dx = [state.rear_x - icx for icx in self.cx]
            dy = [state.rear_y - icy for icy in self.cy]
            d = np.hypot(dx, dy)
            ind = np.argmin(d)
            self.old_nearest_point_index = ind
        else:
            ind = self.old_nearest_point_index
            distance_this_index = state.calc_distance(self.cx[ind], self.cy[ind])
            while True:
                ind = ind + 1 if (ind + 1) < len(self.cx) else ind
                distance_next_index = state.calc_distance(self.cx[ind], self.cy[ind])
                if distance_this_index < distance_next_index:
                    break
                distance_this_index = distance_next_index
            self.old_nearest_point_index = ind

        L = 0.0

        Lf = k * state.v + Lfc

        # search look ahead target point index
        while Lf > L and (ind + 1) < len(self.cx):
            L = state.calc_distance(self.cx[ind], self.cy[ind])
            ind += 1

        return ind

def pure_pursuit_control(state, trajectory, pind):

    ind = trajectory.search_target_index(state)

    if pind >= ind:
        ind = pind

    if ind < len(trajectory.cx):
        tx = trajectory.cx[ind]
        ty = trajectory.cy[ind]
    else:
        tx = trajectory.cx[-1]
        ty = trajectory.cy[-1]
        ind = len(trajectory.cx) - 1

    alpha = math.atan2(ty - state.rear_y, tx - state.rear_x) - state.yaw

    Lf = k * state.v + Lfc

    delta = math.atan2(2.0 * L * math.sin(alpha) / Lf, 1.0)

    return delta, ind

def plot_arrow(x, y, yaw, length=1.0, width=0.5, fc="r", ec="k"):
    """
    Plot arrow
    """

    if not isinstance(x, float):
        for ix, iy, iyaw in zip(x, y, yaw):
            plot_arrow(ix, iy, iyaw)
    else:
        plt.arrow(x, y, length * math.cos(yaw), length * math.sin(yaw),
                  fc=fc, ec=ec, head_width=width, head_length=width)
        plt.plot(x, y)

def main():
    #  target course
    cx = np.arange(0, 50, 0.1)
    cy = [math.sin(ix / 5.0) * ix / 2.0 for ix in cx]

    target_speed = 10.0 / 3.6  # [m/s]

    T = 100.0  # max simulation time

    # initial state
    state = State(x=-0.0, y=-3.0, yaw=0.0, v=0.0)

    lastIndex = len(cx) - 1
    time = 0.0
    states = States()
    states.append(time, state)
    trajectory = Trajectory(cx, cy)
    target_ind = trajectory.search_target_index(state)

    while T >= time and lastIndex > target_ind:
        ai = PIDControl(target_speed, state.v)
        di, target_ind = pure_pursuit_control(state, trajectory, target_ind)
        state.update(ai, di)

        time += dt
        states.append(time, state)

        if show_animation:  # pragma: no cover
            plt.cla()
            # for stopping simulation with the esc key.
            plt.gcf().canvas.mpl_connect('key_release_event',
                    lambda event: [exit(0) if event.key == 'escape' else None])
            plot_arrow(state.x, state.y, state.yaw)
            plt.plot(cx, cy, "-r", label="course")
            plt.plot(states.x, states.y, "-b", label="trajectory")
            plt.plot(cx[target_ind], cy[target_ind], "xg", label="target")
            plt.axis("equal")
            plt.grid(True)
            plt.title("Speed[km/h]:" + str(state.v * 3.6)[:4])
            plt.pause(0.001)

    # Test
    assert lastIndex >= target_ind, "Cannot goal"

    if show_animation:  # pragma: no cover
        plt.cla()
        plt.plot(cx, cy, ".r", label="course")
        plt.plot(states.x, states.y, "-b", label="trajectory")
        plt.legend()
        plt.xlabel("x[m]")
        plt.ylabel("y[m]")
        plt.axis("equal")
        plt.grid(True)

        plt.subplots(1)
        plt.plot(states.t, [iv * 3.6 for iv in states.v], "-r")
        plt.xlabel("Time[s]")
        plt.ylabel("Speed[km/h]")
        plt.grid(True)
        plt.show()

if __name__ == '__main__':
    print("Pure pursuit path tracking simulation start")
    main()