text:
"""
.. module:: payslips
    :platform: Unix, Windows
    :synopsis: Paycheck integrations for custom pdf and APIs.
.. moduleauthor:: Albert Ferguson <albertferguson118@gmail.com>
"""

# core
from src.core import environConfig, graphing, images, timeManips

# third party libs
import pandas as pd
from pandas import Timestamp
import numpy as np
from tabula import read_pdf
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.patches as mpatches

# python core
from datetime import datetime
import math
import os
import traceback
import warnings
import sys

def Payslips_getPayslips(basedir: str, subfolder: str, true_col_header_index=5, filename=str(datetime.strftime(datetime.today(), "%d-%m-%Y")+".pdf")) -> (pd.DataFrame, pd.DataFrame):
        """Retreive the payslip pdf and create aggregate and latest_week frames.

        Convert "structured" pdf to frames for easy use later, this is lots of *icky* scraping/conversion code.

        **Args:**
            payslip_name(str): The name of the file to scrape. Default is the download name default.

			true_col_header_index(int): This value is used to relatively find further dataframes from the pdf.\
				The row index where column titles are actually located. This overrides the default behaviour of\
				tabula guessing where this would be otherwise (and being wrong typically).\

				**Yes this is a magic number.**\
				.. note:: No it isn't tested for everything, only for my example with ADP.\
				Inspect this value yourself if the data is incorrectly parsed.\

        **Notes:**
            .. note:: This function is intended to call up the latest payslip for\
                weekly displays, the stats function for income then aggregates data for\
                longer timeframes.\

        **Returns:**
            latest_week_income(pandas.DataFrame): The the dataframe with hourly data, commissions and deductions.\
                ['Description_Hours', 'Rate', 'Hours', 'Value_Hours', 'Description_Other', 'Tax_Ind', 'Value_Other']\

            aggregate_income(pandas.DataFrame): The aggregate income values.\
                ['Gross', 'Taxable Income', 'Post Tax Allows/Deds', 'Tax', 'NET']\
        """

        ########################################################################
        # Internal Utilities
        ########################################################################

        def _rename_headers(dataframe, header_index, cols_default_headers):
            """ Rename the column headers from default guesses to the correct values.

            .. note:: Also performs some housekeeping by reindexing and dropping the header row.

            .. note:: Due to the nature of separating a frame like this, it is possible to create duplicate\
                header titles if _split_merged_columns is applied next, keep this in mind.\
            """

            try:
                i = 0
                for col in cols_default_headers:
                    dataframe.rename(columns={col:str(dataframe.loc[header_index, col])}, inplace=True, copy=False)
                    i -=- 1

                dataframe = dataframe.drop(header_index)
                row_id = list(range(len(dataframe)))
                dataframe["row_id"] = row_id
                dataframe.set_index("row_id", inplace=True)

                if "Tax Ind" in dataframe.columns:
                    dataframe.rename(columns={"Tax Ind": "Tax_Ind"}, inplace=True, copy=False)

                if np.NaN in dataframe.columns:
                    dataframe = dataframe.drop(np.NaN, axis=1)
                return dataframe

            except TypeError:
                print("The header index was not correctly calculated, please check the header for the frame manually.\n")
                traceback.print_exc()
                traceback.print_stack()
                return

            except Exception as e:
                print("An unknown exception occured in renaming the headers of the frame.\n")
                print(type(e), '\n')
                print('Current frame:\n', dataframe, '\n')
                input()
                return

        def _split_merged_columns(dataframe) -> pd.DataFrame:
            hdr_vals = dataframe.columns.tolist() # check first row 'splittable'
            idxs_added = []

            i = 0
            # start by splitting column names and inserting blank columns ready for data
            for val in hdr_vals:
                if ' ' in str(val):
                    new_hdrs = val.split()
                    # insert a new column at this position with NaN type values
                    try:
                        dataframe.insert(i + 1, new_hdrs[1], np.NaN)
                    except ValueError:
                        # Duplicate case, we had a duplicate to insert!
                        # add _Other to distinguish it
                        dataframe.insert(i + 1, new_hdrs[1] + '_Other', np.NaN)

                    # rename the current column and record the insertion idx
                    # edge case, possible that column renamed to existing name
                    if new_hdrs[0] in dataframe.columns.values:
                        idx = dataframe.columns.values.tolist().index(new_hdrs[0])
                        dataframe.columns.values[idx] = str(new_hdrs[0] + '_Other')
                        dataframe.columns.values[i] = new_hdrs[0]
                    else:
                        dataframe.columns.values[i] = new_hdrs[0]

                    idxs_added.append(i)
                    i -=- 2 # skip the col we just made
                else:
                    i -=- 1 # not splitable, skip

            # now split the vals in cols tracked by idxs_added
            # and perform type conversion and string formatting
            for i in range(len(dataframe)):
                # idxs_added tracked which cols need data_points split to vals
                for idx in idxs_added:
                    data_point = dataframe.iloc[i, idx]
                    if type(data_point) == float:
                        continue # skip nan types
                    vals = data_point.split()

                    try:
                        string_val = ''
                        # val is unknown combination of strings, floats (maybe nan) vals
                        length = len(vals)
                        j = 0
                        while j < length:
                            try:
                                vals[j]= float(vals[j])
                                j -=- 1
                            except ValueError:
                                # val was string, apply string formating
                                string_val += ' ' + vals.pop(j).replace('*', '').lower().capitalize()
                                length -= 1

                        if len(string_val) is not 0:
                            # apply final string formating
                            string_val = string_val.strip()

                        if len(vals) > 2:
                            # we dont know what is there then, RuntimeError and hope for the best
                            raise RuntimeError
                        elif len(vals) is 0:
                            vals =tuple([np.nan, string_val])
                        elif len(string_val) is 0:
                            vals = tuple(vals)
                        else:
                            vals = tuple(vals.append(string_val))

                    except Exception as e:
                        # TODO: add logging to log file here
                        # we dont know error, pass and hope it's caught else where
                        print(e)
                        # traceback.print_stack()
                        pass

                    # add the data to the new column
                    dataframe.iloc[i, idx + 1] = vals[1]
                    # then replace the merged values with the single column value
                    dataframe.iloc[i, idx] = vals[0]

            return dataframe.drop(['nan'], axis= 1)

        ########################################################################

        aggregate_income = pd.DataFrame()
        latest_week_income = pd.DataFrame()
        aggregate_income_header_idx = None

        # retrieve the payslip, uses the tabula pdf_reader
        f_dir = os.path.join(basedir, subfolder)
        fn = filename
        payslip_name = os.path.join(f_dir, fn)
        data = read_pdf(payslip_name)[0] # using Tabula for the pdf read out to data (dataframe)

        try:
            data = data.drop(["Unnamed: 0", "Status"], axis=1)
            cols_default_headers = data.columns.values
        except (KeyError, ValueError) as e:
            raise e

        # split the data into new columns where tabula merged them, this must be
        # dynamic as user could work further combinations of work rates, etc...
        for i in range(true_col_header_index, len(data)):
            row_split_check = data.iloc[[i]].isnull()
            if row_split_check.values.any():
                bool_header_NaN_spacing = row_split_check.sum().tolist()

                if bool_header_NaN_spacing.count(0) == 1:
                    # this is the index where we split the data frame from aggregate
                    # and stat's income values, break after saving the new df
                    latest_week_income = data[true_col_header_index:i]
                    aggregate_income = data[i + 1:len(data)]
                    aggregate_income_header_idx = i + 1
                    break

        # use correct titles in row_id = true_col_header_index for column header values
        if latest_week_income.empty or aggregate_income.empty:
            print("A frame was incorrectly initialised.")
            raise ValueError
        else:
            latest_week_income = _rename_headers(
                latest_week_income, true_col_header_index, cols_default_headers)
            aggregate_income = _rename_headers(
                aggregate_income, aggregate_income_header_idx, cols_default_headers)

        try:
            # now the frames have been split horizontally,
            # split vertically where some columns have been merged
            latest_week_income = _split_merged_columns(latest_week_income)
        except Exception  as e:
            print(e, "\nCould not split merged data values of income stats data.")
            pass

        # manually correct the some column titles
        aggregate_income.rename(
            columns={
                aggregate_income.columns[2]: "Pre_Tax_Deds",
                aggregate_income.columns[2]: "Post_Ta_Deds",
            },
            inplace=True,
            copy=False)
        # now we try to remove NaNs from the frames
        # return our corrected frames
        # first elem of col sets type, assume blanks never form in first row
        t_vals = [type(e) for e in list(latest_week_income.loc[0,:])]
        cols = list(latest_week_income.columns.values)
        # latest_week_income = latest_week_income.astype(dict(zip(cols, t_vals)))
        for i, col in enumerate(cols):
            t_type = t_vals[i]
            # determine k: the first idx with NaN in a col
            # i.e. the first True index in a isnull() test
            try:
                k = list(latest_week_income[col].isnull()).index(True)
            except ValueError:
                k = -1

            if k == -1:
                continue
            elif t_type == str:
                latest_week_income.loc[:,col].values[k:] = ''
            elif t_type == float or t_type == np.float64:
                latest_week_income.loc[:,col].values[k:] = 0.0
            else:
                pass

        try:
            # add summative data from latetst_week_income to aggregate_income
            aggregate_income["Total_Hours"] = sum(latest_week_income.Hours)
        except (AttributeError, Exception) as e:
            pass

        return aggregate_income, latest_week_income


text:
import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

kcomm = comm.Split(rank%2, rank)

experiment = 3

if experiment == 1:
    sendbuf = np.ones(5) + rank
    recvbuf = None
    if rank == 4:
        recvbuf = np.empty((size, 5))
    print('RECEIVING DATA')
    comm.Gather(sendbuf, recvbuf, root=0)
    print('RECEIVED SUCCESSFULLY')

    if rank == 0:
        print(len(recvbuf))
        print(len(recvbuf[0]))
        print(recvbuf)

if experiment == 2:
    data = rank
    data = kcomm.gather(data, root=0)
    if rank in [0, 1]:
        print(data)

if experiment == 3:
    recvbuf = None
    if rank == 0:
        recvbuf = np.empty((6, 5, 5), 'f8')
    if rank < 6:
        sendbuf = np.ones((5, 5)) * rank
    else:
        sendbuf = np.array([], 'f8')
    sendcounts = [25, 25, 25, 25, 25, 25, 0, 0]
    comm.Gatherv(sendbuf, (recvbuf, sendcounts), root=0)

    if rank == 0: print(recvbuf)

#MPI.Finalize()

ext: py

reason: import numpy


text:
from typing import Union, IO, Any
import numpy as np

def save(obj: Any, file: Union[str, IO]):
    np.save(file, obj)

def load(file: Union[str, IO]):
    return np.load(file)[()]

max_stars_repo_name: roosephu/boots

max_issues_repo_name: roosephu/boots

max_forks_repo_name: roosephu/boots


text:
import unittest
import shutil

import tensorflow as tf
import numpy as np

from networks import network_sparse
from utils.pruning_utils import SparseLayer

class TestNetworkGraph(unittest.TestCase):

    def test_shapes(self):

        input_size = 20
        n_classes = 5
        sparse_layers = [SparseLayer(
                         values=np.array([1, 2, 3]).astype(np.float32),
                         indices=np.array([[0, 1], [1, 2], [3, 4]]).astype(np.int16),
                         dense_shape=np.array([5, 10]).astype(np.int64),
                         bias=np.array([1, 1, 1, 1, 1]).astype(np.float32)),
                         SparseLayer(
                         values=np.array([1, 2, 5]).astype(np.float32),
                         indices=np.array([[0, 2], [1, 2], [3, 4]]).astype(np.int16),
                         dense_shape=np.array([10, 5]).astype(np.int64),
                         bias=np.array([0, 0, 0, 0, 0]).astype(np.float32))]

        network = network_sparse.FullyConnectedClassifierSparse(
                                                input_size=input_size,
                                                n_classes=n_classes,
                                                sparse_layers=sparse_layers,
                                                model_path='temp',
                                                verbose=False)

        self.assertEqual(network.logits.get_shape().as_list(), [None, 5])
        self.assertEqual(network.loss.get_shape().as_list(), [])

class TestNetworkSaveRestore(unittest.TestCase):

    def test_save_restore(self):

        input_size = 20
        n_classes = 5
        sparse_layers = [SparseLayer(
                         values=np.array([1, 2, 3]).astype(np.float32),
                         indices=np.array([[0, 1], [1, 2], [3, 4]]).astype(np.int16),
                         dense_shape=np.array([5, 10]).astype(np.int64),
                         bias=np.array([1, 1, 1, 1, 1]).astype(np.float32)),
                         SparseLayer(
                         values=np.array([1, 2, 5]).astype(np.float32),
                         indices=np.array([[0, 2], [1, 2], [3, 4]]).astype(np.int16),
                         dense_shape=np.array([10, 5]).astype(np.int64),
                         bias=np.array([0, 0, 0, 0, 0]).astype(np.float32))]

        network = network_sparse.FullyConnectedClassifierSparse(
                                                input_size=input_size,
                                                n_classes=n_classes,
                                                sparse_layers=sparse_layers,
                                                model_path='temp',
                                                verbose=False)

        variables_original = network.sess.run(network.weight_tensors)
        network.save_model()
        network.load_model()
        variables_restored = network.sess.run(network.weight_tensors)

        for original, restored in zip(variables_original, variables_restored):
            self.assertTrue((original == restored).all())

    def tearDown(self):
        shutil.rmtree('temp')


text:
# Table.py
# William Sexton
# Last Modified: 4/5/19
"""
    This is the table module for the DAS-2018 instance of the DAS-framework.
    It contains a table classes that inherits from the AbstractDASReader.
    The classes must contain a method called process.
"""

from operator import add
import numpy as np
from programs.reader.table_reader import AbstractTable

from exceptions import DASConfigError

from constants import CC

class DenseHistogramTable(AbstractTable):
    """
        This class stores the microdata and metadata for the CEF person records.

        This is an OUTDATED class and should not really be used. Use SQLSparTable (or, if you need for some reason, SparTable)
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.data_shape = None
        tv = set(map(lambda v: v.name, self.variables))
        vv = set(map(lambda v: v.name, self.recode_variables)) if self.recode_variables else set()
        tv_and_vv = tv.union(vv)
        for varname in self.histogram_variables:
            if varname not in tv_and_vv:
                msg = f"Histogram variable {varname} for {self.name} is neither in table variables nor in recode variables, check your config file"
                raise DASConfigError(msg, f"{self.name}.{CC.HISTOGRAM}/{self.name}.{CC.GEOGRAPHY}", CC.READER)

    def printNumPartsAndSchema(self, data):
        print(f"{self.__class__.__name__} received microdata dataframe with {data.rdd.getNumPartitions()} partitions and the following schema")
        data.printSchema()

    def process(self, data):
        """
            args:
                a Spark dataframe containing CEF person (or household, or unit) records

            This function performs the following process:
                (1) Convert the data to a RDD.
                (2) row -> (geo_histogram,1)) (Map the row to a geo_histogram, a tuple of the
                    geography variables plus the histogram variables.)
                (3) Reduce by key.
                (4) (geo_histogram,cnt) -> (geo, (histogram, cnt))
                (5) groupbykey: creates (geo, list of (histogram, cnt))
                (6) (geo, list of (histogram, cnt)) -> (geo, ndarray)

            returns: an rdd of (geo,numpy ndarray) pairs.
        """

        self.printNumPartsAndSchema(data)
        return (data.rdd.map(self.create_key_value_pair)
                        .reduceByKey(add)
                        .map(lambda key_value:
                             (key_value[0][:len(self.geography_variables)],
                              (key_value[0][len(self.geography_variables):],
                               key_value[1])))
                        .groupByKey()
                        .mapValues(self.to_ndarray).repartition(self.reader.num_reader_partitions))

    def to_ndarray(self, prehist_list):
        """
            Input:
                list of (idx, cnt) pairs, where idx consists of histogram variables,
                (e.g (qrel, qage, qsex, cenrace,) for some person table
                     or (gqtype, ) for a unit table)

            This function performs the following process:
                (1) Initialize the ndarray.
                (2) Iterate through the list assigning cnt to
                    idx in ndarray except for error idx.

            Output:
                a ndarray, which is the histogram of detail counts
        """
        hist = np.zeros(self.data_shape).astype(int)
        for idx, cnt in prehist_list:
            hist[idx] = cnt
        return hist

    def create_key_value_pair(self, row):
        """
            This creates a key value pair for geography and histogram variables.
            This function needs to check the validity of the row record.

            Input:
                a row of data from the CEF

            Output:
                (geo_histogram_idx, 1), where geo_histogram is a tuple
                    of the geography variable values and
                    the histogram variable values.
        """
        blk_idx = (tuple([row[var] for var in self.geography_variables]) + tuple([int(row[var]) for var in self.histogram_variables]))

        return blk_idx, 1

class UnitFromPersonTable(DenseHistogramTable):
    """ This is an outdated class and should not really be used. Use SQLSparTable (or, if you need for some reason, SparTable) """
    COUNT = "count"  # This is the name used by Spark SQL for the field produced by count() operation

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.uniqueID = self.gettuplewsec(CC.UNIQUE)
        self.data_shape = None

    def process(self, data):
        """
            Input:
                 data: a Spark dataframe (df)

            Output:
                a RDD with block by block counts of housing units and gqs by type
        """
        self.printNumPartsAndSchema(data)
        all_vars = list(self.geography_variables) + list(self.histogram_variables)
        return (data.select(*(all_vars + list(self.uniqueID)))
                    .distinct()
                    .groupBy(*all_vars)
                    .count() # Up to this point it's an SQL query, and violates partitioning, yielding 200. Repartition here, if you need # This is an outdated class and should not really be used. Use SQLSparTable (or, if you need for some reason, SparTable)
                    .rdd
                    .map(lambda row: (tuple([row[var] for var in self.geography_variables]), (tuple([row[var] for var in self.histogram_variables]), row[self.COUNT])))
                    .groupByKey()
                    .mapValues(self.to_ndarray).repartition(self.reader.num_reader_partitions))

max_stars_repo_name: p-b-j/uscb-das-container-public

max_issues_repo_name: p-b-j/uscb-das-container-public

max_forks_repo_name: p-b-j/uscb-das-container-public


text:
#!/usr/bin/env python

import matplotlib.pyplot as plt
import sys, os, time
import numpy as np
from numpy import linspace, meshgrid
from matplotlib import cm
import collections

#from mesh_utils import *
#from unit_cell import UnitCell
#from earth_params import earthRadius, mantleThickness, cmbRadius

#--------------------------------------------------------------------------
def mainSampleMeshWithTwoVarGroups(workDir, targetPct,
                                   nth, nr, plotting,
                                   thL, thLDeg, thR, thRDeg,
                                   rSurf, rCmb, debugPrint,
                                   L, dth, dr):
  ## create list of cell objects for the full mesh
  # fullMeshCells = createCellsFullMesh(nth, nr, thL, thR, rSurf, rCmb, dth, dr)
  #if plotting != "none":
  #  for it in fullMeshCells: it.plotPoints(ax)

  # # store coordinates and label for each pt of the full grid
  # [coordsVp, coordsSp, dofLabelsSp] = storeCoordsAndLabelsFullMesh(fullMeshCells, nth, nr)

  # #   # total number of points in domain
  # totNumDofs = len(coordsVp) + len(coordsSp)

  # if (debugPrint == 1):
  #   print("----------------------------")
  #   print("--- graph for Vp ---")
  #   printDicPretty(gVpPrint)
  #   print("----------------------------")
  #   print("--- coeff for Vp ---")
  #   printDicPretty(coeffVpPrint)
  #   print("----------------------------")
  #   print("--- graph for Sp ---")
  #   printDicPretty(gSpPrint)

  # # -----------------------------------------------------
  # # number of pts where we compute residual or velocity
  # numResidualPtsVp = len(gVpPrint)
  # numResidualPtsSp = len(gSpPrint)
  # numResidualPts   = numResidualPtsVp + numResidualPtsSp
  # print ("numResidualPtsVp = ", numResidualPtsVp)
  # print ("numResidualPtsSp = ", numResidualPtsSp)
  # print ("numResidualPts = ", numResidualPts,
  #        " which is = ", numResidualPts/totNumDofs*100, " % of full mesh")

  #     if (targetPct < 0 ):
  #       print("For random sample mesh you need to set --pct=n")
  #       print("where n>0 and defines the % of full mesh to take")

  #     # get graphs for the sample mesh only
  #     [gVp, coeffVp, gSp] = buildGraphSampleMesh(fullMeshCells, nth, nr, targetPct)

  #     if (debugPrint == 1):
  #       print("----------------------------")
  #       print("--- graph for Vp ---")
  #       printDicPretty(gVp)
  #       print("----------------------------")
  #       print("--- coeff for Vp ---")
  #       printDicPretty(coeffVp)
  #       print("----------------------------")
  #       print("--- graph for Sp ---")
  #       printDicPretty(gSp)

  #     # Now we have graphs for the sample mesh cells for vp and srp/stp dofs.
  #     # These graphs contain the points where we need to compute the residual
  #     # as well as those where only the state is needed.
  #     # These graphs are defined in terms of the GIDs of the full mesh.
  #     # We need to uniquely enumerate all the sample mesh cells
  #     # separating the vp from srp/stp dofs.

  #     # We need to loop over the vp graph and the sp graph, and extract
  #     # all unique GIDs of the srp/stp nodes, since there might be duplicates
  #     # because of the connectivity.
  #     allVpGids = []
  #     allSpGids = []
  #     # loop over the vp graph first
  #     for k,v in gVp.items():
  #       # for vp, append key since it is a gid of a vp dof
  #       allVpGids.append(np.int64(k))
  #       # for Sp, append gids of the connected points since these in the vp graph
  #       # are either stp or srp dof.
  #       for j in v:
  #         allSpGids.append(np.int64(j))

  #     # loop over the graph for the srp/stp dofs
  #     for k,v in gSp.items():
  #       if (k>=0): allSpGids.append(np.int64(k))
  #       # dealing with srp dof, only take the north and south neighbors
  #       if (dofLabelsSp[k] == 1):
  #         allVpGids.append( np.int64(v[1]) )
  #         allVpGids.append( np.int64(v[3]) )
  #       # dealing with stp dof, only take the west and east neighbors
  #       if (dofLabelsSp[k] == 2):
  #         allVpGids.append( np.int64(v[0]) )
  #         allVpGids.append( np.int64(v[2]) )

  #     # remove duplicates and sort
  #     allVpGids = list(dict.fromkeys(allVpGids)); allVpGids.sort()
  #     allSpGids = list(dict.fromkeys(allSpGids)); allSpGids.sort()
  #     if (debugPrint == 1):
  #       print("----------------------------")
  #       print("Sample mesh unique GIDs for Vp:"); print(allVpGids)
  #       print("----------------------------")
  #       print("Sample mesh unique GIDs for srp/stp:"); print(allSpGids)
  #       print("\n")

  #     # -----------------------------------------------------
  #     # The sample mesh graph contains the GIDs
  #     # wrt the FULL mesh because we simply extracted a subset from the full mesh.
  #     # We enumerate the sample mesh points with new indexing
  #     # and create a map of full-mesh gids to new gids

  #     # fm_to_sm_map is such that:
  #     # - key   = GID wrt full mesh
  #     # - value = GID wrt sample mesh
  #     fm_to_sm_map_vp = collections.OrderedDict()
  #     fm_to_sm_map_vp = {pt:i for (pt,i) in zip(allVpGids, range(len(allVpGids)))}
  #     sm_to_fm_map_vp = {v:k for k,v in fm_to_sm_map_vp.items()}

  #     # sm_to_fm_map is such that:
  #     # - key   = GID wrt sample mesh
  #     # - value = GID wrt full mesh
  #     fm_to_sm_map_sp = collections.OrderedDict()
  #     fm_to_sm_map_sp = {pt:i for (pt,i) in zip(allSpGids, range(len(allSpGids)))}
  #     sm_to_fm_map_sp = {v:k for k,v in fm_to_sm_map_sp.items()}

  #     if (debugPrint == 1):
  #       print("----------------------------")
  #       print("--- fm_to_sm / sm_to_fm map for vp ---")
  #       printDicPretty2(fm_to_sm_map_vp, sm_to_fm_map_vp)
  #       print("----------------------------")
  #       print("--- fm_to_sm / sm_to_fm map for sp ---")
  #       printDicPretty2(fm_to_sm_map_sp, sm_to_fm_map_sp)

  #     # -----------------------------------------------------
  #     # print sm_to_fm file for Vp
  #     # -----------------------------------------------------
  #     f1 = open("sm_to_fm_vp.dat","w+")
  #     for k,v in sm_to_fm_map_vp.items():
  #       f1.write("%11d %11d\n" % (k, v))
  #     f1.close()

  #     # -----------------------------------------------------
  #     # print sm_to_fm file for Sp
  #     # -----------------------------------------------------
  #     f1 = open("sm_to_fm_sp.dat","w+")
  #     for k,v in sm_to_fm_map_sp.items():
  #       f1.write("%11d %11d\n" % (k, v))
  #     f1.close()

  #     # now that we have the fm <-> sm mapping, we can assemble the actual sample mesh
  #     # graph for vp that will be printed to file, which consists of the vp GIDs indexed wrt sample mesh
  #     print("------\n")
  #     gVpPrint     = collections.OrderedDict()
  #     coeffVpPrint = collections.OrderedDict()
  #     # loop over the original gVp graph (which still contains the GIDs wrt full mesh)
  #     for (k1,v1), (k2,v2) in zip(gVp.items(), coeffVp.items()):
  #       # k1 = k2 = vp dof global ID wrt full mesh enumeration
  #       # v1 = global IDs for the connecting dofs wrt full mesh enumeration
  #       # v2 = coeffcients for vp stencil at each point

  #       # use mapping to get the global ID of current vp dof wrt sample mesh indexing
  #       smGID = fm_to_sm_map_vp[k1]

  #       # loop over adjecent nodes (which are either srp or stp),
  #       # map their global ID to the ID wrt sample mesh
  #       # (remember the trick for ghost points which are stored as negative indices)
  #       # and store the converted index
  #       gVpPrint[smGID] = [ fm_to_sm_map_sp[i] for i in v1 ]

  #       # for the coeffs we keep everything as is since we only need to map the keys
  #       coeffVpPrint[smGID] = v2

  #     if (debugPrint == 1):
  #       print("----------------------------")
  #       printDicPretty2(gVpPrint, coeffVpPrint)

  #     # now that we have the fm_sm mapping and viceversa we can assemble the sample
  #     # graph for srp/stp that will be printed to file, which consists of GIDs wrt sample mesh
  #     print("------\n")
  #     gSpPrint = collections.OrderedDict()
  #     for k, v in gSp.items():
  #       # k = srp or stp dof global ID wrt full mesh enumeration
  #       # v = global IDs for the connecting dofs wrt full mesh enumeration

  #       # get the global ID wrt sample mesh indexing of current dof
  #       smGID = fm_to_sm_map_sp[k]

  #       stencilGIDs = v
  #       # if the dof being tested is srp, then we only need the gids of the
  #       # north and south vp nodes
  #       if (dofLabelsSp[k] == 1):
  #         stencilGIDs[1] = fm_to_sm_map_vp[v[1]]
  #         stencilGIDs[3] = fm_to_sm_map_vp[v[3]]

  #       # if the dof being tested is stp, then we only need the gids of the
  #       # west and east vp nodes
  #       if (dofLabelsSp[k] == 2):
  #         stencilGIDs[0] = fm_to_sm_map_vp[v[0]]
  #         stencilGIDs[2] = fm_to_sm_map_vp[v[2]]

  #       gSpPrint[smGID] = stencilGIDs

  #     if (debugPrint == 1):
  #       print("----------------------------")
  #       printDicPretty(gSpPrint)
  #       print("----------------------------")

  #     # -----------------------------------------------------
  #     # number of pts where we compute residual or velocity
  #     numResidualPtsVp = len(gVpPrint)
  #     numResidualPtsSp = len(gSpPrint)
  #     numResidualPts   = numResidualPtsVp + numResidualPtsSp
  #     print ("numResidualPtsVp = ", numResidualPtsVp)
  #     print ("numResidualPtsSp = ", numResidualPtsSp)
  #     print ("numResidualPts = ", numResidualPts,
  #            " which is = ", numResidualPts/totNumDofs*100, " % of full mesh")

  #     # -----------------------------------------------------
  #     # number of pts where we store the state
  #     numStatePtsVp = len(allVpGids)
  #     numStatePtsSp = len(allSpGids)
  #     numStatePts   = numStatePtsVp + numStatePtsSp
  #     print ("numStatePtsVp = ", numStatePtsVp)
  #     print ("numStatePtsSp = ", numStatePtsSp)
  #     print ("numStatePts = ", numStatePts,
  #            " which is = ", numStatePts/totNumDofs*100, " % of full mesh")

  #     # -----------------------------------------------------
  #     if plotting != "none" and samplingType == "random":
  #       th_fm = np.asarray([v[0] for (k,v) in coordsVp.items()])
  #       r_fm  = np.asarray([v[1] for (k,v) in coordsVp.items()])
  #       ax.scatter(th_fm, r_fm, marker='o', s=4, edgecolor='r', facecolor='r')
  #       th_fm = np.asarray([v[0] for (k,v) in coordsSp.items()])
  #       r_fm  = np.asarray([v[1] for (k,v) in coordsSp.items()])
  #       ax.scatter(th_fm, r_fm, marker='s', s=4, edgecolor='k', facecolor='k')

  #       # subselect coords for vp points in sample mesh
  #       th_fm = np.asarray([v[0] for (k,v) in coordsVp.items()])
  #       r_fm  = np.asarray([v[1] for (k,v) in coordsVp.items()])
  #       r_sm_vp, th_sm_vp  = r_fm[ list(fm_to_sm_map_vp.keys()) ], th_fm[ list(fm_to_sm_map_vp.keys()) ]
  #       ax1.scatter(th_sm_vp, r_sm_vp, marker='o', s=4, edgecolor='r', facecolor='r')
  #       # for i in range(len(allVpGids)):
  #       #   ax1.text(th_sm_vp[i], r_sm_vp[i], str( fm_to_sm_map_vp[allVpGids[i]] ),
  #       #          verticalalignment='top', horizontalalignment='center', fontsize=14, color='r')

  #       # subselect coords for srp/stp points in sample mesh
  #       th_fm = np.asarray([v[0] for (k,v) in coordsSp.items()])
  #       r_fm  = np.asarray([v[1] for (k,v) in coordsSp.items()])
  #       r_sm_sp, th_sm_sp  = r_fm[ list(fm_to_sm_map_sp.keys()) ], th_fm[ list(fm_to_sm_map_sp.keys()) ]
  #       ax1.scatter(th_sm_sp, r_sm_sp, marker='s', s=4, edgecolor='k', facecolor='k')
  #       # for i in range(len(allSpGids)):
  #       #   ax1.text(th_sm_sp[i], r_sm_sp[i], str( fm_to_sm_map_sp[allSpGids[i]] ),
  #       #          verticalalignment='top', horizontalalignment='center', fontsize=14, color='k')

max_stars_repo_name: fnrizzi/ElasticShearWaves

max_issues_repo_name: fnrizzi/ElasticShearWaves

max_forks_repo_name: fnrizzi/ElasticShearWaves

reason: import numpy,from numpy


text:
import abc
import time
import typing

import numpy as np
import torch

import mantrap.constants
import mantrap.utility.shaping

from .trajopt import TrajOptSolver

class SearchIntermediate(TrajOptSolver, abc.ABC):

    ###########################################################################
    # Optimization ############################################################
    ###########################################################################
    def optimize_core(
        self,
        z0: torch.Tensor,
        ado_ids: typing.List[str],
        tag: str = mantrap.constants.TAG_OPTIMIZATION,
        max_cpu_time: float = mantrap.constants.SEARCH_MAX_CPU_TIME,
        **solver_kwargs
    ) -> typing.Tuple[torch.Tensor, typing.Dict[str, torch.Tensor]]:
        """Optimization function for single core to find optimal z-vector.

        Given some initial value `z0` find the optimal allocation for z with respect to the internally defined
        objectives and constraints. This function is executed in every thread in parallel, for different initial
        values `z0`. To simplify optimization not all agents in the scene have to be taken into account during
        the optimization but only the ones with ids defined in `ado_ids`.

        In general searching algorithms (at least the ones implemented within this project) have a similar
        "frame", which is some termination constraint, here the computation runtime, and some inner optimization
        loop which repeats until the algorithm has either terminated or is done.

        :param z0: initial value of optimization variables.
        :param tag: name of optimization call (name of the core).
        :param ado_ids: identifiers of ados that should be taken into account during optimization.
        :param max_cpu_time: maximal cpu runtime for optimization.
        :returns: z_opt (optimal values of optimization variable vector)
                  optimization_log (logging dictionary for this optimization = self.log)
        """
        # Start stopping conditions (runtime or number of iterations).
        sampling_start_time = time.time()

        # Back-fall solution to avoid infinite looping.
        z_back_fall = np.random.uniform(*self.z_bounds)

        # Then start searching loop for finding more optimal trajectories.
        z_iteration, obj_iteration, iteration, k = None, np.inf, 0, 0
        while True:
            z_candidate, obj_candidate, iteration_candidate, is_finished = self._optimize_inner(
                z_iteration, obj_iteration, iteration, tag, ado_ids)

            # Update iteration variables (z_best, obj_best, iteration) only if the objective has been
            # improved over the iteration.
            if obj_candidate <= obj_iteration:
                z_iteration = z_candidate
                obj_iteration = obj_candidate
                iteration = iteration_candidate

            # If solver claims to be finished, end the iteration before the runtime has exceeded.
            # Additionally check whether loop has terminated already, then do not reset these values,
            # because somehow the loop iterates one more time after breaking for `RandomSearch`.
            run_time = time.time() - sampling_start_time
            if is_finished or run_time > max_cpu_time:
                if z_iteration is None:
                    z_iteration = z_back_fall
                z_best = z_iteration.copy()
                break
            # Update search iteration count.
            k += 1

        # The best sample is re-evaluated for logging purposes, since the last iteration is always assumed to
        # be the best iteration (logging within objective and constraint function).
        self.evaluate(z=z_best, tag=tag, ado_ids=ado_ids)
        ego_controls = self.z_to_ego_controls(z=z_best)
        assert mantrap.utility.shaping.check_ego_controls(ego_controls, t_horizon=self.planning_horizon)
        return ego_controls, self.logger.log

    @abc.abstractmethod
    def _optimize_inner(self, z_best: typing.Union[None, np.ndarray], obj_best: float, iteration: int,
                        tag: str, ado_ids: typing.List[str]
                        ) -> typing.Tuple[np.ndarray, float, int, bool]:
        """Inner optimization/search function.

        :param z_best: best assignment of optimization vector so far.
        :param obj_best: according objective function value.
        :param iteration: current search (outer loop) iteration.
        :param tag: name of optimization call (name of the core).
        :param ado_ids: identifiers of ados that should be taken into account during optimization.
        :returns: updated best z-values, updated best objective, outer loop iteration, termination flag.
        """
        raise NotImplementedError

    def evaluate(self, z: np.ndarray, ado_ids: typing.List[str], tag: str) -> typing.Tuple[float, float]:
        """Evaluate "value" of z-values, i.e. some choice of optimization variable assignment, by computing
        the overall objective value as well as the constraint violation (aka the feasibility of the choice). """
        objective = self.objective(z, tag=tag, ado_ids=ado_ids)
        _, constraint_violation = self.constraints(z, ado_ids=ado_ids, return_violation=True, tag=tag)
        return objective, constraint_violation

    ###########################################################################
    # Optimization formulation parameters #####################################
    ###########################################################################
    @property
    def z_bounds(self) -> typing.Tuple[np.ndarray, np.ndarray]:
        # Find variable bounds for random sampling during search.
        lb, ub = self.optimization_variable_bounds()
        return np.asarray(lb), np.asarray(ub)

max_stars_repo_name: simon-schaefer/mantrap

max_issues_repo_name: StanfordASL/mantrap

max_forks_repo_name: StanfordASL/mantrap


text:
import logging
import sys

import numpy as np
from pytest import fixture

import rasterio

logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)

def test_masks():
    with rasterio.open('tests/data/RGB.byte.tif') as src:
        rm, gm, bm = src.read_masks()
        r, g, b = src.read(masked=False)
        assert not r[rm==0].any()
        assert not g[gm==0].any()
        assert not b[bm==0].any()

def test_masked_true():
    with rasterio.open('tests/data/RGB.byte.tif') as src:
        r, g, b = src.read(masked=True)
        rm, gm, bm = src.read_masks()
        assert (r.mask==~rm.astype('bool')).all()
        assert (g.mask==~gm.astype('bool')).all()
        assert (b.mask==~bm.astype('bool')).all()

def test_masked_none():
    with rasterio.open('tests/data/RGB.byte.tif') as src:
        r, g, b = src.read(masked=True)
        rm, gm, bm = src.read_masks()
        assert (r.mask==~rm.astype('bool')).all()
        assert (g.mask==~gm.astype('bool')).all()
        assert (b.mask==~bm.astype('bool')).all()

@fixture(scope='function')
def tiffs(tmpdir):
    with rasterio.open('tests/data/RGB.byte.tif') as src:
        kwds = src.meta

        del kwds['nodata']
        with rasterio.open(
                str(tmpdir.join('no-nodata.tif')), 'w',
                **kwds) as dst:
            dst.write(src.read(masked=False))

        with rasterio.open(
                str(tmpdir.join('sidecar-masked.tif')), 'w',
                **kwds) as dst:
            dst.write(src.read(masked=False))
            mask = np.zeros(src.shape, dtype='uint8')
            dst.write_mask(mask)

    return tmpdir

def test_masking_no_nodata(tiffs):
    # if the dataset has no defined nodata values, all data is
    # considered valid data. The GDAL masks bands are arrays of
    # 255 values. ``read()`` returns masked arrays where `mask`
    # is False.
    filename = str(tiffs.join('no-nodata.tif'))
    with rasterio.open(filename) as src:

        rgb = src.read(masked=False)
        assert not hasattr(rgb, 'mask')
        r = src.read(1, masked=False)
        assert not hasattr(r, 'mask')

        rgb = src.read(masked=True)
        assert hasattr(rgb, 'mask')
        assert not rgb.mask.any()
        r = src.read(1, masked=True)
        assert hasattr(r, 'mask')
        assert not r.mask.any()

        rgb = src.read(masked=True)
        assert hasattr(rgb, 'mask')
        assert not r.mask.any()
        r = src.read(1, masked=True)
        assert not r.mask.any()

        masks = src.read_masks()
        assert masks.all()

def test_masking_sidecar_mask(tiffs):
    # If the dataset has a .msk sidecar mask band file, all masks will
    # be derived from that file.
    with rasterio.open(str(tiffs.join('sidecar-masked.tif'))) as src:
        rgb = src.read(masked=True)
        assert rgb.mask.all()
        r = src.read(1, masked=True)
        assert r.mask.all()
        masks = src.read_masks()
        assert not masks.any()

max_stars_repo_name: clembou/rasterio

max_issues_repo_name: clembou/rasterio

max_forks_repo_name: clembou/rasterio


text:
import numpy as np
import argparse

from tensorflow.keras.layers import Input, Flatten, Dense, Reshape, Dropout, Embedding, Multiply, Activation, Conv2D, ZeroPadding2D, LocallyConnected2D, Concatenate, GRU, Lambda
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import LeakyReLU
from tensorflow.keras.models import Model, Sequential
from tensorflow.compat.v1.keras.layers import UpSampling2D

import tensorflow as tf
import time
from tensorflow.keras import backend as K

print(tf.__version__)

loops_to_do = int(3E5)

# List of batch sizes to test
batch_size_i_array = np.logspace(0,np.log10(25000),10)

make_x_particles = False
loops = 100

''' Create the network. '''
# NON-PROMPT
generator = Sequential()

generator.add(Dense(1536, input_shape=(100,)))
generator.add(LeakyReLU(alpha=0.2))
generator.add(BatchNormalization(momentum=0.8))
generator.add(Dense(3072))
generator.add(LeakyReLU(alpha=0.2))
generator.add(BatchNormalization(momentum=0.8))
generator.add(Dense(6, activation='tanh'))
generator.add(Reshape((1, 6, 1)))

# PROMPT
# generator = Sequential()
# generator.add(Dense(512, input_shape=(100,)))
# generator.add(LeakyReLU(alpha=0.2))
# generator.add(BatchNormalization(momentum=0.8))
# generator.add(Dense(1024))
# generator.add(LeakyReLU(alpha=0.2))
# generator.add(BatchNormalization(momentum=0.8))
# generator.add(Dense(4, activation='tanh'))
# generator.add(Reshape((1, 4, 1)))

generator.summary()

for batch_size_i in batch_size_i_array:

	batch_size = int(batch_size_i)
	batchsize = int(batch_size_i)

	print(' ')
	print(batch_size)

	@tf.function(experimental_relax_shapes=True)
	def body(loop_index, output):

		noise = tf.random.normal((batchsize, 100), 0, 1)
		logits = generator(noise, training=False)

		combined_output = tf.squeeze(logits,axis=-1)
		combined_output = tf.squeeze(combined_output,axis=1)

		return [loop_index+1, tf.concat([output, combined_output], axis=0)]

	print(loops_to_do/batchsize)

	print('Warming up...')
	loop_index = tf.constant(0)

	output = tf.zeros([0, 6], dtype=tf.float32)

	condition_func = lambda loop_index, output: loop_index < 1
	t0 = time.time()
	generated_training = tf.while_loop(condition_func, body, loop_vars=[loop_index, output])[1]
	t1 = time.time()
	total_time = t1-t0
	print('warm up1',np.shape(generated_training),'time',total_time)

	print('Starting test...')
	loop_index = tf.constant(0)

	output = tf.zeros([0, 6], dtype=tf.float32)

	if make_x_particles == True:
		condition_func = lambda loop_index, output: tf.shape(output)[0] < loops_to_do
	elif make_x_particles == False:
		condition_func = lambda loop_index, output: loop_index < loops
	t0 = time.time()
	generated_training = tf.while_loop(condition_func, body, loop_vars=[loop_index, output])[1]
	t1 = time.time()
	print(np.shape(generated_training))

	total_time = t1-t0

	points_in_5_mins = (np.shape(generated_training)[0]/total_time)
	print('Time to generate %d points: %.3fs'%(np.shape(generated_training)[0],total_time))

	print('Points generated in 5 mins: %d'%int(points_in_5_mins),'batch_size',batch_size)

	with open("results_shiplarge.txt", "a") as myfile:
		myfile.write('%d, %.2f \n'%(batch_size, points_in_5_mins))


text:
# -*- coding: UTF-8 -*-
import os
import numpy as np
import pandas as pd
from stock.globalvar import FINANCE_DIR, BASIC_DIR, LRB_CH2EN, XJLLB_CH2EN, ZCFZB_CH2EN
from stock.utils.symbol_util import symbol_to_exsymbol

def _set_quarter(df):
    for i in range(len(df)):
        dt = df.index[i]
        if dt.month == 3:
            df.loc[df.index[i], "quarter"] = "Q1"
        elif dt.month == 6:
            df.loc[df.index[i], "quarter"] = "Q2"
        elif dt.month == 9:
            df.loc[df.index[i], "quarter"] = "Q3"
        elif dt.month == 12:
            df.loc[df.index[i], "quarter"] = "Q4"

def _parse_cell(string, parser=None):
    string = string.strip()
    if string == "":
        return np.nan
    if string == "--":
        return np.nan
    if parser == None:
        return string
    return parser(string)

def get_lrb_data(exsymbol):
    filename = "%s_lrb" % exsymbol
    path = os.path.join(FINANCE_DIR["stock"], filename)
    if not os.path.isfile(path):
        msg = "%s has no lrb data" % exsymbol
        raise Exception(msg)
    content = None
    with open(path, "r") as f:
        content = f.read()
    lines = content.splitlines()
    data = {}
    index = []
    for line in lines:
        if line.strip() == "":
            continue
        cells = line.split(",")
        col = cells[0].strip()
        if col == "报告日期":
            index = map(lambda x: _parse_cell(x, str), cells[1:])
        else:
            en = LRB_CH2EN.get(col)
            if en == "":
                raise Exception("en for %s not defined" % cells[0])
            array = data.setdefault(en, [])
            parsed = map(lambda x: _parse_cell(x, float), cells[1:])
            array.extend(parsed)

    df = pd.DataFrame(data=data, index=index).fillna(0.0)
    df = df[pd.notnull(df.index)]
    df.set_index(pd.to_datetime(df.index, format="%Y-%m-%d"), inplace=True)
    _set_quarter(df)
    return df.iloc[::-1]

def get_zcfzb_data(exsymbol):
    filename = "%s_zcfzb" % exsymbol
    path = os.path.join(FINANCE_DIR["stock"], filename)
    if not os.path.isfile(path):
        msg = "%s has no zcfzb data" % exsymbol
        raise Exception(msg)
    content = None
    with open(path, "r") as f:
        content = f.read()
    lines = content.splitlines()
    data = {}
    index = []
    for line in lines:
        if line.strip() == "":
            continue
        cells = line.split(",")
        col = cells[0].strip()
        if col == "报告日期":
            index = map(lambda x: _parse_cell(x, str), cells[1:])
        else:
            en = ZCFZB_CH2EN.get(col)
            if not en:
                raise Exception("en for %s not defined" % cells[0])
            array = data.setdefault(en, [])
            parsed = map(lambda x: _parse_cell(x, float), cells[1:])
            array.extend(parsed)

    df = pd.DataFrame(data=data, index=index).fillna(0.0)
    df = df[pd.notnull(df.index)]
    df.set_index(pd.to_datetime(df.index, format="%Y-%m-%d"), inplace=True)
    _set_quarter(df)
    return df.iloc[::-1]

def get_xjllb_data(exsymbol):
    filename = "%s_xjllb" % exsymbol
    path = os.path.join(FINANCE_DIR["stock"], filename)
    if not os.path.isfile(path):
        msg = "%s has no xjllb data" % exsymbol
        raise Exception(msg)
    content = None
    with open(path, "r") as f:
        content = f.read()
    lines = content.splitlines()
    data = {}
    index = []
    for line in lines:
        if line.strip() == "":
            continue
        cells = line.split(",")
        col = cells[0].strip()
        if col == "报告日期":
            index = map(lambda x: _parse_cell(x, str), cells[1:])
        else:
            en = XJLLB_CH2EN.get(col.strip())
            if not en:
                raise Exception("en for %s not defined" % cells[0])
            array = data.setdefault(en, [])
            parsed = map(lambda x: _parse_cell(x, float), cells[1:])
            array.extend(parsed)

    df = pd.DataFrame(data=data, index=index).fillna(0.0)
    df = df[pd.notnull(df.index)]
    df.set_index(pd.to_datetime(df.index, format="%Y-%m-%d"), inplace=True)
    _set_quarter(df)
    return df.iloc[::-1]

def load_stock_basics():
    filepath = os.path.join(BASIC_DIR, "basics")
    df = pd.read_csv(filepath, encoding="utf-8", dtype={"code": str, "pe": np.float64})
    df["exsymbol"] = list(map(lambda x: symbol_to_exsymbol(x), df["code"]))
    df.set_index("exsymbol", inplace=True)
    return df


text:
from pyomo.environ import ConcreteModel
from pyomo.environ import Set,Param,Var,Objective,Constraint
from pyomo.environ import NonNegativeReals, Reals
from pyomo.environ import SolverFactory, minimize
from pyomo.environ import value
from pyomo.core.base.param import SimpleParam
import numpy as np
from sys import exit

def solve_model(model_instance, solver):
    if 'path' in solver:
        optimizer = SolverFactory(solver['name'], executable=solver['path'])
    else:
        optimizer = SolverFactory(solver['name'])

    optimizer.solve(model_instance, tee=True, keepfiles=False)

    return model_instance

def model_invest(model_data):

    model = ConcreteModel()

    ## SETS
    model.T = Set(dimen=1, ordered=True, initialize=model_data[None]['T']) # Periods
    model.M = Set(dimen=1, ordered=True, initialize=np.array(list(set(model_data[None]['month_order'].values())))) # Months

    ## PARAMETERS
    model.pwr_demand                    = Param(model.T, initialize=model_data[None]['power_demand'])
    model.pwr_gen                       = Param(model.T, initialize=model_data[None]['power_generation_pu'])
    model.ht_demand                     = Param(model.T, initialize=model_data[None]['heat_demand'])

    model.el_price_buy                  = Param(model.T, initialize=model_data[None]['electricity_price_buy'])
    model.el_price_sell                 = Param(model.T, initialize=model_data[None]['electricity_price_sell'])
    model.fl_price                      = Param(model.T, initialize=model_data[None]['fuel_price'])
    model.dh_price                      = Param(model.T, initialize=model_data[None]['district_heat_price'])

    model.grd_fee_fxd                   = Param(initialize=model_data[None]['grid_fee_fixed'])
    model.grd_fee_nrg_in                = Param(model.T, initialize=model_data[None]['grid_fee_energy_import'])
    model.grd_fee_nrg_out               = Param(model.T, initialize=model_data[None]['grid_fee_energy_export'])
    model.grd_fee_pwr_in                = Param(model.T, initialize=model_data[None]['grid_fee_power_import'])
    model.grd_fee_pwr_out               = Param(model.T, initialize=model_data[None]['grid_fee_power_export'])

    model.invest_cost_pv                = Param(initialize=model_data[None]['invest_annual_unit_cost_pv'])
    model.invest_cost_bttr              = Param(initialize=model_data[None]['invest_annual_unit_cost_battery'])
    model.invest_cost_hp                = Param(initialize=model_data[None]['invest_annual_unit_cost_heat_pump'])
    model.invest_cost_fb                = Param(initialize=model_data[None]['invest_annual_unit_cost_fuel_boiler'])
    model.invest_cost_eb                = Param(initialize=model_data[None]['invest_annual_unit_cost_electric_boiler'])
    model.invest_cost_tes               = Param(initialize=model_data[None]['invest_annual_unit_cost_thermal_storage'])
    model.invest_cost_dh                = Param(initialize=model_data[None]['invest_annual_unit_cost_district_heat'])

    model.heat_cap_over                 = Param(initialize=model_data[None]['heat_capacity_oversize'])

    model.pv_cap_max                    = Param(initialize=model_data[None]['pv_capacity_max'])
    model.pv_cap_min                    = Param(initialize=model_data[None]['pv_capacity_min'])

    model.dh_cap_max                    = Param(initialize=model_data[None]['district_heat_capacity_max'])
    model.dh_cap_min                    = Param(initialize=model_data[None]['district_heat_capacity_min'])

    model.hp_cap_max                    = Param(initialize=model_data[None]['heat_pump_capacity_max'])
    model.hp_cap_min                    = Param(initialize=model_data[None]['heat_pump_capacity_min'])
    model.hp_cop                        = Param(initialize=model_data[None]['heat_pump_cop'])

    model.eb_cap_max                    = Param(initialize=model_data[None]['electric_boiler_capacity_max'])
    model.eb_cap_min                    = Param(initialize=model_data[None]['electric_boiler_capacity_min'])
    model.eb_eff                        = Param(initialize=model_data[None]['electric_boiler_efficiency'])

    model.fb_cap_max                    = Param(initialize=model_data[None]['fuel_boiler_capacity_max'])
    model.fb_cap_min                    = Param(initialize=model_data[None]['fuel_boiler_capacity_min'])
    model.fb_eff                        = Param(initialize=model_data[None]['fuel_boiler_efficiency'])

    model.bttr_cap_max                  = Param(initialize=model_data[None]['battery_capacity_max'])
    model.bttr_cap_min                  = Param(initialize=model_data[None]['battery_capacity_min'])
    model.bttr_soc_min                  = Param(initialize=model_data[None]['battery_soc_min'])
    model.bttr_in_max                   = Param(initialize=model_data[None]['battery_charge_max'])
    model.bttr_out_max                  = Param(initialize=model_data[None]['battery_discharge_max'])
    model.bttr_eff_in                   = Param(initialize=model_data[None]['battery_efficiency_charge'])
    model.bttr_eff_out                  = Param(initialize=model_data[None]['battery_efficiency_discharge'])
    model.bttr_soc_ini                  = Param(initialize=model_data[None]['battery_soc_ini'])
    model.bttr_soc_fin                  = Param(initialize=model_data[None]['battery_soc_fin'])
    model.bttr_grid_in                  = Param(initialize=model_data[None]['battery_grid_charging'])

    model.tes_cap_max                   = Param(initialize=model_data[None]['thermal_storage_capacity_max'])
    model.tes_cap_min                   = Param(initialize=model_data[None]['thermal_storage_capacity_min'])
    model.tes_in_max                    = Param(initialize=model_data[None]['thermal_storage_charge_max'])
    model.tes_out_max                   = Param(initialize=model_data[None]['thermal_storage_discharge_max'])
    model.tes_loss                      = Param(initialize=model_data[None]['thermal_storage_losses'])
    model.tes_soc_ini                   = Param(initialize=model_data[None]['thermal_storage_soc_ini'])
    model.tes_soc_fin                   = Param(initialize=model_data[None]['thermal_storage_soc_fin'])

    model.emission_grid                 = Param(model.T, initialize=model_data[None]['specific_emission_grid'])
    model.emission_fuel                 = Param(model.T, initialize=model_data[None]['specific_emission_fuel'])
    model.emission_dh                   = Param(model.T, initialize=model_data[None]['specific_emission_district_heat'])

    model.dt                            = Param(initialize=model_data[None]['dt'])

    ## VARIABLES
    model.COST_INV                      = Var(within=Reals)
    model.COST_ENERGY                   = Var(model.T, within=Reals)
    model.COST_GRID_ENERGY_IMPORT       = Var(model.T, within=Reals)
    model.COST_GRID_ENERGY_EXPORT       = Var(model.T, within=Reals)
    model.COST_GRID_POWER_IMPORT        = Var(model.T, within=NonNegativeReals)
    model.COST_GRID_POWER_EXPORT        = Var(model.T, within=NonNegativeReals)
    model.COST_GRID_POWER_IMPORT_MAX    = Var(model.M, within=Reals)
    model.COST_GRID_POWER_EXPORT_MAX    = Var(model.M, within=Reals)
    model.COST_GRID_FIXED               = Var(within=Reals)
    model.COST_FUEL                     = Var(model.T, within=Reals)
    model.COST_DH                       = Var(model.T, within=Reals)

    model.P_BUY                         = Var(model.T, within=NonNegativeReals)
    model.P_SELL                        = Var(model.T, within=NonNegativeReals)
    model.P_GEN                         = Var(model.T, within=Reals)

    model.P_HP                          = Var(model.T, within=NonNegativeReals)
    model.P_EB                          = Var(model.T, within=NonNegativeReals)
    model.F_FB                          = Var(model.T, within=NonNegativeReals)

    model.Q_DH                          = Var(model.T, within=NonNegativeReals)
    model.Q_HP                          = Var(model.T, within=NonNegativeReals)
    model.Q_EB                          = Var(model.T, within=NonNegativeReals)
    model.Q_FB                          = Var(model.T, within=NonNegativeReals)

    model.BTTR_SOC                      = Var(model.T, within=NonNegativeReals)
    model.BTTR_IN                       = Var(model.T, within=NonNegativeReals)
    model.BTTR_OUT                      = Var(model.T, within=NonNegativeReals)

    model.TES_SOC                       = Var(model.T, within=NonNegativeReals)
    model.TES_IN                        = Var(model.T, within=NonNegativeReals)
    model.TES_OUT                       = Var(model.T, within=NonNegativeReals)

    model.PV_CAP                        = Var(within=NonNegativeReals, bounds=(model.pv_cap_min, model.pv_cap_max))
    model.BTTR_CAP                      = Var(within=NonNegativeReals, bounds=(model.bttr_cap_min, model.bttr_cap_max))
    model.HP_CAP                        = Var(within=NonNegativeReals, bounds=(model.hp_cap_min, model.hp_cap_max))
    model.FB_CAP                        = Var(within=NonNegativeReals, bounds=(model.fb_cap_min, model.fb_cap_max))
    model.EB_CAP                        = Var(within=NonNegativeReals, bounds=(model.eb_cap_min, model.eb_cap_max))
    model.DH_CAP                        = Var(within=NonNegativeReals, bounds=(model.dh_cap_min, model.dh_cap_max))
    model.TES_CAP                       = Var(within=NonNegativeReals, bounds=(model.tes_cap_min, model.tes_cap_max))

    model.CO2_GRID                      = Var(model.T, within=Reals)
    model.CO2_FUEL                      = Var(model.T, within=Reals)
    model.CO2_DH                        = Var(model.T, within=Reals)

    ## OBJECTIVE
    # Minimize cost
    def total_cost(model):
        return  model.COST_INV \
        + sum(model.COST_ENERGY[t] for t in model.T) \
        + sum(model.COST_GRID_ENERGY_IMPORT[t] + model.COST_GRID_ENERGY_EXPORT[t] for t in model.T) \
        + sum(model.COST_GRID_POWER_IMPORT_MAX[m] + model.COST_GRID_POWER_EXPORT_MAX[m] for m in model.M) \
        + model.COST_GRID_FIXED \
        + sum(model.COST_FUEL[t] for t in model.T) \
        + sum(model.COST_DH[t] for t in model.T)
    model.total_cost = Objective(rule=total_cost, sense=minimize)

    ## CONSTRAINTS
    # Investment cost
    def investment_cost(model):
        return model.COST_INV == model.invest_cost_pv*model.PV_CAP \
                                + model.invest_cost_bttr*model.BTTR_CAP \
                                + model.invest_cost_hp*model.HP_CAP \
                                + model.invest_cost_eb*model.EB_CAP \
                                + model.invest_cost_fb*model.FB_CAP \
                                + model.invest_cost_dh*model.DH_CAP \
                                + model.invest_cost_tes*model.TES_CAP
    model.investment_cost = Constraint(rule=investment_cost)

    # Energy cost
    def energy_cost(model, t):
        return model.COST_ENERGY[t] == model.el_price_buy[t]*model.P_BUY[t]*model.dt - model.el_price_sell[t]*model.P_SELL[t]*model.dt
    model.energy_cost = Constraint(model.T, rule=energy_cost)

    # Grid fixed cost
    def grid_fixed_cost(model):
        return model.COST_GRID_FIXED == model.grd_fee_fxd*len(model.M)
    model.grid_fixed_cost = Constraint(rule=grid_fixed_cost)

    # Grid energy import cost
    def grid_energy_import_cost(model, t):
        return model.COST_GRID_ENERGY_IMPORT[t] == model.grd_fee_nrg_in[t]*model.P_BUY[t]*model.dt
    model.grid_energy_import_cost = Constraint(model.T, rule=grid_energy_import_cost)

    # Grid energy export cost
    def grid_energy_export_cost(model, t):
        return model.COST_GRID_ENERGY_EXPORT[t] == model.grd_fee_nrg_out[t]*model.P_SELL[t]*model.dt
    model.grid_energy_export_cost = Constraint(model.T, rule=grid_energy_export_cost)

    # Grid power import cost
    def grid_power_import_cost(model, t):
        return model.COST_GRID_POWER_IMPORT[t] >= model.grd_fee_pwr_in[t]*(model.P_BUY[t]-model.P_SELL[t])
    model.grid_power_import_cost = Constraint(model.T, rule=grid_power_import_cost)

    # Grid power export cost
    def grid_power_export_cost(model, t):
        return model.COST_GRID_POWER_EXPORT[t] >= model.grd_fee_pwr_out[t]*(model.P_SELL[t]-model.P_BUY[t])
    model.grid_power_export_cost = Constraint(model.T, rule=grid_power_export_cost)

    # Max grid import cost
    def max_grid_power_import_cost(model, t):
        return model.COST_GRID_POWER_IMPORT_MAX[model_data[None]['month_order'][t]] >= model.COST_GRID_POWER_IMPORT[t]
    model.max_grid_power_import_cost = Constraint(model.T, rule=max_grid_power_import_cost)

    # Max grid export cost
    def max_grid_power_export_cost(model, t):
        return model.COST_GRID_POWER_EXPORT_MAX[model_data[None]['month_order'][t]] >= model.COST_GRID_POWER_EXPORT[t]
    model.max_grid_power_export_cost = Constraint(model.T, rule=max_grid_power_export_cost)

    # Fuel boiler cost
    def fuel_cost(model, t):
        return model.COST_FUEL[t] == model.fl_price[t]*model.F_FB[t]*model.dt
    model.fuel_cost = Constraint(model.T, rule=fuel_cost)

    # District heating cost
    def dh_cost(model, t):
        return model.COST_DH[t] == model.dh_price[t]*model.Q_DH[t]*model.dt
    model.dh_cost = Constraint(model.T, rule=dh_cost)

    # Power generation
    def power_generation(model, t):
        return model.P_GEN[t] == model.pwr_gen[t]*model.PV_CAP
    model.power_generation = Constraint(model.T, rule=power_generation)

    # Power balance
    def energy_balance(model, t):
        return model.P_SELL[t] - model.P_BUY[t] == model.P_GEN[t] + model.BTTR_OUT[t] - model.BTTR_IN[t] - model.pwr_demand[t] - model.P_HP[t] - model.P_EB[t]
    model.energy_balance = Constraint(model.T, rule=energy_balance)

    # Battery energy balance
    def battery_soc(model, t):
        if t==model.T.first():
            return model.BTTR_SOC[t] - model.bttr_soc_ini*model.BTTR_CAP == model.bttr_eff_in*model.BTTR_IN[t]*model.dt  - (1/model.bttr_eff_out)*model.BTTR_OUT[t]*model.dt
        else:
            return model.BTTR_SOC[t] - model.BTTR_SOC[model.T.prev(t)] == model.bttr_eff_in*model.BTTR_IN[t]*model.dt  - (1/model.bttr_eff_out)*model.BTTR_OUT[t]*model.dt
    model.battery_soc = Constraint(model.T, rule=battery_soc)

    # Battery SOC upper limit
    def battery_soc_upper_limit(model, t):
        return model.BTTR_SOC[t] <= model.BTTR_CAP
    model.battery_soc_upper_limit = Constraint(model.T, rule=battery_soc_upper_limit)

    # Battery SOC lower limit
    def battery_soc_lower_limit(model, t):
        return model.BTTR_SOC[t] >= model.bttr_soc_min*model.BTTR_CAP
    model.battery_soc_lower_limit = Constraint(model.T, rule=battery_soc_lower_limit)

    # Battery charge upper limit
    def battery_charge_upper_limit(model, t):
        return model.BTTR_IN[t] <= model.bttr_in_max*model.BTTR_CAP
    model.battery_charge_upper_limit = Constraint(model.T, rule=battery_charge_upper_limit)

    def battery_discharge_upper_limit(model, t):
        return model.BTTR_OUT[t] <= model.bttr_out_max*model.BTTR_CAP
    model.battery_discharge_upper_limit = Constraint(model.T, rule=battery_discharge_upper_limit)

    # Battery charging from grid
    def no_grid_charging(model, t):
        if value(model.bttr_grid_in) == False:
            return model.P_BUY[t] <= model.pwr_demand[t] + model.P_HP[t] + model.P_EB[t]
        else:
            return model.P_BUY[t] <= model.pwr_demand[t] + model.P_HP[t] + model.P_EB[t] + model.BTTR_IN[t]
    model.no_grid_charging = Constraint(model.T, rule=no_grid_charging)

    # Fuel boiler heat generation - fuel consumption
    def fuel_boiler_gen(model, t):
        return model.Q_FB[t] == model.fb_eff*model.F_FB[t]
    model.fuel_boiler_gen = Constraint(model.T, rule=fuel_boiler_gen)

    # Heat pump heat generation - power consumption
    def heat_pump_gen(model, t):
        return model.Q_HP[t] == model.hp_cop*model.P_HP[t]
    model.heat_pump_gen = Constraint(model.T, rule=heat_pump_gen)

    # Electric boiler heat generation - power consumption
    def electric_boiler_gen(model, t):
        return model.Q_EB[t] == model.eb_eff*model.P_EB[t]
    model.electric_boiler_gen = Constraint(model.T, rule=electric_boiler_gen)

    # Fuel boiler capacity
    def fuel_boiler_cap(model, t):
        return (1 + model.heat_cap_over)*model.Q_FB[t] <= model.FB_CAP
    model.fuel_boiler_cap = Constraint(model.T, rule=fuel_boiler_cap)

    # Heat pump capacity
    def heat_pump_cap(model, t):
        return (1 + model.heat_cap_over)*model.Q_HP[t] <= model.HP_CAP
    model.heat_pump_cap = Constraint(model.T, rule=heat_pump_cap)

    # Electric boiler capacity
    def electric_boiler_cap(model, t):
        return (1 + model.heat_cap_over)*model.Q_EB[t] <= model.EB_CAP
    model.electric_boiler_cap = Constraint(model.T, rule=electric_boiler_cap)

    # District heating capacity
    def district_heating_cap(model, t):
        return (1 + model.heat_cap_over)*model.Q_DH[t] <= model.DH_CAP
    model.district_heating_cap = Constraint(model.T, rule=district_heating_cap)

    # Heat balance
    def heat_balance(model, t):
        return model.Q_FB[t] + model.Q_HP[t] + model.Q_EB[t] + model.Q_DH[t] == model.ht_demand[t] - model.TES_OUT[t] + model.TES_IN[t]
    model.heat_balance = Constraint(model.T, rule=heat_balance)

    # Heat storage energy balance
    def tes_soc(model, t):
        if t==model.T.first():
            return model.TES_SOC[t] - (1-model.tes_loss)*model.tes_soc_ini*model.TES_CAP == model.TES_IN[t]*model.dt - model.TES_OUT[t]*model.dt
        else:
            return model.TES_SOC[t] - (1-model.tes_loss)*model.TES_SOC[model.T.prev(t)] == model.TES_IN[t]*model.dt - model.TES_OUT[t]*model.dt
    model.tes_soc = Constraint(model.T, rule=tes_soc)

	# Heat storage upper limit
    def tes_upper_bound(model, t):
        return model.TES_SOC[t] <= model.TES_CAP
    model.tes_upper_bound = Constraint(model.T, rule=tes_upper_bound)

	# Heat storage charge upper limit
    def tes_charge_limit(model, t):
        return model.TES_IN[t] <= model.tes_in_max*model.TES_CAP
    model.tes_charge_limit = Constraint(model.T, rule=tes_charge_limit)

	# Heat storage discharge upper limit
    def tes_discharge_limit(model, t):
        return model.TES_OUT[t] <= model.tes_out_max*model.TES_CAP
    model.tes_discharge_limit = Constraint(model.T, rule=tes_discharge_limit)

    # Grid CO2 emissions
    def carbon_emissions_grid(model, t):
        return model.CO2_GRID[t] == (model.P_BUY[t] - model.P_SELL[t])*model.emission_grid[t]*model.dt
    model.carbon_emissions_grid = Constraint(model.T, rule=carbon_emissions_grid)

    # Fuel CO2 emissions
    def carbon_emissions_fuel(model, t):
        return model.CO2_FUEL[t] == model.F_FB[t]*model.emission_fuel[t]*model.dt
    model.carbon_emissions_fuel = Constraint(model.T, rule=carbon_emissions_fuel)

    # District heating CO2 emissions
    def carbon_emissions_dh(model, t):
        return model.CO2_DH[t] == model.Q_DH[t]*model.emission_dh[t]*model.dt
    model.carbon_emissions_dh = Constraint(model.T, rule=carbon_emissions_dh)

    # Fix battery soc in the last period
    if value(model.bttr_soc_fin) > 0:
        model.BTTR_SOC[model.T.last()].fix(model.bttr_soc_fin*model.BTTR_CAP)

    # Fix tes soc in the last period
    if value(model.tes_soc_fin) > 0:
        model.TES_SOC[model.T.last()].fix(model.tes_soc_fin*model.TES_CAP)

    return model

def model_invest_input(data):

    def annuity_rate(discount_rate, salvage_rate, life_span):

        anrate = (1-salvage_rate)*discount_rate*(1+discount_rate)**life_span/((1+discount_rate)**life_span-1)
        return anrate

    if 'power_generation_pu' in data:
        periods = np.arange(1, len(data['power_generation_pu'])+1)
    elif 'power_demand' in data:
        periods = np.arange(1, len(data['power_demand'])+1)
    elif 'heat_demand' in data:
        periods = np.arange(1, len(data['heat_demand'])+1)

    power_generation_pu = dict(zip(periods,  data['power_generation_pu'])) if 'power_generation_pu' in data else dict(zip(periods,  [0] * len(periods)))
    power_demand = dict(zip(periods,  data['power_demand'])) if 'power_demand' in data else dict(zip(periods,  [0] * len(periods)))
    heat_demand = dict(zip(periods,  data['heat_demand'])) if 'heat_demand' in data else dict(zip(periods,  [0] * len(periods)))

    electricity_price_buy = dict(zip(periods,  data['electricity_price_buy'])) if 'electricity_price_buy' in data else dict(zip(periods,  [0] * len(periods)))
    electricity_price_sell = dict(zip(periods,  data['electricity_price_sell'])) if 'electricity_price_sell' in data else dict(zip(periods,  [0] * len(periods)))

    grid_fee_fixed = data['grid_fee_fixed'] if 'grid_fee_fixed' in data else 0.0
    grid_fee_energy_import = dict(zip(periods,  data['grid_fee_energy_import'])) if 'grid_fee_energy_import' in data else dict(zip(periods,  [0] * len(periods)))
    grid_fee_energy_export = dict(zip(periods,  data['grid_fee_energy_export'])) if 'grid_fee_energy_export' in data else dict(zip(periods,  [0] * len(periods)))
    grid_fee_power_import = dict(zip(periods,  data['grid_fee_power_import'])) if 'grid_fee_power_import' in data else dict(zip(periods,  [0] * len(periods)))
    grid_fee_power_export = dict(zip(periods,  data['grid_fee_power_export'])) if 'grid_fee_power_export' in data else dict(zip(periods,  [0] * len(periods)))

    fuel_price = dict(zip(periods,  data['fuel_price'])) if 'fuel_price' in data else dict(zip(periods,  [0] * len(periods)))
    district_heat_price = dict(zip(periods,  data['district_heat_price'])) if 'district_heat_price' in data else dict(zip(periods,  [0] * len(periods)))

    discount_rate = data['discount_rate'] if 'discount_rate' in data else 0.0
    salvage_rate = data['salvage_rate'] if 'salvage_rate' in data else 0.0

    pv_lifespan = data['pv_lifespan'] if 'pv_lifespan' in data else 25
    battery_lifespan = data['battery_lifespan'] if 'battery_lifespan' in data else 10
    heat_pump_lifespan = data['heat_pump_lifespan'] if 'heat_pump_lifespan' in data else 25
    electric_boiler_lifespan = data['electric_boiler_lifespan'] if 'electric_boiler_lifespan' in data else 25
    fuel_boiler_lifespan = data['fuel_boiler_lifespan'] if 'fuel_boiler_lifespan' in data else 25
    district_heat_lifespan = data['district_heat_lifespan'] if 'district_heat_lifespan' in data else 50
    thermal_storage_lifespan = data['thermal_storage_lifespan'] if 'thermal_storage_lifespan' in data else 30

    annuity_rate_pv = annuity_rate(discount_rate, salvage_rate, pv_lifespan)
    annuity_rate_battery = annuity_rate(discount_rate, salvage_rate, battery_lifespan)
    annuity_rate_heat_pump = annuity_rate(discount_rate, salvage_rate, heat_pump_lifespan)
    annuity_rate_electric_boiler = annuity_rate(discount_rate, salvage_rate, electric_boiler_lifespan)
    annuity_rate_fuel_boiler = annuity_rate(discount_rate, salvage_rate, fuel_boiler_lifespan)
    annuity_rate_district_heat = annuity_rate(discount_rate, salvage_rate, district_heat_lifespan)
    annuity_rate_thermal_storage = annuity_rate(discount_rate, salvage_rate, thermal_storage_lifespan)

    invest_annual_unit_cost_pv = data['invest_unit_cost_pv']*annuity_rate_pv if 'invest_unit_cost_pv' in data else 0.0
    invest_annual_unit_cost_battery = data['invest_unit_cost_battery']*annuity_rate_battery if 'invest_unit_cost_battery' in data else 0.0
    invest_annual_unit_cost_heat_pump = data['invest_unit_cost_heat_pump']*annuity_rate_heat_pump if 'invest_unit_cost_heat_pump' in data else 0.0
    invest_annual_unit_cost_electric_boiler = data['invest_unit_cost_electric_boiler']*annuity_rate_electric_boiler if 'invest_unit_cost_electric_boiler' in data else 0.0
    invest_annual_unit_cost_fuel_boiler = data['invest_unit_cost_fuel_boiler']*annuity_rate_fuel_boiler if 'invest_unit_cost_fuel_boiler' in data else 0.0
    invest_annual_unit_cost_district_heat = data['invest_unit_cost_district_heat']*annuity_rate_district_heat if 'invest_unit_cost_district_heat' in data else 0.0
    invest_annual_unit_cost_thermal_storage = data['invest_unit_cost_thermal_storage']*annuity_rate_thermal_storage if 'invest_unit_cost_thermal_storage' in data else 0.0

    heat_capacity_oversize = data['heat_capacity_oversize'] if 'heat_capacity_oversize' in data else 0.0

    pv_capacity_max = data['pv_capacity_max'] if 'pv_capacity_max' in data else 0.0
    pv_capacity_min = data['pv_capacity_min'] if 'pv_capacity_min' in data else 0.0

    district_heat_capacity_max = data['district_heat_capacity_max'] if 'district_heat_capacity_max' in data else 0.0
    district_heat_capacity_min = data['district_heat_capacity_min'] if 'district_heat_capacity_min' in data else 0.0

    heat_pump_capacity_max = data['heat_pump_capacity_max'] if 'heat_pump_capacity_max' in data else 0.0
    heat_pump_capacity_min = data['heat_pump_capacity_min'] if 'heat_pump_capacity_min' in data else 0.0
    heat_pump_cop = data['heat_pump_cop'] if 'heat_pump_cop' in data else 3.0

    electric_boiler_capacity_max = data['electric_boiler_capacity_max'] if 'electric_boiler_capacity_max' in data else 0.0
    electric_boiler_capacity_min = data['electric_boiler_capacity_min'] if 'electric_boiler_capacity_min' in data else 0.0
    electric_boiler_efficiency = data['electric_boiler_efficiency'] if 'electric_boiler_efficiency' in data else 0.0

    fuel_boiler_capacity_max = data['fuel_boiler_capacity_max'] if 'fuel_boiler_capacity_max' in data else 0.0
    fuel_boiler_capacity_min = data['fuel_boiler_capacity_min'] if 'fuel_boiler_capacity_min' in data else 0.0
    fuel_boiler_efficiency = data['fuel_boiler_efficiency'] if 'fuel_boiler_efficiency' in data else 0.0

    battery_capacity_max = data['battery_capacity_max'] if 'battery_capacity_max' in data else 0.0
    battery_capacity_min = data['battery_capacity_min'] if 'battery_capacity_min' in data else 0.0
    battery_soc_min = data['battery_soc_min'] if 'battery_soc_min' in data else 0.0
    battery_charge_max = data['battery_charge_max'] if 'battery_charge_max' in data else 1.0
    battery_discharge_max = data['battery_discharge_max'] if 'battery_discharge_max' in data else 1.0
    battery_efficiency_charge = data['battery_efficiency_charge'] if 'battery_efficiency_charge' in data else 0.95
    battery_efficiency_discharge = data['battery_efficiency_discharge'] if 'battery_efficiency_discharge' in data else 0.95
    battery_soc_ini = data['battery_soc_ini'] if 'battery_soc_ini' in data else 0.0
    battery_soc_fin = data['battery_soc_fin'] if 'battery_soc_fin' in data else 0.0
    battery_grid_charging = data['battery_grid_charging'] if 'battery_grid_charging' in data else True

    thermal_storage_capacity_max = data['thermal_storage_capacity_max'] if 'thermal_storage_capacity_max' in data else 0.0
    thermal_storage_capacity_min = data['thermal_storage_capacity_min'] if 'thermal_storage_capacity_min' in data else 0.0
    thermal_storage_charge_max = data['thermal_storage_charge_max'] if 'thermal_storage_charge_max' in data else 1.0
    thermal_storage_discharge_max = data['thermal_storage_discharge_max'] if 'thermal_storage_discharge_max' in data else 1.0
    thermal_storage_losses = data['thermal_storage_losses'] if 'thermal_storage_losses' in data else 0.05
    thermal_storage_soc_ini = data['thermal_storage_soc_ini'] if 'thermal_storage_soc_ini' in data else 0.0
    thermal_storage_soc_fin = data['thermal_storage_soc_fin'] if 'thermal_storage_soc_fin' in data else 0.0

    specific_emission_grid = dict(zip(periods,  data['specific_emission_grid'])) if 'specific_emission_grid' in data else dict(zip(periods,  [0] * len(periods)))
    specific_emission_fuel = dict(zip(periods,  data['specific_emission_fuel'])) if 'specific_emission_fuel' in data else dict(zip(periods,  [0] * len(periods)))
    specific_emission_district_heat = dict(zip(periods,  data['specific_emission_district_heat'])) if 'specific_emission_district_heat' in data else dict(zip(periods,  [0] * len(periods)))

    dt = data['dt'] if 'dt' in data else 1.0
    month_order = dict(zip(periods,  data['month_order'])) if 'month_order' in data else dict(zip(periods,  [1] * len(periods)))

    min_feasible_heat_capacity =  max(heat_demand.values())*(1+heat_capacity_oversize)

    if heat_pump_capacity_max  + electric_boiler_capacity_max + fuel_boiler_capacity_max  + district_heat_capacity_max < min_feasible_heat_capacity:

        print('Total maximum heat capacity is too low. Increase at least to: ', min_feasible_heat_capacity)
        exit()

    # Create model data input dictionary
    model_data = {None: {
        'T': periods,

        'invest_annual_unit_cost_battery': invest_annual_unit_cost_battery,
        'invest_annual_unit_cost_pv': invest_annual_unit_cost_pv,
        'invest_annual_unit_cost_heat_pump': invest_annual_unit_cost_heat_pump,
        'invest_annual_unit_cost_electric_boiler': invest_annual_unit_cost_electric_boiler,
        'invest_annual_unit_cost_fuel_boiler': invest_annual_unit_cost_fuel_boiler,
        'invest_annual_unit_cost_district_heat': invest_annual_unit_cost_district_heat,
        'invest_annual_unit_cost_thermal_storage': invest_annual_unit_cost_thermal_storage,

        'heat_capacity_oversize': heat_capacity_oversize,

        'power_generation_pu': power_generation_pu,
        'power_demand': power_demand,
        'heat_demand': heat_demand,

        'electricity_price_buy': electricity_price_buy,
        'electricity_price_sell': electricity_price_sell,
        'fuel_price': fuel_price,
        'district_heat_price': district_heat_price,

        'grid_fee_fixed': grid_fee_fixed,
        'grid_fee_energy_import': grid_fee_energy_import,
        'grid_fee_energy_export': grid_fee_energy_export,
        'grid_fee_power_import': grid_fee_power_import,
        'grid_fee_power_export': grid_fee_power_export,

        'pv_capacity_max': pv_capacity_max,
        'pv_capacity_min': pv_capacity_min,

        'battery_capacity_max': battery_capacity_max,
        'battery_capacity_min': battery_capacity_min,
        'battery_soc_min': battery_soc_min,
        'battery_charge_max': battery_charge_max,
        'battery_discharge_max': battery_discharge_max,
        'battery_efficiency_charge': battery_efficiency_charge,
        'battery_efficiency_discharge': battery_efficiency_discharge,
        'battery_grid_charging': battery_grid_charging,
        'battery_soc_ini': battery_soc_ini,
        'battery_soc_fin': battery_soc_fin,

        'district_heat_capacity_max': district_heat_capacity_max,
        'district_heat_capacity_min': district_heat_capacity_min,

        'heat_pump_capacity_max': heat_pump_capacity_max,
        'heat_pump_capacity_min': heat_pump_capacity_min,
        'heat_pump_cop': heat_pump_cop,

        'electric_boiler_capacity_max': electric_boiler_capacity_max,
        'electric_boiler_capacity_min': electric_boiler_capacity_min,
        'electric_boiler_efficiency': electric_boiler_efficiency,

        'fuel_boiler_capacity_max': fuel_boiler_capacity_max,
        'fuel_boiler_capacity_min': fuel_boiler_capacity_min,
        'fuel_boiler_efficiency': fuel_boiler_efficiency,

        'thermal_storage_capacity_max': thermal_storage_capacity_max,
        'thermal_storage_capacity_min': thermal_storage_capacity_min,
        'thermal_storage_charge_max': thermal_storage_charge_max,
        'thermal_storage_discharge_max': thermal_storage_discharge_max,
        'thermal_storage_losses': thermal_storage_losses,
        'thermal_storage_soc_ini': thermal_storage_soc_ini,
        'thermal_storage_soc_fin': thermal_storage_soc_fin,

        'specific_emission_grid': specific_emission_grid,
        'specific_emission_fuel': specific_emission_fuel,
        'specific_emission_district_heat': specific_emission_district_heat,

        'month_order': month_order,
        'dt': dt,
    }}

    return model_data

def model_invest_results(solution):

    s = dict()

    s['cost_total'] = solution.total_cost()

    s['cost_investment_annual'] = value(solution.COST_INV)
    s['cost_energy'] = value(solution.COST_ENERGY[:])
    s['cost_grid_energy_import'] = value(solution.COST_GRID_ENERGY_IMPORT[:])
    s['cost_grid_energy_export'] = value(solution.COST_GRID_ENERGY_EXPORT[:])
    s['cost_grid_power_import'] = value(solution.COST_GRID_POWER_IMPORT_MAX[:])
    s['cost_grid_power_export'] = value(solution.COST_GRID_POWER_EXPORT_MAX[:])
    s['cost_grid_power_fixed'] = value(solution.COST_GRID_FIXED)

    s['cost_fuel_heating'] = value(solution.COST_FUEL[:])
    s['cost_district_heating'] = value(solution.COST_DH[:])

    s['power_buy'] = value(solution.P_BUY[:])
    s['power_sell'] = value(solution.P_SELL[:])
    s['power_generation'] = value(solution.P_GEN[:])
    s['power_consumption_heat_pump'] = value(solution.P_HP[:])
    s['power_consumption_electric_boiler'] = value(solution.P_EB[:])
    s['heat_generation_heat_pump'] = value(solution.Q_HP[:])
    s['heat_generation_electric_boiler'] = value(solution.Q_EB[:])
    s['heat_generation_fuel_boiler'] = value(solution.Q_FB[:])
    s['heat_consumption_district'] = value(solution.Q_DH[:])
    s['fuel_consumption_boiler'] = value(solution.F_FB[:])

    s['battery_soc'] = value(solution.BTTR_SOC[:])
    s['battery_charge'] = value(solution.BTTR_IN[:])
    s['battery_discharge'] = value(solution.BTTR_OUT[:])

    s['thermal_storage_soc'] = value(solution.TES_SOC[:])
    s['thermal_storage_charge'] = value(solution.TES_IN[:])
    s['thermal_storage_discharge'] = value(solution.TES_OUT[:])

    s['carbon_emissions_grid'] = value(solution.CO2_GRID[:])
    s['carbon_emissions_fuel'] = value(solution.CO2_FUEL[:])
    s['carbon_emissions_district'] = value(solution.CO2_DH[:])

    s['capacity_pv'] = value(solution.PV_CAP)
    s['capacity_battery'] = value(solution.BTTR_CAP)
    s['capacity_heat_pump'] = value(solution.HP_CAP)
    s['capacity_electric_boiler'] = value(solution.EB_CAP)
    s['capacity_fuel_boiler'] = value(solution.FB_CAP)
    s['capacity_district'] = value(solution.DH_CAP)
    s['capacity_thermal_storage'] = value(solution.TES_CAP)

    return s

reason: import numpy,from pyomo


text:
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy
import copy
import matplotlib
import os
from matplotlib import pyplot as plt
import seaborn as sns
import torch.nn.functional as F
sns.set(style="ticks")
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
import numpy as np
from utils import Str2id
from collections import Counter
from utils import NodeClassifier
#from tqdm import tqdm_notebook as tqdm
from tqdm import tqdm as tqdm
from tqdm import tnrange

class NormalEdgeDataset(Dataset):
    """Edge dataset."""

    def __init__(
        self,
        node_map1,
        node_map2,
        filename,
        neg_sampling_size,
        from_random_walk=False,
        window_size=None
    ):
        """
        Args:
            edge_list (list): A list containts all edges. Each edge is a tuple
        """
        if from_random_walk:
            assert window_size is not None
            self.random_walk_to_edge_list(
                node_map1, node_map2, filename, neg_sampling_size, window_size
            )
        else:
            self.init_from_edge_list_file(
                node_map1, node_map2, filename, neg_sampling_size
            )

    def random_walk_to_edge_list(
        self, node_map1, node_map2, in_file_path, neg_sampling_size, window_size
    ):
        self.edge_list = []
        self.node_list = []
        self.graph = []
        for line in tqdm(open(in_file_path)):
            line = line.strip().split(' ')
            n = len(line)
            for i in range(len(line)):
                left = max(0, i - window_size)
                right = min(i + window_size, n)
                for j in range(left, right):
                    if j == i:
                        continue
                    u = node_map1.str2id(line[i])
                    v = node_map2.str2id(line[j])
                    if u >= len(self.graph):
                        for _ in range(u - len(self.graph) + 1):
                            self.graph.append(set())
                    self.graph[u].add(v)
                    self.edge_list.append([u, v, []])
                    self.node_list.append(v)
        self.init_sample_table()
        self.neg_sampling_size = neg_sampling_size

    def init_from_edge_list_file(
        self, node_map1, node_map2, edge_list_file, neg_sampling_size
    ):
        self.edge_list = []
        self.node_list = []
        self.graph = []
        for line in tqdm(open(edge_list_file)):
            line = line.strip().split(' ')
            line[0] = node_map1.str2id(line[0])
            line[1] = node_map2.str2id(line[1])
            if line[0] >= len(self.graph):
                for _ in range(line[0] - len(self.graph) + 1):
                    self.graph.append(set())
            self.graph[line[0]].add(line[1])
            self.edge_list.append([line[0], line[1], []])
            self.node_list.append(line[1])
        self.init_sample_table()
        self.neg_sampling_size = neg_sampling_size

    def init_sample_table(self):
        self.sample_table = []
        sample_table_size = 1e8
        node_frequency = list(Counter(self.node_list).items())
        nids = [f[0] for f in node_frequency]
        values = [f[1] for f in node_frequency]
        pow_frequency = numpy.array(values)**0.75
        nodes_pow = numpy.sum(pow_frequency)
        ratio = pow_frequency / nodes_pow
        count = numpy.round(ratio * sample_table_size)
        for index, c in enumerate(count):
            self.sample_table += [nids[index]] * int(c)
        self.sample_table = numpy.array(self.sample_table)

    def negative_sampling1(self):
        #numpy.random.seed(41)
        neg = numpy.random.choice(
            self.sample_table,
            size=(len(self.edge_list), self.neg_sampling_size)
        )
        #         neg=numpy.random.choice(list(set(self.node_list)),size=(len(self.edge_list),self.neg_sampling_size))
        for i in range(len(self.edge_list)):
            self.edge_list[i][2] = list(neg[i])

    def negative_sampling2(self):
        #numpy.random.seed(41)
        for i in range(len(self.edge_list)):
            neg = numpy.random.choice(
                self.sample_table, size=(self.neg_sampling_size)
            )
            self.edge_list[i][2] = list(neg)

    def negative_sampling(self):
        #numpy.random.seed(41)
        negs = numpy.random.choice(
            self.sample_table,
            size=(len(self.edge_list), self.neg_sampling_size)
        )
        for i in tqdm(range(len(self.edge_list))):
            u = self.edge_list[i][0]
            v = self.edge_list[i][1]
            neg = list(filter(lambda x: x not in self.graph[u], negs[i]))
            while len(neg) < self.neg_sampling_size:
                nid = numpy.random.choice(self.sample_table)
                if nid in self.graph[u]:
                    continue
                neg.append(nid)
            self.edge_list[i][2] = neg

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

    def __getitem__(self, idx):
        return self.edge_list[idx]

class NodeTextEdgeDataset(Dataset):
    """Edge dataset."""

    def __init__(
        self, node_map1, node_map2, word_map, edge_list_file, text_file,
        neg_sampling_size
    ):
        """
        Args:
            edge_list (list): A list containts all edges. Each edge is a tuple
        """
        self.edge_list = []
        self.node_list = set()
        self.text_list = dict()
        for line in open(edge_list_file):
            line = line.strip().split(' ')
            line[0] = node_map1.str2id(line[0])
            line[1] = node_map2.str2id(line[1])
            self.edge_list.append([line[0], line[1], []])
            self.node_list.add(line[1])
        self.node_list = list(self.node_list)
        for line in open(text_file):
            line = line.replace('||||', ' ').strip().split(' ')
            self.text_list[node_map2.str2id(line[0])] = [
                word_map.str2id(w) for w in line[1:]
            ]
        self.neg_sampling_size = neg_sampling_size
        self.negative_sampling()

    def negative_sampling(self):
        numpy.random.seed(41)
        neg = numpy.random.choice(
            list(self.node_list),
            size=(len(self.edge_list), self.neg_sampling_size)
        )
        for i in range(len(self.edge_list)):
            self.edge_list[i][2] = list(neg[i])

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

    def __getitem__(self, idx):
        node, text, neg_text = self.edge_list[idx]
        text = self.text_list[text]
        neg_text = [self.text_list[t] for t in neg_text]
        return node, text, neg_text

class MultipleLableDataset(Dataset):
    """
    Multiple label dataset
    Input is a file, where each line contains:
        source_node, target_node and some labels
    """

    def __init__(
        self,
        u_map,
        v_map,
        label_filename,
        total_label_count,
        ratio=1.0,
        topk=-1,
        use_all_zeros=False
    ):
        """
        Args:
            edge_list (list): A list containts all edges. Each edge is a tuple
        """
        print('Use %0.2f labels' % ratio)
        self.labels = []
        fin = open(label_filename)
        fin.readline()
        not_count_topk = 0
        for idx, line in enumerate(fin):
            if numpy.random.rand() > ratio:
                continue
            line = line.strip().split(' ')
            line[0] = u_map.str2id(line[0])
            line[1] = v_map.str2id(line[1])
            label = [float(x) for x in line[2:]]
            label = numpy.array(label)
            if (label == 0).all():
                if not use_all_zeros:
                    continue
            if topk > 0:
                if len(set(label) - {0, 1}) > 0:
                    s = numpy.argsort(label)
                    label[s[-topk:]] = 1.0
                    label[s[:-topk]] = 0.0
                else:
                    not_count_topk += 1
            self.labels.append([line[0], line[1], label])
        print('not count topk', not_count_topk)

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

    def __getitem__(self, idx):
        return self.labels[idx]

if __name__ == '__main__':
    edge_list_path = "/mnt/store1/plus1lab/multilabel-data/v1_aminer_10k_1/original-edgelist.txt"
    edge_label_path = "/mnt/store1/plus1lab/multilabel-model/aegcn/v1_aminer_10k_1/prop-mat.txt"
    node_map = Str2id()
    label_map = Str2id()
    edges = []
    NormalEdgeDataset(
        node_map,
        node_map,
        edge_list_path,
        10,
        from_random_walk=True,
        window_size=10
    )
    node_map.freeze = True
    d = MultipleLableDataset(node_map, node_map, edge_label_path, 100, 1.0, 5)
    print(d[999])

max_stars_repo_name: GTmac/edge-enhanced-embeddings

max_issues_repo_name: GTmac/edge-enhanced-embeddings

max_forks_repo_name: GTmac/edge-enhanced-embeddings