text:
"""Generate C code and data for the augmented Lagrangian method."""

import numpy as np

from muaompc._ldt.codegen.codegen import BaseCodeGenerator as BCG
from muaompc._ldt.codegen.former.cvp.codegen import CVPDataGenerator as CVPDG
from muaompc._ldt.codegen.former.cvp.codegen import CDataGenerator as CVPCDG
from muaompc._ldt.codegen.solver.fgm.codegen import CCodeGenerator as FGMCG
from muaompc._ldt.codegen.solver.fgm.codegen import (FGMCVPDataGenerator
                                                     as FGMCVPDG)
from muaompc._ldt.codegen.solver.fgm.codegen import CDataGenerator as FGMCVPCDG

class CCodeGenerator(BCG, object):
    """Generate C code for the augmented Lagrangian method."""

    def __init__(self, ccg):
        """Assign the argument to self, and set the C code paths.

        This depends on common C-code already being generated.

        :param ccg: constains the prefix and the destination directory of the
        generated common C-code.
        :type ccg: muaompc._ldt.codegen.common.codegen.CCodeGenerator
        """
        self.prefix = ccg.prefix
        self.destdir = ccg.destdir
        self.prbname = ccg.prbname
        self.path = self._get_paths('codegen.solver.alm')
        self.fgm = FGMCG(ccg)

    def generate_code(self):
        self.fgm.generate_code()
        self._make_destination_dir_tree()
        self._generate_c_code()

    def _generate_c_code(self):
        self._generate_c_header('alm')
        self._generate_c_body('alm')
        self._generate_getdata_h()
        self._generate_getdata_c()

    def _generate_c_header(self, fname):
        fname += '.h'
        self._replace_prefix(fname)

    def _generate_c_body(self, fname):
        fname += '.c'
        self._replace_prefix(fname)

    def _generate_getdata_h(self):
        fname = 'almdynmem.h'
        self._replace_prefix(fname)

    def _generate_getdata_c(self):
        fname = 'almdynmem.c'
        self._replace_prefix(fname)

class ALMCVPDataGenerator(CVPDG, object):
    """
    Generate json data for the augmented Lagrangian method using a
    condensed vectorial parameters problem formulation.
    """

    def __init__(self, cvpcg):
        """Assign the arguments to self.

        :param cvp: an instance of condensed vectorial parameters formulation
        :type cvp:
        muaompc._ldt.codegen.former.cvp.cvp.CondensedVectorialParameters

        """
        CVPDG.__init__(self, cvpcg)
        self.fgm = FGMCVPDG(cvpcg)

    def _get_data(self, num):
        data = CVPDG._get_data(self, num)
        H, V = self._get_matrices_HV(num)
        try:
            num['mu']
        except KeyError:
            data['mu'] = self._find_penalty_parameter(H, V)
        else:
            data['mu'] = num['mu']
        nu, L = self._compute_extra_step_constant(H, data['mu'], V)
        data_scaled = self.fgm._scale_cost(data, H.flatten(), L)
        data_scaled['nu'] = nu
        data_scaled['Linv'] = 1. / L
        return data_scaled

    def _get_matrices_HV(self, num):
        data = CVPDG._get_data(self, num)
        Ht = data['pmetric']['H']['fac0']
        H = np.reshape(np.array(Ht['data']), (Ht['rows'], Ht['cols']))
        Vt = data['pmetric']['V']['fac0']
        V = np.reshape(np.array(Vt['data']), (Vt['rows'], Vt['cols']))
        return H, V

    def _compute_extra_step_constant(self, H, mu, V):
        from math import sqrt
        L = max(np.linalg.eigvalsh(H + mu * np.dot(V.T, V)))
        phi = min(np.linalg.eigvalsh(H))
        nu = ((sqrt(L) - sqrt(phi)) / (sqrt(L) + sqrt(phi)))
        return (nu, L)

    def _find_penalty_parameter(self, H, V):
        """Returns a penalty parameter that makes the condition number of the
        internal problem (FGM) not much larger than the Hessian's
        condition number.
        """
        msg = '_find_penalty_parameter '

        lower_limit, upper_limit, max_iter = (2, 10, 100)
        mu = 1.
        cnH = self._compute_alg_condnum(H, 0., V)
        cn_fact = self._compute_alg_condnum(H, mu, V) / cnH

        for k in range(max_iter):
            if cn_fact > upper_limit:
                mu /= 2
            elif cn_fact < lower_limit:
                mu *= 2
            else:
                break
            cn_fact = self._compute_alg_condnum(H, mu, V) / cnH
            if k == max_iter - 1:
                msg += 'WARNING: reached maximum number of iterations.'

        return mu

    def _compute_alg_condnum(self, H, mu, V):
        dummy, L = self._compute_extra_step_constant(H, mu, V)
        return L / (min(np.linalg.eigvals(H)))

class CDataGenerator(CCodeGenerator, ALMCVPDataGenerator, object):

    def __init__(self, ccg, cvpcg):
        """Assign the arguments to self.

        :param ccg: an instance of C code generator for the ALM
        :type ccg: muaompc._ldt.codegen.solver.alm.codegen.CCodeGenerator

        :param cvpcg: an instance of C code generator of a
        condensed vectorial parameters formulation
        :type cvpcg: muaompc._ldt.codegen.former.cvp.codegen.CCodeGenerator

        """
        self.cvp = cvpcg.cvp
        self.smb = cvpcg.cvp.smb
        self.prbname = ccg.prbname
        self.prefix = cvpcg.prefix
        self.destdir = ccg.destdir
        self.path = self._get_paths('codegen.solver.alm')
        self.cvpcg = cvpcg
        self.cvpcdg = CVPCDG(cvpcg)
        self.fgmcvpcdg = FGMCVPCDG(ccg, cvpcg)
        self.dict_ = None  # Assignment made during data generation

    def generate_data(self, num, dataname):
        self.dict_ = dict(prefix=self.prefix,
                          PREFIX=self.prefix.upper(),
                          dataname=dataname,
                          DATANAME=dataname.upper())
        self.cvpcdg.dict_ = self.dict_
        self.fgmcvpcdg.dict_ = self.dict_

        almcvpdg = ALMCVPDataGenerator(self.cvpcg)
        data = almcvpdg._get_data(num)
        self.cvpcdg._write_static_data(data, dataname)
        self.fgmcvpcdg._write_static_data(data, dataname)
        self._write_static_data(data, dataname)

    def _write_static_data(self, data, dataname):
        self._generate_h_struct(data, dataname)
        self._generate_c_struct(data, dataname)

    def _generate_h_struct(self, data, dataname):
        fname = 'almdata'
        fname += '.h'
        prefix = self._get_full_fname_prefix(dataname)
        tmp = self.dict_
        self._replace_dict(tmp, prefix, fname)

    def _generate_c_struct(self, data, dataname):
        fname = 'almdata'
        fname += '.c'
        prefix = self._get_full_fname_prefix(dataname)
        zeros_optvar_seqlen = self._np2Carray(
            np.zeros(data['optvar']['seqlen']))
        zeros_lagmul_seqlen = self._np2Carray(
            np.zeros(data['lagmul']['seqlen']))
        zeros_lagmul_optvar_seqlen = self._np2Carray(np.zeros(
            data['lagmul']['seqlen']*data['optvar']['seqlen']))

        tmp = dict(mu=data['mu'],
                   Linv=data['Linv'],
                   zeros_optvar_seqlen=zeros_optvar_seqlen,
                   zeros_lagmul_seqlen=zeros_lagmul_seqlen,
                   zeros_lagmul_optvar_seqlen=zeros_lagmul_optvar_seqlen,
                   lagmul_veclen=data['lagmul']['veclen'])
        tmp.update(self.dict_)
        self._replace_dict(tmp, prefix, fname)

max_stars_repo_name: muaompc/muaompc

max_issues_repo_name: muaompc/muaompc

max_forks_repo_name: muaompc/muaompc


text:
from typing import Optional

import numpy as np

from src.biota_models.vegetation.bio_process.veg_hydro_morphodynamics import (
    Hydro_Morphodynamics,
)
from src.biota_models.vegetation.model.veg_constants import VegetationConstants
from src.biota_models.vegetation.model.veg_lifestages import LifeStages
from src.biota_models.vegetation.model.veg_model import Vegetation
from src.core.base_model import ExtraModel

class Veg_Mortality(ExtraModel):
    """Mortality"""

    def __init__(self):
        super().__init__()
        self.burial_scour_j = None
        self.burial_scour_m = None
        self.scour_j = None
        self.burial_j = None
        self.scour_m = None
        self.burial_m = None
        self.Bl_diff = None
        self.fraction_dead_flood_j = None
        self.fraction_dead_des_j = None
        self.fraction_dead_upr_j = None
        self.fraction_dead_flood_m = None
        self.fraction_dead_des_m = None
        self.fraction_dead_upr_m = None

    # burial_scour: Optional[np.array] = None
    # scour: Optional[np.array] = None
    # burial: Optional[np.array] = None
    # BL_diff: Optional[np.array] = None
    # fraction_dead_flood: Optional[np.array] = None
    # fraction_dead_des: Optional[np.array] = None
    # fraction_dead_upr: Optional[np.array] = None

    def update(
        self,
        veg: Vegetation,
        constants: VegetationConstants,
        ets,
        begin_date,
        end_date,
        period,
    ):
        """Update vegetation characteristics after mortality"""
        Veg_Mortality.drowning_hydroperiod(self, veg, constants, ets)
        Veg_Mortality.uprooting(self, veg, constants)
        Veg_Mortality.erosion_sedimentation(self, veg, ets)

        veg.juvenile.veg_frac = (
            veg.juvenile.veg_frac
            - self.fraction_dead_flood_j
            - self.fraction_dead_des_j
            - self.fraction_dead_upr_j
            - self.burial_scour_j
        )  # update fractions due to mortality
        veg.juvenile.veg_frac[
            veg.juvenile.veg_frac < 0
        ] = 0  # replace negative values with 0
        veg.mature.veg_frac = (
            veg.mature.veg_frac
            - self.fraction_dead_flood_m
            - self.fraction_dead_des_m
            - self.fraction_dead_upr_m
            - self.burial_scour_m
        )  # update fractions due to mortality
        veg.mature.veg_frac[
            veg.mature.veg_frac < 0
        ] = 0  # replace negative values with 0

        veg.juvenile.update_growth(veg.juvenile.veg_frac, period, begin_date, end_date)
        veg.mature.update_growth(veg.mature.veg_frac, period, begin_date, end_date)

    def drowning_hydroperiod(
        self, veg: Vegetation, constants: VegetationConstants, ets
    ):
        flooding_current, drying_current = self.compute_hydroperiod(
            veg.wl_ts, constants
        )

        wet = np.zeros(flooding_current.shape)
        dry = np.zeros(drying_current.shape)
        wet[flooding_current > 0] = 1
        dry[drying_current > 0] = 1
        if ets == 0:
            veg.wl_prev = np.zeros(veg.wl_ts.shape)
        flooding_prev, drying_prev = self.compute_hydroperiod(veg.wl_prev, constants)

        dry = (
            drying_prev * dry
        )  # deleting all cells that have fallen wet during this ETS
        wet = (
            flooding_prev * wet
        )  # deleting all cells that have fallen dry during this ETS

        new_wet = np.where(
            (flooding_prev == 0) & (flooding_current > 0)
        )  # find cells that are newly wet during this ETS
        new_dry = np.where(
            (drying_prev == 0) & (drying_current > 0)
        )  # find cells that are newly dry during this ETS

        wet_j = np.ones(veg.juvenile.veg_frac.shape) * wet.reshape(len(wet), 1)
        # np.repeat(wet.reshape(len(wet), 1), len(veg.juvenile.veg_frac[0]), axis=1)
        dry_j = np.ones(veg.juvenile.veg_frac.shape) * dry.reshape(len(dry), 1)
        wet_j[new_wet] = veg.juvenile.veg_frac[
            new_wet
        ]  # add initial fractions in cells that are newly wet in matrix
        dry_j[new_dry] = veg.juvenile.veg_frac[
            new_dry
        ]  # add initial fractions in cells that are newly dry in matrix
        # determine flooding/drying mortalities based on linear relationship
        mort_flood_j = self.mortality_flood_frequency(
            flooding_current, constants.floMort_thres[0], constants.floMort_slope[0]
        )
        mort_des_j = self.mortality_flood_frequency(
            drying_current, constants.desMort_thres[0], constants.desMort_slope[0]
        )
        self.fraction_dead_flood_j = wet_j * mort_flood_j
        self.fraction_dead_des_j = dry_j * mort_des_j

        wet_m = np.ones(veg.mature.veg_frac.shape) * wet.reshape(len(wet), 1)
        dry_m = np.ones(veg.mature.veg_frac.shape) * dry.reshape(len(dry), 1)
        wet_m[new_wet] = veg.mature.veg_frac[
            new_wet
        ]  # add initial fractions in cells that are newly wet in matrix
        dry_m[new_dry] = veg.mature.veg_frac[
            new_dry
        ]  # add initial fractions in cells that are newly dry in matrix
        # determine flooding/drying mortalities based on linear relationship
        mort_flood_m = self.mortality_flood_frequency(
            flooding_current, constants.floMort_thres[1], constants.floMort_slope[1]
        )
        mort_des_m = self.mortality_flood_frequency(
            drying_current, constants.desMort_thres[1], constants.desMort_slope[1]
        )
        self.fraction_dead_flood_m = wet_m * mort_flood_m
        self.fraction_dead_des_m = dry_m * mort_des_m

    @staticmethod
    def compute_hydroperiod(wl_time, constants: VegetationConstants):
        # determiine cells with water depth > flooding/drying threshold
        fl = np.where(wl_time > constants.fl_dr)
        flood = np.zeros(wl_time.shape)
        flood[fl] = 1
        flood = flood.sum(axis=1)  # sum up for all time steps in the ets

        # compute average flooding and drying period
        flooding_current = flood / constants.ets_duration
        drying_current = (constants.ets_duration - flood) / constants.ets_duration

        return flooding_current, drying_current

    @staticmethod
    def mortality_flood_frequency(fl, th, sl):
        """ "
        calculate mortality fraction due to flooding/drying and velocity
        Function: f(x)=sl*x+b for y= 0: b = -sl*th
        th: Threshold
        sl: slope of function
        fl: matrix with flooded days
        """
        b = -th * sl
        dmax = round((1 - b) / sl, 2)  # no. of days when 100% is died off
        fct = sl * fl + b  # determines all mortality values over the grid

        out_fl = np.zeros(fl.shape)
        B = np.where(fl > dmax)  # cells with 100% mortality
        out_fl[B] = 1
        a = fl > th
        b = fl < dmax
        C = np.nonzero(
            (a == True) & (b == True)
        )  # cells where fct applies to determine mortality
        out_fl[C] = fct[C]
        return out_fl.reshape(len(out_fl), 1)

    def uprooting(self, veg: Vegetation, constants: VegetationConstants):
        """
        Mortality through velocity is determined by lin. Function and multiplied with current fraction.
        """
        mort_flow_j = self.mortality_flood_frequency(
            veg.max_u, constants.vel_thres[0], constants.vel_slope[0]
        )
        self.fraction_dead_upr_j = mort_flow_j * veg.juvenile.veg_frac

        mort_flow_m = self.mortality_flood_frequency(
            veg.max_u, constants.vel_thres[1], constants.vel_slope[1]
        )
        self.fraction_dead_upr_m = mort_flow_m * veg.mature.veg_frac

    def erosion_sedimentation(self, veg: Vegetation, ets):
        """
        For burial/erosion the length of stem/root is compared with sedimentation/erosion. In case of mortality
        the fraction is set to 0. We assume that the plants adapt to the deposition and erosion rates within one
        ETS by resetting them each ETS.
        """
        Veg_Mortality.BedLevel_Dif(self, veg, ets)
        fract_scour_j = np.zeros(veg.juvenile.root_len.shape)
        fract_burial_j = np.zeros(veg.juvenile.veg_height.shape)
        fract_scour_j[
            self.scour_j > veg.juvenile.root_len
        ] = 1  # find cells with mortality (scour > rootlength)
        fract_burial_j[self.burial_j > veg.juvenile.veg_height] = 1

        self.burial_scour_j = (
            fract_scour_j + fract_burial_j
        )  # array with cells where vegetation dies and the fraction of death due to burial and scour

        fract_scour_m = np.zeros(veg.mature.root_len.shape)
        fract_burial_m = np.zeros(veg.mature.veg_height.shape)
        fract_scour_m[
            self.scour_m > veg.mature.root_len
        ] = 1  # find cells with mortality (scour > rootlength)
        fract_burial_m[self.burial_m > veg.mature.veg_height] = 1

        self.burial_scour_m = fract_scour_m + fract_burial_m

    ## TODO make this static method?
    def BedLevel_Dif(self, veg: Vegetation, ets):
        if ets >= 1:  # from second time step onward in each ets
            if ets == 1:
                self.Bl_diff = veg.bl - veg.bl_prev
            else:
                depth_dts = veg.bl - veg.bl_prev
                self.Bl_diff = depth_dts + self.Bl_diff

            self.burial_j = np.zeros(veg.juvenile.veg_height.shape)
            self.scour_j = np.zeros(veg.juvenile.root_len.shape)
            self.burial_m = np.zeros(veg.mature.veg_height.shape)
            self.scour_m = np.zeros(veg.mature.root_len.shape)

            loc_b = np.where(self.Bl_diff < 0)
            self.burial_j[loc_b] = (
                np.ones(self.burial_j.shape)
                * self.Bl_diff[0 : len(self.burial_j)].reshape(
                    len(self.Bl_diff[0 : len(self.burial_j)]), 1
                )
            )[loc_b]
            self.burial_m[loc_b] = (
                np.ones(self.burial_m.shape)
                * self.Bl_diff[0 : len(self.burial_m)].reshape(
                    len(self.Bl_diff[0 : len(self.burial_m)]), 1
                )
            )[loc_b]

            loc_s = np.where(self.Bl_diff > 0)
            self.scour_j[loc_s] = (
                np.ones(self.scour_j.shape)
                * self.Bl_diff[0 : len(self.scour_j)].reshape(
                    len(self.Bl_diff[0 : len(self.scour_j)]), 1
                )
            )[loc_s]
            self.scour_m[loc_s] = (
                np.ones(self.scour_m.shape)
                * self.Bl_diff[0 : len(self.scour_m)].reshape(
                    len(self.Bl_diff[0 : len(self.scour_m)]), 1
                )
            )[loc_s]

        else:
            self.burial_j = np.zeros(veg.juvenile.veg_height.shape)
            self.scour_j = np.zeros(veg.juvenile.root_len.shape)
            self.burial_m = np.zeros(veg.mature.veg_height.shape)
            self.scour_m = np.zeros(veg.mature.root_len.shape)

ext: py

max_stars_repo_name: Deltares/NBSDynamics

max_issues_repo_name: Deltares/NBSDynamics

max_forks_repo_name: Deltares/NBSDynamics

reason: import numpy


text:
#!/usr/bin/env python
##################################################
# Gnuradio Python Flow Graph
# Title: Channel Model
# Generated: Thu Dec  5 19:01:14 2013
##################################################

execfile("./custom_ber.py")
from gnuradio import blocks
from gnuradio import channels
from gnuradio import digital
from gnuradio import eng_notation
from gnuradio import gr
from gnuradio.eng_option import eng_option
from gnuradio.filter import firdes
from optparse import OptionParser
import numpy

class channel_model(gr.top_block):

    def __init__(self, constellation=[-1, 1]):
        gr.top_block.__init__(self, "Channel Model")

        ##################################################
        # Parameters
        ##################################################
        self.constellation = constellation

        ##################################################
        # Variables
        ##################################################
        self.const_points = const_points = constellation
        self.taps = taps = [1.0, 0.25-0.25j, 0.50 + 0.10j, 0.3 + 0.2j]
        self.samp_rate = samp_rate = 8e6
        self.noise_level = noise_level = 0.3
        self.frequency = frequency = 2000000
        self.const_type = const_type = 1
        self.const_dist = const_dist = digital.constellation_calcdist(const_points,[],0,1)
        self.const = const = digital.constellation_8psk().bits_per_symbol()

        ##################################################
        # Blocks
        ##################################################
        self.digital_constellation_decoder_cb_0_0 = digital.constellation_decoder_cb(const_dist.base())
        self.digital_constellation_decoder_cb_0 = digital.constellation_decoder_cb(const_dist.base())
        self.digital_chunks_to_symbols_xx_0 = digital.chunks_to_symbols_bc((const_dist.points()), 1)
        self.custom_ber_0 = custom_ber(
            n_bits=10000,
        )
        self.channels_channel_model_0_0 = channels.channel_model(
        	noise_voltage=0,
        	frequency_offset=0,
        	epsilon=1,
        	taps=([1,0,0,0]),
        	noise_seed=0,
        	block_tags=False
        )
        self.channels_channel_model_0 = channels.channel_model(
        	noise_voltage=noise_level,
        	frequency_offset=0,
        	epsilon=1,
        	taps=(taps),
        	noise_seed=0,
        	block_tags=False
        )
        self.blocks_vector_sink_x_1 = blocks.vector_sink_f(1)
        self.analog_random_source_x_0 = blocks.vector_source_b(map(int, numpy.random.randint(0, len(const_points), 10000)), False)

        ##################################################
        # Connections
        ##################################################
        self.connect((self.custom_ber_0, 0), (self.blocks_vector_sink_x_1, 0))
        self.connect((self.channels_channel_model_0, 0), (self.digital_constellation_decoder_cb_0, 0))
        self.connect((self.digital_chunks_to_symbols_xx_0, 0), (self.channels_channel_model_0, 0))
        self.connect((self.analog_random_source_x_0, 0), (self.digital_chunks_to_symbols_xx_0, 0))
        self.connect((self.digital_constellation_decoder_cb_0_0, 0), (self.custom_ber_0, 0))
        self.connect((self.digital_constellation_decoder_cb_0, 0), (self.custom_ber_0, 1))
        self.connect((self.digital_chunks_to_symbols_xx_0, 0), (self.channels_channel_model_0_0, 0))
        self.connect((self.channels_channel_model_0_0, 0), (self.digital_constellation_decoder_cb_0_0, 0))

# QT sink close method reimplementation

    def get_constellation(self):
        return self.constellation

    def set_constellation(self, constellation):
        self.constellation = constellation
        self.set_const_points(self.constellation)

    def get_const_points(self):
        return self.const_points

    def set_const_points(self, const_points):
        self.const_points = const_points
        self.set_const_dist(digital.constellation_calcdist(self.const_points,[],0,1))

    def get_taps(self):
        return self.taps

    def set_taps(self, taps):
        self.taps = taps
        self.channels_channel_model_0.set_taps((self.taps))

    def get_samp_rate(self):
        return self.samp_rate

    def set_samp_rate(self, samp_rate):
        self.samp_rate = samp_rate

    def get_noise_level(self):
        return self.noise_level

    def set_noise_level(self, noise_level):
        self.noise_level = noise_level
        self.channels_channel_model_0.set_noise_voltage(self.noise_level)

    def get_frequency(self):
        return self.frequency

    def set_frequency(self, frequency):
        self.frequency = frequency

    def get_const_type(self):
        return self.const_type

    def set_const_type(self, const_type):
        self.const_type = const_type

    def get_const_dist(self):
        return self.const_dist

    def set_const_dist(self, const_dist):
        self.const_dist = const_dist

    def get_const(self):
        return self.const

    def set_const(self, const):
        self.const = const

if __name__ == '__main__':
    parser = OptionParser(option_class=eng_option, usage="%prog: [options]")
    (options, args) = parser.parse_args()
    tb = channel_model()
    tb.start()
    tb.wait()

max_stars_repo_name: dkozel/starry-skies

max_issues_repo_name: dkozel/starry-skies

max_forks_repo_name: dkozel/starry-skies


text:
import os
import numpy as np
import unittest
from astropy import units as u
from astropy.coordinates import SkyCoord

import MulensModel as mm

dir_1 = os.path.join(mm.DATA_PATH, 'photometry_files', 'OB140939')
dir_2 = os.path.join(mm.DATA_PATH, 'unit_test_files')
dir_3 = os.path.join(mm.DATA_PATH, 'ephemeris_files')

SAMPLE_FILE_02 = os.path.join(dir_1, 'ob140939_OGLE.dat')  # HJD'
SAMPLE_FILE_02_REF = os.path.join(dir_2, 'ob140939_OGLE_ref_v1.dat')  # HJD'
SAMPLE_FILE_03 = os.path.join(dir_1, 'ob140939_Spitzer.dat')  # HJD'
SAMPLE_FILE_03_EPH = os.path.join(dir_3, 'Spitzer_ephemeris_01.dat')  # UTC
SAMPLE_FILE_03_REF = os.path.join(dir_2, 'ob140939_Spitzer_ref_v1.dat')  # HJD'

# All in HJD':
SAMPLE_ANNUAL_PARALLAX_FILE_01 = os.path.join(dir_2, 'parallax_test_1.dat')
SAMPLE_ANNUAL_PARALLAX_FILE_02 = os.path.join(dir_2, 'parallax_test_2.dat')
SAMPLE_ANNUAL_PARALLAX_FILE_03 = os.path.join(dir_2, 'parallax_test_3.dat')
SAMPLE_ANNUAL_PARALLAX_FILE_04 = os.path.join(dir_2, 'parallax_test_4.dat')
SAMPLE_ANNUAL_PARALLAX_FILE_05 = os.path.join(dir_2, 'parallax_test_5.dat')

class _ParallaxFile(object):
    """
    Private class to allow easy access to the contents of the parallax
    test files.
    """
    def __init__(self, filename):
        """
        Open the file and store parameters.
        """
        self.filename = filename
        self.data = np.genfromtxt(
                self.filename, dtype=None,
                names=['Time', 'Magnification', 'PLflux', 'u', 'qn', 'qe'])

        (self.ulens_params, self.event_params) = self.get_file_params()

    def get_file_params(self):
        """Read in the model parameters used to create the file"""
        with open(self.filename) as data_file:
            lines = data_file.readlines()
        ulens_params = lines[3].split()
        event_params = lines[4].split()
        return (ulens_params, event_params)

    @property
    def parameters(self):
        """Model parameters"""
        model_parameters = mm.ModelParameters({
            't_0': float(self.ulens_params[1])+2450000.,
            'u_0': float(self.ulens_params[3]),
            't_E': float(self.ulens_params[4]),
            'pi_E_N': float(self.ulens_params[5]),
            'pi_E_E': float(self.ulens_params[6]),
            't_0_par': self.t_0_par})
        return model_parameters

    @property
    def coords(self):
        """Coordinates of event"""
        coords = SkyCoord(
            self.event_params[1]+' '+self.event_params[2],
            unit=(u.deg, u.deg))
        return coords

    @property
    def t_0_par(self):
        """Parallax reference time"""
        return float(self.ulens_params[2])+2450000.

    def setup_model(self):
        """Return a model using the parameters of this file"""
        model = mm.Model(parameters=self.parameters,
                         coords=self.coords)
        return model

    def setup_trajectory(self):
        """Return a trajectory using the parameters of this file"""
        trajectory = mm.Trajectory(
            self.data['Time']+2450000., parameters=self.parameters,
            parallax={'earth_orbital': True}, coords=self.coords)
        return trajectory

def test_annual_parallax_calculation():
    """
    This is a high-level unit test for parallax. The "true" values were
    calculated from the sfit routine assuming fs=1.0, fb=0.0.
    """
    t_0 = 2457479.5  # April 1 2016, a time when parallax is large
    times = np.array([t_0-1., t_0, t_0+1., t_0+1.])
    true_with_par = [
        np.array([7.12376832, 10.0386009, 7.13323363, 7.13323363])]

    model_with_par = mm.Model(
        {'t_0': t_0, 'u_0': 0.1, 't_E': 10., 'pi_E': (0.3, 0.5)},
        coords='17:57:05 -30:22:59')
    model_with_par.parallax(satellite=False, earth_orbital=True,
                            topocentric=False)
    ones = np.ones(len(times))
    data = mm.MulensData(data_list=[times, ones, ones])
    model_with_par.set_datasets([data])

    model_with_par.parameters.t_0_par = 2457479.

    np.testing.assert_almost_equal(
        model_with_par.data_magnification, true_with_par, decimal=4)

class test(unittest.TestCase):
    def test_wrong_settings(self):
        """
        make sure that if pi_E is defined, then at least
        one component of parallax is True
        """
        model_no_par = mm.Model(
            {'t_0': 2457479.5, 'u_0': 0.1, 't_E': 10., 'pi_E': (0.3, 0.5)},
            coords='17:57:05 -30:22:59')
        model_no_par.parallax(
            satellite=False, earth_orbital=False, topocentric=False)
        with self.assertRaises(ValueError):
            model_no_par.magnification(2457500.)

def do_get_delta_annual_test(filename):
    """run a test on private method Trajectory._get_delta_annual()"""
    parallax_file = _ParallaxFile(filename)
    trajectory = parallax_file.setup_trajectory()

    result = trajectory._get_delta_annual()

    np.testing.assert_almost_equal(result['N'], parallax_file.data['qn'],
                                   decimal=4)
    np.testing.assert_almost_equal(result['E'], parallax_file.data['qe'],
                                   decimal=4)

def test_get_delta_annual_1():
    """test private method Trajectory._get_delta_annual()"""
    do_get_delta_annual_test(SAMPLE_ANNUAL_PARALLAX_FILE_01)

def test_get_delta_annual_2():
    """test private method Trajectory._get_delta_annual()"""
    do_get_delta_annual_test(SAMPLE_ANNUAL_PARALLAX_FILE_02)

def test_get_delta_annual_3():
    """test private method Trajectory._get_delta_annual()"""
    do_get_delta_annual_test(SAMPLE_ANNUAL_PARALLAX_FILE_03)

def test_get_delta_annual_4():
    """test private method Trajectory._get_delta_annual()"""
    do_get_delta_annual_test(SAMPLE_ANNUAL_PARALLAX_FILE_04)

def test_get_delta_annual_5():
    """test private method Trajectory._get_delta_annual()"""
    do_get_delta_annual_test(SAMPLE_ANNUAL_PARALLAX_FILE_05)

def do_annual_parallax_test(filename):
    """testing functions called by a few unit tests"""
    with open(filename) as data_file:
        lines = data_file.readlines()
    ulens_params = lines[3].split()
    event_params = lines[4].split()
    data = np.loadtxt(filename, dtype=None)
    model = mm.Model({
        't_0': float(ulens_params[1])+2450000.,
        'u_0': float(ulens_params[3]),
        't_E': float(ulens_params[4]),
        'pi_E_N': float(ulens_params[5]),
        'pi_E_E': float(ulens_params[6])},
        coords=SkyCoord(
            event_params[1]+' '+event_params[2], unit=(u.deg, u.deg)))
    model.parameters.t_0_par = float(ulens_params[2])+2450000.

    time = data[:, 0]
    dataset = mm.MulensData([time, 20.+time*0., 0.1+time*0.], add_2450000=True)
    model.set_datasets([dataset])
    model.parallax(satellite=False, earth_orbital=True, topocentric=False)
    return np.testing.assert_almost_equal(
        model.data_magnification[0] / data[:, 1], 1.0, decimal=4)

def test_annual_parallax_calculation_2():
    do_annual_parallax_test(SAMPLE_ANNUAL_PARALLAX_FILE_01)

def test_annual_parallax_calculation_3():
    do_annual_parallax_test(SAMPLE_ANNUAL_PARALLAX_FILE_02)

def test_annual_parallax_calculation_4():
    do_annual_parallax_test(SAMPLE_ANNUAL_PARALLAX_FILE_03)

def test_annual_parallax_calculation_5():
    do_annual_parallax_test(SAMPLE_ANNUAL_PARALLAX_FILE_04)

def test_annual_parallax_calculation_6():
    do_annual_parallax_test(SAMPLE_ANNUAL_PARALLAX_FILE_05)

def test_satellite_and_annual_parallax_calculation():
    """test parallax calculation with Spitzer data"""
    model_with_par = mm.Model({'t_0': 2456836.22, 'u_0': 0.922, 't_E': 22.87,
                               'pi_E_N': -0.248, 'pi_E_E': 0.234},
                              coords="17:47:12.25 -21:22:58.2")
    model_with_par.parallax(satellite=True, earth_orbital=True,
                            topocentric=False)
    model_with_par.parameters.t_0_par = 2456836.2

    data_OGLE = mm.MulensData(file_name=SAMPLE_FILE_02)
    data_Spitzer = mm.MulensData(
        file_name=SAMPLE_FILE_03, ephemerides_file=SAMPLE_FILE_03_EPH)
    model_with_par.set_datasets([data_OGLE, data_Spitzer])

    ref_OGLE = np.loadtxt(SAMPLE_FILE_02_REF, unpack=True, usecols=[5])
    ref_Spitzer = np.loadtxt(SAMPLE_FILE_03_REF, unpack=True, usecols=[5])

    np.testing.assert_almost_equal(model_with_par.data_magnification[0],
                                   ref_OGLE, decimal=2)
    ratio = model_with_par.data_magnification[1] / ref_Spitzer
    np.testing.assert_almost_equal(ratio, [1.]*len(ratio), decimal=4)

def test_satellite_parallax_magnification():
    """
    On a given date, the magnification should be different from the
    ground and from Spitzer. Use OB140939 as a test case and t0 as the
    test time.
    """
    t_0 = 2456836.22
    u_0 = 0.922
    t_E = 22.87
    pi_E_N = -0.248
    pi_E_E = 0.234

    ground_model = mm.Model(
        {'t_0': t_0, 'u_0': u_0, 't_E': t_E, 'pi_E': [pi_E_N, pi_E_E]},
        coords='17:47:12.25 -21:22:58.2')
    space_model = mm.Model(
        {'t_0': t_0, 'u_0': u_0, 't_E': t_E, 'pi_E': [pi_E_N, pi_E_E]},
        ra='17:47:12.25', dec='-21:22:58.2',
        ephemerides_file=SAMPLE_FILE_03_EPH)

    delta = ground_model.magnification(t_0) - space_model.magnification(t_0)
    assert np.abs(delta) > 0.01

def test_horizons_3d():
    """
    Test if Horizons properly reads file with 3D coordinates.
    We use the satellite that has the same position as Earth.
    """
    file_in = os.path.join(dir_2, "earth_position_1.dat")
    file_out = os.path.join(dir_2, "earth_position_2.dat")

    satellite = mm.SatelliteSkyCoord(file_in)
    (times, vec_x, vec_y, vec_z) = np.loadtxt(file_out, unpack=True)
    output = satellite.get_satellite_coords(times).cartesian

    np.testing.assert_almost_equal(vec_x, output.x)
    np.testing.assert_almost_equal(vec_y, output.y)
    np.testing.assert_almost_equal(vec_z, output.z)

reason: import numpy,from astropy


text:
import matplotlib.pyplot as plt
import numpy as np
from typing import Dict, List

from coordinated import configs

plt.rcParams.update({'figure.max_open_warning': 0})

def make_scatter_plot(slot2xy: Dict[str, np.ndarray],
                      what: str,
                      step: int,
                      cat_id2coordinate: Dict[int, List[float]],
                      exclude_slot_y: bool = True,
                      ) -> plt.Figure:

    # fig
    fig, ax = plt.subplots(figsize=configs.Fig.fig_size, dpi=configs.Fig.dpi)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.set_xlabel('Unit 1', fontsize=configs.Fig.axis_fs)
    ax.set_ylabel('Unit 2', fontsize=configs.Fig.axis_fs)
    ax.set_xlim([-1.1, 1.1])
    ax.set_ylim([-1.1, 1.1])
    ax.set_xticks([-1.0, 0, 1.0])
    ax.set_yticks([-1.0, 0, 1.0])

    ax.axhline(y=0, color='grey', linestyle=':', zorder=1)
    ax.axvline(x=0, color='grey', linestyle=':', zorder=1)
    plt.title(f'{what}\nstep={step}\n')

    # plot target locations
    for coordinates in cat_id2coordinate.values():
        x, y = coordinates
        ax.scatter(x, y, s=30, c='black', zorder=3, marker='v')

    # plot embeddings
    for slot, embeddings in slot2xy.items():
        assert embeddings.shape[1] == 2
        if exclude_slot_y and slot == 'y':
            continue
        x, y = embeddings[:, 0], embeddings[:, 1]
        ax.scatter(x, y, label=slot, s=10)

    plt.legend(loc='upper right')
    plt.tight_layout()

    return fig

max_stars_repo_name: phueb/Coordinated

max_issues_repo_name: phueb/Coordinated

max_forks_repo_name: phueb/Coordinated


text:
import torch
import torch.nn.functional as F
import time
import os
import sys
import numpy as np
from imageio import imwrite

from utils import AverageMeter

def normalize_data(data):
    data_min = np.min(data)
    data_max = np.max(data)
    data_norm = np.clip((data - data_min) *
                        (255.0 / (data_max - data_min)),
                0, 255).astype(np.uint8)
    return data_norm

def save_video_results(output_buffer, save_path):
    video_outputs = torch.stack(output_buffer)
    for i in range(video_outputs.size(0)):
        save_name = os.path.join(save_path, 'pred_sal_{0:06d}.jpg'.format(i+9))
        imwrite(save_name, normalize_data(video_outputs[i][0].data.numpy()))

def test(data_loader, model, opt):
    print('test')

    model.eval()

    with torch.no_grad():

        batch_time = AverageMeter()
        data_time = AverageMeter()

        end_time = time.time()
        output_buffer = []
        previous_video_id = ''

        for i, (data, targets, valid) in enumerate(data_loader):
            data_time.update(time.time() - end_time)

            inputs = data['rgb']
            curr_batch_size = inputs.size(0)

            while inputs.size(0) < opt.batch_size:
                inputs = torch.cat((inputs, inputs[0:1, :]), 0)
            while data['audio'].size(0) < opt.batch_size:
                data['audio'] = torch.cat((data['audio'], data['audio'][0:1, :]), 0)

            outputs = model(inputs, data['audio'])

            for ii in range(0, len(outputs['sal'])):
                outputs['sal'][ii] = outputs['sal'][ii][0:curr_batch_size,:]

            if not opt.no_sigmoid_in_test:
                outputs['sal'] = F.sigmoid(outputs['sal'][-1])

            for j in range(outputs['sal'].size(0)):
                if not (i == 0 and j == 0) and targets['video'][j] != previous_video_id:
                    save_path = os.path.join(opt.result_path, opt.dataset, previous_video_id)
                    if not os.path.exists(save_path):
                        os.makedirs(save_path)
                    save_video_results(output_buffer, save_path)
                    output_buffer = []
                output_buffer.append(outputs['sal'][j].data.cpu())
                previous_video_id = targets['video'][j]

            batch_time.update(time.time() - end_time)
            end_time = time.time()

            print('[{}/{}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'.format(
                      i + 1,
                      len(data_loader),
                      batch_time=batch_time,
                      data_time=data_time))

    save_path = os.path.join(opt.result_path, opt.dataset, previous_video_id)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    save_video_results(output_buffer, save_path)

max_stars_repo_name: atsiami/STAViS


text:
#!/usr/bin/python
#
# Copyright 2018-2022 Polyaxon, 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 as np

from hypertune.logger import logger
from hypertune.matrix.utils import get_max, get_min, to_numpy
from polyaxon.polyflow import V1Optimization

class SearchSpace:
    def __init__(self, config):
        self.config = config
        self._dim = 0
        self._bounds = []
        self._features = []
        self._discrete_features = {}
        self._categorical_features = {}
        self._x = []
        self._y = []

        self.set_bounds()

    def is_observations_valid(self):
        len_x = len(self.x)
        len_y = len(self.y)
        if len_x != len_y:
            logger.warning("X and Y observations don't have the same size.")
            return False

        if len_x == 0:
            logger.warning("Space has no observations.")
            return False

        if len_y == 0:
            logger.warning("Space has no observations.")
            return False

        return True

    @property
    def x(self):
        return self._x

    @property
    def y(self):
        return self._y

    @property
    def dim(self):
        return self._dim

    @property
    def features(self):
        return self._features

    @property
    def discrete_features(self):
        return self._discrete_features

    @property
    def categorical_features(self):
        return self._categorical_features

    @property
    def bounds(self):
        return self._bounds

    def set_bounds(self):
        bounds = []
        for key in sorted(self.config.params.keys()):
            value = self.config.params[key]
            self._features.append(key)
            # one hot encoding for categorical type
            if value.is_categorical:
                values = to_numpy(value)
                num_feasible = len(values)
                for _ in range(num_feasible):
                    bounds.append((0, 1))
                self._categorical_features[key] = {
                    "values": values,
                    "number": num_feasible,
                }
                self._dim += num_feasible
            elif value.is_discrete:
                self._dim = self._dim + 1
                discrete_values = to_numpy(value)
                bounds.append((get_min(value), get_max(value)))
                self._discrete_features[key] = {"values": discrete_values}
            elif value.is_uniform:
                self._dim = self._dim + 1
                bounds.append((float(get_min(value)), float(get_max(value))))
        self._bounds = np.asarray(bounds)

    def parse_y(self, metrics):
        if not metrics:
            return metrics
        y_values = []
        for value in metrics:
            if V1Optimization.maximize(self.config.metric.optimization):
                y_values.append(float(value))
            else:
                y_values.append(-float(value))

        return np.array(y_values)

    def parse_x(self, configs):
        if not configs:
            return configs
        x = []
        for config in configs:
            x_config = []
            for feature in self._features:
                if feature in self._categorical_features:
                    x_config += [
                        1 if v == config[feature] else 0
                        for v in self._categorical_features[feature]["values"]
                    ]
                elif feature in self._features:
                    x_config.append(config[feature])
            x.append(x_config)
        return np.array(x)

    def add_observations(self, configs, metrics):
        self._x = self.parse_x(configs=configs)
        self._y = self.parse_y(metrics=metrics)

    def _get_discrete_suggestion(self, feature, suggestion, counter):
        feasible_values = self._discrete_features[feature]["values"]
        current_value = suggestion[counter]

        diff = np.subtract(feasible_values, current_value)
        diff = np.absolute(diff)
        results = feasible_values[np.argmin(diff)]
        return results, counter + 1

    def _get_categorical_suggestion(self, feature, suggestion, counter):
        one_hot_values = suggestion[
            counter : counter + self._categorical_features[feature]["number"]
        ]
        index = np.argmax(one_hot_values)
        feasible_values = self._categorical_features[feature]["values"]
        results = feasible_values[index]
        return results, counter + self._categorical_features[feature]["number"]

    def get_suggestion(self, suggestion):
        if suggestion is None:
            return suggestion
        counter = 0
        results = []
        for feature in self._features:
            if feature in self._discrete_features:
                result, counter = self._get_discrete_suggestion(
                    feature=feature, suggestion=suggestion, counter=counter
                )
                results.append(result)

            elif feature in self._categorical_features:
                result, counter = self._get_categorical_suggestion(
                    feature=feature, suggestion=suggestion, counter=counter
                )
                results.append(result)
            else:
                results.append(suggestion[counter])
                counter = counter + 1
        return dict(zip(self._features, results))

max_stars_repo_name: polyaxon/platform

max_issues_repo_name: polyaxon/hypertune

max_forks_repo_name: polyaxon/hypertune


text:
#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os
import numpy as np
import copy
from local.pso.pso_atoms import PsoAtoms
from local.pso.Pre_Relax import PreRelax
from scipy.optimize import linear_sum_assignment

class Pso(object):
    """Pso object.

    The Pso object contains the main process of PSO.
    It must contains a list of atoms objects for optimization,a parameters
    dictionary and a calculator object for dft calculation.
    The main Pso equation:v(t+1)=w*v(t)+c1*r1*(pbest-x)+c2*r2*(gbest-x)
    The default value of this parameters is:WEIGHT=0.6,c1=c2=2,r1 and r2 is two
    seperately generated random numbers is range [0,1].
    NPAR:number of particles.It increases with the size of the system.Default
    value is 20.
    EDIFF:If the energy difference between the neighbour gbest is lower than
    it,the Pso_evolution will stop.The lower,the more precise.
    Set this parameters' value by a pso_init file or using set_parameters().
    """

    def __init__(self, atoms_list=None, parameters=None, calculator=None):

        self.set_atoms_list(atoms_list)
        self.default_parameters={'UNF':1.,'c1':2.0,'c2':2.0,'COC':0.7298,'NPAR':150,'EDIFF':1e-4,'LSS':3,'ELIR':0.4,'WEIGHT':0.9,'GER':1,'PREX':1,'FIX2':[0],'PATH':0.01,'VMAX':0.5,'SCO':0.05,'DCO':2.,'ZMAX':[1.,1.,1],'ZMIN':[0.,0.,0.],'VAC':[0.,0.,10.],'PBC':[0,0,0],'CUO':[0.8,1.6],'LPAIR':1,'LPRELAX':1}
        if parameters is None:#Use the default value
            parameters=self.default_parameters
        if 'pso_init'  in os.listdir(os.getcwd()):
            from local.pso.Read_Pso_Init import PSO_INIT

            self.pso_init=PSO_INIT()
            tmp=self.default_parameters
            for key in tmp.keys():
                if key in self.pso_init.parameters.keys():
                    tmp[key]=self.pso_init.parameters[key]
            parameters=tmp
        self.set_parameters(parameters)

        if calculator is None:
            from local.pso.nnmodel import Model
            nnmodel = Model()
            self.set_calc(nnmodel)
        self.set_calc(calculator)

    def set_calc(self,calculator):
        self._calc=calculator

    def get_calc(self):
        return self._calc

    def set_atoms_list(self,atoms_list=None):
        self._atoms_list=[]
        if atoms_list!=None:
            for atoms in atoms_list:
                self.add_atoms(atoms)

    def add_atoms(self,atoms=None):
        if not isinstance(atoms,PsoAtoms):
            raise TypeError('The given atoms is not a PsoAtoms object.')
        self._atoms_list.append(atoms)

    def del_atoms(self,atoms=None):
        self._atoms_list.remove(atoms)

    def get_atoms(self):
        return self._atoms

    def set_parameters(self,parameters):
        if not isinstance(parameters,dict):
            raise ValueError('The given format is not a dictionary')
        a=set(parameters.keys())
        b=set(self.default_parameters.keys())
        if not a<=b:
            raise ValueError
        c=self.default_parameters
        for key in parameters.keys():
            c[key]=parameters[key]
        self._parameters=c

    def get_parameters(self):
        return copy.deepcopy(self._parameters)

def pso_evo(pso=None, natoms=None):
        from local.pso.rand_stru import RandStru

        "It is the main function of Pso."
        import os
        import shutil

        def write_pso_data():
            filename='pso_data_%03d'%GER
            atoms=pso._atoms_list[0]
            f=open(filename,'w')
            f.write('System\n')
            f.write(atoms.get_system_name())
            f.write('\n')
            f.write('Subs Atom\n')
            for elem in atoms._subs_elements:
                f.write(elem+'  ')
            f.write('\n')
            for num in atoms._subs_numbers:
                f.write('%3d  '%num)
            f.write('\n')
            f.write('Abso Atom\n')
            for elem in atoms._abso_elements:
                f.write(elem+'  ')
            f.write('\n')
            for num in atoms._abso_numbers:
                f.write('%3d  '%num)
            f.write('\n')
            f.write('Lattice Constant %.16f\n'%atoms.get_lattice_constant())
            f.write('Cell\n')
            for i in atoms.get_cell():
                for j in i:
                    f.write('%.16f  '%j)
                f.write('\n')
            f.write('Pbc    ')
            for i in atoms.get_pbc():
                f.write('%d  '%i)
            f.write('\n')
            f.write('Subs Mass\n')
            for mass in atoms.get_subs_masses():
                f.write('%.4f  '%mass)
            f.write('\n')
            f.write('Abso Mass\n')
            for mass in atoms.get_abso_masses():
                f.write('%.4f   '%mass)
            f.write('\n')
            f.write('Subs Radius\n')
            for radius in atoms.get_subs_radius():
                f.write('%.4f  '%radius)
            f.write('\n')
            f.write('Abso Radius\n')
            for radius in atoms.get_abso_radius():
                f.write('%.4f  '%radius)
            f.write('\n')
            f.write('Constraints\n')
            for i in atoms.get_constraints():
                for j in i:
                   f.write('%d '%j)
            f.write('\n')
            f.write('Subs Structure\n')
            for atom in atoms.get_subs_positions():
                for cord in atom:
                    f.write('%.16f  '%cord)
                f.write('\n')
            f.write('Pso Parameters\n')
            for key in pso._parameters.keys():
                f.write('%s  '%key)
                value=pso._parameters[key]
                if isinstance(value,list):
                    for i in value:
                        f.write('%.8f  '%i)
                    f.write('\n')
                else:
                    f.write('%.8f  \n'%pso._parameters[key])
            f.write('Calculator  '+pso._calc._name+'\n')
            f.write('!!!!!!!!!!!!!!!!!!!!!!!\n')
            f.write('          Generation  %d\n'%GER)
            f.write('\n')
            if GER==1:
                f.write('Last Gbest  %.16f\n'%0)
            else:
                f.write('Last Gbest  %.16f\n'%gbest[1])
                f.write('\n&&&&&&&&&&&&&&&  Number of Eliminated structures  &&&&&&&&&&&&&&&\n')
                for i in elim_list:
                    f.write('%2d  '%i)
                f.write('\n\n')
            for i1,atoms in enumerate(pso._atoms_list):
                stru=atoms.get_abso_positions()
                f.write('----------Particle  %d----------\n'%i1)
                f.write('Positions\n')
                for atom in stru:
                    for cord in atom:
                        f.write('%.16f    '%cord)
                    f.write('\n')
                f.write('Velocities\n')
                for i2 in range(len(stru)):
                    for v in velocities[i1,i2,:]:
                        f.write('%.16f    '%v)
                    f.write('\n')
                f.write('            *******************************\n')
            f.close()

        def new_velocity(atoms,v,pbest,gbest,lpair):
            from math import e
            v1,v2=[],[]
            c1,c2=pso._parameters['c1'],pso._parameters['c2']
            x=pso._parameters['COC']
            w=pso._parameters['WEIGHT']
            r1=np.random.rand()
            r2=np.random.rand()
            w = 0.4 + 0.5 / GER
            if np.abs(pbest[0]-gbest[1])<=5e-2:
                w=0.25
            f2.write('x  %.3f; w  %.3f; c1  %.3f; c2  %.3f; r1  %.3f; r2  %.3f\n'%(x,w,c1,c2,r1,r2))
            f2.write('Last Velocities\n')
            for i in v:
                 for j in i:
                     f2.write('%.16f  '%j)
                 f2.write('\n')
            temp=0
            stru=atoms.get_abso_positions(cord_mod='r')
            pbest=pbest[1].get_abso_positions(cord_mod='r')
            gbest=gbest[2].get_abso_positions(cord_mod='r')
            f2.write('Pbest\n')
            for n in atoms.get_abso_numbers():
                dist=[atoms.get_distance(atom1,atom2)[0] for atom1 in stru[temp:temp+n] for atom2 in pbest[temp:temp+n]]
                dist=np.array(dist)
                dist=dist.reshape(n,n)
                for i in dist:
                    for j in i:
                        f2.write('%.16f    '%j)
                    f2.write('\n')
                if lpair:
                    path1=get_pair(dist)[0]
                else:
                    path1=range(len(dist))
                for i in path1:
                    f2.write('%d    '%i)
                f2.write('\n')
                for i,j in enumerate(path1):
                    v1.append(atoms.get_distance(pbest[temp+j],stru[temp+i])[1])
                temp+=n
            v1=np.array(v1)
            f2.write('v1\n')
            for i in v1:
                for j in i:
                    f2.write('%.16f  '%j)
                f2.write('\n')
            temp=0
            f2.write('Gbest\n')
            for n in atoms.get_abso_numbers():
                dist=[atoms.get_distance(atom1,atom2)[0] for atom1 in stru[temp:temp+n] for atom2 in gbest[temp:temp+n] ]
                dist=np.array(dist)
                dist=dist.reshape(n,n)
                for i in dist:
                    for j in i:
                        f2.write('%.16f    '%j)
                    f2.write('\n')
                if lpair:
                    path1=get_pair(dist)[0]
                else:
                    path1=range(len(dist))
                for i in path1:
                    f2.write('%d    '%i)
                f2.write('\n')
                for i,j in enumerate(path1):
                    v2.append(atoms.get_distance(gbest[temp+j],stru[temp+i])[1])
                temp+=n
            v2=np.array(v2)
            f2.write('v2\n')
            for i in v2:
                for j in i:
                    f2.write('%.16f  '%j)
                f2.write('\n')
            f2.write('\n')
            new_velo=x*(c1*r1*v1+c2*r2*v2)+w*v
            return new_velo

        init_dir=os.getcwd()
        pbest=[]
        gbest=[0,1e5,None]

        # initialization, generate the initial velocity randomly in the range [-0.1,0.1]
        if pso is not None:
            if not isinstance(pso, Pso):
                raise ValueError('NO Pso Object')
            ediff=pso._parameters['EDIFF']
            npar=pso._parameters['NPAR']
            c1=pso._parameters['c1']
            c2=pso._parameters['c2']
            unf=pso._parameters['UNF']
            coc=pso._parameters['COC']
            lss=pso._parameters['LSS']
            elir=pso._parameters['ELIR']
            GER=pso._parameters['GER']
            vmax=pso._parameters['VMAX']
            vac=pso._parameters['VAC']
            dis_cutoff=pso._parameters['DCO']
            lprelax=pso._parameters['LPRELAX']
            lpair=pso._parameters['LPAIR']
            ntsubs=len(pso._atoms_list[0]._subs_symbols)
            ntabso=len(pso._atoms_list[0]._abso_symbols)
            f0 = open('pso_data','a')
            f0.write('%s\n' % pso._atoms_list[0].get_system_name())
            f0.write('Parameters\n')
            for key in pso._parameters.keys():
                f0.write('%s  '%key)
                value = pso._parameters[key]
                if isinstance(value, list):
                    for i in pso._parameters[key]:
                        f0.write('%.3f  ' % i)
                    f0.write('\n')
                else:
                    f0.write('%.3f\n' % pso._parameters[key])
            f0.write("--------Substrate's Atoms' Positions:--------\n")
            for atom in pso._atoms_list[0].get_subs_positions():
                for cord in atom:
                    f0.write('%.16f  '%cord)
                f0.write('\n')
            f0.write('**********************************\n')
            velocities = vmax-2*vmax*np.random.rand(npar, ntabso, 3)
            dirname = 'pso_'+'001'
            if lprelax:
                PreRelax(pso, filename='pre_relax_001', dirt=init_dir)
            if dirname in os.listdir(os.getcwd()):shutil.rmtree(dirname)
            os.mkdir(dirname)
            os.chdir(dirname)
            for n, atoms in enumerate(pso._atoms_list):
                dirname1 = dirname+'_'+'%04d'%n
                if dirname1 in os.listdir(os.getcwd()):shutil.rmtree(dirname1)
                os.mkdir(dirname1)
                os.chdir(dirname1)
                pso._calc.sp_run(atoms=atoms, input_dir=init_dir)
                os.chdir('..')
            write_pso_data()
            os.chdir('..')
        else:
            # Read information from pso_data
            from local.pso.Read_Pso_Init import PSO_INIT
            pso_init=PSO_INIT()
            GER=pso_init.parameters['GER']

            updated = pso_init.parameters['UPDATED']
            f0=open('pso_data','a')
            os.chdir('pso_%03d'%GER)
            pso,last_gbest,last_elim=read_data(filename='pso_data_%03d'%GER)
            ediff=pso._parameters['EDIFF']
            npar=pso._parameters['NPAR']
            c1=pso._parameters['c1']
            c2=pso._parameters['c2']
            unf=pso._parameters['UNF']
            coc=pso._parameters['COC']
            lss=pso._parameters['LSS']
            elir=pso._parameters['ELIR']
            vmax=pso._parameters['VMAX']
            vac=pso._parameters['VAC']
            zmax=pso._parameters['ZMAX']
            zmin=pso._parameters['ZMIN']
            sim_cutoff=pso._parameters['SCO']
            dis_cutoff=pso._parameters['DCO']
            lprelax=pso._parameters['LPRELAX']
            lpair=pso._parameters['LPAIR']
            gen =pso._parameters['GER']
            ntsubs=len(pso._atoms_list[0]._subs_symbols)
            ntabso=len(pso._atoms_list[0]._abso_symbols)

            # Main Loop
            # Read information from result
            os.chdir('..')
            os.chdir('pso_%03d'%GER)
            dir_list=os.listdir(os.getcwd())
            numofupdate = 0
            numofinit = 0
            numoftest = 0

            for n in range(int(npar)):
                print("it's a {0} particle".format(n))
                dirname = 'pso_%03d'%GER+'_%04d'%n
                os.chdir(dirname)

                if not updated:
                    from local.optimize_cluster import structure_optimization
                    from local.error_indicator import read_trajectory
                    if n == 0:
                        if 'update_' + str(int(GER)).zfill(3) in os.listdir("../"):
                            shutil.rmtree('../update_' + str(int(GER)).zfill(3))
                        os.mkdir('../update_' + str(int(GER)).zfill(3))

                        if 'updating_' + str(int(GER)).zfill(3) in os.listdir('../'):
                            os.system('rm -rf ../updating_' + str(int(GER)).zfill(3))
                        if 'test_store_' + str(int(GER)).zfill(3) in os.listdir('../'):
                            os.system('rm -rf ../test_store_' + str(int(GER)).zfill(3))

                    structure_optimization(filename='POSCAR', gen=GER, natoms=natoms)
                    numofupdate, numofinit, numoftest = \
                        read_trajectory(gen=GER, prev_update=numofupdate, prev_init=numofinit, prev_test=numoftest, natoms=natoms)
                    os.system('cp ../../../clustertut/ase_calcs/optimization.poscar ./POSCAR_pbest')
                    os.system('cp ./POSCAR_pbest ../update_' + str(int(GER)).zfill(3) + '/POSCAR_' + str(n).zfill(4))

                energy=pso._calc.get_energy(updated=updated, num=n, gen=GER)
                pso._atoms_list[n].set_atoms_energy(energy)
                pi = [energy, pso._atoms_list[0].copy()]
                abso_stru=pso._calc.get_stru(pi[1])
                pi[1].set_abso_positions(abso_stru, cord_mod='d')
                pbest.append(pi)
                os.chdir('..')

            if updated:
                energies=[i[0] for i in pbest]
                energies=np.array(energies)
                energy_sort=np.argsort(energies)
                gbest=[energies.argmin(),energies.min(),pbest[energies.argmin()][1].copy()]
                velocities=[atoms.get_atoms_velo() for atoms in pso._atoms_list]
                velocities=np.array(velocities)
                filename='pso_data_%03d'%GER
                f2=open(filename,'r')
                f3=open('tmp','w')
                count=0
                for line in f2:
                    if 'Last Gbest' in line:
                        f3.write('Energies sort list\n')
                        for i in energy_sort:
                            f3.write('%d    %.16f\n'%(i,energies[i]))
                    if '            *******************************' in line:
                        f3.write('Pbest Positions and Energy\n')
                        for i in pbest[count][1].get_abso_positions():
                            for j in i:
                                f3.write('%.16f  '%j)
                            f3.write('\n')
                        f3.write('Pbest Free Energy  %.16f\n'%pbest[count][0])
                        count+=1
                    f3.write(line)

                f3.write('----------Gbest Positions and Energy----------\n')
                f3.write('Gbest Positions\n')
                for i in gbest[2].get_abso_positions():
                    for j in i:
                        f3.write('%.16f  '%j)
                    f3.write('\n')
                f3.write('Gbest Free Energy  %.16f\n'%gbest[1])
                f3.write('Gbest Number  %d\n'%gbest[0])
                f3.close()
                f2.close()
                os.rename('tmp',filename)
                os.chdir(init_dir)
                if np.abs(gbest[1]-last_gbest)>=np.abs(ediff):
                    GER+=1
                    pso._parameters['GER']+=1
                    # Update Swarm
                    f2=open('velocities_%03d'%GER,'w')
                    for n,atoms in enumerate(pso._atoms_list):
                        f2.write('*************** Particle  %d ***************\n'%n)
                        velocities[n]=new_velocity(atoms,velocities[n],pbest[n],gbest,lpair)
                    f2.close()
                    #eliminate the high energy structures,and substitute them by new random structures.
                    neli=int(elir*npar)
                    elim_list=energy_sort[-neli:]
                    surv_list=energy_sort[:-neli]
                    elim_list=list(elim_list)

                    surv_list=list(surv_list)

                    surv_list.reverse()
                    tmp1=[]
                    tmp2=[]
                    #if one structure is both in last_elim and elim,we do not eliminate it and keep it oen generation!
                    for n in elim_list:
                        if n in last_elim:
                            for m in surv_list:
                                if m not in last_elim:
                                    surv_list.remove(m)
                                    tmp1.append(m)
                                    tmp2.append(n)
                                    break
                                if m==surv_list[-1]:
                                    tmp1.append(n)
                        else:
                            tmp1.append(n)
                    elim_list=tmp1
                    surv_list.extend(tmp2)
                    for n in elim_list:
                        atoms=pso._atoms_list[n]
                        RandStru(atoms, natoms=natoms)
                        velocities[n]=vmax-vmax*2*np.random.rand(ntabso,3)
                    for n in surv_list:
                        atoms=pso._atoms_list[n]
                        stru=atoms.get_abso_positions(cord_mod='r')
                        stru=stru+velocities[n]
                        atoms.set_abso_positions(stru,cord_mod='r')
                     #Evaluate Swarm
                    if lprelax:
                        PreRelax(pso,filename='pre_relax_%03d'%GER,dirt=init_dir)
                    f0.write('Generation %d\n'%GER)
                    dirname='pso_'+'%03d'%GER
                    if dirname in os.listdir(os.getcwd()):shutil.rmtree(dirname)
                    os.mkdir(dirname)
                    os.chdir(dirname)
                    for n,atoms in enumerate(pso._atoms_list):
                        dirname1=dirname+'_'+'%04d'%n
                        if dirname1 in os.listdir(os.getcwd()):shutil.rmtree(dirname1)
                        os.mkdir(dirname1)
                        os.chdir(dirname1)
                        pso._calc.sp_run(atoms=atoms,input_dir=init_dir)
                        os.chdir('..')
                    temp = pso._atoms_list
                    write_pso_data()
                    print('Done!')
                    os.chdir('..')
                else:  # energy converge
                    print('COMPLETED!')
                    f0.write('\n\n***************Energy Converged!***************\n')
                    return gbest[2]
        f0.close()

def read_data(filename='pso_data_001'):
    # from vasp import Vasp
    from local.pso.nnmodel import Model
    'read information from pso-data file and return a Pso object'
    f=open(filename,'r')
    print(filename)
    data=[line.strip() for line in f ]
    count=0
    last_elim=[]
    last_gbest=0.
    for n,line in enumerate(data):
        if 'System' in line:
            system_name=data[n+1]
            continue
        if 'Subs Atom' in line:
            subs_elements=data[n+1].split()
            subs_numbers=[int(i) for i in data[n+2].split()]
            nsubs=sum(subs_numbers)
            continue
        if 'Abso Atom' in line:
            abso_elements=data[n+1].split()
            abso_numbers=[int(i) for i in data[n+2].split()]
            nabso=sum(abso_numbers)
            continue
        if 'Lattice Constant' in line:
            lattice_constant=float(line.split()[2])
            continue
        if 'Cell' in line:
            cell=[float(j) for i in range(3) for j in data[n+i+1].split()]
            continue
        if 'Pbc' in line:
            pbc=[int(i) for i in line.split()[1:]]
        if 'Subs Mass' in line:
            subs_masses=[float(i) for i in data[n+1].split()]
            continue
        if 'Abso Mass' in line:
            abso_masses=[float(i) for i in data[n+1].split()]
            continue
        if 'Subs Radius' in line:
            subs_radius=[float(i) for i in data[n+1].split()]
            continue
        if 'Abso Radius' in line:
            abso_radius=[float(i) for i in data[n+1].split()]
            continue
        if 'Constraints' in line:
            const=[int(i) for i in data[n+1].split()]
            continue
        if 'Subs Structure' in line:
            subs_positions=[float(j) for i in range(nsubs) for j in data[n+1+i].split()]
            continue
        if 'Parameters' in line:
            parameters={}
            for i in data[n+1:]:
                if 'Calculator' in i:
                    calc=i
                    break
                key=i.split()[0]
                a=[float(j) for j in i.split()[1:]]
                if len(a)==1 and key!='FIX2':a=a[0]
                parameters[key]=a
            cot=1
            atoms_list=[]
            while cot<=parameters['NPAR']:
                subs_symbols=[]
                abso_symbols=[]
                for i,elem in enumerate(subs_elements):
                    subs_symbols.extend([elem]*subs_numbers[i])
                for i,elem in enumerate(abso_elements):
                    abso_symbols.extend([elem]*abso_numbers[i])
                atoms=PsoAtoms(name=system_name,subs_symbols=subs_symbols,                                                         abso_symbols=abso_symbols,subs_positions=subs_positions,                                            subs_masses=subs_masses,abso_masses=abso_masses,
                               subs_radius=subs_radius,abso_radius=abso_radius,cell=cell,
                               lattice_constant=lattice_constant,constraints=const,pbc=pbc)
                atoms_list.append(atoms)
                cot+=1
            nnmodel = Model(lattice_parameter=lattice_constant)
            pso=Pso(atoms_list=atoms_list,calculator=nnmodel,parameters=parameters)
            print('Pso Object Done!')
            print(len(pso._atoms_list))
        if 'Last Gbest' in line:
            last_gbest=float(line.split()[2])
            continue
        if 'Number of Eliminated structures' in line:
            last_elim=[int(i) for i in data[n+1].split()]
            continue
        if 'Positions' in line and 'Pbest' not in line and 'Gbest' not in line:
            abso_positions=[float(j) for i in range(nabso) for j in data[n+1+i].split()]
            pso._atoms_list[count].set_abso_positions(abso_positions)
            continue
        if 'Velocities' in line:
            velocities=[float(j) for i in range(nabso) for j in data[n+1+i].split()]
            pso._atoms_list[count].set_atoms_velo(velocities)
            count+=1
            continue

    #print abso_elements,abso_numbers,abso_symbols
    return pso,last_gbest,last_elim

def get_pair(dist):

    row,col=linear_sum_assignment(dist)
    return col,dist[row,col].sum()

max_stars_repo_name: TongheYing/ML-Au

max_issues_repo_name: TongheYing/ML-Au

max_forks_repo_name: TongheYing/ML-Au

reason: import numpy,from scipy


text:
#encoding: utf-8
from __future__ import print_function
from __future__ import division

from numpy.testing import assert_allclose
from sklearn.linear_model import LinearRegression, Ridge
import numpy as np

from posthoc import Workbench, WorkbenchOptimizer, cov_estimators
from posthoc.utils import gen_data

def _compare_models(wb, base, X, atol=0, rtol=1E-5):
    """Compare a Workbench model with a base model."""
    assert_allclose(wb.coef_, base.coef_, atol, rtol)
    assert_allclose(wb.intercept_, base.intercept_, atol, rtol)

    y_hat_wb = wb.predict(X)
    y_hat_base = base.predict(X)
    assert_allclose(y_hat_wb, y_hat_base, atol, rtol)

def test_pattern_computation():
    """Test computation of the pattern from a linear model."""
    X, y, A = gen_data(noise_scale=0)

    assert (X.std(axis=0) != 0).all()
    assert (X.mean(axis=0) != 0).all()

    assert (y.std(axis=0) != 0).all()
    assert (y.mean(axis=0) != 0).all()

    for normalize in [False, True]:
        # The base model is a simple OLS regressor
        ols = LinearRegression(normalize=normalize)

        # Don't specify any modifications, should trigger shortcut path
        wb = Workbench(ols).fit(X, y)

        # Test pattern computation
        assert_allclose(wb.pattern_, A, atol=1E-7)

def test_inverse_predict():
    """Test inverting a linear model."""
    X, y, _ = gen_data()
    w = Workbench(LinearRegression()).fit(X, y)
    y_hat = w.predict(X)
    m = LinearRegression().fit(y_hat, X)
    assert_allclose(w.pattern_, m.coef_, atol=1E-15)
    assert_allclose(w.inverse_predict(y), m.predict(y))

def test_identity_transform():
    """Test disassembling and re-assembling a model as-is."""
    X, y, A = gen_data()
    train = np.arange(500)  # Samples used as training set
    test = train + 500  # Samples used as test set

    # The base model is a simple OLS regressor
    ols = LinearRegression().fit(X[train], y[train])

    # Specify identity modifications.
    def pattern_modifier(pattern, X, y):
        return pattern

    def normalizer_modifier(normalizer, X_train, y_train, cov_X, coef):
        return normalizer

    # Using a modifier function
    wb = Workbench(LinearRegression(), cov=cov_estimators.Empirical(),
                   pattern_modifier=pattern_modifier,
                   normalizer_modifier=normalizer_modifier)
    wb.fit(X[train], y[train])
    _compare_models(wb, ols, X[test])

def test_ridge_regression():
    """Test post-hoc adaptation of OLS to be a ridge regressor."""
    X, y, A = gen_data()
    train = np.arange(500)  # Samples used as training set
    test = train + 500  # Samples used as test set

    # The shrinkage parameter
    alpha = 5

    # The ridge regressor we're going to imitate
    ridge = Ridge(alpha, normalize=False).fit(X[train], y[train])

    # The OLS regressor we're going to adapt to be a ridge regressor
    ols = LinearRegression(normalize=False).fit(X[train], y[train])

    # Perform the post-hoc adaptation using different code paths
    wb = Workbench(ols, cov=cov_estimators.L2(alpha, scale_by_var=False))
    wb.fit(X[train], y[train])
    _compare_models(wb, ridge, X[test])

    # Try the "kernel" method
    wb = Workbench(ols, cov=cov_estimators.L2Kernel(alpha, scale_by_var=False))
    wb.fit(X[train], y[train])
    _compare_models(wb, ridge, X[test])

def test_post_hoc_modification():
    """Test post-hoc modification of the model."""
    X, y, A = gen_data(noise_scale=10)

    train = np.arange(5)         # Samples used as training set
    test = np.arange(5, len(X))  # Samples used as test set

    # The OLS regressor we're going to try and optimize
    ols = LinearRegression().fit(X[train], y[train])
    ols_score = ols.score(X[test], y[test])

    # Use post-hoc adaptation to swap out the cov matrix estimated on the
    # training data only, with one that was estimated on all data.
    def cov_function(X_train):
        X_ = X - X.mean(axis=0)
        return X_.T.dot(X_)

    # Use post-hoc adaptation to swap out the estimated pattern with the actual
    # pattern.
    def pattern_modifier(pattern, X_train, y_train):
        return A  # The ground truth pattern

    # Use post-hoc adaptation to swap out the normalizer estimated on the
    # training data only, with one that was estimated on all data.
    def normalizer_modifier(normalizer, X_train, y_train, cov_X, coef):
        y_ = y - y.mean(axis=0)
        return y_.T.dot(y_)  # The ground truth normalizer

    wb = Workbench(LinearRegression(),
                   cov=cov_estimators.Function(cov_function),
                   pattern_modifier=pattern_modifier,
                   normalizer_modifier=normalizer_modifier)
    wb.fit(X[train], y[train])

    # The new regressor should perform better
    wb_score = wb.score(X[test], y[test])
    assert wb_score > ols_score
    assert wb_score > 0.9  # The new regression should be awesome

def test_workbench_optimizer():
    """Test using an optimizer to fine-tune parameters."""
    X, y, A = gen_data(noise_scale=50, N=10)

    wbo = WorkbenchOptimizer(LinearRegression(normalize=True),
                             cov=cov_estimators.L2(scale_by_var=False),
                             scoring='r2',
                             cov_param_x0=[0.2],
                             cov_param_bounds=[(0.1, None)],
                             optimizer_options={'maxiter': 100})
    wbo.fit(X, y)

    # Weights of the WorkbenchOptimizer should equal the weights of a
    # Workbench initialized with the optimal parameters.
    wb = Workbench(LinearRegression(normalize=True),
                   cov=cov_estimators.L2(*wbo.cov_params_, scale_by_var=False))
    wb.fit(X, y)
    assert_allclose(wbo.coef_, wb.coef_)
    assert_allclose(wbo.intercept_, wb.intercept_)

    # Weights of the WorkbenchOptimizer should equal the weights of a
    # Ridge initialized with the optimal parameters.
    rr = Ridge(alpha=wbo.cov_params_[0], normalize=True)
    rr.fit(X, y)
    assert_allclose(wbo.coef_, rr.coef_)
    assert_allclose(wbo.intercept_, rr.intercept_)

max_stars_repo_name: wmvanvliet/workbench

max_issues_repo_name: wmvanvliet/posthoc

max_forks_repo_name: wmvanvliet/posthoc

reason: import numpy,from numpy


text:
# Copyright (C) 2020  GreenWaves Technologies, SAS

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from graph.types.input_output import ConstantInputParameters
from generation.helpers.gen_constant import gen_constant
from quantization.qtype import QType
import logging
import numpy as np

from generation.at_types.gen_ctrl import GenCtrl
from generation.bindings import (CommentBindingList, GNodeArgEdge, GNodeArgNode,
                                 NodeBindingList)
from generation.generators.kernels.autotiler_kernel import NewAutoTilerKernel
from generation.new_generators.generator_base import (GeneratorBase, ktype,
                                                      paramstype)
from generation.at_types.constant_info import ConstantInfo
from generation.at_types.tc_arg_info import GlobalArgInfo
from graph.types import SSDDetectorParameters
from quantization.multiplicative.mulbias import set_ssd_scales
from quantization.multiplicative.mult_utils import compute_mul_bias

LOG = logging.getLogger("nntool." + __name__)

@paramstype(SSDDetectorParameters)
@ktype("scaled")
class SSDDetectorSQ8Generator(GeneratorBase):

    @classmethod
    def globals_generator(cls, gen, node, qrec, pnode, fnode) -> bool:
        set_ssd_scales(qrec, node)
        scores_q = qrec.in_qs[1]
        scores_scale, scores_norm = compute_mul_bias(scores_q.scale)

        cname_scales, file_name_scales = gen_constant(gen, node, node, "ssd_scales")
        contents = np.array([qrec.cache['scale_x_q'].qbiases,
                            qrec.cache['scale_x_anc_q'].qbiases,
                            qrec.cache['scale_y_q'].qbiases,
                            qrec.cache['scale_y_anc_q'].qbiases,
                            qrec.cache['scale_h_q'].qbiases,
                            qrec.cache['scale_w_q'].qbiases,
                            qrec.cache['scale_ao_q'].qbiases,
                            scores_scale], dtype=np.int8)
        scale_info = ConstantInfo(file_name_scales, QType.Pow2(
            bits=8, q=0, signed=True), contents=contents)

        cname_norms, file_name_norms = gen_constant(gen, node, node, "ssd_norms")
        contents = np.array([qrec.cache['scale_x_q'].qnorms,
                            qrec.cache['scale_x_anc_q'].qnorms,
                            qrec.cache['scale_y_q'].qnorms,
                            qrec.cache['scale_y_anc_q'].qnorms,
                            qrec.cache['scale_h_q'].qnorms,
                            qrec.cache['scale_w_q'].qnorms,
                            qrec.cache['scale_ao_q'].qnorms,
                            scores_norm], dtype=np.int8)
        norms_info = ConstantInfo(file_name_norms, QType.Pow2(
            bits=8, q=0, signed=True), contents=contents)

        gen.globals.append(GlobalArgInfo(qrec.cache['scale_x_q'].ctype, cname_scales,
                                        gen.opts['default_global_home_location'],
                                        gen.opts['default_global_exec_location'],
                                        const_info=scale_info))

        gen.globals.append(GlobalArgInfo(qrec.cache['scale_x_q'].shift_ctype, cname_norms,
                                        gen.opts['default_global_home_location'],
                                        gen.opts['default_global_exec_location'],
                                        const_info=norms_info))
        return True

    @classmethod
    def bindings_generator(cls, gen, node, qrec, in_eparams, out_eparams, cname) -> bool:
        gen.bindings.append(
            CommentBindingList("Node {} offsetsq {} scoresq {} anchorsq {} outboxesq {}", node.name,
                            str(qrec.in_qs[0]), str(qrec.in_qs[1]), str(qrec.in_qs[2]), str(qrec.out_qs[0]))
        )
        gen.bindings.append(
            NodeBindingList(cname,
                            GNodeArgEdge(in_eparams[0]),
                            GNodeArgEdge(in_eparams[1]),
                            GNodeArgEdge(in_eparams[2]),
                            GNodeArgEdge(out_eparams[0], "GNA_OUT"),
                            GNodeArgEdge(out_eparams[1], "GNA_OUT"),
                            GNodeArgEdge(out_eparams[2], "GNA_OUT"),
                            GNodeArgNode(node, "ssd_scales"),
                            GNodeArgNode(node, "ssd_norms"),)
        )
        return True

    @classmethod
    def kernel_generator(cls, gen, node, qrec, in_eparams, out_eparams, cname) -> bool:
        del in_eparams, out_eparams
        gen.kernels.append(SSDPostProcessKernelSQ8(node.name, cname, node))
        return True

class SSDPostProcessKernelSQ8(NewAutoTilerKernel):
    CALL_TEMPLATE = """
// generator for {node_name}
CNN_SSD_PostProcess_SQ8("{cname}", {gen_ctrl}, {n_anchors}, {n_classes}, {n_outboxes}, {max_bb_before_nms}, {DecScoreThr:.3f}, {NMSThr:.3f});
"""

    def __init__(self, node_name, cname, params, gen_ctrl=None):
        if gen_ctrl is None:
            gen_ctrl = GenCtrl(None, cname=cname)
        else:
            gen_ctrl.cname = cname

        in_dims = params.in_dims
        out_dims = params.out_dims

        attrs = {
            'n_anchors': in_dims[0].shape[0],
            'n_classes': in_dims[1][1],
            'n_outboxes': out_dims[0][0],
            'max_bb_before_nms': params.max_bb_before_nms,
            'DecScoreThr': params.nms_score_threshold,
            'NMSThr': params.nms_iou_threshold
        }

        extra_attrs = {
            'cname': cname,
            'node_name': node_name
        }
        super().__init__(attrs, extra_attrs, gen_ctrl=gen_ctrl)