text:
import torch
import torch.nn as nn

import numpy as np
import torch.nn.functional as f
from torch.autograd import Variable
import torch.nn.functional as F

class CONVKNRM(nn.Module):
    """Class that classifies question pair as duplicate or not."""

    def __init__(self, args):
        """"Constructor of the class."""
        super(CONVKNRM, self).__init__()

        self.wv=args.wv
        self.index_to_word=args.index_to_word

        self.input_dim=args.emsize
        self.device=args.device

        self.nbins = args.nbins

        self.dense_f = nn.Linear(self.nbins * 9, 1, 1)
        self.tanh = nn.Tanh()

        self.conv_uni = nn.Sequential(
            nn.Conv2d(1, 128, (1, self.input_dim)),
            nn.ReLU()
        )

        self.conv_bi = nn.Sequential(
            nn.Conv2d(1, 128, (2, self.input_dim)),
            nn.ReLU()
        )
        self.conv_tri = nn.Sequential(
            nn.Conv2d(1, 128, (3, self.input_dim)),
            nn.ReLU()
        )

        tensor_mu = torch.FloatTensor(args.mu).to(self.device)
        tensor_sigma = torch.FloatTensor(args.sigma).to(self.device)

        self.mu = Variable(tensor_mu, requires_grad=False).view(1, 1, 1, self.nbins)
        self.sigma = Variable(tensor_sigma, requires_grad=False).view(1, 1, 1, self.nbins)

    def get_intersect_matrix(self, q_embed, d_embed):
        sim = torch.bmm(q_embed, d_embed).view(q_embed.size()[0], q_embed.size()[1], d_embed.size()[2], 1)
        pooling_value = torch.exp((- ((sim - self.mu) ** 2) / (self.sigma ** 2) / 2))
        pooling_sum = torch.sum(pooling_value, 2)
        log_pooling_sum = torch.log(torch.clamp(pooling_sum, min=1e-10)) * 0.01
        log_pooling_sum = torch.sum(log_pooling_sum, 1)
        return log_pooling_sum

    def to_embedding(self,input):
        shape_input = list(input.shape)

        em = input.view(-1)
        list_of_embeddings = []
        for key in em:
            list_of_embeddings += self.wv[self.index_to_word[key]]
        list_of_embeddings = torch.Tensor(list_of_embeddings)
        embeds = list_of_embeddings.view(shape_input[0], shape_input[1],
                                         self.input_dim).to(self.device)

        # embeds = self.word_embeddings(input)
        ss = embeds.shape
        # embeds = embeds.view(ss[0],-1, ss[3], self.input_dim)

        #emb = torch.squeeze(embeds)  # batch_size * 1 * seq_len * emb_dim
        return embeds

    def forward(self, batch_queries, batch_docs,batch_semantic):

        qlen = batch_queries.shape[1]
        num_docs, dlen = batch_docs.shape[0], batch_docs.shape[1]

        batch_size=1

        emb_query = self.to_embedding(batch_queries)
        emb_desc = self.to_embedding(batch_docs)

        desc_att_shape = emb_desc.shape
        query_shape = emb_query.shape

        qwu_embed = torch.transpose(
            torch.squeeze(self.conv_uni(emb_query.view(emb_query.size()[0], 1, -1, self.input_dim))), 1,
            2) + 0.000000001

        qwb_embed = torch.transpose(
            torch.squeeze(self.conv_bi(emb_query.view(emb_query.size()[0], 1, -1, self.input_dim))), 1,
            2) + 0.000000001
        qwt_embed = torch.transpose(
            torch.squeeze(self.conv_tri(emb_query.view(emb_query.size()[0], 1, -1, self.input_dim))), 1,
            2) + 0.000000001

        dwu_embed = torch.squeeze(
            self.conv_uni(emb_desc.view(emb_desc.size()[0], 1, -1, self.input_dim))) + 0.000000001
        dwb_embed = torch.squeeze(
            self.conv_bi(emb_desc.view(emb_desc.size()[0], 1, -1, self.input_dim))) + 0.000000001
        dwt_embed = torch.squeeze(
            self.conv_tri(emb_desc.view(emb_desc.size()[0], 1, -1, self.input_dim))) + 0.000000001

        qwu_embed_norm = F.normalize(qwu_embed, p=2, dim=2, eps=1e-10)
        qwb_embed_norm = F.normalize(qwb_embed, p=2, dim=2, eps=1e-10)
        qwt_embed_norm = F.normalize(qwt_embed, p=2, dim=2, eps=1e-10)
        dwu_embed_norm = F.normalize(dwu_embed, p=2, dim=1, eps=1e-10)
        dwb_embed_norm = F.normalize(dwb_embed, p=2, dim=1, eps=1e-10)
        dwt_embed_norm = F.normalize(dwt_embed, p=2, dim=1, eps=1e-10)

        log_pooling_sum_wwuu = self.get_intersect_matrix(qwu_embed_norm, dwu_embed_norm)
        log_pooling_sum_wwut = self.get_intersect_matrix(qwu_embed_norm, dwt_embed_norm)
        log_pooling_sum_wwub = self.get_intersect_matrix(qwu_embed_norm, dwb_embed_norm)
        log_pooling_sum_wwbu = self.get_intersect_matrix(qwb_embed_norm, dwu_embed_norm)
        log_pooling_sum_wwtu = self.get_intersect_matrix(qwt_embed_norm, dwu_embed_norm)

        log_pooling_sum_wwbb = self.get_intersect_matrix(qwb_embed_norm, dwb_embed_norm)
        log_pooling_sum_wwbt = self.get_intersect_matrix(qwb_embed_norm, dwt_embed_norm)
        log_pooling_sum_wwtb = self.get_intersect_matrix(qwt_embed_norm, dwb_embed_norm)
        log_pooling_sum_wwtt = self.get_intersect_matrix(qwt_embed_norm, dwt_embed_norm)

        log_pooling_sum = torch.cat(
            [log_pooling_sum_wwuu, log_pooling_sum_wwut, log_pooling_sum_wwub, log_pooling_sum_wwbu,
             log_pooling_sum_wwtu,log_pooling_sum_wwbb, log_pooling_sum_wwbt, log_pooling_sum_wwtb, log_pooling_sum_wwtt], 1)

        output = torch.squeeze(F.tanh(self.dense_f(log_pooling_sum)), 1)

        return output

max_stars_repo_name: medtray/MultiEm-RGCN

max_issues_repo_name: medtray/MultiEm-RGCN

max_forks_repo_name: medtray/MultiEm-RGCN


text:
# Copyright 2019 Diamond Light Source Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
.. module:: morph_proc
   :platform: Unix
   :synopsis: Wrapper around skimage morphology routines for 2D morphological processing of binary datasets

.. moduleauthor:: Daniil Kazantsev <scientificsoftware@diamond.ac.uk>
"""

from savu.plugins.plugin import Plugin
from savu.plugins.driver.cpu_plugin import CpuPlugin
from savu.plugins.utils import register_plugin

import numpy as np
# importing skimage functions
from skimage.morphology import *

@register_plugin
class MorphProc(Plugin, CpuPlugin):

    def __init__(self):
        super(MorphProc, self).__init__("MorphProc")

    def setup(self):

        in_dataset, out_dataset = self.get_datasets()
        in_pData, out_pData = self.get_plugin_datasets()
        in_pData[0].plugin_data_setup(self.parameters['pattern'], 'single')

        out_dataset[0].create_dataset(in_dataset[0], dtype=np.uint8)
        out_pData[0].plugin_data_setup(self.parameters['pattern'], 'single')

    def pre_process(self):
        # extract given parameters
        self.disk_radius = self.parameters['disk_radius']
        self.morph_operation = self.parameters['morph_operation']
        self.selem = disk(self.disk_radius)

    def process_frames(self, data):
        # run morphological operations here:
        integerMax = np.max(data[0])
        if (np.sum(data[0]) > 0):
            morph_result = eval(self.morph_operation)(data[0], self.selem)
            morph_result = np.uint8(morph_result*integerMax)
        else:
            morph_result = np.uint8(np.zeros(np.shape(data[0])))
        return [morph_result]

    def nInput_datasets(self):
        return 1
    def nOutput_datasets(self):
        return 1

ext: py

max_stars_repo_name: elainehoml/Savu

max_issues_repo_name: elainehoml/Savu

max_forks_repo_name: elainehoml/Savu

reason: import numpy


text:
from __future__ import print_function  ## Force python3-like printing

import os
from collections import OrderedDict

import numpy as np
from astropy.constants import c
from astropy.io import ascii
from astropy.table import Table, Column, vstack, MaskedColumn
from astropy.time import Time

# from .classes import *
# from .functions import *
from . import classes
from . import functions

__all__ = ["LitLightCurveClass",
           "JD_to_MJD",
           "write_dict",
           "load_bianco_phot_table",
           "load_bianco_extinction",
           "get_EBV",
           "make_bianco_AB_phot",
           "translate_filter_names",
           "make_new_CfA_phot_table",
           "load_CfA_phot_table",
           "SN2009jf_read_lit",
           "SN2009jf_read_ap"]

class LitLightCurveClass(classes.BaseLightCurveClass):
    pass

translation_dict = OrderedDict()

translation_dict = {"U" : "BessellU",
                    "B" : "BessellB",
                    "V" : "BessellV",
                    "R" : "BessellR",
                    "I" : "BessellI",
                    "u'": "SDSS_u",
                    "g'": "SDSS_g",
                    "r'": "SDSS_r",
                    "i'": "SDSS_i",
                    "z'": "SDSS_z",
                    "u": "SDSS_u",
                    "g": "SDSS_g",
                    "r": "SDSS_r",
                    "i": "SDSS_i",
                    "z": "SDSS_z",
                    "J" : "2MASS_J",
                    "H" : "2MASS_H",
                    "Ks": "2MASS_Ks"}

def JD_to_MJD(JD_column):
    """
    convert astropy column object of JD floats into MJD's
    """
    MJD_col = Column(Time(JD_column, format = "jd").mjd, name = "MJD")
    return MJD_col

def write_dict(out_dir, sn_dict):
    """
    Each dict entry should be an astropy table. Taking this intermediate step in
    order to make it easier to get the formatting correct
    """
    for key in sn_dict:
        # print(key)
        out_path = os.path.join(out_dir, key + ".dat")
        print(out_path)
        sn_dict[key].write(out_path, format = "ascii.fast_commented_header")
    pass

def load_bianco_phot_table(verbose = False):
    """
    from http://iopscience.iop.org/article/10.1088/0067-0049/213/2/19
    Parameters
    ----------
    Returns
    -------
    """

    fname = '/Users/berto/data/CoreCollapse/phot/apjs496616t6_mrt.txt'
    return ascii.read(fname, data_start = 18)

def load_bianco_extinction(verbose = False):
    """
    from http://iopscience.iop.org/article/10.1088/0067-0049/213/2/19
    Parameters
    ----------
    Returns
    -------
    """
    fname = "/Users/berto/data/CoreCollapse/phot/apjs496616t2_mrt.txt"
    data = ascii.read(fname, format = "commented_header", delimiter = "\t")

    zcol = Column(data["cz"]*1000./c.value, name = "z") ## cz is in km s-1

    data.add_column(zcol)
    return data

def get_EBV(snname):
    data = load_bianco_extinction()
    wsn = np.where(data["snname"] == snname)
    return data["EBV"][wsn]

def make_bianco_AB_phot():
    data = load_bianco_phot_table()

    for sn_id in np.unique(data["SN"]):
        snname = "SN" + sn_id

        print(snname)

def translate_filter_names(data, verbose = False):

    translation_dict = OrderedDict()

    translation_dict = {
    "U" : "BessellU",
    "B" : "BessellB",
    "V" : "BessellV",
    "R" : "BessellR",
    "I" : "BessellI",
    "u'": "SDSS_u",
    "g'": "SDSS_g",
    "r'": "SDSS_r",
    "i'": "SDSS_i",
    "z'": "SDSS_z",
    "u": "SDSS_u",
    "g": "SDSS_g",
    "r": "SDSS_r",
    "i": "SDSS_i",
    "z": "SDSS_z",
    "J" : "2MASS_J",
    "H" : "2MASS_H",
    "Ks": "2MASS_Ks"
    }

    translated_arr = np.array([])
    for i, entry in enumerate(data):
        if verbose: print(i, entry)
        # data["NewFilter"][i] = translation_dict[entry["Filter"]]
        translated_arr = np.append(translated_arr, translation_dict[entry["Filter"]])
    newfilter_Column = Column(name = "Filter", data = translated_arr)
    print(newfilter_Column)
    data.remove_column("Filter")
    data["Filter"] = newfilter_Column

    return data

def make_new_CfA_phot_table(verbose = False):
    """
    from http://iopscience.iop.org/article/10.1088/0067-0049/213/2/19

    Parameters
    ----------
    Returns
    -------
    """

    data = load_bianco_phot_table(verbose = verbose)

    data = translate_filter_names(data, verbose = verbose)

    extinction_table = load_bianco_extinction(verbose = verbose)

    EBV_col = data['mag']*0
    EBV_col.name = "EBV"

    for entry in extinction_table:
        if verbose: print(entry)
        w = np.where(data["SN"] == entry["snname"])
        if verbose: print(w)
        EBV_col[w] = entry["EBV"]

    data.add_column(EBV_col)
    return data

def load_CfA_phot_table(verbose = False):
    fname = "/Users/berto/data/CoreCollapse/phot/apjs496616_combined_translated_table_RF.txt"
    data = ascii.read(fname)
    return data

def SN2009jf_read_lit():
    """
    Reads the tables I have collected for SN2009jf
    """

    sn2009jf_dict = OrderedDict()

    ## Get the Phot from Bianco et al.
    data = load_CfA_phot_table()
    sn2009jf_dict['sn2009jf_CfA'] = data[np.where(data['SN'] == '2009jf')]

    ## Get Valenti et al. 2001
    freadme = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/ReadMe"

    fnamec1 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec1.dat"
    fnamec2 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec2.dat"
    """
    These next few break due to the 'limits' flags in the table. Instead I have downloaded the tables
    in tab-separated-variable VO tables.
    One caveat to this is the the format and the dashed header row get read in too, even with
    data_start set.
    """
    # fnamec3 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec3.dat"
    # fnamec4 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec4.dat"
    # fnamec5 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec5.dat"
    # fnamec6 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/J_MNRAS_416_3138/tablec6.dat"

    fnamec3 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/tablec3.tsv"
    fnamec4 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/tablec4.tsv"
    fnamec5 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/tablec5.tsv"
    fnamec6 = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/tablec6.tsv"

    sn2009jf_dict['standards_UBVRI'] = ascii.read(fnamec1, readme = freadme)
    sn2009jf_dict['standards_ugriz'] = ascii.read(fnamec2, readme = freadme)

    sn2009jf_dict['sn2009jf_UBVRI'] = ascii.read(fnamec3)
    sn2009jf_dict['sn2009jf_UBVRI'].remove_rows([0,1])
    sn2009jf_dict['sn2009jf_ugriz'] = ascii.read(fnamec4)
    sn2009jf_dict['sn2009jf_ugriz'].remove_rows([0,1])
    sn2009jf_dict['sn2009jf_JHK'] = ascii.read(fnamec5)
    sn2009jf_dict['sn2009jf_JHK'].remove_rows([0,1])
    sn2009jf_dict['sn2009jf_swift'] = ascii.read(fnamec6)
    sn2009jf_dict['sn2009jf_swift'].remove_rows([0,1])

    return sn2009jf_dict

def SN2009jf_read_ap(format = "ascii", names = ('MJD', 'flux', 'flux_err', 'filter')):
    """
    Reads the output from write_dict(SN2009jf_read_lit) in order to unpack and
    correct the photometry
    """

    snjf = LitLightCurveClass()
    AB_zp = 48.60

    jfdict = OrderedDict()

    ## Vega Landolt
    sn2009jf_UBVRI = ascii.read("/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/astropy_tables/sn2009jf_UBVRI.dat")
    ## MJDs being calculated twice :@:@
    sn2009jf_UBVRI["mjd"] = Time(sn2009jf_UBVRI["JD"], format = "jd").mjd
    ## Unclear if limits are 3sigma or 5sigma - mask
    filter_names = ["U", "B", "V", "R", "I"]

    for filter_name in filter_names:
        print(filter_name + "mag")

        if filter_name in translation_dict:
            full_filter_name = translation_dict[filter_name]

        magmask = np.logical_not(sn2009jf_UBVRI[filter_name + "mag"].mask)
        vega_mag = sn2009jf_UBVRI[filter_name + "mag"]

        if "l_" + filter_name + "mag" in sn2009jf_UBVRI.keys():
            ulim_flag = sn2009jf_UBVRI["l_" + filter_name + "mag"]
            ## See above sigmas comment!!
            magmask = np.logical_and(ulim_flag.mask, magmask)

        AB_offset = kcorr.calc_offset_AB_minus_Vega(full_filter_name)## This is AB - Vega

        mjd_col = JD_to_MJD(sn2009jf_UBVRI["JD"])

        AB_mag = MaskedColumn(vega_mag + AB_offset, name = "mag")
        # e_AB_mag = Column(vega_mag, name = "e_" + AB_mag.name)
        e_AB_mag = MaskedColumn(sn2009jf_UBVRI["e_" + filter_name + "mag"], name = "emag")
        # print(vega_mag[magmask], AB_mag[magmask])
        # print(AB_mag + AB_zp)
        f_nu = MaskedColumn(np.power(10., (AB_mag + AB_zp)/-2.5), name = "f_nu")
        ef_nu = MaskedColumn(f_nu/(1.086/e_AB_mag), name = "ef_nu")
        # print(f_nu)

        filter_column = Column(np.array([full_filter_name for i in mjd_col]), name = "filter")

        phot_table = Table([mjd_col[magmask], AB_mag[magmask], e_AB_mag[magmask], f_nu[magmask], ef_nu[magmask], filter_column[magmask]])

        jfdict[full_filter_name] = phot_table

    sn2009jf_ugriz = ascii.read("/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/astropy_tables/sn2009jf_ugriz.dat")
    sn2009jf_ugriz["mjd"] = Time(sn2009jf_ugriz["JD"], format = "jd").mjd
    filter_names = ["u", "g", "r", "r", "i"]

    for filter_name in filter_names:
        print(filter_name + "mag")

        if filter_name in translation_dict:
            full_filter_name = translation_dict[filter_name]

        ## No ulims for ugriz phot
    #     ulim_flag = sn2009jf_ugriz["l_" + filter_name + "mag"]

        magmask = np.logical_not(sn2009jf_ugriz[filter_name + "mag"].mask)
        mjd_col = JD_to_MJD(sn2009jf_ugriz["JD"])

        AB_mag = MaskedColumn(sn2009jf_ugriz[filter_name + "mag"], name = "mag")
        e_AB_mag = MaskedColumn(sn2009jf_ugriz["e_" + filter_name + "mag"], name = "emag")

        f_nu = MaskedColumn(np.power(10., (AB_mag + AB_zp)/-2.5), name = "f_nu")
        ef_nu = MaskedColumn(f_nu/(1.086/e_AB_mag), name = "ef_nu")

        filter_column = Column(np.array([full_filter_name for i in mjd_col]), name = "filter")
        print(len(mjd_col), len(AB_mag), len(e_AB_mag), len(filter_column))
        phot_table = Table([mjd_col[magmask], AB_mag[magmask], e_AB_mag[magmask], f_nu[magmask], ef_nu[magmask], filter_column[magmask]])
        jfdict[full_filter_name] = phot_table

    # sn2009jf_JHK = ascii.read("/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/astropy_tables/sn2009jf_JHK.dat")
    # sn2009jf_JHK["mjd"] = Time(sn2009jf_JHK["JD"], format = "jd").mjd
    # sn2009jf_swift = ascii.read("/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/astropy_tables/sn2009jf_swift.dat")
    # sn2009jf_swift["mjd"] = Time(sn2009jf_swift["JD"], format = "jd").mjd

    filter_dir = '/Users/berto/Code/CoCo/data/filters/'

    # return jfdict
    for i, phot_filter in enumerate(jfdict.keys()):
        filter_path = os.path.join(filter_dir, phot_filter + ".dat")
        FilterObj = functions.load_filter(filter_path)
        FilterObj.calculate_frequency()
        FilterObj.calculate_effective_frequency()

        ## REMEBER lambda*f_lambda = nu*f_nu
        f_lambda = MaskedColumn((jfdict[phot_filter]["f_nu"]*FilterObj.nu_effective)/FilterObj.lambda_effective, name = "flux")
        ef_lambda = MaskedColumn((jfdict[phot_filter]["ef_nu"]*FilterObj.nu_effective)/FilterObj.lambda_effective, name = "flux_err")

        jfdict[phot_filter].add_column(f_lambda)
        jfdict[phot_filter].add_column(ef_lambda)

        if i == 0:
            full_phot = jfdict[phot_filter]
        else:
            full_phot = vstack([full_phot, jfdict[phot_filter]])

    # flux = np.power(10., (full_phot["mag"] + AB_zp)/-2.5)
    # return full_phot
    # return jfdict
    ## Write Combined Table
    final_phot_table = Table([full_phot["MJD"].filled(), full_phot["flux"].filled(), full_phot["flux_err"].filled(), full_phot["filter"].filled()], masked = False)
    outpath = "/Users/berto/data/CoreCollapse/phot/rf/SN2009jf/final/SN2009jf.dat"
    final_phot_table.write(outpath, format = "ascii.fast_commented_header")
    pass

#
# if __name__ == "__main__":
#     # print("Foo")
#     # print(get_EBV("2009jf"))
#     # make_bianco_AB_phot()
#
#     ## SN2009jf
#     data = load_bianco_phot_table()
#     snjf = data[np.where(data['SN'] == '2009jf')]
#
# else:
#     pass

max_stars_repo_name: RobFirth/pycoco

max_issues_repo_name: UoS-SNe/pycoco

max_forks_repo_name: RobFirth/pycoco

reason: import numpy,from astropy


text:
#!/usr/bin/env python3

##@package openzgy.impl.compress

import math
import sys
import numpy as np
from collections import namedtuple

from ..impl import enum as impl_enum
from ..exception import *

CompressionResult = namedtuple("Compression", "cdata csize isize signal noise ctime dtime snr_result snr_wanted snr_rounded snr_step")
CompressionResult.__doc__ = """The result returned from _xxx_try_one() functions.
The first member, cdata, is the one that really matters.
The snr_xxx members may be useful for goal seek.
All members except cdata may be useful logging and testing.
"""
CompressionResult.cdata.__doc__ = "bytes-like compressed data"
CompressionResult.csize.__doc__ = "Size in bytes of the compressed data"
CompressionResult.isize.__doc__ = "Size in bytes of the original input data"
CompressionResult.signal.__doc__ = "Sum of all sample values in the input"
CompressionResult.noise.__doc__ = "Sum of all absolute errors in round trip"
CompressionResult.ctime.__doc__ = "Time in seconds used for compression"
CompressionResult.dtime.__doc__ = "Time in seconds used for decompression"
CompressionResult.snr_result.__doc__ = "The SNR that was achieved"
CompressionResult.snr_wanted.__doc__ = "The SNR that was requested by the user"
CompressionResult.snr_rounded.__doc__ = "The SNR that was requested from the algorithm"
CompressionResult.snr_result.__doc__ = "Distance to next or previous quality level"

class CompressPlugin:
    """
    Base class for OpenZGY compression plug-ins.
    If anybody wants to add additional compression algorithms it
    is recommended but not required to use this base class. See
    CompressFactoryImpl.register{Compressor,Decompressor} for how to
    use plain functors (C++) or callables (Python) instead.

    This class performs triple duty as it handles both compression
    and decompression static methods (need not have been together)
    and an instance of the class can be used a compressor functor
    if a lambda is too limiting. To invoke the methods:

        MyCompressPlugin.factory(...)(data)
        MyCompressPlugin.compress(data, ...) (NOT recommended)
        MyCompressPlugin.decompress(cdata,status,shape,file_dtype,user_dtype)

    The following will also work but should only be used for very simple
    compressors that have no parameters. In the first case MyCompressPlugin
    won't have the option to return None for certain parameters, and in the
    second case handling a variable arguent list becomes trickier.

    To register this class:
      CompressFactoryImpl.registerCompressor("My",MyCompressPlugin.factory)
      CompressFactoryImpl.registerDecompressor("My",MyCompressPlugin.decompress)

    To use the compression part from client code:
       compressor = ZgyCompressFactory("My", ...)
    """
    @staticmethod
    def compress(data, *args, **kwargs):
        """
        This is an abstract method.

        Compress a 3d or (TODO-Low 2d) numpy array, returning a bytes-like
        result. If called with a single "data" argument the compression
        will be done with default parameters and no extended logging.

        Additional arguments are specific to the compression type.

        The function can be used directly as the compression hook.
        But you probably want a lambda expression or a real instance
        of this class instead, to be able to specify parameters.

        The compression algorithm is used is assumed to handle big / little
        endian conversion itself. TODO-Worry this is not quite true for ZFP.
        See the documentation. A special compilation flag is needed
        on big endian machines. Also I suspect the optional hedaer
        (which this code uses) might need byte swapping.
        """
        raise ZgyInternalError("Attempt to invoke abstract method")

    @staticmethod
    def decompress(cdata, status, shape, file_dtype, user_dtype):
        """
        This is an abstract method.

        Decompress bytes or similar into a numpy.ndarray.

        Arguments:
          cdata      -- bytes or bytes-like compressed data,
                        possibly with trailing garbage.
          status     -- Currently always BrickStatus.Compressed,
                        in the future the status might be used to
                        distinguish between different compression
                        algorithms instead of relying on magic numbers.
          shape      -- Rank and size of the result in case this is
                        not encoded by the compression algorithm.
          file_dtype -- Original value type before compression,
                        in case the decompressor cannot figure it out.
                        This will exactly match the dtype of the
                        data buffer passed to the compressor.
          user_dtype -- Required value type of returned array.

        Passing an uncompressed brick to this function is an error.
        We don't have enough context to handle uncompressed bricks
        that might require byteswapping and fix for legacy quirks.
        Also cannot handle constant bricks, missing bricks, etc.

        The reason user_dtype is needed is to avoid additional
        quantization noise when the user requests integer compressed data
        to be read as float. the decompressor might need to convert
        float data to int, only to have it converted back to float later.

        Current assumptions made of all candidate algorithms:

            -  The compressed data stream may have trailing garbage;
               this will be silently ignored by the decompressor.

            -  The compressed data stream will never be longer than
               the uncompressed data. This needs to be enforced by
               the compressor. The compressor is allowed to give up
               and tell the caller to not compress this brick.

            -  The reason for the two assumptions above is an
               implementation detail; the reported size of a
               compressed brick is not completely reliable.
               This might change in the next version

            -  The compressed data stream must start with a magic
               number so the decompressor can figure out whether
               this is the correct algorithm to use.

        If the assumptions cannot be met, the compressor / decompressor
        for this particular type could be modified to add an extra header
        with the compressed size and a magic number. Or we might add a
        (size, algorithm number) header to every compressed block to
        relieve the specific compressor / decompressor from worrying
        about this. Or the brick status could be used to encode which
        algorithm was used, picked up from the MSB of the lup entry.
        Which would also require the compressor to return both the
        actual compressed data and the code to identify the decompressor.
        That is the main reason we are also passed the "status" arg.
        Caveat: If adding an extra header, keep in mind that this header
        must be included when checking that the compressed stream is not
        too big.
        """
        raise ZgyInternalError("Attempt to invoke abstract method")

    def __init__(self, *args, **kwargs):
        """
        Create an instance that remembers the arguments it was created with.
        When the instance is called as a function it will invoke compress()
        with those arguments. So you can use either of the following:

            compressor = CompressPlugin.compress # no arguments
            compressor = CompressPlugin(...)
            compressor = lambda x: CompressPlugin.compress(x, ...)

        Derived classes don't need to redefine __init__ and __call__.
        But they might want to in order to get argument checking.
        The __init__ in the base class accepts any arguments so an error
        won't be caught until the first time the compressor is invoked.
        """
        self._args = args
        self._kwargs = kwargs

    def __call__(self, data):
        """
        Invoke the compressor with arguments passed by the constructor.
        """
        return self.compress(data, *args, **kwargs)

    def dump(*args, **kwargs):
        """
        Output statistics to standard output, if possible.
        """
        pass

class CompressFactoryImpl:
    """
    Registry of known compress and decompress algorithms.
    Those two are completely separate but we might as well
    handle both in the same class.
    """
    _compress_registry = {}
    _decompress_registry = []

    @staticmethod
    def registerCompressor(name, fn):
        """
        Register a factory function that will be called to create
        a function that in turn can be used to compress a data block.
        Pass fn = None if you for some reason need to remove a registration.

        The registered function can have any signature; the signature
        needs to include whatever parameters the actual compressor wants.

        The function that is created by the factory must have the signature:
            raw:          bytes.
            brickstatus:  impl.enum.BrickStatus,
            bricksize:    tuple(int,int,int),
            file_dtype:   np.dtype,
            result_dtype: np.dtype

        The function's return value:
            np.ndarray with rank 3, shape bricksize, and dtype result_dtype.

        Example usage of the factory:
            old: with ZgyWriter(snr=snr)
            new: with ZgyWriter(compressor = ZgyCompressFactory("ZFP", snr=30),

        The example shows that this is a bit more inconvenient for the end
        user. But it allows for using different compression plug-ins with
        arbitrary parameters.

        Note that user code doesn't need to use ZgyCompressFactory() and
        its list of known compressors. Instead a compression function
        can be provided directly. But most likely the factory method will
        be simpler to maintain.

        fn() is allowed to return None, which will have the same effect
        as if the user did not specify any compression. E.g. there might be
        a convention that snr<0 means store uncompressed. Allowing the
        factory to return None in that case means the example above would
        still work. Otherwise the snr<0 test would be included in the
        client code. Which is messy.
        """
        CompressFactoryImpl._compress_registry[name] = fn

    @staticmethod
    def registerDecompressor(name, fn):
        """
        Register a factory function that is able to decompress one or more
        types of compressed data. The registered functions will be called
        in reverse order of registration until one of them indicates that
        it has decompressed the data. You cannot remove a registration
        but you can effectively disable it by registering another one
        that recognizes the same input data.
        The supplied name is only for information.

        The function that is created by the factory must have the signature:
            raw:          bytes.
            brickstatus:  impl.enum.BrickStatus,
            bricksize:    tuple(int,int,int),
            file_dtype:   np.dtype,
            result_dtype: np.dtype
        """
        CompressFactoryImpl._decompress_registry.insert(0, (name, fn))

    def knownCompressors():
        """
        Return the names of all compressors known to the system.
        This is primarily for logging, but might in principle be used
        in a GUI to present a list of compressors to choose from.
        The problem with that is how to handle the argument list.
        """
        return list([k for k, v in CompressFactoryImpl._compress_registry.items()])

    def knownDecompressors():
        """
        Return the names of all compressors known to the system.
        This is primarily for logging.
        """
        return list([k for k, v in CompressFactoryImpl._decompress_registry])

    @staticmethod
    def factory(name, *args, **kwargs):
        fn = CompressFactoryImpl._compress_registry.get(name, None)
        if not fn:
            known = ",".join(CompressFactoryImpl.knownDecompressors())
            raise ZgyMissingFeature('Compression algorithm "{0}" not recognized. Must be one of ({1}).'.format(name, known))
        return fn(*args, **kwargs)

    @staticmethod
    def decompress(cdata, status, shape, file_dtype, user_dtype):
        """
        Loop over all registered decompressors and try to find one that
        can handle this particular brick. Raises an error if none found.
        See CompressPlugin.decompress() for parameter descriptions.
        """
        result = None
        if status != impl_enum.BrickStatus.Compressed:
            raise ZgyInternalError("Tried to decompress uncompressed data.")
        for k, v in CompressFactoryImpl._decompress_registry:
            result = v(cdata, status, shape, file_dtype, user_dtype)
            if result is not None:
                break
        if result is None:
            raise ZgyFormatError("Compression algorithm not recognized.")
        elif tuple(result.shape) != tuple(shape):
            raise ZgyFormatError("Decompression returned unexpected data.")
        return result

class CompressStats:
    def __init__(self, details):
        self._details  = details # Parameters etc, just for logging.
        self._signal   = 0
        self._noise    = 0
        self._original = 0
        self._packed   = 0
        self._types    = dict()
        self._lossy    = 0
        self._perfect  = 0
        self._all_info = [] # (measured_snr, measured_compressed_size_%)
        self._timed_bytes = 0
        self._ctime = 0
        self._dtime = 0

    def add_data(self, idata, csize, ddata, *, ctime = None, dtime = None, msg = None):
        signal, noise = self._collect_snr(idata, ddata)
        isize = idata.size * idata.itemsize
        snr = self._compute_snr(signal, noise)
        self.add(signal=signal, noise=noise, snr=snr, isize=isize, csize=csize,
                 ctime=ctime, dtime=dtime, msg=msg)
        #print("signal {0:8.0f} noise {1:8.0f} snr {2:.2f}".format(signal, noise, snr))

    def add(self, signal, noise, snr, isize, csize, *, ctime = None, dtime = None, msg = None):
        self._all_info.append([snr, 100*csize/isize])
        if snr < 99:
            # s/n only logged for lossy compression.
            self._lossy    += 1
            self._signal   += signal
            self._noise    += noise
        else:
            self._perfect  += 1
        self._original += isize
        self._packed   += csize
        if ctime is not None and dtime is not None:
            self._timed_bytes += isize
            self._ctime += ctime
            self._dtime += dtime
        if msg == "compress" and snr >= 99: msg = "compres*"
        if msg:
            self._types[msg] = 1 + self._types.setdefault(msg, 0)
        if False and msg:
            print("@@ yoo-hoo {0} {1} -> {2} factor {3:.1f} snr {4:.1f}".format(
                msg, isize, csize, isize/csize if csize else np.nan, snr))

    @staticmethod
    def _collect_snr(idata, ddata):
        """
        This function along with _compute_snr defines the cost function
        used when trying to quantify noise. The input data should not
        contain NaN or Inf. We won't be applying lossy compression to
        that kind of data anyway.

        There is no perfect solution. We want to avoid having
        neighboring bricks showning significant differences in
        quality. This sounds easy enough, but "similar quality"
        depends very much on the cost function.
        """
        # Computationally inexpensive, but has problems with spikes.
        # A brick containing many spikes will appear to have a higher
        # average amplitude, which allows the compressor to leave
        # more noise in it. It would have been better to use the
        # average signal level of the entire survey instead of the
        # brick. But that information is not available yet.
        signal = np.sum(np.abs(idata), dtype=np.float64)
        noise  = np.sum(np.abs(idata - ddata), dtype=np.float64)
        # Expensive, and has problems with dead traces and bricks
        # that contain more than 50% water. The signal and noise
        # in thise cases would be measuring only the dead and/or
        # water samples.
        #signal = np.median(np.abs(idata)) * idata.size
        #noise  = np.median(np.abs(idata - ddata)) * idata.size
        return signal, noise

    @staticmethod
    def _compute_snr(s, n):
        return 99 if n == 0 else -10 if s/n < 0.315 else min(99, 6 * math.log2(s/n))

    def snr(self):
        return self._compute_snr(self._signal, self._noise)

    def dump(self, msg = None, *, outfile=None, text=True, csv=False):
        outfile = outfile or sys.stdout
        # Measured SNR for each brick, excluding lossless unless everything
        # is lossless. In which case we pretend there is just one brick.
        # The reason is that the median compression makes more sense
        # that way. If I had been plotting all percentiles instead of
        # just printing the median it would have made more sense to include
        # the lossless (SNR 99) bricks as well.
        all_snr = list([x[0] for x in self._all_info if x[0] < 99]) or [99.0]
        all_snr.sort()
        # Measured compression ratio for each brick, excluding uncompressed
        # unless everything is uncompressed. In which case we pretend there
        # is just one brick. Note that lossless can still be compressed.
        ratios = list([ 100.0/e[1] for e in self._all_info if e[1] != 1 and e[1] != 0 ]) or [1.0]
        ratios.sort()
        good = self._original and self._packed
        args = {
            "msg":          msg or "Processed",
            "details":      self._details,
            "orig_mb":      self._original / (1024*1024),
            "pack_mb":      self._packed,
            "factor":       self._original / self._packed if good else -1,
            "percent":      100 * self._packed / self._original if good else 1,
            "n_perfect":    self._perfect,
            "n_lossy":      self._lossy,
            "snr":          self.snr(),
            "snr_min":      all_snr[0],
            "snr_med":      all_snr[len(all_snr)//2],
            "snr_max":      all_snr[-1],
            "factor_min":   ratios[0],
            "factor_med":   ratios[len(ratios)//2],
            "factor_max":   ratios[-1],
            "percent_max":  100 / ratios[0],
            "percent_med":  100 / ratios[len(ratios)//2],
            "percent_min":  100 / ratios[-1],
            "ctime":        0 if not self._ctime else (self._timed_bytes / self._ctime) / (1024*1024),
            "dtime":        0 if not self._dtime else (self._timed_bytes / self._dtime) / (1024*1024),
            "_ctime": self._ctime,
            "_dtime": self._dtime,
            "_cplusdtime": self._ctime + self._dtime,
            "_timed_bytes": self._timed_bytes / (1024*1024),
        }
        if text:
            print("{msg} {orig_mb:.0f} MB, compression factor {factor:.1f} (median {factor_med:.1f} with {n_perfect} bricks lossless and {n_lossy} bricks min snr {snr_min:.1f} median {snr_med:.1f} overall {snr:.1f}) compress {ctime:.1f} MB/s decompress {dtime:.1f} MB/s".format(**args), file=outfile)
            for k in sorted(self._types):
                print("   ", k, self._types[k], file=outfile)
        if csv:
            if csv == "header":
                print(";msg;requested snr;median snr;snr;compressed size%;median size%;compress MB/s;decompress MB/s;compress elapsed;decompress elapsed;sum elapsed;Total data MB", file=outfile)
            else:
                print(";{msg};{details};{snr_med:.1f};{snr:.1f};{percent:.1f};{percent_med:.1f};{ctime:.1f};{dtime:.1f};{_ctime:.1f};{_dtime:.1f};{_cplusdtime:.1f};{_timed_bytes:.1f}".format(**args), file=outfile)

    def empty(self):
            return self._lossy + self._perfect == 0

# Copyright 2017-2020, Schlumberger
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

max_stars_repo_name: equinor/pyzgy

max_issues_repo_name: equinor/pyzgy

max_forks_repo_name: equinor/pyzgy


text:
# -*- coding: utf-8 -*-
# Copyright (c) Hebes Intelligence Private Company

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import dataclasses
import logging
from datetime import datetime, time, timedelta
from itertools import product
from typing import Dict, Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    filename="eevalue.log",
    filemode="w",
)

import numpy as np
import pandas as pd
from pydantic import validator
from pydantic.dataclasses import dataclass
from pyomo.environ import (
    ConcreteModel,
    Constraint,
    NonNegativeReals,
    Objective,
    Param,
    Set,
    SolverFactory,
    Var,
    minimize,
)
from pyomo.opt import TerminationCondition

class ValidationConfig:
    arbitrary_types_allowed = True
    underscore_attrs_are_private = True

@dataclass(config=ValidationConfig)
class ModelData:
    # Required time data
    simulation_period: pd.DatetimeIndex
    # Required market data
    availability_factors: pd.DataFrame
    clustered_plants: pd.DataFrame
    demand: pd.DataFrame
    net_imports: pd.DataFrame
    reserve_technologies: Dict
    # Required cost components
    fuel_price: pd.DataFrame
    no_load_cost: float
    permit_price: pd.DataFrame
    # Optional cost components
    markup: Optional[pd.DataFrame] = None
    voll: Optional[float] = None
    # Renewables
    avail_factor_pv: Optional[pd.Series] = None
    avail_factor_wind: Optional[pd.Series] = None
    capacity_pv: Optional[pd.Series] = None
    capacity_wind: Optional[pd.Series] = None
    cost_curtailment: Optional[float] = None
    res_generation: Optional[pd.DataFrame] = None
    # Initial state
    committed_initial: Optional[Dict] = None
    power_initial: Optional[Dict] = None
    storage_initial: Optional[float] = None
    # Storage
    storage_charge_ef: Optional[float] = None
    storage_discharge_ef: Optional[float] = None
    storage_capacity: Optional[float] = None
    storage_charge_cap: Optional[float] = None
    storage_min: Optional[float] = None
    storage_final_min: Optional[float] = None
    # Other
    calibration: dataclasses.InitVar[bool] = False
    max_load_shaping: Optional[float] = None

    def __post_init__(self, calibration):
        if calibration:
            if self.res_generation is None:
                raise ValueError("RES generation must be provided for calibration")
        else:
            if self.avail_factor_pv is None:
                raise ValueError(
                    "Solar PV availability must be provided for what-if analysis"
                )
            if self.avail_factor_wind is None:
                raise ValueError(
                    "Wind availability must be provided for what-if analysis"
                )
            if self.capacity_pv is None:
                raise ValueError(
                    "Solar PV capacity must be provided for what-if analysis"
                )
            if self.capacity_wind is None:
                raise ValueError("Wind capacity must be provided for what-if analysis")

    @validator("simulation_period")
    def check_no_timestamps_missing(cls, data):
        time_col = data.to_series()
        time_step = time_col.diff().min()

        if time_step != pd.Timedelta("1H"):
            raise ValueError("`simulation_period` must have hourly time step")
        if len(data) < 24:
            raise ValueError(
                "The `simulation_period` range should span at least one full day"
            )
        full_index = pd.date_range(
            start=datetime.combine(time_col.min().date(), time(0, 0)),
            end=datetime.combine(time_col.max().date() + timedelta(days=1), time(0, 0)),
            freq=time_step,
        )[:-1]
        if len(full_index) > len(data):
            raise ValueError("Missing hours in `simulation_period`")
        return data

    @validator("availability_factors", "avail_factor_pv", "avail_factor_wind")
    def check_availability_factor(cls, data):
        if data is not None:
            if (data.values < 0).any() or (data.values > 1).any():
                raise ValueError(
                    "The availability factor's values are outside of the [0,1] interval"
                )
        return data

    @validator("clustered_plants")
    def check_expected_fields(cls, data):
        expected_fields = [
            "Cluster",
            "Nunits",
            "PowerCapacity",
            "PowerMinStable",
            "Efficiency",
            "RampStartUpRate",
            "CO2Intensity",
            "MinDownTime",
            "MinUpTime",
            "RampUpRate",
            "RampDownRate",
            "CostRampUp",
            "Technology",
            "Fuel",
        ]
        if not set(expected_fields).issubset(data.columns):
            raise ValueError("Missing fields from plant cluster dataframe")
        return data

    @validator(
        "availability_factors",
        "avail_factor_pv",
        "avail_factor_wind",
        "capacity_pv",
        "capacity_wind",
        "demand",
        "res_generation",
        "markup",
        "net_imports",
        "fuel_price",
        "permit_price",
    )
    def check_time_index(cls, data, values):
        if data is not None:
            data.index = data.index.map(pd.to_datetime)
            time_step = data.index.to_series().diff().min()

            if time_step != pd.Timedelta("1H"):
                raise ValueError(
                    "The index of the dataframe has not an hourly time step"
                )
            if not data.index.is_monotonic_increasing:
                raise ValueError(
                    "The index of the dataframe is not monotonic increasing"
                )
            if not values["simulation_period"].isin(data.index).all():
                raise ValueError(
                    "The `simulation_period` range is not a subset of the dataframe's index"
                )
        return data

    def __post_init_post_parse__(self, calibration):
        self._sets = dict()
        self._sets["mk"] = ["DA", "2U", "2D"]
        self._sets["cl"] = self.clustered_plants["Cluster"].to_list()
        self._sets["f"] = self.clustered_plants["Fuel"].unique().tolist()
        self._sets["t"] = self.clustered_plants["Technology"].unique().tolist()
        self._sets["h"] = (
            np.array(
                (
                    self.simulation_period - self.simulation_period[0].to_pydatetime()
                ).total_seconds()
                / 3600
            )
            .astype("int")
            .tolist()
        )

    @property
    def sets(self):
        return self._sets

def build_market_params(model_data: ModelData):
    parameters = dict()
    parameters["Reserve"] = model_data.reserve_technologies
    parameters["NoLoadCost"] = model_data.no_load_cost
    parameters["ValueOfLostLoad"] = model_data.voll or 100e3
    parameters["MaxLoadShaping"] = model_data.max_load_shaping or 0
    return parameters

def build_cluster_params(model_data: ModelData):
    parameters = dict()
    plants = model_data.clustered_plants.set_index("Cluster")

    parameters["Nunits"] = plants["Nunits"].to_dict()
    parameters["PowerCapacity"] = plants["PowerCapacity"].to_dict()
    parameters["PowerMinStable"] = plants["PowerMinStable"].to_dict()
    parameters["Efficiency"] = plants["Efficiency"].to_dict()
    parameters["EmissionRate"] = plants["CO2Intensity"].to_dict()
    parameters["TimeDownMinimum"] = plants["MinDownTime"].to_dict()
    parameters["TimeUpMinimum"] = plants["MinUpTime"].to_dict()
    parameters["RampUpMaximum"] = plants["RampUpRate"].to_dict()
    parameters["RampDownMaximum"] = plants["RampDownRate"].to_dict()
    parameters["RampStartUpMaximum"] = plants["RampStartUpRate"].to_dict()
    parameters["CostRampUp"] = plants["CostRampUp"].to_dict()

    parameters["RampShutDownMaximum"] = (
        {cl: 0 for cl in model_data.sets["cl"]}
        if "RampShutDownRate" not in plants.columns
        else plants["RampStartUpRate"].to_dict()
    )
    parameters["CostRampDown"] = (
        {cl: 0 for cl in model_data.sets["cl"]}
        if "CostRampDown" not in plants.columns
        else plants["CostRampDown"].to_dict()
    )
    parameters["Technology"] = {
        (cl, t): 1 if plants.loc[cl, "Technology"] == t else 0
        for (cl, t) in product(model_data.sets["cl"], model_data.sets["t"])
    }
    parameters["Fuel"] = {
        (cl, f): 1 if plants.loc[cl, "Fuel"] == f else 0
        for (cl, f) in product(model_data.sets["cl"], model_data.sets["f"])
    }
    return parameters

def build_time_series_params(model_data: ModelData):
    parameters = dict()

    def align_time_index(data):
        data = data.loc[model_data.simulation_period].reset_index(drop=True)
        if set(data.index).difference(model_data.sets["h"]):
            raise RuntimeError("Dataframe index cannot match `h` index")
        return data

    avail_factors = align_time_index(model_data.availability_factors)
    parameters["AF"] = {}
    for cl in model_data.sets["cl"]:
        parameters["AF"].update(
            {(cl, h): val for h, val in avail_factors[cl].to_dict().items()}
        )

    demand = align_time_index(model_data.demand)
    parameters["Demand"] = {}
    parameters["Demand"].update(
        {("DA", h): val for h, val in demand["DA"].to_dict().items()}
    )
    parameters["Demand"].update(
        {("2U", h): val for h, val in demand["2U"].to_dict().items()}
    )
    parameters["Demand"].update(
        {("2D", h): val for h, val in demand["2D"].to_dict().items()}
    )

    net_imports = align_time_index(model_data.net_imports)
    parameters["NetImports"] = net_imports.iloc[:, 0].to_dict()

    parameters["FuelPrice"] = dict()
    fuel_price = align_time_index(model_data.fuel_price)
    for fuel in model_data.sets["f"]:
        parameters["FuelPrice"].update(
            {(fuel, h): val for h, val in fuel_price[fuel].to_dict().items()}
        )

    permit_price = align_time_index(model_data.permit_price)
    parameters["PermitPrice"] = permit_price.iloc[:, 0].to_dict()

    parameters["Markup"] = dict()
    if model_data.markup is not None:
        markup = align_time_index(model_data.markup)
        for cl in model_data.sets["cl"]:
            parameters["Markup"].update(
                {(cl, h): val for h, val in markup[cl].to_dict().items()}
            )
    else:
        for cl in model_data.sets["cl"]:
            parameters["Markup"].update({(cl, h): 0 for h in model_data.sets["h"]})

    if model_data.avail_factor_pv is not None:
        avail_factor_pv = align_time_index(model_data.avail_factor_pv)
        parameters["AF_Solar"] = avail_factor_pv.iloc[:, 0].to_dict()

    if model_data.avail_factor_wind is not None:
        avail_factor_wind = align_time_index(model_data.avail_factor_wind)
        parameters["AF_Wind"] = avail_factor_wind.iloc[:, 0].to_dict()

    if model_data.res_generation is not None:
        res_generation = align_time_index(model_data.res_generation)
        parameters["RES_Generation"] = res_generation.iloc[:, 0].to_dict()

    return parameters

def build_storage_params(model_data: ModelData):
    parameters = dict()
    parameters["StorageChargingEfficiency"] = model_data.storage_charge_ef or 1
    parameters["StorageDischargeEfficiency"] = model_data.storage_discharge_ef or 1
    parameters["StorageCapacity"] = model_data.storage_capacity or 0
    parameters["StorageChargingCapacity"] = model_data.storage_charge_cap or 0
    parameters["StorageMinimum"] = model_data.storage_min or 0
    parameters["StorageFinalMin"] = model_data.storage_final_min or 0
    return parameters

def build_renewables_params(model_data: ModelData):
    parameters = dict()
    parameters["CostCurtailment"] = model_data.cost_curtailment or 0
    parameters["SolarCapacity"] = model_data.capacity_pv or 0
    parameters["WindCapacity"] = model_data.capacity_wind or 0
    return parameters

def add_initial_state(model_data: ModelData, parameters: Dict):
    try:
        parameters["CommittedInitial"] = (
            model_data.committed_initial
            if model_data.committed_initial is not None
            else {cl: val for cl, val in parameters["Nunits"].items()}
        )
        parameters["PowerInitial"] = (
            model_data.power_initial
            if model_data.power_initial is not None
            else {
                cl: parameters["PowerCapacity"][cl] * parameters["Nunits"][cl]
                for cl in model_data.sets["cl"]
            }
        )
        parameters["StorageInitial"] = model_data.storage_initial or 0
    except KeyError as exc:
        raise ValueError(
            "Build cluster parameters before adding initial state"
        ) from exc
    else:
        return parameters

def build_model_params(model_data: ModelData, calibration=False):
    parameters = build_cluster_params(model_data)
    parameters.update(build_market_params(model_data))
    parameters.update(build_time_series_params(model_data))
    parameters.update(build_storage_params(model_data))
    if not calibration:
        parameters.update(build_renewables_params(model_data))
    parameters = add_initial_state(model_data, parameters)
    return model_data.sets, parameters

def create_model(
    sets,
    params,
    load_shape_mode=None,
    calibration=False,
):
    """
    Creates a concrete Pyomo model for unit commitment optimization.

    Args:
        sets: Dictionary containing the model's sets (defined as a list of strings or integers).
        params: Dictionary containing the parameters of the optimization problem
        load_shape_mode: String indicating the presence and type of the load modifying resources.
            It can be None, 'savings', 'flex'
        calibration: Boolean. Indicates whether the model will be used for calibration or what-if
            analysis

    Returns:
        The Pyomo optimization model instance.

    Adapted from https://github.com/energy-modelling-toolkit/Dispa-SET/tree/v2.3
    """

    # Definition of model:
    model = ConcreteModel("UCM")

    ##############################################################################################
    # Definition of the sets
    ##############################################################################################

    model.h = Set(initialize=sets["h"])  # Hours
    model.cl = Set(initialize=sets["cl"])  # Generation clusters
    model.t = Set(initialize=sets["t"])  # Generation technologies
    model.f = Set(initialize=sets["f"])  # Fuel types
    model.mk = Set(
        initialize=sets["mk"]
    )  # Markets (DA: Day-Ahead, 2U: Reserve up, 2D: Reserve Down)

    ###############################################################################################
    # Definition of the Parameters
    ###############################################################################################

    model.AF = Param(
        model.cl, model.h, initialize=params["AF"]
    )  # [%] Availability factor of dispatchable capacity
    model.CommittedInitial = Param(
        model.cl, initialize=params["CommittedInitial"]
    )  # [MW] Capacity committed at the beginning of the simulation period
    model.CostRampUp = Param(
        model.cl, initialize=params["CostRampUp"]
    )  # [EUR\MW] Ramp-up costs
    model.CostRampDown = Param(
        model.cl, initialize=params["CostRampDown"]
    )  # [EUR\MW] Ramp-down costs
    model.Demand = Param(
        model.mk, model.h, initialize=params["Demand"]
    )  # [MW] Demand per market
    model.Efficiency = Param(
        model.cl, initialize=params["Efficiency"]
    )  # [%] Efficiency per technology cluster
    model.EmissionRate = Param(
        model.cl, initialize=params["EmissionRate"]
    )  # [tCO2/MWh] CO2 emission rate
    model.Fuel = Param(
        model.cl, model.f, initialize=params["Fuel"]
    )  # [n.a.] Fuel type {1 0}
    model.FuelPrice = Param(
        model.f, model.h, initialize=params["FuelPrice"]
    )  # [€/MWh] Fuel price
    model.Markup = Param(model.cl, model.h, initialize=params["Markup"])  # [€/MW]
    model.MaxLoadShaping = Param(
        initialize=params["MaxLoadShaping"]
    )  # [MW] Maximum amount of load shaping resources during the simulation period
    model.NetImports = Param(
        model.h, initialize=params["NetImports"]
    )  # [MW] Net imports
    model.NoLoadCost = Param(
        initialize=params["NoLoadCost"]
    )  # [€/MWh] The cost of fuel required to keep a unit of the cluster running
    model.Nunits = Param(
        model.cl, initialize=params["Nunits"]
    )  # Number of units inside the cluster
    model.PermitPrice = Param(
        model.h, initialize=params["PermitPrice"]
    )  # [EUR/ton] CO2 emission permit price
    model.PowerCapacity = Param(
        model.cl, initialize=params["PowerCapacity"]
    )  # [MW] Installed capacity
    model.PowerInitial = Param(
        model.cl, initialize=params["PowerInitial"]
    )  # [MW] Power output before initial period
    model.PowerMinStable = Param(
        model.cl, initialize=params["PowerMinStable"]
    )  # [MW] Minimum power for stable generation
    model.RampDownMaximum = Param(
        model.cl, initialize=params["RampDownMaximum"]
    )  # [MW/h] Ramp down limit
    model.RampUpMaximum = Param(
        model.cl, initialize=params["RampUpMaximum"]
    )  # [MW\h] Ramp up limit
    model.RampStartUpMaximum = Param(
        model.cl, initialize=params["RampStartUpMaximum"]
    )  # [MW\h] Start-up ramp limit
    model.RampShutDownMaximum = Param(
        model.cl, initialize=params["RampShutDownMaximum"]
    )  # [MW\h] Shut-down ramp limit
    model.Reserve = Param(
        model.t, initialize=params["Reserve"]
    )  # [n.a.] Reserve technology {1 0}
    model.StorageChargingEfficiency = Param(
        initialize=params["StorageChargingEfficiency"]
    )  # [%] Charging efficiency
    model.StorageCapacity = Param(
        initialize=params["StorageCapacity"]
    )  # [MWh] Storage capacity
    model.StorageChargingCapacity = Param(
        initialize=params["StorageChargingCapacity"]
    )  # [MW] Storage capacity
    model.StorageDischargeEfficiency = Param(
        initialize=max(0.0001, params["StorageDischargeEfficiency"])
    )  # [%] Discharge efficiency
    model.StorageFinalMin = Param(initialize=params["StorageFinalMin"])
    model.StorageInitial = Param(
        initialize=params["StorageInitial"]
    )  # [MWh] Storage level before initial period
    model.StorageMinimum = Param(
        initialize=params["StorageMinimum"]
    )  # [MWh] Storage minimum
    model.Technology = Param(
        model.cl, model.t, initialize=params["Technology"]
    )  # [n.a] Technology type {1 0}
    model.TimeDownMinimum = Param(
        model.cl, initialize=params["TimeDownMinimum"]
    )  # [h] Minimum down time
    model.TimeUpMinimum = Param(
        model.cl, initialize=params["TimeUpMinimum"]
    )  # [h] Minimum up time
    model.ValueOfLostLoad = Param(
        initialize=params["ValueOfLostLoad"]
    )  # [€/MWh] Value of lost load

    def _var_cost_init(model, cl, h):
        return (
            model.Markup[cl, h]
            + sum(
                model.Fuel[cl, f] * model.FuelPrice[f, h] / model.Efficiency[cl]
                for f in model.f
            )
            + model.PermitPrice[h] * model.EmissionRate[cl]
        )

    model.CostVariable = Param(
        model.cl, model.h, initialize=_var_cost_init
    )  # [EUR\MWh] Variable costs

    if not calibration:
        model.AF_Solar = Param(model.h, initialize=params["AF_Solar"])  # [%]
        model.AF_Wind = Param(model.h, initialize=params["AF_Wind"])  # [%]
        model.CostCurtailment = Param(initialize=params["CostCurtailment"])  # [EUR\MWh]

        model.SolarCapacity = Param(
            initialize=params["SolarCapacity"]
        )  # [MW] Installed capacity
        model.WindCapacity = Param(
            initialize=params["WindCapacity"]
        )  # [MW] Installed capacity
    else:
        model.RES_Generation = Param(
            model.h, initialize=params["RES_Generation"]
        )  # [MW] RES generation

    ################################################################################################
    # Definition of variables
    ################################################################################################

    def _nunit_rule(model, cl, h):
        return (0, model.Nunits[cl])

    model.Committed = Var(
        model.cl, model.h, within=NonNegativeReals, bounds=_nunit_rule
    )  # Number of units of cluster cl committed at hour h
    model.CostRampUpH = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [EUR/h] Cost of ramping up
    model.CostRampDownH = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [EUR/h] Cost of ramping down
    model.LL_RampDown = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [MW] Deficit in terms of ramping down
    model.LL_RampUp = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [MW] Deficit in terms of ramping up
    model.LL_MaxPower = Var(
        model.h, within=NonNegativeReals
    )  # [MW] Deficit in terms of maximum power
    model.LL_MinPower = Var(
        model.h, within=NonNegativeReals
    )  # [MW] Power exceeding the demand
    model.LL_2U = Var(model.h, within=NonNegativeReals)  # [MW] Deficit in reserve up
    model.LL_2D = Var(model.h, within=NonNegativeReals)  # [MW] Deficit in reserve down
    model.LoadMod_U = Var(
        model.h, within=NonNegativeReals
    )  # [MW] Positive difference from the baseline
    model.LoadMod_D = Var(
        model.h, within=NonNegativeReals
    )  # [MW] Negative difference from the baseline
    model.Power = Var(model.cl, model.h, within=NonNegativeReals)  # [MW] Power output
    model.Reserve_2U = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [MW] Spinning reserve up
    model.Reserve_2D = Var(
        model.cl, model.h, within=NonNegativeReals
    )  # [MW] Spinning reserve down
    model.StartUp = Var(model.cl, model.h, within=NonNegativeReals, bounds=_nunit_rule)
    model.ShutDown = Var(model.cl, model.h, within=NonNegativeReals, bounds=_nunit_rule)
    model.StorageLevel = Var(
        model.h, within=NonNegativeReals
    )  # [MWh] Storage level of charge
    model.StorageInput = Var(
        model.h, within=NonNegativeReals
    )  # [MWh] Charging input for storage units
    model.StorageOutput = Var(
        model.h, within=NonNegativeReals
    )  # [MWh] Discharging output of storage units
    model.SystemCost = Var(model.h, within=NonNegativeReals)  # [EUR] Hourly system cost

    if not calibration:
        model.CurtailedPower = Var(
            model.h, within=NonNegativeReals
        )  # [MW] Curtailed power
        model.Power_RES = Var(
            model.h, within=NonNegativeReals
        )  # [MW] Power consumed from RES

    ###############################################################################################
    # EQUATIONS
    ###############################################################################################

    def EQ_Commitment(model, cl, h):
        if h == 0:
            return (
                model.Committed[cl, h] - model.CommittedInitial[cl]
                == model.StartUp[cl, h] - model.ShutDown[cl, h]
            )
        else:
            return (
                model.Committed[cl, h] - model.Committed[cl, h - 1]
                == model.StartUp[cl, h] - model.ShutDown[cl, h]
            )

    def EQ_MinUpTime(model, cl, h):
        for h in model.h:
            if h > model.TimeUpMinimum[cl]:
                return model.Committed[cl, h] >= sum(
                    model.StartUp[cl, h - i] for i in range(model.TimeUpMinimum[cl])
                )
            else:
                return Constraint.Skip

    def EQ_MinDownTime(model, cl, h):
        for h in model.h:
            if h > model.TimeDownMinimum[cl]:
                return model.Nunits[cl] - model.Committed[cl, h] >= sum(
                    model.ShutDown[cl, h - i] for i in range(model.TimeDownMinimum[cl])
                )
            else:
                return Constraint.Skip

    # Power output with respect to power output in the previous period:
    def EQ_RampUp_TC(model, cl, h):
        if h == 0:
            return (
                model.Power[cl, h] - model.PowerInitial[cl]
                <= (model.Committed[cl, h] - model.StartUp[cl, h])
                * model.RampUpMaximum[cl]
                + model.StartUp[cl, h] * model.RampStartUpMaximum[cl]
                - model.PowerMinStable[cl] * model.AF[cl, h] * model.ShutDown[cl, h]
                + model.LL_RampUp[cl, h]
            )
        else:
            return (
                model.Power[cl, h] - model.Power[cl, h - 1]
                <= (model.Committed[cl, h] - model.StartUp[cl, h])
                * model.RampUpMaximum[cl]
                + model.StartUp[cl, h] * model.RampStartUpMaximum[cl]
                - model.PowerMinStable[cl] * model.AF[cl, h] * model.ShutDown[cl, h]
                + model.LL_RampUp[cl, h]
            )

    def EQ_RampDown_TC(model, cl, h):
        if h == 0:
            return (
                model.PowerInitial[cl] - model.Power[cl, h]
                <= (model.Committed[cl, h] - model.StartUp[cl, h])
                * model.RampDownMaximum[cl]
                + model.ShutDown[cl, h] * model.RampShutDownMaximum[cl]
                - model.PowerMinStable[cl] * model.AF[cl, h] * model.StartUp[cl, h]
                + model.LL_RampDown[cl, h]
            )
        else:
            return (
                model.Power[cl, h - 1] - model.Power[cl, h]
                <= (model.Committed[cl, h] - model.StartUp[cl, h])
                * model.RampDownMaximum[cl]
                + model.ShutDown[cl, h] * model.RampShutDownMaximum[cl]
                - model.PowerMinStable[cl] * model.AF[cl, h] * model.StartUp[cl, h]
                + model.LL_RampDown[cl, h]
            )

    def EQ_CostRampUp(model, cl, h):
        if model.CostRampUp[cl] != 0:
            if h == 0:
                return model.CostRampUpH[cl, h] >= model.CostRampUp[cl] * (
                    model.Power[cl, h] - model.PowerInitial[cl]
                )
            else:
                return model.CostRampUpH[cl, h] >= model.CostRampUp[cl] * (
                    model.Power[cl, h] - model.Power[cl, h - 1]
                )
        else:
            return Constraint.Skip

    def EQ_CostRampDown(model, cl, h):
        if model.CostRampDown[cl] != 0:
            if h == 0:
                return model.CostRampDownH[cl, h] >= model.CostRampDown[cl] * (
                    model.PowerInitial[cl] - model.Power[cl, h]
                )
            else:
                return model.CostRampDownH[cl, h] >= model.CostRampDown[cl] * (
                    model.Power[cl, h - 1] - model.Power[cl, h]
                )
        else:
            return Constraint.Skip

    def EQ_Reserve_2U_capability(model, cl, h):
        return (
            model.Reserve_2U[cl, h]
            <= model.PowerCapacity[cl] * model.AF[cl, h] * model.Committed[cl, h]
            - model.Power[cl, h]
        )

    def EQ_Reserve_2D_capability(model, cl, h):
        return (
            model.Reserve_2D[cl, h]
            <= model.Power[cl, h]
            - model.PowerMinStable[cl] * model.AF[cl, h] * model.Committed[cl, h]
            + model.StorageChargingCapacity
            - model.StorageInput[h]
        )

    def EQ_Power_must_run(model, cl, h):
        return (
            model.PowerMinStable[cl] * model.AF[cl, h] * model.Committed[cl, h]
            <= model.Power[cl, h]
        )

    def EQ_Power_available(model, cl, h):
        return (
            model.Power[cl, h]
            <= model.PowerCapacity[cl] * model.AF[cl, h] * model.Committed[cl, h]
        )

    def EQ_Force_Decommitment(model, cl, h):
        if model.AF[cl, h] == 0:
            return model.Committed[cl, h] == 0
        else:
            return Constraint.Skip

    def EQ_Demand_balance_DA(model, h):
        if calibration:
            return (
                sum(model.Power[cl, h] for cl in model.cl)
                + model.RES_Generation[h]
                + model.StorageOutput[h]
                + model.NetImports[h]
                == model.Demand["DA", h]
                + model.StorageInput[h]
                + model.LoadMod_U[h]
                - model.LoadMod_D[h]
                - model.LL_MaxPower[h]
                + model.LL_MinPower[h]
            )
        else:
            return (
                sum(model.Power[cl, h] for cl in model.cl)
                + model.Power_RES[h]
                + model.StorageOutput[h]
                + model.NetImports[h]
                == model.Demand["DA", h]
                + model.StorageInput[h]
                + model.LoadMod_U[h]
                - model.LoadMod_D[h]
                - model.LL_MaxPower[h]
                + model.LL_MinPower[h]
            )

    def EQ_Demand_balance_2U(model, h):
        return (
            sum(
                model.Reserve_2U[cl, h] * model.Technology[cl, t] * model.Reserve[t]
                for cl in model.cl
                for t in model.t
            )
            >= model.Demand["2U", h] - model.LL_2U[h]
        )

    def EQ_Demand_balance_2D(model, h):
        return (
            sum(
                model.Reserve_2D[cl, h] * model.Technology[cl, t] * model.Reserve[t]
                for cl in model.cl
                for t in model.t
            )
            >= model.Demand["2D", h] - model.LL_2D[h]
        )

    def EQ_Storage_minimum(model, h):
        return model.StorageLevel[h] >= model.StorageMinimum

    def EQ_Storage_level(model, h):
        return model.StorageLevel[h] <= model.StorageCapacity

    def EQ_Storage_input(model, h):
        return model.StorageInput[h] <= model.StorageChargingCapacity

    def EQ_Storage_MaxDischarge(model, h):
        return (
            model.StorageOutput[h]
            <= model.StorageDischargeEfficiency * model.StorageLevel[h]
        )

    def EQ_Storage_MaxCharge(model, h):
        return (
            model.StorageInput[h] * model.StorageChargingEfficiency
            <= model.StorageCapacity - model.StorageLevel[h]
        )

    def EQ_Storage_balance(model, h):
        if h == 0:
            return (
                model.StorageInitial
                + model.StorageInput[h] * model.StorageChargingEfficiency
                == model.StorageLevel[h]
                + model.StorageOutput[h] / model.StorageDischargeEfficiency
            )
        else:
            return (
                model.StorageLevel[h - 1]
                + model.StorageInput[h] * model.StorageChargingEfficiency
                == model.StorageLevel[h]
                + model.StorageOutput[h] / model.StorageDischargeEfficiency
            )

    def EQ_Storage_boundaries(model, h):
        if h == len(model.h):
            return model.StorageLevel[h] >= model.StorageFinalMin
        else:
            return Constraint.Skip

    def EQ_CurtailedPower(model, h):
        return (
            model.AF_Solar[h] * model.SolarCapacity
            + model.AF_Wind[h] * model.WindCapacity
            - model.Power_RES[h]
            <= model.CurtailedPower[h]
        )

    # OBJECTIVE
    def EQ_Objective_function(model):
        return sum(model.SystemCost[h] for h in model.h)

    def EQ_SystemCost(model, h):
        if calibration:
            return model.SystemCost[h] == (
                sum(model.NoLoadCost * model.Committed[cl, h] for cl in model.cl)
                + sum(model.CostVariable[cl, h] * model.Power[cl, h] for cl in model.cl)
                + sum(
                    model.CostRampUpH[cl, h] + model.CostRampDownH[cl, h]
                    for cl in model.cl
                )
                + model.ValueOfLostLoad * (model.LL_MaxPower[h] + model.LL_MinPower[h])
                + 0.8 * model.ValueOfLostLoad * (model.LL_2U[h] + model.LL_2D[h])
                + 0.7
                * model.ValueOfLostLoad
                * sum(
                    model.LL_RampUp[cl, h] + model.LL_RampDown[cl, h] for cl in model.cl
                )
            )
        else:
            return model.SystemCost[h] == (
                sum(model.NoLoadCost * model.Committed[cl, h] for cl in model.cl)
                + sum(model.CostVariable[cl, h] * model.Power[cl, h] for cl in model.cl)
                + sum(
                    model.CostRampUpH[cl, h] + model.CostRampDownH[cl, h]
                    for cl in model.cl
                )
                + model.CostCurtailment * model.CurtailedPower[h]
                + model.ValueOfLostLoad * (model.LL_MaxPower[h] + model.LL_MinPower[h])
                + 0.8 * model.ValueOfLostLoad * (model.LL_2U[h] + model.LL_2D[h])
                + 0.7
                * model.ValueOfLostLoad
                * sum(
                    model.LL_RampUp[cl, h] + model.LL_RampDown[cl, h] for cl in model.cl
                )
            )

    ##########################################################################################
    ######################################## Definition of model #############################
    ##########################################################################################

    model.EQ_Objective_Function = Objective(rule=EQ_Objective_function, sense=minimize)

    model.EQ_Commitment = Constraint(model.cl, model.h, rule=EQ_Commitment)
    model.EQ_CostRampUp = Constraint(model.cl, model.h, rule=EQ_CostRampUp)
    model.EQ_CostRampDown = Constraint(model.cl, model.h, rule=EQ_CostRampDown)
    model.EQ_Demand_balance_DA = Constraint(model.h, rule=EQ_Demand_balance_DA)
    model.EQ_MinUpTime = Constraint(model.cl, model.h, rule=EQ_MinUpTime)
    model.EQ_MinDownTime = Constraint(model.cl, model.h, rule=EQ_MinDownTime)
    model.EQ_RampUp_TC = Constraint(model.cl, model.h, rule=EQ_RampUp_TC)
    model.EQ_RampDown_TC = Constraint(model.cl, model.h, rule=EQ_RampDown_TC)
    model.EQ_Reserve_2U_capability = Constraint(
        model.cl, model.h, rule=EQ_Reserve_2U_capability
    )
    model.EQ_Reserve_2D_capability = Constraint(
        model.cl, model.h, rule=EQ_Reserve_2D_capability
    )
    model.EQ_Power_must_run = Constraint(model.cl, model.h, rule=EQ_Power_must_run)
    model.EQ_Power_available = Constraint(model.cl, model.h, rule=EQ_Power_available)
    model.EQ_Force_Decommitment = Constraint(
        model.cl, model.h, rule=EQ_Force_Decommitment
    )
    model.EQ_Demand_balance_2U = Constraint(model.h, rule=EQ_Demand_balance_2U)
    model.EQ_Demand_balance_2D = Constraint(model.h, rule=EQ_Demand_balance_2D)
    model.EQ_Storage_minimum = Constraint(model.h, rule=EQ_Storage_minimum)
    model.EQ_Storage_level = Constraint(model.h, rule=EQ_Storage_level)
    model.EQ_Storage_input = Constraint(model.h, rule=EQ_Storage_input)
    model.EQ_Storage_MaxDischarge = Constraint(model.h, rule=EQ_Storage_MaxDischarge)
    model.EQ_Storage_MaxCharge = Constraint(model.h, rule=EQ_Storage_MaxCharge)
    model.EQ_Storage_balance = Constraint(model.h, rule=EQ_Storage_balance)
    model.EQ_Storage_boundaries = Constraint(model.h, rule=EQ_Storage_boundaries)
    model.EQ_SystemCost = Constraint(model.h, rule=EQ_SystemCost)

    if not calibration:
        model.EQ_CurtailedPower = Constraint(model.h, rule=EQ_CurtailedPower)

    if calibration or (load_shape_mode is None):
        model.EQ_MaxLoadShaping = Constraint(
            rule=lambda model: sum(
                model.LoadMod_U[h] + model.LoadMod_D[h] for h in model.h
            )
            == 0
        )
    elif load_shape_mode.lower() == "savings":
        model.EQ_MaxLoadShaping_up = Constraint(
            rule=lambda model: sum(model.LoadMod_U[h] for h in model.h) == 0
        )
        model.EQ_MaxLoadShaping_down = Constraint(
            rule=lambda model: sum(model.LoadMod_D[h] for h in model.h)
            <= model.MaxLoadShaping
        )
    elif load_shape_mode.lower() == "flex":
        model.EQ_MaxLoadShaping = Constraint(
            rule=lambda model: sum(
                model.LoadMod_U[h] + model.LoadMod_D[h] for h in model.h
            )
            <= model.MaxLoadShaping
        )

    return model

def run_model(instance, solver="glpk"):
    solver = SolverFactory(solver)
    solver.options["OutputFlag"] = False
    results = solver.solve(instance, tee=True)

    if results.solver.termination_condition != TerminationCondition.optimal:
        # something went wrong
        logging.warn("Solver: %s" % results.solver.termination_condition)
        logging.debug(results.solver)
    else:
        logging.info("Solver: %s" % results.solver.termination_condition)

    instance.solutions.load_from(results)
    return instance

max_stars_repo_name: hebes-io/eevalue

max_issues_repo_name: hebes-io/eevalue

max_forks_repo_name: hebes-io/eevalue

reason: import numpy,from pyomo


text:
from __future__ import division, print_function, absolute_import

import numpy as np
from numpy.testing import assert_equal, assert_allclose

from scipy.special._ufuncs import _sinpi as sinpi
from scipy.special._ufuncs import _cospi as cospi

def test_integer_real_part():
    x = np.arange(-100, 101)
    y = np.hstack((-np.linspace(310, -30, 10), np.linspace(-30, 310, 10)))
    x, y = np.meshgrid(x, y)
    z = x + 1j*y
    # In the following we should be *exactly* right
    res = sinpi(z)
    assert_equal(res.real, 0.0)
    res = cospi(z)
    assert_equal(res.imag, 0.0)

def test_half_integer_real_part():
    x = np.arange(-100, 101) + 0.5
    y = np.hstack((-np.linspace(310, -30, 10), np.linspace(-30, 310, 10)))
    x, y = np.meshgrid(x, y)
    z = x + 1j*y
    # In the following we should be *exactly* right
    res = sinpi(z)
    assert_equal(res.imag, 0.0)
    res = cospi(z)
    assert_equal(res.real, 0.0)

def test_intermediate_overlow():
    # Make sure we avoid overflow in situations where cosh/sinh would
    # overflow but the product with sin/cos would not
    sinpi_pts = [complex(1 + 1e-14, 227),
                 complex(1e-35, 250),
                 complex(1e-301, 445)]
    # Data generated with mpmath
    sinpi_std = [complex(-8.113438309924894e+295, -np.inf),
                 complex(1.9507801934611995e+306, np.inf),
                 complex(2.205958493464539e+306, np.inf)]
    for p, std in zip(sinpi_pts, sinpi_std):
        assert_allclose(sinpi(p), std)

    # Test for cosine, less interesting because cos(0) = 1.
    p = complex(0.5 + 1e-14, 227)
    std = complex(-8.113438309924894e+295, -np.inf)
    assert_allclose(cospi(p), std)

max_stars_repo_name: jeikabu/lumberyard

max_issues_repo_name: jeikabu/lumberyard

max_forks_repo_name: jeikabu/lumberyard

reason: import numpy,from numpy,from scipy


text:
#   Copyright 2015 Ufora Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import numpy
import scipy.special
import pyfora.pure_modules.pure_pandas as PurePandas

class BinaryLogisticRegressionModel(object):
    """
    Represents a fit logit model.

    Attributes:
        coefficients (numpy.array): The regressions coefficients.
        intercept (float): The fitted model's intercept

    Note:
        This class is not intended to be constructed directly. Instances of it are
        returned by :func:`~pyfora.algorithms.BinaryLogisticRegressionFitter.fit`.
    """
    def __init__(
            self,
            coefficients,
            classZeroLabel,
            classOneLabel,
            intercept,
            interceptScale,
            iters):
        self.coefficients = coefficients
        self.classZeroLabel = classZeroLabel
        self.classOneLabel = classOneLabel
        self.intercept = intercept
        self.interceptScale = interceptScale
        self.iters = iters

    def score(self, X, y):
        """
        Returns the mean accuracy on the given test data and labels.

        Args:
            X (DataFrame): Feature vectors
            y (DataFrame): Target labels, corresponding to the vectors in ``X``.

        Returns:
            float: The mean accuracy of
            :func:`~pyfora.algorithms.logistic.BinaryLogisticRegressionModel.BinaryLogisticRegressionModel.predict`
            with respect to ``y``.
        """
        assert len(X) == len(y), "arguments must have the same length"

        if isinstance(y, (PurePandas.PurePythonSeries,
                          PurePandas.PurePythonDataFrame)):
            y = y.as_matrix()
            y = y.reshape((len(y),))

        predictions = self.predict(X)

        def valAtIx(ix):
            if y[ix] == predictions[ix]:
                return 1.0
            return 0.0

        return sum(valAtIx(ix) for ix in xrange(len(predictions))) / len(y)

    def predict(self, X):
        """
        Predict the class labels of ``X``.

        Args:
            X (DataFrame, or numpy.array): a set of feature vectors

        Returns:
            numpy.array: array containing the predicted class labels.
        """
        probabilities = self.predict_probability(X)

        def classForProbability(probability):
            if probability > 0.5:
                return self.classOneLabel
            return self.classZeroLabel

        return numpy.array([
            classForProbability(p) for p in probabilities
            ])

    def predict_probability(self, X):
        """
        Estimate the conditional class-zero probability for the features in ``X``.

        Args:
            X (DataFrame, or numpy.array): a set of feature vectors

        Returns:
            numpy.array: array containing the predicted probabilities.
        """
        res = X.dot(self.coefficients)

        if self.intercept is not None:
            res = res + self.intercept * self.interceptScale

        return numpy.array([
            scipy.special.expit(-val) for val in res
            ])

max_stars_repo_name: ufora/ufora

reason: import numpy,import scipy


text:
"""

Measure the percentage of image files where the whale location and heading has
been computed.

"""

from glob import glob
import numpy as np

# directory containing whale images
datadir = '../../BigData/kaggle-right-whale/right_whale_hunt/imgs/'

# get the number of images in the training set
trainlocs = glob(datadir + 'whale_*/w_*jpg')
ntrain = len(trainlocs)

# get the number of processed images in the training set
trainlocs = glob(datadir + 'whale_*/*colorclass*png')
ntrain_processed = len(trainlocs)

# get the number of images in the test set
trainlocs = glob(datadir + 'w_*jpg')
ntest = len(trainlocs)

# get the number of processed images in the test set
trainlocs = glob(datadir + '*colorclass*png')
ntest_processed = len(trainlocs)

print('Training set statistics')
print('{} {} {:.3f}'.format(ntrain, ntrain_processed,
    np.float(ntrain_processed)/ntrain))

print('Test set statistics')
print('{} {} {:.3f}'.format(ntest, ntest_processed,
    np.float(ntest_processed)/ntest))

max_stars_repo_name: sbussmann/kaggle-right-whale

max_issues_repo_name: sbussmann/kaggle-right-whale

max_forks_repo_name: sbussmann/kaggle-right-whale


text:
# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.

"""Tests that expectation values are correctly computed in the plugin device"""

import pennylane as qml
import pytest
from conftest import A
from pennylane import numpy as np

np.random.seed(42)

@pytest.mark.parametrize("shots", [None, 8192])
class TestExpval:
    """Test expectation values"""

    def test_identity_expectation(self, device, shots, tol):
        """Test that identity expectation value (i.e. the trace) is 1"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Identity(i)) for i in range(2)]

        assert np.allclose(circuit(), np.array([1, 1]), **tol)

    def test_pauliz_expectation(self, device, shots, tol):
        """Test that PauliZ expectation value is correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliZ(i)) for i in range(2)]

        assert np.allclose(circuit(), np.array([np.cos(theta), np.cos(theta) * np.cos(phi)]), **tol)

    def test_paulix_expectation(self, device, shots, tol):
        """Test that PauliX expectation value is correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliX(i)) for i in range(2)]

        assert np.allclose(circuit(), np.array([np.sin(theta) * np.sin(phi), np.sin(phi)]), **tol)

    def test_pauliy_expectation(self, device, shots, tol):
        """Test that PauliY expectation value is correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliY(i)) for i in range(2)]

        assert np.allclose(circuit(), np.array([0, -np.cos(theta) * np.sin(phi)]), **tol)

    def test_hadamard_expectation(self, device, shots, tol):
        """Test that Hadamard expectation value is correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Hadamard(i)) for i in range(2)]

        expected = np.array(
            [np.sin(theta) * np.sin(phi) + np.cos(theta), np.cos(theta) * np.cos(phi) + np.sin(phi)]
        ) / np.sqrt(2)
        assert np.allclose(circuit(), expected, **tol)

    def test_hermitian_expectation(self, device, shots, tol):
        """Test that arbitrary Hermitian expectation values are correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123

        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.Hermitian(A, i)) for i in range(2)]

        a = A[0, 0]
        re_b = A[0, 1].real
        d = A[1, 1]
        ev1 = ((a - d) * np.cos(theta) + 2 * re_b * np.sin(theta) * np.sin(phi) + a + d) / 2
        ev2 = ((a - d) * np.cos(theta) * np.cos(phi) + 2 * re_b * np.sin(phi) + a + d) / 2
        expected = np.array([ev1, ev2])

        assert np.allclose(circuit(), expected, **tol)

    def test_multi_mode_hermitian_expectation(self, device, shots, tol):
        """Test that arbitrary multi-mode Hermitian expectation values are correct"""
        dev = device(2)

        theta = 0.432
        phi = 0.123
        A = np.array(
            [
                [-6, 2 + 1j, -3, -5 + 2j],
                [2 - 1j, 0, 2 - 1j, -5 + 4j],
                [-3, 2 + 1j, 0, -4 + 3j],
                [-5 - 2j, -5 - 4j, -4 - 3j, -6],
            ]
        )

        @qml.qnode(dev)
        def circuit():
            qml.RY(theta, wires=[0])
            qml.RY(phi, wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.Hermitian(A, [0, 1]))

        # below is the analytic expectation value for this circuit with arbitrary
        # Hermitian observable A
        expected = 0.5 * (
            6 * np.cos(theta) * np.sin(phi)
            - np.sin(theta) * (8 * np.sin(phi) + 7 * np.cos(phi) + 3)
            - 2 * np.sin(phi)
            - 6 * np.cos(phi)
            - 6
        )
        assert np.allclose(circuit(), expected, **tol)

    def test_nondiff_param(self, device):
        """Test that the device can be successfully executed with the Autograd
        interface when some of the arguments are tensors with requires_grad=False"""

        dev1 = device(1)
        dev2 = qml.device("default.qubit", wires=1)

        if not dev1.analytic:
            pytest.skip("This test is designed to work in analytic mode")

        def circuit(x, y):
            qml.RX(x[0], wires=0)
            qml.Rot(*x[1:], wires=0)
            qml.RY(y[0], wires=0)
            return qml.expval(qml.PauliZ(0))

        x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False)
        y = np.array([0.5], requires_grad=True)

        qnode1 = qml.QNode(circuit, dev1, interface="autograd")
        qnode2 = qml.QNode(circuit, dev2, interface="autograd")

        r1 = qnode1(x, y)
        r2 = qnode2(x, y)

        assert np.allclose(r1, r2)

@pytest.mark.parametrize("shots", [None, 8192])
class TestTensorExpval:
    """Test tensor expectation values"""

    def test_paulix_pauliy(self, device, shots, tol):
        """Test that a tensor product involving PauliX and PauliY works correctly"""
        dev = device(3)

        theta = 0.432
        phi = 0.123
        varphi = -0.543

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.expval(qml.PauliX(0) @ qml.PauliY(2))

        assert np.allclose(circuit(), np.sin(theta) * np.sin(phi) * np.sin(varphi), **tol)

    def test_pauliz_hadamard(self, device, shots, tol):
        """Test that a tensor product involving PauliZ and PauliY and Hadamard works correctly"""
        dev = device(3)

        theta = 0.432
        phi = 0.123
        varphi = -0.543

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.expval(
                qml.PauliZ(wires=[0]) @ qml.Hadamard(wires=[1]) @ qml.PauliY(wires=[2])
            )

        expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2)
        assert np.allclose(circuit(), expected, **tol)

    def test_hermitian(self, device, shots, tol):
        """Test that a tensor product involving qml.Hermitian works correctly"""
        dev = device(3)

        theta = 0.432
        phi = 0.123
        varphi = -0.543
        A = np.array(
            [
                [-6, 2 + 1j, -3, -5 + 2j],
                [2 - 1j, 0, 2 - 1j, -5 + 4j],
                [-3, 2 + 1j, 0, -4 + 3j],
                [-5 - 2j, -5 - 4j, -4 - 3j, -6],
            ]
        )

        @qml.qnode(dev)
        def circuit():
            qml.RX(theta, wires=[0])
            qml.RX(phi, wires=[1])
            qml.RX(varphi, wires=[2])
            qml.CNOT(wires=[0, 1])
            qml.CNOT(wires=[1, 2])
            return qml.expval(qml.PauliZ(wires=[0]) @ qml.Hermitian(A, wires=[1, 2]))

        expected = 0.5 * (
            -6 * np.cos(theta) * (np.cos(varphi) + 1)
            - 2 * np.sin(varphi) * (np.cos(theta) + np.sin(phi) - 2 * np.cos(phi))
            + 3 * np.cos(varphi) * np.sin(phi)
            + np.sin(phi)
        )
        assert np.allclose(circuit(), expected, **tol)