text:
#!/usr/bin/env python3

import torch
import torch.multiprocessing as mp
import torch.distributed as dist
import torch.nn.functional as F
import torchvision
import albumentations
import albumentations.pytorch
from albumentations.pytorch import ToTensorV2

import numpy as np
import pandas as pd
import cv2
import os
import wandb
import copy

import unitopatho
import utils

import re
import argparse

import torchstain

from tqdm import tqdm
from collections import defaultdict
from sklearn.model_selection import GroupShuffleSplit
from pathlib import Path
from functools import partial
from multiprocessing import Manager

torch.multiprocessing.set_sharing_strategy('file_system')
manager = Manager()

def resnet18(n_classes=2):
    model = torchvision.models.resnet18(pretrained='imagenet')
    model.fc = torch.nn.Linear(in_features=model.fc.in_features, out_features=n_classes, bias=True)
    return model

def preprocess_df(df, label):
    if label == 'norm':
        df.loc[df.grade == 0, 'grade'] = -1
        df.loc[df.type == 'norm', 'grade'] = 0

    df = df[df.grade >= 0].copy()

    if label != 'both' and label != 'norm':
        df = df[df.type == label].copy()
    return df

def main(config):
    checkpoint = None
    if config.test is not None:
        print('=> Loading saved checkpoint')
        checkpoint = torch.hub.load_state_dict_from_url(f'https://api.wandb.ai/files/eidos/UnitoPath-v1/{config.test}/model.pt',
                                                        map_location='cpu', progress=True, check_hash=False)
        test = config.test
        device = config.device
        p = config.path
        config = checkpoint['config']
        config.test = test
        config.device = device
        config.path = p

    utils.set_seed(config.seed)
    scaler = torch.cuda.amp.GradScaler()

    if config.test is None:
        wandb.init(config=config,
                    project=f'unitopatho')

    path = os.path.join(config.path, str(config.size))
    train_df = pd.read_csv(os.path.join(path, 'train.csv'))
    test_df = pd.read_csv(os.path.join(path, 'test.csv'))

    groupby = config.target + ''
    print('=> Raw data (train)')
    print(train_df.groupby(groupby).count())

    print('\n=> Raw data (test)')
    print(test_df.groupby(groupby).count())

    if config.target == 'grade':
        train_df = preprocess_df(train_df, config.label)
        test_df = preprocess_df(test_df, config.label)

        # balance train_df (sample mean size)
        groups = train_df.groupby('grade').count()
        grade_min = int(groups.image_id.idxmin())
        mean_size = int(train_df.groupby('grade').count().mean()['image_id'])

        train_df = pd.concat((
            train_df[train_df.grade == 0].sample(mean_size, replace=(grade_min==0), random_state=config.seed).copy(),
            train_df[train_df.grade == 1].sample(mean_size, replace=(grade_min==1), random_state=config.seed).copy()
        ))

    else:
        # balance train_df (sample 3rd min_size)
        min_size = np.sort(train_df.groupby(groupby).count()['image_id'])[2]
        train_df = train_df.groupby(groupby).apply(lambda group: group.sample(min_size, replace=len(group) < min_size, random_state=config.seed)).reset_index(drop=True)

    print('\n---- DATA SUMMARY ----')
    print('---------------------------------- Train ----------------------------------')
    print(train_df.groupby(groupby).count())
    print(len(train_df.wsi.unique()), 'WSIs')

    print('\n---------------------------------- Test ----------------------------------')
    print(test_df.groupby(groupby).count())
    print(len(test_df.wsi.unique()), 'WSIs')

    im_mean, im_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] # ImageNet
    norm = dict(
        rgb=dict(mean=im_mean,
                    std=im_std),
        he=dict(mean=im_mean,
                std=im_std),
        gray=dict(mean=[0.5],
                    std=[1.0])
    )

    T_aug = albumentations.Compose([
        albumentations.HorizontalFlip(p=0.5),
        albumentations.VerticalFlip(p=0.5),
        albumentations.Rotate(90, p=0.5)
    ])
    T_jitter = albumentations.ColorJitter()

    mean, std = norm[config.preprocess]['mean'], norm[config.preprocess]['std']
    print('=> mean, std:', mean, std)
    T_tensor = ToTensorV2()
    T_post = albumentations.Compose([
        albumentations.Normalize(mean, std),
        T_tensor
    ])

    print('=> Preparing stain normalizer..')
    he_target = cv2.cvtColor(cv2.imread('data/target.jpg'), cv2.COLOR_BGR2RGB)
    normalizer = torchstain.MacenkoNormalizer(backend='torch')
    normalizer.fit(T_tensor(image=he_target)['image']*255)
    print('=> Done')

    def normalize_he(x):
        if config.preprocess == 'he':
            img = x
            try:
                img = T_tensor(image=img)['image']*255
                img, _, _ = normalizer.normalize(img, stains=False)
                img = img.numpy().astype(np.uint8)
            except Exception as e:
                print('Could not normalize image:', e)
                img = x
            return img
        return x

    def apply_transforms(train, img):
        img = normalize_he(img)
        if train:
            img = T_aug(image=img)['image']
            if config.preprocess == 'rgb':
                img = T_jitter(image=img)['image']
        x = img
        return T_post(image=x)['image']

    T_train = partial(apply_transforms, True)
    T_test = partial(apply_transforms, False)

    datasets_kwargs = {
        'path': path,
        'subsample': config.subsample,
        'target': config.target,
        'gray': config.preprocess == 'gray',
        'mock': config.mock
    }

    train_dataset = unitopatho.UTP(train_df, T=T_train, **datasets_kwargs)
    test_dataset = unitopatho.UTP(test_df, T=T_test, **datasets_kwargs)

    # Final loaders
    train_loader = torch.utils.data.DataLoader(train_dataset, shuffle=True,
                                               batch_size=config.batch_size,
                                               num_workers=config.n_workers,
                                               pin_memory=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, shuffle=False,
                                              batch_size=config.batch_size,
                                              num_workers=config.n_workers,
                                              pin_memory=True)

    n_classes = len(train_df[config.target].unique())
    print(f'=> Training for {n_classes} classes')

    n_channels = {
        'rgb': 3,
        'he': 3,
        'gray': 1
    }

    model = resnet18(n_classes=n_classes)
    model.conv1 = torch.nn.Conv2d(n_channels[config.preprocess], 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
    if checkpoint is not None:
        model.load_state_dict(checkpoint['model'])
    model = model.to(config.device)

    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr, weight_decay=config.weight_decay)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)
    criterion = F.cross_entropy

    for epoch in range(config.epochs):
        if config.test is None:
            train_metrics = utils.train(model, train_loader, criterion,
                                        optimizer, config.device, metrics=utils.metrics,
                                        accumulation_steps=config.accumulation_steps, scaler=scaler)
            scheduler.step()

        test_metrics = utils.test(model, test_loader, criterion, config.device, metrics=utils.metrics)

        if config.test is None:
            print(f'Epoch {epoch}: train: {train_metrics}')
            wandb.log({'train': train_metrics,
                       'test': test_metrics})
            torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'config': config},
                    os.path.join(wandb.run.dir, 'model.pt'))

        print(f'test: {test_metrics}')
        if config.test is not None:
            break

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # data config
    parser.add_argument('--path', default=f'{os.path.expanduser("~")}/data/UNITOPATHO', type=str, help='UNITOPATHO dataset path')
    parser.add_argument('--size', default=100, type=int, help='patch size in µm (default 100)')
    parser.add_argument('--subsample', default=-1, type=int, help='subsample size for data (-1 to disable, default -1)')

    # optimizer & network config
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--lr', default=0.0001, type=float, help='learning rate')
    parser.add_argument('--momentum', default=0.99, type=float, help='momentum')
    parser.add_argument('--weight_decay', default=1e-5, type=float, help='weight decay')
    parser.add_argument('--batch_size', default=256, type=int, help='batch size')
    parser.add_argument('--accumulation_steps', default=1, type=int, help='gradient accumulation steps')
    parser.add_argument('--n_workers', default=8, type=int)
    parser.add_argument('--architecture', default='resnet18', help='resnet18, resnet50, densenet121')

    # training config
    parser.add_argument('--preprocess', default='rgb', help='preprocessing type, rgb, he or gray. Default: rgb')
    parser.add_argument('--target', default='grade', help='target attribute: grade, type, top_label (default: grade)')
    parser.add_argument('--label', default='both', type=str, help='only when target=grade; values: ta, tva, norm or both (default: both)')
    parser.add_argument('--test', type=str, help='Run id to test', default=None)

    # misc config
    parser.add_argument('--name', type=str, default=None)
    parser.add_argument('--device', default='cuda', type=str)
    parser.add_argument('--mock', action='store_true', dest='mock', help='mock dataset (random noise)')
    parser.add_argument('--seed', type=int, default=42)
    parser.set_defaults(mock=False)

    config = parser.parse_args()
    config.device = torch.device(config.device)

    main(config)

max_stars_repo_name: EIDOSlab/UNITOPATHO

max_issues_repo_name: EIDOSlab/UNITOPATHO

max_forks_repo_name: EIDOSlab/UNITOPATHO


text:
"""Observations about the action.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import gym

from robovat.observations import observation

class LastActionValidObs(observation.Observation):
    """Observation which indicates if the action is valid."""

    def __init__(self, name=None):
        """Initialize."""
        self.name = name or 'last_action_valid'
        self.env = None

    def initialize(self, env):
        self.env = env

    def get_gym_space(self):
        """Returns gym space of this observation."""
        return gym.spaces.Discrete(1)

    def get_observation(self):
        """Returns the observation data of the current step."""
        return np.array(self.env.last_action_is_valid, dtype=np.int32)

class ModeObs(observation.Observation):
    """Observation of the action mode."""

    def __init__(self, name=None):
        """Initialize."""
        self.name = name or 'mode'
        self.env = None

    def initialize(self, env):
        self.env = env

    def get_gym_space(self):
        """Returns gym space of this observation."""
        return gym.spaces.Discrete(len(self.env.ACTION_MODES))

    def get_observation(self):
        """Returns the observation data of the current step."""
        return np.array(self.env.mode, dtype=np.int32)

ext: py

max_stars_repo_name: leobxpan/robovat

max_issues_repo_name: leobxpan/robovat

max_forks_repo_name: leobxpan/robovat

reason: import numpy


text:
from scipy.interpolate import interp1d

N = [12, 13, 14, 15, 16,
     17, 18, 19, 20, 21,
     22, 24, 26, 28, 30,
     34, 38, 43, 50, 60,
     75, 100, 150, 300, 400]

Y = [0.245, 0.261, 0.277, 0.290, 0.296,
     0.303, 0.309, 0.314, 0.322, 0.328,
     0.331, 0.337, 0.346, 0.353, 0.359,
     0.371, 0.384, 0.397, 0.409, 0.422,
     0.435, 0.447, 0.460, 0.472, 0.480]

def fatorFormaLewis(N_dentes):
    fator = interp1d(N, Y)
    return fator(N_dentes)

max_stars_repo_name: ZaqueuCavalcante/Gears

max_issues_repo_name: ZaqueuCavalcante/Gears

max_forks_repo_name: ZaqueuCavalcante/Gears

reason: from scipy


text:
# Mohammad Saad
# 2/24/2018
# weights.py
# Loads weights into the model
# Pretty much the same as @iapatil's version

import numpy as np
import torch

def load_weights(model, filename, dtype):

	model_params = model.state_dict()
	data_dict = np.load(filename, encoding='latin1').item()

	# initial layer
	model_params["conv1.weight"] = torch.from_numpy(data_dict["conv1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["conv1.bias"] = torch.from_numpy(data_dict["conv1"]["biases"]).type(dtype)
	model_params["bn1.weight"] = torch.from_numpy(data_dict["bn_conv1"]["scale"]).type(dtype)
	model_params["bn1.bias"] = torch.from_numpy(data_dict["bn_conv1"]["offset"]).type(dtype)

	## first projection layer
	model_params["proj1.conv1.weight"] = torch.from_numpy(data_dict["res2a_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj1.bn1.weight"] = torch.from_numpy(data_dict["bn2a_branch2a"]["scale"]).type(dtype)
	model_params["proj1.bn1.bias"] = torch.from_numpy(data_dict["bn2a_branch2a"]["offset"]).type(dtype)

	model_params["proj1.conv2.weight"] = torch.from_numpy(data_dict["res2a_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj1.bn2.weight"] = torch.from_numpy(data_dict["bn2a_branch2b"]["scale"]).type(dtype)
	model_params["proj1.bn2.bias"] = torch.from_numpy(data_dict["bn2a_branch2b"]["offset"]).type(dtype)

	model_params["proj1.conv3.weight"] = torch.from_numpy(data_dict["res2a_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj1.bn3.weight"] = torch.from_numpy(data_dict["bn2a_branch2c"]["scale"]).type(dtype)
	model_params["proj1.bn3.bias"] = torch.from_numpy(data_dict["bn2a_branch2c"]["offset"]).type(dtype)

	model_params["proj1.conv4.weight"] = torch.from_numpy(data_dict["res2a_branch1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj1.bn4.weight"] = torch.from_numpy(data_dict["bn2a_branch1"]["scale"]).type(dtype)
	model_params["proj1.bn4.bias"] = torch.from_numpy(data_dict["bn2a_branch1"]["offset"]).type(dtype)

	#
	model_params["res1_1.conv1.weight"] = torch.from_numpy(data_dict["res2b_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_1.bn1.weight"] = torch.from_numpy(data_dict["bn2b_branch2a"]["scale"]).type(dtype)
	model_params["res1_1.bn1.bias"] = torch.from_numpy(data_dict["bn2b_branch2a"]["offset"]).type(dtype)

	model_params["res1_1.conv2.weight"] = torch.from_numpy(data_dict["res2b_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_1.bn2.weight"] = torch.from_numpy(data_dict["bn2b_branch2b"]["scale"]).type(dtype)
	model_params["res1_1.bn2.bias"] = torch.from_numpy(data_dict["bn2b_branch2b"]["offset"]).type(dtype)

	model_params["res1_1.conv3.weight"] = torch.from_numpy(data_dict["res2b_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_1.bn3.weight"] = torch.from_numpy(data_dict["bn2b_branch2c"]["scale"]).type(dtype)
	model_params["res1_1.bn3.bias"] = torch.from_numpy(data_dict["bn2b_branch2c"]["offset"]).type(dtype)

	#
	model_params["res1_2.conv1.weight"] = torch.from_numpy(data_dict["res2c_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_2.bn1.weight"] = torch.from_numpy(data_dict["bn2c_branch2a"]["scale"]).type(dtype)
	model_params["res1_2.bn1.bias"] = torch.from_numpy(data_dict["bn2c_branch2a"]["offset"]).type(dtype)

	model_params["res1_2.conv2.weight"] = torch.from_numpy(data_dict["res2c_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_2.bn2.weight"] = torch.from_numpy(data_dict["bn2c_branch2b"]["scale"]).type(dtype)
	model_params["res1_2.bn2.bias"] = torch.from_numpy(data_dict["bn2c_branch2b"]["offset"]).type(dtype)

	model_params["res1_2.conv3.weight"] = torch.from_numpy(data_dict["res2c_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res1_2.bn3.weight"] = torch.from_numpy(data_dict["bn2c_branch2c"]["scale"]).type(dtype)
	model_params["res1_2.bn3.bias"] = torch.from_numpy(data_dict["bn2c_branch2c"]["offset"]).type(dtype)

	## second projection layer
	model_params["proj2.conv1.weight"] = torch.from_numpy(data_dict["res3a_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj2.bn1.weight"] = torch.from_numpy(data_dict["bn3a_branch2a"]["scale"]).type(dtype)
	model_params["proj2.bn1.bias"] = torch.from_numpy(data_dict["bn3a_branch2a"]["offset"]).type(dtype)

	model_params["proj2.conv2.weight"] = torch.from_numpy(data_dict["res3a_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj2.bn2.weight"] = torch.from_numpy(data_dict["bn3a_branch2b"]["scale"]).type(dtype)
	model_params["proj2.bn2.bias"] = torch.from_numpy(data_dict["bn3a_branch2b"]["offset"]).type(dtype)

	model_params["proj2.conv3.weight"] = torch.from_numpy(data_dict["res3a_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj2.bn3.weight"] = torch.from_numpy(data_dict["bn3a_branch2c"]["scale"]).type(dtype)
	model_params["proj2.bn3.bias"] = torch.from_numpy(data_dict["bn3a_branch2c"]["offset"]).type(dtype)

	model_params["proj2.conv4.weight"] = torch.from_numpy(data_dict["res3a_branch1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj2.bn4.weight"] = torch.from_numpy(data_dict["bn3a_branch1"]["scale"]).type(dtype)
	model_params["proj2.bn4.bias"] = torch.from_numpy(data_dict["bn3a_branch1"]["offset"]).type(dtype)

	#
	model_params["res2_1.conv1.weight"] = torch.from_numpy(data_dict["res3b_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_1.bn1.weight"] = torch.from_numpy(data_dict["bn3b_branch2a"]["scale"]).type(dtype)
	model_params["res2_1.bn1.bias"] = torch.from_numpy(data_dict["bn3b_branch2a"]["offset"]).type(dtype)

	model_params["res2_1.conv2.weight"] = torch.from_numpy(data_dict["res3b_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_1.bn2.weight"] = torch.from_numpy(data_dict["bn3b_branch2b"]["scale"]).type(dtype)
	model_params["res2_1.bn2.bias"] = torch.from_numpy(data_dict["bn3b_branch2b"]["offset"]).type(dtype)

	model_params["res2_1.conv3.weight"] = torch.from_numpy(data_dict["res3b_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_1.bn3.weight"] = torch.from_numpy(data_dict["bn3b_branch2c"]["scale"]).type(dtype)
	model_params["res2_1.bn3.bias"] = torch.from_numpy(data_dict["bn3b_branch2c"]["offset"]).type(dtype)

	#
	model_params["res2_2.conv1.weight"] = torch.from_numpy(data_dict["res3c_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_2.bn1.weight"] = torch.from_numpy(data_dict["bn3c_branch2a"]["scale"]).type(dtype)
	model_params["res2_2.bn1.bias"] = torch.from_numpy(data_dict["bn3c_branch2a"]["offset"]).type(dtype)

	model_params["res2_2.conv2.weight"] = torch.from_numpy(data_dict["res3c_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_2.bn2.weight"] = torch.from_numpy(data_dict["bn3c_branch2b"]["scale"]).type(dtype)
	model_params["res2_2.bn2.bias"] = torch.from_numpy(data_dict["bn3c_branch2b"]["offset"]).type(dtype)

	model_params["res2_2.conv3.weight"] = torch.from_numpy(data_dict["res3c_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_2.bn3.weight"] = torch.from_numpy(data_dict["bn3c_branch2c"]["scale"]).type(dtype)
	model_params["res2_2.bn3.bias"] = torch.from_numpy(data_dict["bn3c_branch2c"]["offset"]).type(dtype)

	#
	model_params["res2_3.conv1.weight"] = torch.from_numpy(data_dict["res3d_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_3.bn1.weight"] = torch.from_numpy(data_dict["bn3d_branch2a"]["scale"]).type(dtype)
	model_params["res2_3.bn1.bias"] = torch.from_numpy(data_dict["bn3d_branch2a"]["offset"]).type(dtype)

	model_params["res2_3.conv2.weight"] = torch.from_numpy(data_dict["res3d_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_3.bn2.weight"] = torch.from_numpy(data_dict["bn3d_branch2b"]["scale"]).type(dtype)
	model_params["res2_3.bn2.bias"] = torch.from_numpy(data_dict["bn3d_branch2b"]["offset"]).type(dtype)

	model_params["res2_3.conv3.weight"] = torch.from_numpy(data_dict["res3d_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res2_3.bn3.weight"] = torch.from_numpy(data_dict["bn3d_branch2c"]["scale"]).type(dtype)
	model_params["res2_3.bn3.bias"] = torch.from_numpy(data_dict["bn3d_branch2c"]["offset"]).type(dtype)

	## third projection layer
	model_params["proj3.conv1.weight"] = torch.from_numpy(data_dict["res4a_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj3.bn1.weight"] = torch.from_numpy(data_dict["bn4a_branch2a"]["scale"]).type(dtype)
	model_params["proj3.bn1.bias"] = torch.from_numpy(data_dict["bn4a_branch2a"]["offset"]).type(dtype)

	model_params["proj3.conv2.weight"] = torch.from_numpy(data_dict["res4a_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj3.bn2.weight"] = torch.from_numpy(data_dict["bn4a_branch2b"]["scale"]).type(dtype)
	model_params["proj3.bn2.bias"] = torch.from_numpy(data_dict["bn4a_branch2b"]["offset"]).type(dtype)

	model_params["proj3.conv3.weight"] = torch.from_numpy(data_dict["res4a_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj3.bn3.weight"] = torch.from_numpy(data_dict["bn4a_branch2c"]["scale"]).type(dtype)
	model_params["proj3.bn3.bias"] = torch.from_numpy(data_dict["bn4a_branch2c"]["offset"]).type(dtype)

	model_params["proj3.conv4.weight"] = torch.from_numpy(data_dict["res4a_branch1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj3.bn4.weight"] = torch.from_numpy(data_dict["bn4a_branch1"]["scale"]).type(dtype)
	model_params["proj3.bn4.bias"] = torch.from_numpy(data_dict["bn4a_branch1"]["offset"]).type(dtype)

	#
	model_params["res3_1.conv1.weight"] = torch.from_numpy(data_dict["res4b_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_1.bn1.weight"] = torch.from_numpy(data_dict["bn4b_branch2a"]["scale"]).type(dtype)
	model_params["res3_1.bn1.bias"] = torch.from_numpy(data_dict["bn4b_branch2a"]["offset"]).type(dtype)

	model_params["res3_1.conv2.weight"] = torch.from_numpy(data_dict["res4b_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_1.bn2.weight"] = torch.from_numpy(data_dict["bn4b_branch2b"]["scale"]).type(dtype)
	model_params["res3_1.bn2.bias"] = torch.from_numpy(data_dict["bn4b_branch2b"]["offset"]).type(dtype)

	model_params["res3_1.conv3.weight"] = torch.from_numpy(data_dict["res4b_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_1.bn3.weight"] = torch.from_numpy(data_dict["bn4b_branch2c"]["scale"]).type(dtype)
	model_params["res3_1.bn3.bias"] = torch.from_numpy(data_dict["bn4b_branch2c"]["offset"]).type(dtype)

	#
	model_params["res3_2.conv1.weight"] = torch.from_numpy(data_dict["res4c_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_2.bn1.weight"] = torch.from_numpy(data_dict["bn4c_branch2a"]["scale"]).type(dtype)
	model_params["res3_2.bn1.bias"] = torch.from_numpy(data_dict["bn4c_branch2a"]["offset"]).type(dtype)

	model_params["res3_2.conv2.weight"] = torch.from_numpy(data_dict["res4c_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_2.bn2.weight"] = torch.from_numpy(data_dict["bn4c_branch2b"]["scale"]).type(dtype)
	model_params["res3_2.bn2.bias"] = torch.from_numpy(data_dict["bn4c_branch2b"]["offset"]).type(dtype)

	model_params["res3_2.conv3.weight"] = torch.from_numpy(data_dict["res4c_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_2.bn3.weight"] = torch.from_numpy(data_dict["bn4c_branch2c"]["scale"]).type(dtype)
	model_params["res3_2.bn3.bias"] = torch.from_numpy(data_dict["bn4c_branch2c"]["offset"]).type(dtype)

	#
	model_params["res3_3.conv1.weight"] = torch.from_numpy(data_dict["res4d_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_3.bn1.weight"] = torch.from_numpy(data_dict["bn4d_branch2a"]["scale"]).type(dtype)
	model_params["res3_3.bn1.bias"] = torch.from_numpy(data_dict["bn4d_branch2a"]["offset"]).type(dtype)

	model_params["res3_3.conv2.weight"] = torch.from_numpy(data_dict["res4d_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_3.bn2.weight"] = torch.from_numpy(data_dict["bn4d_branch2b"]["scale"]).type(dtype)
	model_params["res3_3.bn2.bias"] = torch.from_numpy(data_dict["bn4d_branch2b"]["offset"]).type(dtype)

	model_params["res3_3.conv3.weight"] = torch.from_numpy(data_dict["res4d_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_3.bn3.weight"] = torch.from_numpy(data_dict["bn4d_branch2c"]["scale"]).type(dtype)
	model_params["res3_3.bn3.bias"] = torch.from_numpy(data_dict["bn4d_branch2c"]["offset"]).type(dtype)

	#
	model_params["res3_4.conv1.weight"] = torch.from_numpy(data_dict["res4e_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_4.bn1.weight"] = torch.from_numpy(data_dict["bn4e_branch2a"]["scale"]).type(dtype)
	model_params["res3_4.bn1.bias"] = torch.from_numpy(data_dict["bn4e_branch2a"]["offset"]).type(dtype)

	model_params["res3_4.conv2.weight"] = torch.from_numpy(data_dict["res4e_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_4.bn2.weight"] = torch.from_numpy(data_dict["bn4e_branch2b"]["scale"]).type(dtype)
	model_params["res3_4.bn2.bias"] = torch.from_numpy(data_dict["bn4e_branch2b"]["offset"]).type(dtype)

	model_params["res3_4.conv3.weight"] = torch.from_numpy(data_dict["res4e_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_4.bn3.weight"] = torch.from_numpy(data_dict["bn4e_branch2c"]["scale"]).type(dtype)
	model_params["res3_4.bn3.bias"] = torch.from_numpy(data_dict["bn4e_branch2c"]["offset"]).type(dtype)

	#
	model_params["res3_5.conv1.weight"] = torch.from_numpy(data_dict["res4f_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_5.bn1.weight"] = torch.from_numpy(data_dict["bn4f_branch2a"]["scale"]).type(dtype)
	model_params["res3_5.bn1.bias"] = torch.from_numpy(data_dict["bn4f_branch2a"]["offset"]).type(dtype)

	model_params["res3_5.conv2.weight"] = torch.from_numpy(data_dict["res4f_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_5.bn2.weight"] = torch.from_numpy(data_dict["bn4f_branch2b"]["scale"]).type(dtype)
	model_params["res3_5.bn2.bias"] = torch.from_numpy(data_dict["bn4f_branch2b"]["offset"]).type(dtype)

	model_params["res3_5.conv3.weight"] = torch.from_numpy(data_dict["res4f_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res3_5.bn3.weight"] = torch.from_numpy(data_dict["bn4f_branch2c"]["scale"]).type(dtype)
	model_params["res3_5.bn3.bias"] = torch.from_numpy(data_dict["bn4f_branch2c"]["offset"]).type(dtype)

	## fourth projection layer
	model_params["proj4.conv1.weight"] = torch.from_numpy(data_dict["res5a_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj4.bn1.weight"] = torch.from_numpy(data_dict["bn5a_branch2a"]["scale"]).type(dtype)
	model_params["proj4.bn1.bias"] = torch.from_numpy(data_dict["bn5a_branch2a"]["offset"]).type(dtype)

	model_params["proj4.conv2.weight"] = torch.from_numpy(data_dict["res5a_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj4.bn2.weight"] = torch.from_numpy(data_dict["bn5a_branch2b"]["scale"]).type(dtype)
	model_params["proj4.bn2.bias"] = torch.from_numpy(data_dict["bn5a_branch2b"]["offset"]).type(dtype)

	model_params["proj4.conv3.weight"] = torch.from_numpy(data_dict["res5a_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj4.bn3.weight"] = torch.from_numpy(data_dict["bn5a_branch2c"]["scale"]).type(dtype)
	model_params["proj4.bn3.bias"] = torch.from_numpy(data_dict["bn5a_branch2c"]["offset"]).type(dtype)

	model_params["proj4.conv4.weight"] = torch.from_numpy(data_dict["res5a_branch1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["proj4.bn4.weight"] = torch.from_numpy(data_dict["bn5a_branch1"]["scale"]).type(dtype)
	model_params["proj4.bn4.bias"] = torch.from_numpy(data_dict["bn5a_branch1"]["offset"]).type(dtype)

	#
	model_params["res4_1.conv1.weight"] = torch.from_numpy(data_dict["res5b_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_1.bn1.weight"] = torch.from_numpy(data_dict["bn5b_branch2a"]["scale"]).type(dtype)
	model_params["res4_1.bn1.bias"] = torch.from_numpy(data_dict["bn5b_branch2a"]["offset"]).type(dtype)

	model_params["res4_1.conv2.weight"] = torch.from_numpy(data_dict["res5b_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_1.bn2.weight"] = torch.from_numpy(data_dict["bn5b_branch2b"]["scale"]).type(dtype)
	model_params["res4_1.bn2.bias"] = torch.from_numpy(data_dict["bn5b_branch2b"]["offset"]).type(dtype)

	model_params["res4_1.conv3.weight"] = torch.from_numpy(data_dict["res5b_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_1.bn3.weight"] = torch.from_numpy(data_dict["bn5b_branch2c"]["scale"]).type(dtype)
	model_params["res4_1.bn3.bias"] = torch.from_numpy(data_dict["bn5b_branch2c"]["offset"]).type(dtype)

	#
	model_params["res4_2.conv1.weight"] = torch.from_numpy(data_dict["res5c_branch2a"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_2.bn1.weight"] = torch.from_numpy(data_dict["bn5c_branch2a"]["scale"]).type(dtype)
	model_params["res4_2.bn1.bias"] = torch.from_numpy(data_dict["bn5c_branch2a"]["offset"]).type(dtype)

	model_params["res4_2.conv2.weight"] = torch.from_numpy(data_dict["res5c_branch2b"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_2.bn2.weight"] = torch.from_numpy(data_dict["bn5c_branch2b"]["scale"]).type(dtype)
	model_params["res4_2.bn2.bias"] = torch.from_numpy(data_dict["bn5c_branch2b"]["offset"]).type(dtype)

	model_params["res4_2.conv3.weight"] = torch.from_numpy(data_dict["res5c_branch2c"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["res4_2.bn3.weight"] = torch.from_numpy(data_dict["bn5c_branch2c"]["scale"]).type(dtype)
	model_params["res4_2.bn3.bias"] = torch.from_numpy(data_dict["bn5c_branch2c"]["offset"]).type(dtype)

	## Middle convolution layer
	model_params["conv2.weight"] = torch.from_numpy(data_dict["layer1"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["conv2.bias"] = torch.from_numpy(data_dict["layer1"]["biases"]).type(dtype)
	model_params["bn2.weight"] = torch.from_numpy(data_dict["layer1_BN"]["scale"]).type(dtype)
	model_params["bn2.bias"] = torch.from_numpy(data_dict["layer1_BN"]["offset"]).type(dtype)

	## Up projection layer 1
	model_params["UpProj1.UpConv1.conv1.weight"] = torch.from_numpy(data_dict["layer2x_br1_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv1.conv1.bias"] = torch.from_numpy(data_dict["layer2x_br1_ConvA"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv1.conv2.weight"] = torch.from_numpy(data_dict["layer2x_br1_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv1.conv2.bias"] = torch.from_numpy(data_dict["layer2x_br1_ConvB"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv1.conv3.weight"] = torch.from_numpy(data_dict["layer2x_br1_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv1.conv3.bias"] = torch.from_numpy(data_dict["layer2x_br1_ConvC"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv1.conv4.weight"] = torch.from_numpy(data_dict["layer2x_br1_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv1.conv4.bias"] = torch.from_numpy(data_dict["layer2x_br1_ConvD"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv1.bn.weight"] = torch.from_numpy(data_dict["layer2x_br1_BN"]['scale']).type(dtype)
	model_params["UpProj1.UpConv1.bn.bias"] = torch.from_numpy(data_dict["layer2x_br1_BN"]['offset']).type(dtype)

	#
	model_params["UpProj1.UpConv2.conv1.weight"] = torch.from_numpy(data_dict["layer2x_br2_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv2.conv1.bias"] = torch.from_numpy(data_dict["layer2x_br2_ConvA"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv2.conv2.weight"] = torch.from_numpy(data_dict["layer2x_br2_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv2.conv2.bias"] = torch.from_numpy(data_dict["layer2x_br2_ConvB"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv2.conv3.weight"] = torch.from_numpy(data_dict["layer2x_br2_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv2.conv3.bias"] = torch.from_numpy(data_dict["layer2x_br2_ConvC"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv2.conv4.weight"] = torch.from_numpy(data_dict["layer2x_br2_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.UpConv2.conv4.bias"] = torch.from_numpy(data_dict["layer2x_br2_ConvD"]["biases"]).type(dtype)

	model_params["UpProj1.UpConv2.bn.weight"] = torch.from_numpy(data_dict["layer2x_br2_BN"]['scale']).type(dtype)
	model_params["UpProj1.UpConv2.bn.bias"] = torch.from_numpy(data_dict["layer2x_br2_BN"]['offset']).type(dtype)

	model_params["UpProj1.conv1.weight"] = torch.from_numpy(data_dict["layer2x_Conv"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj1.conv1.bias"] = torch.from_numpy(data_dict["layer2x_Conv"]["biases"]).type(dtype)

	#
	model_params["UpProj1.bn.weight"] = torch.from_numpy(data_dict["layer2x_BN"]['scale']).type(dtype)
	model_params["UpProj1.bn.bias"] = torch.from_numpy(data_dict["layer2x_BN"]['offset']).type(dtype)

	## Up projection layer 2
	model_params["UpProj2.UpConv1.conv1.weight"] = torch.from_numpy(data_dict["layer4x_br1_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv1.conv1.bias"] = torch.from_numpy(data_dict["layer4x_br1_ConvA"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv1.conv2.weight"] = torch.from_numpy(data_dict["layer4x_br1_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv1.conv2.bias"] = torch.from_numpy(data_dict["layer4x_br1_ConvB"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv1.conv3.weight"] = torch.from_numpy(data_dict["layer4x_br1_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv1.conv3.bias"] = torch.from_numpy(data_dict["layer4x_br1_ConvC"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv1.conv4.weight"] = torch.from_numpy(data_dict["layer4x_br1_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv1.conv4.bias"] = torch.from_numpy(data_dict["layer4x_br1_ConvD"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv1.bn.weight"] = torch.from_numpy(data_dict["layer4x_br1_BN"]['scale']).type(dtype)
	model_params["UpProj2.UpConv1.bn.bias"] = torch.from_numpy(data_dict["layer4x_br1_BN"]['offset']).type(dtype)

	#
	model_params["UpProj2.UpConv2.conv1.weight"] = torch.from_numpy(data_dict["layer4x_br2_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv2.conv1.bias"] = torch.from_numpy(data_dict["layer4x_br2_ConvA"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv2.conv2.weight"] = torch.from_numpy(data_dict["layer4x_br2_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv2.conv2.bias"] = torch.from_numpy(data_dict["layer4x_br2_ConvB"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv2.conv3.weight"] = torch.from_numpy(data_dict["layer4x_br2_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv2.conv3.bias"] = torch.from_numpy(data_dict["layer4x_br2_ConvC"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv2.conv4.weight"] = torch.from_numpy(data_dict["layer4x_br2_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.UpConv2.conv4.bias"] = torch.from_numpy(data_dict["layer4x_br2_ConvD"]["biases"]).type(dtype)

	model_params["UpProj2.UpConv2.bn.weight"] = torch.from_numpy(data_dict["layer4x_br2_BN"]['scale']).type(dtype)
	model_params["UpProj2.UpConv2.bn.bias"] = torch.from_numpy(data_dict["layer4x_br2_BN"]['offset']).type(dtype)

	model_params["UpProj2.conv1.weight"] = torch.from_numpy(data_dict["layer4x_Conv"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj2.conv1.bias"] = torch.from_numpy(data_dict["layer4x_Conv"]["biases"]).type(dtype)

	#
	model_params["UpProj2.bn.weight"] = torch.from_numpy(data_dict["layer4x_BN"]['scale']).type(dtype)
	model_params["UpProj2.bn.bias"] = torch.from_numpy(data_dict["layer4x_BN"]['offset']).type(dtype)

	## Up projection layer 3
	model_params["UpProj3.UpConv1.conv1.weight"] = torch.from_numpy(data_dict["layer8x_br1_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv1.conv1.bias"] = torch.from_numpy(data_dict["layer8x_br1_ConvA"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv1.conv2.weight"] = torch.from_numpy(data_dict["layer8x_br1_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv1.conv2.bias"] = torch.from_numpy(data_dict["layer8x_br1_ConvB"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv1.conv3.weight"] = torch.from_numpy(data_dict["layer8x_br1_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv1.conv3.bias"] = torch.from_numpy(data_dict["layer8x_br1_ConvC"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv1.conv4.weight"] = torch.from_numpy(data_dict["layer8x_br1_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv1.conv4.bias"] = torch.from_numpy(data_dict["layer8x_br1_ConvD"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv1.bn.weight"] = torch.from_numpy(data_dict["layer8x_br1_BN"]['scale']).type(dtype)
	model_params["UpProj3.UpConv1.bn.bias"] = torch.from_numpy(data_dict["layer8x_br1_BN"]['offset']).type(dtype)

	#
	model_params["UpProj3.UpConv2.conv1.weight"] = torch.from_numpy(data_dict["layer8x_br2_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv2.conv1.bias"] = torch.from_numpy(data_dict["layer8x_br2_ConvA"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv2.conv2.weight"] = torch.from_numpy(data_dict["layer8x_br2_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv2.conv2.bias"] = torch.from_numpy(data_dict["layer8x_br2_ConvB"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv2.conv3.weight"] = torch.from_numpy(data_dict["layer8x_br2_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv2.conv3.bias"] = torch.from_numpy(data_dict["layer8x_br2_ConvC"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv2.conv4.weight"] = torch.from_numpy(data_dict["layer8x_br2_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.UpConv2.conv4.bias"] = torch.from_numpy(data_dict["layer8x_br2_ConvD"]["biases"]).type(dtype)

	model_params["UpProj3.UpConv2.bn.weight"] = torch.from_numpy(data_dict["layer8x_br2_BN"]['scale']).type(dtype)
	model_params["UpProj3.UpConv2.bn.bias"] = torch.from_numpy(data_dict["layer8x_br2_BN"]['offset']).type(dtype)

	model_params["UpProj3.conv1.weight"] = torch.from_numpy(data_dict["layer8x_Conv"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj3.conv1.bias"] = torch.from_numpy(data_dict["layer8x_Conv"]["biases"]).type(dtype)

	#
	model_params["UpProj3.bn.weight"] = torch.from_numpy(data_dict["layer8x_BN"]['scale']).type(dtype)
	model_params["UpProj3.bn.bias"] = torch.from_numpy(data_dict["layer8x_BN"]['offset']).type(dtype)

	## Up projection layer 4
	model_params["UpProj4.UpConv1.conv1.weight"] = torch.from_numpy(data_dict["layer16x_br1_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv1.conv1.bias"] = torch.from_numpy(data_dict["layer16x_br1_ConvA"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv1.conv2.weight"] = torch.from_numpy(data_dict["layer16x_br1_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv1.conv2.bias"] = torch.from_numpy(data_dict["layer16x_br1_ConvB"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv1.conv3.weight"] = torch.from_numpy(data_dict["layer16x_br1_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv1.conv3.bias"] = torch.from_numpy(data_dict["layer16x_br1_ConvC"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv1.conv4.weight"] = torch.from_numpy(data_dict["layer16x_br1_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv1.conv4.bias"] = torch.from_numpy(data_dict["layer16x_br1_ConvD"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv1.bn.weight"] = torch.from_numpy(data_dict["layer16x_br1_BN"]['scale']).type(dtype)
	model_params["UpProj4.UpConv1.bn.bias"] = torch.from_numpy(data_dict["layer16x_br1_BN"]['offset']).type(dtype)

	#
	model_params["UpProj4.UpConv2.conv1.weight"] = torch.from_numpy(data_dict["layer16x_br2_ConvA"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv2.conv1.bias"] = torch.from_numpy(data_dict["layer16x_br2_ConvA"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv2.conv2.weight"] = torch.from_numpy(data_dict["layer16x_br2_ConvB"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv2.conv2.bias"] = torch.from_numpy(data_dict["layer16x_br2_ConvB"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv2.conv3.weight"] = torch.from_numpy(data_dict["layer16x_br2_ConvC"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv2.conv3.bias"] = torch.from_numpy(data_dict["layer16x_br2_ConvC"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv2.conv4.weight"] = torch.from_numpy(data_dict["layer16x_br2_ConvD"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.UpConv2.conv4.bias"] = torch.from_numpy(data_dict["layer16x_br2_ConvD"]["biases"]).type(dtype)

	model_params["UpProj4.UpConv2.bn.weight"] = torch.from_numpy(data_dict["layer16x_br2_BN"]['scale']).type(dtype)
	model_params["UpProj4.UpConv2.bn.bias"] = torch.from_numpy(data_dict["layer16x_br2_BN"]['offset']).type(dtype)

	model_params["UpProj4.conv1.weight"] = torch.from_numpy(data_dict["layer16x_Conv"]["weights"]).type(dtype).permute(3,2,0,1)
	model_params["UpProj4.conv1.bias"] = torch.from_numpy(data_dict["layer16x_Conv"]["biases"]).type(dtype)

	#
	model_params["UpProj4.bn.weight"] = torch.from_numpy(data_dict["layer16x_BN"]['scale']).type(dtype)
	model_params["UpProj4.bn.bias"] = torch.from_numpy(data_dict["layer16x_BN"]['offset']).type(dtype)

	# final conv layer
	model_params['conv3.weight'] = torch.from_numpy(data_dict['ConvPred']['weights']).type(dtype).permute(3,2,0,1)
	model_params['conv3.bias'] = torch.from_numpy(data_dict['ConvPred']['biases']).type(dtype)

	return model_params

max_stars_repo_name: mohsaad/Deeper-Depth-Prediction


text:
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys

sys.path.append("../../common")

import test_util as tu
import tritonclient.http as httpclient
from tritonclient.utils import *
import numpy as np
import unittest

class ExplicitModelTest(tu.TestResultCollector):

    def send_identity_request(self, client, model_name):
        inputs = []
        inputs.append(httpclient.InferInput('INPUT0', [1, 16], "FP32"))
        input0_data = np.arange(start=0, stop=16, dtype=np.float32)
        input0_data = np.expand_dims(input0_data, axis=0)
        inputs[0].set_data_from_numpy(input0_data)
        result = client.infer(
            model_name=model_name,
            inputs=inputs,
            outputs=[httpclient.InferRequestedOutput('OUTPUT0')])
        output_numpy = result.as_numpy('OUTPUT0')
        self.assertTrue(np.all(input0_data == output_numpy))

    def test_model_reload(self):
        model_name = "identity_fp32"
        ensemble_model_name = 'simple_' + "identity_fp32"
        with httpclient.InferenceServerClient("localhost:8000") as client:
            for _ in range(5):
                self.assertFalse(client.is_model_ready(model_name))
                # Load the model before the ensemble model to make sure reloading the
                # model works properly in Python backend.
                client.load_model(model_name)
                client.load_model(ensemble_model_name)
                self.assertTrue(client.is_model_ready(model_name))
                self.assertTrue(client.is_model_ready(ensemble_model_name))
                self.send_identity_request(client, model_name)
                self.send_identity_request(client, ensemble_model_name)
                client.unload_model(ensemble_model_name)
                client.unload_model(model_name)
                self.assertFalse(client.is_model_ready(model_name))
                self.assertFalse(client.is_model_ready(ensemble_model_name))

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: galv/server

max_issues_repo_name: galv/server

max_forks_repo_name: galv/server


text:
import numpy as np
from scipy import spatial

def cosine_similarity(A, B):
    if len(A.shape) == 1 and len(B.shape) == 1:
        return 1 - spatial.distance.cosine(A, B)
    elif len(A.shape) == 2 and B.shape == A.shape:
        return np.array([cosine_similarity(A[i], B[i])
            for i in xrange(len(A))])
    else:
        raise ValueError

max_stars_repo_name: orestis-z/track-rcnn

max_issues_repo_name: orestis-z/track-rcnn

max_forks_repo_name: orestis-z/track-rcnn

reason: import numpy,from scipy


text:
#!/usr/bin/python
#author: zhaofeng-shu33

import numpy as np
from ace_cream import ace_cream

def pearson_correlation(X,Y):
    return (np.mean(X*Y, axis=0) -np.mean(X, axis = 0)* np.mean(Y, axis = 0)) / ( np.std(X, axis = 0) * np.std(Y, axis = 0))

if __name__ == '__main__':
    N_SIZE = 1000
    ERROR_PROBABILITY = 0.1
    x = np.random.choice([0,1],size=N_SIZE)
    y = np.random.uniform(size=N_SIZE)
    for i in range(len(x)):
        if(y[i] < ERROR_PROBABILITY):
            y[i] = 2
        else:
            y[i] = x[i]
    dic_Y = {0:6, 1:8, 2:3}
    dic_X = {0:7, 1:9}
    for i in range(len(y)):
        y[i] = dic_Y[y[i]]
        x[i] = dic_X[x[i]]
    print('rho(x,y)',pearson_correlation(x,y))
    # use fortran ace by 1985 article author
    tx, ty = ace_cream(x, y, cat = [-1,0])
    print('mapped X symbol list: ')
    print(np.unique(tx))
    print('mapped Y symbol list: ')
    print(np.unique(ty))

    print('mean(tx) = %f, std(tx) = %f'%(np.mean(tx), np.std(tx)))
    print('mean(ty) = %f, std(ty) = %f'%(np.mean(ty), np.std(ty)))
    print('rho(tx,ty)',pearson_correlation(tx,ty))
    # matches theoretical result: np.sqrt(1-ERROR_PROBABILITY)


text:
from tests.testing import *
from reem.datatypes import KeyValueStore
from reem.connection import RedisInterface
from reem import ships
import logging
import numpy as np

# Logging Configuration
log_file_name = "logs/reem_testing_kvs_timed.log"
FORMAT = "%(asctime)20s %(filename)30s:%(lineno)3s  %(funcName)20s() %(levelname)10s     %(message)s"
logging.basicConfig(format=FORMAT, filename=log_file_name, filemode='w')
logger = logging.getLogger("reem.datatypes")
logger.setLevel(logging.DEBUG)

image_array = np.random.rand(640, 480, 3)

flat_data = get_flat_data()
nested_data = get_nested_data()
image_dict = {"image": image_array}
hundred_key_dict = single_level_dictionary()
layered_dictionary = nested_level_dictionary(levels=20)

interface = RedisInterface(host="localhost", ships=[ships.NumpyShip()])
interface.initialize()

server = KeyValueStore(interface)

def test_write_hundred_key_dict():
    for i in range(100):
        server["hundred_key_dict"] = hundred_key_dict
        logger.debug("\n")

def test_write_multilevel_dict():
    for i in range(100):
        server["layered_dict"] = layered_dictionary
        logger.debug("\n")

def test_read_hundred_key_dict():
    server["hundred_key_dict"] = hundred_key_dict
    for i in range(100):
        # server["hundred_key_dict"].read()
        # server["hundred_key_dict"]["single_key"] = {"number":5}
        server["image_dictionary"] = image_dict
        logger.debug("\n")

def test_list_set():
    list_dict = {"list": ["sdfdsfdfsfdsfsdfsd" for i in range(100)]}
    for i in range(100):
        server["list_dict"] = list_dict
        logger.debug("\n")

# New Tests
"""
Reset Whole Dictionary vs Reset each key
Come up with workarounds for slow sets and warnings about performance
"""

max_stars_repo_name: joaomcm/reem

max_issues_repo_name: joaomcm/reem

max_forks_repo_name: joaomcm/reem


text:
#!/usr/bin/env python

import numpy as np
from utils import utils

ORIGIN = np.array((0, 0))

def manhattan_dist(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    return abs(x1 - x2) + abs(y1 - y2)

def build_wire(line):
    directions = {
        'R': np.array((1, 0)),
        'U': np.array((0, 1)),
        'L': np.array((-1, 0)),
        'D': np.array((0, -1))
    }

    wire = {}
    current = np.array(ORIGIN)
    step = 0
    for part in line.split(','):
        direction = part[0]
        x = int(part[1:])
        for _ in range(x):
            current += directions[direction]
            step += 1
            wire[tuple(current)] = step
    return wire

def find_intersections(wire1, wire2):
    unique_points_w1 = set(wire1.keys())
    unique_points_w2 = set(wire2.keys())
    intersections = unique_points_w1 & unique_points_w2
    return sorted(intersections, key=lambda p: manhattan_dist(p, ORIGIN))

def calculate(input_txt: str) -> (int, int):
    input_txt = input_txt.strip()
    line1, line2 = input_txt.split('\n')
    wire1 = build_wire(line1)
    wire2 = build_wire(line2)
    intersections = find_intersections(wire1, wire2)
    lowest_dist_p = intersections[0]
    distance = manhattan_dist(lowest_dist_p, ORIGIN)
    min_signal_delay = min([wire1[p] + wire2[p] for p in intersections])

    return distance, min_signal_delay

def main():
    input_txt = utils.get_input(3)
    p1, p2 = calculate(input_txt)

    print(f"Part 1: {p1}")
    print(f"Part 2: {p2}")

if __name__ == "__main__":
    main()

max_stars_repo_name: max-f/advent-of-code

max_issues_repo_name: max-f/advent-of-code

max_forks_repo_name: max-f/advent-of-code


text:
import numpy as np
import pytest

from mchap.assemble.likelihood import log_likelihood
from mchap.assemble import mcmc
from mchap.jitutils import seed_numba
from mchap.testing import simulate_reads
from mchap.encoding import integer
from mchap import mset

def test_point_beta_probabilities():

    probs = mcmc._point_beta_probabilities(6, 1, 1)
    assert np.allclose(probs.sum(), 1.0)
    assert np.allclose(probs[0], probs)

    probs1 = mcmc._point_beta_probabilities(6, 1, 3)
    probs2 = mcmc._point_beta_probabilities(6, 3, 1)
    assert np.allclose(probs1.sum(), 1.0)
    assert np.allclose(probs2.sum(), 1.0)
    assert np.allclose(probs1, probs2[::-1])

def test_read_mean_dist():
    reads = np.array(
        [
            [[0.9, 0.1], [0.8, 0.2], [0.8, 0.2]],
            [[0.9, 0.1], [0.8, 0.2], [0.8, 0.2]],
            [[0.9, 0.1], [0.2, 0.8], [np.nan, np.nan]],
            [[0.9, 0.1], [0.2, 0.8], [np.nan, np.nan]],
        ]
    )
    expect = np.array(
        [
            [0.9, 0.1],
            [0.5, 0.5],
            [0.8, 0.2],
        ]
    )
    actual = mcmc._read_mean_dist(reads)
    np.testing.assert_array_equal(actual, expect)

def test_homozygosity_probabilities():
    haplotypes = np.array(
        [
            [0, 0, 0, 0, 1, 1],
            [0, 1, 0, 0, 1, 1],
            [0, 1, 0, 1, 1, 1],
            [0, 1, 1, 1, 1, 1],
        ]
    )
    ploidy = len(haplotypes)
    n_alleles = [2] * 6

    # 16 reads
    reads = simulate_reads(
        haplotypes,
        n_reads=16,
        uniform_sample=True,
        errors=False,
    )
    actual = mcmc._homozygosity_probabilities(reads, n_alleles, ploidy) > 0.999
    expect = np.zeros((6, 2), dtype=bool)
    np.testing.assert_array_equal(actual, expect)

    # 32 reads
    reads = simulate_reads(
        haplotypes,
        n_reads=32,
        uniform_sample=True,
        errors=False,
    )
    actual = mcmc._homozygosity_probabilities(reads, n_alleles, ploidy) > 0.999
    expect = np.array(
        [
            [True, False],
            [False, False],
            [False, False],
            [False, False],
            [False, True],
            [False, True],
        ],
        dtype=bool,
    )
    np.testing.assert_array_equal(actual, expect)

def test_DenovoMCMC__zero_reads():
    n_chains = 2
    ploidy, n_base = 4, 6
    n_steps = 1000
    n_burn = 500
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    # zero reads
    reads = np.empty((0, n_base, 2), dtype=float)

    trace = model.fit(reads).burn(n_burn)
    posterior = trace.posterior()
    assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
    assert posterior.probabilities[0] < 0.05

def test_DenovoMCMC__zero_snps():
    n_chains = 2
    ploidy, n_base = 4, 0  # no snps in locus
    n_steps = 1000
    n_burn = 500
    n_alleles = []  # no snps in locus
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    reads = np.empty((10, n_base, 0), dtype=float)

    trace = model.fit(reads).burn(n_burn)
    posterior = trace.posterior()
    assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, 0)
    assert np.all(np.isnan(trace.llks))
    assert posterior.probabilities[0] == 1  # no variability

def test_DenovoMCMC__zero_reads_or_snps():
    n_chains = 2
    ploidy, n_base = 4, 0  # no snps in locus
    n_steps = 1000
    n_burn = 500
    n_alleles = []  # no snps in locus
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    reads = np.empty((0, n_base, 0), dtype=float)

    trace = model.fit(reads).burn(n_burn)
    posterior = trace.posterior()
    assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, 0)
    assert np.all(np.isnan(trace.llks))
    assert posterior.probabilities[0] == 1  # no variability

def test_DenovoMCMC__all_nans():
    n_chains = 2
    ploidy, n_base = 4, 6
    n_steps = 1000
    n_burn = 500
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    # high read depth
    reads = np.empty((10, n_base, 2), dtype=float)
    reads[:] = np.nan

    trace = model.fit(reads).burn(n_burn)
    posterior = trace.posterior()
    assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
    assert posterior.probabilities[0] < 0.05

def test_DenovoMCMC__nans_ignored():
    # nan-reads should have no affect on mcmc
    haplotypes = np.array(
        [
            [0, 0, 0, 0, 0, 0],
            [0, 1, 0, 1, 1, 1],
            [0, 1, 0, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
        ]
    )
    n_chains = 2
    ploidy, n_base = haplotypes.shape
    n_steps = 1000
    n_burn = 500
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    reads1 = simulate_reads(
        haplotypes,
        n_reads=16,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )

    nan_reads = np.empty((4, n_base, 2), dtype=reads1.dtype)
    nan_reads[:] = np.nan
    reads2 = np.concatenate([reads1, nan_reads])

    np.random.seed(42)
    seed_numba(42)
    trace1 = model.fit(reads1).burn(n_burn)

    np.random.seed(42)
    seed_numba(42)
    trace2 = model.fit(reads2).burn(n_burn)

    np.testing.assert_array_equal(trace1.genotypes, trace2.genotypes)

def test_DenovoMCMC__non_variable():
    haplotypes = np.array(
        [
            [0, 0, 0, 1, 1, 1],
            [0, 0, 0, 1, 1, 1],
            [0, 0, 0, 1, 1, 1],
        ]
    )
    n_chains = 2
    ploidy, n_base = haplotypes.shape
    n_steps = 100
    n_burn = 50
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    # high read depth
    reads = simulate_reads(
        haplotypes,
        n_reads=40,
        errors=False,
        qual=(60, 60),
    )

    trace = model.fit(reads).burn(n_burn)
    posterior = trace.posterior()

    assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
    # nan likelihoods for non-variable chain
    assert np.isnan(trace.llks).all()
    assert len(posterior.genotypes) == 1
    np.testing.assert_array_equal(haplotypes, posterior.genotypes[0])
    assert posterior.probabilities[0] == 1.0

def test_DenovoMCMC__diploid():

    np.random.seed(42)

    haplotypes = np.array(
        [
            [0, 0, 0, 1, 1, 1],
            [1, 0, 0, 0, 0, 0],
        ]
    )
    n_chains = 2
    ploidy, n_base = haplotypes.shape
    n_steps = 1000
    n_burn = 500
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
    )

    reads = simulate_reads(
        haplotypes,
        n_reads=2,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )

    for i in range(10):
        trace = model.fit(reads).burn(n_burn)
        posterior = trace.posterior()
        assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
        assert posterior.probabilities[0] > 0.90
        np.testing.assert_array_equal(haplotypes, posterior.genotypes[0])

def test_DenovoMCMC__tetraploid():

    haplotypes = np.array(
        [
            [0, 0, 0, 0, 0, 0],
            [0, 1, 0, 1, 1, 1],
            [0, 1, 0, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
        ]
    )
    n_chains = 2
    ploidy, n_base = haplotypes.shape
    n_steps = 1000
    n_burn = 500
    n_alleles = [2] * n_base
    model = mcmc.DenovoMCMC(
        ploidy=ploidy,
        n_alleles=n_alleles,
        steps=n_steps,
        chains=n_chains,
        random_seed=42,
    )

    # high read depth
    reads = simulate_reads(
        haplotypes,
        n_reads=40,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )
    for _ in range(10):
        trace = model.fit(reads).burn(n_burn)
        posterior = trace.posterior()
        assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
        assert posterior.probabilities[0] > 0.90
        np.testing.assert_array_equal(haplotypes, posterior.genotypes[0])
        assert np.any((trace.genotypes[0] != trace.genotypes[1]))

    # medium read depth
    reads = simulate_reads(
        haplotypes,
        n_reads=16,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )
    for _ in range(10):
        trace = model.fit(reads).burn(n_burn)
        posterior = trace.posterior()
        assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
        assert 0.30 < posterior.probabilities[0] < 0.90
        np.testing.assert_array_equal(haplotypes, posterior.genotypes[0])
        assert np.any((trace.genotypes[0] != trace.genotypes[1]))

    # low read depth
    reads = simulate_reads(
        haplotypes,
        n_reads=8,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )
    for _ in range(10):
        trace = model.fit(reads).burn(n_burn)
        posterior = trace.posterior()
        assert trace.genotypes.shape == (n_chains, n_steps - n_burn, ploidy, n_base)
        assert posterior.probabilities[0] < 0.30
        assert np.any((trace.genotypes[0] != trace.genotypes[1]))

def test_DenovoMCMC__seed():

    haplotypes = np.array(
        [
            [0, 0, 0, 0, 0, 0],
            [0, 1, 0, 1, 1, 1],
            [0, 1, 0, 1, 1, 1],
            [1, 1, 1, 1, 1, 1],
        ]
    )
    n_chains = 2
    ploidy, n_base = haplotypes.shape
    n_steps = 1000
    n_alleles = [2] * n_base

    # medium read depth
    reads = simulate_reads(
        haplotypes,
        n_reads=16,
        uniform_sample=True,
        errors=False,
        qual=(60, 60),
    )

    model_1 = mcmc.DenovoMCMC(
        ploidy, n_alleles, steps=n_steps, chains=n_chains, random_seed=42
    )
    model_2 = mcmc.DenovoMCMC(
        ploidy, n_alleles, steps=n_steps, chains=n_chains, random_seed=33
    )
    model_3 = mcmc.DenovoMCMC(
        ploidy, n_alleles, steps=n_steps, chains=n_chains, random_seed=42
    )

    trace_1 = model_1.fit(reads)
    trace_2 = model_2.fit(reads)
    trace_3 = model_3.fit(reads)

    # compair among runs
    assert np.any((trace_1.genotypes != trace_2.genotypes))
    np.testing.assert_array_equal(trace_1.genotypes, trace_3.genotypes)

    # compair within runs (among chains)
    assert np.any((trace_1.genotypes[0] != trace_1.genotypes[1]))
    assert np.any((trace_2.genotypes[0] != trace_2.genotypes[1]))

def test_DenovoMCMC__fuzz():
    for _ in range(10):
        n_chains = np.random.randint(1, 4)
        ploidy = np.random.randint(2, 5)
        n_base = np.random.randint(4, 8)
        n_steps = np.random.randint(100, 400)
        n_reads = np.random.randint(1, 50)

        haplotypes = np.random.choice(
            [0, 1, 2],  # small chance of triallelic base
            p=[0.45, 0.45, 0.1],
            size=(ploidy, n_base),
        )
        n_alleles = [np.max(haplotypes) + 1] * n_base
        model = mcmc.DenovoMCMC(
            ploidy=ploidy, n_alleles=n_alleles, steps=n_steps, chains=n_chains
        )
        reads = simulate_reads(haplotypes, n_reads=n_reads)
        trace = model.fit(reads)
        assert trace.genotypes.shape == (n_chains, n_steps, ploidy, n_base)
        assert trace.llks.shape == (n_chains, n_steps)

@pytest.mark.parametrize(
    "temperatures",
    [
        [1.0],
        [0.1, 1.0],
        [0.01, 0.1, 1.0],
        [0.001, 0.01, 0.1, 1.0],
    ],
)
def test_DenovoMCMC__temperatures_bias(temperatures):
    """Test to ensure the implementation of parallel-tempering does
    not bias the MCMC results when tempering is not needed.
    """
    reads = np.array(
        [
            [[0.9, 0.1], [0.1, 0.9]],
            [[0.9, 0.1], [0.1, 0.9]],
            [[0.9, 0.1], [0.9, 0.1]],
            [[0.9, 0.1], [0.9, 0.1]],
        ]
    )
    genotypes = np.array(
        [
            [[0, 0], [0, 0]],  # 2
            [[0, 0], [0, 1]],  # 1:1
            [[0, 0], [1, 0]],  # 1:1
            [[0, 0], [1, 1]],  # 1:1
            [[0, 1], [0, 1]],  # 2
            [[0, 1], [1, 0]],  # 1:1
            [[0, 1], [1, 1]],  # 1:1
            [[1, 0], [1, 0]],  # 2
            [[1, 0], [1, 1]],  # 1:1
            [[1, 1], [1, 1]],  # 2
        ],
        dtype=np.int8,
    )

    # calculate exact posteriors
    llks = np.array([log_likelihood(reads, g) for g in genotypes])
    priors = np.array([1, 2, 2, 2, 1, 2, 2, 1, 2, 1])
    priors = priors / priors.sum()
    exact_posteriors = np.exp(llks + np.log(priors))
    exact_posteriors = exact_posteriors / exact_posteriors.sum()

    # simulate posteriors
    model = mcmc.DenovoMCMC(
        ploidy=2,
        n_alleles=[2, 2],
        steps=50000,
        chains=1,
        random_seed=11,
        temperatures=temperatures,
    )
    trace = model.fit(reads)
    posterior = trace.burn(100).posterior()
    simulation_posteriors = {g.tobytes(): 0 for g in genotypes}
    simulation_posteriors.update(
        {g.tobytes(): p for g, p in zip(posterior.genotypes, posterior.probabilities)}
    )
    simulation_posteriors = [simulation_posteriors[g.tobytes()] for g in genotypes]

    # check posteriors are similar
    np.testing.assert_array_almost_equal(
        exact_posteriors,
        simulation_posteriors,
        decimal=2,
    )

def test_DenovoMCMC__temperatures_submode():
    """This test is taken from real world example in which the mcmc
    can become stuck in a sub-optimal alternative mode when parallel
    tempering is not used. Note that the optimal mode is not necessarily
    the true genotype and the data there may be additional haplotypes due
    to duplication.
    """
    read_counts = [
        ("000000", 22),
        ("-00000", 3),
        ("0000-0", 2),
        ("00000-", 1),
        ("000-00", 1),
        ("001000", 16),
        ("-01000", 2),
        ("001---", 2),
        ("--1000", 3),
        ("000110", 39),
        ("-00110", 1),
        ("----10", 1),
        ("100101", 3),
        ("-00101", 3),
        ("----01", 1),
        ("110100", 1),
        ("110---", 2),
        ("1101--", 1),
        ("11----", 1),
        ("0--000", 1),
        ("---000", 3),
        ("----00", 4),
        ("0-----", 2),
        ("1-----", 2),
        ("001110", 1),
    ]
    strings, counts = zip(*read_counts)
    calls = integer.from_strings(strings)
    counts = np.array(list(counts))
    dists = integer.as_probabilistic(
        calls,
        n_alleles=2,
        p=0.999,
    )
    reads = mset.repeat(dists, counts)
    n_alleles = [2] * 6

    alt_mode = np.array(
        [
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 1, 0, 0, 0],
            [1, 0, 0, 1, 0, 1],
        ]
    )

    opt_mode = np.array(
        [
            [0, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 0, 1, 1, 0],
            [0, 0, 1, 0, 0, 0],
            [1, 0, 0, 1, 0, 1],
            [1, 1, 0, 1, 0, 0],
        ]
    )

    # initialise all chains within sub-optimal mode
    initial = np.tile(alt_mode, (2, 1, 1))

    # without tempering
    model = mcmc.DenovoMCMC(
        ploidy=6, n_alleles=n_alleles, steps=1100, temperatures=[1.0], random_seed=1
    )
    posterior = model.fit(reads, initial=initial).burn(100).posterior()
    np.testing.assert_array_equal(posterior.genotypes[0], alt_mode)  # sub-optimal mode
    assert posterior.probabilities[0] > 0.95

    # with tempering
    model = mcmc.DenovoMCMC(
        ploidy=6,
        n_alleles=n_alleles,
        steps=1100,
        temperatures=[0.1, 1.0],
        random_seed=1,
    )
    posterior = model.fit(reads, initial=initial).burn(100).posterior()
    np.testing.assert_array_equal(posterior.genotypes[0], opt_mode)  # optimal mode
    assert posterior.probabilities[0] > 0.95

max_stars_repo_name: PlantandFoodResearch/MCHap

max_issues_repo_name: PlantandFoodResearch/MCHap

max_forks_repo_name: PlantandFoodResearch/MCHap


text:
"""
Plotting for the Huys task. The functions have been written
for the object in resourceAllocator.py (gradient-free
optimization - either CMAES or Bayesian Optimization for
the equal precision model).
Also plots comparison of gradient-based with gradient-free
after the relevant simulation results have been produced.
"""

import itertools
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib, matplotlib.pyplot as plt
sns.set()

def plot_table(obj, tableColour='mean'):
    # Setting up memory columns other than resourceAllocVector
    a   = np.array([' left','right'])
    s   = np.arange(6) + 1
    r   = obj.env.rewards[:12]
    s1  = ((obj.env.transitions[:12]+1)[:,1] - 6)
    sa  = [sa for sa in itertools.product(s,a)] * (obj.depth-1)
    rs  = list(zip(r,s1)) * (obj.depth-1)
    q   = obj.q[:-12]
    d   = {'(s,a)': sa, '(r,s\')': rs, 'q_mu(s,a)':q}

    # Setting up resource allocation vectors for table of memories
    nMems = obj.results.shape[-1]
    means = np.transpose( np.mean(obj.results, axis=1) )
    stds  = np.transpose( np.std(obj.results, axis=1) )
    cols  = np.around( (np.array(obj.searchBudget) / \
            (obj.depth * 2**obj.depth) * 100), decimals=1)

    # Table of optimal resource allocation vectors for different search budgets
    df = pd.DataFrame(d)
    if len(means)==1:       # if equal precision => broadcast
        means = np.array(list(means)*len(df))
        nMems = len(df)
    dm = pd.DataFrame(means.astype(int), columns=cols)
    df = pd.merge(df, dm, left_index=True, right_index=True)
    vals = dm.values
    extraCols = 3

    # Setting up figure
    fig = plt.figure()
    ax = fig.add_subplot(111, frameon=False, xticks=[], yticks=[])

    # Depth dependent parameters:
    if obj.depth==3:
        fh = 6
        cbx= 0.42
        sc = 1.2
        tx = 0.4
    elif obj.depth>=4:
        fh = obj.depth*3 - 2
        cbx= 0.32
        sc = 1.35
        tx = 0.2

    fig.set_figheight(fh)
    fig.set_figwidth(8)

    # setting colourmap either based on means only, or both means & std
    if tableColour=='mean':
        norm = plt.Normalize(0,100)
        colours = plt.cm.viridis(norm(vals))
        colours = np.concatenate((np.array([[[0]*4]*extraCols]*nMems),colours), axis=1)
        fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.viridis), ax=ax\
            , cax= fig.add_axes([0.42, 0.04, 0.54, 0.02]), orientation='horizontal')
    elif tableColour=='std':
        norm = plt.Normalize(0,50)
        colours = plt.cm.Greens_r(norm(stds))
        colours = np.concatenate((np.array([[[0]*4]*extraCols]*nMems),colours), axis=1)
        fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.Greens_r), ax=ax\
            , cax= fig.add_axes([cbx, 0.04, 0.54, 0.02]), orientation='horizontal')

    # Plotting the table of memories according to the colourmap
    the_table=matplotlib.table.table(ax, cellText=df.values, rowLabels=None, \
                colLabels=df.columns, loc='center', cellColours=colours)
    fig.text(tx,0.925,'Resource allocation, i.e. q_sigma(s,a), for search budget (%)')
    the_table.scale(1,sc)  # depth 4: 1.35
    plt.show()
    return fig, ax

# Plotting <reward>, cost, and overall objective vs search budget
def plot_curves(obj):
    # Initialising arrays
    reward      = np.nan * np.ones( len(obj.searchBudget) )
    reward_std  = np.nan * np.ones( len(obj.searchBudget) )
    cost        = np.nan * np.ones( len(obj.searchBudget) )
    cost_std    = np.nan * np.ones( len(obj.searchBudget) )
    objective   = np.nan * np.ones( len(obj.searchBudget) )
    obj_std     = np.nan * np.ones( len(obj.searchBudget) )

    # Simulating to get expected rewards, cost, and obj as funcs of sigma_opt
    for idx, budget in enumerate(obj.searchBudget):
        sigma = obj.results[idx,:,:]    # size (n_restarts x n_mems)
        r = np.zeros(len(sigma))
        c = np.zeros(len(sigma))
        o = np.zeros(len(sigma))
        for i in range(len(sigma)):
            r[i] = obj.expectedRewards(sigma=sigma[i], searchBudget=budget)
            c[i] = obj.cost(sigma=sigma[i])
            o[i] = r[i] - c[i]
        reward[idx]     = np.mean(r)
        reward_std[idx] = np.std(r)
        cost[idx]       = np.mean(c)
        cost_std[idx]   = np.std(c)
        objective[idx]  = np.mean(o)
        obj_std[idx]    = np.std(o)

    # Storing the arrays in a pandas dataframe
    d = {'Reward': reward, 'Cost': cost, 'Objective': objective,\
        'err_reward': reward_std, 'err_cost': cost_std, 'err_obj': obj_std}
    obj.dfResults = pd.DataFrame(d)

    # Setting up the figure subplots and axes
    f, (ax,ax2) = plt.subplots(nrows=2,ncols=1,sharex=True)#,facecolor='w')
    xaxis = np.array(obj.searchBudget) / (obj.depth * 2**obj.depth) * 100
    ax2.set_xlabel('Search Budget (% tree explored)', fontsize = 16)
    ax.set_ylabel('(a.u.)', fontsize = 16)

    # Plot same data on both axis (hack to get different colour for each curve)
    ax.errorbar(xaxis, reward, yerr = reward_std, label = 'Reward')
    ax.errorbar(xaxis, objective, yerr = obj_std, label = 'Objective')
    ax.errorbar(xaxis, cost, yerr = cost_std, label = 'Cost')
    ax2.errorbar(xaxis, reward, yerr = reward_std, label = 'Reward')
    ax2.errorbar(xaxis, objective, yerr = obj_std, label = 'Objective')
    ax2.errorbar(xaxis, cost, yerr = cost_std, label = 'Cost')

    # Set axis limits
    ax.set_ylim(50,80)
    ax2.set_ylim(0,15)

    # Hide the spines between ax and ax2
    ax.spines['bottom'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax.xaxis.tick_top()
    ax.tick_params(labeltop='off')

    # Making the kinks or diagonal break-lines between the subplots
    d = .01  # size of diagonal kinks
    # arguments to pass to plot, just so we don't keep repeating them
    kwargs = dict(transform=ax.transAxes, color='k', clip_on=False)
    ax.plot((-d, +d), (-d, +d), **kwargs)        # top-left diagonal
    ax.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # top-right diagonal
    kwargs.update(transform=ax2.transAxes)  # switch to the bottom axes
    ax2.plot((-d, +d), (1 - d, 1 + d), **kwargs)  # bottom-left diagonal
    ax2.plot((1 - d, 1 + d), (1 - d, 1 + d), **kwargs)  # bottom-right diagonal

    # Show plot with legend
    ax2.legend()
    plt.show()
    return f, ax, ax2

# Plotting the table of memories
def plot_dprime(obj, table=True):
    # Setting up memory columns other than resourceAllocVector
    states = np.arange(6) + 1
    q = obj.q[:-12]

    # Calculating d' and setting column names
    cols = np.around((np.array(obj.searchBudget) / \
            (obj.depth * 2**obj.depth) * 100), decimals=1)
    m1 = (q[1:len(q):2]).reshape(int(len(q)/2), 1)      # mean for action 'right'
    m2 = (q[0:len(q)+1:2]).reshape(int(len(q)/2), 1)    # mean for action 'left'
    dprime = np.zeros( (len(m1),len(obj.searchBudget),obj.results.shape[1]) )
    for i in range(obj.results.shape[1]):
        s1 = np.transpose(obj.results[:, i, 1:len(q):2])
        s2 = np.transpose(obj.results[:, i, 0:len(q)+1:2])
        dprime[:,:,i] = abs(m1-m2) / np.sqrt( s1**2 + s2**2 )
    dprime_mean = np.mean(dprime, axis=2)
    nMems = len(m1)

    # Creating dataframe for plotting
    d = {'State': list(states)*(obj.depth-1), 'Reward diff.': abs(m1-m2).ravel().astype(str)}
    dm = pd.DataFrame(np.around(dprime_mean, decimals=1), columns=cols)
    vals = dm.values
    df = pd.merge(pd.DataFrame(d), dm, left_index=True, right_index=True)
    obj.dprime = df

    if table:
        # setting plotting values right
        extraCols = 2

        # setting up figure
        fig = plt.figure()
        ax = fig.add_subplot(111, frameon=False, xticks=[], yticks=[])

        # Depth dependent parameters:
        if obj.depth==3:
            fh = 4
            cbx= 0.42
            sc = 1.2
            tx = 0.3
        elif obj.depth>=4:
            fh = 2*obj.depth - 2
            cbx= 0.32
            sc = 1.35
            tx = 0.2

        fig.set_figheight(fh)
        fig.set_figwidth(8)

        # setting colourmap based on mean dprime values
        norm = plt.Normalize(0, 4)
        colours = plt.cm.viridis_r(norm(vals))
        colours = np.concatenate((np.array([[[0]*4]*extraCols]*nMems),colours), axis=1)
        fig.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.viridis_r), ax=ax\
            , cax= fig.add_axes([cbx, 0.1, 0.54, 0.02]), orientation='horizontal')

        # Plotting the table of memories according to the colourmap
        the_table=matplotlib.table.table(ax, cellText=df.values, rowLabels=None, \
                    colLabels=df.columns, loc='center', cellColours=colours)
        fig.text(tx,0.875,'d\', i.e., discriminability between q-distributions, for search budget (%)')
        the_table.scale(1,sc)
        plt.show()
        return fig, ax
    else:
        # Dataframe of all dprime values
        dprime_full = dprime.transpose(1,0,2).reshape(dprime.shape[1],-1).transpose()
        d = {'State': list(np.repeat(states, dprime.shape[2])) * (obj.depth-1), \
            'Reward diff.': np.repeat(abs(m1-m2).ravel(), dprime.shape[2])}
        dm = pd.DataFrame(np.around(dprime_full, decimals=1), columns=cols)
        dfull = pd.merge(pd.DataFrame(d), dm, left_index=True, right_index=True)

        # Figure settings
        palette = sns.color_palette("mako_r", len(obj.searchBudget))
        sns.set()

        # separating different depths
        df_depth = list(range(obj.depth-1))
        for i in range(obj.depth-1):
            blockLen = 6 * obj.n_restarts
            df_depth[i] = pd.melt(dfull.iloc[i*blockLen:(i+1)*blockLen,1:], id_vars=['Reward diff.'], \
                        value_vars=cols, var_name='Search budget', value_name='d-prime')
            # Plotting figure for each depth
            f = plt.figure()
            ax1 = f.subplots()
            ax1 = sns.lineplot(x='Reward diff.', y='d-prime', hue='Search budget', \
                            data=df_depth[i], palette=palette).set_title(f'Depth {i+1} memories')
            plt.show()
        return ax1

############### Plotting comparison with CMAES ##################
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import itertools
import pickle
sns.set()

# Defining the cost functions
def kl_mvn(m0, S0, m1, S1):
    """
    Kullback-Liebler divergence from Gaussian pm,pv to Gaussian qm,qv.
    Diagonal covariances are assumed.  Divergence is expressed in nats.
    """
    # store inv diag covariance of S1 and diff between means
    N = m0.shape[0]
    iS1 = np.linalg.inv(S1)
    diff = m1 - m0

    # kl is made of three terms
    tr_term   = np.trace(iS1 @ S0)
    det_term  = np.log(np.linalg.det(S1)/np.linalg.det(S0))
    quad_term = diff.T @ np.linalg.inv(S1) @ diff
    return .5 * (tr_term + det_term + quad_term - N)

# Initialising the lists
x               = []
method          = []
expectedRewards = []
cost            = []
objective       = []

# Initiliazing path and other variables
nRuns   = 5
path    = "./figures/"  # initialResults/gradientBasedLearning/q-learning_adv

# Gradient-based optimization for all depths
depth = 3
for lmda,N_traj in itertools.product([0.15,0.3,1],[1,10]):
    method += [f'$\\nabla_\sigma$, $N_{{traj}}$={N_traj}, $q_{{*}}$']*nRuns
    x      += [f'Depth {depth}, lambda={lmda}']*nRuns
    expectedRewards += list(np.load(path+f"expR_depth{depth}_{int(lmda*100)}_{N_traj}.npy"))
    # computing the cost
    df_allRuns = pickle.load(open(path +\
            f"df_allRuns_depth{depth}_{int(lmda*100)}_{N_traj}",'rb'))
    for df in df_allRuns:
        q          = df['q'][:-12]
        sigma      = df['sigma'][:-12]
        cost      += [lmda * kl_mvn(q, np.diag(np.square(sigma)), \
                            q, np.diag(np.square( [100]*(len(q)) )))]
    objective += list(np.array(expectedRewards)[-nRuns:] - np.array(cost)[-nRuns:])

lmda = 1
for depth,N_traj in itertools.product([4,5],[1,10]):
    method += [f'$\\nabla_\sigma$, $N_{{traj}}$={N_traj}, $q_{{*}}$']*nRuns
    x      += [f'Depth {depth}, lambda={lmda}']*nRuns
    expectedRewards += list(np.load(path+f"expR_depth{depth}_{int(lmda*100)}_{N_traj}.npy"))
    # computing the cost
    df_allRuns = pickle.load(open(path+f"df_allRuns_depth{depth}_{int(lmda*100)}_{N_traj}",'rb'))
    for df in df_allRuns:
        q          = df['q'][:-12]
        sigma      = df['sigma'][:-12]
        cost      += [lmda * kl_mvn(q, np.diag(np.square(sigma)), \
                            q, np.diag(np.square( [100]*(len(q)) )))]
    # Objective = Expected rewards - cost
    objective += list(np.array(expectedRewards)[-nRuns:] - np.array(cost)[-nRuns:])

# Depth 3, lambda = {1, 0.3, 0.15}; gradient-free (CMAES) optimization
x       +=  ['Depth 3, lambda=1']*1000 #+\
            # ['Depth 3, lambda=0.3']*1000 +\
            # ['Depth 3, lambda=0.15']*1000
method  += ['Gradient-free (CMAES)']*1000

for i in [100]: #,30,15]:
    dfResults = pickle.load(open(f'./initialResults/Depth3/DKL_lmda{i}_variablePrecision/dfResults','rb'))
    meanR = dfResults.iloc[0]['Reward']
    stdR  = dfResults.iloc[0]['err_reward']
    meanC = dfResults.iloc[0]['Cost']
    stdC  = dfResults.iloc[0]['err_cost']
    meanO = dfResults.iloc[0]['Objective']
    stdO  = dfResults.iloc[0]['err_obj']
    expectedRewards  += list(meanR + np.random.randn(1000)*stdR)
    cost             += list(meanC + np.random.randn(1000)*stdC)
    objective        += list(meanO + np.random.randn(1000)*stdO)

# Depth 4, lambda = 1; gradient-free (CMAES) optimization
x      += ['Depth 4, lambda=1']*1000
method += ['Gradient-free (CMAES)']*1000
model   = pickle.load(open("./initialResults/Depth4/varPrecModel.pkl",'rb'))
meanR   = model.dfResults.iloc[0]['Reward']
stdR    = model.dfResults.iloc[0]['err_reward']
meanC   = model.dfResults.iloc[0]['Cost']
stdC    = model.dfResults.iloc[0]['err_cost']
meanO   = model.dfResults.iloc[0]['Objective']
stdO    = model.dfResults.iloc[0]['err_obj']
expectedRewards  += list(meanR + np.random.randn(1000)*stdR)
cost             += list(meanC + np.random.randn(1000)*stdC)
objective        += list(meanO + np.random.randn(1000)*stdO)

# Creating final dataframe for plotting
df = pd.DataFrame({'x':x, 'Method':method, 'Expected Rewards':expectedRewards,\
                   'Cost':cost, 'Objective':objective})

# Plotting
toPlot  = ["Expected Rewards", "Cost", "Objective"]
color   = ["Blues_d", "Greens_d", "Reds_d"]
for y,colr in zip(toPlot,color):
    fig, ax = plt.subplots()
    fig.set_figwidth(10)
    ax = sns.barplot(x="x", y=y, hue="Method",data=df, ci="sd", palette=colr)
    plt.show()
# ax.set_ylim([])

######################## Comparison #####################
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
import itertools
import pickle
sns.set()

# Initialising the lists
x               = []
method          = []
expectedRewards = []
cost            = []
objective       = []

# Initiliazing path and other variables
nRuns   = 5
basePath   = f"./initialResults/gradientBasedLearning/onPolicy-learning"
subPath    = [f"/", f"_return/", f"_qpiSamples/", f"_qpiMean/", f"_adv/"]
gradUpdate = [f"$r(\\tau)$", f"$G_t$", f"$q_\pi$", f"$meanq_\pi$", f"A"]

for subP, label in zip(subPath, gradUpdate):
    depth = 3
    for lmda,N_traj in itertools.product([0.15,0.3,1],[10]):
        method += [label]*nRuns
        x      += [f'Depth {depth}, lambda={lmda}']*nRuns
        expectedRewards += list(np.load(basePath+subP+f"expR_depth{depth}_{int(lmda*100)}_{N_traj}.npy"))
        # computing the cost
        df_allRuns = pickle.load(open(basePath + subP +\
                f"df_allRuns_depth{depth}_{int(lmda*100)}_{N_traj}",'rb'))
        for df in df_allRuns:
            q          = df['q'][:-12]
            sigma      = df['sigma'][:-12]
            cost      += [lmda * kl_mvn(q, np.diag(np.square(sigma)), \
                                q, np.diag(np.square( [100]*(len(q)) )))]
        objective += list(np.array(expectedRewards)[-nRuns:] - np.array(cost)[-nRuns:])

    lmda = 1
    for depth,N_traj in itertools.product([4,5],[10]):
        method += [label]*nRuns
        x      += [f'Depth {depth}, lambda={lmda}']*nRuns
        expectedRewards += list(np.load(basePath+subP+f"expR_depth{depth}_{int(lmda*100)}_{N_traj}.npy"))
        # computing the cost
        df_allRuns = pickle.load(open(basePath+subP+f"df_allRuns_depth{depth}_{int(lmda*100)}_{N_traj}",'rb'))
        for df in df_allRuns:
            q          = df['q'][:-12]
            sigma      = df['sigma'][:-12]
            cost      += [lmda * kl_mvn(q, np.diag(np.square(sigma)), \
                                q, np.diag(np.square( [100]*(len(q)) )))]
        # Objective = Expected rewards - cost
        objective += list(np.array(expectedRewards)[-nRuns:] - np.array(cost)[-nRuns:])

    depth = 10
    if subP in [f"/", f"_qpiMean/", f"_adv/"]:
        method += [label]*nRuns
        x      += [f'Depth {depth}, lambda={lmda}']*nRuns
        expectedRewards += list(np.load(basePath+subP+f"expR_depth{depth}_{int(lmda*100)}_{N_traj}.npy"))
        # computing the cost
        df_allRuns = pickle.load(open(basePath+subP+f"df_allRuns_depth{depth}_{int(lmda*100)}_{N_traj}",'rb'))
        for df in df_allRuns:
            q          = df['q'][:-12]
            sigma      = df['sigma'][:-12]
            cost      += [lmda * kl_mvn(q, np.diag(np.square(sigma)), \
                                q, np.diag(np.square( [100]*(len(q)) )))]
        # Objective = Expected rewards - cost
        objective += list(np.array(expectedRewards)[-nRuns:] - np.array(cost)[-nRuns:])

# Creating final dataframe for plotting
df = pd.DataFrame({'x':x, 'Method':method, 'Expected Rewards':expectedRewards,\
                   'Cost':cost, 'Objective':objective})

# Plotting
toPlot  = ["Expected Rewards", "Cost", "Objective"]
color   = ["Blues_d", "Greens_d", "Reds_d"]
for y,colr in zip(toPlot,color):
    fig, ax = plt.subplots()
    fig.set_figwidth(10)
    ax = sns.barplot(x="x", y=y, hue="Method",data=df, ci="sd", palette=colr)
    plt.show()
# ax.set_ylim([])

max_stars_repo_name: nisheetpatel/DynamicResourceAllocator

max_issues_repo_name: nisheetpatel/DynamicResourceAllocator

max_forks_repo_name: nisheetpatel/DynamicResourceAllocator


text:
#!/usr/bin/env nemesis
#
# ======================================================================
#
# Brad T. Aagaard, U.S. Geological Survey
#
# This code was developed as part of the Computational Infrastructure
# for Geodynamics (http://geodynamics.org).
#
# Copyright (c) 2010-2017 University of California, Davis
#
# See COPYING for license information.
#
# ======================================================================
#

import unittest

import numpy
from spatialdata.geocoords.CSCart import CSCart
from spatialdata.spatialdb.SimpleGridAscii import SimpleGridAscii
from spatialdata.spatialdb.SimpleGridDB import SimpleGridDB

# ----------------------------------------------------------------------------------------------------------------------
class TestSimpleGridDB(unittest.TestCase):

    def setUp(self):
        db = SimpleGridDB()
        db.inventory.label = "test"
        db.inventory.queryType = "nearest"
        db.inventory.filename = "data/grid.spatialdb"
        db._configure()
        self._db = db

    def test_database(self):
        locs = numpy.array([[3.5, 0.1, 0.8], [8.0, -0.4, 1.8]], numpy.float64)
        cs = CSCart()
        cs._configure()
        queryVals = ["two", "one"]
        dataE = numpy.array([[0.0220, 1.04], [0.0222, 1.02]], numpy.float64)
        errE = [0, 0]

        db = self._db
        db.open()
        db.setQueryValues(queryVals)
        data = numpy.zeros(dataE.shape, dtype=numpy.float64)
        err = []
        nlocs = locs.shape[0]
        for i in range(nlocs):
            e = db.query(data[i, :], locs[i, :], cs)
            err.append(e)
        db.close()

        self.assertEqual(len(errE), len(err))
        for vE, v in zip(errE, err):
            self.assertEqual(vE, v)

        self.assertEqual(len(dataE.shape), len(data.shape))
        for dE, d in zip(dataE.shape, data.shape):
            self.assertEqual(dE, d)
        for vE, v in zip(numpy.reshape(dataE, -1), numpy.reshape(data, -1)):
            self.assertAlmostEqual(vE, v, 6)

    def test_databasemulti(self):
        locs = numpy.array([[2.3, -0.1, 1.8], [2.0,  0.4, -2.0]], numpy.float64)
        cs = CSCart()
        cs._configure()
        queryVals = ["two", "one"]
        dataE = numpy.array([[0.0207, 1.17], [0.0209, 1.15]], numpy.float64)
        errE = numpy.array([0, 0], numpy.int32)

        db = self._db
        db.open()
        db.setQueryValues(queryVals)
        data = numpy.zeros(dataE.shape, dtype=numpy.float64)
        err = numpy.zeros(errE.shape, dtype=numpy.int32)
        db.multiquery(data, err, locs, cs)
        db.close()

        self.assertEqual(len(errE), len(err))
        for vE, v in zip(errE, err):
            self.assertEqual(vE, v)

        self.assertEqual(len(dataE.shape), len(data.shape))
        for dE, d in zip(dataE.shape, data.shape):
            self.assertEqual(dE, d)
        for vE, v in zip(numpy.reshape(dataE, -1), numpy.reshape(data, -1)):
            self.assertAlmostEqual(vE, v, 6)

    def test_io_3d(self):
        from spatialdata.spatialdb.SimpleGridAscii import createWriter

        filename = "data/gridio3d.spatialdb"
        x = numpy.array([-2.0, 0.0, 3.0], dtype=numpy.float64)
        y = numpy.array([0.0, 1.0], dtype=numpy.float64)
        z = numpy.array([-2.0, -1.0, 2.0], dtype=numpy.float64)

        points = numpy.array([
            [-2.0,  0.0, -2.0],
            [-2.0,  1.0, -2.0],
            [-2.0,  0.0, -1.0],
            [-2.0,  1.0, -1.0],
            [-2.0,  0.0,  2.0],
            [-2.0,  1.0,  2.0],
            [0.0,  0.0, -2.0],
            [0.0,  1.0, -2.0],  # query (5.7, 8.2)
            [0.0,  0.0, -1.0],
            [0.0,  1.0, -1.0],
            [0.0,  0.0,  2.0],
            [0.0,  1.0,  2.0],
            [3.0,  0.0, -2.0],
            [3.0,  1.0, -2.0],
            [3.0,  0.0, -1.0],
            [3.0,  1.0, -1.0],
            [3.0,  0.0,  2.0],
            [3.0,  1.0,  2.0],
        ], dtype=numpy.float64)
        one = numpy.array([6.6, 5.5, 2.3, 5.7, 6.3, 3.4, 7.2, 5.7, 3.4, 5.7, 9.4,
                           7.2, 4.8, 9.2, 5.8, 4.7, 7.8, 2.9], dtype=numpy.float64)
        two = numpy.array([3.4, 6.7, 4.1, 2.0, 6.7, 6.4, 6.8, 8.2, 9.8, 2.3, 8.5,
                           9.3, 7.5, 8.3, 8.5, 8.9, 6.2, 8.3], dtype=numpy.float64)

        cs = CSCart()

        writer = createWriter(filename)
        writer.write({'points': points,
                      'x': x,
                      'y': y,
                      'z': z,
                      'coordsys': cs,
                      'data_dim': 3,
                      'values': [{'name': "one",
                                  'units': "m",
                                  'data': one},
                                 {'name': "two",
                                  'units': "m",
                                  'data': two},
                                 ]})

        db = SimpleGridDB()
        db.inventory.label = "test"
        db.inventory.queryType = "nearest"
        db.inventory.filename = filename
        db._configure()
        self._db = db

        locs = numpy.array([[0.1, 0.95, -1.8]], numpy.float64)
        cs = CSCart()
        cs._configure()
        queryVals = ["two", "one"]
        dataE = numpy.array([[8.2, 5.7]], numpy.float64)
        errE = [0]

        db = self._db
        db.open()
        db.setQueryValues(queryVals)
        data = numpy.zeros(dataE.shape, dtype=numpy.float64)
        err = []
        nlocs = locs.shape[0]
        for i in range(nlocs):
            e = db.query(data[i, :], locs[i, :], cs)
            err.append(e)
        db.close()

        self.assertEqual(len(errE), len(err))
        for vE, v in zip(errE, err):
            self.assertEqual(vE, v)

        self.assertEqual(len(dataE.shape), len(data.shape))
        for dE, d in zip(dataE.shape, data.shape):
            self.assertEqual(dE, d)
        for vE, v in zip(numpy.reshape(dataE, -1), numpy.reshape(data, -1)):
            self.assertAlmostEqual(vE, v, 6)

    def test_io_2d(self):
        from spatialdata.spatialdb.SimpleGridAscii import createWriter

        filename = "data/gridio2d.spatialdb"
        x = numpy.array([-2.0, 0.0, 3.0], dtype=numpy.float64)
        y = numpy.array([0.0, 1.0], dtype=numpy.float64)

        points = numpy.array([
            [-2.0,  0.0],
            [-2.0,  1.0],
            [0.0,  0.0],
            [0.0,  1.0],  # query (5.7, 8.2)
            [3.0,  0.0],
            [3.0,  1.0],
        ], dtype=numpy.float64)
        one = numpy.array([6.6, 5.5, 7.2, 5.7, 4.8, 9.2], dtype=numpy.float64)
        two = numpy.array([3.4, 6.7, 6.8, 8.2, 7.5, 8.3], dtype=numpy.float64)

        cs = CSCart()
        cs.inventory.spaceDim = 2
        cs._configure()

        writer = createWriter(filename)
        writer.write({'points': points,
                      'x': x,
                      'y': y,
                      'coordsys': cs,
                      'data_dim': 2,
                      'values': [{'name': "one",
                                  'units': "m",
                                  'data': one},
                                 {'name': "two",
                                  'units': "m",
                                  'data': two},
                                 ]})

        db = SimpleGridDB()
        db.inventory.label = "test"
        db.inventory.queryType = "nearest"
        db.inventory.filename = filename
        db._configure()
        self._db = db

        locs = numpy.array([[0.1, 0.95]], numpy.float64)
        queryVals = ["two", "one"]
        dataE = numpy.array([[8.2, 5.7]], numpy.float64)
        errE = [0]

        db = self._db
        db.open()
        db.setQueryValues(queryVals)
        data = numpy.zeros(dataE.shape, dtype=numpy.float64)
        err = []
        nlocs = locs.shape[0]
        for i in range(nlocs):
            e = db.query(data[i, :], locs[i, :], cs)
            err.append(e)
        db.close()

        self.assertEqual(len(errE), len(err))
        for vE, v in zip(errE, err):
            self.assertEqual(vE, v)

        self.assertEqual(len(dataE.shape), len(data.shape))
        for dE, d in zip(dataE.shape, data.shape):
            self.assertEqual(dE, d)
        for vE, v in zip(numpy.reshape(dataE, -1), numpy.reshape(data, -1)):
            self.assertAlmostEqual(vE, v, 6)

# ----------------------------------------------------------------------------------------------------------------------
if __name__ == '__main__':
    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(TestSimpleGridDB))
    unittest.TextTestRunner(verbosity=2).run(suite)

# End of file

max_stars_repo_name: rwalkerlewis/spatialdata

max_issues_repo_name: rwalkerlewis/spatialdata

max_forks_repo_name: rwalkerlewis/spatialdata


text:
#!/usr/bin/python
# -*- coding: UTF-8 -*-
import io
import numpy as np
import os
import sys
import yaml

from pathlib import Path

ROOT_DIR = os.environ.get('DEFUND12_EMAIL_DIR')
if ROOT_DIR is None:
  print('$DEFUND12_EMAIL_DIR is not set')
  sys.exit(1)

ALLOWLISTED_KEYS = [
  'title',
  'permalink',
  'name',
  'state',
  'city',
  'layout',
  'recipients',
  'subject',
  'body'
]

def success(test_name):
  print('✅' + ' ' + test_name)

def fail(error_message):
  print('❌' + ' ' + error_message)
  sys.exit(1)

def test_files_exist():
  files = [file for file in Path(ROOT_DIR).rglob('*.md')]
  _, counts = np.unique([file.parent for file in files ], return_counts=True)
  if counts.sum() == 0:
    fail('test received no files at at path %s' % ROOT_DIR)

def validate_document_has_allowlisted_keys(doc, filepath):
  for allowlisted_key in ALLOWLISTED_KEYS:
    if allowlisted_key not in doc:
      fail('allowlisted_key key %s not found in file %s' % (allowlisted_key, filepath))

def get_markdown_files():
  markdown_files = []
  for subdir, dirs, files in os.walk(ROOT_DIR):
    for file in files:
      if file.endswith('.md'):
        markdown_files.append(os.path.join(subdir, file))
  return markdown_files

def test_files_contain_allowlisted_keys():
  for filepath in get_markdown_files():
    with open(filepath, 'r') as stream:
      docs = yaml.load_all(stream)
      for doc in docs:
        if doc is None:
            continue
        validate_document_has_allowlisted_keys(doc, filepath)

def test_files_contain_allowlisted_keys():
  for filepath in get_markdown_files():
    with open(filepath, 'r') as stream:
      docs = yaml.load_all(stream)
      for doc in filter(None, docs):
        validate_document_has_allowlisted_keys(doc, filepath)

def test_files_contain_unique_permalinks():
  permalinks = []
  filepaths = []
  for filepath in get_markdown_files():
    with open(filepath, 'r') as stream:
      docs = yaml.load_all(stream)
      for doc in filter(None, docs):
        permalink = doc['permalink']
        if permalink in permalinks:
          index = permalinks.index(permalink)
          fail('permalink %s in file %s already exists in file %s' % (permalink, filepath, filepaths[index]))
        permalinks.append(permalink)
        filepaths.append(filepath)

def main():
  print('Running markdown file tests...')

  test_files_exist()
  success('test_files_exist')

  test_files_contain_allowlisted_keys()
  success('test_files_contain_allowlisted_keys')

  test_files_contain_unique_permalinks()
  success('test_files_contain_unique_permalinks')

  print('All tests pass!')
  sys.exit(0)

if __name__ == "__main__":
    main()