text:
import numpy as np

import torch
import torch.nn as nn

def get_mlp(layer_dims, activation=nn.ReLU, last_layer_activation=False):
    """
    Simple helper to create a fully connected network, given a list of
    number of hidden units in each layer.
    """
    layers = []
    for i in range(len(layer_dims) - 1):
        layers.append(nn.Linear(layer_dims[i], layer_dims[i + 1]))
        layers.append(activation())
    if not last_layer_activation:  # if not using activation in the last (output) layer
        layers = layers[:-1]

    return nn.Sequential(*layers)

from configs import float_type

if float_type == 'float32':
    float_type = torch.float32
else:
    float_type = torch.float

torch_softplus = torch.nn.Softplus()

class CondNoiseMLP(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_layer_dims, stoch,
                 activation=nn.ReLU, noise_activation=nn.ReLU, input_stats=None,
                 inv_masses=None, output_stats=None, output_raw=False,
                 sigma_fun='exp'):
        """
        Generic model of an implicit and sample-able conditional distribution, p(output | input), parameterized by MLP.
        In the deterministic case (stoch=False), output = nn(input).
        In the stochastic case, we first sample a conditional noise from p(\epsilon|input) = N(mu(input), sigma2(input))
        where mu() and sigma() are neural nets, then we deterministically transform [\epsilon, input] into output with
        another nn. We set \epsilon to have the same dimension as the output, as in a conditional normalizing flow.
        As neural nets tend to be easier to optimize with SGD when the input and output are standardized, this class
        provides the option to either train on standardized data (with input_stats=None, output_raw=False), or raw data
        (with input_stats, output_stats specified as np/torch arrays, and output_raw=True). When training on raw data,
        the input is standardized by input_stats before feeding to the NN, and the model output is the NN output
        un-standardized by output_stats.
        :param input_dim:
        :param output_dim:
        :param hidden_layer_dims:
        :param stoch:
        :param activation: type of actiavtion function in the NNs
        :param noise_activation: overriding option for the actiavtion function in the conditional noise NN
        :param input_stats: if specified, we assume the input is 'raw' (not standardized), and input_stats=[mean, std]
        of the raw input in the form of a 2 x D torch tensor, with inputs_stats[0]=mean, input_stats[1]=std; raw input
        will then be standardized by input_stats before being fed to the NN
        :param inv_masses: a vector of invariant masses of particles. If provided, then output_stats must also be
        provided, and we assume the (raw) model output consists of the 4-momentum (px,py,pz,E) of K particles concatenated
        together, and inv_masses is a vector of length K, such that for each particle, the invariant mass relation
        E^2 = px^2 + py^2 + pz^2 + inv_mass^2 holds.
        :param output_stats: [mean, std] of the raw output space in the form of a 2 x D torch tensor.
        :param output_raw: boolean; if True, the model output will be "raw"; otherwise, the NN output will be directly
        taken to be the model output.
        :param sigma_fun: activation ('squash') function that converts the unrestricted NN output to positively valued
        sigma; default is 'exp' (so the NN parameterizes log sigma), can also use 'softplus' for more numerical stability
        """
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim  # dimensionality of the output space
        self.hidden_layer_dims = hidden_layer_dims
        if sigma_fun == 'exp':
            sigma_fun = torch.exp
        elif sigma_fun == 'softplus':
            sigma_fun = torch_softplus
        else:
            raise NotImplementedError
        self.sigma_fun = sigma_fun

        if isinstance(input_stats, np.ndarray):
            input_stats = torch.Tensor(input_stats)
        self.input_stats = input_stats
        # self.inv_masses = torch.Tensor(inv_masses)
        self.inv_masses = inv_masses
        if inv_masses is not None:
            assert output_stats is not None, "Must provide output stats in order to enforce invariant mass relation."
            num_particles = len(inv_masses)
            assert 4 * num_particles == self.output_dim, "Output space dimensionality != 4 times the number of particles"
            self.num_particles = num_particles
            output_dim -= num_particles  # these dimensions correspond to energy and will be manually computed
        if isinstance(output_stats, np.ndarray):
            output_stats = torch.Tensor(output_stats)
        self.output_stats = output_stats
        self.output_raw = output_raw

        self.stoch = stoch
        if not stoch:
            self.output_nn = get_mlp([input_dim, *hidden_layer_dims, output_dim], activation=activation,
                                     last_layer_activation=False)
        else:
            noise_dim = output_dim
            self.noise_dim = noise_dim
            self.output_nn = get_mlp([input_dim + noise_dim, *hidden_layer_dims, output_dim], activation=activation,
                                     last_layer_activation=False)
            # currently using the reverse architecture of the output_nn;
            # 2 * output_dim for mu and logsigma
            self.cond_noise_nn = get_mlp([input_dim,
                                          *list(reversed(hidden_layer_dims)), 2 * output_dim],
                                         activation=noise_activation, last_layer_activation=False)
            # self.cond_noise_nn = get_mlp([input_dim, *hidden_layer_dims, 2 * output_dim], activation=activation, last_layer_activation=False)

            # custom (hack) initialization to make sure logsigma won't be too
            # small at the beginning of training to cause NaN issues
            for layer in self.cond_noise_nn.modules():
                if isinstance(layer, nn.Linear):
                    final_bias = layer.bias
            final_bias.data.fill_(1.0)

    def forward(self, input):
        """
        Given input batch vector, produce samples from the conditional distribution p(output|input).
        :param input:
        :return:
        """
        if self.input_stats is not None:
            input_mean, input_std = self.input_stats.to(input.device)
            input = (input - input_mean) / input_std
        if self.stoch:
            noise_dim = self.noise_dim
            mean, logsigma = self.cond_noise_nn(input).split(noise_dim, dim=-1)
            # sigma = torch.exp(logsigma)
            # sigma = torch_softplus(logsigma)
            sigma = self.sigma_fun(logsigma)
            input_shape = input.shape
            eps = torch.randn([*input_shape[:-1], noise_dim], dtype=input.dtype, device=input.device)
            eps = mean + eps * sigma
            input = torch.cat((input, eps), dim=-1)
        X = self.output_nn(input)
        if self.inv_masses is not None:
            num_particles = self.num_particles
            batch = len(X)
            output = torch.zeros([batch, self.output_dim], dtype=X.dtype, device=X.device)
            output_mean, output_std = self.output_stats.to(X.device)
            idx_noE = np.hstack([np.arange(4 * i, 4 * i + 3) for i in range(num_particles)])  # NN output doesn't have E
            X = (X * output_std[idx_noE]) + output_mean[idx_noE]  # unstandardize
            for i in range(num_particles):
                # for the ith particle
                p = X[:, 3 * i:3 * (i + 1)]  # 3-momentum, (px, py, pz)
                m = self.inv_masses[i]
                E = ((p ** 2).sum(axis=-1) + m ** 2) ** 0.5
                output[:, 4 * i:4 * i + 3] = p
                output[:, 4 * i + 3] = E

            if self.output_raw:  # output in "raw" (data observation) space
                pass
            else:
                output = (output - output_mean) / output_std
        else:
            output = X
            if self.output_raw:
                output_mean, output_std = self.output_stats.to(X.device)
                output = (output * output_std) + output_mean

        return output

class CondNoiseAutoencoder(nn.Module):
    """
    Consists of a probabilistic encoder, q(z|x), and decoder p_G(x|z). In case they're deterministic (dirac delta
    functions), they're implemented by neural nets z -> x and x -> z. In the stochastic case, they're defined implicitly
    as (z, eps) -> x and (x, eps) -> z.
    """

    def __init__(self, x_dim, z_dim, hidden_layer_dims, stoch_enc=False, stoch_dec=False, activation=nn.ReLU,
                 raw_io=False, x_inv_masses=None, x_stats=None,
                 z_inv_masses=None, z_stats=None, **kwargs):
        """

        :param x_dim:
        :param z_dim:
        :param hidden_layer_dims:
        :param stoch_enc:
        :param stoch_dec:
        :param activation:
        :param raw_io: whether or not the model inputs/outputs raw observations; if False, the model assumes standardized
        input and output.
        :param x_inv_masses:
        :param x_stats:
        :param z_inv_masses:
        :param z_stats:
        """
        super().__init__()
        self.x_dim = x_dim
        self.z_dim = z_dim
        self.stoch_enc = stoch_enc
        self.stoch_dec = stoch_dec
        if raw_io:
            assert x_stats is not None and z_stats is not None

        self.encoder = CondNoiseMLP(x_dim, z_dim, hidden_layer_dims, stoch_enc, activation,
                                    input_stats=x_stats if raw_io else None,
                                    inv_masses=z_inv_masses, output_stats=z_stats,
                                    output_raw=raw_io, **kwargs)
        self.decoder = CondNoiseMLP(z_dim, x_dim, list(reversed(hidden_layer_dims)), stoch_dec, activation,
                                    input_stats=z_stats if raw_io else None,
                                    inv_masses=x_inv_masses, output_stats=x_stats,
                                    output_raw=raw_io, **kwargs)

    def encode(self, x):
        """
        Draw a sample from p_G(z|x) (which might be a dirac delta fn in case of det encoder)
        :param x:
        :return:
        """
        return self.encoder(x)

    def decode(self, z):
        """
        Draw a sample from q(x|z) (which might be a dirac delta fn in case of det deccoder)
        :param z:
        :return:
        """
        return self.decoder(z)

    def forward(self, x):
        """
        Given data batch, return latent code and reconstruction.
        :param x:
        :return:
        """
        z_tilde = self.encode(x)
        x_tilde = self.decode(z_tilde)
        return z_tilde, x_tilde

class BasicAutoencoder(nn.Module):
    """
    Basic (encoder, decoder) implemented by MLPs.
    Consists of a probabilistic encoder, q(z|x), and decoder p_G(x|z). In case they're deterministic (dirac delta
    functions), they're implemented by neural nets z -> x and x -> z. In the stochastic case, they're defined implicitly
    as (z, eps) -> x and (x, eps) -> z.
    """

    def __init__(self, x_dim, z_dim, hidden_layer_dims, stoch_enc=False, stoch_dec=False, activation=nn.ReLU):
        super().__init__()
        self.x_dim = x_dim
        self.z_dim = z_dim
        self.stoch_enc = stoch_enc
        self.stoch_dec = stoch_dec
        layer_dims = [x_dim, *hidden_layer_dims, z_dim]
        if stoch_enc:
            layer_dims[0] *= 2  # for additional noise input
        self.encoder = get_mlp(layer_dims, activation=activation, last_layer_activation=False)
        layer_dims = [z_dim, *hidden_layer_dims, x_dim]
        if stoch_dec:
            layer_dims[0] *= 2  # for additional noise input
        self.decoder = get_mlp(layer_dims, activation=activation, last_layer_activation=False)

    def sample_nosie(self, shape, dtype=float_type):
        eps = torch.randn(shape, dtype=dtype)
        return eps

    def encode(self, x):
        """
        Draw a sample from p_G(z|x) (which might be a dirac delta fn in case of det encoder)
        :param x:
        :return:
        """
        if self.stoch_enc:
            eps = self.sample_nosie(x.shape)
            eps = eps.to(dtype=x.dtype, device=x.device)
            x = torch.cat((x, eps), dim=1)
        return self.encoder(x)

    def decode(self, z):
        """
        Draw a sample from q(x|z) (which might be a dirac delta fn in case of det deccoder)
        :param z:
        :return:
        """
        if self.stoch_dec:
            eps = self.sample_nosie(z.shape)
            eps = eps.to(dtype=z.dtype, device=z.device)
            z = torch.cat((z, eps), dim=1)
        return self.decoder(z)

    def forward(self, x):
        """
        Given data batch, return latent code and reconstruction.
        :param x:
        :return:
        """
        z_tilde = self.encode(x)
        x_tilde = self.decode(z_tilde)
        return z_tilde, x_tilde

class ResBlock(torch.nn.Module):
    def __init__(self, input_dim, activation=nn.ReLU, depth=1):
        """
        A residual block that implements f(x) = MLP(x) + x, where the MLP consists of a series of (BN, activation, linear)
        operations. The output has the same dimensionality as input.
        Architecture (e) 'full pre-activation' of https://towardsdatascience.com/an-overview-of-resnet-and-its-variants-5281e2f56035
        """
        super().__init__()
        layers = []
        for _ in range(depth):
            layers += [
                nn.BatchNorm1d(num_features=input_dim),
                activation(),
                nn.Linear(input_dim, input_dim)  # linear layer preserves dimensionality for residual to work
            ]
        self.module = nn.Sequential(*layers)

    def forward(self, x):
        return self.module(x) + x

class StochasticResNet(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_layer_dims, stoch,
                 activation=nn.ReLU, input_stats=None,
                 inv_masses=None, output_stats=None, output_raw=False,
                 io_residual=False, res_mlp_depth=1):
        """
        Generic model of an implicit and sample-able conditional distribution, p(output | input), implemented by a ResNet.
        In the deterministic case (stoch=False), output = nn(input).
        In the stochastic case, we sample \epsilon from standard Normal, concatenate it with input, and pass
        [\epsilon, input] through a ResNet to get output vector.
        We set \epsilon to have the same dimension as the output, as in a conditional normalizing flow.
        As neural nets tend to be easier to optimize with SGD when the input and output are standardized, this class
        provides the option to either train on standardized data (with input_stats=None, output_raw=False), or raw data
        (with input_stats, output_stats specified as np/torch arrays, and output_raw=True). When training on raw data,
        the input is standardized by input_stats before feeding to the NN, and the model output is the NN output
        un-standardized by output_stats.
        :param input_dim:
        :param output_dim:
        :param hidden_layer_dims: list of dimensions of the ResBlocks; whenever two adjacent numbers are equal, a ResBlock
        with that number of dimensions is created (since a ResBlock must preserve input output dimensionality); otherwise,
        if m!=n, a simple linear layer R^m -> R^n is created to operate between the two dimensions.
        :param stoch:
        :param activation: type of actiavtion function in the NNs
        :param noise_activation: overriding option for the actiavtion function in the conditional noise NN
        :param input_stats: if specified, we assume the input is 'raw' (not standardized), and input_stats=[mean, std]
        of the raw input in the form of a 2 x D torch tensor, with inputs_stats[0]=mean, input_stats[1]=std; raw input
        will then be standardized by input_stats before being fed to the NN
        :param inv_masses: a vector of invariant masses of particles. If provided, then output_stats must also be
        provided, and we assume the (raw) model output consists of the 4-momentum (px,py,pz,E) of K particles concatenated
        together, and inv_masses is a vector of length K, such that for each particle, the invariant mass relation
        E^2 = px^2 + py^2 + pz^2 + inv_mass^2 holds.
        :param output_stats: [mean, std] of the raw output space in the form of a 2 x D torch tensor.
        :param output_raw: boolean; if True, the model output will be "raw"; otherwise, the NN output will be directly
        taken to be the model output.
        :param io_residual: if True, will add a residual connection from input directly to the output.
        If inv-mass constraint is enforced to output so that the NN output is only 18 dimensional, the corresponding 18
        elements of the input (instead of the entire input vector) will be added to the NN output.
        :param res_mlp_depth: depth of MLP within each ResBlock.

        """
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim  # dimensionality of the output space
        self.hidden_layer_dims = hidden_layer_dims
        self.io_residual = io_residual

        if isinstance(input_stats, np.ndarray):
            input_stats = torch.Tensor(input_stats)
        self.input_stats = input_stats
        # self.inv_masses = torch.Tensor(inv_masses)
        self.inv_masses = inv_masses
        if inv_masses is not None:
            assert output_stats is not None, "Must provide output stats in order to enforce invariant mass relation."
            num_particles = len(inv_masses)
            assert 4 * num_particles == self.output_dim, "Output space dimensionality != 4 times the number of particles"
            self.num_particles = num_particles
            output_dim -= num_particles  # these dimensions correspond to energy and will be manually computed
        if isinstance(output_stats, np.ndarray):
            output_stats = torch.Tensor(output_stats)
        self.output_stats = output_stats
        self.output_raw = output_raw

        self.stoch = stoch

        if stoch:
            self.noise_dim = output_dim
            input_dim += self.noise_dim
        layer_dims = [input_dim, *hidden_layer_dims, output_dim]
        nn_layers = []
        for i in range(len(layer_dims) - 1):
            in_dim, out_dim = layer_dims[i], layer_dims[i + 1]
            if in_dim == out_dim:  # use residual block
                layer = ResBlock(in_dim, activation, depth=res_mlp_depth)
            else:  # just apply a simple linear transform to change dimensionality
                layer = nn.Linear(in_dim, out_dim)
            nn_layers.append(layer)
        self.nn = nn.Sequential(*nn_layers)

    def forward(self, input):
        """
        Given input batch vector, produce samples from the conditional distribution p(output|input).
        :param input:
        :return:
        """
        if self.input_stats is not None:
            raw_input = input.clone()
            input_mean, input_std = self.input_stats.to(input.device)
            input = (input - input_mean) / input_std
        if self.stoch:
            noise_dim = self.noise_dim
            input_shape = input.shape
            eps = torch.randn([*input_shape[:-1], noise_dim], dtype=input.dtype, device=input.device)
            input = torch.cat((input, eps), dim=-1)
        X = self.nn(input)
        if self.inv_masses is not None:
            num_particles = self.num_particles
            batch = len(X)
            output = torch.zeros([batch, self.output_dim], dtype=X.dtype, device=X.device)
            output_mean, output_std = self.output_stats.to(X.device)
            idx_noE = np.hstack([np.arange(4 * i, 4 * i + 3) for i in range(num_particles)])  # NN output doesn't have E
            X = (X * output_std[idx_noE]) + output_mean[idx_noE]  # unstandardize
            if self.io_residual:  # residual connection from the corresponding elements of the raw input
                X = X + raw_input[:, idx_noE]  # X is [batch_size, 18], raw_input is [batch_size, 24]
            for i in range(num_particles):
                # for the ith particle
                p = X[:, 3 * i:3 * (i + 1)]  # 3-momentum, (px, py, pz)
                m = self.inv_masses[i]
                E = ((p ** 2).sum(axis=-1) + m ** 2) ** 0.5
                output[:, 4 * i:4 * i + 3] = p
                output[:, 4 * i + 3] = E

            if self.output_raw:  # output in "raw" (data observation) space
                pass
            else:
                output = (output - output_mean) / output_std
        else:
            output = X
            if self.output_raw:
                output_mean, output_std = self.output_stats.to(X.device)
                output = (output * output_std) + output_mean
                if self.io_residual:  # residual connection from the corresponding elements of the input
                    output = output + raw_input

        return output

class MaskedStochasticResNet(StochasticResNet):
    """
    A wrapper around StochasticResNet that allows optionally applying a mask/filter to the output.
    """

    def __init__(self, mask_fun, *args, **kwargs):
        """

        :param mask_fun: given N x D matrix of "raw" data, returns N dimensional vector indicating whether each row is
        "valid". This was used for filtering out ppttbar data that don't pass the detector threshold.
        :param args:
        :param kwargs:
        """
        super().__init__(*args, **kwargs)
        self.mask_fun = mask_fun

    def forward(self, input, mask_output=False, return_mask=True):
        """
        :param input:
        :param mask_output: whether to apply mask to output
        :param return_mask: whether to also return the boolean that indicats which input passed the threshold
        :return:
        """
        output = super().forward(input)
        if (not mask_output) and (not return_mask):
            return output

        if self.output_raw:  # output is in raw space already
            raw_output = output
        else:
            output_mean, output_std = self.output_stats.to(input.device)
            raw_output = (output * output_std) + output_mean

        mask = self.mask_fun(raw_output)

        if mask_output:  # apply mask
            output = output[mask]
        if return_mask:
            return output, mask
        else:
            return output

class Autoencoder(nn.Module):
    """
    A generic autoencoder module that wraps around an (encoder, decoder) pair, with encoder/decoder implemented by
        a user defined ConditionalModel class. Subsumes the legacy CondNoiseAutoencoder (specialized to CondNoiseMLP).
    """

    def __init__(self, x_dim, z_dim, ConditionalModel, encoder_hidden_layer_dims, stoch_enc=True, stoch_dec=True,
                 activation=nn.ReLU,
                 raw_io=False, x_inv_masses=None, x_stats=None,
                 z_inv_masses=None, z_stats=None, **kwargs):
        """
        :param x_dim:
        :param z_dim:
        :param ConditionalModel: a class that implements the (stochastic) encoder and decoder model.
        :param encoder_hidden_layer_dims: list of hidden layer dims for the encoder; the reverse architecture is used
        for decoder.
        :param stoch_enc:
        :param stoch_dec:
        :param activation:
        :param raw_io: whether or not the model inputs/outputs raw observations; if False, the model assumes standardized
        input and output.
        :param x_inv_masses:
        :param x_stats:
        :param z_inv_masses:
        :param z_stats:
        """
        super().__init__()
        self.x_dim = x_dim
        self.z_dim = z_dim
        self.stoch_enc = stoch_enc
        self.stoch_dec = stoch_dec
        assert ConditionalModel in (StochasticResNet, CondNoiseMLP)
        if raw_io:
            assert x_stats is not None and z_stats is not None

        self.encoder = ConditionalModel(x_dim, z_dim, encoder_hidden_layer_dims, stoch_enc, activation,
                                        input_stats=x_stats if raw_io else None,
                                        inv_masses=z_inv_masses, output_stats=z_stats,
                                        output_raw=raw_io, **kwargs)
        self.decoder = ConditionalModel(z_dim, x_dim, list(reversed(encoder_hidden_layer_dims)), stoch_dec, activation,
                                        input_stats=z_stats if raw_io else None,
                                        inv_masses=x_inv_masses, output_stats=x_stats,
                                        output_raw=raw_io, **kwargs)

    def encode(self, x, *args, **kwargs):
        """
        Draw a sample from p_G(z|x) (which might be a dirac delta fn in case of det encoder)
        :param x:
        :return:
        """
        return self.encoder(x, *args, **kwargs)

    def decode(self, z, *args, **kwargs):
        """
        Draw a sample from q(x|z) (which might be a dirac delta fn in case of det deccoder)
        :param z:
        :return:
        """
        return self.decoder(z, *args, **kwargs)

    def forward(self, x):
        """
        Given data batch, return latent code and reconstruction.
        :param x:
        :return:
        """
        z_tilde = self.encode(x)
        x_tilde = self.decode(z_tilde)
        return z_tilde, x_tilde

max_stars_repo_name: yiboyang/otus

max_issues_repo_name: yiboyang/otus

max_forks_repo_name: yiboyang/otus


text:
import logging
import typing

import pandas as pd
import numpy as np

logger = logging.getLogger(__name__)
# logger.setLevel(logging.DEBUG)

# define time constants
SECONDS_PER_MINUTE = 60
MINUTES_PER_HOUR = 60
HOURS_PER_DAY = 24
DAYS_PER_WEEK = 7
DAYS_PER_MONTH_AVG = [28, 30, 31]
DAYS_PER_YEAR_AVG = [365, 366]
DAYS_PER_YEAR = 365.25
DAYS_PER_MONTH = DAYS_PER_YEAR / 12

S_PER_YEAR = SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_YEAR
S_PER_YEAR_0 = (
    SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_YEAR_AVG[0]
)
S_PER_YEAR_1 = (
    SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_YEAR_AVG[1]
)
S_PER_MONTH = SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_MONTH
S_PER_MONTH_28 = (
    SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_MONTH_AVG[0]
)
S_PER_MONTH_30 = (
    SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_MONTH_AVG[1]
)
S_PER_MONTH_31 = (
    SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_MONTH_AVG[2]
)
S_PER_WEEK = SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY * DAYS_PER_WEEK
S_PER_DAY = SECONDS_PER_MINUTE * MINUTES_PER_HOUR * HOURS_PER_DAY
S_PER_HR = SECONDS_PER_MINUTE * MINUTES_PER_HOUR

def calculate_time_frequency(time_diff, model="var"):
    """method that calculates the frequency of a datetime difference (for prediction slices)

    Arguments:
        time_diff {timedelta or float} -- difference between two instances

    Returns:
        str -- string alias representing granularity of pd.datetime object
    """

    return_strings = {
        "var": ["YS", "MS", "M", "W-MON", "D", "H", "S"],
        "gluon": [
            ("12M", "YS"),
            ("M", "MS"),
            ("M", "M"),
            ("W", "W-MON"),
            ("D", "D"),
            ("H", "H"),
            ("S", "S"),
        ],
    }

    # convert to seconds representation
    if type(time_diff) is pd._libs.tslibs.timedeltas.Timedelta:
        time_diff = time_diff.total_seconds()

    if time_diff % S_PER_YEAR_0 == 0:
        logger.debug("granularity is years")
        return return_strings[model][0]
    elif time_diff % S_PER_YEAR_1 == 0:
        logger.debug("granularity is years")
        return return_strings[model][0]
    elif time_diff % S_PER_MONTH_31 == 0:
        # Sunspots monthly
        logger.debug("granularity is months")
        return return_strings[model][1]
    elif time_diff % S_PER_MONTH_30 == 0:
        logger.debug("granularity is months 30")
        return return_strings[model][1]
    elif time_diff % S_PER_MONTH_28 == 0:
        # PHEM monthly
        logger.debug("granularity is months 28")
        return return_strings[model][2]
    elif time_diff % S_PER_WEEK == 0:
        logger.debug("granularity is weeks")
        return return_strings[model][3]
    elif time_diff % S_PER_DAY == 0:
        logger.debug("granularity is days")
        return return_strings[model][4]
    elif time_diff % S_PER_HR == 0:
        logger.debug("granularity is hours")
        return return_strings[model][5]
    else:
        logger.debug("granularity is seconds")
        return return_strings[model][6]

def discretize_time_difference(
    times, initial_time, frequency, integer_timestamps=False, zero_index=False
) -> typing.Sequence[int]:
    """method that discretizes sequence of datetimes (for prediction slices)

    Arguments:
        times {Sequence[datetime] or Sequence[float]} -- sequence of datetime objects
        initial_time {datetime or float} -- last datetime instance from training set
            (to offset test datetimes)
        frequency {str} -- string alias representing granularity of pd.datetime object

    Keyword Arguments:
        integer_timestamps {bool} -- whether timestamps are integers or datetime values
        zero_index {bool} -- whether to subtract 1 from each index to account for 0-indexing

    Returns:
        typing.Sequence[int] -- prediction intervals expressed at specific time granularity

    """

    # take differences to convert to deltas
    time_differences = times - initial_time

    # edge case for integer timestamps
    if integer_timestamps:
        time_differences = time_differences.values.astype(int)
        if zero_index:
            time_differences = [diff - 1 for diff in time_differences]
        return time_differences

    # convert to seconds representation
    if type(time_differences.iloc[0]) is pd._libs.tslibs.timedeltas.Timedelta:
        time_differences = time_differences.apply(lambda t: t.total_seconds())

    if frequency == "YS" or frequency == "12M":
        time_differences = [round(x / S_PER_YEAR) for x in time_differences]
    elif frequency == "MS" or frequency == "M":
        time_differences = [round(x / S_PER_MONTH) for x in time_differences]
        #     round(x / ((1/12 * S_PER_MONTH_28) * (1/3 * S_PER_MONTH_30) + ( * S_PER_MONTH_31)))
        #     for x in time_differences
        # ]
    elif frequency == "W" or frequency == "W-MON":
        time_differences = [round(x / S_PER_WEEK) for x in time_differences]
    elif frequency == "D":
        time_differences = [round(x / S_PER_DAY) for x in time_differences]
    elif frequency == "H":
        time_differences = [round(x / S_PER_HR) for x in time_differences]
    else:
        time_differences = [round(x / SECONDS_PER_MINUTE) for x in time_differences]

    if zero_index:
        time_differences = [diff - 1 for diff in time_differences]

    return time_differences

ext: py

reason: import numpy


text:
"""
retrieved from https://github.com/princewen/tensorflow_practice/blob/master/RL/Basic-MADDPG-Demo/three_agent_maddpg.py
modified the code to make it work with our environment
test the env
"""

import numpy as np
import tensorflow as tf
from MADDPG import *
from collections import deque
import os
import copy
from replay_buffer import *
from make_env import *
import matplotlib.pyplot as plt

tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

def create_init_update(online_name, target_name, tau=0.99):
    """
    This creates an init operation, which will initialize the weights of both online and target models of an agent to be
    equal. It also creates the update operation, which when called, will adjust the weights of the target model to be closer
    to that of the online model.
    :param online_name: The scope of the online model, so we can retrieve its trainable variables (weights)
    :param target_name: The scope of the target model
    :param tau: A parameter which specifies how much to adjust the target weights in the direction of the online weights
    :return:
    """

    online_var = [i for i in tf.trainable_variables() if online_name in i.name]

    target_var = [i for i in tf.trainable_variables() if target_name in i.name]

    target_init = [tf.assign(target, online) for online, target in zip(online_var, target_var)]

    target_update = [tf.assign(target, (1 - tau) * online + tau * target) for online, target in zip(online_var, target_var)]

    return target_init, target_update

def get_agents_action(obs, a1, a2, a3, sess, noise_rate=0.0):

    # debugged here
    obs = np.asarray(obs)

    o_1 = np.expand_dims(obs,axis=0)

    o_1 = np.reshape(o_1, [1, 18 * 3])

    act1 = a1.action(state=o_1, sess=sess)

    act2 = a2.action(state=o_1, sess=sess)

    act3 = a3.action(state=o_1, sess=sess)

    return act1, act2, act3

def train_agent(agent_ddpg, agent_ddpg_target, agent_memory, agent_actor_target_update, agent_critic_target_update, sess, other_actors, num_agents=3):
    """
    This is an important function, which runs a single train step for a single agent.
    :param agent_ddpg: The online part of agent which we will be training. This is the object which represents the deep network that we choose actions from
    :param agent_ddpg_target: The target part of the agent. We base our update values on this agent's output
    :param agent_memory: This is the agent's memory. It includes up to 2000 tuples of (all agents' obs, all agents' actions, agent's reward, all agents' next obs, done)
    :param agent_actor_target_update: The update operation for the actor network. Will update actor network's target weights to be more equal to the online weights
    :param agent_critic_target_update: The update operation for the critic network.
    :param sess: Session object to run tensorflow operations
    :param other_actors: A list of target models for the other agents in the environment.
    :return:
    """

    total_obs_batch, act_batch, other_act_batch, rew_batch, total_next_obs_batch, done_mask = agent_memory.sample(32)

    #EXPAND DIMS WHEN ACT_BATCH IS ACTION TAKEN
    #act_batch = np.expand_dims(act_batch, axis=1)

    obs_batch = np.squeeze(total_obs_batch)

    next_obs_batch = np.squeeze(total_next_obs_batch)

    other_act_batch = np.reshape(other_act_batch, [32, num_agents-1])

    #next_other_action = np.hstack([other_actors[0].action(next_other_actor1_o, sess), other_actors[1].action(next_other_actor2_o, sess)])
    next_act = agent_ddpg.action(next_obs_batch, sess)

    """
    next_a = np.asarray([np.random.choice(np.arange(5), p=next_act[i, :]) for i in range(batch_size)])
    next_a = np.expand_dims(next_a, axis=1)
    """

    #action used to equal agent_ddpg.action(next_obs_batch, sess)
    target = rew_batch.reshape(-1, 1) + 0.95 * agent_ddpg_target.Q(state=next_obs_batch, action=next_act,

                                                                     other_action=other_act_batch, sess=sess)

    agent_ddpg.train_actor(state=obs_batch, action=act_batch, other_action=other_act_batch, sess=sess)

    agent_ddpg.train_critic(state=obs_batch, action=act_batch, other_action=other_act_batch, target=target, sess=sess)

    #sess.run([agent_actor_target_update, agent_critic_target_update])

def load_weights(a1, a1_tar, name1, session, a2=None, a2_tar=None, name2=None, a3=None, a3_tar=None, name3=None):
    a1.load_weights(name1 + "_online", session)
    a1_tar.load_weights(name1 + "_target", session)

    if a2 is not None:
        a2.load_weights(name2 + "_online", session)
        a2_tar.load_weights(name2 + "_target", session)

    if a3 is not None:
        a3.load_weights(name3 + "_online", session)
        a3_tar.load_weights(name3 + "_target", session)

def save_weights(a1, a1_tar, a2, a2_tar, a3, a3_tar, name1, name2, name3, session, episode=None):
    a1.save_weights(name1 + "_online", session, episode)
    a1_tar.save_weights(name1 + "_target", session, episode)
    a2.save_weights(name2 + "_online", session, episode)
    a2_tar.save_weights(name2 + "_target", session, episode)
    a3.save_weights(name3 + "_online", session, episode)
    a3_tar.save_weights(name3 + "_target", session, episode)

def run():
    """
    Weights will save to Weights_save. After training, pick which episode
    you would like to retrieve the weights from. Sort the files in that directory by
    date modified, and then select all the weights with the ending corresponding to the episode.
    This should be around 24 files.
    Move those weights to Weights_final, and delete the ending from all the files (tedious I know).
    Weights might look like "weightname-1800.data", delete just the "-1800" part. Then, turn testing
    to true.
    """

    save_dir = "saves_training"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    env, scenario, world = make_env('simple_spread')

    agent1 = MADDPG('agent1', nb_actions=5, nb_input=54, nb_other_action=2)

    agent1_target = MADDPG('agent1_target', nb_actions=5, nb_input=54, nb_other_action=2)

    agent2 = MADDPG('agent2', nb_actions=5, nb_input=54, nb_other_action=2)

    agent2_target = MADDPG('agent2_target', nb_actions=5, nb_input=54, nb_other_action=2)

    agent3 = MADDPG('agent3', nb_actions=5, nb_input=54, nb_other_action=2)

    agent3_target = MADDPG('agent3_target', nb_actions=5, nb_input=54, nb_other_action=2)

    #saver = tf.train.Saver()

    agent1_actor_target_init, agent1_actor_target_update = create_init_update('agent1_actor', 'agent1_target_actor')

    agent1_critic_target_init, agent1_critic_target_update = create_init_update('agent1_critic', 'agent1_target_critic')

    agent2_actor_target_init, agent2_actor_target_update = create_init_update('agent2_actor', 'agent2_target_actor')

    agent2_critic_target_init, agent2_critic_target_update = create_init_update('agent2_critic', 'agent2_target_critic')

    agent3_actor_target_init, agent3_actor_target_update = create_init_update('agen3_actor', 'agent3_target_actor')

    agent3_critic_target_init, agent3_critic_target_update = create_init_update('agent3_critic', 'agent3_target_critic')

    sess = tf.Session()

    sess.run(tf.global_variables_initializer())

    sess.run([agent1_actor_target_init, agent1_critic_target_init,

              agent2_actor_target_init, agent2_critic_target_init,

              agent3_actor_target_init, agent3_critic_target_init])

    save_weight_dir = "Weights_save"
    load_weight_dir = "Weights_final"
    if not os.path.exists(save_weight_dir):
        os.makedirs(save_weight_dir)

    if not os.path.exists(load_weight_dir):
        os.makedirs(load_weight_dir)

    weights_fname = load_weight_dir + "/weights_new"
    testing = True
    if testing:
        load_weights(agent1, agent1_target, weights_fname + "_1", sess, agent2, agent2_target, weights_fname + "_2", agent3, agent3_target, weights_fname + "_3")

    weights_fname = save_weight_dir + "/weights_new"

    num_episodes = 10

    rewards = []
    average_over = int(num_episodes / 10)
    average_rewards = []
    average_r = deque(maxlen=average_over)

    agent1_memory = ReplayBuffer(10**6) #This was at 100

    agent2_memory = ReplayBuffer(10**6)

    agent3_memory = ReplayBuffer(10**6)

    # e = 1

    batch_size = 32

    num_steps = 50

    transition = []

    num_samples = 0

    for i in range(num_episodes):
    # make graph with reward

        t_rew = 0
        t_collisions = 0
        t_min_dist = 0
        t_occ_landmarks = 0
        a_rew = 0
        a_collisions = 0
        a_min_dist = 0
        a_occ_landmarks = 0

        #Reset the environment at the start of each episode
        o_n = env.reset()

        total_ep_reward = 0

        steps = 0
        """
        render = False
        if i % 10 == 0:
            render = True
        """
        render = True

        for t in range(num_steps):

            if render:
                env.render()

            o_n = np.reshape(o_n, [1, 54])

            a1_action, a2_action, a3_action = get_agents_action(o_n, agent1, agent2, agent3, sess, noise_rate=0.2)

            a1_action = np.squeeze(a1_action)
            a2_action = np.squeeze(a2_action)
            a3_action = np.squeeze(a3_action)

            #Sample from probabilities
            action1 = np.random.choice(np.arange(len(a1_action)), p=a1_action)
            action2 = np.random.choice(np.arange(len(a2_action)), p=a2_action)
            action3 = np.random.choice(np.arange(len(a3_action)), p=a3_action)

            a_n = [action1, action2, action3]

            #print("action of agent is " + str(a))
            # global reward as the reward for each agent

            #Get global reward

            o_n_next, r_n, done_n, _ = env.step(a_n)

            #r_n is actually the individual reward of each agent multiplied by 3

            #print("Reward: " + str(r_n))
            total_ep_reward += r_n[0]

            done = True
            for m in range(3):
                if not done_n[m]:
                    done = False

            o_n = np.asarray(o_n)
            o_n_next = np.asarray(o_n_next)
            o_n_next = np.reshape(o_n_next, [1, 54])

            # WITH ACTION TAKEN FOR OTHER ACTIONS, BUT PROBS FOR THE MAIN AGENT
            agent1_memory.add(o_n, a1_action, np.vstack([action2, action3]), r_n[0]/3,
                              o_n_next, done_n[0])

            agent2_memory.add(o_n, a2_action, np.vstack([action1, action3]), r_n[1]/3,
                              o_n_next, done_n[1])

            agent3_memory.add(o_n, a3_action, np.vstack([action1, action2]), r_n[2]/3,
                              o_n_next, done_n[2])

            num_samples += 1

            if done:
                break

            #print((o_n, action1, action2, o_n_next))
            transition.append((o_n, action1, action2, action3, o_n_next))

            #print(o_n.shape)
            #print(o_n_next.shape)
            #Add to agents' memories the state, actions, reward, next state, done tuples

            """
            ORIGINAL
            agent1_memory.add(np.vstack([o_n, o_n, o_n]), np.vstack([action1, action2, action3]), r_n[0], np.vstack([o_n_next, o_n_next, o_n_next]), False)

            agent2_memory.add(np.vstack([o_n, o_n, o_n]), np.vstack([action2, action1, action3]), r_n[1], np.vstack([o_n_next, o_n_next, o_n_next]), False)

            agent3_memory.add(np.vstack([o_n, o_n, o_n]), np.vstack([action3, action1, action2]), r_n[1],
                              np.vstack([o_n_next, o_n_next, o_n_next]), False)
            """

            """
            #WITH ACTION AS INTEGER ACT TAKEN
            agent1_memory.add(o_n, action1, np.vstack([action2, action3]), r_n[0],
                              o_n_next, False)

            agent2_memory.add(o_n, action2, np.vstack([action1, action3]), r_n[1],
                              o_n_next, False)

            agent3_memory.add(o_n, action3, np.vstack([action1, action2]), r_n[1],
                              o_n_next, False)
            """

            """
            #WITH ACTION PROBABILITIES FOR EACH AGENT
            agent1_memory.add(o_n, a1_action, np.vstack([a2_action, a3_action]), r_n[0],
                              o_n_next, False)

            agent2_memory.add(o_n, a2_action, np.vstack([a1_action, a3_action]), r_n[1],
                              o_n_next, False)

            agent3_memory.add(o_n, a3_action, np.vstack([a1_action, a2_action]), r_n[1],
                              o_n_next, False)
            """

            rew, collisions, min_dist, occ_landmarks = scenario.benchmark_data(world.agents[0], world)

            t_rew += rew
            t_collisions += collisions - 1
            t_min_dist += min_dist
            t_occ_landmarks += occ_landmarks
            a_rew = t_rew / (t+1)
            a_min_dist = t_min_dist / (t+1)
            a_occ_landmarks = t_occ_landmarks / (t+1)

            # original is 50000
            if agent1_memory.__len__() > 32 and not testing and num_samples % 100 == 0:

                # e *= 0.9999

                # agent1 train

                #Run a single train step for each agent
                train_agent(agent1, agent1_target, agent1_memory, agent1_actor_target_update,
                            agent1_critic_target_update, sess, [agent2_target, agent3_target])

                train_agent(agent2, agent2_target, agent2_memory, agent2_actor_target_update,
                            agent2_critic_target_update, sess, [agent1_target, agent3_target])

                train_agent(agent3, agent3_target, agent3_memory, agent3_actor_target_update,
                            agent3_critic_target_update, sess, [agent1_target, agent2_target])

                #print("step " + str(t))
                #print("observation by agents are  " + str(env.get_obs()) + "\n\n")
                #print("fire level at step " + str(t) + " is " + str(env.firelevel))

            o_n = copy.copy(o_n_next)

        if i % 1024 == 0:
            sess.run([agent1_actor_target_update, agent1_critic_target_update])
            sess.run([agent2_actor_target_update, agent2_critic_target_update])
            sess.run([agent3_actor_target_update, agent3_critic_target_update])

        print("Episode: {}. Global Reward: {}.".format(i+1, total_ep_reward))
        print("Avg Rew per t: {}. Episode Collisions: {}. Avg min dist per t: {}. Avg landmarks occ per t: {}".format(a_rew, t_collisions, a_min_dist, a_occ_landmarks))

        rewards.append(total_ep_reward)
        average_r.append(total_ep_reward)

        if i < average_over:
            r = 0
            for j in range(i):
                r += average_r[j]
            r /= (i + 1)
            average_rewards.append(r)
        else:
            average_rewards.append(sum(average_r) / average_over)

        if i % average_over == 0:
            plt.plot(rewards)
            plt.xlabel('Episode')
            plt.ylabel('Reward')
            plt.savefig(save_dir + "/reward.png")
            plt.clf()

            plt.plot(average_rewards)
            plt.xlabel('Episode')
            plt.ylabel('Moving average')
            plt.savefig(save_dir + "/moving_avg.png")
            plt.clf()

            save_weights(agent1, agent1_target, agent2, agent2_target, agent3, agent3_target,
                         weights_fname + "_1", weights_fname + "_2", weights_fname + "_3",
                         sess, i)

    """
    transition = np.asarray(transition)
    print(transition.shape)
    np.save('Transition', transition)
    """

    sess.close()


text:
"""
This is the GPU implementation of regular CPAC-CNN model using pytorch.
python train_mnist.py [number of filters] [filter_size] [rank] [epochs] [device]; python train_mnist.py 8 3 6 10 cuda:0
"""

import torch
import torchvision
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable as V
import numpy as np
np.random.seed(48)
torch.manual_seed(48)
from conv_decomp_torch import Conv_Decomp
from cnn_torch import CNN_Decomp
from cnn_torch import CNN

import time
import sys
import os
import warnings
warnings.filterwarnings("ignore")

##MNIST
##get parameters from command line
num_filters = int(sys.argv[1])
filter_h = int(sys.argv[2])
filter_w = int(sys.argv[2])
rank = int(sys.argv[3])
n_epochs = int(sys.argv[4])
cuda_device = str(sys.argv[5])

image_channels = 1
num_class = 10
devices = [torch.device(cuda_device),torch.device("cpu")]
learning_rate = 0.0001
batch_size_train = 1
batch_size_test = 1
log_interval = 100

##data loader
def load_dataset():
    data_path = 'data/'

    train_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST('data/', train=True, download=True,
                                   transform=torchvision.transforms.Compose([
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize(
                                           (0.5,), (0.5,))
                                   ])),
        batch_size=batch_size_train, shuffle=True)

    test_loader = torch.utils.data.DataLoader(
        torchvision.datasets.MNIST('data/', train=False, download=True,
                                   transform=torchvision.transforms.Compose([
                                       torchvision.transforms.ToTensor(),
                                       torchvision.transforms.Normalize(
                                           (0.5,), (0.5,))
                                   ])),
        batch_size=batch_size_test, shuffle=True)

    return train_loader, test_loader

##model training
def train(epoch,file_root,file_name):
    Net.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = Net(data[0,0:1,:,:]) ##the reshape process need to be done on CPU
        #output = Net(data[:,0:1,:,:].to(devices[0]))
        loss = F.nll_loss(output.to(devices[1]), target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            log = 'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                   epoch, batch_idx * len(data), len(train_loader.dataset),
                   100. * batch_idx / len(train_loader), loss.item())
            print(log)
            train_losses.append(loss.item())
            train_counter.append(
               (batch_idx*1) + ((epoch-1)*len(train_loader.dataset)))

##model testing
def test(test_acc,file_root,file_name):
    Net.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
      for data, target in test_loader:
          output = Net(data[0,0:1,:,:]) ##the reshape process need to be done on CPU
          test_loss += F.nll_loss(output.to(devices[1]), target, size_average=False).item()
          pred = output.data.max(1, keepdim=True)[1].to(devices[1])
          correct += pred.eq(target.data.view_as(pred)).sum()
    test_loss /= len(test_loader.dataset)
    test_losses.append(test_loss)
    log = '\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
          test_loss, correct, len(test_loader.dataset),
          100. * correct / len(test_loader.dataset))
    with open(file_root+'/'+file_name, 'a') as f:
          f.write(log+'\n')
    print(log)
    ##save the model with best testing accuracy
    if 100. * correct / len(test_loader.dataset) > test_acc:
        test_acc = 100. * correct / len(test_loader.dataset)
        torch.save(Net.state_dict(), file_root + '/' + 'model.pth')
        torch.save(optimizer.state_dict(), file_root + '/' + 'optimizer.pth')
    return test_acc

##result directory
par_root = "./result"
file_root = par_root + '/' + 'mnist_2layer_'+str(num_filters) + '_' + str(rank) + '_' + str(n_epochs) + '_' + time.strftime("%Y%m%d-%H%M%S")
os.mkdir(file_root)
file_name = 'mnist_2layer_'+str(num_filters) + '_' + str(rank) + '_' + str(n_epochs) + '.txt'

##initialize data loader
train_loader, test_loader = load_dataset()

##get image size
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
input_shape = example_data.shape

##initialize model CPAC-CNN, the detailed model structure can be modified in cnn_torch.py, currently there are two CPAC-conv layers
Net = CNN_Decomp(num_filters, filter_h, filter_w, image_channels, rank, devices, num_class, input_shape)
#Net = CNN(num_filters, filter_h, filter_w, image_channels, rank, devices, num_class, input_shape)
Net.to(devices[0])
optimizer = optim.Adam(Net.parameters(), lr=learning_rate)

train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

test_acc = 0
##train and test
for epoch in range(1, n_epochs + 1):
  train(epoch,file_root,file_name)
  test_acc = test(test_acc,file_root,file_name)


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import math
import matplotlib.patheffects as path_effects
import cv2

from matplotlib import pyplot as plt
from matplotlib.patches import Ellipse
from matplotlib.pyplot import arrow
from matplotlib.collections import PatchCollection, LineCollection
from matplotlib import colors

from scipy.signal import savgol_filter
from scipy.spatial import Voronoi

from shapely.geometry import Polygon

from moviepy import editor as mpy
from moviepy.video.io.bindings import mplfig_to_npimage

"""

Football field vizualization function cloned from https://github.com/Friends-of-Tracking-Data-FoTD
with some minor modifications to:
    * add velocity vectors
    * ball and possession markers

"""

X_SIZE = 105
Y_SIZE = 68

BOX_HEIGHT = (16.5 * 2 + 7.32) / Y_SIZE * 100
BOX_WIDTH = 16.5 / X_SIZE * 100

GOAL = 7.32 / Y_SIZE * 100

GOAL_AREA_HEIGHT = 5.4864 * 2 / Y_SIZE * 100 + GOAL
GOAL_AREA_WIDTH = 5.4864 / X_SIZE * 100

SCALERS = np.array([X_SIZE / 100, Y_SIZE / 100])
pitch_polygon = Polygon(((0, 0), (0, 100), (100, 100), (100, 0)))

def visualize(**images):
    """PLot images in one row.

    Arguments:
        **images: images to plot
    Returns:

    Raises:

    """
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(" ".join(name.split("_")).title())
        plt.imshow(image)
    plt.show()

def draw_pitch(dpi=100, pitch_color="#a8bc95"):
    """Sets up field.
    Arguments:
      dpi: Dots per inch in the field
      pitch_color: Color of the field
    Returns:
      fig,axes: matplotlib fig and axes objects.
    Raises:
    """
    fig = plt.figure(figsize=(12.8, 7.2), dpi=dpi)
    fig.patch.set_facecolor(pitch_color)

    axes = fig.add_subplot(1, 1, 1)
    axes.set_axis_off()
    axes.set_facecolor(pitch_color)
    axes.xaxis.set_visible(False)
    axes.yaxis.set_visible(False)

    axes.set_xlim(0, 100)
    axes.set_ylim(0, 100)

    plt.xlim([-13.32, 113.32])
    plt.ylim([-5, 105])

    fig.tight_layout(pad=3)

    draw_patches(axes)

    return fig, axes

def draw_patches(axes):
    """Draws basic field shapes on an axes
    Arguments:
      axes: matplotlib axes objects.
    Returns:
      axes: matplotlib axes objects.
    Raises:
    """
    # pitch
    axes.add_patch(plt.Rectangle((0, 0), 100, 100, edgecolor="white", facecolor="none"))

    # half-way line
    axes.add_line(plt.Line2D([50, 50], [100, 0], c="w"))

    # penalty areas
    axes.add_patch(
        plt.Rectangle(
            (100 - BOX_WIDTH, (100 - BOX_HEIGHT) / 2),
            BOX_WIDTH,
            BOX_HEIGHT,
            ec="w",
            fc="none",
        )
    )
    axes.add_patch(
        plt.Rectangle(
            (0, (100 - BOX_HEIGHT) / 2), BOX_WIDTH, BOX_HEIGHT, ec="w", fc="none"
        )
    )

    # goal areas
    axes.add_patch(
        plt.Rectangle(
            (100 - GOAL_AREA_WIDTH, (100 - GOAL_AREA_HEIGHT) / 2),
            GOAL_AREA_WIDTH,
            GOAL_AREA_HEIGHT,
            ec="w",
            fc="none",
        )
    )
    axes.add_patch(
        plt.Rectangle(
            (0, (100 - GOAL_AREA_HEIGHT) / 2),
            GOAL_AREA_WIDTH,
            GOAL_AREA_HEIGHT,
            ec="w",
            fc="none",
        )
    )

    # goals
    axes.add_patch(plt.Rectangle((100, (100 - GOAL) / 2), 1, GOAL, ec="w", fc="none"))
    axes.add_patch(plt.Rectangle((0, (100 - GOAL) / 2), -1, GOAL, ec="w", fc="none"))

    # halfway circle
    axes.add_patch(
        Ellipse(
            (50, 50),
            2 * 9.15 / X_SIZE * 100,
            2 * 9.15 / Y_SIZE * 100,
            ec="w",
            fc="none",
        )
    )

    return axes

def draw_frame(
    df,
    t,
    dpi=100,
    fps=20,
    add_vector=False,
    display_num=False,
    display_time=False,
    show_players=True,
    highlight_color=None,
    highlight_player=None,
    shadow_player=None,
    text_color="white",
    flip=False,
    **anim_args
):
    """
    Draws players from time t (in seconds) from a DataFrame df
    """
    fig, ax = draw_pitch(dpi=dpi)

    dfFrame = get_frame(df, t, fps=fps)

    if show_players:
        for pid in dfFrame.index:
            if pid == 0:
                # se for bola
                try:
                    z = dfFrame.loc[pid]["z"]
                except:
                    z = 0
                size = 1.2 + z
                lw = 0.9
                color = "black"
                edge = "white"
                zorder = 100
            else:
                # se for jogador
                size = 3
                lw = 2
                edge = dfFrame.loc[pid]["edgecolor"]

                if pid == highlight_player:
                    color = highlight_color
                else:
                    color = dfFrame.loc[pid]["bgcolor"]
                if dfFrame.loc[pid]["team"] == "attack":
                    zorder = 21
                else:
                    zorder = 20

            ax.add_artist(
                Ellipse(
                    (dfFrame.loc[pid]["x"], dfFrame.loc[pid]["y"]),
                    size / X_SIZE * 100,
                    size / Y_SIZE * 100,
                    edgecolor=edge,
                    linewidth=lw,
                    facecolor=color,
                    alpha=0.8,
                    zorder=zorder,
                )
            )

            if add_vector:

                arrow_length = (
                    math.sqrt(
                        (dfFrame.loc[pid]["dx"] ** 2 + dfFrame.loc[pid]["dy"] ** 2)
                    )
                    * 20
                )
                color_arrow = "white" if pid == 0 else color
                plt.arrow(
                    x=dfFrame.loc[pid]["x"],
                    y=dfFrame.loc[pid]["y"],
                    dx=dfFrame.loc[pid]["dx"] * 20,
                    dy=dfFrame.loc[pid]["dy"] * 20,
                    length_includes_head=True,
                    color=color_arrow,
                    edgecolor=edge,
                    head_width=1,
                    head_length=arrow_length / 4.0,
                )

            try:
                s = str(int(dfFrame.loc[pid]["player_num"]))
            except ValueError:
                s = ""
            text = plt.text(
                dfFrame.loc[pid]["x"],
                dfFrame.loc[pid]["y"],
                s,
                horizontalalignment="center",
                verticalalignment="center",
                fontsize=8,
                color=text_color,
                zorder=22,
                alpha=0.8,
            )

            text.set_path_effects(
                [
                    path_effects.Stroke(linewidth=1, foreground=text_color, alpha=0.8),
                    path_effects.Normal(),
                ]
            )

    return fig, ax, dfFrame

def add_voronoi_to_fig(fig, ax, dfFrame):
    """ Adds a voronoi diagram to the field, according to the players positions
    Arguments:
      fig,ax: matplotlib fig and axes objects.
      dfFrame: pd.DataFrame with player tracking data
    Returns:
      fig,ax: matplotlib fig and axes objects.
      dfFrame: pd.DataFrame with player tracking data
    Raises:
    """
    polygons = {}
    vor, dfVor = calculate_voronoi(dfFrame)
    for index, region in enumerate(vor.regions):
        if not -1 in region:
            if len(region) > 0:
                try:
                    pl = dfVor[dfVor["region"] == index]
                    polygon = Polygon(
                        [vor.vertices[i] for i in region] / SCALERS
                    ).intersection(pitch_polygon)
                    polygons[pl.index[0]] = polygon
                    color = pl["bgcolor"].values[0]
                    x, y = polygon.exterior.xy
                    plt.fill(x, y, c=color, alpha=0.30)
                except IndexError:
                    pass
                except AttributeError:
                    pass

    plt.scatter(dfVor["x"], dfVor["y"], c=dfVor["bgcolor"], alpha=0.2)

    return fig, ax, dfFrame

def calculate_voronoi(dfFrame):
    """ Computes the voronoi diagram for the players positions
    Arguments:
      dfFrame: pd.DataFrame with player tracking data
    Returns:
      vor: Voronoi dataframe (region for each coordinates)
      dfFrame: pd.DataFrame with player tracking data
    Raises:
    """
    dfTemp = dfFrame.copy().drop(0, errors="ignore")

    values = np.vstack(
        (
            dfTemp[["x", "y"]].values * SCALERS,
            [-1000, -1000],
            [+1000, +1000],
            [+1000, -1000],
            [-1000, +1000],
        )
    )

    vor = Voronoi(values)

    dfTemp["region"] = vor.point_region[:-4]

    return vor, dfTemp

def get_frame(df, t, fps=20):
    """Gets the player data from the right frame
    Arguments:
      df: pd.DataFrame with player tracking data
      t: timestamp of the play
      fps: frame per second
    Returns:
      dfFrame: pd.DataFrame with player tracking data from timestamp t
    Raises:
    """
    dfFrame = df.loc[int(t * fps)].set_index("player")
    dfFrame.player_num = dfFrame.player_num.fillna("")
    return dfFrame

def draw_frame_x(df, t, fps, voronoi=True):
    """Draw field, player and voronoi on the same image
    Arguments:
      df: pd.DataFrame with player tracking data
      t: timestamp of the play
      fps: frame per second
      voronoi : If we draw the voronoi diagram or not
    Returns:
      image: Image of the field, players, ball, and eventually voronoi diagram.
    Raises:
    """
    fig, ax, dfFrame = draw_frame(df, t=t, fps=fps, add_vector=True)
    if voronoi:
        fig, ax, dfFrame = add_voronoi_to_fig(fig, ax, dfFrame)
    image = mplfig_to_npimage(fig)
    plt.close()
    return image

def make_animation(df, fps=20, voronoi=True):
    """Makes a clip from the entire dataset
    Arguments:
      df: pd.DataFrame with player tracking data
      fps: frame per second
      voronoi : If we draw the voronoi diagram or not
    Returns:
      clip: mpy Clip object
    Raises:
    """
    # calculated variables
    length = (df.index.max() + 20) / fps
    clip = mpy.VideoClip(
        lambda x: draw_frame_x(df, t=x, fps=fps, voronoi=voronoi), duration=length - 1
    ).set_fps(fps)
    return clip

def draw_line(df_value, t, fps, smooth=True, show=True):
    """ Draw the value function overtime, and add a marker at the wanted frame
    Arguments:
      df_value: pd.Dataframe with the value of each frame
      t: timestamp of the play to mark the value function
      fps: frame per second
      smooth : If we smooth the value function or not
      show: If we show the plot (otherwise return it as an image)
    Returns:
      image, fig,ax: show or image of the plot
    Raises:
    """
    fig = plt.figure(figsize=(19.2, 10.8), dpi=100)
    ax = fig.add_subplot(111)
    if smooth:
        df_value["Value_smooth"] = savgol_filter(df_value["value"], 21, 3)
    else:
        df_value["Value_smooth"] = df_value["value"]
    x = df_value["frame_count"].values
    y = df_value["Value_smooth"].values

    vmin = df_value["Value_smooth"].min()
    vmax = df_value["Value_smooth"].max()
    norm = colors.Normalize(vmin=vmin, vmax=vmax)
    points = np.array([x, y]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    lc = LineCollection(segments, cmap="cool", norm=norm)
    lc.set_array(y)
    lc.set_linewidth(2)
    line = ax.add_collection(lc)
    fig.colorbar(line, ax=ax)
    ax.scatter(
        df_value.values[int(t * fps)][0],
        df_value.values[int(t * fps)][2],
        color="green",
        marker="o",
        linewidths=4,
    )
    if show:
        return fig, ax
    else:
        image = mplfig_to_npimage(fig)
        plt.close()
        return image

def make_animation_value(df, fps=20):
    """Makes a clip from the entire dataset of the value function
    Arguments:
      df: pd.DataFrame with player tracking data
      fps: frame per second
    Returns:
      clip: mpy Clip object
    Raises:
    """
    # calculated variables
    length = (df.index.max() + 20) / fps
    clip = mpy.VideoClip(
        lambda x: draw_line(df, t=x, fps=fps), duration=length - 1
    ).set_fps(fps)
    return clip

def add_edg_to_fig(fig, ax, edg_map, vmin=None, vmax=None):
    """Adds an edg_map to a field
    Arguments:
      fig,ax : Matplotlib object from draw_frame
      edg_map: edg_map from agent.py
      vmin,vmax : min and max value of the edg_map
    Returns:
      fig,ax: Matplotlib object from draw_frame, with the map on top
      edg_map: edg_map from agent.py
    Raises:
    """
    cmap = "bwr"
    if vmin is None:
        vmin = np.min(edg_map)
    if vmax is None:
        vmax = np.max(edg_map)
    ax = ax.imshow(
        edg_map,
        extent=(0, 100, 0, 100),
        aspect="auto",
        interpolation="mitchell",
        vmin=vmin,
        vmax=vmax,
        cmap=cmap,
        alpha=0.5,
    )
    cbar = plt.colorbar(ax)
    cbar.set_label("Value")
    return fig, ax, edg_map

def get_color(idx):
    idx = idx * 3
    color = ((17 * idx) % 255, (37 * idx) % 255, (29 * idx) % 255)

    return color

def plot_tracking(image, tlwhs, obj_ids, scores=None, frame_id=0, fps=0.0, ids2=None):
    im = np.ascontiguousarray(np.copy(image))
    im_h, im_w = im.shape[:2]

    text_scale = max(1, image.shape[1] / 1600.0)
    text_thickness = 1 if text_scale > 1.1 else 1
    line_thickness = max(1, int(image.shape[1] / 500.0))

    radius = max(5, int(im_w / 140.0))
    cv2.putText(
        im,
        "frame: %d fps: %.2f num: %d" % (frame_id, fps, len(tlwhs)),
        (0, int(15 * text_scale)),
        cv2.FONT_HERSHEY_PLAIN,
        text_scale,
        (0, 0, 255),
        thickness=2,
    )

    for i, tlwh in enumerate(tlwhs):
        x1, y1, w, h = tlwh
        intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h)))
        obj_id = int(obj_ids[i])
        id_text = "{}".format(int(obj_id))
        if ids2 is not None:
            id_text = id_text + ", {}".format(int(ids2[i]))
        _line_thickness = 1 if obj_id <= 0 else line_thickness
        color = get_color(abs(obj_id))
        cv2.rectangle(
            im, intbox[0:2], intbox[2:4], color=color, thickness=line_thickness
        )
        cv2.putText(
            im,
            id_text,
            (intbox[0], intbox[1] + 30),
            cv2.FONT_HERSHEY_PLAIN,
            text_scale,
            (0, 0, 255),
            thickness=text_thickness,
        )
    return im

def rgb_template_to_coord_conv_template(rgb_template):
    assert isinstance(rgb_template, np.ndarray)
    assert rgb_template.min() >= 0.0
    assert rgb_template.max() <= 1.0
    rgb_template = np.mean(rgb_template, 2)
    x_coord, y_coord = np.meshgrid(
        np.linspace(0, 1, num=rgb_template.shape[1]),
        np.linspace(0, 1, num=rgb_template.shape[0]),
    )
    coord_conv_template = np.stack((rgb_template, x_coord, y_coord), axis=2)
    return coord_conv_template

def merge_template(img, warped_template):
    valid_index = warped_template[:, :, 0] > 0.0
    overlay = (
        img[valid_index].astype("float32")
        + warped_template[valid_index].astype("float32")
    ) / 2
    new_image = np.copy(img)
    new_image[valid_index] = overlay
    return new_image

max_stars_repo_name: JefeDryden/narya

max_issues_repo_name: JefeDryden/narya

max_forks_repo_name: JefeDryden/narya

reason: import numpy,from scipy