text:
# Author: Simon Blanke
# Email: simon.blanke@yahoo.com
# License: MIT License

import random
import numpy as np

from ..base_optimizer import BaseOptimizer
from ...search import Search

def sort_list_idx(list_):
    list_np = np.array(list_)
    idx_sorted = list(list_np.argsort()[::-1])
    return idx_sorted

def centeroid(array_list):
    centeroid = []
    for idx in range(array_list[0].shape[0]):
        center_dim_pos = []
        for array in array_list:
            center_dim_pos.append(array[idx])

        center_dim_mean = np.array(center_dim_pos).mean()
        centeroid.append(center_dim_mean)

    return centeroid

class DownhillSimplexOptimizer(BaseOptimizer, Search):
    name = "Downhill Simplex"
    _name_ = "downhill_simplex"

    def __init__(self, *args, alpha=1, gamma=2, beta=0.5, sigma=0.5, **kwargs):
        super().__init__(*args, **kwargs)

        self.alpha = alpha
        self.gamma = gamma
        self.beta = beta
        self.sigma = sigma

        self.n_simp_positions = len(self.conv.search_space) + 1
        self.simp_positions = []

        self.simplex_step = 0

    def finish_initialization(self):
        idx_sorted = sort_list_idx(self.scores_valid)
        self.simplex_pos = [self.positions_valid[idx] for idx in idx_sorted]
        self.simplex_scores = [self.scores_valid[idx] for idx in idx_sorted]

        n_inits = len(self.positions_valid)
        if n_inits < self.n_simp_positions:
            print("\n Error: Not enough initial positions to form simplex")
            print("\n Increase number of initial positions")

        self.simplex_step = 1

        self.i_x_0 = 0
        self.i_x_N_1 = -2
        self.i_x_N = -1

    @BaseOptimizer.track_nth_iter
    def iterate(self):
        simplex_stale = all(
            [np.array_equal(self.simplex_pos[0], array) for array in self.simplex_pos]
        )

        if simplex_stale:
            idx_sorted = sort_list_idx(self.scores_valid)
            self.simplex_pos = [self.positions_valid[idx] for idx in idx_sorted]
            self.simplex_scores = [self.scores_valid[idx] for idx in idx_sorted]

            self.simplex_step = 1

        if self.simplex_step == 1:
            idx_sorted = sort_list_idx(self.simplex_scores)
            self.simplex_pos = [self.simplex_pos[idx] for idx in idx_sorted]
            self.simplex_scores = [self.simplex_scores[idx] for idx in idx_sorted]

            self.center_array = centeroid(self.simplex_pos[:-1])

            r_pos = self.center_array + self.alpha * (
                self.center_array - self.simplex_pos[-1]
            )
            self.r_pos = self.conv2pos(r_pos)
            return self.r_pos

        elif self.simplex_step == 2:
            e_pos = self.center_array + self.gamma * (
                self.center_array - self.simplex_pos[-1]
            )
            self.e_pos = self.conv2pos(e_pos)
            self.simplex_step = 1

            return self.e_pos

        elif self.simplex_step == 3:
            # iter Contraction
            c_pos = self.h_pos + self.beta * (self.center_array - self.h_pos)
            c_pos = self.conv2pos(c_pos)

            return c_pos

        elif self.simplex_step == 4:
            # iter Shrink
            pos = self.simplex_pos[self.compress_idx]
            pos = pos + self.sigma * (self.simplex_pos[0] - pos)

            return self.conv2pos(pos)

    def evaluate(self, score_new):
        self.score_new = score_new

        if self.simplex_step != 0:
            self.prev_pos = self.positions_valid[-1]

        if self.simplex_step == 1:
            # self.r_pos = self.prev_pos
            self.r_score = score_new

            if self.r_score > self.simplex_scores[0]:
                self.simplex_step = 2

            elif self.r_score > self.simplex_scores[-2]:
                # if r is better than x N-1
                self.simplex_pos[-1] = self.r_pos
                self.simplex_scores[-1] = self.r_score
                self.simplex_step = 1

            if self.simplex_scores[-1] > self.r_score:
                self.h_pos = self.simplex_pos[-1]
                self.h_score = self.simplex_scores[-1]
            else:
                self.h_pos = self.r_pos
                self.h_score = self.r_score

            self.simplex_step = 3

        elif self.simplex_step == 2:
            self.e_score = score_new

            if self.e_score > self.r_score:
                self.simplex_scores[-1] = self.e_pos
            elif self.r_score > self.e_score:
                self.simplex_scores[-1] = self.r_pos
            else:
                self.simplex_scores[-1] = random.choice([self.e_pos, self.r_pos])[0]

        elif self.simplex_step == 3:
            # eval Contraction
            self.c_pos = self.prev_pos
            self.c_score = score_new

            if self.c_score > self.simplex_scores[-1]:
                self.simplex_scores[-1] = self.c_score
                self.simplex_pos[-1] = self.c_pos

                self.simplex_step = 1

            else:
                # start Shrink
                self.simplex_step = 4
                self.compress_idx = 0

        elif self.simplex_step == 4:
            # eval Shrink
            self.simplex_scores[self.compress_idx] = score_new
            self.simplex_pos[self.compress_idx] = self.prev_pos

            self.compress_idx += 1

            if self.compress_idx == self.n_simp_positions:
                self.simplex_step = 1

max_stars_repo_name: SimonBlanke/derivative-free-optimizers

max_issues_repo_name: SimonBlanke/derivative-free-optimizers

max_forks_repo_name: SimonBlanke/derivative-free-optimizers


text:
import cv2
import numpy as np
from read_config import read_config
from tqdm import tqdm
import accv_utils
import os

'''
    Function to load avi file from disc in required frame resolution
'''

def load_avi_into_nparray(fileName, frameHeight, frameWidth, startFrame, endFrame):
    print fileName
    video = cv2.VideoCapture(fileName)
    numFrames = 0
    debugMode = read_config('debugMode')

    arr = fileName.split('/')
    arrLen = len(arr)
    cacheName = ''
    cachePath = ''
    if arrLen == 5:
        cacheName = arr[3] + '_' + arr[4].split('.')[0] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 6:
        cacheName = arr[3] + '_' + arr[4] + '_' + arr[5].split('.')[0] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 3:
        cacheName = arr[2].split('.')[0] + '.dat'
        cachePath = '../cache/' + cacheName

    pbar = tqdm(total=(endFrame-startFrame+1))

    videoNP = np.empty((0, 1, frameHeight, frameWidth), dtype='float32')

    if os.path.isfile(cachePath):
        print('Restoring from cache')
        videoNP = accv_utils.restore_data(cachePath)
    else:
        while(video.isOpened()):
            ret, frame = video.read()
            if startFrame <= numFrames and endFrame >= numFrames:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                resized = cv2.resize(frame, (frameWidth, frameHeight), cv2.INTER_LINEAR)
                resized = resized.reshape(1, frameHeight, frameWidth)
                #if debugMode:
                #    print('resized shape:{}'.format(resized.shape))
                videoNP = np.append(videoNP, [resized], axis=0)
                #if debugMode:
                #    print('videoNP shape:{}'.format(videoNP.shape))
                pbar.update(1)
            if endFrame < numFrames:
                break
            numFrames += 1
        accv_utils.cache_data(videoNP, cachePath)

    pbar.close()
    print('Num frames loaded:{}'.format(endFrame - startFrame + 1))

    return videoNP

def detectFace(frame):
    #debugMode = read_config('debugMode')
    cascPath = '../input/cascades/haarcascades/haarcascade_frontalface_alt2.xml'
    faceCascade = cv2.CascadeClassifier(cascPath)
    faces = faceCascade.detectMultiScale(
        frame,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(30, 30)
    )
    if len(faces) == 1:
        face = faces[0]
        x = face[0]
        y = face[1]
        w = face[2]
        h = face[3]
        detectedFace = frame[y:y+h, x:x+w]
        #if debugMode:
        #    cv2.imshow('Face Detected', detectedFace)
        #    cv2.waitKey(0)
        #    cv2.destroyAllWindows()
        return detectedFace, 1
    else:
        return frame, len(faces)

def load_processed_avi_into_nparray(fileName, frameHeight, frameWidth, startFrame, endFrame):
    arr = fileName.split('/')
    arrLen = len(arr)
    if arrLen == 5:
        cacheName = arr[3] + '_' + arr[4].split('.')[0] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 6:
        cacheName = arr[3] + '_' + arr[4] + '_' + arr[5].split('.')[0] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 3:
        cacheName = arr[2].split('.')[0] + '.dat'
        cachePath = '../cache/' + cacheName

    video = cv2.VideoCapture(fileName)
    numFrames = 0
    debugMode = read_config('debugMode')

    videoNP = np.empty((0, 1, frameHeight, frameWidth), dtype='float32')

    if os.path.isfile(cachePath):
        print('Restoring from cache')
        videoNP = accv_utils.restore_data(cachePath)
    else:
	pbar = tqdm(total=(endFrame-startFrame+1))
        while (video.isOpened()):
            ret, frame = video.read()
            if startFrame <= numFrames and endFrame >= numFrames:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                frame, numFaces = detectFace(frame)
                resized = cv2.resize(frame, (frameWidth, frameHeight), cv2.INTER_LINEAR)
                resized = resized.reshape(1, frameHeight, frameWidth)
                # if debugMode:
                #    print('resized shape:{}'.format(resized.shape))
                videoNP = np.append(videoNP, [resized], axis=0)
                # if debugMode:
                #    print('videoNP shape:{}'.format(videoNP.shape))
                pbar.update(1)
            if endFrame < numFrames:
                break
            numFrames += 1
        pbar.close()
        print('Creating cache')
        accv_utils.cache_data(videoNP, cachePath)

    print('Num frames loaded:{}'.format(endFrame - startFrame + 1))
    return videoNP

def load_face_into_nparray(filePath, frameHeight, frameWidth, startFrame, endFrame):
    print('{}-{} to {}'.format(filePath, startFrame, endFrame))
    debugMode = read_config('debugMode')

    if not os.path.isdir(filePath):
        print('FilePath:{} does not exist'.format(filePath))
        return

    arr = filePath.split('/')
    arrLen = len(arr)
    cacheName = ''
    cachePath = ''
    if arrLen == 5:
        cacheName = 'face_' + arr[3] + '_' + arr[4] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 6:
        cacheName = 'face_' + arr[3] + '_' + arr[4] + '_' + arr[5] + '_' + str(startFrame) + '_' + str(endFrame) + '.dat'
        cachePath = '../cache/' + cacheName
    elif arrLen == 3:
        cacheName = arr[2].split('.')[0] + '.dat'
        cachePath = '../cache/' + cacheName

    pbar = tqdm(total=(endFrame-startFrame+1))
    videoNP = np.empty((0, 1, frameHeight, frameWidth), dtype='float32')
    if os.path.isfile(cachePath):
        print('Restoring from cache')
        videoNP = accv_utils.restore_data(cachePath)
    else:
        for i in range(startFrame, endFrame + 1):
            fileName = os.path.join(filePath, str(i) + '.png')
            frame = cv2.imread(fileName)
            if len(frame.shape) == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            # detected region is too big or small to be a face
            if frame.shape[0] > 420 and frame.shape[1] > 420 and frame.shape[0] < 150 and frame.shape[1] < 150:
                frame = np.zeros((250, 250), dtype='float32')

            resized = cv2.resize(frame, (frameWidth, frameHeight), cv2.INTER_LINEAR)
            resized = resized.reshape(1, frameHeight, frameWidth)
            #if debugMode:
            #    print('resized shape:{}'.format(resized.shape))
            videoNP = np.append(videoNP, [resized], axis=0)
            #if debugMode:
            #    print('videoNP shape:{}'.format(videoNP.shape))
            pbar.update(1)
        accv_utils.cache_data(videoNP, cachePath)
    pbar.close()
    print('Num frames loaded:{}'.format(endFrame - startFrame + 1))

    return videoNP

def load_frame(fileName, frameHeight, frameWidth, frameNumber):
    video = cv2.VideoCapture(fileName)
    numFrames = 0
    frameNP = np.empty((0, frameHeight, frameWidth))
    while(video.isOpened()):
        if numFrames == frameNumber:
            ret, frame = video.read()
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            frameNP = cv2.resize(frame, (frameWidth, frameHeight), cv2.INTER_LINEAR)
            break
        numFrames += 1

    return frameNP

'''
    TODO: Function to create jittered images for Data augmentation
    eg. images with rotation, stretches, crops, flips, ZCA whitening
'''

ext: py

max_stars_repo_name: SuriyaNitt/DDD

max_issues_repo_name: SuriyaNitt/DDD

max_forks_repo_name: SuriyaNitt/DDD

reason: import numpy


text:
import numpy as np
import matplotlib.pyplot as plt

def image_and_background(data, meta, n):

    intstart, subdata, submask, subbg = data.intstart, data.subdata, data.submask, data.subbg

    plt.figure(3301)
    plt.clf()
    plt.suptitle(str(intstart + n))
    plt.subplot(211)
    max = np.max(subdata[n] * submask[n])
    plt.imshow(subdata[n] * submask[n], origin='lower', aspect='auto', vmin=0, vmax=max / 10)
    # plt.imshow(subdata[n], origin='lower', aspect='auto', vmin=0, vmax=10000)
    plt.subplot(212)
    # plt.imshow(submask[i], origin='lower', aspect='auto', vmax=1)
    median = np.median(subbg[n])
    std = np.std(subbg[n])
    plt.imshow(subbg[n], origin='lower', aspect='auto', vmin=median - 3 * std, vmax=median + 3 * std)
    # plt.imshow(submask[n], origin='lower', aspect='auto', vmin=0, vmax=1)
    plt.savefig(meta.workdir + '/figs/fig3301-' + str(intstart + n) + '-Image+Background.png')
    # plt.pause(0.1)

def optimal_spectrum(data, meta, n):

    intstart, subnx, stdspec, optspec, opterr = data.intstart, meta.subnx, data.stdspec, data.optspec, data.opterr

    plt.figure(3302)
    plt.clf()
    plt.suptitle(str(intstart + n))
    plt.plot(range(subnx), stdspec[n], '-', color='C1', label='Std Spec')
    # plt.errorbar(range(subnx), stdspec[n], yerr=np.sqrt(stdvar[n]), fmt='-', color='C1', ecolor='C0', label='Std Spec')
    plt.errorbar(range(subnx), optspec[n], opterr[n], fmt='-', color='C2', ecolor='C2', label='Optimal Spec')
    plt.legend(loc='best')
    plt.savefig(meta.workdir + '/figs/fig3302-' + str(intstart + n) + '-Spectrum.png')
    # plt.pause(0.1)

def lc_nodriftcorr(meta, wave_1d, optspec):
    plt.figure(3101, figsize=(8, 8))  # ev.n_files/20.+0.8))
    plt.clf()
    wmin = wave_1d.min()
    wmax = wave_1d.max()
    n_int, nx = optspec.shape
    # iwmin       = np.where(ev.wave[src_ypos]>wmin)[0][0]
    # iwmax       = np.where(ev.wave[src_ypos]>wmax)[0][0]
    vmin = 0.97
    vmax = 1.03
    # normspec    = np.mean(ev.optspec,axis=1)/np.mean(ev.optspec[ev.inormspec[0]:ev.inormspec[1]],axis=(0,1))
    normspec = optspec / np.mean(optspec, axis=0)
    plt.imshow(normspec, origin='lower', aspect='auto', extent=[wmin, wmax, 0, n_int], vmin=vmin, vmax=vmax,
               cmap=plt.cm.RdYlBu_r)
    ediff = np.zeros(n_int)
    for m in range(n_int):
        ediff[m] = 1e6 * np.median(np.abs(np.ediff1d(normspec[m])))
        # plt.scatter(ev.wave[src_ypos], np.zeros(nx)+m, c=normspec[m],
        #             s=14,linewidths=0,vmin=vmin,vmax=vmax,marker='s',cmap=plt.cm.RdYlBu_r)
    plt.title("MAD = " + str(np.round(np.mean(ediff), 0).astype(int)) + " ppm")
    # plt.xlim(wmin,wmax)
    # plt.ylim(0,n_int)
    plt.ylabel('Integration Number')
    plt.xlabel(r'Wavelength ($\mu m$)')
    plt.colorbar(label='Normalized Flux')
    plt.tight_layout()
    plt.savefig(meta.workdir + '/figs/fig3101-2D_LC.png')

max_stars_repo_name: ladsantos/Eureka

max_issues_repo_name: ladsantos/Eureka

max_forks_repo_name: ladsantos/Eureka


text:
"""Geometric harmonics module.

This module implements out-of-sample evaluation of functions using the
Geometric Harmonics method introduced in:

Coifman, R. R., & Lafon, S. (2006). Geometric harmonics: A novel tool for
multiscale out-of-sample extension of empirical functions. Applied and
Computational Harmonic Analysis, 21(1), 31–52. DOI:10.1016/j.acha.2005.07.005

"""

__all__ = ['GeometricHarmonicsInterpolator']

from typing import Optional, Dict

import numpy as np
import scipy.spatial
from scipy.interpolate.interpnd import (NDInterpolatorBase,
                                        _ndim_coords_from_arrays)

from diffusion_maps.diffusion_maps import DiffusionMaps

class GeometricHarmonicsInterpolator(NDInterpolatorBase):
    """Geometric Harmonics interpolator.

    """
    def __init__(self, points: np.array, values: np.array, epsilon: float,
                 diffusion_maps_options: Optional[Dict] = None,
                 diffusion_maps: Optional[DiffusionMaps] = None) -> None:
        NDInterpolatorBase.__init__(self, points, values,
                                    need_contiguous=False, need_values=True)
        self.epsilon = epsilon
        if diffusion_maps:
            self.diffusion_maps = diffusion_maps
        else:
            if not diffusion_maps_options:
                diffusion_maps_options = dict()
            diffusion_maps_options['normalize_kernel'] = False
            self.diffusion_maps = DiffusionMaps(self.points, epsilon,
                                                **diffusion_maps_options)

    def __call__(self, *args):
        """Evaluate interpolator at the given points.

        """
        xi = _ndim_coords_from_arrays(args, ndim=self.points.shape[1])
        xi = self._check_call_shape(xi)

        kdtree = scipy.spatial.cKDTree(xi)
        dmaps_kdtree = self.diffusion_maps._kdtree
        radius = self.diffusion_maps._cut_off

        ew = self.diffusion_maps.eigenvalues
        ev = self.diffusion_maps.eigenvectors
        aux = ev.T @ np.diag(1.0 / ew) @ ev @ self.values

        distance_matrix \
            = kdtree.sparse_distance_matrix(dmaps_kdtree, radius,
                                            output_type='coo_matrix')
        kernel_matrix \
            = self.diffusion_maps.compute_kernel_matrix(distance_matrix)

        return np.squeeze(kernel_matrix @ aux)

    def gradient(self, *args):
        """Evaluate gradient of interpolator at the given points.

        """
        ps = self.points
        xi = _ndim_coords_from_arrays(args, ndim=ps.shape[1])
        xi = self._check_call_shape(xi)

        kdtree = scipy.spatial.cKDTree(xi)
        dmaps_kdtree = self.diffusion_maps._kdtree
        radius = self.diffusion_maps._cut_off

        grad = np.zeros_like(xi)

        distance_matrix \
            = kdtree.sparse_distance_matrix(dmaps_kdtree, radius,
                                            output_type='coo_matrix')
        kernel_matrix \
            = self.diffusion_maps.compute_kernel_matrix(distance_matrix)

        I, J = xi.shape[0], ps.shape[0]

        V = np.zeros((ps.shape[1], I, J))
        for i in range(I):
            for j in range(J):
                V[:, i, j] = xi[i, :] - ps[j, :]

        for i in range(I):
            ri = np.squeeze(kernel_matrix.getrow(i).toarray())
            ki = scipy.sparse.spdiags(ri, 0, J, J)
            ki_psi = ki @ np.squeeze(self.values)
            grad[i, :] = V[:, i, :] @ ki_psi

        return grad

max_stars_repo_name: vishalbelsare/diffusion-maps

max_issues_repo_name: vishalbelsare/diffusion-maps

max_forks_repo_name: vishalbelsare/diffusion-maps

reason: import numpy,import scipy,from scipy


text:
# relabel_imagenet
# Copyright (c) 2021-present NAVER Corp.
# Apache License v2.0

import torch.optim as optim
from adamp import AdamP
import numpy as np

def set_init_lr(param_groups):
    for group in param_groups:
        group['init_lr'] = group['lr']

def load_optimizer(args, param_group):

    if args.optim.optimizer.lower() == 'sgd':
        optimizer = optim.SGD(
            param_group,
            args.optim.lr.init,
            momentum=args.optim.momentum,
            weight_decay=args.optim.wd.base,
            nesterov=args.optim.nesterov,
        )
    elif args.optim.optimizer.lower() == 'adamp':
        optimizer = AdamP(
            param_group,
            args.optim.lr.init,
            betas=(args.optim.momentum, 0.999),
            weight_decay=args.optim.wd.base,
            nesterov=args.optim.nesterov,
        )
    else:
        raise ValueError("Unknown optimizer : {}".format(args.optim.optimizer))

    set_init_lr(optimizer.param_groups)

    return optimizer

def warmup_learnig_rate(init_lr, warmup_lr, warmup_epochs,
                        iteration, epoch, dataset_len):
    lr = warmup_lr + (init_lr - warmup_lr) * \
         float(iteration + epoch * dataset_len) / (warmup_epochs * dataset_len)
    return lr

def adjust_learning_rate_cosine(epoch, iteration, dataset_len,
                                epochs, warmup_epochs):
    total_iter = (epochs - warmup_epochs) * dataset_len
    current_iter = iteration + (epoch - warmup_epochs) * dataset_len

    lr = 1 / 2 * (np.cos(np.pi * current_iter / total_iter) + 1)

    return lr

def adjust_learning_rate_linear(epoch, iteration, dataset_len,
                                epochs, warmup_epochs):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    total_iter = (epochs - warmup_epochs) * dataset_len
    current_iter = iteration + (epoch - warmup_epochs) * dataset_len

    lr = 1 - current_iter / total_iter

    return lr

def adjust_learning_rate_default(epoch, iteration, dataset_len,
                                 epochs, warmup_epochs):
    """
    LR schedule that should yield 76% converged accuracy
    with batch size 256"""
    factor = epoch // 30

    lr = 0.1 ** factor

    return lr

def adjust_learning_rate(optimizer, epoch, iteration, lr_decay_type,
                         epochs, train_len, warmup_lr, warmup_epochs):
    if epoch < warmup_epochs:
        for param_group in optimizer.param_groups:
            lr = warmup_learnig_rate(init_lr=param_group['init_lr'],
                                     warmup_lr=warmup_lr,
                                     warmup_epochs=warmup_epochs,
                                     epoch=epoch,
                                     iteration=iteration,
                                     dataset_len=train_len)
            param_group['lr'] = lr
    else:
        if lr_decay_type == 'cos':
            lr_function = adjust_learning_rate_cosine
        elif lr_decay_type == 'linear':
            lr_function = adjust_learning_rate_linear
        elif lr_decay_type == 'default':
            lr_function = adjust_learning_rate_default
        else:
            raise ValueError("Unknown lr decay type {}."
                             .format(lr_decay_type))

        lr_factor = lr_function(epoch=epoch,
                                iteration=iteration,
                                dataset_len=train_len,
                                epochs=epochs,
                                warmup_epochs=warmup_epochs)

        for param_group in optimizer.param_groups:
            param_group['lr'] = param_group['init_lr'] * lr_factor

    return optimizer.param_groups[0]['lr']


text:
# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2004-2018 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/
"""Widget displaying curves legends and allowing to operate on curves.

This widget is meant to work with :class:`PlotWindow`.
"""

__authors__ = ["V.A. Sole", "T. Rueter", "T. Vincent"]
__license__ = "MIT"
__data__ = "16/10/2017"

import logging
import weakref

import numpy

from .. import qt, colors
from . import items

_logger = logging.getLogger(__name__)

# Build all symbols
# Courtesy of the pyqtgraph project
Symbols = dict([(name, qt.QPainterPath())
                for name in ['o', 's', 't', 'd', '+', 'x', '.', ',']])
Symbols['o'].addEllipse(qt.QRectF(.1, .1, .8, .8))
Symbols['.'].addEllipse(qt.QRectF(.3, .3, .4, .4))
Symbols[','].addEllipse(qt.QRectF(.4, .4, .2, .2))
Symbols['s'].addRect(qt.QRectF(.1, .1, .8, .8))

coords = {
    't': [(0.5, 0.), (.1, .8), (.9, .8)],
    'd': [(0.1, 0.5), (0.5, 0.), (0.9, 0.5), (0.5, 1.)],
    '+': [(0.0, 0.40), (0.40, 0.40), (0.40, 0.), (0.60, 0.),
          (0.60, 0.40), (1., 0.40), (1., 0.60), (0.60, 0.60),
          (0.60, 1.), (0.40, 1.), (0.40, 0.60), (0., 0.60)],
    'x': [(0.0, 0.40), (0.40, 0.40), (0.40, 0.), (0.60, 0.),
          (0.60, 0.40), (1., 0.40), (1., 0.60), (0.60, 0.60),
          (0.60, 1.), (0.40, 1.), (0.40, 0.60), (0., 0.60)]
}
for s, c in coords.items():
    Symbols[s].moveTo(*c[0])
    for x, y in c[1:]:
        Symbols[s].lineTo(x, y)
    Symbols[s].closeSubpath()
tr = qt.QTransform()
tr.rotate(45)
Symbols['x'].translate(qt.QPointF(-0.5, -0.5))
Symbols['x'] = tr.map(Symbols['x'])
Symbols['x'].translate(qt.QPointF(0.5, 0.5))

NoSymbols = (None, 'None', 'none', '', ' ')
"""List of values resulting in no symbol being displayed for a curve"""

LineStyles = {
    None: qt.Qt.NoPen,
    'None': qt.Qt.NoPen,
    'none': qt.Qt.NoPen,
    '': qt.Qt.NoPen,
    ' ': qt.Qt.NoPen,
    '-': qt.Qt.SolidLine,
    '--': qt.Qt.DashLine,
    ':': qt.Qt.DotLine,
    '-.': qt.Qt.DashDotLine
}
"""Conversion from matplotlib-like linestyle to Qt"""

NoLineStyle = (None, 'None', 'none', '', ' ')
"""List of style values resulting in no line being displayed for a curve"""

class LegendIcon(qt.QWidget):
    """Object displaying a curve linestyle and symbol.

    :param QWidget parent: See :class:`QWidget`
    :param Union[~silx.gui.plot.items.Curve,None] curve:
        Curve with which to synchronize
    """

    def __init__(self, parent=None, curve=None):
        super(LegendIcon, self).__init__(parent)
        self._curveRef = None

        # Visibilities
        self.showLine = True
        self.showSymbol = True

        # Line attributes
        self.lineStyle = qt.Qt.NoPen
        self.lineWidth = 1.
        self.lineColor = qt.Qt.green

        self.symbol = ''
        # Symbol attributes
        self.symbolStyle = qt.Qt.SolidPattern
        self.symbolColor = qt.Qt.green
        self.symbolOutlineBrush = qt.QBrush(qt.Qt.white)

        # Control widget size: sizeHint "is the only acceptable
        # alternative, so the widget can never grow or shrink"
        # (c.f. Qt Doc, enum QSizePolicy::Policy)
        self.setSizePolicy(qt.QSizePolicy.Fixed,
                           qt.QSizePolicy.Fixed)

        self.setCurve(curve)

    def sizeHint(self):
        return qt.QSize(50, 15)

    # Synchronize with a curve

    def getCurve(self):
        """Returns curve associated to this widget

        :rtype: Union[~silx.gui.plot.items.Curve,None]
        """
        return None if self._curveRef is None else self._curveRef()

    def setCurve(self, curve):
        """Set the curve with which to synchronize this widget.

        :param curve: Union[~silx.gui.plot.items.Curve,None]
        """
        assert curve is None or isinstance(curve, items.Curve)

        previousCurve = self.getCurve()
        if curve == previousCurve:
            return

        if previousCurve is not None:
            previousCurve.sigItemChanged.disconnect(self._curveChanged)

        self._curveRef = None if curve is None else weakref.ref(curve)

        if curve is not None:
            curve.sigItemChanged.connect(self._curveChanged)

        self._update()

    def _update(self):
        """Update widget according to current curve state.
        """
        curve = self.getCurve()
        if curve is None:
            _logger.error('Curve no more exists')
            self.setEnabled(False)
            return

        style = curve.getCurrentStyle()

        self.setEnabled(curve.isVisible())
        self.setSymbol(style.getSymbol())
        self.setLineWidth(style.getLineWidth())
        self.setLineStyle(style.getLineStyle())

        color = style.getColor()
        if numpy.array(color, copy=False).ndim != 1:
            # array of colors, use transparent black
            color = 0., 0., 0., 0.
        color = colors.rgba(color)  # Make sure it is float in [0, 1]
        alpha = curve.getAlpha()
        color = qt.QColor.fromRgbF(
            color[0], color[1], color[2], color[3] * alpha)
        self.setLineColor(color)
        self.setSymbolColor(color)
        self.update()  # TODO this should not be needed

    def _curveChanged(self, event):
        """Handle update of curve item

        :param event: Kind of change
        """
        if event in (items.ItemChangedType.VISIBLE,
                     items.ItemChangedType.SYMBOL,
                     items.ItemChangedType.SYMBOL_SIZE,
                     items.ItemChangedType.LINE_WIDTH,
                     items.ItemChangedType.LINE_STYLE,
                     items.ItemChangedType.COLOR,
                     items.ItemChangedType.ALPHA,
                     items.ItemChangedType.HIGHLIGHTED,
                     items.ItemChangedType.HIGHLIGHTED_STYLE):
            self._update()

    # Modify Symbol
    def setSymbol(self, symbol):
        symbol = str(symbol)
        if symbol not in NoSymbols:
            if symbol not in Symbols:
                raise ValueError("Unknown symbol: <%s>" % symbol)
        self.symbol = symbol
        # self.update() after set...?
        # Does not seem necessary

    def setSymbolColor(self, color):
        """
        :param color: determines the symbol color
        :type style: qt.QColor
        """
        self.symbolColor = qt.QColor(color)

    # Modify Line

    def setLineColor(self, color):
        self.lineColor = qt.QColor(color)

    def setLineWidth(self, width):
        self.lineWidth = float(width)

    def setLineStyle(self, style):
        """Set the linestyle.

        Possible line styles:

        - '', ' ', 'None': No line
        - '-': solid
        - '--': dashed
        - ':': dotted
        - '-.': dash and dot

        :param str style: The linestyle to use
        """
        if style not in LineStyles:
            raise ValueError('Unknown style: %s', style)
        self.lineStyle = LineStyles[style]

    # Paint

    def paintEvent(self, event):
        """
        :param event: event
        :type event: QPaintEvent
        """
        painter = qt.QPainter(self)
        self.paint(painter, event.rect(), self.palette())

    def paint(self, painter, rect, palette):
        painter.save()
        painter.setRenderHint(qt.QPainter.Antialiasing)
        # Scale painter to the icon height
        # current -> width = 2.5, height = 1.0
        scale = float(self.height())
        ratio = float(self.width()) / scale
        painter.scale(scale,
                      scale)
        symbolOffset = qt.QPointF(.5 * (ratio - 1.), 0.)
        # Determine and scale offset
        offset = qt.QPointF(float(rect.left()) / scale, float(rect.top()) / scale)

        # Override color when disabled
        if self.isEnabled():
            overrideColor = None
        else:
            overrideColor = palette.color(qt.QPalette.Disabled,
                                          qt.QPalette.WindowText)

        # Draw BG rectangle (for debugging)
        # bottomRight = qt.QPointF(
        #    float(rect.right())/scale,
        #    float(rect.bottom())/scale)
        # painter.fillRect(qt.QRectF(offset, bottomRight),
        #                 qt.QBrush(qt.Qt.green))
        llist = []
        if self.showLine:
            linePath = qt.QPainterPath()
            linePath.moveTo(0., 0.5)
            linePath.lineTo(ratio, 0.5)
            # linePath.lineTo(2.5, 0.5)
            lineBrush = qt.QBrush(
                self.lineColor if overrideColor is None else overrideColor)
            linePen = qt.QPen(
                lineBrush,
                (self.lineWidth / self.height()),
                self.lineStyle,
                qt.Qt.FlatCap
            )
            llist.append((linePath, linePen, lineBrush))
        if (self.showSymbol and len(self.symbol) and
                self.symbol not in NoSymbols):
            # PITFALL ahead: Let this be a warning to others
            # symbolPath = Symbols[self.symbol]
            # Copy before translate! Dict is a mutable type
            symbolPath = qt.QPainterPath(Symbols[self.symbol])
            symbolPath.translate(symbolOffset)
            symbolBrush = qt.QBrush(
                self.symbolColor if overrideColor is None else overrideColor,
                self.symbolStyle)
            symbolPen = qt.QPen(
                self.symbolOutlineBrush,  # Brush
                1. / self.height(),       # Width
                qt.Qt.SolidLine           # Style
            )
            llist.append((symbolPath,
                          symbolPen,
                          symbolBrush))
        # Draw
        for path, pen, brush in llist:
            path.translate(offset)
            painter.setPen(pen)
            painter.setBrush(brush)
            painter.drawPath(path)
        painter.restore()

class LegendModel(qt.QAbstractListModel):
    """Data model of curve legends.

    It holds the information of the curve:

    - color
    - line width
    - line style
    - visibility of the lines
    - symbol
    - visibility of the symbols
    """
    iconColorRole = qt.Qt.UserRole + 0
    iconLineWidthRole = qt.Qt.UserRole + 1
    iconLineStyleRole = qt.Qt.UserRole + 2
    showLineRole = qt.Qt.UserRole + 3
    iconSymbolRole = qt.Qt.UserRole + 4
    showSymbolRole = qt.Qt.UserRole + 5

    def __init__(self, legendList=None, parent=None):
        super(LegendModel, self).__init__(parent)
        if legendList is None:
            legendList = []
        self.legendList = []
        self.insertLegendList(0, legendList)
        self._palette = qt.QPalette()

    def __getitem__(self, idx):
        if idx >= len(self.legendList):
            raise IndexError('list index out of range')
        return self.legendList[idx]

    def rowCount(self, modelIndex=None):
        return len(self.legendList)

    def flags(self, index):
        return (qt.Qt.ItemIsEditable |
                qt.Qt.ItemIsEnabled |
                qt.Qt.ItemIsSelectable)

    def data(self, modelIndex, role):
        if modelIndex.isValid:
            idx = modelIndex.row()
        else:
            return None
        if idx >= len(self.legendList):
            raise IndexError('list index out of range')

        item = self.legendList[idx]
        isActive = item[1].get("active", False)
        if role == qt.Qt.DisplayRole:
            # Data to be rendered in the form of text
            legend = str(item[0])
            return legend
        elif role == qt.Qt.SizeHintRole:
            # size = qt.QSize(200,50)
            _logger.warning('LegendModel -- size hint role not implemented')
            return qt.QSize()
        elif role == qt.Qt.TextAlignmentRole:
            alignment = qt.Qt.AlignVCenter | qt.Qt.AlignLeft
            return alignment
        elif role == qt.Qt.BackgroundRole:
            # Background color, must be QBrush
            if isActive:
                brush = self._palette.brush(qt.QPalette.Normal, qt.QPalette.Highlight)
            elif idx % 2:
                brush = qt.QBrush(qt.QColor(240, 240, 240))
            else:
                brush = qt.QBrush(qt.Qt.white)
            return brush
        elif role == qt.Qt.ForegroundRole:
            # ForegroundRole color, must be QBrush
            if isActive:
                brush = self._palette.brush(qt.QPalette.Normal, qt.QPalette.HighlightedText)
            else:
                brush = self._palette.brush(qt.QPalette.Normal, qt.QPalette.WindowText)
            return brush
        elif role == qt.Qt.CheckStateRole:
            return bool(item[2])  # item[2] == True
        elif role == qt.Qt.ToolTipRole or role == qt.Qt.StatusTipRole:
            return ''
        elif role == self.iconColorRole:
            return item[1]['color']
        elif role == self.iconLineWidthRole:
            return item[1]['linewidth']
        elif role == self.iconLineStyleRole:
            return item[1]['linestyle']
        elif role == self.iconSymbolRole:
            return item[1]['symbol']
        elif role == self.showLineRole:
            return item[3]
        elif role == self.showSymbolRole:
            return item[4]
        else:
            _logger.info('Unkown role requested: %s', str(role))
            return None

    def setData(self, modelIndex, value, role):
        if modelIndex.isValid:
            idx = modelIndex.row()
        else:
            return None
        if idx >= len(self.legendList):
            # raise IndexError('list index out of range')
            _logger.warning(
                'setData -- List index out of range, idx: %d', idx)
            return None

        item = self.legendList[idx]
        try:
            if role == qt.Qt.DisplayRole:
                # Set legend
                item[0] = str(value)
            elif role == self.iconColorRole:
                item[1]['color'] = qt.QColor(value)
            elif role == self.iconLineWidthRole:
                item[1]['linewidth'] = int(value)
            elif role == self.iconLineStyleRole:
                item[1]['linestyle'] = str(value)
            elif role == self.iconSymbolRole:
                item[1]['symbol'] = str(value)
            elif role == qt.Qt.CheckStateRole:
                item[2] = value
            elif role == self.showLineRole:
                item[3] = value
            elif role == self.showSymbolRole:
                item[4] = value
        except ValueError:
            _logger.warning('Conversion failed:\n\tvalue: %s\n\trole: %s',
                            str(value), str(role))
        # Can that be right? Read docs again..
        self.dataChanged.emit(modelIndex, modelIndex)
        return True

    def insertLegendList(self, row, llist):
        """
        :param int row: Determines after which row the items are inserted
        :param llist: Carries the new legend information
        :type llist: List
        """
        modelIndex = self.createIndex(row, 0)
        count = len(llist)
        super(LegendModel, self).beginInsertRows(modelIndex,
                                                 row,
                                                 row + count)
        head = self.legendList[0:row]
        tail = self.legendList[row:]
        new = []
        for (legend, icon) in llist:
            linestyle = icon.get('linestyle', None)
            if linestyle in NoLineStyle:
                # Curve had no line, give it one and hide it
                # So when toggle line, it will display a solid line
                showLine = False
                icon['linestyle'] = '-'
            else:
                showLine = True

            symbol = icon.get('symbol', None)
            if symbol in NoSymbols:
                # Curve had no symbol, give it one and hide it
                # So when toggle symbol, it will display 'o'
                showSymbol = False
                icon['symbol'] = 'o'
            else:
                showSymbol = True

            selected = icon.get('selected', True)
            item = [legend,
                    icon,
                    selected,
                    showLine,
                    showSymbol]
            new.append(item)
        self.legendList = head + new + tail
        super(LegendModel, self).endInsertRows()
        return True

    def insertRows(self, row, count, modelIndex=qt.QModelIndex()):
        raise NotImplementedError('Use LegendModel.insertLegendList instead')

    def removeRow(self, row):
        return self.removeRows(row, 1)

    def removeRows(self, row, count, modelIndex=qt.QModelIndex()):
        length = len(self.legendList)
        if length == 0:
            # Nothing to do..
            return True
        if row < 0 or row >= length:
            raise IndexError('Index out of range -- ' +
                             'idx: %d, len: %d' % (row, length))
        if count == 0:
            return False
        super(LegendModel, self).beginRemoveRows(modelIndex,
                                                 row,
                                                 row + count)
        del(self.legendList[row:row + count])
        super(LegendModel, self).endRemoveRows()
        return True

    def setEditor(self, event, editor):
        """
        :param str event: String that identifies the editor
        :param editor: Widget used to change data in the underlying model
        :type editor: QWidget
        """
        if event not in self.eventList:
            raise ValueError('setEditor -- Event must be in %s' %
                             str(self.eventList))
        self.editorDict[event] = editor

class LegendListItemWidget(qt.QItemDelegate):
    """Object displaying a single item (i.e., a row) in the list."""

    # Notice: LegendListItem does NOT inherit
    # from QObject, it cannot emit signals!

    def __init__(self, parent=None, itemType=0):
        super(LegendListItemWidget, self).__init__(parent)

        # Dictionary to render checkboxes
        self.cbDict = {}
        self.labelDict = {}
        self.iconDict = {}

        # Keep checkbox and legend to get sizeHint
        self.checkbox = qt.QCheckBox()
        self.legend = qt.QLabel()
        self.icon = LegendIcon()

        # Context Menu and Editors
        self.contextMenu = None

    def paint(self, painter, option, modelIndex):
        """
        Here be docs..

        :param QPainter painter:
        :param QStyleOptionViewItem option:
        :param QModelIndex modelIndex:
        """
        painter.save()
        rect = option.rect

        # Calculate the icon rectangle
        iconSize = self.icon.sizeHint()
        # Calculate icon position
        x = rect.left() + 2
        y = rect.top() + int(.5 * (rect.height() - iconSize.height()))
        iconRect = qt.QRect(qt.QPoint(x, y), iconSize)

        # Calculate label rectangle
        legendSize = qt.QSize(rect.width() - iconSize.width() - 30,
                              rect.height())
        # Calculate label position
        x = rect.left() + iconRect.width()
        y = rect.top()
        labelRect = qt.QRect(qt.QPoint(x, y), legendSize)
        labelRect.translate(qt.QPoint(10, 0))

        # Calculate the checkbox rectangle
        x = rect.right() - 30
        y = rect.top()
        chBoxRect = qt.QRect(qt.QPoint(x, y), rect.bottomRight())

        # Remember the rectangles
        idx = modelIndex.row()
        self.cbDict[idx] = chBoxRect
        self.iconDict[idx] = iconRect
        self.labelDict[idx] = labelRect

        # Draw background first!
        if option.state & qt.QStyle.State_MouseOver:
            backgroundBrush = option.palette.highlight()
        else:
            backgroundBrush = modelIndex.data(qt.Qt.BackgroundRole)
        painter.fillRect(rect, backgroundBrush)

        # Draw label
        legendText = modelIndex.data(qt.Qt.DisplayRole)
        textBrush = modelIndex.data(qt.Qt.ForegroundRole)
        textAlign = modelIndex.data(qt.Qt.TextAlignmentRole)
        painter.setBrush(textBrush)
        painter.setFont(self.legend.font())
        painter.setPen(textBrush.color())
        painter.drawText(labelRect, textAlign, legendText)

        # Draw icon
        iconColor = modelIndex.data(LegendModel.iconColorRole)
        iconLineWidth = modelIndex.data(LegendModel.iconLineWidthRole)
        iconLineStyle = modelIndex.data(LegendModel.iconLineStyleRole)
        iconSymbol = modelIndex.data(LegendModel.iconSymbolRole)
        icon = LegendIcon()
        icon.resize(iconRect.size())
        icon.move(iconRect.topRight())
        icon.showSymbol = modelIndex.data(LegendModel.showSymbolRole)
        icon.showLine = modelIndex.data(LegendModel.showLineRole)
        icon.setSymbolColor(iconColor)
        icon.setLineColor(iconColor)
        icon.setLineWidth(iconLineWidth)
        icon.setLineStyle(iconLineStyle)
        icon.setSymbol(iconSymbol)
        icon.symbolOutlineBrush = backgroundBrush
        icon.paint(painter, iconRect, option.palette)

        # Draw the checkbox
        if modelIndex.data(qt.Qt.CheckStateRole):
            checkState = qt.Qt.Checked
        else:
            checkState = qt.Qt.Unchecked

        self.drawCheck(
            painter, qt.QStyleOptionViewItem(), chBoxRect, checkState)

        painter.restore()

    def editorEvent(self, event, model, option, modelIndex):
        # From the docs:
        # Mouse events are sent to editorEvent()
        # even if they don't start editing of the item.
        if event.button() == qt.Qt.RightButton and self.contextMenu:
            self.contextMenu.exec_(event.globalPos(), modelIndex)
            return True
        elif event.button() == qt.Qt.LeftButton:
            # Check if checkbox was clicked
            idx = modelIndex.row()
            cbRect = self.cbDict[idx]
            if cbRect.contains(event.pos()):
                # Toggle checkbox
                model.setData(modelIndex,
                              not modelIndex.data(qt.Qt.CheckStateRole),
                              qt.Qt.CheckStateRole)
            event.ignore()
            return True
        else:
            return super(LegendListItemWidget, self).editorEvent(
                event, model, option, modelIndex)

    def createEditor(self, parent, option, idx):
        _logger.info('### Editor request ###')

    def sizeHint(self, option, idx):
        # return qt.QSize(68,24)
        iconSize = self.icon.sizeHint()
        legendSize = self.legend.sizeHint()
        checkboxSize = self.checkbox.sizeHint()
        height = max([iconSize.height(),
                      legendSize.height(),
                      checkboxSize.height()]) + 4
        width = iconSize.width() + legendSize.width() + checkboxSize.width()
        return qt.QSize(width, height)

class LegendListView(qt.QListView):
    """Widget displaying a list of curve legends, line style and symbol."""

    sigLegendSignal = qt.Signal(object)
    """Signal emitting a dict when an action is triggered by the user."""

    __mouseClickedEvent = 'mouseClicked'
    __checkBoxClickedEvent = 'checkBoxClicked'
    __legendClickedEvent = 'legendClicked'

    def __init__(self, parent=None, model=None, contextMenu=None):
        super(LegendListView, self).__init__(parent)
        self.__lastButton = None
        self.__lastClickPos = None
        self.__lastModelIdx = None
        # Set default delegate
        self.setItemDelegate(LegendListItemWidget())
        # Set default editors
        # self.setSizePolicy(qt.QSizePolicy.MinimumExpanding,
        #                    qt.QSizePolicy.MinimumExpanding)
        # Set edit triggers by hand using self.edit(QModelIndex)
        # in mousePressEvent (better to control than signals)
        self.setEditTriggers(qt.QAbstractItemView.NoEditTriggers)

        # Control layout
        # self.setBatchSize(2)
        # self.setLayoutMode(qt.QListView.Batched)
        # self.setFlow(qt.QListView.LeftToRight)

        # Control selection
        self.setSelectionMode(qt.QAbstractItemView.NoSelection)

        if model is None:
            model = LegendModel(parent=self)
        self.setModel(model)
        self.setContextMenu(contextMenu)

    def setLegendList(self, legendList, row=None):
        self.clear()
        if row is None:
            row = 0
        model = self.model()
        model.insertLegendList(row, legendList)
        _logger.debug('LegendListView.setLegendList(legendList) finished')

    def clear(self):
        model = self.model()
        model.removeRows(0, model.rowCount())
        _logger.debug('LegendListView.clear() finished')

    def setContextMenu(self, contextMenu=None):
        delegate = self.itemDelegate()
        if isinstance(delegate, LegendListItemWidget) and self.model():
            if contextMenu is None:
                delegate.contextMenu = LegendListContextMenu(self.model())
                delegate.contextMenu.sigContextMenu.connect(
                    self._contextMenuSlot)
            else:
                delegate.contextMenu = contextMenu

    def __getitem__(self, idx):
        model = self.model()
        try:
            item = model[idx]
        except ValueError:
            item = None
        return item

    def _contextMenuSlot(self, ddict):
        self.sigLegendSignal.emit(ddict)

    def mousePressEvent(self, event):
        self.__lastButton = event.button()
        self.__lastPosition = event.pos()
        super(LegendListView, self).mousePressEvent(event)
        # call _handleMouseClick after editing was handled
        # If right click (context menu) is aborted, no
        # signal is emitted..
        self._handleMouseClick(self.indexAt(self.__lastPosition))

    def mouseDoubleClickEvent(self, event):
        self.__lastButton = event.button()
        self.__lastPosition = event.pos()
        super(LegendListView, self).mouseDoubleClickEvent(event)
        # call _handleMouseClick after editing was handled
        # If right click (context menu) is aborted, no
        # signal is emitted..
        self._handleMouseClick(self.indexAt(self.__lastPosition))

    def mouseMoveEvent(self, event):
        # LegendListView.mouseMoveEvent is overwritten
        # to suppress unwanted behavior in the delegate.
        pass

    def mouseReleaseEvent(self, event):
        # LegendListView.mouseReleaseEvent is overwritten
        # to subpress unwanted behavior in the delegate.
        pass

    def _handleMouseClick(self, modelIndex):
        """
        Distinguish between mouse click on Legend
        and mouse click on CheckBox by setting the
        currentCheckState attribute in LegendListItem.

        Emits signal sigLegendSignal(ddict)

        :param QModelIndex modelIndex: index of the clicked item
        """
        _logger.debug('self._handleMouseClick called')
        if self.__lastButton not in [qt.Qt.LeftButton,
                                     qt.Qt.RightButton]:
            return
        if not modelIndex.isValid():
            _logger.debug('_handleMouseClick -- Invalid QModelIndex')
            return
        # model = self.model()
        idx = modelIndex.row()

        delegate = self.itemDelegate()
        cbClicked = False
        if isinstance(delegate, LegendListItemWidget):
            for cbRect in delegate.cbDict.values():
                if cbRect.contains(self.__lastPosition):
                    cbClicked = True
                    break

        # TODO: Check for doubleclicks on legend/icon and spawn editors

        ddict = {
            'legend': str(modelIndex.data(qt.Qt.DisplayRole)),
            'icon': {
                'linewidth': str(modelIndex.data(
                    LegendModel.iconLineWidthRole)),
                'linestyle': str(modelIndex.data(
                    LegendModel.iconLineStyleRole)),
                'symbol': str(modelIndex.data(LegendModel.iconSymbolRole))
            },
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data())
        }
        if self.__lastButton == qt.Qt.RightButton:
            _logger.debug('Right clicked')
            ddict['button'] = "right"
            ddict['event'] = self.__mouseClickedEvent
        elif cbClicked:
            _logger.debug('CheckBox clicked')
            ddict['button'] = "left"
            ddict['event'] = self.__checkBoxClickedEvent
        else:
            _logger.debug('Legend clicked')
            ddict['button'] = "left"
            ddict['event'] = self.__legendClickedEvent
        _logger.debug('  idx: %d\n  ddict: %s', idx, str(ddict))
        self.sigLegendSignal.emit(ddict)

class LegendListContextMenu(qt.QMenu):
    """Contextual menu associated to items in a :class:`LegendListView`."""

    sigContextMenu = qt.Signal(object)
    """Signal emitting a dict upon contextual menu actions."""

    def __init__(self, model):
        super(LegendListContextMenu, self).__init__(parent=None)
        self.model = model

        self.addAction('Set Active', self.setActiveAction)
        self.addAction('Map to left', self.mapToLeftAction)
        self.addAction('Map to right', self.mapToRightAction)

        self._pointsAction = self.addAction(
            'Points', self.togglePointsAction)
        self._pointsAction.setCheckable(True)

        self._linesAction = self.addAction('Lines', self.toggleLinesAction)
        self._linesAction.setCheckable(True)

        self.addAction('Remove curve', self.removeItemAction)
        self.addAction('Rename curve', self.renameItemAction)

    def exec_(self, pos, idx):
        self.__currentIdx = idx

        # Set checkable action state
        modelIndex = self.currentIdx()
        self._pointsAction.setChecked(
            modelIndex.data(LegendModel.showSymbolRole))
        self._linesAction.setChecked(
            modelIndex.data(LegendModel.showLineRole))

        super(LegendListContextMenu, self).popup(pos)

    def currentIdx(self):
        return self.__currentIdx

    def mapToLeftAction(self):
        _logger.debug('LegendListContextMenu.mapToLeftAction called')
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
            'event': "mapToLeft"
        }
        self.sigContextMenu.emit(ddict)

    def mapToRightAction(self):
        _logger.debug('LegendListContextMenu.mapToRightAction called')
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
            'event': "mapToRight"
        }
        self.sigContextMenu.emit(ddict)

    def removeItemAction(self):
        _logger.debug('LegendListContextMenu.removeCurveAction called')
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
            'event': "removeCurve"
        }
        self.model.removeRow(modelIndex.row())
        self.sigContextMenu.emit(ddict)

    def renameItemAction(self):
        _logger.debug('LegendListContextMenu.renameCurveAction called')
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
            'event': "renameCurve"
        }
        self.sigContextMenu.emit(ddict)

    def toggleLinesAction(self):
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
        }
        linestyle = modelIndex.data(LegendModel.iconLineStyleRole)
        visible = not modelIndex.data(LegendModel.showLineRole)
        _logger.debug('toggleLinesAction -- lines visible: %s', str(visible))
        ddict['event'] = "toggleLine"
        ddict['line'] = visible
        ddict['linestyle'] = linestyle if visible else ''
        self.model.setData(modelIndex, visible, LegendModel.showLineRole)
        self.sigContextMenu.emit(ddict)

    def togglePointsAction(self):
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
        }
        flag = modelIndex.data(LegendModel.showSymbolRole)
        symbol = modelIndex.data(LegendModel.iconSymbolRole)
        visible = not flag or symbol in NoSymbols
        _logger.debug(
            'togglePointsAction -- Symbols visible: %s', str(visible))

        ddict['event'] = "togglePoints"
        ddict['points'] = visible
        ddict['symbol'] = symbol if visible else ''
        self.model.setData(modelIndex, visible, LegendModel.showSymbolRole)
        self.sigContextMenu.emit(ddict)

    def setActiveAction(self):
        modelIndex = self.currentIdx()
        legend = str(modelIndex.data(qt.Qt.DisplayRole))
        _logger.debug('setActiveAction -- active curve: %s', legend)
        ddict = {
            'legend': legend,
            'label': legend,
            'selected': modelIndex.data(qt.Qt.CheckStateRole),
            'type': str(modelIndex.data()),
            'event': "setActiveCurve",
        }
        self.sigContextMenu.emit(ddict)

class RenameCurveDialog(qt.QDialog):
    """Dialog box to input the name of a curve."""

    def __init__(self, parent=None, current="", curves=()):
        super(RenameCurveDialog, self).__init__(parent)
        self.setWindowTitle("Rename Curve %s" % current)
        self.curves = curves
        layout = qt.QVBoxLayout(self)
        self.lineEdit = qt.QLineEdit(self)
        self.lineEdit.setText(current)
        self.hbox = qt.QWidget(self)
        self.hboxLayout = qt.QHBoxLayout(self.hbox)
        self.hboxLayout.addStretch(1)
        self.okButton = qt.QPushButton(self.hbox)
        self.okButton.setText('OK')
        self.hboxLayout.addWidget(self.okButton)
        self.cancelButton = qt.QPushButton(self.hbox)
        self.cancelButton.setText('Cancel')
        self.hboxLayout.addWidget(self.cancelButton)
        self.hboxLayout.addStretch(1)
        layout.addWidget(self.lineEdit)
        layout.addWidget(self.hbox)
        self.okButton.clicked.connect(self.preAccept)
        self.cancelButton.clicked.connect(self.reject)

    def preAccept(self):
        text = str(self.lineEdit.text())
        addedText = ""
        if len(text):
            if text not in self.curves:
                self.accept()
                return
            else:
                addedText = "Curve already exists."
        text = "Invalid Curve Name"
        msg = qt.QMessageBox(self)
        msg.setIcon(qt.QMessageBox.Critical)
        msg.setWindowTitle(text)
        text += "\n%s" % addedText
        msg.setText(text)
        msg.exec_()

    def getText(self):
        return str(self.lineEdit.text())

class LegendsDockWidget(qt.QDockWidget):
    """QDockWidget with a :class:`LegendSelector` connected to a PlotWindow.

    It makes the link between the LegendListView widget and the PlotWindow.

    :param parent: See :class:`QDockWidget`
    :param plot: :class:`.PlotWindow` instance on which to operate
    """

    def __init__(self, parent=None, plot=None):
        assert plot is not None
        self._plotRef = weakref.ref(plot)
        self._isConnected = False  # True if widget connected to plot signals

        super(LegendsDockWidget, self).__init__("Legends", parent)

        self._legendWidget = LegendListView()

        self.layout().setContentsMargins(0, 0, 0, 0)
        self.setWidget(self._legendWidget)

        self.visibilityChanged.connect(
            self._visibilityChangedHandler)

        self._legendWidget.sigLegendSignal.connect(self._legendSignalHandler)

    @property
    def plot(self):
        """The :class:`.PlotWindow` this widget is attached to."""
        return self._plotRef()

    def renameCurve(self, oldLegend, newLegend):
        """Change the name of a curve using remove and addCurve

        :param str oldLegend: The legend of the curve to be changed
        :param str newLegend: The new legend of the curve
        """
        is_active = self.plot.getActiveCurve(just_legend=True) == oldLegend
        curve = self.plot.getCurve(oldLegend)
        self.plot.remove(oldLegend, kind='curve')
        self.plot.addCurve(curve.getXData(copy=False),
                           curve.getYData(copy=False),
                           legend=newLegend,
                           info=curve.getInfo(),
                           color=curve.getColor(),
                           symbol=curve.getSymbol(),
                           linewidth=curve.getLineWidth(),
                           linestyle=curve.getLineStyle(),
                           xlabel=curve.getXLabel(),
                           ylabel=curve.getYLabel(),
                           xerror=curve.getXErrorData(copy=False),
                           yerror=curve.getYErrorData(copy=False),
                           z=curve.getZValue(),
                           selectable=curve.isSelectable(),
                           fill=curve.isFill(),
                           resetzoom=False)
        if is_active:
            self.plot.setActiveCurve(newLegend)

    def _legendSignalHandler(self, ddict):
        """Handles events from the LegendListView signal"""
        _logger.debug("Legend signal ddict = %s", str(ddict))

        if ddict['event'] == "legendClicked":
            if ddict['button'] == "left":
                self.plot.setActiveCurve(ddict['legend'])

        elif ddict['event'] == "removeCurve":
            self.plot.removeCurve(ddict['legend'])

        elif ddict['event'] == "renameCurve":
            curveList = self.plot.getAllCurves(just_legend=True)
            oldLegend = ddict['legend']
            dialog = RenameCurveDialog(self.plot, oldLegend, curveList)
            ret = dialog.exec_()
            if ret:
                newLegend = dialog.getText()
                self.renameCurve(oldLegend, newLegend)

        elif ddict['event'] == "setActiveCurve":
            self.plot.setActiveCurve(ddict['legend'])

        elif ddict['event'] == "checkBoxClicked":
            self.plot.hideCurve(ddict['legend'], not ddict['selected'])

        elif ddict['event'] in ["mapToRight", "mapToLeft"]:
            legend = ddict['legend']
            curve = self.plot.getCurve(legend)
            yaxis = 'right' if ddict['event'] == 'mapToRight' else 'left'
            self.plot.addCurve(x=curve.getXData(copy=False),
                               y=curve.getYData(copy=False),
                               legend=curve.getLegend(),
                               info=curve.getInfo(),
                               yaxis=yaxis)

        elif ddict['event'] == "togglePoints":
            legend = ddict['legend']
            curve = self.plot.getCurve(legend)
            symbol = ddict['symbol'] if ddict['points'] else ''
            self.plot.addCurve(x=curve.getXData(copy=False),
                               y=curve.getYData(copy=False),
                               legend=curve.getLegend(),
                               info=curve.getInfo(),
                               symbol=symbol)

        elif ddict['event'] == "toggleLine":
            legend = ddict['legend']
            curve = self.plot.getCurve(legend)
            linestyle = ddict['linestyle'] if ddict['line'] else ''
            self.plot.addCurve(x=curve.getXData(copy=False),
                               y=curve.getYData(copy=False),
                               legend=curve.getLegend(),
                               info=curve.getInfo(),
                               linestyle=linestyle)

        else:
            _logger.debug("unhandled event %s", str(ddict['event']))

    def updateLegends(self, *args):
        """Sync the LegendSelector widget displayed info with the plot.
        """
        legendList = []
        for curve in self.plot.getAllCurves(withhidden=True):
            legend = curve.getLegend()
            # Use active color if curve is active
            isActive = legend == self.plot.getActiveCurve(just_legend=True)
            style = curve.getCurrentStyle()
            color = style.getColor()
            if numpy.array(color, copy=False).ndim != 1:
                # array of colors, use transparent black
                color = 0., 0., 0., 0.

            curveInfo = {
                'color': qt.QColor.fromRgbF(*color),
                'linewidth': style.getLineWidth(),
                'linestyle': style.getLineStyle(),
                'symbol': style.getSymbol(),
                'selected': not self.plot.isCurveHidden(legend),
                'active': isActive}
            legendList.append((legend, curveInfo))

        self._legendWidget.setLegendList(legendList)

    def _visibilityChangedHandler(self, visible):
        if visible:
            self.updateLegends()
            if not self._isConnected:
                self.plot.sigContentChanged.connect(self.updateLegends)
                self.plot.sigActiveCurveChanged.connect(self.updateLegends)
                self._isConnected = True
        else:
            if self._isConnected:
                self.plot.sigContentChanged.disconnect(self.updateLegends)
                self.plot.sigActiveCurveChanged.disconnect(self.updateLegends)
                self._isConnected = False

    def showEvent(self, event):
        """Make sure this widget is raised when it is shown
        (when it is first created as a tab in PlotWindow or when it is shown
        again after hiding).
        """
        self.raise_()

max_stars_repo_name: vincefn/silx

max_issues_repo_name: vincefn/silx

max_forks_repo_name: vincefn/silx


text:
from mjrl.utils.gym_env import GymEnv
from mjrl.policies.gaussian_linear import LinearPolicy
from mjrl.policies.gaussian_linear_lpg_ftw import LinearPolicyLPGFTW
from mjrl.baselines.mlp_baseline import MLPBaseline
from mjrl.algos.npg_cg import NPG
from mjrl.algos.npg_cg_ftw import NPGFTW
from mjrl.utils.train_agent import train_agent
import time as timer
import numpy as np
import gym

from mjrl.utils.make_train_plots import make_multitask_train_plots, make_multitask_test_plots

SEED = 500

# MTL policy
# ==================================
e = {}
baseline_mtl = {}
baseline_stl = {}
policy_stl = {}
agent_stl = {}
num_tasks = 20
num_cpu = 1

np.random.seed(SEED)
for task_id in range(num_tasks):
      try_again = True
      while try_again:
            gravity_factor = np.random.rand() + 0.5
            env_id = 'HalfCheetah{:.4f}g-v0'.format(gravity_factor)
            try:
                  gym.envs.register(
                        id=env_id,
                        entry_point='gym_extensions.continuous.mujoco.modified_half_cheetah:HalfCheetahGravityEnv',
                        max_episode_steps=1000,
                        reward_threshold=3800.0,
                        kwargs=dict(gravity=-gravity_factor*9.81)
                        )
                  try_again = False
            except:
                  try_again = True

      # e[task_id] =  GymEnv('mjrl_swimmer-v0')
      e[task_id] = GymEnv(env_id)   # only do the environment here, so different files can create the same tasks
      print(env_id)

for task_id in range(num_tasks):
      # e = GymEnv('mjrl_point_mass-v0')
      # e = GymEnv('Swimmer-v2')
      baseline_mtl[task_id] = MLPBaseline(e[task_id].spec, reg_coef=1e-3, batch_size=64, epochs=2, learn_rate=1e-3)

      baseline_stl[task_id] = MLPBaseline(e[task_id].spec, reg_coef=1e-3, batch_size=64, epochs=2, learn_rate=1e-3)
      policy_stl[task_id] = LinearPolicy(e[task_id].spec, seed=SEED)
      agent_stl[task_id] = NPG(e[task_id], policy_stl[task_id], baseline_stl[task_id], normalized_step_size=1, seed=SEED, save_logs=True)

policy_mtl = LinearPolicyLPGFTW(e[0].spec, k=1, max_k=2, seed=SEED)
agent_mtl = NPGFTW(e, policy_mtl, baseline_mtl, normalized_step_size=1, seed=SEED, save_logs=True, new_col_mode='max_k')

# agent = BatchREINFORCE(e, policy, baseline, learn_rate=0.0001, seed=SEED, save_logs=True)
job_name_mtl = 'halfcheetah_lpgftw_gravity_exp1'
job_name_stl = 'halfcheetah_stl_gravity_exp1'
loggers_stl = {}
for task_id in range(num_tasks):
      ts = timer.time()
      train_agent(job_name=job_name_mtl,
                  agent=agent_mtl,
                  seed=SEED,
                  niter=50,
                  gamma=0.995,
                  gae_lambda=0.97,
                  num_cpu=num_cpu,
                  sample_mode='trajectories',
                  num_traj=10,
                  save_freq=5,
                  evaluation_rollouts=5,
                  task_id=task_id)
      agent_mtl.add_approximate_cost(N=10,
            task_id=task_id,
            num_cpu=1)

      train_agent(job_name=job_name_stl,
                  agent=agent_stl[task_id],
                  seed=SEED,
                  niter=50,
                  gamma=0.995,
                  gae_lambda=0.97,
                  num_cpu=num_cpu,
                  sample_mode='trajectories',
                  num_traj=10,
                  save_freq=5,
                  evaluation_rollouts=5,
                  task_id=task_id)
      loggers_stl[task_id] = agent_stl[task_id].logger
      print("time taken for linear policy training = %f" % (timer.time()-ts))

make_multitask_train_plots(loggers=agent_mtl.logger, keys=['stoc_pol_mean'], save_loc=job_name_mtl+'/logs/')
make_multitask_train_plots(loggers=loggers_stl, keys=['stoc_pol_mean'], save_loc=job_name_stl+'/logs/')

mean_test_perf = agent_mtl.test_tasks(test_rollouts=10,
                    num_cpu=num_cpu)
result = np.mean(list(mean_test_perf.values()))
print(result)
make_multitask_test_plots(mean_test_perf, save_loc=job_name_mtl+'/')

max_stars_repo_name: optimass/LPG-FTW

max_issues_repo_name: optimass/LPG-FTW

max_forks_repo_name: optimass/LPG-FTW


text:
"""
Useful set of tools
"""

import numpy as np
import pandas as pd
from functools import partial
from scipy.optimize import fmin
from abc import ABC, abstractmethod
import sklearn as sk

class MultiStratifiedKFold():
    """Multi Stratified K-Fold cross validator. Indices split happens at creation time

    Parameters
    ----------
    n_splits : int
        Number of splits
    df : pd.DataFrame
        DataFrame to split
    features : list of str
        List of features to be considered when stratifying
    seed : int, defaults to 0
        Random seed
    """

    def __init__(self, n_splits, df, features, seed=0):
        self.n_splits = n_splits
        self.features = features

        index = []
        np.random.seed(seed)

        for g, (key, idx) in enumerate(df.groupby(self.features).groups.items()):
            arr = idx.values.copy()
            np.random.shuffle(arr)
            k = (1 / self.n_splits)
            splits = np.split(arr, [int(i * k * len(arr)) for i in range(1, self.n_splits)])

            if g > 0:
                # Calculate length of each split
                len_splits = np.array([len(i) for i in splits])

                # sort indexes by length
                len_idx = np.argsort(len_splits)

                tst = []
                # Calculate and sort by length of fold so far
                for j in range(len(index)):
                    tst.append([len(index[j][i]) for i in range(self.n_splits)])
                tst = np.array(tst)
                len_folds = np.flip(np.argsort(tst.sum(axis=0)))

                new_splits = [[]] * self.n_splits

                for j in range(self.n_splits):
                    source = len_idx[j]
                    dest = len_folds[j]
                    new_splits[dest] = splits[source]
            else:
                new_splits = splits

            index.append(new_splits)

        self.indices = list(map(np.concatenate, zip(*index)))  # list of tuples containing the

    def get_indices(self, fold):
        """Returns indices for a given fold in the range [0,folds)

        Parameters
        ----------
        fold : int
            fold index to retrieve indices from

        Returns
        -------
        train_idx : ndarray
            The training set indices for that split.
        test_idx : ndarray
            The testing set indices for that split.
        """
        train_idx = np.concatenate([self.indices[i] for i in range(self.n_splits) if i != fold])
        val_idx = self.indices[fold]
        return train_idx, val_idx

    def split(self, X=None, y=None, groups=None):
        """Generate indices to split data into training and tests set.

        Parameters
        ----------
            All parameters are included for compatibility with SK-Learn only

        Yields
        ------
        train : ndarray
            The training set indices for that split.
        tests : ndarray
            The testing set indices for that split.
        """
        for i in range(self.n_splits):
            yield self.get_indices(i)

    def as_list(self):
        """Returns a list with the indices of each fold

        Returns
        -------
        List of tuples of ndarray, [(train_idx, val_idx)]
            List containing each split
        """
        return [self.get_indices(i) for i in range(self.n_splits)]

    def get_n_splits(self, X=None, y=None, groups=None):
        """Returns the number of splitting iterations in the cross-validator

        Parameters
        ----------
        All parameters are included for compatibility with SK-Learn only

        Returns
        -------
        int :
            The number of splitting iterations in the cross-validator.
        """
        return self.n_splits

class BaseOptimizeBlend(ABC):
    """
    Base class for Optimizer used in blending results of models.

    Parameters
    ----------
    maxiter : int
        Maximum number of iterations. Defaults to 1000
    maxfun : int
        Maximum number of function calls. Defaults to 1000
    """
    def __init__(self, maxiter=1000, maxfun=1000):
        self._coef = 0
        self._maxiter = maxiter
        self._maxfun = maxfun

    @abstractmethod
    def metric(self, coef, X, y):
        """
        This is an abstract method that shall return the metric to be minimized.
        If what is needed is to maximize the metric, then return -1 * metric.

        Parameters
        ----------
        coef : (ndarray with shape (n_models, ))
            Array of coefficients to be optimized.
        X : ( ndarray with shape (n_samples, n_models))
            Predicted values by each estimator.
        y : ( ndarray with shape (n_samples, ))
            True values.

        Returns
        -------
        float
            The score metric of all estimators with coefficients applied
        """
        coef = coef.reshape([X.shape[0]]  + [1 for i in range(len(X.shape)-1)])
        x_coef = X * coef
        predictions = np.sum(x_coef, axis=1)
        score = sk.metrics.mean_squared_error(y, predictions)
        return score

    def fit(self, X, y):
        """
        Fit the results ``X`` to the true values ``y`` by calling ``fmin`` function

        Parameters
        ----------
        X : (ndarray with shape (n_samples, n_models))
            Predicted values by the estimators
        y : (ndarray with shape (n_samples, ))
            true values
        """
        partial_loss = partial(self.metric, X=X, y=y)
        init_coef = np.random.dirichlet(np.ones( X.shape[0] ))
        self._coef = fmin(partial_loss, init_coef, disp=True, maxiter=self._maxiter, maxfun=self._maxfun)

    def predict(self, X):
        """
        Once the Optimizer is fitted, blend the predictions from n_models using the calculated coefficients

        Parameters
        ----------
        X : (ndarray with shape (n_samples, n_models))
            predicted values by the estimators.

        Returns
        -------
        ndarray with shape (n_samples, )
            Array with the X values blended applying coefficients
        """
        x_coef = X * self._coef
        predictions = np.sum(x_coef, axis=1)
        return predictions

    def get_coef(self):
        """
        Returns the number of splitting iterations in the cross-validator

        Returns
        -------
        ndarray with shape (n_models,)
            The calculated coefficients.
        """
        return self._coef

def read_df(file):
    """
    Read Dataframe from file, based on file extension

    Parameters
    ----------
    file : str
        Path with an extension of type ``.csv``, ``.feather`` or ``.parquet``
    """
    extension = file.split('.')[1]
    if extension == 'csv':
        return pd.read_csv(file)
    if extension == 'feather':
        return pd.read_feather(file)
    if extension == 'parquet':
        return pd.read_parquet(file)

reason: import numpy,from scipy


text:
import numpy as np
import cv2 as cv
import sys
import ctypes
def Mbox(title, text, style):
    return ctypes.windll.user32.MessageBoxW(0, text, title, style)

Mbox('OpenCV Version: ', cv.getBuildInformation(), 1)


text:
#     Copyright 2016-present CERN – European Organization for Nuclear Research
#
#     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.
from typing import Sequence, Dict, Any

import blpapi
import numpy as np
from blpapi import DataType
from pandas import to_datetime

from qf_lib.common.tickers.tickers import BloombergTicker
from qf_lib.common.utils.logging.qf_parent_logger import qf_logger
from qf_lib.containers.dataframe.qf_dataframe import QFDataFrame
from qf_lib.data_providers.bloomberg.bloomberg_names import REF_DATA_SERVICE_URI, SECURITY, FIELD_DATA
from qf_lib.data_providers.bloomberg.exceptions import BloombergError
from qf_lib.data_providers.bloomberg.helpers import set_tickers, set_fields, get_response_events, \
    check_event_for_errors, extract_security_data, check_security_data_for_errors

class ReferenceDataProvider:
    """ Used for providing current price data from Bloomberg. """

    def __init__(self, session):
        self._session = session
        self.logger = qf_logger.getChild(self.__class__.__name__)

    def get(self, tickers: Sequence[BloombergTicker], fields,  override_name: str = None, override_value: Any = None):
        ref_data_service = self._session.getService(REF_DATA_SERVICE_URI)
        request = ref_data_service.createRequest("ReferenceDataRequest")

        tickers_str = [t.as_string() for t in tickers]
        set_tickers(request, tickers_str)
        set_fields(request, fields)
        if override_name is not None:
            self._set_override(request, override_name, override_value)

        self._session.sendRequest(request)
        return self._receive_reference_response(tickers, fields)

    def _receive_reference_response(self, tickers, fields):
        ticker_str_to_ticker: Dict[str, BloombergTicker] = {t.as_string(): t for t in tickers}

        response_events = get_response_events(self._session)
        tickers_fields_container = QFDataFrame(index=tickers, columns=fields)

        for ev in response_events:
            try:
                check_event_for_errors(ev)
                security_data_array = extract_security_data(ev)
                check_security_data_for_errors(security_data_array)

                for security_data in security_data_array.values():
                    try:
                        check_security_data_for_errors(security_data)
                        field_data_array = security_data.getElement(FIELD_DATA)

                        security_name = security_data.getElementAsString(SECURITY) if \
                            security_data.hasElement(SECURITY, True) else None

                        ticker = ticker_str_to_ticker[security_name]
                        for field_name in fields:
                            value = self._parse_value(field_data_array, field_name)
                            tickers_fields_container.loc[ticker, field_name] = value

                    except KeyError:
                        self.logger.warning("Received data for a ticker which was not present in the request. "
                                            "The data for that ticker will be excluded from parsing.")
                    except BloombergError as e:
                        self.logger.error(e)

            except BloombergError as e:
                self.logger.error(e)

        return tickers_fields_container.infer_objects()

    @classmethod
    def _set_override(cls, request, override_name, override_value):
        overrides = request.getElement("overrides")
        override = overrides.appendElement()
        override.setElement("fieldId", override_name)
        override.setElement("value", override_value)

    @staticmethod
    def _parse_value(field_data_array, field_name):
        if not field_data_array.hasElement(field_name, True):
            return None
        field_element = field_data_array.getElement(field_name)
        element_data_type = field_element.datatype()
        try:
            if element_data_type in (DataType.FLOAT32, DataType.FLOAT64):
                value = field_element.getValueAsFloat()
            elif element_data_type in (DataType.INT32, DataType.INT64):
                value = field_element.getValueAsInteger()
            elif element_data_type is DataType.BOOL:
                value = field_element.getValueAsBool()
            elif element_data_type in (DataType.DATETIME, DataType.TIME, DataType.DATE):
                value = to_datetime(field_element.getValueAsDatetime())
            else:
                value = field_element.getValueAsString()
        except blpapi.exception.NotFoundException:
            value = np.nan

        return value

max_stars_repo_name: quarkfin/QF-Lib

max_issues_repo_name: quarkfin/QF-Lib

max_forks_repo_name: quarkfin/QF-Lib


text:
'''
HRFI target detection.

Calling it tracker for self-consistency with other modules,
however time is not currently factored into tracking
'''
import json
import os
import logging

import numpy              as np

from skimage.filters      import threshold_otsu
from skimage.segmentation import clear_border
from skimage.measure      import label, regionprops
from skimage.morphology   import closing, square

from utils.file_manipulation import tiff_read

def hrfi_tracker(hrfi_filepath, experiment, config):

    rgb = tiff_read(hrfi_filepath)

    # Verify band weights in config
    sum_band_w = np.round(config["instrument"]["red_band_weight"] + \
                          config["instrument"]["green_band_weight"] + \
                          config["instrument"]["blue_band_weight"], 3)

    if sum_band_w != 1.0:
        logging.warning(f"Instrument band weights don't sum to 1 ({sum_band_w}), normalizing.")
        config["instrument"]['red_band_weight'] /= sum_band_w
        config["instrument"]['green_band_weight'] /= sum_band_w
        config["instrument"]['blue_band_weight'] /= sum_band_w

    # Reduce 3 banded image to 1, with band weights being set by detector calibration
    gray = ((config["instrument"]["red_band_weight"]   * rgb[:,:,0]) +
            (config["instrument"]["green_band_weight"] * rgb[:,:,1]) +
            (config["instrument"]["blue_band_weight"]  * rgb[:,:,2]))

    # Scale data 0 to 1
    gray = np.clip(gray, 0, config["instrument"]["max_dn"])
    gray = gray / config["instrument"]["max_dn"]

    # Determine threshold
    t1 = threshold_otsu(gray)
    t2 = np.percentile(gray, config["instrument"]["min_perc"])
    t = max(t1, t2)

    # Determine connected regions of whitespace and mark them as potential critters
    image = np.zeros((rgb.shape[0],rgb.shape[1]))
    image[gray > t] = 1
    bw = closing(image==1, square(3))
    cleared = clear_border(bw)

    label_image = label(cleared)

    # Filter to only save organisms larger than min_organism_size pixels
    bboxes = []
    [bboxes.append(region.bbox) for region in regionprops(label_image) if region.area >= config["tracker"]["min_bbox_area"] and region.area <= config["tracker"]["max_bbox_area"]]

    track_folder = os.path.join(experiment, config['experiment_dirs']['track_dir'])
    if not os.path.exists(track_folder):
        os.makedirs(track_folder)

    file_tag = hrfi_filepath.split("/")[-1].split(".")[0]
    with open(os.path.join(track_folder, f"{file_tag}_bboxes.json"), 'w') as f:
        json.dump(bboxes, f, indent=2)

max_stars_repo_name: JPLMLIA/OWLS-Autonomy

max_issues_repo_name: JPLMLIA/OWLS-Autonomy

max_forks_repo_name: JPLMLIA/OWLS-Autonomy


text:
#
# imports
#

import os
import sys
import glob
import copy
import numpy as np
import argparse

#
# main function
#

if __name__ == '__main__':

    # initialize arg parser
    parser = argparse.ArgumentParser(description='Split/merge the MS/MS data set into/from partitions')

    # data root directory
    parser.add_argument('-i', '--idir', dest='inpath', type=str, required=True,
                        help='Path to data set files')

    # data file extensions
    parser.add_argument('-e', '--ext', dest='extension', type=str, required=True,
                        help='Data set file extension e.g. ms2, mzML, mzXML, mgf')

    # number of nodes/partitions
    parser.add_argument('-N', dest='nodes', type=int, required=True,
                        help='Number of partitions: >= 2 if -m not set, else anything')

    # merge operation?
    parser.add_argument('-m', '--merge', dest='merge', action='store_true',
                        help='Merge all existing partitions: -N <anything is fine>, default: False')

    # parse arguments
    args = parser.parse_args()

    # input path to the MS2 data directory
    inpath = args.inpath.lstrip(' ')
    inpath = os.path.expanduser(inpath)

    # check if directory exists
    if not os.path.isdir(inpath):
        print ('ERROR: Directory does not exist\n')
        sys.exit (-1)

    # get file extension
    extension = args.extension.lstrip(' ')

    # extract extra fields from spectra
    merge = args.merge

    # if merge operation
    if (merge == True):

        # get all files in part_*/*.ext
        files = np.array(glob.glob(inpath + '/part_*/*.' + extension))

        # move data parts back to the root
        for file in files:
            os.rename(file, inpath + '/' + os.path.split(file)[1])

    # scatter operation
    else:
        # get number of nodes
        nodes = args.nodes

        # check for number of nodes
        if nodes < 2:
            print ('ERROR: N >= 2 required')
            sys.exit(-1)

        # get all files with *.ext
        files = np.array(glob.glob(inpath + '/*.' + extension))

        # check if files found
        if len(files) < 1:
            print ('ERROR: No ' + extension + ' files found\n')
            sys.exit(-1)

        # check number of nodes and files
        if nodes > len(files):
            print ('WARNING: # partitions > # files')
            print ('Some partitions may be empty\n')

        # shuffle the files
        np.random.shuffle(files)

        # split files in 'nodes' partitions
        splits = np.array_split(files, nodes)

        # create new directories
        for i in range(nodes):
            os.makedirs(inpath + '/part_' + str(i+1), exist_ok=True)

            # move data to parts
            for file in splits[i]:
                os.rename(file, inpath + '/part_' +  str(i+1) + '/' + os.path.split(file)[1])

    # print final status
    print ('DONE\n')

max_forks_repo_name: hicops/hicops


text:
# Example 1-1. Monte Carlo valuation of European call option
#
# Monte Carlo valuation of European call option
# in Black-Scholes-Merton model
# bsm_mcs_euro.py
#
# Python for Finance, 2nd ed.
# (c) Dr. Yves J. Hilpisch
#
import math
import numpy as np

# Parameter Values
S0 = 100.  # initial index level
K = 105.  # strike price
T = 1.0  # time-to-maturity
r = 0.05  # riskless short rate
sigma = 0.2  # volatility

I = 100000  # number of simulations

# Valuation Algorithm
z = np.random.standard_normal(I)  # pseudo-random numbers
# index values at maturity
ST = S0 * np.exp((r - 0.5 * sigma ** 2) * T + sigma * math.sqrt(T) * z)
hT = np.maximum(ST - K, 0)  # payoff at maturity
C0 = math.exp(-r * T) * np.mean(hT)  # Monte Carlo estimator

# Result Output
print('Value of the European call option %5.3f.' % C0);

max_stars_repo_name: JackyCSer/Python-for-Finance

max_issues_repo_name: JackyCSer/Python-for-Finance

max_forks_repo_name: JackyCSer/Python-for-Finance


text:
from __future__ import division
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yaml
import scipy as sp

# Conversion constants

k_B = 1.38065e-23 #[J/K]
N_A = 6.02214e23 #[1/mol]
m3_to_nm3 = 1e27
gm_to_kg = 1./1000
J_to_kJ = 1./1000
J_per_m3_to_kPA = 1./1000

class LennardJones_2C():
    def __init__(self,M_w):

        self.M_w = M_w

        with open("DCLJQ_fluid.yaml") as yfile:
            yfile = yaml.load(yfile)

        self.T_c_star_params = np.array(yfile["correlation_parameters"]["Stoll"]["T_c_star_params"])
        self.rho_c_star_params = np.array(yfile["correlation_parameters"]["Stoll"]["rho_c_star_params"])
        self.b_C1 = np.array(yfile["correlation_parameters"]["Stoll"]["rho_L_star_params"]["C1_params"])
        self.b_C2_L = np.array(yfile["correlation_parameters"]["Stoll"]["rho_L_star_params"]["C2_params"])
        self.b_C3_L = np.array(yfile["correlation_parameters"]["Stoll"]["rho_L_star_params"]["C3_params"])
        self.b_C2_v = np.array(yfile["correlation_parameters"]["Stoll"]["rho_v_star_params"]["C2_params"])
        self.b_C3_v = np.array(yfile["correlation_parameters"]["Stoll"]["rho_v_star_params"]["C3_params"])
        self.b_c1 = np.array(yfile["correlation_parameters"]["Stoll"]["P_v_star_params"]["c1_params"])
        self.b_c2 = np.array(yfile["correlation_parameters"]["Stoll"]["P_v_star_params"]["c2_params"])
        self.b_c3 = np.array(yfile["correlation_parameters"]["Stoll"]["P_v_star_params"]["c3_params"])

    def T_c_star_hat(self,q,l):
        b=self.T_c_star_params
        x = np.array([1, q**2, q**3, 1./(0.1+l**2), 1./(0.1+l**5), q**2/(0.1+l**2), q**2/(0.1+l**5), q**3/(0.1+l**2), q**3/(0.1+l**5)])
        T_c_star = x*b
        T_c_star = T_c_star.sum()
        return T_c_star

    def rho_c_star_hat(self,q,l):
        b=self.rho_c_star_params
        x = np.array([1, q**2, q**3, l**2/(0.11+l**2), l**5/(0.11+l**5), l**2*q**2/(0.11+l**2), l**5*q**2/(0.11+l**5), l**2*q**3/(0.11+l**2), l**5*q**3/(0.11+l**5)])
        rho_c_star = x*b
        rho_c_star = rho_c_star.sum()
        return rho_c_star

    def C1_hat(self,q,l,b):
        x_C1 = np.array([1, q**2, q**3, l**3/(l+0.4)**3, l**4/(l+0.4)**5, q**2*l**2/(l+0.4), q**2*l**3/(l+0.4)**7, q**3*l**2/(l+0.4), q**3*l**3/(l+0.4)**7])
        C1 = x_C1*b
        C1 = C1.sum()
        return C1

    def C2_hat(self,q,l,b):
        x_C2 = np.array([1, q**2, q**3, l**2, l**3, q**2*l**2, q**2*l**3, q**3*l**2])
        C2 = x_C2*b
        C2 = C2.sum()
        return C2

    def C3_hat(self,q,l,b):
        x_C3 = np.array([1, q**2, q**3, l, l**4, q**2*l, q**2*l**4, q**3*l**4])
        C3 = x_C3*b
        C3 = C3.sum()
        return C3

    def rho_star_hat_2CLJQ(self,T_star,q,l,phase):
        b_C1, b_C2_L, b_C3_L, b_C2_v, b_C3_v = self.b_C1, self.b_C2_L, self.b_C3_L, self.b_C2_v, self.b_C3_v
        T_c_star = self.T_c_star_hat(q,l)
        rho_c_star = self.rho_c_star_hat(q,l)
        tau = T_c_star - T_star # T_c_star - T_star
        if all(tau>0):
            x = np.ones([len(tau),4]) # First column is supposed to be all ones
            x[:,1] = tau**(1./3)
            x[:,2] = tau
            x[:,3] = tau**(3./2)
            C1 = self.C1_hat(q,l,b_C1)
            if phase == 'liquid':
                C2 = self.C2_hat(q,l,b_C2_L)
                C3 = self.C3_hat(q,l,b_C3_L)
                b = np.array([rho_c_star, C1, C2, C3])
            elif phase == 'vapor':
                C2 = self.C2_hat(q,l,b_C2_v)
                C3 = self.C3_hat(q,l,b_C3_v)
                b = np.array([rho_c_star, -C1, C2, C3])
            else:
                return 0
            #rho_star = b[0]+b[1]*tau**(1./3)+b[2]*tau+b[3]*tau**(3./2) #The brute force approach
            rho_star = x*b
            rho_star = rho_star.sum(axis=1) # To add up the rows (that pertain to a specific T_star)
        else:
            rho_star = np.zeros([len(tau)])
        return rho_star

    def rho_hat_2CLJQ(self,Temp,eps,sig,Lbond,Qpole,phase):
        M_w = self.M_w
        T_star = Temp/eps
        Q2_star = Qpole**2/(eps*sig**5)
        L_star = Lbond/sig
        rho_star = self.rho_star_hat_2CLJQ(T_star,Q2_star,L_star,phase)
        rho = rho_star *  M_w  / sig**3 / N_A * m3_to_nm3 * gm_to_kg #[kg/m3]
        return rho

    def rhol_hat_2CLJQ(self,Temp,eps,sig,Lbond,Qpole):
        rhol = self.rho_hat_2CLJQ(Temp,eps,sig,Lbond,Qpole,'liquid')
        return rhol #[kg/m3]

    def rhov_hat_2CLJQ(self,Temp,eps,sig,Lbond,Qpole):
        rhov = self.rho_hat_2CLJQ(Temp,eps,sig,Lbond,Qpole,'vapor')
        return rhov #[kg/m3]

    def Psat_star_hat_2CLJQ(self,T_star, q,l):
        b_c1, b_c2, b_c3 = self.b_c1, self.b_c2, self.b_c3
        x_c1 = [1., q**2, q**3, l**2/(l**2+0.75), l**3/(l**3+0.75), l**2*q**2/(l**2+0.75), l**3*q**2/(l**3+0.75), l**2*q**3/(l**2+0.75), l**3*q**3/(l**3+0.75)]
        x_c2 = [1., q**2, q**3, l**2/(l+0.75)**2, l**3/(l+0.75)**3, l**2*q**2/(l+0.75)**2, l**3*q**2/(l+0.75)**3, l**2*q**3/(l+0.75)**2, l**3*q**3/(l+0.75)**3]
        x_c3 = [q**2, q**5, l**0.5]
        c1 = (x_c1*b_c1).sum()
        c2 = (x_c2*b_c2).sum()
        c3 = (x_c3*b_c3).sum()
        Psat_star = np.exp(c1 + c2/T_star + c3/(T_star**4))
        return Psat_star

    def Psat_hat_2CLJQ(self,Temp,eps,sig,Lbond,Qpole):
        T_star = Temp/eps
        Q2_star = Qpole**2/(eps*sig**5)
        L_star = Lbond/sig
        Psat_star = self.Psat_star_hat_2CLJQ(T_star,Q2_star,L_star)
        Psat = Psat_star *  eps  / sig**3 * k_B * m3_to_nm3 * J_per_m3_to_kPA #[kPa]
        return Psat

    def LJ_model(self,r,eps,sig):
        r_star = r/sig
        U = 4 * eps * (r_star**(-12) - r_star**(-6))
        return U

reason: import numpy,import scipy