text:
import numpy as np
import tqdm
from losses.dsm import anneal_dsm_score_estimation, anneal_dsm_score_estimation_TraceTrick
from losses.sliced_sm import anneal_sliced_score_estimation_vr
from losses.sliced_sm import anneal_ESM_scorenet, anneal_ESM_scorenet_VR
import torch.nn.functional as F
import logging
import torch
import os
import shutil
import tensorboardX
import torch.optim as optim
from torchvision.datasets import MNIST, CIFAR10, SVHN
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset
from datasets.celeba import CelebA
from models.cond_refinenet_dilated import CondRefineNetDilated
from torchvision.utils import save_image, make_grid
from PIL import Image
import time
__all__ = ['AnnealRunner']

class AnnealRunner():
    def __init__(self, args, config):
        self.args = args
        self.config = config

    def get_optimizer(self, parameters):
        if self.config.optim.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.config.optim.lr, weight_decay=self.config.optim.weight_decay,
                              betas=(self.config.optim.beta1, 0.999), amsgrad=self.config.optim.amsgrad)
        elif self.config.optim.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.config.optim.lr, weight_decay=self.config.optim.weight_decay)
        elif self.config.optim.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.config.optim.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.config.optim.optimizer))

    def logit_transform(self, image, lam=1e-6):
        image = lam + (1 - 2 * lam) * image
        return torch.log(image) - torch.log1p(-image)

    def train(self):
        if self.config.data.random_flip is False:
            tran_transform = test_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.ToTensor()
            ])
        else:
            tran_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ToTensor()
            ])
            test_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.ToTensor()
            ])

        if self.config.data.dataset == 'CIFAR10':
            dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10'), train=True, download=True,
                              transform=tran_transform)
            test_dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10_test'), train=False, download=True,
                                   transform=test_transform)
        elif self.config.data.dataset == 'MNIST':
            dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist'), train=True, download=True,
                            transform=tran_transform)
            test_dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist_test'), train=False, download=True,
                                 transform=test_transform)

        elif self.config.data.dataset == 'CELEBA':
            if self.config.data.random_flip:
                dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
                                 transform=transforms.Compose([
                                     transforms.CenterCrop(140),
                                     transforms.Resize(self.config.data.image_size),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                 ]), download=False)
            else:
                dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
                                 transform=transforms.Compose([
                                     transforms.CenterCrop(140),
                                     transforms.Resize(self.config.data.image_size),
                                     transforms.ToTensor(),
                                 ]), download=False)

            test_dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba_test'), split='test',
                                  transform=transforms.Compose([
                                      transforms.CenterCrop(140),
                                      transforms.Resize(self.config.data.image_size),
                                      transforms.ToTensor(),
                                  ]), download=False)

        elif self.config.data.dataset == 'SVHN':
            dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn'), split='train', download=True,
                           transform=tran_transform)
            test_dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn_test'), split='test', download=True,
                                transform=test_transform)

        dataloader = DataLoader(dataset, batch_size=self.config.training.batch_size, shuffle=True, num_workers=4)
        test_loader = DataLoader(test_dataset, batch_size=self.config.training.batch_size, shuffle=True,
                                 num_workers=4, drop_last=True)

        test_iter = iter(test_loader)
        self.config.input_dim = self.config.data.image_size ** 2 * self.config.data.channels

        tb_path = os.path.join(self.args.run, 'tensorboard', self.args.doc)
        if os.path.exists(tb_path):
            shutil.rmtree(tb_path)

        tb_logger = tensorboardX.SummaryWriter(log_dir=tb_path)
        score = CondRefineNetDilated(self.config).to(self.config.device)

        score = torch.nn.DataParallel(score)

        optimizer = self.get_optimizer(score.parameters())

        if self.args.resume_training:
            states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'))
            score.load_state_dict(states[0])
            optimizer.load_state_dict(states[1])

        step = 0

        sigmas = torch.tensor(
            np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                               self.config.model.num_classes))).float().to(self.config.device)

        time_record = []
        for epoch in range(self.config.training.n_epochs):
            for i, (X, y) in enumerate(dataloader):
                step += 1
                score.train()
                X = X.to(self.config.device)
                X = X / 256. * 255. + torch.rand_like(X) / 256.

                if self.config.data.logit_transform:
                    X = self.logit_transform(X)

                labels = torch.randint(0, len(sigmas), (X.shape[0],), device=X.device)
                if self.config.training.algo == 'dsm':
                    t = time.time()
                    loss = anneal_dsm_score_estimation(score, X, labels, sigmas, self.config.training.anneal_power)
                elif self.config.training.algo == 'dsm_tracetrick':
                    t = time.time()
                    loss = anneal_dsm_score_estimation_TraceTrick(score, X, labels, sigmas, self.config.training.anneal_power)
                elif self.config.training.algo == 'ssm':
                    t = time.time()
                    loss = anneal_sliced_score_estimation_vr(score, X, labels, sigmas,
                                                             n_particles=self.config.training.n_particles)
                elif self.config.training.algo == 'esm_scorenet':
                    t = time.time()
                    loss = anneal_ESM_scorenet(score, X, labels, sigmas,
                                                             n_particles=self.config.training.n_particles)
                elif self.config.training.algo == 'esm_scorenet_VR':
                    t = time.time()
                    loss = anneal_ESM_scorenet_VR(score, X, labels, sigmas,
                                                             n_particles=self.config.training.n_particles)

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                t = time.time() - t
                time_record.append(t)

                if step >= self.config.training.n_iters:
                    return 0

                if step % 100 == 0:
                    tb_logger.add_scalar('loss', loss, global_step=step)
                    logging.info("step: {}, loss: {}, time per step: {:.3f} +- {:.3f} ms".format(step, loss.item(),
                                                                                np.mean(time_record) * 1e3, np.std(time_record) * 1e3))

                    # if step % 2000 == 0:
                    #     score.eval()
                    #     try:
                    #         test_X, test_y = next(test_iter)
                    #     except StopIteration:
                    #         test_iter = iter(test_loader)
                    #         test_X, test_y = next(test_iter)

                    #     test_X = test_X.to(self.config.device)
                    #     test_X = test_X / 256. * 255. + torch.rand_like(test_X) / 256.

                    #     if self.config.data.logit_transform:
                    #         test_X = self.logit_transform(test_X)

                    #     test_labels = torch.randint(0, len(sigmas), (test_X.shape[0],), device=test_X.device)

                    #     #if self.config.training.algo == 'dsm':
                    #     with torch.no_grad():
                    #         test_dsm_loss = anneal_dsm_score_estimation(score, test_X, test_labels, sigmas,
                    #                                                         self.config.training.anneal_power)

                    #     tb_logger.add_scalar('test_dsm_loss', test_dsm_loss, global_step=step)
                    #     logging.info("step: {}, test dsm loss: {}".format(step, test_dsm_loss.item()))

                        # elif self.config.training.algo == 'ssm':
                        #     test_ssm_loss = anneal_sliced_score_estimation_vr(score, test_X, test_labels, sigmas,
                        #                                          n_particles=self.config.training.n_particles)

                        #     tb_logger.add_scalar('test_ssm_loss', test_ssm_loss, global_step=step)
                        #     logging.info("step: {}, test ssm loss: {}".format(step, test_ssm_loss.item()))

                if step >= 140000 and step % self.config.training.snapshot_freq == 0:
                    states = [
                        score.state_dict(),
                        optimizer.state_dict(),
                    ]
                    torch.save(states, os.path.join(self.args.log, 'checkpoint_{}.pth'.format(step)))
                    torch.save(states, os.path.join(self.args.log, 'checkpoint.pth'))

    def Langevin_dynamics(self, x_mod, scorenet, n_steps=200, step_lr=0.00005):
        images = []

        labels = torch.ones(x_mod.shape[0], device=x_mod.device) * 9
        labels = labels.long()

        with torch.no_grad():
            for _ in range(n_steps):
                images.append(torch.clamp(x_mod, 0.0, 1.0).to('cpu'))
                noise = torch.randn_like(x_mod) * np.sqrt(step_lr * 2)
                grad = scorenet(x_mod, labels)
                x_mod = x_mod + step_lr * grad + noise
                x_mod = x_mod
                print("modulus of grad components: mean {}, max {}".format(grad.abs().mean(), grad.abs().max()))

            return images

    def anneal_Langevin_dynamics(self, x_mod, scorenet, sigmas, n_steps_each=100, step_lr=0.00002):
        images = []

        with torch.no_grad():
            for c, sigma in tqdm.tqdm(enumerate(sigmas), total=len(sigmas), desc='annealed Langevin dynamics sampling'):
                labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
                labels = labels.long()
                step_size = step_lr * (sigma / sigmas[-1]) ** 2
                for s in range(n_steps_each):
                    images.append(torch.clamp(x_mod, 0.0, 1.0).to('cpu'))
                    noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                    grad = scorenet(x_mod, labels)
                    x_mod = x_mod + step_size * grad + noise
                    # print("class: {}, step_size: {}, mean {}, max {}".format(c, step_size, grad.abs().mean(),
                    #                                                          grad.abs().max()))

            return images

    def test(self):
        states = torch.load(os.path.join(self.args.log, 'checkpoint_199000.pth'), map_location=self.config.device)
        score = CondRefineNetDilated(self.config).to(self.config.device)
        score = torch.nn.DataParallel(score)

        score.load_state_dict(states[0])

        if not os.path.exists(self.args.image_folder):
            os.makedirs(self.args.image_folder)

        sigmas = np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                                    self.config.model.num_classes))

        score.eval()
        grid_size = 5

        imgs = []
        if self.config.data.dataset == 'MNIST':
            samples = torch.rand(grid_size ** 2, 1, 28, 28, device=self.config.device)
            all_samples = self.anneal_Langevin_dynamics(samples, score, sigmas, 100, 0.00002)

            for i, sample in enumerate(tqdm.tqdm(all_samples, total=len(all_samples), desc='saving images')):
                sample = sample.view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
                                     self.config.data.image_size)

                if self.config.data.logit_transform:
                    sample = torch.sigmoid(sample)

                image_grid = make_grid(sample, nrow=grid_size)
                if i % 10 == 0:
                    im = Image.fromarray(image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
                    imgs.append(im)

                save_image(image_grid, os.path.join(self.args.image_folder, 'image_{}.png'.format(i)))
                torch.save(sample, os.path.join(self.args.image_folder, 'image_raw_{}.pth'.format(i)))

        # else:
        #     samples = torch.rand(grid_size ** 2, 3, 32, 32, device=self.config.device)

        #     all_samples = self.anneal_Langevin_dynamics(samples, score, sigmas, 100, 0.00002)

        #     for i, sample in enumerate(tqdm.tqdm(all_samples, total=len(all_samples), desc='saving images')):
        #         sample = sample.view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
        #                              self.config.data.image_size)

        #         if self.config.data.logit_transform:
        #             sample = torch.sigmoid(sample)

        #         image_grid = make_grid(sample, nrow=grid_size)
        #         if i % 10 == 0:
        #             im = Image.fromarray(image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
        #             imgs.append(im)

        #         save_image(image_grid, os.path.join(self.args.image_folder, 'image_{}.png'.format(i)), nrow=10)
        #         torch.save(sample, os.path.join(self.args.image_folder, 'image_raw_{}.pth'.format(i)))

        # imgs[0].save(os.path.join(self.args.image_folder, "movie.gif"), save_all=True, append_images=imgs[1:], duration=1, loop=0)

        else:
            grid_size = 10
            for i in range(10):
                samples = torch.rand(grid_size ** 2, 3, 32, 32, device=self.config.device)
                all_samples = self.anneal_Langevin_dynamics(samples, score, sigmas, 100, 0.00002)
                sample = all_samples[len(all_samples)-1].view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
                                         self.config.data.image_size)

                image_grid = make_grid(sample, nrow=grid_size)
                im = Image.fromarray(image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
                im.save(os.path.join(self.args.image_folder, 'image_final_{}.png'.format(i)))

    def anneal_Langevin_dynamics_inpainting(self, x_mod, refer_image, scorenet, sigmas, n_steps_each=100,
                                            step_lr=0.000008):
        images = []

        refer_image = refer_image.unsqueeze(1).expand(-1, x_mod.shape[1], -1, -1, -1)
        refer_image = refer_image.contiguous().view(-1, 3, 32, 32)
        x_mod = x_mod.view(-1, 3, 32 ,32)
        half_refer_image = refer_image[..., :16]
        with torch.no_grad():
            for c, sigma in tqdm.tqdm(enumerate(sigmas), total=len(sigmas), desc="annealed Langevin dynamics sampling"):
                labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
                labels = labels.long()
                step_size = step_lr * (sigma / sigmas[-1]) ** 2

                corrupted_half_image = half_refer_image + torch.randn_like(half_refer_image) * sigma
                x_mod[:, :, :, :16] = corrupted_half_image
                for s in range(n_steps_each):
                    images.append(torch.clamp(x_mod, 0.0, 1.0).to('cpu'))
                    noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                    grad = scorenet(x_mod, labels)
                    x_mod = x_mod + step_size * grad + noise
                    x_mod[:, :, :, :16] = corrupted_half_image
                    # print("class: {}, step_size: {}, mean {}, max {}".format(c, step_size, grad.abs().mean(),
                    #                                                          grad.abs().max()))

            return images

    def test_inpainting(self):
        states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'), map_location=self.config.device)
        score = CondRefineNetDilated(self.config).to(self.config.device)
        score = torch.nn.DataParallel(score)

        score.load_state_dict(states[0])

        if not os.path.exists(self.args.image_folder):
            os.makedirs(self.args.image_folder)

        sigmas = np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                                    self.config.model.num_classes))
        score.eval()

        imgs = []
        if self.config.data.dataset == 'CELEBA':
            dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='test',
                             transform=transforms.Compose([
                                 transforms.CenterCrop(140),
                                 transforms.Resize(self.config.data.image_size),
                                 transforms.ToTensor(),
                             ]), download=True)

            dataloader = DataLoader(dataset, batch_size=20, shuffle=True, num_workers=4)
            refer_image, _ = next(iter(dataloader))

            samples = torch.rand(20, 20, 3, self.config.data.image_size, self.config.data.image_size,
                                 device=self.config.device)

            all_samples = self.anneal_Langevin_dynamics_inpainting(samples, refer_image, score, sigmas, 100, 0.00002)
            torch.save(refer_image, os.path.join(self.args.image_folder, 'refer_image.pth'))

            for i, sample in enumerate(tqdm.tqdm(all_samples)):
                sample = sample.view(400, self.config.data.channels, self.config.data.image_size,
                                     self.config.data.image_size)

                if self.config.data.logit_transform:
                    sample = torch.sigmoid(sample)

                image_grid = make_grid(sample, nrow=20)
                if i % 10 == 0:
                    im = Image.fromarray(image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
                    imgs.append(im)

                save_image(image_grid, os.path.join(self.args.image_folder, 'image_completion_{}.png'.format(i)))
                torch.save(sample, os.path.join(self.args.image_folder, 'image_completion_raw_{}.pth'.format(i)))

        else:
            transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.ToTensor()
            ])

            if self.config.data.dataset == 'CIFAR10':
                dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10'), train=True, download=True,
                                  transform=transform)
            elif self.config.data.dataset == 'SVHN':
                dataset = SVHN(os.path.join(self.args.run, 'datasets', 'svhn'), split='train', download=True,
                               transform=transform)

            dataloader = DataLoader(dataset, batch_size=20, shuffle=True, num_workers=4)
            data_iter = iter(dataloader)
            refer_image, _ = next(data_iter)

            torch.save(refer_image, os.path.join(self.args.image_folder, 'refer_image.pth'))
            samples = torch.rand(20, 20, self.config.data.channels, self.config.data.image_size,
                                 self.config.data.image_size).to(self.config.device)

            all_samples = self.anneal_Langevin_dynamics_inpainting(samples, refer_image, score, sigmas, 100, 0.00002)

            for i, sample in enumerate(tqdm.tqdm(all_samples)):
                sample = sample.view(400, self.config.data.channels, self.config.data.image_size,
                                     self.config.data.image_size)

                if self.config.data.logit_transform:
                    sample = torch.sigmoid(sample)

                image_grid = make_grid(sample, nrow=20)
                if i % 10 == 0:
                    im = Image.fromarray(image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
                    imgs.append(im)

                save_image(image_grid, os.path.join(self.args.image_folder, 'image_completion_{}.png'.format(i)))
                torch.save(sample, os.path.join(self.args.image_folder, 'image_completion_raw_{}.pth'.format(i)))

        imgs[0].save(os.path.join(self.args.image_folder, "movie.gif"), save_all=True, append_images=imgs[1:], duration=1, loop=0)

### add by tianyu ###
    def anneal_Langevin_dynamics_GenerateImages(self, x_mod, scorenet, sigmas, n_steps_each=100, step_lr=0.00002):

        with torch.no_grad():
            for c, sigma in enumerate(sigmas):
                labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
                labels = labels.long()
                step_size = step_lr * (sigma / sigmas[-1]) ** 2
                for s in range(n_steps_each):
                    noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                    grad = scorenet(x_mod, labels)
                    x_mod = x_mod + step_size * grad + noise
                    # print("class: {}, step_size: {}, mean {}, max {}".format(c, step_size, grad.abs().mean(),
                    #                                                          grad.abs().max()))

            return torch.clamp(x_mod, 0.0, 1.0).to('cpu')

    def save_sampled_images(self):
        num_of_step = 100
        bs = 100

        sigmas = np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                                        self.config.model.num_classes))

        if not os.path.exists(self.args.image_folder):
            os.makedirs(self.args.image_folder)

        print('Load checkpoint from' + self.args.log)

        for epochs in [100000,170000]:
            states = torch.load(os.path.join(self.args.log, 'checkpoint_'+str(epochs)+'.pth'), map_location=self.config.device)
            score = CondRefineNetDilated(self.config).to(self.config.device)
            score = torch.nn.DataParallel(score)

            score.load_state_dict(states[0])
            score.eval()

            if not os.path.exists(os.path.join(self.args.image_folder, 'epochs'+str(epochs))):
                os.makedirs(os.path.join(self.args.image_folder, 'epochs'+str(epochs)))

            save_index = 0
            print("Begin epochs", epochs)
            for j in range(num_of_step):
                samples = torch.rand(bs, 3, 32, 32, device=self.config.device)
                all_samples = self.anneal_Langevin_dynamics_GenerateImages(samples, score, sigmas, 100, 0.00002)
                images_new = all_samples.mul_(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 1).to('cpu', torch.uint8).numpy()
                for k in range(len(images_new)):
                    ims = Image.fromarray(images_new[k])
                    ims.save(os.path.join(self.args.image_folder, 'epochs'+str(epochs), 'img_'+str(save_index)+'.png'))
                    print('Save images ', k)
                    save_index+=1

    def calculate_fid(self):
        import fid
        import tensorflow as tf

        num_of_step = 500
        bs = 100

        sigmas = np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                                        self.config.model.num_classes))
        stats_path = 'fid_stats_cifar10_train.npz' # training set statistics
        inception_path = fid.check_or_download_inception(None) # download inception network

        print('Load checkpoint from' + self.args.log)
        #for epochs in range(140000, 200001, 1000):
        for epochs in [149000]:
            states = torch.load(os.path.join(self.args.log, 'checkpoint_'+str(epochs)+'.pth'), map_location=self.config.device)
            #states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'), map_location=self.config.device)
            score = CondRefineNetDilated(self.config).to(self.config.device)
            score = torch.nn.DataParallel(score)

            score.load_state_dict(states[0])

            score.eval()

            if self.config.data.dataset == 'MNIST':
                print("Begin epochs", epochs)
                samples = torch.rand(bs, 1, 28, 28, device=self.config.device)
                all_samples = self.anneal_Langevin_dynamics_GenerateImages(samples, score, sigmas, 100, 0.00002)
                images = all_samples.mul_(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 1).to('cpu').numpy()
                for j in range(num_of_step-1):
                    samples = torch.rand(bs, 3, 32, 32, device=self.config.device)
                    all_samples = self.anneal_Langevin_dynamics_GenerateImages(samples, score, sigmas, 100, 0.00002)
                    images_new = all_samples.mul_(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 1).to('cpu').numpy()
                    images = np.concatenate((images, images_new), axis=0)

            else:
                print("Begin epochs", epochs)
                samples = torch.rand(bs, 3, 32, 32, device=self.config.device)
                all_samples = self.anneal_Langevin_dynamics_GenerateImages(samples, score, sigmas, 100, 0.00002)
                images = all_samples.mul_(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 1).to('cpu').numpy()
                for j in range(num_of_step-1):
                    samples = torch.rand(bs, 3, 32, 32, device=self.config.device)
                    all_samples = self.anneal_Langevin_dynamics_GenerateImages(samples, score, sigmas, 100, 0.00002)
                    images_new = all_samples.mul_(255).add_(0.5).clamp_(0, 255).permute(0, 2, 3, 1).to('cpu').numpy()
                    images = np.concatenate((images, images_new), axis=0)

            # load precalculated training set statistics
            f = np.load(stats_path)
            mu_real, sigma_real = f['mu'][:], f['sigma'][:]
            f.close()

            fid.create_inception_graph(inception_path)  # load the graph into the current TF graph
            with tf.Session() as sess:
                sess.run(tf.global_variables_initializer())
                mu_gen, sigma_gen = fid.calculate_activation_statistics(images, sess, batch_size=100)

            fid_value = fid.calculate_frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real)
            print("FID: %s" % fid_value)

max_stars_repo_name: taufikxu/FD-ScoreMatching

max_issues_repo_name: taufikxu/FD-ScoreMatching

max_forks_repo_name: taufikxu/FD-ScoreMatching


text:
import numpy as np

def rmse(a, b):
    """
    A function which calculates the root mean squared error
    between two vectors.

    Notes
    ---------
    .. math::

        RMSE = \\sqrt{\\frac{1}{n}(a-b)^2}
    """
    assert isinstance(a, np.ndarray), \
        'This function expects a numpy array as input'
    assert isinstance(b, np.ndarray), \
        'This function expects a numpy array as input'
    assert len(a) == len(b), \
        'a and b should be the same length, but got %i, %i' % (len(a), len(b))
    return np.sqrt(np.mean(np.square(a-b)))

ext: py

max_stars_repo_name: nealde/Ampere

max_issues_repo_name: nealde/battery

max_forks_repo_name: nealde/battery

reason: import numpy


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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import time

from hyperdopamine.agents.dqn import dqn_agent
from hyperdopamine.agents.rainbow import rainbow_agent
from hyperdopamine.agents.bdqn import bdqn_agent
from hyperdopamine.agents.hgqn_r1 import hgqn_r1_agent
from hyperdopamine.interfaces import atari_lib
from hyperdopamine.interfaces import checkpointer
from hyperdopamine.interfaces import iteration_statistics
from hyperdopamine.interfaces import logger

import random
import numpy as np
import tensorflow as tf
import gin.tf

def load_gin_configs(gin_files, gin_bindings):
  gin.parse_config_files_and_bindings(gin_files,
                                      bindings=gin_bindings,
                                      skip_unknown=False)

@gin.configurable
def create_agent(sess, environment, agent_name=None, summary_writer=None,
                 debug_mode=False):
  assert agent_name is not None
  if not debug_mode:
    summary_writer = None
  if agent_name == 'dqn':
    return dqn_agent.DQNAgent(sess, num_actions=environment.action_space.n,
        num_sub_actions=[sub_action_space.n for sub_action_space in \
            environment.sub_action_space],
        summary_writer=summary_writer)
  elif agent_name == 'rainbow':
    return rainbow_agent.RainbowAgent(sess, num_actions=environment.action_space.n,
        num_sub_actions=[sub_action_space.n for sub_action_space in \
            environment.sub_action_space],
        summary_writer=summary_writer)
  elif agent_name == 'hgqn_r1':
    return hgqn_r1_agent.HGQNr1Agent(
        sess, num_sub_actions=[sub_action_space.n for sub_action_space in \
            environment.action_space],
        summary_writer=summary_writer)
  elif agent_name == 'bdqn':
    return bdqn_agent.BDQNAgent(
        sess, num_sub_actions=[sub_action_space.n for sub_action_space in \
            environment.action_space],
        summary_writer=summary_writer)
  else:
    raise ValueError('Unknown agent: {}'.format(agent_name))

@gin.configurable
def create_runner(base_dir, schedule='continuous_train_and_eval'):
  assert base_dir is not None
  # Continuously runs training and evaluation until max num_iterations is hit.
  if schedule == 'continuous_train_and_eval':
    return Runner(base_dir, create_agent)
  # Continuously runs training until max num_iterations is hit.
  elif schedule == 'continuous_train':
    return TrainRunner(base_dir, create_agent)
  else:
    raise ValueError('Unknown schedule: {}'.format(schedule))

@gin.configurable
class Runner(object):
  ''' Object that handles running train-and-evaluate experiments. '''

  def __init__(self,
               base_dir,
               create_agent_fn,
               create_environment_fn=atari_lib.create_atari_environment,
               agent_seed=None,
               checkpoint_file_prefix='ckpt',
               checkpoint_frequency=1,
               logging_file_prefix='log',
               log_every_n=1,
               num_iterations=200,
               training_steps=250000,
               evaluation_steps=125000,
               max_steps_per_episode=27000,
               render=False,
               reward_clipping=True):
    assert base_dir is not None
    self._checkpoint_frequency = checkpoint_frequency
    self._logging_file_prefix = logging_file_prefix
    self._log_every_n = log_every_n
    self._num_iterations = num_iterations
    self._training_steps = training_steps
    self._evaluation_steps = evaluation_steps
    self._max_steps_per_episode = max_steps_per_episode
    self._base_dir = base_dir
    self._create_directories()
    self._summary_writer = tf.summary.FileWriter(self._base_dir)
    self._render = render
    self._reward_clipping = reward_clipping
    self._environment = create_environment_fn()

    print('Seeding agent with seed', agent_seed)
    tf.set_random_seed(agent_seed)
    np.random.seed(agent_seed)
    random.seed(agent_seed)

    config = tf.ConfigProto(allow_soft_placement=True)
    self._sess = tf.Session('', config=config)
    self._agent = create_agent_fn(self._sess, self._environment,
                                  summary_writer=self._summary_writer)
    self._summary_writer.add_graph(graph=tf.get_default_graph())
    self._sess.run(tf.global_variables_initializer())
    self._initialize_checkpointer_and_maybe_resume(checkpoint_file_prefix)

  def _create_directories(self):
    self._checkpoint_dir = os.path.join(self._base_dir, 'checkpoints')
    self._logger = logger.Logger(os.path.join(self._base_dir, 'logs'))

  def _initialize_checkpointer_and_maybe_resume(self, checkpoint_file_prefix):
    self._checkpointer = checkpointer.Checkpointer(self._checkpoint_dir,
                                                   checkpoint_file_prefix,
                                                   self._checkpoint_frequency)
    self._start_iteration = 0
    latest_checkpoint_version = checkpointer.get_latest_checkpoint_number(
        self._checkpoint_dir)
    if latest_checkpoint_version >= 0:
      experiment_data = self._checkpointer.load_checkpoint(
          latest_checkpoint_version)
      if self._agent.unbundle(
          self._checkpoint_dir, latest_checkpoint_version, experiment_data):
        assert 'logs' in experiment_data
        assert 'current_iteration' in experiment_data
        self._logger.data = experiment_data['logs']
        self._start_iteration = experiment_data['current_iteration'] + 1
        tf.logging.info('Reloaded checkpoint and will start from iteration %d',
                        self._start_iteration)

  def _initialize_episode(self):
    initial_observation = self._environment.reset()
    return self._agent.begin_episode(initial_observation)

  def _run_one_step(self, action):
    observation, reward, is_terminal, info = self._environment.step(action)
    if self._render and self._agent.eval_mode:
      self._environment.render()
    return observation, reward, is_terminal

  def _end_episode(self, reward, is_timeout):
    self._agent.end_episode(reward, is_timeout)

  def _run_one_episode(self):
    step_number = 0
    total_reward = 0.
    action = self._initialize_episode()
    is_terminal = False

    # Keep interacting until we reach a terminal state or the time limit.
    while True:
      observation, reward, is_terminal = self._run_one_step(action)
      total_reward += reward
      step_number += 1

      if self._reward_clipping:
        reward = np.clip(reward, -1, 1)

      if self._environment.game_over:
        # Stop the run loop once we reach the true end of episode.
        is_timeout = False
        break
      elif step_number == self._max_steps_per_episode:
        # Stop the run loop once we reach the timeout of episode.
        is_timeout = True
        break
      elif is_terminal:
        # If we lose a life but the episode is not over, signal an artificial
        # end of episode to the agent. Only meaningful for Atari domains,
        # otherwise will not reach here.
        self._agent.end_episode(reward)
        action = self._agent.begin_episode(observation)
      else:
        action = self._agent.step(reward, observation)

    self._end_episode(reward, is_timeout)
    return step_number, total_reward

  def _run_one_phase(self, min_steps, statistics, run_mode_str):
    step_count = 0
    num_episodes = 0
    sum_returns = 0.
    while step_count < min_steps:
      episode_length, episode_return = self._run_one_episode()
      statistics.append({
          '{}_episode_lengths'.format(run_mode_str): episode_length,
          '{}_episode_returns'.format(run_mode_str): episode_return
      })
      step_count += episode_length
      sum_returns += episode_return
      num_episodes += 1
      sys.stdout.write('Steps executed: {} '.format(step_count) +
                       'Episode length: {} '.format(episode_length) +
                       'Return: {}\r'.format(episode_return))
      sys.stdout.flush()
    return step_count, sum_returns, num_episodes

  def _run_train_phase(self, statistics):
    self._agent.eval_mode = False
    start_time = time.time()
    number_steps, sum_returns, num_episodes = self._run_one_phase(
        self._training_steps, statistics, 'train')
    average_return = sum_returns / num_episodes if num_episodes > 0 else 0.0
    statistics.append({'train_average_return': average_return})
    time_delta = time.time() - start_time
    tf.logging.info('Average undiscounted return per training episode: %.2f',
                    average_return)
    tf.logging.info('Average training steps per second: %.2f',
                    number_steps / time_delta)
    return num_episodes, average_return

  def _run_eval_phase(self, statistics):
    self._agent.eval_mode = True
    _, sum_returns, num_episodes = self._run_one_phase(
        self._evaluation_steps, statistics, 'eval')
    average_return = sum_returns / num_episodes if num_episodes > 0 else 0.0
    tf.logging.info('Average undiscounted return per evaluation episode: %.2f',
                    average_return)
    statistics.append({'eval_average_return': average_return})
    return num_episodes, average_return

  def _run_one_iteration(self, iteration):
    statistics = iteration_statistics.IterationStatistics()
    tf.logging.info('Starting iteration %d', iteration)
    num_episodes_train, average_reward_train = self._run_train_phase(
        statistics)
    num_episodes_eval, average_reward_eval = self._run_eval_phase(
        statistics)
    self._save_tensorboard_summaries(iteration, num_episodes_train,
                                     average_reward_train, num_episodes_eval,
                                     average_reward_eval)
    return statistics.data_lists

  def _save_tensorboard_summaries(self, iteration,
                                  num_episodes_train,
                                  average_reward_train,
                                  num_episodes_eval,
                                  average_reward_eval):
    summary = tf.Summary(value=[
        tf.Summary.Value(tag='Train/NumEpisodes',
                         simple_value=num_episodes_train),
        tf.Summary.Value(tag='Train/AverageReturns',
                         simple_value=average_reward_train),
        tf.Summary.Value(tag='Eval/NumEpisodes',
                         simple_value=num_episodes_eval),
        tf.Summary.Value(tag='Eval/AverageReturns',
                         simple_value=average_reward_eval)
    ])
    self._summary_writer.add_summary(summary, iteration)

  def _log_experiment(self, iteration, statistics):
    self._logger['iteration_{:d}'.format(iteration)] = statistics
    if iteration % self._log_every_n == 0:
      self._logger.log_to_file(self._logging_file_prefix, iteration)

  def _checkpoint_experiment(self, iteration):
    experiment_data = self._agent.bundle_and_checkpoint(self._checkpoint_dir,
                                                        iteration)
    if experiment_data:
      experiment_data['current_iteration'] = iteration
      experiment_data['logs'] = self._logger.data
      self._checkpointer.save_checkpoint(iteration, experiment_data)

  def run_experiment(self):
    tf.logging.info('Beginning training...')
    if self._num_iterations <= self._start_iteration:
      tf.logging.warning('num_iterations (%d) < start_iteration(%d)',
                         self._num_iterations, self._start_iteration)
      return

    # @TODO Render PyBullet only during evaluation.
    if self._render and 'Bullet' in self._environment.spec.id:
        self._environment.render()

    for iteration in range(self._start_iteration, self._num_iterations):
      statistics = self._run_one_iteration(iteration)
      self._log_experiment(iteration, statistics)
      if iteration % self._checkpoint_frequency == 0:
        self._checkpoint_experiment(iteration)

@gin.configurable
class TrainRunner(Runner):
  ''' Object that handles running train-only experiments. '''

  def __init__(self, base_dir, create_agent_fn,
               create_environment_fn=atari_lib.create_atari_environment):
    tf.logging.info('Creating TrainRunner...')
    super(TrainRunner, self).__init__(base_dir, create_agent_fn,
                                      create_environment_fn)
    self._agent.eval_mode = False

  def _run_one_iteration(self, iteration):
    statistics = iteration_statistics.IterationStatistics()
    num_episodes_train, average_reward_train = self._run_train_phase(
        statistics)
    self._save_tensorboard_summaries(iteration, num_episodes_train,
                                     average_reward_train)
    return statistics.data_lists

  def _save_tensorboard_summaries(self, iteration, num_episodes,
                                  average_reward):
    summary = tf.Summary(value=[
        tf.Summary.Value(tag='Train/NumEpisodes', simple_value=num_episodes),
        tf.Summary.Value(
            tag='Train/AverageReturns', simple_value=average_reward),
    ])
    self._summary_writer.add_summary(summary, iteration)

max_stars_repo_name: atavakol/action-hypergraph-networks

max_issues_repo_name: atavakol/action-hypergraph-networks

max_forks_repo_name: atavakol/action-hypergraph-networks


text:
from tensorflow import gfile
from utils import logged_timer as timer
from gtp_wrapper import make_gtp_instance, MCTSPlayer
from sgf_wrapper import sgf_prop
from tqdm import tqdm
import pdb
import sgf
import tensorflow as tf
import pandas as pd
import coords
import matplotlib.pyplot as plt
import itertools
import numpy as np
import main
import shipname
import os
import rl_loop
import utils
import evaluation
import selfplay_mcts
import preprocessing
import dual_net
import go
import sgf_wrapper
"""
Used to plot the accuracy of the policy and value networks in
predicting professional game moves and results over the course
of training. Check FLAGS for default values for what models to
load and what sgf files to parse.

Usage:
python training_curve.py

Sample 3 positions from each game
python training_curve.py --num_positions=3

Only grab games after 2005 (default is 2000)
python training_curve.py --min_year=2005
"""
import sys
sys.path.insert(0, '.')

tf.app.flags.DEFINE_string("sgf_dir", "sgf/baduk_db/", "sgf database")

tf.app.flags.DEFINE_string("model_dir", "saved_models",
                           "Where the model files are saved")
tf.app.flags.DEFINE_string("plot_dir", "data", "Where to save the plots.")
tf.app.flags.DEFINE_integer("min_year", "2000",
                            "Only take sgf games with date >= min_year")
tf.app.flags.DEFINE_string("komi", "7.5",
                           "Only take sgf games with given komi")
tf.app.flags.DEFINE_integer("idx_start", 150,
                            "Only take models after given idx")
tf.app.flags.DEFINE_integer("num_positions", 1,
                            "How many positions from each game to sample from.")
tf.app.flags.DEFINE_integer("eval_every", 5,
                            "Eval every k models to generate the curve")

FLAGS = tf.app.flags.FLAGS

def get_model_paths(model_dir):
    '''Returns all model paths in the model_dir.'''
    all_models = gfile.Glob(os.path.join(model_dir, '*.meta'))
    model_filenames = [os.path.basename(m) for m in all_models]
    model_numbers_names = [
        (shipname.detect_model_num(m), shipname.detect_model_name(m))
        for m in model_filenames]
    model_names = sorted(model_numbers_names)
    return [os.path.join(model_dir, name[1]) for name in model_names]

def load_player(model_path):
    print("Loading weights from %s ... " % model_path)
    with timer("Loading weights from %s ... " % model_path):
        network = dual_net.DualNetwork(model_path)
        network.name = os.path.basename(model_path)
    player = MCTSPlayer(network, verbosity=2)
    return player

def restore_params(model_path, player):
    with player.network.sess.graph.as_default():
        player.network.initialize_weights(model_path)

def batch_run_many(player, positions, batch_size=100):
    """Used to avoid a memory oveflow issue when running the network
    on too many positions. TODO: This should be a member function of
    player.network?"""
    prob_list = []
    value_list = []
    for idx in range(0, len(positions), batch_size):
        probs, values = player.network.run_many(positions[idx:idx+batch_size])
        prob_list.append(probs)
        value_list.append(values)
    return np.concatenate(prob_list, axis=0), np.concatenate(value_list, axis=0)

def eval_player(player, positions, moves, results):
    probs, values = batch_run_many(player, positions)
    policy_moves = [coords.from_flat(c) for c in np.argmax(probs, axis=1)]
    top_move_agree = [moves[idx] == policy_moves[idx]
                      for idx in range(len(moves))]
    square_err = (values - results)**2/4
    return top_move_agree, square_err

def get_sgf_props(sgf_path):
    with open(sgf_path) as f:
        sgf_contents = f.read()
    collection = sgf.parse(sgf_contents)
    game = collection.children[0]
    props = game.root.properties
    return props

def parse_sgf(sgf_path):
    with open(sgf_path) as f:
        sgf_contents = f.read()

    collection = sgf.parse(sgf_contents)
    game = collection.children[0]
    props = game.root.properties
    assert int(sgf_prop(props.get('GM', ['1']))) == 1, "Not a Go SGF!"

    result = utils.parse_game_result(sgf_prop(props.get('RE')))

    positions, moves = zip(*[(p.position, p.next_move)
                             for p in sgf_wrapper.replay_sgf(sgf_contents)])
    return positions, moves, result, props

def check_year(props, year):
    if year is None:
        return True
    if props.get('DT') is None:
        return False

    try:
        # Most sgf files in this database have dates of the form
        # "2005-01-15", but there are some rare exceptions like
        # "Broadcasted on 2005-01-15.
        year_sgf = int(props.get('DT')[0][:4])
    except:
        return False
    # TODO: better to use datetime comparison here?
    return year_sgf >= year

def check_komi(props, komi_str):
    if komi_str is None:
        return True
    if props.get('KM') is None:
        return False
    return props.get('KM')[0] == komi_str

def find_and_filter_sgf_files(base_dir, min_year=None, komi=None):
    sgf_files = []
    count = 0
    print("Finding all sgf files in {} with year >= {} and komi = {}".format(
        base_dir, min_year, komi))
    for i, (dirpath, dirnames, filenames) in enumerate(tqdm(os.walk(base_dir))):
        for filename in filenames:
            count += 1
            if count % 5000 == 0:
                print("Parsed {}, Found {}".format(count, len(sgf_files)))
            if filename.endswith('.sgf'):
                path = os.sep.join([dirpath, filename])
                props = get_sgf_props(path)
                if check_year(props, min_year) and check_komi(props, komi):
                    sgf_files.append(path)
    print("Found {} sgf files matching filters".format(len(sgf_files)))
    return sgf_files

def sample_positions_from_games(sgf_files, num_positions=1):
    pos_data = []
    move_data = []
    result_data = []
    move_idxs = []

    fail_count = 0
    for i, path in enumerate(tqdm(sgf_files, desc="loading sgfs", unit="games")):
        try:
            positions, moves, result, props = parse_sgf(path)
        except KeyboardInterrupt:
            raise
        except:
            fail_count += 1
            continue

        # add entire game
        if num_positions == -1:
            pos_data.extend(positions)
            move_data.extend(moves)
            move_idxs.extend(range(len(positions)))
            result_data.extend([result for i in range(len(positions))])
        else:
            for idx in np.random.choice(len(positions), num_positions):
                pos_data.append(positions[idx])
                move_data.append(moves[idx])
                result_data.append(result)
                move_idxs.append(idx)
    print("Sampled {} positions, failed to parse {} files".format(
        len(pos_data), fail_count))
    return pos_data, move_data, result_data, move_idxs

def get_training_curve_data(model_dir, pos_data, move_data, result_data, idx_start=150, eval_every=10):
    model_paths = get_model_paths(model_dir)
    df = pd.DataFrame()
    player = None

    print("Evaluating models {}-{}, eval_every={}".format(idx_start,
                                                          len(model_paths), eval_every))
    for idx in tqdm(range(idx_start, len(model_paths), eval_every)):
        if player:
            restore_params(model_paths[idx], player)
        else:
            player = load_player(model_paths[idx])

        correct, squared_errors = eval_player(
            player=player, positions=pos_data, moves=move_data, results=result_data)

        avg_acc = np.mean(correct)
        avg_mse = np.mean(squared_errors)
        print("Model: {}, acc: {:.4f}, mse: {:.4f}".format(
            model_paths[idx], avg_acc, avg_mse))
        df = df.append({"num": idx, "acc": avg_acc,
                        "mse": avg_mse}, ignore_index=True)
    return df

def save_plots(data_dir, df):
    plt.plot(df["num"], df["acc"])
    plt.xlabel("Model idx")
    plt.ylabel("Accuracy")
    plt.title("Accuracy in Predicting Professional Moves")
    plot_path = os.sep.join([data_dir, "move_acc.pdf"])
    plt.savefig(plot_path)

    plt.figure()

    plt.plot(df["num"], df["mse"])
    plt.xlabel("Model idx")
    plt.ylabel("MSE/4")
    plt.title("MSE in predicting outcome")
    plot_path = os.sep.join([data_dir, "value_mse.pdf"])
    plt.savefig(plot_path)

def main(unusedargv):
    sgf_files = find_and_filter_sgf_files(
        FLAGS.sgf_dir, FLAGS.min_year, FLAGS.komi)
    pos_data, move_data, result_data, move_idxs = sample_positions_from_games(
        sgf_files=sgf_files, num_positions=FLAGS.num_positions)
    df = get_training_curve_data(FLAGS.model_dir, pos_data, move_data,
                                 result_data, FLAGS.idx_start, eval_every=FLAGS.eval_every)
    save_plots(FLAGS.plot_dir, df)

FLAGS = tf.app.flags.FLAGS

if __name__ == "__main__":
    tf.app.run(main)


text:
# MIT License
#
# Copyright (c) 2019 Tuomas Halvari, Juha Harviainen, Juha Mylläri, Antti Röyskö, Juuso Silvennoinen
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import numpy as np
from dpemu.nodes import Array
from dpemu.filters.time_series import Gap, SensorDrift

def test_seed_determines_result_for_gap_filter():
    a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])
    x_node = Array()
    x_node.addfilter(Gap("prob_break", "prob_recover", "missing_value"))
    params = {"prob_break": 0.1, "prob_recover": 0.1, "missing_value": 1337}
    out1 = x_node.generate_error(a, params, np.random.RandomState(seed=42))
    out2 = x_node.generate_error(a, params, np.random.RandomState(seed=42))
    assert np.array_equal(out1, out2)

def test_sensor_drift():
    drift = SensorDrift("magnitude")
    params = {"magnitude": 2}
    drift.set_params(params)
    y = np.full((100), 1)
    drift.apply(y, np.random.RandomState(), named_dims={})

    increases = np.arange(1, 101)

    assert len(y) == len(increases)
    for i, val in enumerate(y):
        assert val == increases[i]*2 + 1

def test_one_gap():
    gap = Gap("prob_break", "prob_recover", "missing")
    y = np.arange(10000.0)
    params = {"prob_break": 0.0, "prob_recover": 1, "missing": np.nan}
    gap.set_params(params)
    gap.apply(y, np.random.RandomState(), named_dims={})

    for _, val in enumerate(y):
        assert not np.isnan(val)

def test_two_gap():
    gap = Gap("prob_break", "prob_recover", "missing")
    params = {"prob_break": 1.0, "prob_recover": 0.0, "missing": np.nan}
    y = np.arange(10000.0)
    gap.set_params(params)
    gap.apply(y, np.random.RandomState(), named_dims={})

    for _, val in enumerate(y):
        assert np.isnan(val)

max_forks_repo_name: thalvari/dpEmu-AutoML


text:
#-----------IMPORTS-----------------------------------------------------------------------------
import pandas as pd
import numpy as np
import tensorflow as tf
import os
import matplotlib.pyplot as plt
import seaborn as sns
import PIL
from typing import List

# EfficientNet
from tensorflow.keras.applications import EfficientNetB4, ResNet50
from tensorflow.keras.applications.efficientnet import preprocess_input

# Data Augmentation
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# Model Layers
from tensorflow.keras import Model, Input
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv2D, MaxPool2D , Flatten, Dropout, BatchNormalization

# Compiling and Callbacks
from tensorflow.keras.optimizers import SGD,Adam
from tensorflow.keras.losses import BinaryCrossentropy
from tensorflow.keras.callbacks import ReduceLROnPlateau, ModelCheckpoint, EarlyStopping

# Sklearn
from sklearn.metrics import roc_auc_score
#-----------------------------------------------------------------------------------------------------
# Competition Directory
comp_dir="/kaggle/input/ranzcr-clip-catheter-line-classification/"

# Get Training Data Labels
df_train=pd.read_csv(comp_dir+"train.csv").sample(frac=1).reset_index(drop=True)

# Get Training/Testing Data Paths
test_files = os.listdir(comp_dir+"test")

df_test = pd.DataFrame({"StudyInstanceUID": test_files})

image_size = 512
batch_size = 16       #32 changed
num_epochs = 8        #12 changed
learn_rate = 1e-03
df_train.StudyInstanceUID += ".jpg"

# Train-Val = [:21000], [21000:], test_files
# Train-Val-Test (for tuning model) = [:18000], [18000:24000], [24000:]

#----------------EfficientNetB4------------------------------------------------------------------------------

# base_model = ResNet50(include_top=False,
#                       weights=None,
#                       input_shape=(image_size, image_size, 3))
# base_model.load_weights("../input/keras-pretrained-models/resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5", by_name=True)

# for layer in base_model.layers[:-6]:
#     layer.trainable = False

base_model = EfficientNetB4(include_top=False,
                      weights="imagenet",
                      input_shape=(image_size, image_size, 3))
# base_model.load_weights("../input/tfkeras-22-pretrained-and-vanilla-efficientnet/TF2.2_EfficientNetB4_NoTop_ImageNet.h5", by_name=True)
# base_model.trainable = False

for layer in base_model.layers[:-6]:
        if not isinstance(layer, BatchNormalization):
            layer.trainable = False

#----------------IMAGE DATA GENERATOR------------------------------------------------------------------
label_cols=df_train.columns.tolist()
label_cols.remove("StudyInstanceUID")
label_cols.remove("PatientID")
datagen=ImageDataGenerator(rescale=1./255.)
test_datagen=ImageDataGenerator(rescale=1./255.)

#----------------TEST GENERATOR------------------------------------------------------------------------
test_generator=test_datagen.flow_from_dataframe(
    dataframe=df_train[24000:],          # changed from df_test
    directory=comp_dir+"train",          # change this
    x_col="StudyInstanceUID",
    batch_size=1,
    seed=42,
    shuffle=False,
    color_mode="rgb",
    class_mode=None,
    target_size=(image_size,image_size),
    interpolation="bilinear")
#-----------------------------------------------------------------------------------------------------

#----------------GENERATOR WRAPPER--------------------------------------------------------------------

def generator_wrapper(generator, start, end):
    for batch_x,batch_y in generator:
        yield (batch_x,[batch_y[:,i] for i in range(start, end)])
#-----------------------------------------------------------------------------------------------------

#---------# CVC DF GENERATORS (TRAIN, VALID, TEST [UNIVERSAL/NON-UNIQUE]) #---------#

## DATFRAME FILTER
CVC_df_train = df_train[["StudyInstanceUID", "CVC - Abnormal", "CVC - Borderline", "CVC - Normal"]]

# CVC GENERATORS
CVC_dftrain_generator=datagen.flow_from_dataframe(
    dataframe=CVC_df_train[:18000],
    directory=comp_dir+"train",
    x_col="StudyInstanceUID",
    y_col=["CVC - Abnormal", "CVC - Borderline", "CVC - Normal"],
    batch_size=batch_size,
    seed=42,
    shuffle=True,
    color_mode="rgb",
    class_mode="raw",
    target_size=(image_size,image_size),
    interpolation="bilinear")

CVC_dfvalid_generator=test_datagen.flow_from_dataframe(
    dataframe=CVC_df_train[18000:24000],
    directory=comp_dir+"train",
    x_col="StudyInstanceUID",
    y_col=["CVC - Abnormal", "CVC - Borderline", "CVC - Normal"],
    batch_size=batch_size,
    seed=42,
    shuffle=True,
    color_mode="rgb",
    class_mode="raw",
    target_size=(image_size,image_size),
    interpolation="bilinear")

#---------# CVC CATHETER (ABNORMAL, BORDERLINE, NORMAL) #---------#

inp = Input(shape = (image_size,image_size,3))
x = base_model(inp)
x = Flatten()(x)

# OUTPUT FUNNELLING TO CVC_model
# “binary_crossentropy” as loss function and “sigmoid” as the final layer activation
output4 = Dense(1, activation = 'sigmoid')(x)
output5 = Dense(1, activation = 'sigmoid')(x)
output6 = Dense(1, activation = 'sigmoid')(x)

# CVC_MODEL
CVC_model = Model(inp,[output4,output5,output6])

# STOCHASTIC GRADIENT DESCENT
sgd = SGD(lr=learn_rate, momentum=.9, nesterov=False)

CVC_model.compile(optimizer=sgd,
              loss = ["binary_crossentropy" for i in range(3)],
              metrics = ["accuracy"])

STEP_SIZE_TRAIN_CVC = CVC_dftrain_generator.n//CVC_dfvalid_generator.batch_size
STEP_SIZE_VALID_CVC = CVC_dfvalid_generator.n//CVC_dfvalid_generator.batch_size
STEP_SIZE_TEST_CVC = test_generator.n//test_generator.batch_size

CVC_history = CVC_model.fit(generator_wrapper(CVC_dftrain_generator, 0, 3),
                    steps_per_epoch=STEP_SIZE_TRAIN_CVC,
                    validation_data=generator_wrapper(CVC_dfvalid_generator, 0, 3),
                    validation_steps=STEP_SIZE_VALID_CVC,
                    epochs=num_epochs,verbose=2)

test_generator.reset()
CVC_pred = CVC_model.predict_generator(test_generator,
                             steps=STEP_SIZE_TEST_CVC,
                             verbose=1)

#==========================FINAL SUBMISSION============================#

# Get AUC Score    # TEST THIS

#ETT_model.save     save weights- debug to load file with weights (google)
# y_pred = np.transpose(np.squeeze(pred))

# Get AUC Score
# DOUBLE CHECK
y_pred = np.transpose(np.squeeze(CVC_pred))

y_true = df_train.loc[24000:, label_cols].to_numpy()
aucs = roc_auc_score(y_true, y_pred, average=None)  # predictions per row

print("AUC Scores: ", aucs)
print("Average AUC: ", np.mean(aucs))

# CATHETER COLUMNS
CVC_COLUMNS = ["CVC - Abnormal", "CVC - Borderline",
         "CVC - Normal"]

# Catheter Dataframes
CVC_df_submission = pd.DataFrame(np.squeeze(CVC_pred).transpose(), columns = CVC_COLUMNS)

# StudyInstanceUID DataFrame
SUID_df_submission = pd.DataFrame({"StudyInstanceUID": test_files})

# concatenated dataframes
catheter_df = [SUID_df_submission, CVC_df_submission]

# FINAL SUBMISSION DF

# puts IDs in first dataframe index
df_submission = pd.concat(catheter_df)

# SUBMISSION FILE
df_submission.to_csv("submission.csv", index=False)

#==========================CATHETER RESULTS GRAPHS============================#

# CVC GRAPH
epochs = range(1,num_epochs)
plt.plot(CVC_history.CVC_history['loss'], label='Training Set')
plt.plot(CVC_history.CVC_history['val_loss'], label='Validation Data)')
plt.title('CVC Training and Validation Loss')
plt.ylabel('MAE')
plt.xlabel('Num Epochs')
plt.legend(loc="upper left")
plt.show()
plt.savefig("CVC_loss.png")

max_stars_repo_name: YasPHP/RANZCR-CLiP

max_issues_repo_name: YasPHP/RANZCR-CLiP

max_forks_repo_name: YasPHP/RANZCR-CLiP


text:
#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
This module contains some useful hard-coded data and equations for calculations
in the other modules.
"""

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)
import numpy as np
import astropy.units as u
from warnings import warn

__all__ = ["hydrogen_cross_section", "helium_singlet_cross_section",
           "helium_triplet_cross_section", "he_3_properties"]

# Photoionization cross-section of hydrogen
def hydrogen_cross_section(wavelength=None, energy=None):
    """
    Compute the photoionization cross-section of hydrogen in function of
    wavelength or energy.

    Parameters
    ----------
    wavelength (``float`` or ``numpy.ndarray``, optional):
        Wavelength in unit of angstrom. Default is ``None``. If ``None``,
        ``energy`` cannot be ``None``.

    energy (``float`` or ``numpy.ndarray``, optional):
        Energy in unit of electron-volt. Default is ``None``. If ``None``,
        ``wavelength`` cannot be ``None``.

    Returns
    -------
    a_lambda (``float`` or ``numpy.ndarray``):
        Cross-section in function of wavelength and in unit of cm ** 2. Only
        returned if wavelength was input.

    a_nu (``float`` or ``numpy.ndarray``):
        Cross-section in function of energy and in unit of cm ** 2. Only
        returned if energy was input.
    """

    if wavelength is not None:
        epsilon = (911.65 / wavelength - 1) ** 0.5
        # Photoionization cross-section in function of wavelength
        a_lambda = \
            (6.3E-18 * np.exp(4 - (4 * np.arctan(epsilon)) / epsilon) /
             (1 - np.exp(-2 * np.pi / epsilon)) * (wavelength / 911.65)
             ** 4)
        if isinstance(a_lambda, np.ndarray):
            a_lambda[np.isnan(a_lambda)] = 0.0  # Remove NaNs if necessary
        return a_lambda

    elif energy is not None:
        epsilon = (energy / 13.6 - 1) ** 0.5
        # Photoionization cross-section in function of energy
        a_nu = (6.3E-18 * np.exp(4 - (4 * np.arctan(epsilon)) / epsilon) /
                (1 - np.exp(-2 * np.pi / epsilon)) * (13.6 / energy) ** 4)
        if isinstance(a_nu, np.ndarray):
            a_nu[np.isnan(a_nu)] = 0.0  # Remove NaNs if necessary
        return a_nu

    else:
        raise ValueError('Either the wavelength or energy has to be provided.')

# Total He photoionization cross-section
def helium_total_cross_section(wavelength):
    """
    Compute the total photoionization cross-section of helium in function of
    wavelength.

    Source: Yan, Sadeghpour, & Dalgarno (1998, ApJ)

    Parameters
    ----------
    wavelength (``numpy.ndarray``):
        Wavelength in unit of angstrom.

    Returns
    -------
    a_lambda_1 (``numpy.ndarray``):
        Cross-section in function of wavelength and in unit of cm ** 2.
    """
    nu_init = (wavelength * u.AA).to(u.Hz, equivalencies=u.spectral())
    threshold_energy = 24.58 * u.eV
    nu = nu_init.to(u.eV, equivalencies=u.spectral())
    x = nu / threshold_energy
    e_term = (nu / u.eV / 1000.) ** (7 / 2)

    coefs = [-4.7416, 14.8200, -30.8678, 37.3584, -23.4585, 5.9133]
    sum_term = 0.
    for i, coef in enumerate(coefs):
        sum_term += coef / x**((i+1)/2)

    num = 733 * u.barn
    out = num.to(u.cm * u.cm) / e_term * (1 + sum_term)
    out[nu < threshold_energy] = 0. * (u.cm ** 2)

    return out.to(u.cm ** 2).value

# Photoionization cross-section of helium singlet
def helium_singlet_cross_section(wavelength):
    """
    Compute the photoionization cross-section of helium singlet in function of
    wavelength.

    Parameters
    ----------
    wavelength (``float`` or ``numpy.ndarray``):
        Wavelength in unit of angstrom.

    Returns
    -------
    a_lambda_1 (``float`` or ``numpy.ndarray``):
        Cross-section in function of wavelength and in unit of cm ** 2.
    """
    energy = 12398.41984332 / wavelength  # Energy in unit of eV

    # The cross-sections for helium are partially hard-coded and based on those
    # of hydrogen
    a_lambda_h = hydrogen_cross_section(wavelength=wavelength)

    # For the singlet, we simply scale the cross-sections of H following the
    # results from Brown (1971; ADS:1971ApJ...164..387B)
    scale = np.ones_like(wavelength)
    scale *= 37.0 - 19.1 * (energy / 65.4) ** (-0.76)
    # Clip negative values of scale
    scale[scale < 0] = 0.0
    a_lambda_1 = a_lambda_h * scale

    return a_lambda_1

# Photoionization cross-section of helium triplet
def helium_triplet_cross_section():
    """
    Compute the photoionization cross-section of helium triplet in function of
    wavelength.

    Returns
    -------
    wavelength (``numpy.ndarray``):
        Wavelength in which the cross-section was sampled in Norcross (1971).

    a_lambda_3 (``numpy.ndarray``):
        Cross-section in function of wavelength and in unit of cm ** 2.
    """
    # The photoionization cross-section of He triplet is hard-coded with the
    # values calculated by Norcross 1971 (ADS:1971JPhB....4..652N). The
    # differential oscillator strength is calculated for bins of wavelength that
    # are not necessarily the same as the stellar spectrum wavelength bins.

    # Helium 2^3S differential oscillator strength
    data_array = np.array([[2593.01,     0.605],
                           [2528.27,     0.589],
                           [2275.74,     0.537],
                           [2023.15,     0.501],
                           [1655.63,     0.435],
                           [1214.41,     0.247],
                           [958.87,     0.1572],
                           [792.18,     0.1138],
                           [674.86,     0.0780],
                           [587.81,     0.0620],
                           [520.65,     0.0557],
                           [467.27,     0.0461],
                           [423.81,     0.0358],
                           [387.75,     0.0310],
                           [357.34,     0.0325],
                           [331.36,     0.0520],
                           [271.94,     0.343],
                           [271.21,     0.338],
                           [256.70,     0.274],
                           [243.01,     0.231],
                           [230.71,     0.200],
                           [219.59,     0.1750],
                           [209.49,     0.1537]
                           ])
    wavelength = np.flip(data_array[:, 0])
    differential_oscillator_strength = np.flip(data_array[:, 1])
    a_lambda_3 = 8.0670E-18 * differential_oscillator_strength
    return wavelength, a_lambda_3

# Collisional strengths for He
def he_collisional_strength():
    """
    Returns a hard-coded array containing the helium collisional strengths in
    function of temperature.

    Returns
    -------
    array (``numpy.ndarray``):
        Collisional strengths array. Columns: 0 = temperature, 1 = gamma_13,
        2 = gamma_31a, 3 = gamma_31b.
    """
    # log(T) [K]    gamma_13    gamma_31a   gamma_31b
    array = np.array([
        [3.75,            6.198E-2,    2.389,       7.965E-1],
        [4.00,            6.458E-2,    2.456,       9.579E-1],
        [4.25,            6.387E-2,    2.275,       1.042],
        [4.50,            6.157E-2,    1.916,       1.015],
        [4.75,            5.832E-2,    1.496,       8.950E-1],
        [5.00,            5.320E-2,    1.111,       7.265E-1],
        [5.25,            4.787E-2,    8.003E-1,    5.516E-1],
        [5.50,            4.018E-2,    5.660E-1,    3.948E-1],
        [5.75,            3.167E-2,    3.944E-1,    2.677E-1],
    ])
    return array

# Line properties of the 1.083 microns He triplet taken from the NIST database
# https://www.nist.gov/pml/atomic-spectra-database
def he_3_properties():
    """
    Returns the central wavelengths in air, oscillator strengths and the
    Einstein coefficient of the helium triplet in 1.083 microns. The values
    were taken from the NIST database:
    https://www.nist.gov/pml/atomic-spectra-database

    Returns
    -------
    lambda_0 (``float``):
        Central wavelength in air of line 0 in unit of m.

    lambda_1 (``float``):
        Central wavelength in air of line 1 in unit of m.

    lambda_2 (``float``):
        Central wavelength in air of line 2 in unit of m.

    f_0 (``float``):
        Oscillator strength of line 0 (unitless).

    f_1 (``float``):
        Oscillator strength of line 1 (unitless).

    f_2 (``float``):
        Oscillator strength of line 2 (unitless).

    a_ij (``float``):
        Einstein coefficient of the whole triplet in unit of 1 / s.
    """
    # Central wavelengths in units of m
    lambda_0 = 1.082909114 * 1E-6
    lambda_1 = 1.083025010 * 1E-6
    lambda_2 = 1.083033977 * 1E-6
    # Oscillator strengths
    f_0 = 5.9902e-02
    f_1 = 1.7974e-01
    f_2 = 2.9958e-01
    # Einstein coefficient in units of s ** (-1)
    a_ij = 1.0216e+07

    return lambda_0, lambda_1, lambda_2, f_0, f_1, f_2, a_ij

# Parametrized photoionization cross-section of atoms and ions from Verner+1996
def general_cross_section(energy, species):
    """
    Calculates the photoionization cross-section of an atomic or ion species
    using the parametrization of Verner et al. (1996)
    [https://ui.adsabs.harvard.edu/abs/1996ApJ...465..487V/abstract].

    Parameters
    ----------
    energy (``float`` or ``numpy.ndarray``)
        Photon energy in eV.

    species (``str``):
        String containing the species for which you request the cross-section in
        the format ``'X N'`` where ``X`` is the element and ``N`` is the
        ionization number. Example: ``'C II'``.

    Returns
    -------

    """
    parameters_dict = sigma_properties_v1996()
    energy_threshold, energy_max, energy_0, sigma_0, y_a, p, y_w, y_0, y_1 = \
        parameters_dict[species]

    # Check if the requested energy is within the energy range reported by
    # Verner et al. (1996)
    if isinstance(energy, float):
        if energy < energy_threshold or energy > energy_max:
            warn('The requested energy is outside the range of validity: '
                 '[{:.1f}, {:.1f}] eV'.format(energy_threshold, energy_max))
        else:
            pass
    elif isinstance(energy, np.ndarray):
        if energy[0] < energy_threshold or energy[-1] > energy_max:
            warn('The requested energy is outside the range of validity: '
                 '[{:.1f}, {:.1f}] eV'.format(energy_threshold, energy_max))
        else:
            pass
    else:
        raise ValueError('`energy` must be either `float` or `numpy.ndarray`.')

    mb = 1E-18  # cm ** (-2)
    x = energy / energy_0 - y_0
    y = (x ** 2 + y_1 ** 2) ** 0.5

    term1 = (x - 1) ** 2 + y_w ** 2
    term2 = y ** (0.5 * p - 5.5)
    term3 = (1 + (y / y_a) ** 0.5) ** (-p)
    function_y = term1 * term2 * term3

    cross_section = sigma_0 * function_y * mb  # cm ** (-2)

    return cross_section

def sigma_properties_v1996():
    """
    Function that hard-codes the cross-section parameters from Verner et al.
    (1996). We include only the species important for exoplanetary atmospheres.

    Returns
    -------
    parameters_dict (``dict``):
        Dictionary containing the parameters.
    """
    parameters_dict = {
        #         E_thresold, E_max, energy_0, sigma_0, y_a, p, y_w, y_0, y_1
        'C I':    [1.126E1, 2.910E2, 2.144E0, 5.027E2, 6.126E1, 5.101E0, 9.157E-2, 1.133E0, 1.607E0],
        'C II':   [2.438E1, 3.076E2, 4.058E-1, 8.709E0, 1.261E2, 8.578E0, 2.093E0, 4.929E1, 3.234E0],
        'C III':  [4.789E1, 3.289E2, 4.614E0, 1.539E4, 1.737E0, 1.593E1, 5.922E0, 4.378E-3, 2.528E-2],
        'N I':    [1.453E1, 4.048E2, 4.034E0, 8.235E2, 8.033E1, 3.928E0, 9.097E-2, 8.598E-1, 2.325E0],
        'O I':    [1.362E1, 5.380E2, 1.240E0, 1.745E3, 3.784E0, 1.764E1, 7.589E-2, 8.698E0, 1.271E-1],
        'O II':   [3.512E1, 5.581E2, 1.386E0, 5.967E1, 3.175E1, 8.943E0, 1.934E-2, 2.131E1, 1.503E-2],
        'O III':  [5.494E1, 5.840E2, 1.723E-1, 6.753E2, 3.852E2, 6.822E0, 1.191E-1, 3.839E-3, 4.569E-1],
        'Mg I':   [7.646E0, 5.490E1, 1.197E1, 1.372E8, 2.228E-1, 1.574E1, 2.805E-1, 0, 0],
        'Mg II':  [1.504E1, 6.569E1, 8.139E0, 3.278E0, 4.341E7, 3.610E0, 0, 0, 0],
        'Si I':   [8.152E0, 1.060E2, 2.317E1, 2.506E1, 2.057E1, 3.546E0, 2.837E-1, 1.672E-5, 4.207E-1],
        'Si II':  [1.635E1, 1.186E2, 2.556E0, 4.140E0, 1.337E1, 1.191E1, 1.570E0, 6.634E0, 1.272E-1],
        'Si III': [3.349E1, 1.311E2, 1.659E-1, 5.790E-4, 1.474E2, 1.336E1, 8.626E-1, 9.613E1, 6.442E-1],
        'Si IV':  [4.514E1, 1.466E2, 1.288E1, 6.083E0, 1.356E6, 3.353E0, 0, 0, 0],
        'Ca I':   [6.113E0, 3.443E1, 1.278E1, 5.370E5, 3.162E-1, 1.242E1, 4.477E-1, 1.012E-3, 1.851E-2],
        'Ca II':  [1.187E1, 4.090E1, 1.553E1, 1.064E7, 7.790E-1, 2.130E1, 6.453E-1, 2.161E-3, 6.706E-2],
        'Fe I':   [7.902E0, 6.600E1, 5.461E-2, 3.062E-1, 2.671E7, 7.923E0, 2.069E1, 1.382E2, 2.481E-1],
        'Fe II':  [1.619E1, 7.617E1, 1.761E-1, 4.365E3, 6.298E3, 5.204E0, 1.141E1, 9.272E1, 1.075E2],
    }
    return parameters_dict

max_stars_repo_name: astroshrey/p-winds

max_issues_repo_name: astroshrey/p-winds

max_forks_repo_name: astroshrey/p-winds

reason: import numpy,import astropy


text:
from dataclasses import dataclass
from multiprocessing import Queue
from typing import List, Optional, Dict

import numpy as np
from simpy import Environment, PriorityResource

from .simulation import Person
from .lab import Lab

@dataclass
class SimulationEnvironment:
    env: Environment
    sim_params: 'Parameters'
    duration: int
    sim_number: int
    scaling: float
    simulate_capacity: bool
    isolation_factor: float
    attention: PriorityResource
    hospital_bed: PriorityResource
    ventilator: PriorityResource
    icu: PriorityResource
    stats: np.ndarray
    street_expositions_interval: float
    social_group_expositions_interval: float
    creation_queue: Optional[Queue]
    simulation_queue: Optional[Queue]
    d0: Optional[float] = None
    populations: Dict[str, List[Person]] = None
    people: List[Person] = None
    lab: Optional[Lab] = None

max_stars_repo_name: victorhunguyen/covidsimulation


text:
import numpy as np

from keras.layers import Dense
from keras.models import Sequential
from keras.datasets import mnist
from keras.utils import to_categorical

from evolute import GeneticPopulation
from evolute.evaluation import SimpleFitness
from evolute.initialization import NormalRandom
from evolute.operators import RandomPickMate, Operators, UniformLocuswiseMutation, ScatterMateWrapper
from evolute.utility.keras_utility import WeightFolding

def pull_mnist():
    learning, testing = mnist.load_data()
    Xs, Ys = (learning[0], testing[0]), (learning[1], testing[1])
    Xs = map(lambda x: (x.reshape(-1, 784) - 127.5) / 255., Xs)
    Ys = map(lambda y: to_categorical(y, num_classes=10), Ys)
    return tuple(Xs), tuple(Ys)

def fitness_callback(phenotype, model: Sequential, w_folder, X, Y):
    model.set_weights(w_folder(phenotype))
    cost, acc = model.evaluate(X, Y, verbose=0)
    return cost

(lX, tX), (lY, tY) = pull_mnist()

ann = Sequential([
    Dense(64, activation="tanh", input_dim=lX.shape[1]),
    Dense(lY.shape[1], activation="softmax")
])
ann.compile(optimizer="sgd", loss="categorical_crossentropy", metrics=["acc"])

w_shapes = [w.shape for w in ann.get_weights()]
w_flat = np.concatenate([w.flat for w in ann.get_weights()])
w_folder = WeightFolding(ann)

fitness = SimpleFitness(fitness_function=fitness_callback,
                        constants={"model": ann, "w_folder": w_folder})

population = GeneticPopulation(
    limit=100,
    loci=w_flat.size,
    fitness_wrapper=fitness,
    initializer=NormalRandom(mean=w_flat),
    operators=Operators(mate_op=ScatterMateWrapper(RandomPickMate(), stdev=2.),
                        mutate_op=UniformLocuswiseMutation(low=-3., high=3.))
)

BATCH_SIZE = 128
batch_stream = ((lX[start:start+BATCH_SIZE], lY[start:start+BATCH_SIZE])
                for start in range(0, len(lX), BATCH_SIZE))

population.operators.selection.set_selection_rate(0.98)
population.operators.mutation.set_rate(0.0)
for i, (x, y) in enumerate(batch_stream, start=1):
    population.epoch(X=x, Y=y, verbosity=0)
    ann.set_weights(w_folder(population.get_best()))
    cost, acc = ann.evaluate(tX, tY, verbose=0)
    print("\rBatch: {} Acc: {:.2%}, Cost: {:.4f}".format(i, acc, cost))

max_stars_repo_name: ysglh/evolute

max_issues_repo_name: ysglh/evolute

max_forks_repo_name: ysglh/evolute


text:
import pandas as pd
import numpy as np

class Team():
    """Team"""

    def __init__(self, name: str, games: pd.DataFrame):
        self.name = name
        self.games_df = games.loc[
            (games.home_team == name) |
            (games.away_team == name)]

    def winning_percentage_for_season(self, season: int):
        """Gives percentage of games won"""

        games = self.games_df.loc[(self.games_df.season == season)]

        home_wins = len(games.loc[
            (games.home_team == self.name) &
            (games.home_score >
            games.away_score)].index)

        away_wins = len(games.loc[
            (games.away_team == self.name) &
            (games.away_score >
            games.home_score)].index)

        total = len(self.games_df.loc[(self.games_df.season == season)].index)

        if total == 0:
            return -1
        return (home_wins + away_wins)/total

    def get_list_of_opponents(self, season: int):

        games = self.games_df.loc[(self.games_df.season == season)]
        if len(games) == 0:
            return []
        print(self.name + " getting list of opponents")
        print(games[['home_team','away_team']])

        return games[['home_team','away_team']] \
            .apply(lambda  x: x[0] if x[1] == self.name else x[1], axis=1) \
            .tolist()

max_stars_repo_name: timvinc/nfl-stat-explorer

max_issues_repo_name: timvinc/nfl-stat-explorer

max_forks_repo_name: timvinc/nfl-stat-explorer


text:
import logging
import matplotlib.pyplot as plt
import matplotlib # color maps
import time
import numpy as np
import pandas as pd
from pathlib import Path
from itertools import chain
#from importlib import resources # py3.7+
from .progress_bar import * # context tqdm

LOGGER = logging.getLogger(__name__)
LOGGER.setLevel(logging.INFO)

pkg_namespace = 'methylize.data'
probe2chr = None
probe2map = None
def load_probe_chr_map():
    """__deprecated__ -- manhattan plot function will detect the array type and load this data when needed.
    runs inside manhattan plot."""
    global probe2chr
    global probe2map
    if probe2chr != None:
        return
    # maps probes to chromosomes for all known probes in major array types.
    probe2chr_path = Path(Path(__file__).parent,'data','probe2chr.csv')
    probe2map = pd.read_csv(probe2chr_path)
    #with resources.path(pkg_namespace, 'probe2chr.pkl') as probe2chr_filepath:
    #    with open(probe2chr_filepath,'rb') as f:
    #        probe2map = pickle.load(f)
        # structure is dataframe with 'CGidentifier, CHR, MAPINFO' columns -- bumphunter uses MAPINFO (chromosome position of probes)
    # dict for volcano plots, with a sortable order for chart
    probe2map['CHR'] = probe2map['CHR'].apply(lambda i: f"CHR-0{i}" if str(i).isdigit() and int(i) < 10 else f"CHR-{i}")
    probe2map = probe2map[ ~probe2map.CHR.isna() ] # drops control probes or any not linked to a chromosome
    probe2chr = dict(zip(probe2map.CGidentifier, probe2map.CHR)) # computationally fasted conversion method
    # OLD v0.9.3
    #probes = probe2map[['CGidentifier','CHR']].to_dict('records') #SLOW!!!
    #probe2chr = {probe['CGidentifier']:f"CHR-0{probe['CHR']}" if probe['CHR'] not in ('X','Y') and type(probe['CHR']) is str and int(probe['CHR']) < 10 else f"CHR-{probe['CHR']}" for probe in probes}
    #OLD pre v0.9: probe2chr = {k:f"CH-0{v}" if v not in ('X','Y') and type(v) is str and int(v) < 10 else f"CH-{v}" for k,v in probes.items()}

def create_mapinfo(manifest, genome_build=None, include_sex=True):
    """ convert a manifest.data_frame into a dictionary that maps probe_ids to chromosomes, for manhattan plots.
    - genome_build: use NEW or OLD to toggle which genome build to use from manifest.
    - only used by manhattan_plot()"""
    mapinfo = 'OLD_MAPINFO' if genome_build == 'OLD' else 'MAPINFO'
    chr = 'OLD_CHR' if genome_build == 'OLD' else 'CHR'
    sex_chromosomes = ['X','Y']
    probe2map = manifest.data_frame[[mapinfo, chr]]
    probe2map = probe2map[ ~probe2map[mapinfo].isna() ] # drops control probes or any not linked to a chromosome
    if include_sex:
        probe2map = probe2map[ (probe2map[chr].str.isdigit() | probe2map[chr].isin(sex_chromosomes)) ] #manifests have 'X,Y,M,*' omitting these.
    else:
        probe2map = probe2map[ probe2map[chr].str.isdigit() ] #manifests have 'X,Y,M,*' omitting these.
    probe2map[mapinfo] = probe2map[mapinfo].apply(lambda i: f"CHR-0{i}" if str(i).isdigit() and int(i) < 10 else f"CHR-{i}")
    #probe2chr = dict(zip(probe2map.index, probe2map[use])) # computationally fastest conversion method
    return probe2map # dataframe with CHR and MAPINFO columns and probe_names in index.

def create_probe_chr_map(manifest, genome_build=None, include_sex=True):
    """ convert a manifest.data_frame into a dictionary that maps probe_ids to chromosomes, for manhattan plots.
    - genome_build: use NEW or OLD to toggle which genome build to use from manifest.
    - only used by manhattan_plot()"""
    sex_chromosomes = ['X','Y']
    use = 'OLD_CHR' if genome_build == 'OLD' else 'CHR'

    probe2map = manifest.data_frame[[use]]
    probe2map = probe2map[ ~probe2map[use].isna() ] # drops control probes or any not linked to a chromosome
    #other_map = probe2map[ ~probe2map[use].str.isdigit() ]
    if include_sex:
        probe2map = probe2map[ (probe2map[use].str.isdigit() | probe2map[use].isin(sex_chromosomes)) ] #manifests have 'X,Y,M,*' omitting these.
    else:
        probe2map = probe2map[ probe2map[use].str.isdigit() ] #manifests have 'X,Y,M,*' omitting these.
    probe2map[use] = probe2map[use].apply(lambda i: f"CHR-0{i}" if str(i).isdigit() and int(i) < 10 else f"CHR-{i}")
    #probe2map['CHR'] = probe2map['CHR'].copy().apply(lambda i: f"CHR-0{i}" if i not in ('X','Y') and type(i) is str and int(i) < 10 else f"CHR-{i}")
    probe2chr = dict(zip(probe2map.index, probe2map[use])) # computationally fastest conversion method
    return probe2chr

color_schemes = {}
def load_color_schemes():
    global color_schemes
    if color_schemes != {}:
        return
    color_schemes = {cmap: matplotlib.cm.get_cmap(cmap) for cmap in
        ['Pastel1', 'Pastel2', 'Paired', 'Accent',
        'Dark2', 'Set1', 'Set2', 'Set3',
        'tab10', 'tab20', 'tab20b', 'tab20c']
    }
    color_schemes['Gray'] = matplotlib.colors.ListedColormap(['darkgrey','black'])
    color_schemes['Gray2'] = matplotlib.colors.ListedColormap(['darkgrey','gray'])
    color_schemes['Gray3'] = matplotlib.colors.ListedColormap(['whitesmoke','lightgray','silver','darkgray','gray','dimgray','black'])
    color_schemes['Gray4'] = matplotlib.colors.ListedColormap(['slategrey','silver'])
    color_schemes['Volcano'] = matplotlib.colors.ListedColormap(['tab:red','tab:blue','silver'])
    color_schemes['default'] = matplotlib.colors.ListedColormap(['mistyrose', 'navajowhite', 'palegoldenrod', 'yellowgreen', 'mediumseagreen', 'powderblue', 'skyblue',  'lavender', 'plum', 'palevioletred'])
load_color_schemes()

def to_genome(df, rgset): # pragma: no cover
    """__deprecated__ Maps dataframe to genome locations

Parameters:

    df: dataframe
            Dataframe containing methylation, unmethylation, M or Beta
            values for each sample at each site
    rgset: rg channel set instance
            RG channel set instance related to provided df
Returns:

    df: dataframe
            Dataframe containing the original values with the addition
            of genomic locations for each site
    """
    # This (mapToGenome) was in an old copy of methpype, but moved to methylize where it might be used.
    if 'Name' in df.columns:
        mani = rgset.manifest
        chromosomes = dict(
            zip(list(mani['Name'].values), list(mani['CHR'].values)))
        strands = dict(
            zip(list(mani['Name'].values), list(mani['Strand'].values)))
        build = dict(zip(list(mani['Name'].values),
                         list(mani['Genome_Build'].values)))
        mapinfo = dict(
            zip(list(mani['Name'].values), list(mani['MAPINFO'].values)))
        df['Chr'] = df['Name'].map(chromosomes)
        df['Strand'] = df['Name'].map(strands)
        df['GenomeBuild'] = df['Name'].map(build)
        df['MapInfo'] = df['Name'].map(mapinfo)
        return df
    else:
        print("No 'Name' column in dataframe.")
        return

def to_BED(stats, manifest_or_array_type, save=True, filename='', genome_build=None, columns=None):
    """Converts & exports manifest and probe p-value dataframe to BED format.
    - https://en.wikipedia.org/wiki/BED_(file_format)
    - BED format: [ chromosome number | start position | end position | p-values]
    Where p-values are the output from diff_meth_pos() comparing probes across two or more
    groups of samples for genomic differences in methylation.

    This output is required for combined-pvalues library to read and annotate manhattan plots
    with the nearest Gene(s) for each significant CpG cluster.

    manifest_or_array_type:
        either pass in a Manifest instance from methylprep, or a string that defines which
        manifest to load. One of {'27k', '450k', 'epic', 'epic+', 'mouse'}.
    genome_build:
        pass in 'OLD' to use the older genome build for each respective manifest array type.

    note: if manifest has probes that aren't mapped to genome, they are omitted in BED file.

    TODO: incorporate STRAND and OLD_STRAND in calculations.

    returns a BED formatted dataframe if save is False, or the saved filename if save is True.
    """
    array_types = {'27k', '450k', 'epic', 'epic+', 'mouse'}
    manifest = None
    if isinstance(manifest_or_array_type, str) and manifest_or_array_type not in array_types:
        raise ValueError(f"Specify array type as one of: {array_types}")
    if isinstance(manifest_or_array_type, str) and manifest_or_array_type in array_types:
        import methylprep
        manifest = methylprep.Manifest(methylprep.ArrayType(manifest_or_array_type))
    if not manifest and hasattr(manifest_or_array_type, 'data_frame'):
        manifest = manifest_or_array_type
    if not manifest:
        raise ValueError("Either provide a manifest or specify array_type.")
    if not isinstance(stats, pd.DataFrame):
        raise TypeError("stats should be a dataframe with either a PValue or a FDR_QValue column")
    if not isinstance(manifest.data_frame, pd.DataFrame):
        raise AttributeError("Expected manifest_or_array_type to be a methylprep manifest with a data_frame attribute but this does not have one.")
    if "FDR_QValue" in stats:
        pval = stats['FDR_QValue']
    elif "PValue" in stats:
        pval = stats['PValue']
    else:
        raise IndexError("stats did not contain either a PValue or a FDR_QValue column.")

    # an unfinished, internal undocumented way to change the column names, if exactly 5 columns in list provided in same order.
    if columns is None:
        columns = ['chrom','chromStart','chromEnd','pvalue','name']
        renamer = {}
    else:
        renamer = dict(zip(['chrom','chromStart','chromEnd','pvalue','name'],columns))

    pval = pval.rename("pvalue")
    genes = manifest_gene_map(manifest, genome_build=genome_build)
    # finally, inner join and save/return the combined BED data frame.
    BED = pd.merge(genes[['chrom','chromStart','chromEnd']], pval, left_index=True, right_index=True, how='inner')
    BED = BED.sort_values(['chrom','chromStart'], ascending=True)
    BED = BED.reset_index().rename(columns={'index':'name'})
    BED = BED[['chrom','chromStart','chromEnd','pvalue','name']] # order matters, so be explicit
    # omit unmapped CpGs
    unmapped = len(BED[ BED['chromStart'].isna() ])
    BED = BED[ ~BED['chromStart'].isna() ]
    if renamer != {}:
        BED = BED.rename(columns=renamer)
    # cpv / combined-pvalues needs a tab-separated .bed file
    timestamp = int(time.time())
    if save:
        if isinstance(filename, type(None)):
            BED.to_csv(f"{timestamp}.bed", index=False, sep='\t')
            return f"{timestamp}.bed"
        if not isinstance(filename, Path):
            filename = f"{filename}.bed"
        # otherwise, use as is, assuming it is a complete path/filename
        BED.to_csv(filename, index=False, sep='\t')
        return filename
    return BED

def manifest_gene_map(manifest, genome_build='NEW'):
    """returns 3 columns from manifest for chromosome/gene mapping. Used in >2 functions.
    genome_build: NEW or OLD"""
    ## sort probes by CHR,MAPINFO
    genome_prefix = 'OLD_' if genome_build == 'OLD' else '' # defaults to NEW, no prefix
    if genome_prefix+'MAPINFO' not in manifest.data_frame.columns:
        raise IndexError(f"{genome_prefix+'MAPINFO'} not in manifest")
    genes = manifest.data_frame[[genome_prefix+'CHR', genome_prefix+'MAPINFO']] # Strand and Genome_Build also available, but not needed.
    genes = genes.rename(columns={genome_prefix+'MAPINFO': 'chromStart', genome_prefix+'CHR': 'chrom'}) # standardize, regardless of genome build used
    genes = genes.astype({'chromStart':float})
    genes['chromEnd'] = genes['chromStart'] + 50 # all probes are 50 base pairs long. Ignoring that strand might affect the direction that each probe extends within genome for now. test then fix.
    return genes

max_stars_repo_name: FOXOBioScience/methylize

max_issues_repo_name: FOXOBioScience/methylize

max_forks_repo_name: FOXOBioScience/methylize


text:
import itertools

import numpy as np
from numpy.testing import assert_
import pytest

from qutip import *
from qutip.legacy.ptrace import _ptrace as _pt

@pytest.fixture(params=[True, False], ids=['sparse', 'dense'])
def sparse(request):
    return request.param

@pytest.fixture(params=[True, False], ids=['dm', 'ket'])
def dm(request):
    return request.param

@pytest.fixture
def state(dm):
    dims = [2, 3, 4]
    state = qutip.rand_ket(np.prod(dims), dims=[dims, [1]*len(dims)])
    if dm:
        state = state.proj()
    return state

def test_ptrace_noncompound_rand(sparse, dm):
    """Test `A.ptrace(0) == A` when `A` is in a non-tensored Hilbert space."""
    for _ in range(5):
        state = qutip.rand_ket(5)
        if dm:
            state = state.proj()
        assert state.ptrace(0, sparse=sparse) == state

@pytest.mark.parametrize('pair', list(itertools.combinations(range(3), 2)))
def test_ptrace_unsorted_selection_subset(state, sparse, pair):
    """
    Regression test for gh-1325.  ptrace should work the same independently of
    the order of the input; no transposition in done in the trace operation.
    """
    # pair is always sorted.
    state_ordered = state.ptrace(pair, sparse=sparse)
    state_reversed = state.ptrace(pair[::-1], sparse=sparse)
    assert state_ordered.dims == state_reversed.dims
    assert state_ordered == state_reversed

@pytest.mark.parametrize('permutation', list(itertools.permutations(range(3))))
def test_ptrace_unsorted_selection_all(state, sparse, permutation):
    state_ptraced = state.ptrace(permutation, sparse=sparse)
    assert state.dims == state_ptraced.dims
    assert state == state_ptraced

@pytest.mark.parametrize(['selection', 'exception'], [
    pytest.param(4, IndexError, id='too big'),
    pytest.param(-1, IndexError, id='too small'),
    pytest.param([0, 0], ValueError, id='duplicate'),
    # 'too many' may throw either from duplication or invalid index.
    pytest.param([0, 1, 2, 3], Exception, id='too many'),
])
def test_ptrace_fails_on_invalid_input(state, sparse, selection, exception):
    with pytest.raises(exception):
        state.ptrace(selection, sparse=sparse)

def test_ptrace_rand():
    'ptrace : randomized tests, sparse'
    for k in range(5):
        A = tensor(rand_ket(5), rand_ket(2), rand_ket(3))
        B = A.ptrace([1,2], True)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_dm(2), thermal_dm(10,1), rand_unitary(3))
        B = A.ptrace([1,2], True)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_ket(2),rand_ket(2),rand_ket(2),
                    rand_ket(2),rand_ket(2),rand_ket(2))
        B = A.ptrace([3,2], True)
        bdat,bd,bs = _pt(A, [3,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], True)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = rand_dm(64,0.5,dims=[[4,4,4],[4,4,4]])
        B = A.ptrace([0], True)
        bdat,bd,bs = _pt(A, [0])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([1], True)
        bdat,bd,bs = _pt(A, [1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], True)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

def test_ptrace_rand():
    'ptrace : randomized tests, dense'
    for k in range(5):
        A = tensor(rand_ket(5), rand_ket(2), rand_ket(3))
        B = A.ptrace([1,2], False)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_dm(2), thermal_dm(10,1), rand_unitary(3))
        B = A.ptrace([1,2], False)
        bdat,bd,bs = _pt(A, [1,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = tensor(rand_ket(2),rand_ket(2),rand_ket(2),
                    rand_ket(2),rand_ket(2),rand_ket(2))
        B = A.ptrace([3,2], False)
        bdat,bd,bs = _pt(A, [3,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,1], False)
        bdat,bd,bs = _pt(A, [0,1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

    for k in range(5):
        A = rand_dm(64,0.5,dims=[[4,4,4],[4,4,4]])
        B = A.ptrace([0], False)
        bdat,bd,bs = _pt(A, [0])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([1], False)
        bdat,bd,bs = _pt(A, [1])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

        B = A.ptrace([0,2], False)
        bdat,bd,bs = _pt(A, [0,2])
        C = Qobj(bdat,dims=bd)
        assert_(B==C)

max_stars_repo_name: camponogaraviera/qutip

max_issues_repo_name: camponogaraviera/qutip

max_forks_repo_name: camponogaraviera/qutip

reason: import numpy,from numpy