text:
#!/usr/bin/env python
import argparse
import concurrent.futures
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import json
import os
import shutil
import subprocess
import sys

INVALID_ENERGY = 1e100
RANKING_COLORS = [
    'aqua',
    'lime',
    'yellow',
    ]

def main():
    parser = argparse.ArgumentParser(description='update_result')
    parser.add_argument('--default_info_directory_path', required=True,
                        help='Directory path containing info files for default traces.')
    parser.add_argument('--byhand_info_directory_path', required=False,
                        help='Directory path containing info files for hand written traces.',
                        default='../data/by_hand/info')
    parser.add_argument('--info_directory_path_base', required=True,
                        help='Parent directory path containing info files.')
    parser.add_argument('--engines', required=True,
                        help='Space-separated engine names. ex) "stupid_engine stupid_engine_v2 stupid_bbox_engine"')
    parser.add_argument('--input_model_directory_path', required=True,
                        help='Directory path containing input model files.')
    parser.add_argument('--output_html_file_path', required=True,
                        help='Directory path containing input model files.')
    parser.add_argument('--output_csv_file_path', default='compare_engine.csv',
                        help='CSV output path.')
    args = parser.parse_args()

    engines = args.engines.split()

    headers = []
    rows = []

    with open(args.output_html_file_path, 'wt') as f:
        print('''

<title>Engine comparison result</title>

model name
default
registered
''', file=f)
        headers.append('model_name')
        headers.append('default')
        headers.append('registered')
        for engine in engines:
            print('{}	'.format(engine), file=f)
            headers.append(engine)
        print('
'.format(engine), file=f)
        headers.append('best')

        info_directory_paths = list()
        info_directory_paths.append(args.default_info_directory_path)
        if args.byhand_info_directory_path:
            info_directory_paths.append(args.byhand_info_directory_path)
        for engine in engines:
            info_directory_paths.append(os.path.join(args.info_directory_path_base, engine))

        model_titles = sorted(list({f[:f.index('_')]
                                    for f
                                    in os.listdir(args.input_model_directory_path)
                                    if os.path.splitext(f)[1] == '.mdl'}))
        for model_title in model_titles:
            row = []
            print('
', file=f)
            print('{}	'.format(model_title), file=f)
            row.append(model_title)
            energies = list()
            best_energy = 1e100
            for info_directory_path in info_directory_paths:
                info_file_path = os.path.join(info_directory_path, model_title + '.json')
                successful = False
                try:
                    with open(info_file_path, 'rt') as json_file:
                        info = json.load(json_file)
                    successful = info['successful']
                except IOError:
                    pass

                if not successful:
                    energies.append(INVALID_ENERGY)
                    continue

                energy = info['energy']
                energies.append(energy)
                if best_energy > energy:
                    best_energy = energy
            default_energy = energies[0]

            energy_to_ranking = {v: k for k, v in reversed(list(enumerate(sorted(energies))))}

            for energy in energies:
                if energy == INVALID_ENERGY:
                    print('	', file=f)
                    row.append(np.nan)
                    continue

                rank = energy_to_ranking[energy]
                if default_energy < energy:
                    color = 'red'
                elif rank < len(RANKING_COLORS):
                    color = RANKING_COLORS[rank]
                else:
                    color = 'transparent'

                print('{1}	'.format(color, energy), file=f)
                row.append(energy)
            row.append(best_energy)
            rows.append(row)
            print('
', file=f)
        print('''

''', file=f)

    df = pd.DataFrame(rows, columns=headers)
    df.to_csv(args.output_csv_file_path, index=False)

if __name__ == '__main__':
	main()


text:
import numpy as np
from testing_module import things_to_test, run_test
import matplotlib.pyplot as plt
import pandas as pd
#######################################################
# go to the bottom section of the code to run this file
#######################################################
def run_episodes(filename):
    kp_test = np.array([700, 700, 500, 100, 100, 50])
    kd_test = np.array([1.5, 1.5, 2, 5, 5, 5])

    perception_error_test = things_to_test('perception_error',
                                            testing_min = np.array([0.0]),
                                            testing_max = np.array([0.005]),
                                            amount_of_tests = 40)

    # testing code (no need to change anything below)
    data = []
    test_num = 1
    for perception_error in iter(perception_error_test):
        result = run_test(kp_test, kd_test, perception_error)
        data.append([kp_test,
                    kd_test,
                    perception_error,
                    result['run_time'],
                    result['inserting_eeff_xy_max'],
                    result['inserting_eeff_z_max'],
                    result['actuation_torque']])

        print('episode number: {a:3d}, run_time: {b:2.2f}'.format(a=test_num,b=result['run_time']))
        print('          kp: ',end="")
        print(kp_test)
        print('          kd: ',end="")
        print(kd_test)
        test_num += 1

    df = pd.DataFrame(data, columns=['kp', 'kd', 'perception error','run time', 'xy error', 'z error','torque'])
    import os.path
    df.to_csv(filename,mode='a',header= not os.path.exists(filename),index=False)
    print('Saved result to ' + filename)

def draw_perception_error(filename):
   # read file
    df = pd.read_csv(filename)
    # draw some figures with the test results.
    # Options:
        # run time
        # xy error
        # z error
    data_to_plot = 'run time'
    # drawing
    x = df['perception error'].apply(lambda x: float(x.replace("[","").replace("]",""))*1000)
    y = df[data_to_plot]

    fig = plt.figure(figsize=(10,6))
    plt.plot(x,y, '.')
    plt.xlabel('Perception error [mm]')
    plt.ylabel(data_to_plot)
    plt.grid()
    plt.show()

if __name__ == '__main__':
    filename = './results/percp700.csv' # the file to read/write

    # runs the simulation and save result file
    # run_episodes(filename)

    # draw the results from the file
    draw_perception_error(filename)

ext: py

max_stars_repo_name: garlicbutter/Jonathan-Tom

max_issues_repo_name: garlicbutter/Tom-Jonathan

max_forks_repo_name: garlicbutter/Tom-Jonathan

reason: import numpy


text:
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patheffects as PathEffects
import json
import cv2
from PIL import Image, ImageDraw, ImageFont
import numpy as np
from adjustText import adjust_text
from blend_modes import multiply
from datetime import datetime
import requests

def matPlotGenerate(mapType, dpiRes):

    fig = plt.figure()

    # Allow for plotting to a single subplot, so multiple windows for nation don't pop up
    # figsize is a little weird when working in conjunction with the padding output:
    #    "... pad_inches = -0.685)".
    # This is to align and overlay the claims later, figsize=(15,15) seems to align right.
    # Idk why, but something tells me "Metric is better"

    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(15,15))

    texts = []

    def parseJSONData(alpha):

        response = requests.get("https://raw.githubusercontent.com/ccmap/data/master/land_claims.civmap.json")
        data = json.loads(response.text)

        for item in data["features"]:

            # try, except to catch errors in data, will pass if so

            try:

                # Default size for names on the graph.
                nameSize = 5
                colour = str(item["color"]).replace(" ","")

                for polygon in item["polygon"]:
                    coord = polygon
                    coord.append(coord[0])
                    xs, ys = zip(*coord)

                    # Polygon filling, if alpha is True land will be filled,
                    # otherwise text will be generated separately, on transparent land ... (Yes,I know)

                    if alpha:
                        ax.fill(xs, ys, colour, alpha=1)

                    else:

                        # Name sizing is arbitrary atm and based on land mass,
                        # dividing by 80,000 seemed to produce decent results

                        area = 0.5*np.abs(np.dot(xs, np.roll(ys, 1))-np.dot(ys,np.roll(xs, 1)))
                        nameSize = int(round(area/(nameSize*80000)))

                        # Some names are still either too big or too small

                        if nameSize <= 5:
                            nameSize = 5
                        if nameSize >= 30:
                            nameSize = 30

                        matplotlib.rcParams["font.serif"] = "Times"
                        matplotlib.rcParams["font.family"] = "serif"

                        ax.fill(xs, ys, color=colour, alpha=0)
                        centerOfPolygon = (max(xs)+min(xs))/2., (max(ys)+min(ys))/2.
                        txt = ax.text(centerOfPolygon[0],centerOfPolygon[1],item["name"], size=nameSize, color="black", horizontalalignment="center", verticalalignment="center")
                        texts.append(txt)

            except Exception as ex:
                print(ex)
                print(item['name'])
                pass

    def setMatplotLib():

        # Invert y axis, otherwise using Minecraft's coordinates will be invert the map vertically

        plt.gca().invert_yaxis()

        # Matplotlib settings to remove axis' and whitespace from image

        fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
        ax.get_xaxis().set_visible(False)
        ax.axis("tight")
        ax.axis("off")
        ax.set_aspect(1)

    if mapType is "textMap":

        parseJSONData(alpha=False)
        setMatplotLib()

        # TODO: Better text generation, currently this keeps them from overlapping,
        # it's also really resource hungry to do this.

        adjust_text(texts,
                    expand_text=(0.5, 0.5),
                    force_text=(0.01, 0.1))

        plt.savefig("textMap.png",
                    dpi=dpiRes,
                    transparent=True)

    elif mapType is "polygons":

        parseJSONData(alpha=True)

        setMatplotLib()
        plt.savefig("polygons.png",
                    dpi=dpiRes,
                    transparent=True,
                    bbox_inches="tight",
                    pad_inches = -0.685)

def resizeImage(inputPath, outputPath, width, height):

    # Set max amount pixels to None, otherwise pillow will log a warning.

    Image.MAX_IMAGE_PIXELS = None
    img = Image.open(inputPath)
    img = img.convert("RGBA")
    img = img.resize((width, height), Image.ANTIALIAS)
    img.save(outputPath)

def removeWater(inputPath, outputPath):

    img = Image.open(inputPath)
    if "A" in img.getbands():

        # Just checking, and converting to rgba if it isn't already
        img = img.convert("RGBA")

    # Then convert to HSV, as its way easier to work with

    HSVim = img.convert("HSV")
    RGBna = np.array(img)
    HSVna = np.array(HSVim)
    H = HSVna[:,:,0]

    # HSV values for blue

    lo,hi = 170, 250

    # Convert HSV to back to RGBA

    lo = int((lo * 255) / 360)
    hi = int((hi * 255) / 360)
    blue = np.where((H>lo) & (H<hi))
    RGBna[blue] = [0, 0, 0, 0]

    Image.fromarray(RGBna).save(outputPath)

def multiplyMaps(backgroundImage, foregroundImage, opacity=0.6):

    # From the blend_modes library.

    background_img_float = cv2.imread(backgroundImage,-1).astype(float)
    foreground_img_float = cv2.imread(foregroundImage,-1).astype(float)
    blended_img_float = multiply(background_img_float, foreground_img_float, opacity)
    blended_img_uint8 = blended_img_float.astype(np.uint8)
    cv2.imwrite("overlayMap.png", blended_img_uint8)

def createImage(height, width, colour, typeOfimage, textColour="white", font="Times"):

    img = Image.new("RGBA", (width, height), (255, 0, 0, 0))
    x, y =  img.size
    eX, eY = width, height
    bbox =  (x/2 - eX/2, y/2 - eY/2, x/2 + eX/2, y/2 + eY/2)

    draw = ImageDraw.Draw(img)

    if typeOfimage is "ocean":
        draw.ellipse(bbox, fill=colour)

    elif typeOfimage is "background":
        draw.rectangle(bbox, fill=colour)

        # TODO: Convert text, and text placements to percentages

        smallFont = ImageFont.truetype(font, 90)
        mediumFont = ImageFont.truetype(font, 140)
        fntSymbols = ImageFont.truetype(font, 350)

        draw.text((100,100), "Claims map organised by\nThe Unified map project", font=mediumFont, fill=textColour)
        draw.text((100, 350), "discord.gg/hKJ9hQXB5S", font=smallFont, fill=textColour)
        draw.text((width-1200, 100), "Base map by Gjum", font=mediumFont, fill=textColour)
        draw.text((width-700, 250), "ccmap.github.io", font=smallFont, fill=textColour)

        draw.text((200,height-300), "Made by Camokool", font=mediumFont, fill=textColour)

        datem = datetime.now().strftime("%h")
        datey = datetime.now().strftime("%y")
        dateString = datem.capitalize() + " " + datey
        draw.text((width-900, height-300), dateString, font=mediumFont, fill=textColour)

        draw.text((550, 550), "~,~", font=fntSymbols, fill=textColour)
        draw.text((550, height-850), "-,+", font=fntSymbols, fill=textColour)

        draw.text((width-1000, 550), "+,-", font=fntSymbols, fill=textColour)
        draw.text((width-1000, height-850), "+,+", font=fntSymbols, fill=textColour)

    del draw
    img.save(typeOfimage + ".png", "PNG")

def mergeImage(backgroundImage, overlayImage, outputPath):

    img = Image.open(overlayImage)
    background = Image.open(backgroundImage)
    img_w, img_h = img.size
    bg_w, bg_h = background.size
    offset = ((bg_w - img_w) // 2, (bg_h - img_h) // 2)
    background.paste(img, offset, img)
    background.save(outputPath, "PNG")

def main():

    start = datetime.now()

    # Current resolution and dpi are important to the text generation, meaning...
    # The values at this current time can't be changed. See createImage().

    dpiRes = 500
    width = 5000
    height = 5000

    oceanColour = "#A6BBEB"

    backgroundBorder = 500
    backgroundColour = "#363636"

    print("Creating polygons, and painting my nails... 💅")
    matPlotGenerate("polygons", dpiRes)

    print("Generating text map, this may take awhile... 😴")
    matPlotGenerate("textMap", dpiRes)

    print("Resizing a few files... 🗃")
    resizeImage("baseMap.png", "resizedBaseMap.png", width, height)
    resizeImage("polygons.png", "resizedPolygons.png", width, height)
    resizeImage("textMap.png", "resizedText.png", width, height)

    print("Removing 💦 from the 🌏... ")
    removeWater("resizedBaseMap.png", "transparentBaseMap.png")

    print("Overlaying polygons to map... 🤓 ☝ 🗺")
    multiplyMaps("transparentBaseMap.png", "resizedPolygons.png")

    print("Creating a serene sea... 🌅")
    createImage(width, height, oceanColour, "ocean")

    print("Merging earth and sea... ⛰💥🌊")
    mergeImage("ocean.png", "overlayMap.png", "mapNoBackground.png")

    # We'll leave these out for now

    # print("Merging text and map layers")
    # mergeImage("mapNoBackground.png", "resizedText.png", "mapWithText.png")

    # print("Creating a nice background... 🧑💻")
    # createImage((width+backgroundBorder), (height+backgroundBorder), backgroundColour, "background")

    # print("Merging final layers... 🍆💦🍑")
    # #mergeImage("background.png", "mapWithText.png", "mapWithBackground.png")
    # mergeImage("background.png", "mapNoBackground.png", "mapWithBackground.png")

    print("Done! ❤  Generation 🕓 taken: " + str(datetime.now()-start) +" ❤")

if __name__ == "__main__":
    main()

max_stars_repo_name: CivCamo/Civclassics-Map-Generator

max_issues_repo_name: CivCamo/Civclassics-Map-Generator

max_forks_repo_name: CivCamo/Civclassics-Map-Generator


text:
from typing import cast, List

import numpy as np
from numpy.linalg import norm

from preprocessor_meta import CORR_NUM_WINDOWS, CORR_WINDOW_LEN, CORR_MAX_LAG

def lagged_cosine(x: np.ndarray, y: np.ndarray, lag_y: int) -> float:
    xi0 = max(0, lag_y)
    xi1 = len(x) + min(0, lag_y)

    yi0 = max(0, -lag_y)
    yi1 = len(y) + min(0, -lag_y)

    x_slice = x[xi0:xi1]
    y_slice = y[yi0:yi1]

    x_normed = x_slice / norm(x_slice)
    y_normed = y_slice / norm(y_slice)

    dot_product = np.dot(x_normed, y_normed)
    return cast(float, dot_product)

def scan_lagged_cosine(x: np.ndarray, y: np.ndarray, max_lag: int) -> np.ndarray:
    cs = [lagged_cosine(x, y, lag_y) for lag_y in range(-max_lag, max_lag + 1)]
    return np.array(cs)

def correlate(sig1: np.ndarray, sig2: np.ndarray) -> np.ndarray:
    """
    Computes an np.ndarray shaped (NUM_LAGS, WINDOWS) with dimensions as defined for `process`
    """
    if len(sig1.shape) != 1 or sig1.shape != sig2.shape:
        raise ValueError(f"Invalid shapes: {sig1.size} and {sig2.size}")
    window_lags_list: List[np.ndarray] = []
    for window in range(CORR_NUM_WINDOWS):
        i0 = window * CORR_WINDOW_LEN
        i1 = (window + 1) * CORR_WINDOW_LEN
        window_lags_list.append(
            scan_lagged_cosine(sig1[i0:i1], sig2[i0:i1], CORR_MAX_LAG)
        )
    return np.stack(window_lags_list, axis=1)

def process(sigs: np.ndarray) -> np.ndarray:
    """
    Computes an np.ndarray of type float, shaped (NUM_LAGS, 2, WINDOWS), where

    * the indices represent (channel, signal pair, width)

    * the signal pair dimension represents comparing signals (0, 1) and then (0, 2)

    :param sigs: output from filter_sig
    """
    results = [correlate(sigs[0], sigs[i2]) for i2 in (1, 2)]
    return np.stack(results, axis=1)

max_stars_repo_name: wisdom-parts/kaggle-gw

max_issues_repo_name: wisdom-parts/kaggle-gw

max_forks_repo_name: wisdom-parts/kaggle-gw

reason: import numpy,from numpy


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : tql-ANN.
# @File         : ANN
# @Time         : 2019-12-04 20:05
# @Author       : yuanjie
# @Email        : yuanjie@xiaomi.com
# @Software     : PyCharm
# @Description  : https://shenchunxu.blog.csdn.net/article/details/102881115?utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-1.no_search_link&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-1.no_search_link

# todo: 增加idmapping并提供在线服务

import faiss
import numpy as np

class ANN(object):
    """Flat支持: https://github.com/Jie-Yuan/faiss_note/blob/master/4.Faiss%20indexes%20%E8%BF%9B%E9%98%B6%E6%93%8D%E4%BD%9C.ipynb
    恢复原数据
    从index中移除向量
    搜索距离范围内的向量
    拆分/合并index

    cpu_index.make_direct_map()
    cpu_index.reconstruct(0)
    ann.index.reconstruct
    ann.index.reconstruct_c
    ann.index.reconstruct_n
    ann.index.reconstruct_n_c
    ann.index.search_and_reconstruct
    ann.index.search_and_reconstruct_c

    数据集的大小
    在高效检索的index中，聚类是其中的基础操作，数据量的大小主要影响聚类过程。

    如果小于1M， 使用"...,IVFx,..."
    N是数据集中向量个数，x一般取值[4sqrt(N),16sqrt(N)],需要30x ～ 256x个向量的数据集去训练。

    如果在1M-10M，使用"...,IMI2x10,..."
    使用k-means将训练集聚类为2^10个类，但是执行过程是在数据集的两半部分独立执行，即聚类中心有2^(2*10)个。

    如果在10M-100M，使用"...,IMI2x12,..."

    一个簇至少39个样本：一般2的N次方，我们一般用2**14或者更小一点 >> topk
    经验公式 2 ** (np.log2(n) // 2 + 2)
    """

    def __init__(self):
        self.nogpu_index_factory = {'HNSW', 'SQ'}

    def train(self, data, index_factory='Flat', metric=0, noramlize=False):
        """

        :param data:
        :param index_factory: 'IVF1024,Flat'
            https://www.cnblogs.com/houkai/p/9316172.html
            https://blog.csdn.net/xiaoxu2050/article/details/84982478
            https://github.com/facebookresearch/faiss/wiki/Faiss-indexes
            https://github.com/liqima/faiss_note/blob/master/4.Faiss%20indexes%20IO%E5%92%8Cindex%20factory.ipynb
        :param metric:
            faiss.METRIC_BrayCurtis
            faiss.METRIC_Canberra
            faiss.METRIC_INNER_PRODUCT: L2之后内积≈cosine 0 默认内积
            faiss.METRIC_JensenShannon
            faiss.METRIC_L1
            faiss.METRIC_L2
            faiss.METRIC_Linf
            faiss.METRIC_Lp
        :return:
        """
        if noramlize:
            data = self.noramlize(data)
        assert data.dtype == 'float32', "TODO: np.array([]).astype('float32')"
        dim = data.shape[1]
        args = [dim, index_factory, metric] if metric is not None else [dim, index_factory]

        self.index = faiss.index_factory(*args)

        if faiss.get_num_gpus() > 0:
            if any(index_factory.__contains__(i) for i in self.nogpu_index_factory):
                pass
                print(f"Donot Support GPU: {index_factory}")
            else:
                # gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
                self.index = faiss.index_cpu_to_all_gpus(self.index)

        print("Train ...")
        self.index.train(data)
        self.index.add(data)

        print(f"Ntotal: {self.index.ntotal}")

    def search(self, data, topK=10, nprobe=1, k_factor=1):
        """

        :param data:
        :param topK:
        :param nprobe: nprobe参数始终是调整速度和结果精度之间权衡的一种方式。
        :param k_factor:
        :return:
        """
        self.index.k_factor = k_factor  # 搜索阶段会首先选取 k_factor*topK，重排序
        self.index.nprobe = nprobe  # default nprobe is 1, try a few more
        return self.index.search(data, topK)

    def write_index(self, file_name="index_file.index"):
        # faiss.index_cpu_to_all_gpus(index)
        if faiss.get_num_gpus() > 0:
            index = faiss.index_gpu_to_cpu(self.index)  # gpu版本？
            faiss.write_index(index, str(file_name))
        else:
            faiss.write_index(self.index, str(file_name))

    def read_index(self, file_name="index_file.index"):
        index = faiss.read_index(str(file_name))
        # if faiss.get_num_gpus() > 0:
        #     index = faiss.index_cpu_to_gpu()
        # index_new = faiss.clone_index(index) # 复制索引
        return index

    def noramlize(self, x):
        if len(x.shape) > 1:
            return x / np.clip(x ** 2, 1e-12, None).sum(axis=1).reshape((-1, 1) + x.shape[2:]) ** 0.5
        else:
            return x / np.clip(x ** 2, 1e-12, None).sum() ** 0.5

if __name__ == '__main__':
    import numpy as np

    data = np.random.random((1000, 128)).astype('float32')

    ann = ANN()
    ann.train(data, noramlize=True)

    dis, idx = ann.search(data[:10])

    print(dis)
    print(idx)

    ann.write_index()

max_stars_repo_name: Jie-Yuan/MeUtils

max_issues_repo_name: Jie-Yuan/MeUtils

max_forks_repo_name: Jie-Yuan/MeUtils


text:
import glob
import json
import os.path
import math
from functools import reduce
from typing import Union
from collections.abc import Iterable
from datetime import datetime

import numpy as np
import scipy.interpolate
from matplotlib import transforms, rcParams
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, Ellipse
from matplotlib.widgets import Button
import seaborn as sns
import pint
from icecream import ic

from scripts.data_path import *

rcParams['figure.constrained_layout.use'] = True
rcParams['figure.dpi'] = 100
rcParams['font.size'] = 12
sns.set_style('darkgrid')

inf = float('inf')

reg = pint.UnitRegistry()

def unit_converter(m, src=reg.inch, dst=reg.meter):
    return (m * src).to(dst).magnitude

def json_load(fnm):
    f = open(fnm, 'r')
    scans = json.load(f)
    f.close()
    return scans

def now(as_str=True):
    d = datetime.now()
    return d.strftime('%Y-%m-%d %H:%M:%S') if as_str else d

def get(dic, ks):
    # return reduce(lambda acc, elm: acc[elm], keys, dic)
    return reduce(lambda acc, elm: acc[elm] if elm in acc else None, ks.split('.'), dic)

def config(attr):
    """
    Retrieves the queried attribute value from the config file.

    Loads the config file on first call.
    """
    if not hasattr(config, 'config'):
        with open(f'{PATH_BASE}/{DIR_PROJ}/config.json') as f:
            config.config = json.load(f)
    return get(config.config, attr)
    # ic('here', attr)
    # if isinstance(ret, dict) and list(ret.keys())[0].isdigit():  # For JSON write
    #     ic('here')
    #     ret = {int(k): v for k, v in ret.items()}
    # return ret

def eg_hsr_scan(k1=0, k2=77):
    """
    :return: Example HSR laser scan as 2D coordinates, given file name and measurement number
    """
    path = os.path.join(PATH_BASE, DIR_DATA)
    fls = sorted(glob.iglob(f'{path}/{config(f"{DIR_DATA}.eg.HSR.fmt")}', recursive=True))
    hsr_scans = json_load(fls[k1])
    s = hsr_scans[k2]
    return laser_polar2planar(s['angle_max'], s['angle_min'])(np.array(s['ranges']))

def pts2max_dist(pts):
    """
    :param pts: List of 2d points
    :return: The maximum distance between any two pairs of points
    """
    assert pts.shape[1] == 2

    def dist(a, b):
        return (a[0] - b[0])**2 + (a[1] - b[1])**2
    n = pts.shape[0]
    idxs = ((i, j) for i in range(n-1) for j in range(i, n))
    # ic(list(idxs))
    return math.sqrt(max(dist(pts[a], pts[b]) for a, b in idxs))

def clipper(low, high):
    """
    :return: A clipping function for range [low, high]
    """
    return lambda x: max(min(x, high), low)

def get_3rd_side(a, b):
    """
    Returns hypotenuse of a right-angled triangle, given its other sides
    """
    # return np.sqrt(np.sum(np.square(mags)))
    return math.sqrt(a**2 + b**2)

class JsonWriter:
    """
    Each time the object is called, the data is appended to the end of a list which is serialized to a JSON file
    """

    def __init__(self, fnm):
        self.fnm = fnm
        self.data = []
        self.fnm_ext = f'data/{self.fnm}.json'
        open(self.fnm_ext, 'a').close()  # Create file in OS

    def __call__(self, data):
        f = open(self.fnm_ext, 'w')
        self.data.append(data)
        json.dump(self.data, f, indent=4)
        # ic(self.data)

def laser_scan2dict(data):
    """
    :param data: Of type [sensor_msgs/LaserScan](https://docs.ros.org/en/noetic/api/sensor_msgs/html/msg/LaserScan.html)
    """
    h = data.header
    d_h = dict(
        seq=h.seq,
        stamp=dict(
            secs=h.stamp.secs,
            nsecs=h.stamp.nsecs
        ),
        frame_id=h.frame_id
    )
    return dict(
        header=d_h,
        angle_min=data.angle_min,
        angle_max=data.angle_max,
        angle_increment=data.angle_increment,
        time_increment=data.time_increment,
        scan_time=data.scan_time,
        range_min=data.range_min,
        ranges=data.ranges,
        intensities=data.intensities
    )

def extend_1s(arr):
    """
    Return array with column of 1's appended
    :param arr: 2D array
    """
    return np.hstack([arr, np.ones([arr.shape[0], 1])])

def cartesian(arrs: list, out=None):
    """
    :param arrs: list of 1D arrays
    :param out: Array to place the cartesian product in.
    :return: Cartesian product of `arrs` of shape

    Modified from https://stackoverflow.com/a/1235363/10732321
    """
    arrs = [np.asarray(x) for x in arrs]
    n = np.prod([x.size for x in arrs])
    if out is None:
        out = np.zeros([n, len(arrs)], dtype=arrs[0].dtype)

    m = int(n / arrs[0].size)
    out[:, 0] = np.repeat(arrs[0], m)
    if arrs[1:]:
        cartesian(arrs[1:], out=out[0:m, 1:])
        for j in range(1, arrs[0].size):
            out[j*m:(j+1)*m, 1:] = out[0:m, 1:]
    return out

def polar2planar(dist, angle):
    return (
        dist * np.cos(angle),
        dist * np.sin(angle)
    )

def laser_polar2planar(a_max, a_min, split=False):
    """
    :param a_max: Maximum angle
    :param a_min: Minimum angle
    :param split: If True, the function returns a 2-tuple of x and y coordinates
    :return: A function that returns an array of 2D points

    Assumes the angles are [a_min, a_max)
    """
    def _get(ranges):
        """
        :param ranges: Array of laser scan ranges

        Number of beams taken from size of `range`;
        """
        theta = np.linspace(a_min, a_max, num=ranges.size + 1)[:-1]
        x, y = polar2planar(ranges, theta)
        return (x, y) if split else np.vstack([x, y]).T
    return _get

def rot_mat(theta):
    c, s = np.cos(theta), np.sin(theta)
    return np.array([
        [c, -s],
        [s, c]
    ])

def tsl_n_angle2tsf(tsl: Iterable = np.array([0, 0]), theta: Union[int, float] = 0):
    """
    Converts translation in 2D & an angle into matrix transformation

    :param tsl: 3-array of (translation_x, translation_y, theta),
        or 2-array of (translation_x, translation_y)
    :param theta: Angle in radians
    """
    tsl = np.asarray(tsl)
    tsf = np.identity(3)
    tsf[:2, 2] = tsl[:2]
    # ic(tsl[-1] if tsl.size == 3 else theta)
    tsf[:2, :2] = rot_mat(tsl[-1] if tsl.size == 3 else theta)
    return tsf

def tsf2tsl_n_angle(tsf):
    """
    :return: 2-tuple of 2D translation and angle in radians from transformation matrix
    """
    return tsf[:2, 2], math.acos(tsf[0][0])

def apply_tsf_2d(arr, tsf):
    """
    Syntactic sugar
    :param arr: Array of 2D points
    :param tsf: Transformation matrix in R^{3 x 3}
    :return: Array of 2D points with transformation matrix applied
    """
    return (extend_1s(arr[:, :2]) @ tsf.T)[:, :2]

def get_kuka_pointcloud():
    # d_dim = config('dimensions.KUKA')
    # return get_rect_pointcloud(d_dim['length'], d_dim['width'])
    return get_rect_pointcloud(config('dimensions.KUKA'))

def get_rect_pointcloud(dim, n=240, visualize=False):
    """
    :param dim: 2-tuple of (length, width) of dict with keys `length` and `width`
    :param n: Number of points/beams
    :param visualize: If True, shows an illustration of the process
    :return: Array of 2D points of a rectangular contour, as if by a 360 degree of beams
    """
    ln, wd = dim['length'], dim['width'] if isinstance(dim, dict) else dim
    r = max(ln, wd)
    r = np.full(n, r)
    theta = np.linspace(0, 2 * math.pi, num=n+1)[:-1]
    x, y = polar2planar(r, theta)
    boundaries = (-ln/2, -wd/2, ln/2, wd/2)

    def intersec_rect(left, bot, right, top):
        """ :return: function that returns the intersection of point relative to a rectangle """
        def _get(x_, y_):
            """
            x, y should be outside of the rectangle
            """
            ct_x = (left + right) / 2
            ct_y = (bot + top) / 2
            slope = (ct_y - y_) / (ct_x - x_)

            if x_ <= ct_x:
                y__ = slope * (left - x_) + y_
                if bot <= y__ <= top:
                    return left, y__
            if x_ >= ct_x:
                y__ = slope * (right - x_) + y_
                if bot <= y__ <= top:
                    return right, y__
            if y_ <= ct_y:
                x__ = (bot - y_) / slope + x_
                if left <= x__ <= right:
                    return x__, bot
            if y_ >= ct_y:
                x__ = (top - y_) / slope + x_
                if left <= x__ <= right:
                    return x__, top
            if x_ == ct_x and y_ == ct_y:
                return x_, y_
        return _get

    if visualize:
        fig, ax = plt.subplots(figsize=(16, 9), constrained_layout=True)
        for x_i, y_i in zip(x, y):
            x_int, y_int = intersec_rect(*boundaries)(x_i, y_i)
            ax.add_patch(Rectangle((-ln/2, -wd/2), ln, wd, edgecolor='b', fill=False))
            ax.plot((0, x_int), (0, y_int), marker='o', c='c', ms=2, lw=0.5, ls='dotted')
            ax.plot((x_i, x_int), (y_i, y_int), marker='o', c='orange', ms=2, ls='dotted')
        plt.gca().set_aspect('equal')
        plt.show()
    intersec = intersec_rect(*boundaries)
    return np.apply_along_axis(lambda i: intersec(*i), 1, np.vstack([x, y]).T)

def save_fig(save, title):
    if save:
        fnm = f'{title}.png'
        plt.savefig(os.path.join(PATH_BASE, DIR_PROJ, 'plot', fnm), dpi=300)

def plot_points(arr, **kwargs):
    """
    :param arr: Array of 2d points to plot
    :param kwargs: Arguments are forwarded to `matplotlib.axes.Axes.plot`
    """
    arr = np.asarray(arr)
    kwargs_ = dict(
        marker='.', lw=0.5, ms=1,
        c='orange',
    )
    plt.plot(arr[:, 0], arr[:, 1], **(kwargs_ | kwargs))

def plot_2d(arr, label=None, title=None, save=False, show=True, **kwargs):
    """ Plot potentially list pf 2D points """
    def _plot(a, lb):
        plt.plot(a[:, 0], a[:, 1], marker='o', ms=0.3, lw=0.25, label=lb, **kwargs)
    plt.figure(figsize=(16, 9), constrained_layout=True)
    if not isinstance(arr, list):
        arr = [arr]
    lbl = [None for _ in arr] if label is None else label
    _ = [_plot(a, lb) for a, lb in zip(arr, lbl)]  # Execute
    if label:
        plt.legend()
    if title:
        plt.title(title)
    plt.gca().set_aspect('equal')
    save_fig(save, title)
    if show:
        plt.show()

def plot_points3d(arr, **kwargs):
    """
    :param arr: Array of 3d points to plot
    :param kwargs: Arguments are forwarded to `matplotlib.axes.Axes.plot`
    """
    arr = np.asarray(arr)
    kwargs_ = dict(
        marker='.', lw=0.5, ms=1,
        c='orange',
    )
    plt.plot(arr[:, 0], arr[:, 1], arr[:, 2], **(kwargs_ | kwargs))

def plot_line_seg(c1, c2, with_arrow=True, **kwargs):
    kwargs_ = dict(
        ls='dotted', marker='o', lw=1, ms=2,
        c='orange',
    )
    kwargs = kwargs_ | kwargs
    plt.plot((c1[0], c2[0]), (c1[1], c2[1]), **kwargs)
    if with_arrow:
        plot_line_seg_arrow(c1, c2, color=get(kwargs, 'c'), alpha=get(kwargs, 'alpha'))

def plot_line_seg_arrow(c1, c2, r=0.01, **kwargs):
    coords = np.array([c1, c2])
    mean = coords.mean(axis=0)
    mags = (coords[1] - coords[0]) * r

    width = 5 * get_3rd_side(*mags)
    if not hasattr(plot_icp_result, 'clp'):
        plot_icp_result.clp = clipper(0.01, 0.05)
    width = plot_icp_result.clp(width)

    kwargs_ = dict(
        alpha=0.5,
        # head_width=0.05,
        head_width=width,
        length_includes_head=True,
        lw=0,
        overhang=0.2,
    )
    plt.arrow(
        *(mean-mags/2), *mags,
        **(kwargs_ | kwargs)
    )

def plot_icp_result(
        src, tgt, tsf,
        title=None, save=False, states=None, xlim=None, ylim=None, with_arrow=True, show=True,
        init_tsf=np.identity(3), mode='static', scl=1
):
    """
    :param src: Source coordinates
    :param tgt: Target coordinates
    :param tsf: ICP result transformation
    :param title: Plot title
    :param save: If true, plot saved as image
    :param states: A list of source-target matched points & transformation for each iteration
    :param xlim: X limit for plot, inferred if not given
    :param ylim: Y limit for plot, inferred if not given
    :param with_arrow: If true, matched points are shown with arrows
    :param show: If true, show the plot
    :param init_tsf: Initial transformation guess for ICP
    :param mode: Plotting mode, one of [`static`, `animate`, `control`]
    :param scl: Plot window scale

    .. note:: Assumes 2d data
    """

    def _plot_matched_points(stt, **kwargs):
        src__, tgt__ = stt[0], stt[1]
        for s_, t_ in zip(src__, tgt__):
            plot_line_seg(s_, t_, with_arrow=with_arrow, **kwargs)

    N_STT = (states and len(states)) or 0

    x_rang, y_rang = (
        abs(xlim[0] - xlim[1]), abs(ylim[0] - ylim[1])
    ) if xlim and ylim else (
        np.ptp(tgt[:, 0]), np.ptp(tgt[:, 1])
    )
    ratio = 1 / x_rang * y_rang
    d = 8 * scl
    plt.figure(figsize=(d, d * ratio), constrained_layout=True)
    # ic((d, d * ratio))
    plt.xlabel('Target space dim 1 (m)')
    plt.ylabel('Target space dim 2 (m)')
    t = 'ICP results'
    if title:
        t = f'{t}, {title}'
    plt.suptitle(t)
    plt.gca().set_aspect('equal')

    def _step(idx_, t_):
        tsf_ = states[idx_][2] if states else tsf
        plt.cla()
        if xlim:
            plt.xlim(xlim)
        if ylim:
            plt.ylim(ylim)
        if mode == 'control':
            t_ = f'{t_}, iteration {idx_}'
            plt.suptitle(t_)

        tsl, theta = tsf2tsl_n_angle(tsf_)
        ic(tsf_, tsl, math.degrees(theta))

        unit_sqr = np.array([
            [0, 0],
            [0, 1],
            [1, 1],
            [1, 0],
            [0, 0]
        ])
        unit_sqr_tsf = (extend_1s(unit_sqr) @ tsf_.T)[:, :2]

        cs = iter(sns.color_palette(palette='husl', n_colors=5))
        c = next(cs)
        plot_points(src, c=c, alpha=0.5, label='Source points')
        if not np.array_equal(init_tsf, np.identity(3)):
            plot_points(apply_tsf_2d(src, init_tsf), c=c, alpha=0.5, label='Source points, initial guess')
        plot_points(apply_tsf_2d(src, tsf_), c=c, label='Source points, transformed')
        plot_points(tgt, c=next(cs), label='Target points')

        c = next(cs)
        if states:
            _plot_matched_points(states[0], c=c, ms=1, ls='solid', alpha=0.5, label='Matched points, initial')
            _plot_matched_points(states[idx_], c=c, ms=1, ls='solid', label='Matched points, final')

        c = next(cs)
        plt.plot(0, 0, marker='o', c=c, ms=4)
        plot_points(unit_sqr, ms=0, marker=None, c=c, alpha=0.6, label='Unit square')
        plot_points(unit_sqr_tsf, ms=0.5, marker=None, c=c, alpha=0.9, label='Unit square, transformed')
        for i in zip(unit_sqr, unit_sqr_tsf):
            plot_line_seg(*i, with_arrow=with_arrow, marker=None, c=c, alpha=0.5)

        handles, labels_ = plt.gca().get_legend_handles_labels()  # Distinct labels
        by_label = dict(zip(labels_, handles))
        plt.legend(by_label.values(), by_label.keys())

        save_fig(save, t_)

        if mode != 'static':
            plt.pause(1 if mode == 'animate' else 0.1)  # 'control'

    if mode == 'animate':
        plt.ion()
        for idx in range(N_STT):
            _step(idx, t)
    elif mode == 'control':
        class PlotFrame:
            def __init__(self, i=0):
                self.idx = i
                self.clp = clipper(0, N_STT-1)

            def next(self, event):
                prev_idx = self.idx
                self.idx = self.clp(self.idx+1)
                if prev_idx != self.idx:
                    _step(self.idx, t)

            def prev(self, event):
                prev_idx = self.idx
                self.idx = self.clp(self.idx-1)
                if prev_idx != self.idx:
                    _step(self.idx, t)
        init = 0
        pf = PlotFrame(i=init)
        ax = plt.gca()
        btn_next = Button(plt.axes([0.81, 0.05, 0.1, 0.075]), 'Next')
        btn_next.on_clicked(pf.next)
        btn_prev = Button(plt.axes([0.7, 0.05, 0.1, 0.075]), 'Previous')
        btn_prev.on_clicked(pf.prev)
        plt.sca(ax)

        _step(init, t)
    else:
        _step(N_STT-1, t)

    plt.ioff()  # So that window doesn't close
    if show:
        plt.show()

def plot_cluster(
        data, data_labels,
        title=None, save=False, new_fig=True, show_eclipse=True, line_kwargs=None,
        cls_kwargs=None
):
    d_clusters = {lb: data[np.where(data_labels == lb)] for lb in np.unique(data_labels)}

    cs = iter(sns.color_palette(palette='husl', n_colors=len(d_clusters) + 1))
    x, y = data[:, 0], data[:, 1]
    if new_fig:
        fig, ax = plt.subplots(figsize=(12, 12 / np.ptp(x) * np.ptp(y)), constrained_layout=True)
    else:
        ax = plt.gca()
    # ic(ex)
    if line_kwargs is None:
        line_kwargs = dict()
    plt.plot(x, y, **(dict(marker='o', ms=0.3, lw=0.25, c=next(cs), alpha=0.5, label='Whole') | line_kwargs))
    line_kwargs = dict(
        marker='o', ms=0.4, lw=0.25
    ) | line_kwargs

    # if cls_kwargs is not None:  # Kwargs for each cluster

    # labels = None, colors = None
    # if labels is not None and not isinstance(labels, list):
    #     labels = [labels] * len(data_labels)
    # labels = None, colors = None
    if cls_kwargs is not None and not isinstance(cls_kwargs, list):
        cls_kwargs = [cls_kwargs] * len(data_labels)

    for idx, (lb, d) in enumerate(d_clusters.items()):
        x_, y_ = d[:, 0], d[:, 1]
        c = next(cs)

        def confidence_ellipse(n_std=1., **kwargs):
            """
            Modified from https://matplotlib.org/stable/gallery/statistics/confidence_ellipse.html
            Create a plot of the covariance confidence ellipse of x and y

            :param n_std: number of standard deviations to determine the ellipse's radius'
            :return matplotlib.patches.Ellipse
            """
            cov = np.cov(x_, y_)
            pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
            ell_radius_x = np.sqrt(1 + pearson)
            ell_radius_y = np.sqrt(1 - pearson)
            ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2,
                              **(dict(fc='none') | kwargs))

            tsf = transforms.Affine2D().rotate_deg(45)
            tsf = tsf.scale(
                np.sqrt(cov[0, 0]) * n_std,
                np.sqrt(cov[1, 1]) * n_std
            )
            tsf = tsf.translate(np.mean(x_), np.mean(y_))

            ellipse.set_transform(tsf + ax.transData)
            return ax.add_patch(ellipse)

        if show_eclipse and lb != -1:  # Noise as in DBSCAN
            confidence_ellipse(n_std=1.25, fc=c, alpha=0.25)

        # lb = labels[idx] if labels is not None else f'Cluster {lb + 1}'
        lb = f'Cluster {lb + 1}'
        cls_kwarg = (cls_kwargs is not None and cls_kwargs[idx]) or dict()
        ax.plot(x_, y_, **(dict(c=c, label=lb) | line_kwargs | cls_kwarg))

    if new_fig:
        t = 'Clustering results'
        if title:
            t = f'{t}, {title}'
        plt.title(t)
        plt.legend()
        if not hasattr(ax, 'get_zlim'):  # Not supported in 3D projection
            ax.set_aspect('equal')
        save_fig(save, t)
        plt.show()

def scale(arr, factor=4, as_sz=False):
    """
    :param arr: 1D array of uniform values
    :param factor: Factor to sample
    :param as_sz: If true, return the size of scaled aray
    :return: 1D array of `arr` with a finer sample,
        in particular, the distance between two adjacent points is `factor` times smaller
    """
    num = (arr.size - 1) * factor + 1
    return num if as_sz else np.linspace(arr[0], arr[-1], num=num)

def interpolate(X, Y, Z, x_coords, y_coords, factor=4, method='cubic'):
    """
    :return: Interpolated X, Y, Z coordinate tuples, given by a factor, and a method
        in `scipy.interpolate.griddata`
    """
    X_f, Y_f, Z_f = X.flatten(), Y.flatten(), Z.flatten()
    x_inter = scale(x_coords, factor=factor).reshape(1, -1)
    y_inter = scale(y_coords, factor=factor).reshape(-1, 1)
    X_, Y_ = np.meshgrid(x_inter, y_inter)
    Z_ = scipy.interpolate.griddata((X_f, Y_f), Z_f, (x_inter, y_inter), method=method)
    return X_, Y_, Z_

def get_offset(arr, frac=2**4):
    """
    :param arr: Array-like
    :param frac: Difference between range of `arr` and the offset
    :return: 2-tuple offset value pair for `arr`, with a relative gap factor,
        in the order of (`down`, `up`)
    """
    ma, mi = arr.max(), arr.min()
    diff = (ma-mi) / frac
    return mi-diff, ma+diff

def pts2bins(pts, prec=0.25):
    """
    :param pts: List of 2d points
    :param prec: Width of each bin
    :return: A dictionary of bin centers and points in the bin
    """
    # for pt in pts:
        # def val2bin_int(val):
        #     return math.ceil(val / prec) if val > 0 else math.floor(pt[1] / prec)
        # f =
        # loc = math.floor(pt[0] / prec), math.floor(pt[1] / prec)
        # ic(loc, pt)
    locs_pt = {tuple(pt.tolist()): (math.floor(pt[0] / prec), math.floor(pt[1] / prec)) for pt in pts}
    d_bin = dict()
    for pt, loc in locs_pt.items():
        # ic(pt, loc)
        if loc in d_bin:
            d_bin[loc].append(pt)
        else:
            d_bin[loc] = [pt]

    def unit2val(u):
        return (u+0.5)*prec
    return {(unit2val(k1), unit2val(k2)): v for (k1, k2), v in d_bin.items()}

def plot_grid_search(
        pcr, pts, opns_x, opns_y, opns_ang, errs, labels=None,
        interp=True, inverse_loss=False, inverse_pts=False,
        title=None, save=False, zlabel='Loss', offset_frac=2**3,
        tsf_ideal=None,
        interp_kwargs=None,
        plot3d_kwargs=None
):
    """
    Plot grid search result per `PoseEstimator.FusePose`, i.e. for the get-pose first approach

    The 3 search spaces, `opns_x`, `opns_y`, `opns_ang`, should be uniformly sampled and increasing numbers

    Plots loss against translation (x, y pair), for each setup, the angle/rotation with the lowest loss is picked
    """
    label_idxs = None
    if labels is not None:
        assert len(errs) == 2
        label_idxs, errs = errs
    if inverse_pts:  # Illustrate as if ICP ran in the reversed order
        opns_x = -opns_x[::-1]
        opns_y = -opns_y[::-1]
        errs = errs[::-1]
        pcr, pts = pts, pcr

    errs_by_tsl = np.min(errs.reshape(-1, opns_ang.size * len(label_idxs) if labels is not None else 1), axis=-1)
    errs_by_tsl = errs_by_tsl.reshape(-1, opns_x.size)  # Shape flipped for `np.meshgrid`
    d = 12
    fig, ax = plt.subplots(figsize=(d, d), subplot_kw=dict(projection='3d'))
    [X, Y], Z = np.meshgrid(opns_x, opns_y), errs_by_tsl  # Negated cos lower error = better

    if inverse_loss:
        Z = -Z

    interp_kwargs = dict(
        factor=2**3
    ) | (dict() if interp_kwargs is None else interp_kwargs)
    if interp:
        X, Y, Z = interpolate(X, Y, Z, opns_x, opns_y, **interp_kwargs)

    bot, top = get_offset(Z, frac=offset_frac)  # Level/Height of contour plot & 2d point plots, respectively
    if plot3d_kwargs is None:
        plot3d_kwargs = dict()
    ord_3d, ord_2d = 1, 20
    kwargs_surf = dict(
        zorder=ord_3d, antialiased=True,
        alpha=0.9, cmap='Spectral_r', edgecolor='black', lw=0.3
    ) | plot3d_kwargs
    kwargs_cont = dict(
        zorder=ord_3d, antialiased=True,
        linewidths=1, levels=np.linspace(Z.min(), Z.max(), 2 ** 4), offset=bot, zdir='z',
        cmap='Spectral_r'
    ) | plot3d_kwargs
    surf = ax.plot_surface(X, Y, Z, **kwargs_surf)
    ax.contour(X, Y, Z, **kwargs_cont)

    cp = sns.color_palette(palette='husl', n_colors=7)
    cp = list(reversed(cp)) if inverse_loss else cp
    cs = iter(cp)

    lb_tgt = 'Laser scan, target'
    plot_points([[0, 0]], zs=top, zorder=ord_2d, ms=10, alpha=0.5)
    c = next(cs)
    plot_points(pcr, zorder=ord_2d, zs=top, c=c, alpha=0.5, label='Point cloud representation, source')
    if labels is not None:
        plot_cluster(pts, labels, new_fig=False, show_eclipse=False, line_kwargs=dict(
            zorder=ord_2d, zs=top, label=lb_tgt
        ))
    else:
        plot_points(pts, zorder=ord_2d, zs=top, c=next(cs), label=lb_tgt)
    if tsf_ideal is not None:  # Illustrate the ideal translation+-
        plot_points(
            apply_tsf_2d(pcr, tsf_ideal),
            zorder=ord_2d, zs=top, c=c, alpha=0.7,
            label='Point cloud representation at actual pose'
        )

        bot, top = get_offset(Z, frac=offset_frac/1.25)
        ran = max(opns_x[-1] - opns_x[0], opns_y[-1] - opns_y[0]) / 2**4
        wd = pts2max_dist(pcr)/2
        wd = min(ran, wd)

        segs = np.array([
            [wd/2, 0, bot],
            [0, 0, bot],
            [0, 0, top],
            [wd/2, 0, top],
        ])
        api = 'Actual pose indicator'
        kwargs_origin = dict(zorder=ord_2d, ms=10)
        plot_points([tsf2tsl_n_angle(tsf_ideal)[0]], zs=top, **kwargs_origin)
        plot_points([tsf2tsl_n_angle(tsf_ideal)[0]], zs=bot, **kwargs_origin)
        segs[:, :2] = apply_tsf_2d(segs, tsf_ideal)
        plot_points3d(segs, zorder=ord_2d, c='black', ls='dashed', lw=1, label=api)

    fig.colorbar(surf, shrink=0.5, aspect=2**5, pad=2**-4)
    plt.xlabel('Translation in X (m)')
    plt.ylabel('Translation in y (m)')
    ax.set_zlabel(zlabel + ', inverted' if inverse_loss else '')
    plt.legend()
    handles, labels_ = plt.gca().get_legend_handles_labels()  # Distinct labels
    by_label = dict(zip(labels_, handles))
    plt.legend(by_label.values(), by_label.keys())

    def prec_pair(nm, mag):
        return r'$\Delta ' + nm + r' = ' + str(mag) + r'$'

    # prec = str(tuple([opns_x[1]-opns_x[0], opns_y[1]-opns_y[0], opns_ang[1]-opns_ang[0]]))
    prec = f'{prec_pair("x", opns_x[1]-opns_x[0])}, ' \
           f'{prec_pair("y", opns_y[1]-opns_y[0])}'
    prec += ', ' + prec_pair(r'\theta', round(opns_ang[1]-opns_ang[0], 2))
    t = r'Loss against translation grid search, by best $\theta$'
    t = f'{t} for {prec}'
    if title:
        t = f'{t}, {title}'
    plt.title(t)
    save_fig(save, t)
    plt.show()

if __name__ == '__main__':
    def _create():
        jw = JsonWriter('jw')

        for d in [1, 'as', {1: 2, "a": "b"}, [12, 4]]:
            jw(d)

    def _check():
        with open('jw.json') as f:
            l = json.load(f)
            ic(l)

    # _create()
    # _check()

    def _kuka_pointcloud():
        # pc = get_rect_pointcloud(2, 0.8, visualize=False)
        ptc = get_kuka_pointcloud()
        plt.figure(figsize=(16, 9), constrained_layout=True)
        plt.plot(ptc[:, 0], ptc[:, 1], marker='o', ms=1, lw=0.5)
        plt.show()

    # ic(config('dimensions.KUKA.length'))
    # ic(cartesian([[1, 2, 3], [4, 5], [6, 7]]))

max_stars_repo_name: StefanHeng/Multi-robot-Pose-Estimation

max_issues_repo_name: StefanHeng/Multi-robot-Pose-Estimation

max_forks_repo_name: StefanHeng/Multi-robot-Pose-Estimation

reason: import numpy,import scipy


text:
import numpy as np
from scipy import signal as signal

"""
 ECG R-peak detector and heart rate detector

 The detector uses the matched filter approach by creating
 an IIR bandpass filter which looks like an R peak and
 thus is a recursive matched filter. One could also say
 it's a causal wavelet or perhaps just a bandpass filter
 which covers the frequency range of the R peak. It's all
 the same in different forms!
 As an input the detector just gets the data samples
 at a given sampling rate and then it detects the r-peak and
 heart rate from it.
 It also has a callback function which is called when
 a heartbeat is detected (implemented as a listener).
"""

class heartrate_detector:

    def __init__(self,_fs):
        # how fast the adaptive threshold follows changes in ECG
        # amplitude. Realisic values: 0.1 .. 1.0
        # 0.1 = slow recovery after an artefact but no wrong detections
        # 1 = fast recovery after an artefact but possibly wrong detections
        self.adaptive_threshold_decay_constant = 0.25

        # the threshold for the detection is 0.6 times smaller than the amplitude
        self.threshold_factor = 0.6

        # 0.5mV as the thereshold the bandpass filtered ECG is an artefact
        self.artefact_threshold = 1

        # ignores 1000 samples to let the filter settle
        self.ignoreECGdetector = 1000

        #adaptive amplitude value of the detector output
        self.amplitude = 0.0

        def getAmplitude(self):
            return amplitude
        self.timestamp = 0

        # previous timestamp
        self.t2 = 0

        # timewindow not to detect an R peak
        self.doNotDetect = 0

        self.ignoreRRvalue = 2

        # create a 2nd order order bandpass filter
        center = 20
        width = 15
        f1 = center-width/2
        f2 = center+width/2
        self.bp_b, self.bp_a = signal.butter(2, [f1/_fs*2.0, f2/_fs*2.0 ], 'bandpass')
        # create the memory of the filter (=delay lines) so that
        # incoming data can be shifted through sample by sample
        self.bp_z = signal.lfiltic(self.bp_b, self.bp_a, [0])

        # sampling rate in Hz
        self.samplingRateInHz = _fs

        # heartrate in BPM
        self.bpm = 0;

    # detect r peaks
    # input: ECG samples at the specified sampling rate and in V
    def detect(self,v):
        #print('v=',v)
        h,self.bp_z = signal.lfilter(self.bp_b, self.bp_a, [v], zi = self.bp_z)
        if (self.ignoreECGdetector > 0):
            self.ignoreECGdetector = self.ignoreECGdetector - 1
            return
        h = h * h
        if (np.sqrt(h) > self.artefact_threshold):
            # ignore signal for 1 sec
            ignoreECGdetector = samplingRateInHz;
            ignoreRRvalue = 2
            return
        if (h > self.amplitude):
            self.amplitude = h

        self.amplitude = self.amplitude - self.adaptive_threshold_decay_constant * self.amplitude / self.samplingRateInHz

        if (self.doNotDetect > 0):
            self.doNotDetect = self.doNotDetect - 1
        else:
            self.threshold = self.threshold_factor * self.amplitude
            if (h > self.threshold):
                t = (self.timestamp - self.t2) / self.samplingRateInHz;
                if t>0:
                    tbpm = 1 / t * 60
                else:
                    tbpm = 0
                if ((tbpm > 30) and (tbpm < 250)):
                    if (self.ignoreRRvalue > 0):
                        self.ignoreRRvalue = self.ignoreRRvalue - 1
                    else:
                        self.bpm = tbpm
                        print(self.bpm," BPM")
                else:
                    self.ignoreRRvalue = 3
                self.t2 = self.timestamp
                # advoid 1/5 sec
                self.doNotDetect = self.samplingRateInHz / 5;
        self.timestamp = self.timestamp + 1

max_stars_repo_name: glasgowneuro/attys-python-examples

max_issues_repo_name: glasgowneuro/attys-python-examples

max_forks_repo_name: glasgowneuro/attys-python-examples

reason: import numpy,from scipy


text:
import collections
import copy

import numpy as np

from threeML.plugins.XYLike import XYLike
from threeML.utils.photometry.filter_set import FilterSet

__instrument_name = "Generic photometric data"

class PhotometryLike(XYLike):
    def __init__(self, name, filters, **data):
        """
        The photometry plugin is desinged to fit optical/IR/UV photometric data from a given
        filter system. Filters are given in the form a speclite (http://speclite.readthedocs.io)
        FitlerResponse or FilterSequence objects. 3ML contains a vast number of filters via the SVO
        VO service: http://svo2.cab.inta-csic.es/svo/theory/fps/ and can be accessed via:

        from threeML.plugins.photometry.filter_library import threeML_filter_library

        One can also construct their own filters with speclite.

        Example:

        grond = PhotometryLike('GROND',
                       filters=threeML_filter_library.ESO.GROND,
                       g=(20.93,.23),
                       r=(20.6,0.12),
                       i=(20.4,.07),
                       z=(20.3,.04),
                       J=(20.0,.03),
                       H=(19.8,.03),
                       K=(19.7,.04))

        Magnitudes and errors are entered as keyword arguments where the key is the filter name and
        the argument is a tuple containing the data. You can exclude data for individual filters and
        they will be ignored during the fit.

        NOTE: PhotometryLike expects apparent AM magnitudes. Please calibrate your data to this system

        :param name: plugin name
        :param filters: speclite filters
        :param data: keyword args of band name and tuple(mag, mag error)
        """

        # convert names so that only the filters are present
        # speclite uses '-' to separate instrument and filter

        try:

            # we have a filter sequence

            names = [fname.split('-')[1] for fname in filters.names]

        except(AttributeError):

            # we have a filter response

            names = [filters.name.split('-')[1]]

        # since we may only have a few of the  filters in use
        # we will mask the filters not needed. The will stay fixed
        # during the life of the plugin

        starting_mask = np.zeros(len(names), dtype=bool)

        for band in data.keys():

            assert band in names, 'band %s is not a member of the filter set %s'%(band,'blah')
            starting_mask[ names.index(band)] = True

        # create a filter set and use only the bands that were specified

        self._filter_set = FilterSet(filters, starting_mask)

        self._magnitudes = np.zeros(self._filter_set.n_bands)

        self._magnitude_errors = np.zeros(self._filter_set.n_bands)

        # we want to fill the magnitudes in the same order as the
        # the filters

        for i, band in enumerate(self._filter_set.filter_names):

            self._magnitudes[i] = data[band][0]
            self._magnitude_errors[i] = data[band][1]

        # pass thru to XYLike

        super(PhotometryLike, self).__init__(name=name,
                                             x=self._filter_set.effective_wavelength, # dummy x values
                                             y=self._magnitudes,
                                             yerr=self._magnitude_errors,
                                             poisson_data=False)

    @property
    def magnitudes(self):
        return self._magnitudes

    @property
    def magnitude_errors(self):
        return self._magnitude_errors

    # def set_active_filters(self, *filter_names):
    #     """
    #     set the active filters to be used in the fit
    #     :param filter_names: filter names ot be set active
    #     :return:
    #     """
    #
    #     # scroll through the known filter names
    #
    #     for i, name in enumerate(self._filter_set.filter_names):
    #
    #         for select_name in filter_names:
    #
    #             # if one of the filters is hit, then activate it
    #
    #             if name == select_name:
    #                 self._mask[i] = True
    #
    #
    #     print("Now using %d of %d filters:\n\tActive Filters: %s", (sum(self._mask),
    #                                                                 len(self._mask),
    #                                                                 ', '.join(
    #                                                                     self._filter_set.filter_names[self._mask])))
    #
    #     # reconstruct the plugin with selected data
    #
    #     super(PhotometryLike, self).__init__(name=self.name,
    #                                          x=self._filter_set.effective_wavelength[self._mask],  # dummy x values
    #                                          y=self._magnitudes[self._mask],
    #                                          yerr=self._magnitude_errors[self._mask],
    #                                          poisson_data=False)
    #
    # def set_inactive_filters(self, *filter_names):
    #     """
    #     set filters to be excluded from the fit
    #     :param filter_names: filter names ot be set inactive
    #     :return:
    #     """
    #
    #     # scroll through the known filter names
    #
    #
    #     for i, name in enumerate(self._filter_set.filter_names):
    #
    #         for select_name in filter_names:
    #
    #             if name == select_name:
    #                 self._mask[i] = False
    #
    #     print("Now using %d of %d filters:\n\tActive Filters: %s", (sum(self._mask),
    #                                                                 len(self._mask),
    #                                                                 ', '.join(
    #                                                                     self._filter_set.filter_names[self._mask])))
    #
    #     # reconstruct the plugin with selected data
    #
    #     super(PhotometryLike, self).__init__(name=self.name,
    #                                          x=self._filter_set.effective_wavelength[self._mask],  # dummy x values
    #                                          y=self._magnitudes[self._mask],
    #                                          yerr=self._magnitude_errors[self._mask],
    #                                          poisson_data=False)

    def set_model(self, likelihood_model):
        """
        set the likelihood model
        :param likelihood_model:
        :return:
        """

        super(PhotometryLike, self).set_model(likelihood_model)

        n_point_sources = self._likelihood_model.get_number_of_point_sources()

        # sum up the differential

        def differential_flux(energies):

            fluxes = self._likelihood_model.get_point_source_fluxes(0, energies, tag=self._tag)

            # If we have only one point source, this will never be executed
            for i in range(1, n_point_sources):

                fluxes += self._likelihood_model.get_point_source_fluxes(i, energies, tag=self._tag)

            return fluxes

        self._filter_set.set_model(differential_flux)

    def _get_total_expectation(self):

        return self._filter_set.ab_magnitudes()[self._mask]#.as_matrix()

    def display_filters(self):
        """
        display the filter transmission curves

        :return:
        """

        return self._filter_set.plot_filters()

    def _new_plugin(self, name, x, y, yerr):
        """
        construct a new PhotometryLike plugin. allows for returning a new plugin
        from simulated data set while customizing the constructor
        further down the inheritance tree

        :param name: new name
        :param x: new x
        :param y: new y
        :param yerr: new yerr
        :return: new XYLike

        """

        bands = collections.OrderedDict()

        for i, band in enumerate(self._filter_set.filter_names):

            bands[band] = (y[i],yerr[i])

        new_photo = PhotometryLike(name,filters=self._filter_set.speclite_filters, **bands)

        # apply the current mask

        new_photo._mask = copy.copy(self._mask)

        return new_photo

max_stars_repo_name: BjoernBiltzinger/threeML

max_issues_repo_name: BjoernBiltzinger/threeML

max_forks_repo_name: BjoernBiltzinger/threeML


text:
import functools
import warnings

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp

from tensorflow_probability import distributions as tfd
from tensorflow_probability import bijectors as tfb

from scipy.stats import multivariate_normal

tf.enable_v2_behavior()

warnings.filterwarnings('ignore')

NUM_VI_ITERS = 500
LEARNING_RATE_VI = 0.05

# ------- Specify model ---------

def clvm(data_dim, num_datapoints, counts_per_cell, dummy, is_H0=False):

    mu = yield tfd.Normal(loc=tf.zeros([data_dim, 1]),
                 scale=tf.ones([data_dim, 1]),
                 name="mu")

    beta = yield tfd.Normal(loc=tf.zeros([data_dim, 1]),
                                        scale=tf.ones([data_dim, 1]),
                                        name="beta")

    # size_factor = yield tfd.LogNormal(loc=np.mean(np.log(counts_per_cell)) * tf.ones([1, num_datapoints]),
    #                                     scale=np.std(np.log(counts_per_cell)) * tf.ones([1, num_datapoints]),
    #                                     name="size_factor")

    data = yield tfd.Poisson(rate=tf.math.exp(tf.matmul(beta, dummy) + mu), # + tf.math.log(size_factor)),
                                          name="x")

def fit_model(X, Y, compute_size_factors=True, is_H0=False, sf_x=None, sf_y=None):

    assert X.shape[0] == Y.shape[0]
    data_dim = X.shape[0]
    num_datapoints_x, num_datapoints_y = X.shape[1], Y.shape[1]
    n = num_datapoints_x + num_datapoints_y

    dummy = np.zeros(n)
    dummy[num_datapoints_x:] = 1
    dummy = np.expand_dims(dummy, 0)

    data = np.concatenate([X, Y], axis=1)

    if (sf_x is not None) and (sf_y is not None):
        counts_per_cell_X = sf_x
        counts_per_cell_Y = sf_y
        counts_per_cell_X = np.expand_dims(counts_per_cell_X, 0)
        counts_per_cell_Y = np.expand_dims(counts_per_cell_Y, 0)
        counts_per_cell = np.concatenate([counts_per_cell_X, counts_per_cell_Y], axis=1)
        assert counts_per_cell.shape[1] == X.shape[1] + Y.shape[1]
    else:
        if compute_size_factors:
            counts_per_cell_X = np.sum(X, axis=0)
            counts_per_cell_X = np.expand_dims(counts_per_cell_X, 0)
            counts_per_cell_Y = np.sum(Y, axis=0)
            counts_per_cell_Y = np.expand_dims(counts_per_cell_Y, 0)
            counts_per_cell = np.concatenate([counts_per_cell_X, counts_per_cell_Y], axis=1)
            assert counts_per_cell.shape[1] == X.shape[1] + Y.shape[1]
        else:
            counts_per_cell = 1.0

    # ------- Specify model ---------

    concrete_clvm_model = functools.partial(clvm,
                                            data_dim=data_dim,
                                            num_datapoints=n,
                                            counts_per_cell=counts_per_cell,
                                            dummy=dummy,
                                            is_H0=is_H0)

    model = tfd.JointDistributionCoroutineAutoBatched(concrete_clvm_model)

    if is_H0:

        def target_log_prob_fn(mu, beta, size_factor): return model.log_prob(
            (mu, beta, size_factor, data))

    else:

        # def target_log_prob_fn(mu, beta, size_factor): return model.log_prob(
        #     (mu, beta, size_factor, data))
        def target_log_prob_fn(mu, beta): return model.log_prob(
            (mu, beta, data))

    # ------- Specify variational families -----------

    # mu = tf.Variable(tf.random.normal([data_dim, 1]))
    # beta = tf.Variable(tf.random.normal([data_dim, 1]))
    # size_factor = tfp.util.TransformedVariable(tf.math.exp(tf.random.normal([1, n])), bijector=tfb.Softplus())

    # losses = tfp.math.minimize(
    #     lambda: -target_log_prob_fn(mu, beta, size_factor),
    #     optimizer=tf.optimizers.Adam(learning_rate=0.05),
    #     num_steps=500)

    # d = np.log(data + 1)
    # d = d / d.sum(0)
    # from sklearn.linear_model import LinearRegression
    # plt.scatter(np.squeeze(LinearRegression().fit(dummy.T, d.T).coef_), np.squeeze(beta.numpy()))
    # plt.show()
    # import ipdb
    # ipdb.set_trace()

    # Variational parmater means

    # mu
    qmu_mean = tf.Variable(tf.random.normal([data_dim, 1]))
    qmu_stddv = tfp.util.TransformedVariable(
        1e-4 * tf.ones([data_dim, 1]),
        bijector=tfb.Softplus())

    # beta
    qbeta_mean = tf.Variable(tf.random.normal([data_dim, 1]))
    qbeta_stddv = tfp.util.TransformedVariable(
        1e-4 * tf.ones([data_dim, 1]),
        bijector=tfb.Softplus())

    # size factor
    # qsize_factor_mean = tf.Variable(np.mean(np.log(counts_per_cell)) * tf.ones([1, n]))
    # qsize_factor_stddv = tfp.util.TransformedVariable(
    #     np.std(np.log(counts_per_cell)) * tf.ones([1, n]), # + tf.math.exp(tf.random.normal([1, n])),
    #     # 1e-4 * tf.ones([1, n]),
    #     bijector=tfb.Softplus())
    qsize_factor_mean = tf.Variable(tf.random.normal([1, n]))
    qsize_factor_stddv = tfp.util.TransformedVariable(
        1e-4 * tf.ones([1, n]),
        bijector=tfb.Softplus())

    def factored_normal_variational_model():

        qmu = yield tfd.Normal(loc=qmu_mean,
                 scale=qmu_stddv,
                 name="qmu")

        qbeta = yield tfd.Normal(loc=qbeta_mean,
                     scale=qbeta_stddv,
                     name="qbeta")

        # qsize_factor = yield tfd.LogNormal(loc=qsize_factor_mean,
        #              scale=qsize_factor_stddv,
        #              name="qsize_factor")

    # Surrogate posterior that we will try to make close to p
    surrogate_posterior = tfd.JointDistributionCoroutineAutoBatched(
        factored_normal_variational_model)

    # --------- Fit variational inference model using MC samples and gradient descent ----------

    losses = tfp.vi.fit_surrogate_posterior(
        target_log_prob_fn,
        surrogate_posterior=surrogate_posterior,
        optimizer=tf.optimizers.Adam(learning_rate=LEARNING_RATE_VI),
        num_steps=NUM_VI_ITERS)
    # ipdb.set_trace()

    if is_H0:
        return_dict = {
            'loss_trace': losses,
        }
    else:
        return_dict = {
            'loss_trace': losses,
            'qmu_mean': qmu_mean,
            'qmu_stddv': qmu_stddv,
            'qbeta_mean': qbeta_mean,
            'qbeta_stddv': qbeta_stddv,
            'qsize_factor_mean': qsize_factor_mean,
            'qsize_factor_stddv': qsize_factor_stddv
        }

    return return_dict

if __name__ == "__main__":

    data_dim = 10
    n = 5000
    dummy = np.zeros(n)
    dummy[3000:] = 1
    dummy = np.expand_dims(dummy, 0)
    mimosca_model = functools.partial(clvm,
                                            data_dim=data_dim,
                                            dummy=dummy,
                                            num_datapoints=n,
                                            counts_per_cell=1,
                                            is_H0=False)

    model = tfd.JointDistributionCoroutineAutoBatched(mimosca_model)

    samples = model.sample()
    data_sampled = samples[-1]

    data_sampled = data_sampled.numpy()
    import ipdb; ipdb.set_trace()

max_stars_repo_name: vishalbelsare/cplvm

max_issues_repo_name: vishalbelsare/cplvm

max_forks_repo_name: vishalbelsare/cplvm


text:
"""Pytorch Dataset object that loads perfectly balanced MNIST dataset in bag form."""

import numpy as np
import torch
import torch.utils.data as data_utils
from sklearn.model_selection import KFold
from torch.utils.data import ConcatDataset
from torchvision import datasets, transforms

class MnistBags(data_utils.Dataset):
    def __init__(self, target_number=9, bag_length_mean=200, bag_length_std=150, bag_length_min=50, bag_length_max=600,
                 positive_samples_in_bag_ratio_mean=0.1, positive_samples_in_bag_ratio_std=0.01,
                 num_bags_train=300, num_bags_test=300, folds=10, fold_id=0, random_state=3, train=True, push=False,
                 test=False, all_labels=False):
        self.target_number = target_number
        self.mean_bag_length = bag_length_mean
        self.var_bag_length = bag_length_std
        self.num_bag = num_bags_train if train else num_bags_test
        self.train = train
        self.push = push
        self.test = test
        self.negative_bags = self.num_bag // 2
        self.min_bag_size = bag_length_min
        self.max_bag_size = bag_length_max
        self.folds = folds
        self.fold_id = fold_id
        self.random_state = random_state
        self.all_labels = all_labels

        self.target_numbers_in_pos_bag_mean = positive_samples_in_bag_ratio_mean
        self.target_numbers_in_pos_bag_std = positive_samples_in_bag_ratio_std

        self.r = np.random.RandomState(random_state)

        self.pil_to_rgb_tensor = transforms.Compose([
            transforms.Lambda(lambda x: x.convert('RGB')),
            transforms.ToTensor(),
        ])

        self.normalize_to_tensor_transform = transforms.Compose([
            # transforms.Normalize((0.1307,), (0.3081,)),
        ])

        self.bags_list, self.labels_list = self._form_bags()

    def _transform_single(self, x):
        return self.pil_to_rgb_tensor(x)

    def _form_bags(self):
        mnist_train = datasets.MNIST('data', train=True, download=True)
        mnist_test = datasets.MNIST('data', train=False, download=True)
        mnist = ConcatDataset([mnist_train, mnist_test])
        folds = list(KFold(n_splits=self.folds, shuffle=True, random_state=self.random_state).split(mnist))

        if self.test:
            indices = set(folds[self.fold_id][1])
        else:
            if self.train:
                val_indices = self.r.choice(folds[self.fold_id][0], len(folds[self.fold_id][1]))
                indices = set(folds[self.fold_id][0]) - set(val_indices)
            else:  # valid
                indices = self.r.choice(folds[self.fold_id][0], len(folds[self.fold_id][1]))

        negative_samples = torch.stack(
            [self._transform_single(d) for idx, (d, t) in enumerate(mnist) if
             idx in indices and t != self.target_number])
        positive_samples = torch.stack(
            [self._transform_single(d) for idx, (d, t) in enumerate(mnist) if
             idx in indices and t == self.target_number])

        bags_list = []
        labels_list = []

        for i in range(self.num_bag):
            bag_length = np.int(self.r.normal(self.mean_bag_length, self.var_bag_length, 1))
            bag_length = min(max(bag_length, self.min_bag_size), self.max_bag_size)

            if i < self.negative_bags:
                indices = self.r.randint(0, len(negative_samples), bag_length)
                bags_list.append(negative_samples[indices])
                labels_list.append(torch.zeros(bag_length, dtype=torch.long))
            else:
                positive_ratio = float(
                    self.r.normal(self.target_numbers_in_pos_bag_mean, self.target_numbers_in_pos_bag_std, 1))
                positive_num = min(max(int(bag_length * positive_ratio), 1), bag_length - 1)
                negative_num = bag_length - positive_num
                pos_indices = self.r.randint(0, len(positive_samples), positive_num)
                neg_indices = self.r.randint(0, len(negative_samples), negative_num)
                bag = torch.cat([positive_samples[pos_indices], negative_samples[neg_indices]])
                labels = torch.cat(
                    [torch.ones(positive_num, dtype=torch.long), torch.zeros(negative_num, dtype=torch.long)])
                shuffle_indices = self.r.permutation(bag_length)
                bag = bag[shuffle_indices]
                labels = labels[shuffle_indices]
                bags_list.append(bag)
                labels_list.append(labels)

        return bags_list, labels_list

    def __len__(self):
        if self.train:
            return len(self.labels_list)
        else:
            return len(self.labels_list)

    def __getitem__(self, index):
        bag = self.bags_list[index]
        if self.all_labels:
            label = self.labels_list[index]
        else:
            label = self.labels_list[index].max().unsqueeze(0)
        if self.push:
            return bag, self.normalize_to_tensor_transform(bag), label
        else:
            return self.normalize_to_tensor_transform(bag), label

max_stars_repo_name: apardyl/ProtoPNet

max_issues_repo_name: apardyl/ProtoPNet

max_forks_repo_name: apardyl/ProtoPNet


text:
# Copyright 2019 NREL

# Licensed under the Apache License, Version 2.0 (the "License"); you may not use
# this file except in compliance with the License. You may obtain a copy of the
# License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import pickle
from floris.utilities import wrap_180

class PowerRose():
    """
    PowerRose object class used to parse data and generate figures.
    """

    def __init__(self):
        """
        Init Function of WindRose Object.
        No explicit arguments required.
        """

    def _all_combine(self):
        df_power = self.df_power.copy(deep=True)
        df_yaw = self.df_yaw.copy(deep=True)
        df_turbine_power_no_wake = self.df_turbine_power_no_wake.copy(
            deep=True)
        df_turbine_power_baseline = self.df_turbine_power_baseline.copy(
            deep=True)
        df_turbine_power_opt = self.df_turbine_power_opt.copy(deep=True)

        # Adjust the column names for uniqunes
        df_yaw.columns = [
            'yaw_%d' % c if type(c) is int else c for c in df_yaw.columns
        ]
        df_turbine_power_no_wake.columns = [
            'tnw_%d' % c if type(c) is int else c
            for c in df_turbine_power_no_wake.columns
        ]
        df_turbine_power_baseline.columns = [
            'tb_%d' % c if type(c) is int else c
            for c in df_turbine_power_baseline.columns
        ]
        df_turbine_power_opt.columns = [
            'topt_%d' % c if type(c) is int else c
            for c in df_turbine_power_opt.columns
        ]

        # Merge
        df_combine = df_power.merge(df_yaw, on=['ws', 'wd'])
        df_combine = df_combine.merge(df_turbine_power_no_wake,
                                      on=['ws', 'wd'])
        df_combine = df_combine.merge(df_turbine_power_baseline,
                                      on=['ws', 'wd'])
        df_combine = df_combine.merge(df_turbine_power_opt, on=['ws', 'wd'])

        return df_combine

    def _norm_frequency(self, df):
        print('Norming frequency total of %.2f to 1.0' % df.freq_val.sum())
        df['freq_val'] = df.freq_val / df.freq_val.sum()
        return df

    def _compute_energy(self):
        self.df_power[
            'energy_no_wake'] = self.df_power.freq_val * self.df_power.power_no_wake
        self.df_power[
            'energy_baseline'] = self.df_power.freq_val * self.df_power.power_baseline
        self.df_power[
            'energy_opt'] = self.df_power.freq_val * self.df_power.power_opt

    def _compute_totals(self):
        df = self.df_power.copy(deep=True)
        df = df.sum()

        # Get total energy amounts
        self.total_no_wake = df.energy_no_wake
        self.total_baseline = df.energy_baseline
        self.total_opt = df.energy_opt

        # Get wake loss amounts
        self.baseline_percent = self.total_baseline / self.total_no_wake
        self.baseline_wake_loss = 1 - self.baseline_percent
        self.opt_percent = self.total_opt / self.total_no_wake
        self.opt_wake_loss = 1 - self.opt_percent

        # Percent gain
        self.percent_gain = (self.total_opt -
                             self.total_baseline) / self.total_baseline
        self.reduction_in_wake_loss = -1 * (
            self.opt_wake_loss -
            self.baseline_wake_loss) / self.baseline_wake_loss

    def report(self):
        """
        Print operational data information to screen.
        """

        print('=============================================')
        print('Case %s has results:' % self.name)
        print('=============================================')
        print('-\tNo-Wake\t\tBaseline\tOpt ')
        print('---------------------------------------------')
        print('Energy\t%.1E\t\t%.1E\t\t%.1E' %
              (self.total_no_wake, self.total_baseline, self.total_opt))
        print('%%\t--\t\t%.1f%%\t\t%.1f%%' %
              (100. * self.baseline_percent, 100. * self.opt_percent))
        print('Wk Loss\t--\t\t%.1f%%\t\t%.1f%%' %
              (100. * self.baseline_wake_loss, 100. * self.opt_wake_loss))
        print('AEP Gain --\t\t--\t\t%.1f%%' % (100. * self.percent_gain))
        print('Loss Red --\t\t--\t\t%.1f%%' %
              (100. * self.reduction_in_wake_loss))

    def initialize(self, name, df_power, df_yaw, df_turbine_power_no_wake,
                   df_turbine_power_baseline, df_turbine_power_opt):
        """
        Additional initialization of the PowerRose object.
        #TODO Shouldn't this just be part of the __init__ method?

        Args:
            name (str): name of PowerRose object.
            df_power (pd.DataFrame): plant power data.
            df_yaw (pd.DataFrame): yaw data.
            df_turbine_power_no_wake (pd.DataFrame): wind turbine power
                data without wake losses.
            df_turbine_power_baseline (pd.DataFrame): wind turbine
                baseline power data.
            df_turbine_power_opt (pd.DataFrame): optimal wind turbine
                power data.
        """
        self.name = name
        self.df_power = self._norm_frequency(df_power)
        self.df_yaw = df_yaw
        self.df_turbine_power_no_wake = df_turbine_power_no_wake
        self.df_turbine_power_baseline = df_turbine_power_baseline
        self.df_turbine_power_opt = df_turbine_power_opt

        # Make a single combined frame in case it's useful
        self.df_combine = self._all_combine()

        # Compute energies
        self._compute_energy()

        # Compute totals
        self._compute_totals()

    def load(self, filename):
        """
        Load PowerRose object from pickle file.

        Args:
            filename (str): Read-from path to pickled PowerRose object.
        """

        self.name, self.df_power, self.df_yaw, self.df_turbine_power_no_wake, self.df_turbine_power_baseline, self.df_turbine_power_opt, self.df_combine = pickle.load(
            open(filename, "rb"))

        # Compute energies
        self._compute_energy()

        # Compute totals
        self._compute_totals()

    def save(self, filename):
        """
        Pickle a PowerRose Object

        Args:
            filename (str): Write-to path for PowerRose pickle.
        """
        pickle.dump([
            self.name, self.df_power, self.df_yaw,
            self.df_turbine_power_no_wake, self.df_turbine_power_baseline,
            self.df_turbine_power_opt, self.df_combine
        ], open(filename, "wb"))

    def plot_by_direction(self):
        """
        Bin Power data by direction and plot.

        Returns:
            tuple: tuple containing:

                -   **fig** (*plt.figure*): Figure handle.
                -   **axarr** (*list*): list of axis handles.
        """

        df = self.df_power.copy(deep=True)
        df = df.groupby('wd').sum().reset_index()

        fig, axarr = plt.subplots(2, 1, sharex=True)
        ax = axarr[0]
        ax.plot(df.wd,
                df.energy_baseline / df.energy_no_wake,
                label='Baseline',
                color='k')
        ax.axhline(np.mean(df.energy_baseline) / np.mean(df.energy_no_wake),
                   color='k',
                   ls='--')
        ax.plot(df.wd,
                df.energy_opt / df.energy_no_wake,
                label='Optimized',
                color='r')
        ax.axhline(np.mean(df.energy_opt) / np.mean(df.energy_no_wake),
                   color='r',
                   ls='--')
        ax.set_ylabel('Wake Loss')
        ax.grid(True)
        ax.legend()

        ax = axarr[1]
        ax.plot(
            df.wd,
            100. * (df.energy_opt - df.energy_baseline) / df.energy_baseline,
            'r')
        ax.axhline(100. * (df.energy_opt.mean() - df.energy_baseline.mean()) /
                   df.energy_baseline.mean(),
                   color='r',
                   ls='--')
        ax.set_ylabel('Percent Gain')

        return fig, axarr

    def wake_loss_at_direction(self, wd):
        """
        Calculate wake losses for a given direction. Plot rose figures
        for Power, Energy, Baseline power, Optimal gain, Total Gain,
        Percent Gain, etc.

        Args:
            wd (float): Wind direction of interest.

        Returns:
            tuple: tuple containing:

                -   **fig** (*plt.figure*): Figure handle.
                -   **axarr** (*list*): list of axis handles.
        """

        df = self.df_power.copy(deep=True)

        # Choose the nearest direction
        # Find nearest wind direction
        df['dist'] = np.abs(wrap_180(df.wd - wd))
        wd_select = df[df.dist == df.dist.min()]['wd'].unique()[0]
        print('Nearest wd to %.1f is %.1f' % (wd, wd_select))
        df = df[df.wd == wd_select]

        df = df.groupby('ws').sum().reset_index()

        fig, axarr = plt.subplots(4, 2, sharex=True, figsize=(14, 12))

        ax = axarr[0, 0]
        ax.set_title('Power')
        ax.plot(df.ws, df.power_no_wake, 'k', label='No Wake')
        ax.plot(df.ws, df.power_baseline, 'b', label='Baseline')
        ax.plot(df.ws, df.power_opt, 'r', label='Opt')
        ax.set_ylabel('Total')
        ax.grid()

        ax = axarr[0, 1]
        ax.set_title('Energy')
        ax.plot(df.ws, df.energy_no_wake, 'k', label='No Wake')
        ax.plot(df.ws, df.energy_baseline, 'b', label='Baseline')
        ax.plot(df.ws, df.energy_opt, 'r', label='Opt')
        ax.legend()
        ax.grid()

        ax = axarr[1, 0]
        ax.plot(df.ws,
                df.power_baseline / df.power_no_wake,
                'b',
                label='Baseline')
        ax.plot(df.ws, df.power_opt / df.power_no_wake, 'r', label='Opt')
        ax.set_ylabel('Percent')
        ax.grid()

        ax = axarr[1, 1]
        ax.plot(df.ws,
                df.energy_baseline / df.energy_no_wake,
                'b',
                label='Baseline')
        ax.plot(df.ws, df.energy_opt / df.energy_no_wake, 'r', label='Opt')
        ax.grid()

        ax = axarr[2, 0]
        ax.plot(df.ws, (df.power_opt - df.power_baseline), 'r')
        ax.set_ylabel('Total Gain')
        ax.grid()

        ax = axarr[2, 1]
        ax.plot(df.ws, (df.energy_opt - df.energy_baseline), 'r')
        ax.grid()

        ax = axarr[3, 0]
        ax.plot(df.ws, (df.power_opt - df.power_baseline) / df.power_baseline,
                'r')
        ax.set_ylabel('Percent Gain')
        ax.grid()
        ax.set_xlabel('Wind Speed (m/s)')

        ax = axarr[3, 1]
        ax.plot(df.ws,
                (df.energy_opt - df.energy_baseline) / df.energy_baseline, 'r')
        ax.grid()
        ax.set_xlabel('Wind Speed (m/s)')

        return fig, axarr