text:
import argparse
import math
import random
import os
import cv2
from functools import partial

import numpy as np
import torch
from torch import nn, autograd, optim
from torch.nn import functional as F
from torch.utils import data
import torch.distributed as dist
from torchvision import transforms, utils
from tqdm import tqdm

try:
    import wandb
except ImportError:
    wandb = None

from model import PyramidEncoder, Generator, Discriminator, Cooccurv2Discriminator, MultiProjectors, PatchNCELoss
from dataset import CXR14maskDataset
from stylegan2.distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)

def get_N(W, H):
    """N that maps from unnormalized to normalized coordinates"""
    N = np.zeros((3, 3), dtype=np.float64)
    N[0, 0] = 2.0 / W
    N[0, 1] = 0
    N[1, 1] = 2.0 / H
    N[1, 0] = 0
    N[0, -1] = -1.0
    N[1, -1] = -1.0
    N[-1, -1] = 1.0
    return N

def get_N_inv(W, H):
    """N that maps from normalized to unnormalized coordinates"""
    # TODO: do this analytically maybe?
    N = get_N(W, H)
    return np.linalg.inv(N)

def cvt_MToTheta(M, w, h):
    """convert affine warp matrix `M` compatible with `opencv.warpAffine` to `theta` matrix
    compatible with `torch.F.affine_grid`
    Note:
    M works with `opencv.warpAffine`.
    To transform a set of bounding box corner points using `opencv.perspectiveTransform`, M^-1 is required
    Parameters
    ----------
    M : np.ndarray
        affine warp matrix shaped [2, 3]
    w : int
        width of image
    h : int
        height of image
    Returns
    -------
    np.ndarray
        theta tensor for `torch.F.affine_grid`, shaped [2, 3]
    """
    M_aug = np.concatenate([M, np.zeros((1, 3))], axis=0)
    M_aug[-1, -1] = 1.0
    N = get_N(w, h)
    N_inv = get_N_inv(w, h)
    theta = N @ M_aug @ N_inv
    theta = np.linalg.inv(theta)
    return theta[:2, :]

def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)

    if shuffle:
        return data.RandomSampler(dataset)

    else:
        return data.SequentialSampler(dataset)

def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag

def accumulate(model1, model2, decay=0.999):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())

    for k in par1.keys():
        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)

def sample_data(loader):
    while True:
        for batch in loader:
            yield batch

def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)

    return real_loss.mean() + fake_loss.mean()

def d_r1_loss(real_pred, real_img):
    grad_real_list = autograd.grad(
        outputs=real_pred.sum(), inputs=real_img, create_graph=True
    )
    grad_penalty = grad_real_list[0].pow(2).reshape(grad_real_list[0].shape[0], -1).sum(1).mean()
    if len(grad_real_list) > 1:
        for grad_real in grad_real_list[1:]:
            grad_penalty += grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()

    return grad_penalty

def g_nonsaturating_loss(fake_pred):
    loss = F.softplus(-fake_pred).mean()

    return loss

def set_grad_none(model, targets):
    for n, p in model.named_parameters():
        if n in targets:
            p.grad = None

def differential_crop_w_random_rotation(image, max_angle):
    """
    max_angle defines the maximum rotation angles from [-max_angle, max_angle]
    """
    thetas = []
    N, D, H, W = image.shape
    for i in range(N):
        angle = random.uniform(-max_angle, max_angle)
        # angle = 60
        center = (W//2, H//2)
        scale = math.cos(math.pi*abs(angle)/180) + math.sin(math.pi*abs(angle)/180)*H/W
        # scale = 1.6
        # print(scale)
        affine_trans = cv2.getRotationMatrix2D(center, angle, scale)
        theta = cvt_MToTheta(affine_trans, W, H)
        thetas.append(torch.from_numpy(theta))
    thetas = torch.stack(thetas).to(image.device)
    grid = F.affine_grid(thetas, image.size(), align_corners=False)
    rotated = F.grid_sample(image, grid.float(), align_corners=False)

    return rotated

def raw_patchify_image(img, n_crop, mask=None, min_size=1 / 16, max_size=1 / 8, max_angle=60):
    batch, channel, height, width = img.shape

    def compute_corner(cent_x, cent_y, c_w, c_h):
        c_x = cent_x - c_w//2
        c_y = cent_y - c_h//2
        return c_x, c_y

    default_mask = torch.zeros(batch, height, width)
    default_mask[:, height//4:3*height//4, width//4:3*width//4] = 1

    if mask is None:
        mask = default_mask
    mask = mask.squeeze()

    crop_size = torch.rand(batch, n_crop) * (max_size - min_size) + min_size
    target_h = int(height * max_size)
    target_w = int(width * max_size)
    crop_h = (crop_size * height).type(torch.int64).tolist()
    crop_w = (crop_size * width).type(torch.int64).tolist()
    min_h = int(height * min_size)
    min_w = int(width * min_size)

    patches = []
    for b in range(batch):
        indices = torch.nonzero(mask[b]) # [height, width]
        if indices.size(0) == 0:
            indices = torch.nonzero(default_mask[b])
        for _ in range(n_crop):
            # random sample origin
            ind = random.randrange(0, indices.size(0))
            cent_y, cent_x = indices[ind].tolist()
            # random sample crop size
            crop_ratio = random.uniform(0,1) * (max_size - min_size) + min_size
            c_h, c_w = int(height * crop_ratio), int(width * crop_ratio)
            # recompute corners and area
            c_x, c_y = compute_corner(cent_x, cent_y, c_w, c_h)
            # clip the coordinates
            if c_y < 0:
                c_y = 0
            if c_x < 0:
                c_x = 0
            if c_y + c_h >= height:
                c_y = height - c_h - 1
            if c_x + c_w >= width:
                c_x = width - c_w - 1

            init_patch = img[b, :, c_y : c_y + c_h, c_x : c_x + c_w].view(1, channel, c_h, c_w)
            intp_patch = F.interpolate(init_patch, size=(target_h, target_w), mode="bilinear", align_corners=False)
            patches.append(intp_patch)
    # patches shape: [batch*n_crop, channel, target_h, target_w]
    patches = torch.cat(patches, dim=0)
    rotated = differential_crop_w_random_rotation(patches, max_angle)

    return rotated

def sample_patches(feat_list, n_crop, mask=None, coords=None, inv=False):
    if inv and mask is not None:
        mask = 1 - mask
    if mask is not None:
        mask = mask.squeeze()

    # collect info
    batchSize = feat_list[0].size(0)
    channels = []
    spt_dims = []
    for feat in feat_list:
        assert(feat.size(0) == batchSize), "Batch size of features should be consistent"
        channels.append(feat.size(1))
        spt_dims.append(feat.shape[2:])

    if coords is None:
        # sample coords from mask
        batch_coords = []
        for b in range(batchSize):
            # height and width of mask
            h, w = mask[b].shape
            # get valid candidates from mask
            indices = torch.nonzero(mask[b])
            if indices.size(0) == 0:
                indices = torch.nonzero(torch.ones(h, w))
            # sample points
            normed_coords = []
            for _ in range(n_crop):
                ind = random.randrange(0, indices.size(0))
                cent_y, cent_x = indices[ind].tolist()
                normed_coords.append((cent_y / h, cent_x / w))
            batch_coords.append(normed_coords)
    else:
        batch_coords = coords

    # extract features according to sampled points
    scale_feats = []
    for i, feat in enumerate(feat_list):
        h, w = spt_dims[i]
        sampled_feats = []
        for b in range(batchSize):
            for j in range(n_crop):
                cent_y, cent_x = min(h-1, int(batch_coords[b][j][0] * h)), min(w-1, int(batch_coords[b][j][1] * w))
                sampled_feats.append(feat[b, :, cent_y, cent_x])
        # [batchSize*n_crop, channel]
        sampled_feats = torch.stack(sampled_feats, dim=0)
        scale_feats.append(sampled_feats)

    return scale_feats, batch_coords

def train(
    args,
    loaders,
    encoder,
    generator,
    str_projectors,
    discriminator,
    cooccur,
    g_optim,
    d_optim,
    e_ema,
    g_ema,
    device,
):
    patchify_image = partial(raw_patchify_image, min_size=args.min_patch, max_size=args.max_patch)

    loader = sample_data(loaders[0])
    ts_loader = sample_data(loaders[1])

    patchnce_loss = PatchNCELoss(nce_T=0.07, batch=args.batch//2)

    pbar = range(args.iter)

    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)

    d_loss_val = 0
    r1_loss = torch.tensor(0.0, device=device)
    cooccur_r1_loss = torch.tensor(0.0, device=device)
    feat_recon_loss = torch.tensor(0.0, device=device)
    g_loss_val = 0
    loss_dict = {}

    # create output folders
    if args.proj_name != "":
        sample_dir = f"outputs/sample-{args.proj_name}"
        ckpt_dir = f"outputs/ckpt-{args.proj_name}"
    else:
        sample_dir = "outputs/sample"
        ckpt_dir = "outputs/ckpt"
    if not os.path.exists(sample_dir):
        os.makedirs(sample_dir)
    if not os.path.exists(ckpt_dir):
        os.makedirs(ckpt_dir)

    if args.distributed:
        e_module = encoder.module
        g_module = generator.module
        s_module = str_projectors.module
        d_module = discriminator.module
        c_module = cooccur.module

    else:
        e_module = encoder
        g_module = generator
        s_module = str_projectors
        d_module = discriminator
        c_module = cooccur

    accum = 0.5 ** (32 / (10 * 1000))

    for idx in pbar:
        i = idx + args.start_iter

        if i > args.iter:
            print("Done!")

            break

        real_img, real_msk, _ = next(loader)
        real_img = real_img.to(device)

        requires_grad(encoder, False)
        requires_grad(generator, False)
        requires_grad(str_projectors, False)
        requires_grad(discriminator, True)
        requires_grad(cooccur, True)

        real_img1, real_img2 = real_img.chunk(2, dim=0)
        real_msk1, real_msk2 = real_msk.chunk(2, dim=0)

        structure1, texture1 = encoder(real_img1, multi_tex=False)
        _, texture2 = encoder(real_img2, run_str=False, multi_tex=False)

        # image adversarial loss
        fake_img1 = generator(structure1, texture1)
        fake_img2 = generator(structure1, texture2)
        fake_pred = discriminator(torch.cat((fake_img1, fake_img2), 0))
        real_pred = discriminator(real_img)
        d_loss = d_logistic_loss(real_pred, fake_pred)

        # texture adversarial loss
        fake_patch = patchify_image(fake_img2, args.n_crop, mask=real_msk1)
        real_patch = patchify_image(real_img2, args.n_crop, mask=real_msk2)
        ref_patch = patchify_image(real_img2, args.ref_crop * args.n_crop, mask=real_msk2)
        fake_patch_pred, ref_input = cooccur(
            fake_patch, args.n_crop, reference=ref_patch, ref_batch=args.ref_crop
        )
        real_patch_pred, _ = cooccur(real_patch, args.n_crop, ref_input=ref_input)
        cooccur_loss = d_logistic_loss(real_patch_pred, fake_patch_pred)

        loss_dict["d"] = d_loss
        loss_dict["cooccur"] = cooccur_loss

        d_optim.zero_grad()
        (d_loss + cooccur_loss).backward()
        d_optim.step()

        d_regularize = i % args.d_reg_every == 0

        if d_regularize:
            real_img.requires_grad = True
            real_pred = discriminator(real_img)
            r1_loss = d_r1_loss(real_pred, real_img)

            real_patch.requires_grad = True
            real_patch_pred, _ = cooccur(real_patch, args.n_crop, reference=ref_patch, ref_batch=args.ref_crop)
            cooccur_r1_loss = d_r1_loss(real_patch_pred, real_patch)

            d_optim.zero_grad()

            r1_loss_sum = args.r1 / 2 * r1_loss * args.d_reg_every
            r1_loss_sum += args.cooccur_r1 / 2 * cooccur_r1_loss * args.d_reg_every
            r1_loss_sum += 0 * real_pred[0, 0] + 0 * real_patch_pred[0, 0]
            r1_loss_sum.backward()

            d_optim.step()
            # real_img.requires_grad = False

        loss_dict["r1"] = r1_loss
        loss_dict["cooccur_r1"] = cooccur_r1_loss

        requires_grad(encoder, True)
        requires_grad(generator, True)
        requires_grad(str_projectors, True)
        requires_grad(discriminator, False)
        requires_grad(cooccur, False)

        structure1_list, texture1 = encoder(real_img1, multi_str=True, multi_tex=False)
        _, texture2 = encoder(real_img2, run_str=False, multi_tex=False)
        structure1 = structure1_list[-1]

        fake_img1 = generator(structure1, texture1)
        fake_img2 = generator(structure1, texture2)

        # reconstruction loss
        recon_loss = F.l1_loss(fake_img1, real_img1.detach())

        # image adversarial loss
        fake_pred = discriminator(torch.cat((fake_img1, fake_img2), 0))
        g_loss = g_nonsaturating_loss(fake_pred)

        # texture adversarial loss
        fake_patch = patchify_image(fake_img2, args.n_crop, mask=real_msk1)
        ref_patch = patchify_image(real_img2, args.ref_crop * args.n_crop, mask=real_msk2)
        fake_patch_pred, _ = cooccur(fake_patch, args.n_crop, reference=ref_patch, ref_batch=args.ref_crop)
        g_cooccur_loss = g_nonsaturating_loss(fake_patch_pred)

        # Patch NCE loss
        # re-encode
        fake_structure1_list, fake_texture2 = encoder(fake_img2, multi_str=True, multi_tex=False)
        fake_patch_vectors, coords = sample_patches(fake_structure1_list[:-1], args.n_crop, mask=real_msk1, inv=True)
        real_patch_vectors, _ = sample_patches(structure1_list[:-1], args.n_crop, coords=coords)
        str_qs = str_projectors(fake_patch_vectors)
        str_ks = str_projectors(real_patch_vectors)
        ## compute loss
        g_pnce_loss = 0
        num_scales = len(str_qs)
        for str_q, str_k in zip(str_qs, str_ks):
            g_pnce_loss += patchnce_loss(str_q, str_k) / num_scales
            # g_pnce_loss += patchnce_loss(str_k, str_k)

        # feature reconstruction loss
        feat_recon_loss = 0.5 * F.mse_loss(fake_structure1_list[-1], structure1.detach()) + 0.5 * F.mse_loss(fake_texture2, texture2.detach())

        loss_dict["recon"] = recon_loss
        loss_dict["g"] = g_loss
        loss_dict["g_cooccur"] = g_cooccur_loss
        loss_dict["g_pnce"] = g_pnce_loss
        loss_dict["feat_recon"] = feat_recon_loss

        g_optim.zero_grad()
        (recon_loss + 0.5*g_loss + g_pnce_loss + g_cooccur_loss + 0.5*feat_recon_loss).backward() #  + 10*
        g_optim.step()

        accumulate(e_ema, e_module, accum)
        accumulate(g_ema, g_module, accum)

        loss_reduced = reduce_loss_dict(loss_dict)

        d_loss_val = loss_reduced["d"].mean().item()
        cooccur_val = loss_reduced["cooccur"].mean().item()
        recon_val = loss_reduced["recon"].mean().item()
        g_loss_val = loss_reduced["g"].mean().item()
        g_cooccur_val = loss_reduced["g_cooccur"].mean().item()
        g_pnce_val = loss_reduced["g_pnce"].mean().item()
        feat_recon_val = loss_reduced["feat_recon"].mean().item()
        r1_val = loss_reduced["r1"].mean().item()
        cooccur_r1_val = loss_reduced["cooccur_r1"].mean().item()

        if get_rank() == 0:
            pbar.set_description(
                (
                    f"d: {d_loss_val:.4f}; c: {cooccur_val:.4f} g: {g_loss_val:.4f}; g_cooccur: {g_cooccur_val:.4f}; "
                    f"recon: {recon_val:.4f}; feat_recon: {feat_recon_val:.4f}; pnce: {g_pnce_val:.4f}; "
                    f"r1: {r1_val:.4f}; r1_cooccur: {cooccur_r1_val:.4f}"
                )
            )

            if wandb and args.wandb and i % 10 == 0:
                wandb.log(
                    {
                        "Generator": g_loss_val,
                        "Discriminator": d_loss_val,
                        "Cooccur": cooccur_val,
                        "Recon": recon_val,
                        "Feat Recon": feat_recon_val,
                        "Generator Cooccur": g_cooccur_val,
                        "R1": r1_val,
                        "Cooccur R1": cooccur_r1_val,
                        "PNCE": g_pnce_val,
                    },
                    step=i,
                )

            if i % 200 == 0:
                with torch.no_grad():
                    # read test image
                    real_img, _, _ = next(ts_loader)
                    real_img = real_img.to(device)
                    real_img1, real_img2 = real_img.chunk(2, dim=0)

                    e_ema.eval()
                    g_ema.eval()

                    structure1, texture1 = e_ema(real_img1, multi_tex=False)
                    _, texture2 = e_ema(real_img2, run_str=False, multi_tex=False)

                    fake_img1 = g_ema(structure1, texture1)
                    fake_img2 = g_ema(structure1, texture2)

                    sample = torch.cat((real_img1, fake_img1, fake_img2, real_img2), 0)

                    utils.save_image(
                        sample,
                        f"{sample_dir}/{str(i).zfill(6)}.png",
                        nrow=sample.shape[0] // 4,
                        normalize=True,
                        range=(-1, 1),
                    )

            if i % 10000 == 0:
                torch.save(
                    {
                        "e": e_module.state_dict(),
                        "g": g_module.state_dict(),
                        "s": s_module.state_dict(),
                        "d": d_module.state_dict(),
                        "cooccur": c_module.state_dict(),
                        "e_ema": e_ema.state_dict(),
                        "g_ema": g_ema.state_dict(),
                        "g_optim": g_optim.state_dict(),
                        "d_optim": d_optim.state_dict(),
                        "args": args,
                    },
                    f"{ckpt_dir}/{str(i).zfill(6)}.pt",
                )

if __name__ == "__main__":
    device = "cuda"
    torch.backends.cudnn.benchmark = True

    parser = argparse.ArgumentParser()

    parser.add_argument("path", type=str, nargs="+")
    parser.add_argument("--trlist", type=str)
    parser.add_argument("--tslist", type=str)
    parser.add_argument("--iter", type=int, default=800000)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--min_patch", type=float, default=1/16)
    parser.add_argument("--max_patch", type=float, default=1/4)
    parser.add_argument("--r1", type=float, default=10)
    parser.add_argument("--cooccur_r1", type=float, default=1)
    parser.add_argument("--ref_crop", type=int, default=4)
    parser.add_argument("--n_crop", type=int, default=8)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--ckpt", type=str, default=None)
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--channel", type=int, default=32)
    parser.add_argument("--channel_multiplier", type=int, default=1)
    parser.add_argument("--wandb", action="store_true")
    parser.add_argument("--local_rank", type=int, default=0)
    # a new param, weights of scales
    parser.add_argument("--weights", type=float, nargs="+")
    parser.add_argument("--proj_name", type=str, default="")

    args = parser.parse_args()

    n_gpu = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    args.distributed = n_gpu > 1

    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
        synchronize()

    args.latent = 512
    args.n_mlp = 8

    args.start_iter = 0

    encoder = PyramidEncoder(args.channel, gray=True).to(device)
    generator = Generator(args.channel, gray=True).to(device)
    str_projectors = MultiProjectors([args.channel, args.channel * 2, args.channel * 8], use_mlp=True).to(device)

    discriminator = Discriminator(
        args.size, channel_multiplier=args.channel_multiplier, gray=True
    ).to(device)
    cooccur = Cooccurv2Discriminator(args.channel, size=args.size*args.max_patch, gray=True).to(device)

    e_ema = PyramidEncoder(args.channel, gray=True).to(device)
    g_ema = Generator(args.channel, gray=True).to(device)
    e_ema.eval()
    g_ema.eval()
    accumulate(e_ema, encoder, 0)
    accumulate(g_ema, generator, 0)

    d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)

    g_optim = optim.Adam(
        list(encoder.parameters()) + list(generator.parameters()) + list(str_projectors.parameters()),
        lr=args.lr,
        betas=(0, 0.99),
    )
    d_optim = optim.Adam(
        list(discriminator.parameters()) + list(cooccur.parameters()),
        lr=args.lr * d_reg_ratio,
        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
    )

    if args.ckpt is not None:
        print("load model:", args.ckpt)

        ckpt = torch.load(args.ckpt, map_location=lambda storage, loc: storage)

        try:
            ckpt_name = os.path.basename(args.ckpt)
            args.start_iter = int(os.path.splitext(ckpt_name)[0])

        except ValueError:
            pass

        encoder.load_state_dict(ckpt["e"])
        generator.load_state_dict(ckpt["g"])
        str_projectors.load_state_dict(ckpt['s'])
        discriminator.load_state_dict(ckpt["d"])
        cooccur.load_state_dict(ckpt["cooccur"])
        e_ema.load_state_dict(ckpt["e_ema"])
        g_ema.load_state_dict(ckpt["g_ema"])

        g_optim.load_state_dict(ckpt["g_optim"])
        d_optim.load_state_dict(ckpt["d_optim"])

    if args.distributed:
        encoder = nn.parallel.DistributedDataParallel(
            encoder,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        generator = nn.parallel.DistributedDataParallel(
            generator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        str_projectors = nn.parallel.DistributedDataParallel(
            str_projectors,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        discriminator = nn.parallel.DistributedDataParallel(
            discriminator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

        cooccur = nn.parallel.DistributedDataParallel(
            cooccur,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
        )

    transform = transforms.Compose(
        [
            # transforms.RandomHorizontalFlip(),
            transforms.Resize(args.size),
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,), inplace=True),
        ]
    )

    dataset = CXR14maskDataset(args.path[0], args.path[1], args.trlist, transform, gray=True)

    loader = data.DataLoader(
        dataset,
        batch_size=args.batch,
        sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
        drop_last=True,
        num_workers=4,
    )

    # test dataset
    ts_dataset = CXR14maskDataset(args.path[0], args.path[1], args.tslist, transform, gray=True)
    ts_loader = data.DataLoader(
        ts_dataset,
        batch_size=args.batch,
        sampler=data_sampler(ts_dataset, shuffle=True, distributed=args.distributed),
        drop_last=True,
    )

    if get_rank() == 0 and wandb is not None and args.wandb:
        wandb.init(project=f"proj_{args.proj_name}")

    train(
        args,
        [loader, ts_loader],
        encoder,
        generator,
        str_projectors,
        discriminator,
        cooccur,
        g_optim,
        d_optim,
        e_ema,
        g_ema,
        device,
    )

max_stars_repo_name: cvlab-stonybrook/LSAE

max_issues_repo_name: cvlab-stonybrook/LSAE

max_forks_repo_name: cvlab-stonybrook/LSAE


text:
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from utils import flowlib
from utils.resize_utils import resize_tensor
from utils.image_utils import rgb2gray, warp_image, torchify
# from utils.torch_utils import forward_and_eval
plt.rcParams["axes.grid"] = False
EPS = 1e-5

################# Flow and images ##################
def get_batch_flow_images(flow_batch):
    res = np.stack([flowlib.flow_to_image(input.detach().cpu().numpy().transpose(1, 2, 0))
                    for input in flow_batch])
    return res

def show_velocity(flow, velocity):
    plt.subplot(2, 1, 1)
    plt.title('Flow')
    show_flow(flow, subplote=True)
    plt.subplot(2, 1, 2)
    plt.title('Velocity')
    show_image(velocity, subplote=True)
    plt.show()

def show_flow(flow, subplote=False, occ=None):
    if flow is None:
        print("Flow is None")
        return

    if occ is not None:
        occ = occ.transpose(1, 2, 0)

    # Take first image
    if len(flow.shape) == 4:
        flow = flow[0]

    # Move to numpy
    if type(flow) is torch.Tensor:
        flow = flow.detach().cpu().numpy()

    # Change to channel last
    if flow.shape[0] == 2:
        flow = flow.transpose(1, 2, 0)

    flowlib.visualize_flow(flow,
                           subplote=subplote,
                           occ=occ)

def show_image(im, subplote=False, overlay=False, ax=True):
    if type(im) is not np.ndarray:
        im = im.detach().cpu().numpy()

    if overlay:
        alpha = 0.5
    else:
        alpha = 1

    if im.max() > 1.:
        im = im / 255.

    if len(im.shape) == 4:  # Print first
        im = im[0]

    if len(im.shape) == 3:
        if im.shape[0] == 1:  # Grayscale
            plt.imshow(im[0], alpha=alpha)
        else:  # RGB
            if im.shape[2] > 3:  # 3'rd channel is W
                im = im.transpose(1, 2, 0)
            plt.imshow(im, alpha=alpha)
    else:
        plt.imshow(im, alpha=alpha)
    if not ax:
        plt.axis('off')

    if not subplote:
        plt.show()

def plot_occ_error(occ_prediction, gt_occ):
    plt.subplot(121)
    plt.title('Pred')
    show_image(occ_prediction, subplote=True)
    plt.subplot(122)
    plt.title('GT')
    show_image(gt_occ)

def plot_inputs(data):
    plt.subplot(331)
    plt.title('input_im1')
    show_image(data['input1'], subplote=True)
    plt.subplot(332)
    plt.title('input_im2')
    show_image(data['input2'], subplote=True)
    plt.subplot(333)
    plt.title('input_flow_gt')
    show_flow(data['target1'], subplote=True)
    plt.subplot(334)
    plt.title('input_occ_gt')
    show_image(data['target_occ1'], subplote=True)
    if 'invalid' in data:
        plt.subplot(335)
        plt.title('input_invalid_pixels')
        show_image(data['invalid'], subplote=True)
    plt.subplot(336)
    plt.title('input_occ_orig')
    show_image(data['target_occ1_orig'], subplote=True)
    plt.subplot(337)
    plt.title('input_im1_orig')
    show_image(data['input1_orig'], subplote=True)
    plt.subplot(338)
    plt.title('input_im2_orig')
    show_image(data['input2_orig'], subplote=True)
    plt.subplot(339)
    plt.title('input_flow_orig')
    show_flow(data['target1_orig'], subplote=True)

    plt.subplots_adjust(hspace=0.5)
    mng = plt.get_current_fig_manager()
    mng.resize(1200, 800)
    plt.show()

def plot_occ_inpainting(im1, im2, occ, gt_flow, res, occ_th=0.0):
    occ = torchify(occ).unsqueeze(0)
    occ_est = torch.Tensor(np.array(res['occs'][0], dtype=np.uint8)).cuda()
    flow_orig = torchify(res['flows'][0])
    flow_inp = torchify(res['interp_flows'][0])
    diff = (flow_inp - flow_orig).norm(2, keepdim=True, dim=1)
    occ = occ
    if occ.shape[-1] <= 3:
        occ = occ.transpose(2, 3).transpose(1, 2)
    occ = resize_tensor(occ, occ_est.shape[2:])[:, 0, :, :].unsqueeze(1) > 0
    occ_diff = (occ.float() - occ_est.float())

    plt.subplot(331)
    plt.title('Image 1')
    show_image(im1, subplote=True)

    plt.subplot(332)
    plt.title('Image 2')
    show_image(im2, subplote=True)

    plt.subplot(333)
    plt.title('GT flow')
    show_flow(gt_flow, subplote=True)

    plt.subplot(334)
    plt.title('Flow before refinement')
    show_flow(flow_orig, subplote=True)

    plt.subplot(335)
    plt.title('Flow after refinement')
    show_flow(flow_inp, subplote=True)

    plt.subplot(336)
    plt.title('Inpainting change magnitude')
    show_image(diff, subplote=True)

    plt.subplot(337)
    plt.title('Occlusions')
    show_image(occ, subplote=True)

    plt.subplot(338)
    plt.title('Occlusions est (th={})'.format(occ_th))
    show_image(occ_est, subplote=True)

    plt.subplot(339)
    plt.title('Occlusions diff')
    show_image(occ_diff)

def show_heatmaps(im, flow, prediction, grad_im):
    """

    :param im:
    :param flow:
    :param prediction:
    :param grad_im:
    :return:
    """
    plt.subplot(2, 2, 1)
    plt.title('Image')
    show_image(im, subplote=True)

    plt.subplot(2, 2, 2)
    plt.title('GT Flow')
    show_flow(flow, subplote=True)

    plt.subplot(2, 2, 3)
    plt.title('Flow error')
    show_flow(flow - prediction, subplote=True)

    plt.subplot(2, 2, 4)
    plt.title('Gradient')
    plt.imshow(rgb2gray(grad_im))

    plt.show()

def plot_summary(im1, im2, occ, flow, prediction, prediction_back, mepe, show_occ=False):
    """

    :param im1:
    :param im2:
    :param occ:
    :param flow:
    :param prediction:
    :return:
    """
    plt.subplot(2, 3, 1)
    plt.title('First image')
    show_image(im1, subplote=True)

    plt.subplot(2, 3, 2)
    plt.title('Second image')
    show_image(im2, subplote=True)

    plt.subplot(2, 3, 3)
    occ_percent = occ.sum() / (occ.shape[1] * occ.shape[2])
    if occ_percent < 1.:
        occ_percent *= 100.
    plt.title('Occlusions ({:.3}% of the image)'.format(occ_percent))
    show_image(occ, subplote=True)

    plt.subplot(2, 3, 4)
    plt.title('True flow')
    show_flow(flow, subplote=True)

    plt.subplot(2, 3, 5)
    plt.title('Predicted flow (MEPE: {:.3} px)'.format(float(mepe)))
    show_flow(prediction, subplote=True)

    plt.subplot(2, 3, 6)
    flow_diff = flow - prediction
    normed = np.sqrt(flow_diff[0, 0].__pow__(2) + flow_diff[0, 1].__pow__(2))
    plt.title('Flow error norm')
    show_image(normed, subplote=True)

    plt.show()

def plot_warped_images(warp_func, im1, im2, tprediction, flow, occ):
    occ = occ.transpose(1, 2, 0)
    flow = torchify(flow)
    warped_pred = warp_image(warp_func, im2, tprediction, im2compare=im1,
                             show_image=False)
    warped_true = warp_image(warp_func, im2, flow, im2compare=im1,
                             show_image=False)

    plt.subplot(2, 3, 1)
    plt.title('Predicted warping')
    plt.imshow(warped_pred)

    plt.subplot(2, 3, 2)
    plt.title('Ground Truth warping')
    plt.imshow(warped_true)

    plt.subplot(2, 3, 3)
    plt.title('Ground Truth warping masked with non occluded')
    plt.imshow(warped_true * (1 - occ))

    plt.subplot(2, 3, 4)
    plt.title('Predicted flow')
    show_flow(tprediction, subplote=True)

    plt.subplot(2, 3, 5)
    plt.title('GT flow')
    show_flow(flow, subplote=True)

    plt.subplot(2, 3, 6)
    plt.title('Flow error norm')
    flow_diff = flow - tprediction
    normed = torch.sqrt(flow_diff[0, 0].__pow__(2) + flow_diff[0, 1].__pow__(2))
    show_image(normed, subplote=True)

    plt.show()

def plot_occlusions_summary(im1, im2, occ, flow, filled_flow, edges):
    plt.subplot(2, 3, 1).set_axis_off()
    plt.title('First image')
    show_image(im1, subplote=True)
    plt.subplot(2, 3, 2).set_axis_off()
    plt.title('Second image')
    show_image(im2, subplote=True)
    plt.subplot(2, 3, 3).set_axis_off()
    occ_percent = 100 * (occ.sum() / (occ.shape[1] * occ.shape[2]))
    plt.title('GT Occlusions ({:.3}% of the image)'.format(occ_percent))
    show_image(occ, subplote=True)
    plt.subplot(2, 3, 4).set_axis_off()
    plt.title('True flow')
    show_flow(flow, subplote=True)
    plt.subplot(2, 3, 5).set_axis_off()
    plt.title('Occlusions masked true flow')
    show_flow(flow, subplote=True, occ=occ)
    plt.subplot(2, 3, 6).set_axis_off()
    plt.title('Occlusions inpainted flow')
    show_flow(filled_flow, subplote=True)

    plt.show()

def plot_refinement_results(out_interponet, pwc_flow, pwc_occ):
    res = out_interponet[0]
    show_flow(pwc_flow)
    if pwc_occ is not None:
        show_flow(pwc_flow * (1 - (pwc_occ / 255.)))
    show_flow(res)

def plot_refinement_iterations(res, refinements, gt_flow, args, rows=2):

    if gt_flow is not None:
        num_plots = args.refinement_passes + 2
    else:
        num_plots = args.refinement_passes + 1

    cols = np.ceil(num_plots / rows)
    plt.subplot(rows, cols, 1)
    plt.title('Backbone_output')
    show_flow(res['flows'], subplote=True)

    for ix, r in enumerate(refinements):
        offset = 2
        pos = ((ix + offset) % cols) + cols * ((ix + offset) // cols)
        plt.subplot(rows, cols, pos)
        plt.title('Refinement number {}'.format(ix))
        show_flow(r, subplote=True)

    if gt_flow is not None:
        offset = 3
        pos = ((ix + offset) % cols) + cols * ((ix + offset) // cols)
        plt.subplot(rows, cols, pos)
        plt.title('GT Flow')
        show_flow(gt_flow, subplote=True)

    plt.show()

def get_heatmaps(im1, im2, prediction, flow, tflow):
    raise Exception("Get heatmap is deprecated")
    # im1.requires_grad = True
    # predicted_flow, dur = forward_and_eval(im1, im2)
    # loss = torch.norm(predicted_flow - tflow[0], p=2, dim=1).mean()
    # loss = (predicted_flow - tflow[0]).pow(2).mean()
    # loss.backward()
    # grad_im = im1.grad.detach().cpu().numpy()
    # show_heatmaps(im2, flow, prediction, grad_im)
    # im1.requires_grad = False

def stitch_images(inputs, outputs, img_per_row=1):
    gap = 5
    height, width = inputs[0][0, :, :, 0].shape
    images = [*inputs, *outputs]
    columns = len(images)
    img = Image.new('RGB', (width * img_per_row * columns + gap * (img_per_row - 1),
                            height * (int(len(inputs) / 2*img_per_row))))

    for ix in range(0, 1):
        xoffset = int(ix % img_per_row) * width * columns + int(ix % img_per_row) * gap
        yoffset = int(ix / img_per_row) * height

        for cat in range(len(images)):
            cur_im = np.array(images[cat][ix])
            if cur_im.max() <= 1.:
                cur_im *= 255
            im = cur_im.astype(np.uint8).squeeze()
            im = Image.fromarray(im)
            img.paste(im, (xoffset + cat * width, yoffset))

    return img

################# GRAPHS ##################
def plot_disponorm(flow_norm, disp_norm, title_suffix='', xlabel='', ylabel='',
                   subplot=[1, 2, 1]):
    """

    :param flow_norm:
    :param disp_norm:
    :param title_suffix:
    :param xlabel:
    :param ylabel:
    :param subplot:
    :return:
    """

    # Flatten
    flow_norm_vector = flow_norm.flatten().astype(int)
    displacement_vector = disp_norm.flatten().astype(int)

    # Get mean displacement per flow
    unique_flows = np.array([i for i in range(0, np.amax(flow_norm_vector) + 1)])
    flows_count = np.histogram(flow_norm_vector, bins=np.array(
        [i for i in range(0, np.amax(flow_norm_vector) + 2)]))[0]

    # Plot
    binned_error = np.bincount(flow_norm_vector, weights=displacement_vector)
    mean_disp_per_flow = binned_error / (flows_count + EPS)
    plot_graph(x=unique_flows, y=binned_error,
               title=('Sum of all errors ' + title_suffix),
               xlabel=xlabel or 'GT speed norm (Pixel per frame)',
               ylabel='Sum of prediction error norms (Pixel per frame)',
               subplot=subplot)
    plot_graph(x=unique_flows, y=mean_disp_per_flow,
               title='Mean error ' + title_suffix,
               xlabel=xlabel or 'GT speed norm (Pixel per frame)',
               ylabel='Mean of prediction error norms (Pixel per frame)',
               subplot=[subplot[0], subplot[1], subplot[2] + 1])

    # normed_mean_disp_per_flow = mean_disp_per_flow / (unique_flows + EPS)
    # plot_graph(x=unique_flows, y=normed_mean_disp_per_flow,
    #            title='Normalized (by GT flow value) mean error ' + title_suffix,
    #            xlabel=xlabel or 'GT speed norm (Pixel per frame)',
    #            ylabel='Normalized (by GT flow value) mean of prediction error norms
    #            (Pixel per frame)')

def plot_graph(x, y, title, xlabel, ylabel, subplot=None, show=True):
    """

    :param x:
    :param y:
    :param title:
    :param xlabel:
    :param ylabel:
    :param subplot:
    :param show:
    :return:
    """
    if subplot is not None:
        plt.subplot(subplot[0], subplot[1], subplot[2])
    plt.title(title)
    plt.plot(x, y)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    if show and subplot is None or (not subplot[0] * subplot[1] != subplot[2]):
        plt.show()

def plot_graphs(metrics, disponorm=False, dispoangle=False, dispolap=False,
                mepe_max_flow=False, mepe_occ_perc=False):
    if disponorm:
        # Plot regular
        plot_disponorm(metrics['flow_norms'], metrics['displacements_norm'],
                       title_suffix='per GT speed over all pixels')

        # Plot norm by occluded
        occ_mask = (metrics['occ'].flatten() > 0).astype(int)
        speed_no_occ = metrics['flow_norms'].flatten()[occ_mask > 0]
        disp_err_no_occ = metrics['displacements_norm'].flatten()[occ_mask > 0]
        plot_disponorm(speed_no_occ, disp_err_no_occ,
                       title_suffix='per GT speed over all pixels (without occluded pixels)')

    if dispoangle:
        # Plot regular
        plot_disponorm(metrics['flow_angles'], metrics['displacements_norm'],
                       title_suffix='per GT angle over all pixels',
                       xlabel='GT flow angle', ylabel='Predicted speed error (norm)')

        # Plot norm by occluded
        occ_mask = (metrics['occ'].flatten() > 0).astype(int)
        speed_no_occ = metrics['flow_angles'].flatten()[occ_mask > 0]
        disp_err_no_occ = metrics['displacements_norm'].flatten()[occ_mask > 0]
        plot_disponorm(speed_no_occ, disp_err_no_occ,
                       title_suffix='per GT angle over all pixels (without occluded pixels)',
                       xlabel='GT flow angle', ylabel='Predicted speed error (norm)')

    if dispolap:
        pass

ext: py

max_stars_repo_name: Pavelrst/ScopeFlow

max_issues_repo_name: Pavelrst/ScopeFlow

max_forks_repo_name: Pavelrst/ScopeFlow

reason: import numpy


text:
import argparse
import torch
import json
from torch.autograd import Variable
import train
from PIL import Image
import numpy as np

def load_checkpoint(checkpoint):
    state = torch.load(checkpoint)

    arch = state['arch']
    lr = float(state['learning_rate'])
    hidden_units = int(state['hidden_units'])

    model, optimizer, criterion = \
        train.build_model(arch, hidden_units, lr)

    model.class_to_idx = state['class_to_idx']
    model.load_state_dict(state['state_dict'])
    optimizer.load_state_dict(state['optimizer'])
    return model

def process_image(image):
    size = 224
    width, height = image.size

    if height > width:
        height = int(max(height * size / width, 1))
        width = int(size)
    else:
        width = int(max(width * size / height, 1))
        height = int(size)

    # resized_image = image.resize((width, height))

    x0 = (width - size) / 2
    y0 = (height - size) / 2
    x1 = x0 + size
    y1 = y0 + size
    cropped_image = image.crop((x0, y0, x1, y1))
    np_image = np.array(cropped_image) / 255.
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    np_image_array = (np_image - mean) / std
    np_image_array = np_image.transpose((2, 0, 1))

    return np_image_array

def predict(input_path, model, use_gpu, results_to_show, top_k):
    model.eval()
    image = Image.open(input_path)
    np_array = process_image(image)
    tensor = torch.from_numpy(np_array)

    if use_gpu:
        var_inputs = Variable(tensor.float().cuda(), volatile=True)
    else:
        var_inputs = Variable(tensor, volatile=True).float()

    var_inputs = var_inputs.unsqueeze(0)
    output = model.forward(var_inputs)
    ps = torch.exp(output).data.topk(top_k)
    probs = ps[0].cpu() if use_gpu else ps[0]
    classes = ps[1].cpu() if use_gpu else ps[1]
    class_to_idx_inverted = {
        model.class_to_idx[k]: k for k in model.class_to_idx}
    classes_list = list()
    for label in classes.numpy()[0]:
        classes_list.append(class_to_idx_inverted[label])
    return probs.numpy()[0], classes_list

def get_command_line_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('input', type=str, help='Image file')

    parser.add_argument('checkpoint', type=str, help='Saved model checkpoint')

    parser.add_argument('--top_k', type=int, help='Return the top K most likely classes')

    parser.set_defaults(top_k=1)

    parser.add_argument('--category_names', type=str, help='File of category names')

    parser.add_argument('--gpu', dest='gpu', action='store_true', help='Use GPU')

    parser.set_defaults(gpu=False)

    return parser.parse_args()

def main():
    args = get_command_line_args()
    use_gpu = torch.cuda.is_available() and args.gpu
    print("Input file is: {}".format(args.input))
    print("Checkpoint file is: {}".format(args.checkpoint))
    if args.top_k:
        print("Returning {} most likely classes".format(args.top_k))
    if args.category_names:
        print("Category names file: {}".format(args.category_names))
    if use_gpu:
        print("Using GPU.")
    else:
        print("Using CPU.")

    model = load_checkpoint(args.checkpoint)
    print("Checkpoint loaded.")

    if use_gpu:
        model.cuda()

    if args.category_names:
        with open(args.category_names, 'r') as f:
            categories = json.load(f)
            print("Category names loaded")

    results_to_show = args.top_k if args.top_k else 1

    print("Processing image")
    probabilities, classes = predict(args.input, model, use_gpu, results_to_show, args.top_k)

    if results_to_show > 1:
        print("Top {} Classes for '{}':".format(len(classes), args.input))

        if args.category_names:
            print("{:<30} {}".format("Flower", "Probability"))
            print("------------------------------------------")
        else:
            print("{:<10} {}".format("Class", "Probability"))
            print("----------------------")

        for i in range(0, len(classes)):
            if args.category_names:
                print("{:<30} {:.2f}".format(categories[classes[i]], probabilities[i]))
            else:
                print("{:<10} {:.2f}".format(classes[i], probabilities[i]))
    else:
        print("The most likely class is '{}': probability: {:.2f}".format(categories[classes[0]] if args.category_names else classes[0], probabilities[0]))

if __name__ == "__main__":
    main()


text:
""" cbvcorrector.py module unit tests
"""

import pytest
from numpy.testing import (
    assert_almost_equal,
    assert_array_equal,
    assert_allclose,
    assert_raises,
)

import warnings
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import astropy.units as u
import pandas as pd
from astropy.table import Table
from astropy.time import Time

from lightkurve import TessLightCurve, KeplerLightCurve
from lightkurve import search_lightcurve
from lightkurve import LightkurveWarning
from lightkurve.correctors.designmatrix import DesignMatrix
from lightkurve.correctors.cbvcorrector import (
    load_kepler_cbvs,
    load_tess_cbvs,
    CotrendingBasisVectors,
    KeplerCotrendingBasisVectors,
    TessCotrendingBasisVectors,
)
from lightkurve.correctors.cbvcorrector import CBVCorrector
from .. import TESTDATA

# *******************************************************************************
# *******************************************************************************
# *******************************************************************************
# CotrendingBasisVectors unit tests
def test_CotrendingBasisVectors_nonretrieval():
    """Tests CotrendingBasisVectors class without requiring remote data"""

    # ***
    # Constructor
    # Create some generic CotrendingBasisVectors objects

    # Generic CotrendingBasisVectors object
    dataTbl = Table(
        [[1, 2, 3], [False, True, False], [2.0, 3.0, 4.0], [3.0, 4.0, 5.0]],
        names=("CADENCENO", "GAP", "VECTOR_1", "VECTOR_3"),
    )
    cbvTime = Time([443.51090033, 443.53133457, 443.55176891], format="bkjd")
    cbvs = CotrendingBasisVectors(data=dataTbl, time=cbvTime)
    assert isinstance(cbvs, CotrendingBasisVectors)
    assert cbvs.cbv_indices == [1, 3]
    assert np.all(cbvs.time.value == [443.51090033, 443.53133457, 443.55176891])

    # Auto-initiate 'GAP' and 'CADENCENO'
    dataTbl = Table([[2.0, 3.0, 4.0], [3.0, 4.0, 5.0]], names=("VECTOR_3", "VECTOR_12"))
    cbvTime = Time([443.51090033, 443.53133457, 443.55176891], format="bkjd")
    cbvs = CotrendingBasisVectors(data=dataTbl, time=cbvTime)
    assert isinstance(cbvs, CotrendingBasisVectors)
    assert cbvs.cbv_indices == [3, 12]
    assert np.all(cbvs.gap_indicators == [False, False, False])
    assert np.all(cbvs.cadenceno == [0, 1, 2])

    # ***
    # _to_designmatrix
    # Make sure CBVs are the columns in the returned 2-dim array
    dataTbl = Table(
        [
            [1, 2, 3],
            [False, True, False],
            [1.0, 2.0, 3.0],
            [4.0, 5.0, 6.0],
            [7.0, 8.0, 9.0],
        ],
        names=("CADENCENO", "GAP", "VECTOR_1", "VECTOR_2", "VECTOR_3"),
    )
    cbvTime = Time([1569.44053967, 1569.44192856, 1569.44331746], format="btjd")
    cbvs = CotrendingBasisVectors(dataTbl, cbvTime)
    cbv_dm_name = "test cbv set"
    # CBV index 5 does not exists and should be ingored
    cbv_designmatrix = cbvs.to_designmatrix(cbv_indices=[1, 3, 5], name=cbv_dm_name)
    assert cbv_designmatrix.shape == (3, 2)
    assert np.all(cbv_designmatrix["VECTOR_1"] == np.array([1.0, 2.0, 3.0]))
    assert np.all(cbv_designmatrix["VECTOR_3"] == np.array([7.0, 8.0, 9.0]))
    assert cbv_designmatrix.name == cbv_dm_name
    # CBV #2 was not requested, so make sure it is not present
    with pytest.raises(KeyError):
        cbv_designmatrix["VECTOR_2"]

    # ***
    # plot
    ax = cbvs.plot(cbv_indices=[1, 2], ax=None)
    assert isinstance(ax, matplotlib.axes.Axes)

    # There is no CBV # 5 so the third cbv_indices entry will be ignored
    ax = cbvs.plot(cbv_indices=[1, 2, 5], ax=ax)
    assert isinstance(ax, matplotlib.axes.Axes)

    # CBVs use 1-based indexing. Throw error if requesting CBV index 0
    with pytest.raises(ValueError):
        ax = cbvs.plot(cbv_indices=[0, 1, 2], ax=ax)

    # Only 'all' or specific CBV indices can be requested
    with pytest.raises(ValueError):
        ax = cbvs.plot("Doh!")

    # ***
    # align
    # Set up some cadenceno such that both CBV is trimmed and NaNs inserted
    sample_lc = TessLightCurve(
        time=[1, 2, 3, 4, 6, 7],
        flux=[1, 2, 3, 4, 6, 7],
        flux_err=[0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
        cadenceno=[1, 2, 3, 4, 6, 7],
    )
    dataTbl = Table(
        [
            [1, 2, 3, 5, 6],
            [False, True, False, False, False],
            [1.0, 2.0, 3.0, 5.0, 6.0],
        ],
        names=("CADENCENO", "GAP", "VECTOR_1"),
    )
    cbvTime = Time(
        [1569.43915078, 1569.44053967, 1569.44192856, 1569.44470635, 1569.44609524],
        format="btjd",
    )
    cbvs = CotrendingBasisVectors(dataTbl, cbvTime)
    cbvs = cbvs.align(sample_lc)
    assert np.all(sample_lc.cadenceno == cbvs.cadenceno)
    assert len(cbvs.cadenceno) == 6
    assert len(sample_lc.flux) == 6
    assert np.all(cbvs.gap_indicators.value[[1, 3, 5]])
    # Ignore the warning in to_designmatric due to a low rank matrix
    with warnings.catch_warnings():
        # Instantiating light curves with NaN times will yield a warning
        warnings.simplefilter("ignore", LightkurveWarning)
        cbv_designmatrix = cbvs.to_designmatrix(cbv_indices=[1])
    assert np.all(cbv_designmatrix["VECTOR_1"][[0, 1, 2, 4]] == [1.0, 2.0, 3.0, 6.0])
    assert np.all(np.isnan(cbv_designmatrix["VECTOR_1"][[3, 5]]))

    # ***
    # interpolate
    nLcCadences = 20
    xLc = np.linspace(0.0, 2 * np.pi, num=nLcCadences)
    sample_lc = TessLightCurve(
        time=xLc,
        flux=np.sin(xLc),
        flux_err=np.full(nLcCadences, 0.1),
        cadenceno=np.arange(nLcCadences),
    )
    nCbvCadences = 10
    xCbv = np.linspace(0.0, 2 * np.pi, num=nCbvCadences)
    dataTbl = Table(
        [
            np.arange(nCbvCadences),
            np.full(nCbvCadences, False),
            np.cos(xCbv),
            np.sin(xCbv + np.pi * 0.125),
        ],
        names=("CADENCENO", "GAP", "VECTOR_1", "VECTOR_2"),
    )
    cbvTime = Time(xCbv, format="btjd")
    cbvs = CotrendingBasisVectors(dataTbl, cbvTime)
    cbv_interpolated = cbvs.interpolate(sample_lc, extrapolate=False)
    assert np.all(cbv_interpolated.time.value == sample_lc.time.value)
    # Extrapolation test
    # If extrapolate=False then all outside values set to 0.0
    xCbv = np.linspace(0.0, 1.5 * np.pi, num=nCbvCadences)
    dataTbl = Table(
        [
            np.arange(nCbvCadences),
            np.full(nCbvCadences, False),
            np.cos(xCbv),
            np.sin(xCbv + np.pi * 0.125),
        ],
        names=("CADENCENO", "GAP", "VECTOR_1", "VECTOR_2"),
    )
    cbvTime = Time(xCbv, format="btjd")
    cbvs = CotrendingBasisVectors(dataTbl, cbvTime)
    cbv_interpolated = cbvs.interpolate(sample_lc, extrapolate=False)
    assert np.all(
            cbv_interpolated["VECTOR_1"].value[
                np.nonzero(cbv_interpolated.time.value > 1.5 * np.pi)[0]
            ] == 0.0
    )
    # extrapolate
    cbv_interpolated = cbvs.interpolate(sample_lc, extrapolate=True)
    assert np.all(
                cbv_interpolated["VECTOR_1"].value[
                    np.nonzero(cbv_interpolated.time.value > 1.5 * np.pi)[0]
                ] != 0.0
    )

@pytest.mark.remote_data
def test_cbv_retrieval():
    """Tests reading in some CBVs from MAST

    This indirectly tests the classes KeplerCotrendingBasisVectors and
    TessCotrendingBasisVectors

    """

    cbvs = load_tess_cbvs(sector=10, camera=2, ccd=4, cbv_type="SingleScale")
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot([1, 2, 4, 6, 8])
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "TESS"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.band is None
    assert cbvs.sector == 10
    assert cbvs.camera == 2
    assert cbvs.ccd == 4

    cbvs = load_tess_cbvs(sector=10, camera=2, ccd=4, cbv_type="MultiScale", band=2)
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.band == 2

    cbvs = load_tess_cbvs(sector=8, camera=3, ccd=1, cbv_type="Spike")
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)

    # No band specified for MultiScale, this should error
    with pytest.raises(AssertionError):
        cbvs = load_tess_cbvs(sector=10, camera=2, ccd=4, cbv_type="MultiScale")
    # Band specified for SingleScale, this should also error
    with pytest.raises(AssertionError):
        cbvs = load_tess_cbvs(
            sector=10, camera=2, ccd=4, cbv_type="SingleScale", band=2
        )
    # Improper CBV type request
    with pytest.raises(Exception):
        cbvs = load_tess_cbvs(
            sector=10, camera=2, ccd=4, cbv_type="SuperSingleScale"
        )

    cbvs = load_kepler_cbvs(mission="Kepler", quarter=8, module=16, output=4)
    assert isinstance(cbvs, KeplerCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "Kepler"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.quarter == 8
    assert cbvs.campaign is None
    assert cbvs.module == 16
    assert cbvs.output == 4

    cbvs = load_kepler_cbvs(mission="K2", campaign=15, channel=24)
    assert isinstance(cbvs, KeplerCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "K2"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.quarter is None
    assert cbvs.campaign == 15
    assert cbvs.module == 8
    assert cbvs.output == 4

def test_cbv_local():
    """Tests loading local CBVs as above

    This indirectly tests the classes KeplerCotrendingBasisVectors and
    TessCotrendingBasisVectors

    """
    cbv_dir = TESTDATA
    cbvs = load_tess_cbvs(cbv_dir=cbv_dir,sector=10, camera=2, ccd=4, cbv_type="SingleScale")
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot([1, 2, 4, 6, 8])
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "TESS"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.band is None
    assert cbvs.sector == 10
    assert cbvs.camera == 2
    assert cbvs.ccd == 4

    cbvs = load_tess_cbvs(cbv_dir=cbv_dir,sector=10, camera=2, ccd=4, cbv_type="MultiScale", band=2)
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.band == 2

    cbvs = load_tess_cbvs(cbv_dir=cbv_dir,sector=10, camera=2, ccd=4, cbv_type="Spike")
    assert isinstance(cbvs, TessCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)

    # No band specified for MultiScale, this should error
    with pytest.raises(AssertionError):
        cbvs = load_tess_cbvs(sector=10, camera=2, ccd=4, cbv_type="MultiScale")
    # Band specified for SingleScale, this should also error
    with pytest.raises(AssertionError):
        cbvs = load_tess_cbvs(
            sector=10, camera=2, ccd=4, cbv_type="SingleScale", band=2
        )
    # Improper CBV type request
    with pytest.raises(Exception):
        cbvs = load_tess_cbvs(
            sector=10, camera=2, ccd=4, cbv_type="SuperSingleScale"
        )

    cbvs = load_kepler_cbvs(cbv_dir=cbv_dir,mission="Kepler", quarter=8, module=16, output=4)
    assert isinstance(cbvs, KeplerCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "Kepler"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.quarter == 8
    assert cbvs.campaign is None
    assert cbvs.module == 16
    assert cbvs.output == 4

    cbvs = load_kepler_cbvs(cbv_dir=cbv_dir,mission="K2", campaign=15, channel=24)
    assert isinstance(cbvs, KeplerCotrendingBasisVectors)
    ax = cbvs.plot("all")
    assert isinstance(ax, matplotlib.axes.Axes)
    assert cbvs.mission == "K2"
    assert cbvs.cbv_type == "SingleScale"
    assert cbvs.quarter is None
    assert cbvs.campaign == 15
    assert cbvs.module == 8
    assert cbvs.output == 4

# *******************************************************************************
# *******************************************************************************
# *******************************************************************************
# CBVCorrector Unit Tests

def test_CBVCorrector():

    # Create a CBVCorrector without reading CBVs from MAST
    sample_lc = TessLightCurve(
        time=[1, 2, 3, 4, 5],
        flux=[1, 2, np.nan, 4, 5],
        flux_err=[0.1, 0.1, 0.1, 0.1, 0.1],
        cadenceno=[1, 2, 3, 4, 5],
        flux_unit=u.Unit("electron / second"),
    )

    cbvCorrector = CBVCorrector(sample_lc, do_not_load_cbvs=True)
    # Check that Nan was removed
    assert len(cbvCorrector.lc.flux) == 4
    # Check that the median flux value is preserved
    assert_allclose(
        np.nanmedian(cbvCorrector.lc.flux).value, np.nanmedian(sample_lc.flux).value
    )

    dm = DesignMatrix(pd.DataFrame({"a": np.ones(4), "b": [1, 2, 4, 5]}))

    # ***
    # RegressionCorrector.correct passthrough method
    lc = cbvCorrector.correct_regressioncorrector(dm)
    # Check that returned lc is in absolute flux units
    assert isinstance(lc, TessLightCurve)
    # The design matrix should have completely zeroed the flux around the median
    lc_median = np.nanmedian(lc.flux)
    assert_allclose(lc.flux, lc_median)

    # ***
    # Gaussian Prior fit
    lc = cbvCorrector.correct_gaussian_prior(
        cbv_type=None, cbv_indices=None, alpha=1e-9, ext_dm=dm
    )
    assert isinstance(lc, TessLightCurve)
    # Check that returned lc is in absolute flux units
    assert lc.flux.unit == u.Unit("electron / second")
    # The design matrix should have completely zeroed the flux around the median
    lc_median = np.nanmedian(lc.flux)
    assert_allclose(lc.flux, lc_median)
    ax = cbvCorrector.diagnose()
    assert len(ax) == 2 and isinstance(ax[0], matplotlib.axes._subplots.Axes)

    # Now add a strong regularization term and under-fit the data
    lc = cbvCorrector.correct_gaussian_prior(
        cbv_type=None, cbv_indices=None, alpha=1e9, ext_dm=dm
    )
    # There should be virtually no change in the flux
    assert_allclose(lc.flux, sample_lc.remove_nans().flux)

    # This should error because the dm has incorrect number of cadences
    dm_err = DesignMatrix(pd.DataFrame({"a": np.ones(5), "b": [1, 2, 4, 5, 6]}))
    with pytest.raises(ValueError):
        lc = cbvCorrector.correct_gaussian_prior(
            cbv_type=None, cbv_indices=None, alpha=1e-2, ext_dm=dm_err
        )

    # ***
    # ElasticNet fit
    lc = cbvCorrector.correct_elasticnet(
        cbv_type=None, cbv_indices=None, alpha=1e-20, l1_ratio=0.5, ext_dm=dm
    )
    assert isinstance(lc, TessLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")
    # The design matrix should have completely zeroed the flux around the median
    lc_median = np.nanmedian(lc.flux)
    assert_allclose(lc.flux, lc_median, rtol=1e-3)
    ax = cbvCorrector.diagnose()
    assert len(ax) == 2 and isinstance(ax[0], matplotlib.axes._subplots.Axes)
    # Now add a strong regularization term and under-fit the data
    lc = cbvCorrector.correct_elasticnet(
        cbv_type=None, cbv_indices=None, alpha=1e9, l1_ratio=0.5, ext_dm=dm
    )
    # There should be virtually no change in the flux
    assert_allclose(lc.flux, sample_lc.remove_nans().flux)

    # ***
    # Correction optimizer
    # The optimizer cannot be run without downloading targest from MAST for use
    # within the under-fitting metric.
    # So let's just verify it fails as expected (not much else we can do)
    dm_err = DesignMatrix(pd.DataFrame({"a": np.ones(5), "b": [1, 2, 4, 5, 6]}))
    with pytest.raises(ValueError):
        lc = cbvCorrector.correct(
            cbv_type=None,
            cbv_indices=None,
            alpha_bounds=[1e-4, 1e4],
            ext_dm=dm_err,
            target_over_score=0.5,
            target_under_score=0.8,
        )

@pytest.mark.remote_data
def test_CBVCorrector_retrieval():
    """Tests CBVCorrector by retrieving some sample Kepler/TESS light curves
    and correcting them
    """

    # ***
    # A good TESS example of both over- and under-fitting
    # The "over-fitted" curve looks better to the eye, but eyes can be deceiving!
    lc = search_lightcurve(
        "TIC 357126143", mission="tess", author="spoc", sector=10
    ).download(flux_column="sap_flux")
    cbvCorrector = CBVCorrector(lc)
    assert isinstance(cbvCorrector, CBVCorrector)

    cbv_type = ["SingleScale", "Spike"]
    cbv_indices = [np.arange(1, 9), "ALL"]

    # Gaussian Prior correction
    lc = cbvCorrector.correct_gaussian_prior(
        cbv_type=cbv_type, cbv_indices=cbv_indices, alpha=1e-2
    )
    assert isinstance(lc, TessLightCurve)
    # Check that returned lightcurve is in flux units
    assert lc.flux.unit == u.Unit("electron / second")
    ax = cbvCorrector.diagnose()
    assert len(ax) == 2 and isinstance(ax[0], matplotlib.axes._subplots.Axes)

    # ElasticNet corrections
    lc = cbvCorrector.correct_elasticnet(
        cbv_type=cbv_type, cbv_indices=cbv_indices, alpha=1e1, l1_ratio=0.5
    )
    assert isinstance(lc, TessLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")
    ax = cbvCorrector.diagnose()
    assert len(ax) == 2 and isinstance(ax[0], matplotlib.axes._subplots.Axes)

    # Correction optimizer
    lc = cbvCorrector.correct(
        cbv_type=cbv_type,
        cbv_indices=cbv_indices,
        alpha_bounds=[1e-4, 1e4],
        target_over_score=0.5,
        target_under_score=0.8,
    )
    assert isinstance(lc, TessLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")
    ax = cbvCorrector.diagnose()
    assert len(ax) == 2 and isinstance(ax[0], matplotlib.axes._subplots.Axes)

    # Goodness metric scan plot
    ax = cbvCorrector.goodness_metric_scan_plot(
        cbv_type=cbv_type, cbv_indices=cbv_indices
    )
    assert isinstance(ax, matplotlib.axes.Axes)

    # Try multi-scale basis vectors
    cbv_type = ["MultiScale.1", "MultiScale.2", "MultiScale.3"]
    cbv_indices = ["ALL", "ALL", "ALL"]
    lc = cbvCorrector.correct_gaussian_prior(
        cbv_type=cbv_type, cbv_indices=cbv_indices, alpha=1e-2
    )
    assert isinstance(lc, TessLightCurve)

    # ***
    # A Kepler and K2 example
    lc = search_lightcurve(
        "KIC 6508221", mission="kepler", author="kepler", quarter=5
    ).download(flux_column="sap_flux")
    cbvCorrector = CBVCorrector(lc)
    lc = cbvCorrector.correct_gaussian_prior(alpha=1.0)
    assert isinstance(lc, KeplerLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")

    lc = search_lightcurve("EPIC 247887989", mission="k2", author="k2").download(
        flux_column="sap_flux"
    )
    cbvCorrector = CBVCorrector(lc)
    lc = cbvCorrector.correct_gaussian_prior(alpha=1.0)
    assert isinstance(lc, KeplerLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")

    lc = cbvCorrector.correct()
    assert isinstance(lc, KeplerLightCurve)
    assert lc.flux.unit == u.Unit("electron / second")

    # ***
    # Try some expected failures

    # cbv_type and cbv_indices not the same list lengths
    with pytest.raises(AssertionError):
        lc = cbvCorrector.correct_gaussian_prior(
            cbv_type=["SingleScale", "Spike"], cbv_indices=["all"], alpha=1e-2
        )

    # cbv_type is not a list
    with pytest.raises(AssertionError):
        lc = cbvCorrector.correct_gaussian_prior(
            cbv_type="SingleScale", cbv_indices=["all"], alpha=1e-2
        )

    # cbv_indices is not a list
    with pytest.raises(AssertionError):
        lc = cbvCorrector.correct_gaussian_prior(
            cbv_type=["SingleScale"], cbv_indices="all", alpha=1e-2
        )

    #***
    # Test the need for extrapolation with Kepler data
    lc = search_lightcurve("KIC 2437317", mission="Kepler", author="kepler", cadence='long',
            quarter=6).download(flux_column="sap_flux")
    cbv_type = ['SingleScale']
    cbv_indices = [np.arange(1,9)]
    # This will generate an warning about the need for extrapolation
    cbvCorrector = CBVCorrector(lc, interpolate_cbvs=True, extrapolate_cbvs=False)
    # This will generate a light curve with the last value well below the median corrected flux
    cbvCorrector.correct_gaussian_prior(cbv_type=cbv_type, cbv_indices=cbv_indices, alpha=1e-4)
    assert (cbvCorrector.corrected_lc.flux[-1] - np.median(cbvCorrector.corrected_lc.flux)).value < -300
    # This will generate a light curve with the last value at about the median corrected flux
    cbvCorrector = CBVCorrector(lc, interpolate_cbvs=True, extrapolate_cbvs=True)
    cbvCorrector.correct_gaussian_prior(cbv_type=cbv_type, cbv_indices=cbv_indices, alpha=1e-4)
    assert ((cbvCorrector.corrected_lc.flux[-1] - np.median(cbvCorrector.corrected_lc.flux)).value > 0.0 and
            (cbvCorrector.corrected_lc.flux[-1] - np.median(cbvCorrector.corrected_lc.flux)).value < 20)

reason: import numpy,from numpy,import astropy,from astropy


text:
from copy import deepcopy
from math import ceil
from multiprocessing import Pool
from shutil import copyfile
from sys_simulator.a2c.parallel import env_step, env_step_test
from sys_simulator.a2c.agent import A2CAgent, A2CCentralAgent
from sys_simulator.a2c.framework import A2CDiscreteFramework
from sys_simulator.general.actions_discretizations import db_five, db_six, db_ten
from sys_simulator.dqn.externalDQNFramework import ExternalDQNFramework
from sys_simulator.dqn.agents.dqnAgent import ExternalDQNAgent
from sys_simulator.noises.decaying_gauss_noise import DecayingGaussNoise
from time import time

import numpy as np
import torch
from torch.utils.tensorboard import SummaryWriter

from sys_simulator.channels import BANChannel, UrbanMacroNLOSWinnerChannel
from sys_simulator.ddpg.parameters_noise import (
    AdaptiveParamNoiseSpec,
)
import sys_simulator.general as gen
from sys_simulator.general import (
    print_evaluating,
    print_stuff_ddpg,
    random_seed,
    save_with_pickle,
)
from sys_simulator.noises.ou_noise import OUNoise2, SysSimOUNoise
from sys_simulator.parameters.parameters import DQNAgentParameters, EnvironmentParameters
from sys_simulator.plots import plot_env_states, plot_trajectories
from sys_simulator.q_learning.environments.completeEnvironment12 import (
    CompleteEnvironment12,
)

# Uses CompleteEnvironment12
# Centralized Learning-Distributed Execution
# NN Parameters noise.
# Similar to script5.
# Trains a central agent for a fixed amount of devices.
# There are different channels to the BS and to the devices.
# There are multiple episodes, where the devices are distributed
# in different positions, and there are different motion models.
# Single episodes convergence. The states are in linear scale.
n_mues = 1  # number of mues
n_d2d = 2  # number of d2d pairs
n_rb = n_mues   # number of RBs
carrier_frequency = 2.4  # carrier frequency in GHz
bs_radius = 1000  # bs radius in m
rb_bandwidth = 180*1e3  # rb bandwidth in Hz
device_height = 1.5  # mobile devices height in m
bs_height = 25  # BS antenna height in m
p_max = 40  # max tx power in dBm
noise_power = -116  # noise power per RB in dBm
bs_gain = 17    # macro bs antenna gain in dBi
user_gain = 4   # user antenna gain in dBi
sinr_threshold_train = 6  # mue sinr threshold in dB for training
mue_margin = 6  # mue margin in dB
MIN_D2D_PAIR_DISTANCE = 1.5
MAX_D2D_PAIR_DISTANCE = 15
# conversions from dBm to dB
p_max = p_max - 30
noise_power = noise_power - 30
p_min = -90
# env parameters
RND_SEED = True
SEED = 42
CHANNEL_RND = True
C = 8  # C constant for the improved reward function
ENVIRONMENT_MEMORY = 2
MAX_NUMBER_OF_AGENTS = 2
REWARD_PENALTY = 1.5
N_STATES_BINS = 100
DELTA_T = 1e-3
# q-learning parameters
# training
ALGO_NAME = 'a2c'
OPTIMIZERS = 'adam'
NUM_ENVS = 4
NUM_POOLS = 4
REWARD_FUNCTION = 'multi_agent_continuous'
STATES_OPTIONS = ['sinrs', 'positions', 'channels']
MOTION_MODEL = 'random'
STATES_FUNCTION = 'multi_agent'
# MAX_STEPS = 1000
MAX_STEPS = 40000
EVAL_STEPS = 1000
# MAX_STEPS = 1000
STEPS_PER_EPISODE = 5
EVAL_STEPS_PER_EPISODE = 10
TEST_NUM_EPISODES = ceil(1000 / NUM_POOLS)
REPLAY_MEMORY_SIZE = int(10E3)
ACTOR_LEARNING_RATE = 2E-4
CRITIC_LEARNING_RATE = 2E-3
HIDDEN_SIZE = 64
N_HIDDEN_LAYERS = 1
GAMMA = .5
LBDA = .95
BETA = 1e-2
PRINT_EVERY = int(MAX_STEPS/100)
EVAL_EVERY = int(MAX_STEPS/20)
# EVAL_EVERY = int(MAX_STEPS / 1)
OU_DECAY_PERIOD = 100000
# OU_DECAY_PERIOD = STEPS_PER_EPISODE
# ou noise params
OU_MU = 0.2
OU_THETA = .25
OU_MAX_SIGMA = .5
OU_MIN_SIGMA = .05
OU_DIM = 1e-2
# adaptive noise params
INITIAL_STDDEV = 0.1
DESIRED_ACTION_STDDEV = 0.3
ADAPTATION_COEFFICIENT = 1.01
UPDATE_PERTURBERD_EVERY = 5
# normal actions noise
NORMAL_LOC = 0
NORMAL_SCALE = 4
NORMAL_T = MAX_STEPS
NORMAL_MIN_SCALE = .01

if RND_SEED:
    random_seed(SEED)
torch_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
env_params = EnvironmentParameters(
    rb_bandwidth, None, p_max, noise_power,
    bs_gain, user_gain, sinr_threshold_train,
    n_mues, MAX_NUMBER_OF_AGENTS, n_rb, bs_radius,
    c_param=C, mue_margin=mue_margin,
    min_d2d_pair_distance=MIN_D2D_PAIR_DISTANCE,
    max_d2d_pair_distance=MAX_D2D_PAIR_DISTANCE
)
channel_to_devices = BANChannel(rnd=CHANNEL_RND)
channel_to_bs = UrbanMacroNLOSWinnerChannel(
    rnd=CHANNEL_RND, f_c=carrier_frequency, h_bs=bs_height, h_ms=device_height,
    small_sigma=4.0, sigma=8.0
)
ref_env = CompleteEnvironment12(
    env_params,
    channel_to_bs,
    channel_to_devices,
    reward_penalty=REWARD_PENALTY,
    memory=ENVIRONMENT_MEMORY,
    bs_height=bs_height,
    reward_function=REWARD_FUNCTION,
    states_options=STATES_OPTIONS,
    memories_capacity=int(5e2),
    dt=DELTA_T,
    states_function=STATES_FUNCTION
)
envs = [deepcopy(ref_env) for _ in range(NUM_ENVS)]
a_min = -60
a_max = 60
a_offset = -10
# a_min = 0 + 1e-9
# a_max = db_to_power(p_max - 10)
action_size = MAX_NUMBER_OF_AGENTS
env_state_size = ref_env.state_size()
# actions = db_five(p_min, p_max)
actions = db_six(p_min, p_max)
# actions = [-90, -60, -40, -30, -20, -10]
# actions = db_ten(p_min, p_max)
NUMBER_OF_ACTIONS = len(actions)
framework = A2CDiscreteFramework(
    env_state_size,
    len(actions),
    HIDDEN_SIZE,
    N_HIDDEN_LAYERS,
    STEPS_PER_EPISODE,
    NUM_ENVS*MAX_NUMBER_OF_AGENTS,
    ACTOR_LEARNING_RATE,
    CRITIC_LEARNING_RATE,
    BETA,
    GAMMA,
    LBDA,
    OPTIMIZERS,
    torch_device
)
best_framework = deepcopy(framework)
param_noise = AdaptiveParamNoiseSpec(
    initial_stddev=INITIAL_STDDEV,
    desired_action_stddev=DESIRED_ACTION_STDDEV,
    adaptation_coefficient=ADAPTATION_COEFFICIENT
)
ou_noise = SysSimOUNoise(
    action_size,
    a_min, a_max,
    OU_MU, OU_THETA,
    OU_MAX_SIGMA,
    OU_MIN_SIGMA,
    OU_DECAY_PERIOD
)
ou_noise2 = OUNoise2(
    OU_MU,
    action_size,
    OU_MAX_SIGMA,
    OU_THETA,
    OU_DIM
)
decaying_noise = DecayingGaussNoise(
    NORMAL_LOC, NORMAL_SCALE, NORMAL_T, NORMAL_MIN_SCALE, action_size
)
agents = [[
    A2CAgent(torch_device, actions)
    for _ in range(MAX_NUMBER_OF_AGENTS)
] for _ in range(NUM_ENVS)]
central_agent = A2CCentralAgent(torch_device)

def train(start: float, writer: SummaryWriter, pool):
    global envs
    global agents
    mue_sinrs_bag = list()
    d2d_sinrs_bag = list()
    rewards_bag = list()
    mue_avail_bag = list()
    step = 0
    collected_states = list()
    best_avg_reward = float('-inf')
    while step < MAX_STEPS:
        for e, a in zip(envs, agents):
            e.reset()
            e.build_scenario(a, d2d_limited_power=False,
                             motion_model=MOTION_MODEL)
        obs, _, _, envs, agents = env_step(pool, envs, agents)
        i = 0
        total_entropy = 0
        while i < STEPS_PER_EPISODE:
            action, log_prob, entropy, value = \
                central_agent.act(obs, framework)
            a_index = 0
            for ag in agents:
                for a in ag:
                    a.action_index = action[a_index].item()
                    a.action = actions[a.action_index]
                    a_index += 1
            next_obs, reward, done, envs, agents = env_step(pool, envs, agents)
            total_entropy += entropy
            framework.push_experience(log_prob, value,
                                      reward, done, obs, action)
            obs = next_obs
            i += 1
            step += 1
            if step % PRINT_EVERY == 0:
                now = (time() - start) / 60
                print_stuff_ddpg(step, now, MAX_STEPS, 'standard')
            # testing
            if step % EVAL_EVERY != 0:
                continue
            t_bags = test(framework, pool)
            t_rewards = t_bags['rewards']
            t_mue_sinrs = t_bags['mue_sinrs']
            t_d2d_sinrs = t_bags['d2d_sinrs']
            t_availability = t_bags['mue_availability']
            # mue spectral eff
            mue_sinrs_bag.append(t_mue_sinrs)
            # average d2d spectral eff
            d2d_sinrs_bag.append(t_d2d_sinrs)
            rewards_bag.append(t_rewards)
            mue_avail_bag.append(t_availability)
            # write metrics
            t_mue_sinrs = np.mean(t_mue_sinrs, axis=0)
            t_d2d_sinrs = np.mean(t_d2d_sinrs, axis=0)
            t_d2d_sinrs = {f'device {i}': a for i, a in enumerate(t_d2d_sinrs)}
            # t_d2d_sinrs = np.mean(t_d2d_sinrs)
            t_availability = np.mean(t_availability, axis=0)
            t_rewards = np.mean(t_rewards)
            if t_rewards > best_avg_reward:
                best_framework.a2c.actor.load_state_dict(
                    framework.a2c.actor.state_dict())
                best_avg_reward = t_rewards
            writer.add_scalar(
                '2. Testing - Average MUE SINRs [dB]',
                t_mue_sinrs,
                step
            )
            writer.add_scalars(
                '2. Testing - Average D2D SINRs [dB]',
                t_d2d_sinrs,
                step
            )
            writer.add_scalar(
                '2. Testing - Aggregated Rewards', t_rewards, step)
            writer.add_scalar('2. Testing - Average MUE Availability',
                              np.mean(t_availability), step)
        _, _, _, next_value = central_agent.act(next_obs, framework)
        framework.push_next(next_obs, next_value, total_entropy)
        a_loss, c_loss = framework.learn()
        writer.add_scalar('1. Training - Actor Losses', a_loss, step)
        writer.add_scalar('1. Training - Critic Losses', c_loss, step)
        # adaptive param noise
    all_bags = {
        'mue_sinrs': mue_sinrs_bag,
        'd2d_sinrs': d2d_sinrs_bag,
        'collected_states': collected_states,
        'rewards': rewards_bag,
    }
    return all_bags

def test(framework: A2CDiscreteFramework, pool):
    global envs
    global agents
    mue_availability = []
    mue_sinrs = []
    d2d_sinrs = []
    rewards_bag = []
    for _ in range(TEST_NUM_EPISODES):
        for e, a in zip(envs, agents):
            e.reset()
            e.build_scenario(a, d2d_limited_power=False,
                             motion_model=MOTION_MODEL)
        obs, _, _, envs, agents = env_step(pool, envs, agents)
        i = 0
        done = False
        ep_availability = []
        ep_rewards = []
        ep_mue_sinrs = []
        ep_d2d_sinrs = []
        while i < EVAL_STEPS_PER_EPISODE:
            action, log_prob, entropy, value = \
                central_agent.act(obs, framework)
            a_index = 0
            for ag in agents:
                for a in ag:
                    a.action_index = action[a_index].item()
                    a.action = actions[a.action_index]
                    a_index += 1
            next_obs, reward, done, envs, agents = \
                env_step_test(pool, envs, agents)
            obs = next_obs
            st_mue_sinrs = np.array([e.mue.sinr for e in envs])
            avails = st_mue_sinrs > envs[0].params.sinr_threshold
            env_pairs = [p for p in [e.d2d_pairs for e in envs]]
            st_d2d_sinrs = []
            for ep in env_pairs:
                sinrs = [p[0].sinr for p in ep]
                st_d2d_sinrs.append(sinrs)
            ep_availability += avails.tolist()
            ep_rewards += reward.tolist()
            ep_mue_sinrs += [*st_mue_sinrs]
            ep_d2d_sinrs += st_d2d_sinrs
            i += 1
        rewards_bag += ep_rewards
        mue_sinrs += ep_mue_sinrs
        d2d_sinrs += ep_d2d_sinrs
        mue_availability += ep_availability
    all_bags = {
        'rewards': rewards_bag,
        'mue_sinrs': mue_sinrs,
        'd2d_sinrs': d2d_sinrs,
        'mue_availability': mue_availability
    }
    return all_bags

def evaluate(start: float, writer: SummaryWriter):
    step = 0
    mue_availability = []
    env.reset()
    env.build_scenario(agents, motion_model=MOTION_MODEL)
    env.reset_devices_positions()
    devices = env.get_devices()
    trajectories = {d.id: [d.position] for d in devices}
    d2d_sinrs = []
    while step < EVAL_STEPS:
        obs, _, _, _ = env.step(agents)
        now = (time() - start) / 60
        print_evaluating(step, now, EVAL_STEPS)
        done = False
        i = 0
        past_actions = np.zeros(len(agents))
        with torch.no_grad():
            for j, agent in enumerate(agents):
                agent.get_action(framework, obs[j])
                past_actions[j] = agent.action_index
        next_obs, reward, _, _ = env.step(agents)
        obs = next_obs
        i += 1
        step += 1
        rewards = {f'device {i}': r for i, r in enumerate(reward)}
        writer.add_scalars('3. Eval - Rewards', rewards, step)
        sinrs = [a.d2d_tx.sinr for a in agents]
        d2d_sinrs.append(sinrs)
        sinrs = {f'device {i}': s for i, s in enumerate(sinrs)}
        writer.add_scalars('3. Eval - SINRs [dB]', sinrs, step)
        writer.add_scalars(
            '3. Eval - Transmission powers [dBW]',
            {f'device {i}': a for i, a in enumerate(past_actions)},
            step
        )
        mue_success = int(env.mue.sinr > env.params.sinr_threshold)
        mue_availability.append(mue_success)
        writer.add_scalar('3. Eval - MUE success', mue_success, step)
        for d in env.get_devices():
            trajectories[d.id].append(d.position)
    mue_availability = np.mean(mue_availability)
    d2d_sinrs = np.mean(d2d_sinrs, axis=0)
    writer.add_text('3. Eval - Average MUE availability',
                    str(mue_availability), step)
    for i, s in enumerate(d2d_sinrs):
        writer.add_text(f'3. Eval - Average D2D {i} SINR', str(s), step)
    writer.add_figure('3. Eval - Trajectories',
                      plot_trajectories(env, trajectories), step)

def run():
    # multiprocessing pool
    pool = Pool(NUM_POOLS)
    # make data dir
    filename = gen.path_leaf(__file__)
    filename = filename.split('.')[0]
    dir_path = f'data/{ALGO_NAME}/{filename}'
    data_path, _ = gen.make_dir_timestamp(dir_path)
    writer = SummaryWriter(f'{data_path}/tensorboard')
    start = time()
    # set environment up
    train_bags = train(start, writer, pool)
    # states = train_bags['collected_states']
    # plot_env_states(states, N_STATES_BINS, f'{data_path}/env_states.png')
    test_bags = test(framework, pool)
    # evaluate(start, writer)
    writer.close()
    # save stuff
    now = (time() - start) / 60
    data_file_path = f'{data_path}/log.pickle'
    data = {
        'train_bags': train_bags,
        'test_bags': test_bags,
        'elapsed_time': now,
        'eval_every': EVAL_EVERY,
        'mue_sinr_threshold': sinr_threshold_train,
    }
    save_with_pickle(envs[0], f'{data_path}/env.pickle')
    save_with_pickle(data, data_file_path)
    copyfile(__file__, f'{data_path}/{filename}.py')
    torch.save(framework, f'{data_path}/last_model.pt')
    torch.save(best_framework, f'{data_path}/best_model.pt')
    print(f'done. Elapsed time: {now} minutes.')
    pool.close()

if __name__ == '__main__':
    run()


text:
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d
from multiprocessing import Pool
import numpy as np
from joblib import Parallel, delayed
from tqdm import trange
import warnings as warn

signal_path = 'output\\signal\\signal_{0}.dat'
network_path = 'output\\Complete_Network.dat'
neuron_path = 'output\\Complete_Network_Neurons.dat'

def f(i, fig, ax, signals=0):
    ax.view_init(30, float(i) / float(545)* 360)

    ax.text2D(0.0, 0.0, 't  = {0}'.format(i), transform=ax.transAxes)

    if signals:
        lines = plot_scatt(signal_path.format(i), ax, c='r', s=70)
        if lines:
            fig.canvas.draw()
            save_fig(i)
            lines.remove()
        else:
            save_fig(i)
    else:
        save_fig(i)

def save_fig(i):
    plt.savefig('output\\frames\\frame' + str(i) + ".png", format='png')

def load_points(path):
    with warn.catch_warnings():
        # To catch "Empty input file" warning
        warn.simplefilter('ignore', category=UserWarning)
        return np.loadtxt(path)

def plot_scatt(path, ax, **kwargs):

    points = load_points(path).T

    if points.size:
        x, y, z = points
        return ax.scatter(x, y, z, **kwargs)
    else:
        return None

def main():
    fig = plt.figure()
    ax = fig.add_subplot(111, projection ='3d')
    ax.set_axis_off()

    plot_scatt(network_path, ax, c='b', marker='.', s=5, alpha=0.02)
    plot_scatt(neuron_path, ax, c='b', marker='.', s=50, alpha=0.5)

    limits = 40
    ax.set_xlim3d(-limits,limits)
    ax.set_ylim3d(-limits,limits)
    ax.set_zlim3d(-limits,limits)

    ax.view_init(30,0)

    Parallel(n_jobs=4)(delayed(f)(i, fig, ax, True) for i in trange(545))

if __name__ == '__main__':
    main()


text:
import getpass
import sys
from traceback import format_exc
import os
import base64
import stat
import argparse

class Colors:
    __black_o_green = '\x1b[5;30;42m'
    __black_o_yellow = '\x1b[5;30;43m'
    __black_o_red = '\x1b[5;30;43m'
    __end = '\x1b[0m'
    error = __black_o_red + " [ERROR] " + __end
    warning = __black_o_yellow + " [ADVERTENCIA] " + __end
    ok = __black_o_green + " [OK] " + __end

git = None

print("Comprobando paquetes necesarios...")
try:
    import pip
except Exception:
    print(Colors.error + "El paquete \'pip\' no se encuentra instalado")
    exit(-1)

try:
    from github import Github, GithubException
except Exception:
    print(Colors.warning + "Instalando paquete necesario: pygithub")
    pip.main(['install', '--user', 'pygithub'])

try:
    import numpy as np
except Exception:
    print("Instalando paquete necesario: numpy")
    pip.main(['install', '--user', 'numpy'])

print(Colors.ok + "Todos los paquetes necesarios instalados")

def status_ok(github):
    s = github.get_api_status()
    return True if s.status == 'good' else False

def download_repository(repository, path_to_download):
    """
    Download all contents at server_path with commit tag sha in
    the repository.
    """

    sha = "master"  # r.get_branch("master").commit.sha
    contents = repository.get_dir_contents(path_to_download)
    base_dir = os.path.join('test', repository.name)

    for content in contents:
        print("\tProcesando {}".format(content.path))
        if content.type == 'dir':
            download_repository(repository, content.path)
        else:
            try:
                if not os.path.exists(os.path.join(base_dir, path_to_download)):
                    os.makedirs(os.path.join(base_dir, path_to_download))
                path = content.path
                file_content = repository.get_contents(path, ref=sha)
                file_data = base64.b64decode(file_content.content)
                file_out = open(os.path.join(base_dir, path_to_download, content.name), "w")
                file_out.write(file_data.decode())
                file_out.close()
            except (GithubException, IOError) as exc:
                print('\tError al procesar {}: {}'.format(content.path, exc))

def permissions(repository):
    path = os.path.join('test', repository.name)
    # os.chmod('somefile', stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)

    for file in os.listdir(os.path.join(path, 'nodes')):
        if file.endswith(".py") and '__' not in file:
            try:
                print("Dando permisos de ejecución a {} ...".format(file))
                full_path = os.path.join(path, 'nodes', file)
                current_permissions = stat.S_IMODE(os.lstat(full_path).st_mode)
                os.chmod(full_path, current_permissions | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
                print("{} ejecutable".format(file) + Colors.ok)

            except Exception:
                print(Colors.warning + "NO se pudo dar permiso al fichero {}  :(".format(file))

def get_user_organization():
    global git

    try:
        # user = input("Usuario Github: ")
        # password = getpass.getpass(prompt='Contraseña: ', stream=sys.stderr)
        user = "Ait0r"
        password = "22593712aitor."
        git = Github(user, password)
        print("Comprobando estado de github")
        if status_ok(git):
            print("\tGithub: " + Colors.ok)
        else:
            print("\t" + Colors.error + "Api de Github no disponible")
            exit(-1)

        print("Conectando a github...")
        u = git.get_user()
        orgs = list(u.get_orgs())

        print("Entrando en organizaciones...")
        if len(orgs) > 1:
            for idx, o in enumerate(orgs):
                print("{} - {}".format(idx, o.name))
            org = input("Organización a utilizar: ")
            print("Seleccionada la organización {}".format(orgs[int(org)].name))
            org = orgs[0]
        else:
            print("Detectada solo 1 organización ({})".format(orgs[0].name))
            org = orgs[0]

        return org, u

    except Exception:
        print(format_exc())

def get_repositories(org):
    print("Accediendo a repositorios de {}".format(org.name))

    repos = list(org.get_repos())

    for idx, repo in enumerate(repos):
        print("{} - {}".format(idx, repo.name))

    not_correct = True
    to_migrate = None
    while not_correct:
        try:
            to_migrate = input("Repositorios a migrar (Separados por \" \"): ")
            to_migrate = list(map(int, to_migrate.split(" ")))
            not_correct = False
        except Exception:
            print("Error, introduce el número de repositorio de nuevo")

    if len(to_migrate) > 0:
        repos = np.array(repos)
        selected_repos = repos[to_migrate]
        print("Repositorios seleccionados: ")
        for r in selected_repos:
            print("\t{}".format(r.name))

        print("-" * 50 + "\n\n")
        for r in selected_repos:
            print("\n\nIniciando descarga de {}".format(r.name))
            download_repository(r, '')
            print(Colors.ok + "Repositorio descargado")
            print("Iniciando asignación de permisos")
            permissions(r)
    else:
        print(Colors.error + "Ningun repositorio seleccionado")
        exit(0)

def edit_bash(d):
    packages = os.listdir(d)
    bashrc = "/home/{}/.bashrc".format(getpass.getuser())
    file = open(bashrc, mode='a')
    export = "export PYTHONPATH=$PYTHONPATH:"
    file.writelines("\n\n\n#Configuracion ROS\n\n")
    for p in packages:
        path_to_add = os.path.join(os.path.abspath(p), 'src')
        line = "{}={}\n".format(p.split('-')[-1], path_to_add)
        file.writelines(line)
        export += '${}:'.format(p.split('-')[-1])

    line = "ROS=/opt/ros/kinetic/lib/python2.7/dist-packages\n"
    file.writelines(line)
    export += '${}:'.format('Ros')

    file.writelines(export[:-1])

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Directorio para realizar la migración')
    parser.add_argument('-o', action="store", dest="output_path", type=str, required=True)

    args = parser.parse_args()
    BASE_DIR = os.path.abspath(args.output_path)
    edit_bash(BASE_DIR)
    exit(0)
    print('Directorio seleccionado: {!r}'.format(BASE_DIR))
    git_org, git_user = get_user_organization()
    get_repositories(git_org)

max_stars_repo_name: SIMCA-USI/ROS-Miragration

max_issues_repo_name: SIMCA-USI/ROS-Miragration

max_forks_repo_name: SIMCA-USI/ROS-Miragration


text:
import pandas as pd
import numpy as np
from app import db
import json

with open(db.data_dir + 'supported.json', 'r') as fp:
    supported = json.load(fp)

attr_id_map = pd.read_sql('attr_id_map', db.db_engine)
attr_id_map.set_index("id", drop=True, inplace=True)
attr_id_map = attr_id_map.to_dict(orient="index")

class Data(object):
    def __init__(self, table_name, place_type):
        self.engine = db.db_engine
        self.table_name = table_name
        self.place_type = place_type
        self.supported_attributes = supported[table_name]

    def create_pivoted_table(self):
        print("Creating pivot table for {}".format(self.table_name))
        self.df = pd.read_sql(self.table_name, self.engine)
        values = list(set(self.df.keys())-set(['year', 'id']))
        self.df = None
        pivoted = pd.pivot_table(pd.read_sql(self.table_name, self.engine),
                             index='id',
                             columns='year',
                             values=values)
        keys = list(set([key[0] for key in pivoted.keys()]))
        for i, key in enumerate(keys, 1):
            pivoted[key].iloc[:, 0] = pivoted[key].iloc[:, 0].fillna(0)
            pivoted[key].iloc[:, -1] = pivoted[key].iloc[:, -1].fillna(0)
            print(np.round(i/len(keys), 2), end='\r')

        self.df = pivoted.interpolate(axis=1)

        print("Successfully created pivot table for {}".format(self.table_name))


text:
from dnnv.properties import *
import numpy as np

N = Network("N")
input_lb = np.full(N.input_shape[0], -10, dtype=N.input_details[0].dtype)
input_ub = np.full(N.input_shape[0], 10, dtype=N.input_details[0].dtype)

Forall(x, Implies(input_lb <= x <= input_ub, N(x) > 0))

max_stars_repo_name: dlshriver/dnnv-benchmarks

max_issues_repo_name: dlshriver/dnnv-benchmarks

max_forks_repo_name: dlshriver/dnnv-benchmarks