text:
#!/usr/bin/env python3
# Copyright 2019 Alexander Meulemans
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
In here, we define classes for fully connected layers in a multilayer perceptron
network that will be trained by difference target propagation.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tensorboardX import SummaryWriter
import lib.utils as utils
import warnings

class DTPLayer(nn.Module):
    """ An abstract base class for a layer of an MLP that will be trained by the
    differece target propagation method. Child classes should specify which
    activation function is used.

    Attributes:
        weights (torch.Tensor): The forward weight matrix :math:`W` of the layer
        bias (torch.Tensor): The forward bias :math:`\mathbf{b}`
            of the layer.
            Attribute is ``None`` if argument ``bias`` was passed as ``None``
            in the constructor.
        feedback_weights (torch.Tensor): The feedback weight matrix :math:`Q`
            of the layer. Warning: if we use the notation of the theoretical
            framework, the feedback weights are actually $Q_{i-1}$ from the
            previous layer!! We do this because this makes the implementation
            of the reconstruction loss and training the feedback weights much
            easier (as g_{i-1} and hence Q_{i-1} needs to approximate
            f_i^{-1}). However for the direct feedback connection layers, it
            might be more logical to let the feedbackweights represent Q_i
            instead of Q_{i-1}, as now only direct feedback connections exist.
        feedback_bias (torch.Tensor): The feedback bias :math:`\mathbf{b}`
            of the layer.
            Attribute is ``None`` if argument ``bias`` was passed as ``None``
            in the constructor.
        forward_requires_grad (bool): Flag indicating whether the computational
            graph with respect to the forward parameters should be saved. This
            flag should be True if you want to compute BP or GN updates. For
            TP updates, computational graphs are not needed (custom
            implementation by ourselves)
        reconstruction_loss (float): The reconstruction loss of this layer
            evaluated at the current mini-batch.
        forward_activation (str): String indicating the forward nonlinear
            activation function used by the layer. Choices: 'tanh', 'relu',
            'linear'.
        feedback_activation (str): String indicating the feedback nonlinear
            activation function used by the layer. Choices: 'tanh', 'relu',
            'linear'.
        target (torch.Tensor or None): The target for this layer on the current
            minibatch. During normal training, it is not needed
            to save the targets so this attribute will stay None. If the user
            wants to compute the angle between (target - activation) and a
            BP update or GN update, the target needs to be saved in the layer
            object to use later on to compute the angles. The saving happens in
            the backward method of the network.

    Args:
        in_features (int): Size of each input sample.
        out_features (int): Size of each output sample.
        bias (bool): If set to ``False``, the layer will not learn an additive
            bias.
        forward_requires_grad (bool): Flag indicating whether the forward
            parameters require gradients that can be computed with autograd.
            This might be needed when comparing the DTP updates with BP updates
            and GN updates.
        forward_activation (str): String indicating the forward nonlinear
            activation function used by the layer. Choices: 'tanh', 'relu',
            'linear'.
        feedback_activation (str): String indicating the feedback nonlinear
            activation function used by the layer. Choices: 'tanh', 'relu',
            'linear'.
    """

    def __init__(self, in_features, out_features, bias=True,
                 forward_requires_grad=False, forward_activation='tanh',
                 feedback_activation='tanh', initialization='orthogonal'):
        nn.Module.__init__(self)

        self._weights = nn.Parameter(torch.Tensor(out_features, in_features),
                                     requires_grad=forward_requires_grad)
        self._feedbackweights = nn.Parameter(torch.Tensor(in_features,
                                                          out_features),
                                             requires_grad=False)
        if bias:
            self._bias = nn.Parameter(torch.Tensor(out_features),
                                      requires_grad=forward_requires_grad)
            self._feedbackbias = nn.Parameter(torch.Tensor(in_features),
                                              requires_grad=False)
        else:
            self._bias = None
            self._feedbackbias = None

        # Initialize the weight matrices following Lee 2015
        # TODO: try other initializations, such as the special one to mimic
        # batchnorm
        if initialization == 'orthogonal':
            gain = np.sqrt(6. / (in_features + out_features))
            nn.init.orthogonal_(self._weights, gain=gain)
            nn.init.orthogonal_(self._feedbackweights, gain=gain)
        elif initialization == 'xavier':
            nn.init.xavier_uniform_(self._weights)
            nn.init.xavier_uniform_(self._feedbackweights)
        elif initialization == 'xavier_normal':
            nn.init.xavier_normal_(self._weights)
            nn.init.xavier_normal_(self._feedbackweights)
        elif initialization == 'teacher':
            nn.init.xavier_normal_(self._weights, gain=3.)
            nn.init.xavier_normal_(self._feedbackweights)
        else:
            raise ValueError('Provided weight initialization "{}" is not '
                             'supported.'.format(initialization))

        if bias:
            nn.init.constant_(self._bias, 0)
            nn.init.constant_(self._feedbackbias, 0)

        self._activations = None
        self._linearactivations = None
        self._reconstruction_loss = None
        self._forward_activation = forward_activation
        self._feedback_activation = feedback_activation
        self._target = None

    @property
    def weights(self):
        """Getter for read-only attribute :attr:`weights`."""
        return self._weights

    @property
    def bias(self):
        """Getter for read-only attribute :attr:`bias`."""
        return self._bias

    @property
    def feedbackweights(self):
        """Getter for read-only attribute :attr:`feedbackweights`."""
        return self._feedbackweights

    @property
    def feedbackbias(self):
        """Getter for read-only attribute :attr:`feedbackbias`."""
        return self._feedbackbias

    @property
    def activations(self):
        """Getter for read-only attribute :attr:`activations` """
        return self._activations

    @activations.setter
    def activations(self, value):
        """ Setter for the attribute activations"""
        self._activations = value

    @property
    def linearactivations(self):
        """Getter for read-only attribute :attr:`linearactivations` """
        return self._linearactivations

    @linearactivations.setter
    def linearactivations(self, value):
        """Setter for the attribute :attr:`linearactivations` """
        self._linearactivations = value

    @property
    def reconstruction_loss(self):
        """ Getter for attribute reconstruction_loss."""
        return self._reconstruction_loss

    @reconstruction_loss.setter
    def reconstruction_loss(self, value):
        """ Setter for attribute reconstruction_loss."""
        self._reconstruction_loss = value

    @property
    def forward_activation(self):
        """ Getter for read-only attribute forward_activation"""
        return self._forward_activation

    @property
    def feedback_activation(self):
        """ Getter for read-only attribute feedback_activation"""
        return self._feedback_activation

    @property
    def target(self):
        """ Getter for attribute target"""
        return self._target

    @target.setter
    def target(self, value):
        """ Setter for attribute target"""
        self._target = value

    def get_forward_parameter_list(self):
        """ Return forward weights and forward bias if applicable"""
        parameterlist = []
        parameterlist.append(self.weights)
        if self.bias is not None:
            parameterlist.append(self.bias)
        return parameterlist

    def forward_activationfunction(self, x):
        """ Element-wise forward activation function"""
        if self.forward_activation == 'tanh':
            return torch.tanh(x)
        elif self.forward_activation == 'relu':
            return F.relu(x)
        elif self.forward_activation == 'linear':
            return x
        elif self.forward_activation == 'leakyrelu':
            return F.leaky_relu(x, 0.2)
        elif self.forward_activation == 'sigmoid':
            return torch.sigmoid(x)
        else:
            raise ValueError('The provided forward activation {} is not '
                             'supported'.format(self.forward_activation))

    def feedback_activationfunction(self, x):
        """ Element-wise feedback activation function"""
        if self.feedback_activation == 'tanh':
            return torch.tanh(x)
        elif self.feedback_activation == 'relu':
            return F.relu(x)
        elif self.feedback_activation == 'linear':
            return x
        elif self.feedback_activation == 'leakyrelu':
            return F.leaky_relu(x, 5)
        elif self.feedback_activation == 'sigmoid':
            if torch.sum(x < 1e-12) > 0 or torch.sum(x > 1-1e-12) > 0:
                warnings.warn('Input to inverse sigmoid is out of'
                                 'bound: x={}'.format(x))
            inverse_sigmoid = torch.log(x/(1-x))
            if utils.contains_nan(inverse_sigmoid):
                raise ValueError('inverse sigmoid function outputted a NaN')
            return torch.log(x/(1-x))
        else:
            raise ValueError('The provided feedback activation {} is not '
                             'supported'.format(self.feedback_activation))

    def compute_vectorized_jacobian(self, a):
        """ Compute the vectorized Jacobian of the forward activation function,
        evaluated at a. The vectorized Jacobian is the vector with the diagonal
        elements of the real Jacobian, as it is a diagonal matrix for element-
        wise functions. As a is a minibatch, the output will also be a
        mini-batch of vectorized Jacobians (thus a matrix).
        Args:
            a (torch.Tensor): linear activations
        """
        if self.forward_activation == 'tanh':
            return 1. - torch.tanh(a)**2
        elif self.forward_activation == 'relu':
            J = torch.ones_like(a)
            J[a < 0.] = 0.
            return J
        elif self.forward_activation == 'leakyrelu':
            J = torch.ones_like(a)
            J[a < 0.] = 0.2
            return J
        elif self.forward_activation == 'linear':
            return torch.ones_like(a)
        elif self.forward_activation == 'sigmoid':
            s = torch.sigmoid(a)
            return s * (1 - s)
        else:
            raise ValueError('The provided forward activation {} is not '
                             'supported'.format(self.forward_activation))

    def requires_grad(self):
        """ Set require_grad attribute of the activations of this layer to
        True, such that the gradient will be saved in the activation tensor."""
        self._activations.requires_grad = True

    def forward(self, x):
        """Compute the output activation of the layer.

        This method applies first a linear mapping with the
        parameters ``weights`` and ``bias``, after which it applies the
        forward activation function.

        Args:
            x: A mini-batch of size B x in_features with input activations from
            the previous layer or input.

        Returns:
            The mini-batch of output activations of the layer.
        """

        h = x.mm(self.weights.t())
        if self.bias is not None:
            h += self.bias.unsqueeze(0).expand_as(h)
        self.linearactivations = h

        self.activations = self.forward_activationfunction(h)
        return self.activations

    def dummy_forward(self, x):
        """ Same as the forward method, besides that the activations and
        linear activations are not saved in self."""
        h = x.mm(self.weights.t())
        if self.bias is not None:
            h += self.bias.unsqueeze(0).expand_as(h)
        h = self.forward_activationfunction(h)
        return h

    def dummy_forward_linear(self, x):
        """ Propagate the input of the layer forward to the linear activation
        of the current layer (so no nonlinearity applied), without saving the
        linear activations."""
        a = x.mm(self.weights.t())
        if self.bias is not None:
            a += self.bias.unsqueeze(0).expand_as(a)

        return a

    def propagate_backward(self, h):
        """ Propagate the activation h backward through the backward mapping
        function g(h) = t(Q*h + d)
        Args:
            h (torch.Tensor): a mini-batch of activations
        """
        h = h.mm(self.feedbackweights.t())
        if self.feedbackbias is not None:
            h += self.feedbackbias.unsqueeze(0).expand_as(h)
        return self.feedback_activationfunction(h)

    def backward(self, h_target, h_previous, h_current):
        """Compute the target activation for the previous layer, based on the
        provided target.

        Args:
            h_target: a mini-batch of the provided targets for this layer.
            h_previous: the activations of the previous layer, used for the
                DTP correction term.
            h_current: the activations of the current layer, used for the
                DTP correction term.

        Returns:
            h_target_previous: The mini-batch of target activations for
                the previous layer.
        """

        h_target_previous = self.propagate_backward(h_target)
        h_tilde_previous = self.propagate_backward(h_current)
        h_target_previous = h_target_previous + h_previous - h_tilde_previous

        return h_target_previous

    def compute_forward_gradients(self, h_target, h_previous, norm_ratio=1.):
        """ Compute the gradient of the forward weights and bias, based on the
        local mse loss between the layer activation and layer target.
        The gradients are saved in the .grad attribute of the forward weights
        and forward bias.
        Args:
            h_target (torch.Tensor): the DTP target of the current layer
            h_previous (torch.Tensor): the rate activation of the previous
                layer
            norm_ratio (float): Depreciated.
        """

        if self.forward_activation == 'linear':
            teaching_signal = 2 * (self.activations - h_target)
        else:
            vectorized_jacobians = self.compute_vectorized_jacobian(
                self.linearactivations)
            teaching_signal = 2 * vectorized_jacobians * (
                    self.activations - h_target)
        batch_size = h_target.shape[0]
        bias_grad = teaching_signal.mean(0)
        weights_grad = 1./batch_size * teaching_signal.t().mm(h_previous)

        if self.bias is not None:
            self._bias.grad = bias_grad.detach()
        self._weights.grad = weights_grad.detach()

    def set_feedback_requires_grad(self, value):
        """
        Sets the 'requires_grad' attribute of the feedback weights and bias to
        the given value
        Args:
            value (bool): True or False
        """
        if not isinstance(value, bool):
            raise TypeError('The given value should be a boolean.')
        self._feedbackweights.requires_grad = value
        if self._feedbackbias is not None:
            self._feedbackbias.requires_grad = value

    def compute_feedback_gradients(self, h_previous_corrupted, sigma):
        """ Compute the gradient of the backward weights and bias, based on the
        local reconstruction loss of a corrupted sample of the previous layer
        activation. The gradients are saved in the .grad attribute of the
        feedback weights and feedback bias."""

        self.set_feedback_requires_grad(True)

        h_current = self.dummy_forward(h_previous_corrupted)
        h = self.propagate_backward(h_current)

        if sigma < 0.02:
            scale = 1/0.02**2
        else:
            scale = 1/sigma**2
        reconstruction_loss = scale * F.mse_loss(h, h_previous_corrupted)

        self.save_feedback_gradients(reconstruction_loss)

        self.set_feedback_requires_grad(False)

    def save_feedback_gradients(self, reconstruction_loss):
        """
        Compute the gradients of the reconstruction_loss with respect to the
        feedback parameters by help of autograd and save them in the .grad
        attribute of the feedback parameters
        Args:
            reconstruction_loss: the reconstruction loss

        """
        self.reconstruction_loss = reconstruction_loss.item()
        if self.feedbackbias is not None:
            grads = torch.autograd.grad(reconstruction_loss, [
                self.feedbackweights, self.feedbackbias], retain_graph=False)
            self._feedbackbias.grad = grads[1].detach()
        else:
            grads = torch.autograd.grad(reconstruction_loss,
                                        self.feedbackweights,
                                        retain_graph=False
                                        )
        self._feedbackweights.grad = grads[0].detach()

    def compute_bp_update(self, loss, retain_graph=False):
        """ Compute the error backpropagation update for the forward
        parameters of this layer, based on the given loss.
        Args:
            loss (nn.Module): network loss
            retain_graph (bool): flag indicating whether the graph of the
                network should be retained after computing the gradients or
                jacobians. If the graph will not be used anymore for the current
                minibatch afterwards, retain_graph should be False.
        """

        if self.bias is not None:
            grads = torch.autograd.grad(loss, [self.weights, self.bias],
                                        retain_graph=retain_graph)
        else:
            grads = torch.autograd.grad(loss, self.weights,
                                        retain_graph=retain_graph)

        return grads

    def compute_gn_update(self, output_activation, loss, damping=0.,
                          retain_graph=False):
        """
        Compute the Gauss Newton update for the parameters of this layer based
        on the current minibatch.
        Args:
            output_activation (torch.Tensor): The tensor containing the output
                activations of the network for the current mini-batch
            loss (torch.Tensor): The 0D tensor containing the loss value of the
                current mini-batch.
            damping (float): the damping coefficient to damp the GN curvature
                matrix J^TJ. Default: 0.
            retain_graph (bool): flag indicating whether the graph of the
                network should be retained after computing the gradients or
                jacobians. If the graph will not be used anymore for the current
                minibatch afterwards, retain_graph should be False.

        Returns (tuple): A tuple containing the Gauss Newton updates for the
            forward parameters (at index 0 the weight updates, at index 1
            the bias updates if the layer has a bias)

        """
        output_error = torch.autograd.grad(loss, output_activation,
                                           retain_graph=True)[0].detach()
        parameters = self.get_forward_parameters()
        jacobian = utils.compute_jacobian(parameters, output_activation,
                                          retain_graph=retain_graph)

        gn_updates = utils.compute_damped_gn_update(jacobian, output_error,
                                                    damping)

        if self.bias is not None:
            weight_update_flattened = gn_updates[:self.weights.numel(), :]
            bias_update_flattened = gn_updates[self.weights.numel():, :]
            weight_update = weight_update_flattened.view_as(self.weights)
            bias_update = bias_update_flattened.view_as(self.bias)
            return (weight_update, bias_update)
        else:
            weight_update = gn_updates.view(self.weights.shape)
            return (weight_update, )

    def compute_gn_activation_updates(self, output_activation, loss,
                                      damping=0., retain_graph=False,
                                      linear=False):
        """
        Compute the Gauss Newton update for activations of the layer. Target
        propagation tries to approximate these updates by the difference between
        the layer targets and the layer activations.
        Args:
            output_activation (torch.Tensor): The tensor containing the output
                activations of the network for the current mini-batch
            loss (torch.Tensor): The 0D tensor containing the loss value of the
                current mini-batch.
            damping (float): the damping coefficient to damp the GN curvature
                matrix J^TJ. Default: 0.
            retain_graph (bool): flag indicating whether the graph of the
                network should be retained after computing the gradients or
                jacobians. If the graph will not be used anymore for the current
                minibatch afterwards, retain_graph should be False.
            linear (bool): Flag indicating whether the GN update for the
                linear activations should be computed instead of for the
                nonlinear activations.

        Returns (torch.Tensor): A tensor containing the Gauss-Newton updates
            for the layer activations of the current mini-batch. The size is
            minibatchsize x layersize

        """
        output_error = torch.autograd.grad(loss, output_activation,
                                           retain_graph=True)[0].detach()
        if linear:
            activations = self.linearactivations
        else:
            activations = self.activations
        activations_updates = torch.Tensor(activations.shape)
        layersize = activations.shape[1]

        # compute the GN update for each batch sample separate, as we are now
        # computing 'updates' for the activations of the layer instead of the
        # parameters of the layers
        for batch_idx in range(activations.shape[0]):
            # print(batch_idx)
            #  compute jacobian for one batch sample:
            if batch_idx == activations.shape[0] - 1:
                retain_graph_flag = retain_graph
            else:
                # if not yet at the end of the batch, we should retain the graph
                # used for computing the jacobian, as the graph needs to be
                # reused for the computing the jacobian of the next batch sample
                retain_graph_flag = True
            jacobian = utils.compute_jacobian(activations,
                                              output_activation[batch_idx,
                                              :],
                                            retain_graph=retain_graph_flag)
            # torch.autograd.grad only accepts the original input tensor,
            # not a subpart of it. Thus we compute the jacobian to all the
            # batch samples from activations and then select the correct
            # part of it
            jacobian = jacobian[:, batch_idx*layersize:
                                   (batch_idx+1)*layersize]

            gn_updates = utils.compute_damped_gn_update(jacobian,
                                                output_error[batch_idx, :],
                                                        damping)
            activations_updates[batch_idx, :] = gn_updates.view(-1)
        return activations_updates

    def compute_bp_activation_updates(self, loss, retain_graph=False,
                                      linear=False):
        """ Compute the error backpropagation teaching signal for the
        activations of this layer, based on the given loss.
        Args:
            loss (nn.Module): network loss
            retain_graph (bool): flag indicating whether the graph of the
                network should be retained after computing the gradients or
                jacobians. If the graph will not be used anymore for the current
                minibatch afterwards, retain_graph should be False.
            linear (bool): Flag indicating whether the GN update for the
                linear activations should be computed instead of for the
                nonlinear activations.
        Returns (torch.Tensor): A tensor containing the BP updates for the layer
            activations for the current mini-batch.

                """

        if linear:
            activations = self.linearactivations
        else:
            activations = self.activations
        grads = torch.autograd.grad(loss, activations,
                                    retain_graph=retain_graph)[0].detach()
        return grads

    def compute_gnt_updates(self, output_activation, loss, h_previous, damping=0.,
                            retain_graph=False, linear=False):
        """ Compute the angle with the GNT updates for the parameters of the
        network."""
        gn_activation_update = self.compute_gn_activation_updates(output_activation=output_activation,
                                                                  loss=loss,
                                                                  damping=damping,
                                                                  retain_graph=retain_graph,
                                                                  linear=linear)

        if not linear:
            vectorized_jacobians = self.compute_vectorized_jacobian(
                self.linearactivations)
            teaching_signal = 2 * vectorized_jacobians * (
                        gn_activation_update)
        else:
            teaching_signal = 2 * gn_activation_update

        batch_size = self.activations.shape[0]
        bias_grad = teaching_signal.mean(0)
        weights_grad = 1. / batch_size * teaching_signal.t().mm(h_previous)

        if self.bias is not None:
            return (weights_grad, bias_grad)
        else:
            return (weights_grad, )

    def compute_nullspace_relative_norm(self, output_activation, retain_graph=False):
        """ Compute the norm of the components of weights.grad that are in the nullspace
        of the jacobian of the output with respect to weights, relative to the norm of
        weights.grad."""
        J = utils.compute_jacobian(self.weights, output_activation,
                                   structured_tensor=False,
                                   retain_graph=retain_graph)
        weights_update_flat = self.weights.grad.view(-1)
        relative_norm = utils.nullspace_relative_norm(J, weights_update_flat)
        return relative_norm

    def save_logs(self, writer, step, name, no_gradient=False,
                  no_fb_param=False):
        """
        Save logs and plots of this layer on tensorboardX
        Args:
            writer (SummaryWriter): summary writer from tensorboardX
            step (int): the global step used for the x-axis of the plots
            name (str): The name of the layer
            no_gradient (bool): flag indicating whether we should skip saving
                the gradients of the feedback weights.
            no_fb_param (bool): don't log the feedback parameters

        """
        forward_weights_norm = torch.norm(self.weights)
        writer.add_scalar(tag='{}/forward_weights_norm'.format(name),
                          scalar_value=forward_weights_norm,
                          global_step=step)
        if self.weights.grad is not None:
            forward_weights_gradients_norm = torch.norm(self.weights.grad)
            writer.add_scalar(tag='{}/forward_weights_gradients_norm'.format(name),
                              scalar_value=forward_weights_gradients_norm,
                              global_step=step)
        if self.bias is not None:
            forward_bias_norm = torch.norm(self.bias)

            writer.add_scalar(tag='{}/forward_bias_norm'.format(name),
                              scalar_value=forward_bias_norm,
                              global_step=step)
            if self.bias.grad is not None:
                forward_bias_gradients_norm = torch.norm(self.bias.grad)
                writer.add_scalar(tag='{}/forward_bias_gradients_norm'.format(name),
                                  scalar_value=forward_bias_gradients_norm,
                                  global_step=step)
        if not no_fb_param:
            feedback_weights_norm = torch.norm(self.feedbackweights)
            writer.add_scalar(tag='{}/feedback_weights_norm'.format(name),
                              scalar_value=feedback_weights_norm,
                              global_step=step)
            if self.feedbackbias is not None:
                feedback_bias_norm = torch.norm(self.feedbackbias)
                writer.add_scalar(tag='{}/feedback_bias_norm'.format(name),
                                  scalar_value=feedback_bias_norm,
                                  global_step=step)

            if not no_gradient and self.feedbackweights.grad is not None:
                feedback_weights_gradients_norm = torch.norm(
                    self.feedbackweights.grad)
                writer.add_scalar(
                    tag='{}/feedback_weights_gradients_norm'.format(name),
                    scalar_value=feedback_weights_gradients_norm,
                    global_step=step)
                if self.feedbackbias is not None:
                    feedback_bias_gradients_norm = torch.norm(
                        self.feedbackbias.grad)
                    writer.add_scalar(
                        tag='{}/feedback_bias_gradients_norm'.format(name),
                        scalar_value=feedback_bias_gradients_norm,
                        global_step=step)

    def save_feedback_batch_logs(self, writer, step, name, no_gradient=False,
                                 init=False):
        """
        Save logs for one minibatch.
        Args:
            writer (SummaryWriter): summary writer from tensorboardX
            step (int): the global step used for the x-axis of the plots
            name (str): The name of the layer
            no_gradient (bool): flag indicating whether we should skip saving
                the gradients of the feedback weights.
            init (bool): flag indicating that the training is in the
                initialization phase (only training the feedback weights).
        """
        if not init:
            if not no_gradient and self.reconstruction_loss is not None:
                writer.add_scalar(
                    tag='{}/reconstruction_loss'.format(name),
                    scalar_value=self.reconstruction_loss,
                    global_step=step)
        else:
            if not no_gradient and self.reconstruction_loss is not None:
                writer.add_scalar(
                    tag='{}/reconstruction_loss_init'.format(name),
                    scalar_value=self.reconstruction_loss,
                    global_step=step)

    def get_forward_parameters(self):
        """ Return a list containing the forward parameters."""
        if self.bias is not None:
            return [self.weights, self.bias]
        else:
            return [self.weights]

    def get_forward_gradients(self):
        """ Return a tuple containing the gradients of the forward
        parameters."""

        if self.bias is not None:
            return (self.weights.grad, self.bias.grad)
        else:
            return (self.weights.grad, )


text:
"""Train the model"""
"""This is from the pytorch_shuffle dir"""

import argparse
import logging
import os
import sys
import numpy as np
import torch
import torch.optim as optim
from sklearn.utils import check_random_state
from sklearn.model_selection import train_test_split
import time
import pickle
import utils
import model.data_loader as dl
import model.dataset as dataset
from model import recNet as net
from model import preprocess

##----------------------------------------------------------------------------------------------------------
if __name__=='__main__':

  print('Preprocessing jet trees ...')
  print('==='*20)
  ##----------------------------------------------------------------------------------------------------------
  # Global variables
  ##-------------------
  data_dir='../data/'
  os.system('mkdir -p '+data_dir)

  # Select the right dir for jets data
  trees_dir='preprocessed_trees/'
  os.system('mkdir -p '+data_dir+'/'+trees_dir)

  ##-------------------
  #If true the preprocessed trees are generated and saved. Do it only once and then turn in off
#   make_preprocess=True
#   make_preprocess=False

#   pT_order=True
  pT_order=False

  # Select the input sample
#   nyu=True
  nyu=False

  toptag_reference_dataset=True # If True, specify below the location of the train, val, test input trees
  test_train_val_datasets=False

  ##--------------------
  algo=''
  sample_name=''
#   if nyu==True:
    #Directory with the input trees
#     sample_name='nyu_jets'

  #   algo='antikt-antikt-delphes'
#     algo='antikt-kt-delphes'
#     algo='antikt-antikt'

    # file_name=algo+'-train.pickle'

#   if toptag_reference_dataset
#     sample_name='top_tag_reference_dataset'

#   else:
#     algo=''

    #Directory with the input trees
    ### CHECK THAT SEARCH_HYPERPARAMS.PY HAS THE SAME SAMPLE NAME

#     sample_name='top_qcd_jets_antikt_antikt'
#     sample_name='top_qcd_jets_antikt_kt'
#     sample_name='top_qcd_jets_antikt_CA'

    #labels to look for the input files
  #   sg='tt'

  sg='ttbar'
  bg='qcd'

  ##------------------------------------------------------------
  parser = argparse.ArgumentParser()
  parser.add_argument('--data_dir', default='../data/inputTrees/'+sample_name, help="Directory containing the raw datasets")
  parser.add_argument('--model_dir', default='experiments/base_model', help="Directory containing params.json")
#   parser.add_argument('--restore_file', default=None,
#                       help="Optional, name of the file in --model_dir containing weights to reload before \
#                       training")  # 'best' or 'last'

  parser.add_argument('--jet_algorithm', default=algo, help="jet algorithm")
#   parser.add_argument('--transformer', default='../data/preprocessed_trees/', help="Transformen from running RobustScaler on the train dataset")

  # Load the parameters from json file
  args = parser.parse_args()
  json_path = os.path.join(args.model_dir, 'params.json')
  assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)
  params = utils.Params(json_path)

  ##-------------------
  # Set the logger
  utils.set_logger(os.path.join(args.model_dir, 'preprocess.log'))

  dir_jets_subjets= args.data_dir
  algo=args.jet_algorithm
  file_name=algo+'-train.pickle'
  ##-------------------
  # Define output file names with the batches of data. We rewrite the sample name if running from search_hyperparam.py
  sample_name=str(args.data_dir).split('/')[-1]
  logging.info('sample_name={}'.format(sample_name))
  logging.info('----'*20)
  # sample_filename=sample_name+'_'+algo+'_'+str(params.myN_jets)+'_Njets_'+str(params.batch_size)+'_batch'+'_'+str(params.info)
  sample_filename=sample_name+'_'+algo+'_'+str(params.myN_jets)+'_Njets_'+str(params.info)
  logging.info('sample_filename={}'.format(sample_filename))

  train_data=data_dir+trees_dir+'train_'+sample_filename+'.pkl'
  val_data=data_dir+trees_dir+'dev_'+sample_filename+'.pkl'
  test_data=data_dir+trees_dir+'test_'+sample_filename+'.pkl'

  transformer_filename=sample_name+'_'+algo+'_'+str(params.myN_jets)+'_Njets_'
  transformer_data=data_dir+trees_dir+'transformer_'+transformer_filename+'.pkl'

  start_time = time.time()
  ##-----------------------------------------------------------------------------------------------------------
  data_loader=dl.DataLoader # Main class with the methods to load the raw data, create and preprocess the trees

  # Create the batches
#   if make_preprocess==True:

  # FOR NYU SAMPLES (from arXiv:1702.00748)
  if nyu==True:

    nyu_train=dir_jets_subjets+'/'+file_name
    # loading dataset_params and make trees
    logging.info('nyu_train={}'.format(nyu_train))
    fd = open(nyu_train, "rb")
    X, Y = pickle.load(fd,encoding='latin-1')
    fd.close()

    X=np.asarray(X)
    Y=np.asarray(Y)

    logging.info('Training data size={}'.format(len(X)))
    logging.info('---'*20)

    # Shuffle the sets
    indices = check_random_state(1).permutation(len(X))
    X = X[indices]
    Y = Y[indices]
    X=np.asarray(X)
    Y=np.asarray(Y)

    # Preprocessing steps: Ensure that the left sub-jet has always a larger pt than the right. Change the input variables (features)
    X = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in X]
#
#       # Apply RobustScaler (remove outliers, center and scale data)
#       X=data_loader.scale_features(X)

    # Split into train+validation
    logging.info("Splitting into train and validation...")

    train_x, dev_x, train_y, dev_y = train_test_split(X, Y, test_size=5000, random_state=0)

    # Apply RobustScaler (remove outliers, center and scale data)
    transformer=data_loader.get_transformer(train_x)
    # Save transformer
    with open(transformer_data, "wb") as f: pickle.dump(transformer, f)

    #Scale features using the training set transformer
    train_x = data_loader.transform_features(transformer,train_x)
    dev_x = data_loader.transform_features(transformer,dev_x)

    #----------------------------------------------------------
    # Preprocess test set

    # Create the input data pipeline
    logging.info("Preprocessing the test dataset...")

#     batches_dir='input_batches_pad/'
#     test_data=data_dir+batches_dir+'test_'+sample_filename+'.pkl'

    #Load test smaple
    nyu_test=dir_jets_subjets+'/'+file_name
    print('nyu_test=',nyu_test)
    fd = open(nyu_test, "rb")
    X, y = pickle.load(fd,encoding='latin-1')
    fd.close()

    X=np.asarray(X)
    y=np.asarray(y)

    logging.info('Training data size= '+str(len(X)))
    logging.info('---'*20)

    #------------------
    # Shuffle the sets
    indices = check_random_state(1).permutation(len(X))
    X = X[indices]
    y = y[indices]
    X=np.asarray(X)
    y=np.asarray(y)

  # Preprocessing steps: Ensure that the left sub-jet has always a larger pt than the right. Change the input variables (features)
    X = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in X]

    # Apply RobustScaler (remove outliers, center and scale data) with train set transformer
#     transformer_filename=sample_name+'_'+algo+'_'+str(params.myN_jets)+'_Njets_'
#     transformer_data=data_dir+batches_dir+'transformer_'+transformer_filename+'.pkl'
    with open(transformer_data, "rb") as f: transformer =pickle.load(f)

    #Scale features using the training set transformer
    X = data_loader.transform_features(transformer,X)

#
#       # Apply RobustScaler (remove outliers, center and scale data)
#       X=data_loader.scale_features(X)

    #------------------
    # Cropping
    X_ = [j for j in X if 250 < j["pt"] < 300 and 50 < j["mass"] < 110]
    y_ = [y[i] for i, j in enumerate(X) if 250 < j["pt"] < 300 and 50 < j["mass"] < 110]

    X = X_
    y = y_
    X=np.asarray(X)
    y = np.asarray(y)

    logging.info('Lenght X= '+str(len(X)))
    logging.info('Length y= '+str(len(y)))

    #------------------
    # Weights for flatness in pt
    w = np.zeros(len(y))

    logging.info('Length w before='+str(len(w)))

    X0 = [X[i] for i in range(len(y)) if y[i] == 0]
    pdf, edges = np.histogram([j["pt"] for j in X0], density=True, range=[250, 300], bins=50)
    pts = [j["pt"] for j in X0]
    indices = np.searchsorted(edges, pts) - 1
    inv_w = 1. / pdf[indices]
    inv_w /= inv_w.sum()
    w[y==0] = inv_w

    X1 = [X[i] for i in range(len(y)) if y[i] == 1]
    pdf, edges = np.histogram([j["pt"] for j in X1], density=True, range=[250, 300], bins=50)
    pts = [j["pt"] for j in X1]
    indices = np.searchsorted(edges, pts) - 1
    inv_w = 1. / pdf[indices]
    inv_w /= inv_w.sum()
    w[y==1] = inv_w

    logging.info('Length w after='+str(len(w)))

    test_x=X
    test_y=y

    # Save trees
    with open(train_data, "wb") as f: pickle.dump(zip(train_x,train_y), f)
    with open(val_data, "wb") as f: pickle.dump(zip(dev_x,dev_y), f)
    with open(test_data, "wb") as f: pickle.dump((zip(test_x,test_y),w), f) #We save the weights and test set

#
#     #-------------------
#     # Generate dataset batches.
#     test_batches=dl.batch_array(X, y, params.batch_size, params.features)
#     logging.info('Number of test_batches='+str(len(test_batches)))
#
#     # Save batches
#     with open(test_data, "wb") as f: pickle.dump((test_batches,w), f)

    logging.info("- done.")

  ##-------------------------------------------------------------------------------------------------
  ##-------------------------------------------------------------------------------------------------
  if toptag_reference_dataset:

    if test_train_val_datasets:
    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_1200000_R_0.3_rot_boost_rot_flip.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_400000_R_0.3_rot_boost_rot_flip.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_400000_R_0.3_rot_boost_rot_flip.pkl'

    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_120001.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_40001.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_40001.pkl'

        toptag_reference_train=dir_jets_subjets+'/tree_train_jets.pkl'
        toptag_reference_val=dir_jets_subjets+'/tree_val_jets.pkl'
        toptag_reference_test=dir_jets_subjets+'/tree_test_jets.pkl'

    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_1001.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_1001.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_1001.pkl'

        # loading dataset_params and make trees
        logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_val))
        with open(toptag_reference_val, "rb") as f: toptag_reference_val =pickle.load(f,encoding='latin-1')

        logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_train))
        with open(toptag_reference_train, "rb") as f: toptag_reference_train =pickle.load(f,encoding='latin-1')

        logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_test))
        with open(toptag_reference_test, "rb") as f: toptag_reference_test =pickle.load(f,encoding='latin-1')
    #

        toptag_reference_train_x=np.asarray([x for (x,y) in toptag_reference_train])
        toptag_reference_train_y=np.asarray([y for (x,y) in toptag_reference_train])

        toptag_reference_val_x=np.asarray([x for (x,y) in toptag_reference_val])
        toptag_reference_val_y=np.asarray([y for (x,y) in toptag_reference_val])

        toptag_reference_test_x=np.asarray([x for (x,y) in toptag_reference_test])
        toptag_reference_test_y=np.asarray([y for (x,y) in toptag_reference_test])

        logging.info('Training data size={}'.format(len(toptag_reference_train_x)))
        logging.info('---'*20)

        # Shuffle the training set
        indices = check_random_state(1).permutation(len(toptag_reference_train_x))
        toptag_reference_train_x = toptag_reference_train_x[indices]
        toptag_reference_train_y = toptag_reference_train_y[indices]

        toptag_reference_train_x=np.asarray(toptag_reference_train_x)
        toptag_reference_train_y=np.asarray(toptag_reference_train_y)

    #     print('toptag_reference_train_y=',toptag_reference_train_y[0:20])

        # Preprocess
        toptag_reference_train_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_train_x]

        toptag_reference_val_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_val_x]

        toptag_reference_test_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_test_x]

        # Apply RobustScaler (remove outliers, center and scale data)
        transformer=data_loader.get_transformer(toptag_reference_train_x)
        # Save transformer
        with open(transformer_data, "wb") as f: pickle.dump(transformer, f)

        #Scale features using the training set transformer
        toptag_reference_train_x = data_loader.transform_features(transformer,toptag_reference_train_x)
        toptag_reference_val_x = data_loader.transform_features(transformer,toptag_reference_val_x)
        toptag_reference_test_x = data_loader.transform_features(transformer,toptag_reference_test_x)

        ##---------------------------------

        # Save trees
        with open(train_data, "wb") as f: pickle.dump(zip(toptag_reference_train_x,toptag_reference_train_y), f)
        with open(val_data, "wb") as f: pickle.dump(zip(toptag_reference_val_x,toptag_reference_val_y), f)
        with open(test_data, "wb") as f: pickle.dump(zip(toptag_reference_test_x,toptag_reference_test_y), f)

    else:
        toptag_reference_test=dir_jets_subjets+'/tree_test_jets.pkl'
    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_1200000_R_0.3_rot_boost_rot_flip.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_400000_R_0.3_rot_boost_rot_flip.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_400000_R_0.3_rot_boost_rot_flip.pkl'

    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_120001.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_40001.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_40001.pkl'

    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets.pkl'

    #     toptag_reference_train=dir_jets_subjets+'/tree_train_jets_1001.pkl'
    #     toptag_reference_val=dir_jets_subjets+'/tree_val_jets_1001.pkl'
    #     toptag_reference_test=dir_jets_subjets+'/tree_test_jets_1001.pkl'

        # loading dataset_params and make trees
    #     logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_val))
    #     with open(toptag_reference_val, "rb") as f: toptag_reference_val =pickle.load(f,encoding='latin-1')
    #
    #     logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_train))
    #     with open(toptag_reference_train, "rb") as f: toptag_reference_train =pickle.load(f,encoding='latin-1')

        logging.info('Loading toptag_reference_dataset={}'.format(toptag_reference_test))
        with open(toptag_reference_test, "rb") as f: toptag_reference_test =pickle.load(f,encoding='latin-1')
    #

    #     toptag_reference_train_x=np.asarray([x for (x,y) in toptag_reference_train])
    #     toptag_reference_train_y=np.asarray([y for (x,y) in toptag_reference_train])
    #
    #     toptag_reference_val_x=np.asarray([x for (x,y) in toptag_reference_val])
    #     toptag_reference_val_y=np.asarray([y for (x,y) in toptag_reference_val])

        toptag_reference_test_x=np.asarray([x for (x,y) in toptag_reference_test])
        toptag_reference_test_y=np.asarray([y for (x,y) in toptag_reference_test])

    #     logging.info('Training data size={}'.format(len(toptag_reference_train_x)))
        logging.info('---'*20)

        # Shuffle the training set
    #     indices = check_random_state(1).permutation(len(toptag_reference_train_x))
    #     toptag_reference_train_x = toptag_reference_train_x[indices]
    #     toptag_reference_train_y = toptag_reference_train_y[indices]
    #
    #     toptag_reference_train_x=np.asarray(toptag_reference_train_x)
    #     toptag_reference_train_y=np.asarray(toptag_reference_train_y)

    #     print('toptag_reference_train_y=',toptag_reference_train_y[0:20])

        # Preprocess
    #     toptag_reference_train_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_train_x]
    #
    #     toptag_reference_val_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_val_x]

        toptag_reference_test_x = [preprocess.extract_nyu_samples(preprocess.permute_by_pt(preprocess.rewrite_content(jet))) for jet in toptag_reference_test_x]

        # Apply RobustScaler (remove outliers, center and scale data)
    #     transformer=data_loader.get_transformer(toptag_reference_train_x)
        # Save transformer
    #     with open(transformer_data, "wb") as f: pickle.dump(transformer, f)

        #Load transformer
        with open(transformer_data, "rb") as f: transformer =pickle.load(f)

        #Scale features using the training set transformer
    #     toptag_reference_train_x = data_loader.transform_features(transformer,toptag_reference_train_x)
    #     toptag_reference_val_x = data_loader.transform_features(transformer,toptag_reference_val_x)
        toptag_reference_test_x = data_loader.transform_features(transformer,toptag_reference_test_x)

        ##---------------------------------

        # Save trees
    #     with open(train_data, "wb") as f: pickle.dump(zip(toptag_reference_train_x,toptag_reference_train_y), f)
    #     with open(val_data, "wb") as f: pickle.dump(zip(toptag_reference_val_x,toptag_reference_val_y), f)
        with open(test_data, "wb") as f: pickle.dump(zip(toptag_reference_test_x,toptag_reference_test_y), f)

  ##-------------------------------------------------------------------------------------------------
  # FOR OUR OWN SAMPLES
  else:

    # loading dataset_params and make trees
    sig_list=data_loader.makeTrees(dir_jets_subjets,sg,params.myN_jets,1)
    bkg_list=data_loader.makeTrees(dir_jets_subjets,bg,params.myN_jets,0)
    elapsed_time=time.time()-start_time
    logging.info('Tree generation time (minutes) ={}'.format(elapsed_time/60))

    ##----------------------------------------------
    # Preprocessing steps:
    #  - Ensure that the left sub-jet has always a larger pt than the right (or more leaves).
    #  - Change the input variables (features)

    if pT_order==True:
      sig_list = [preprocess.extract(preprocess.sequentialize_by_pt(preprocess.permute_by_pt(jet)), params.features) for jet in sig_list]
      bkg_list = [preprocess.extract(preprocess.sequentialize_by_pt(preprocess.permute_by_pt(jet)), params.features) for jet in bkg_list]

    else:

      ##----
      ## a) Add number of jet constituents in each branch and assign the branch with the largest number as the left node
#         sig_list = [preprocess.extract(preprocess.permute_by_n_leaves(preprocess.rewrite_content_leaves(jet)), params.features,kappa=0.4) for jet in sig_list]
#         bkg_list = [preprocess.extract(preprocess.permute_by_n_leaves(preprocess.rewrite_content_leaves(jet)), params.features,kappa=0.4) for jet in bkg_list]

      ##----
      ## b) Assign the biggest pT subjet as the left node

      sig_list = [preprocess.extract(preprocess.permute_by_pt(jet), params.features,kappa=0.4) for jet in sig_list]
      bkg_list = [preprocess.extract(preprocess.permute_by_pt(jet), params.features,kappa=0.4) for jet in bkg_list]

    ##-------------------
    # Split into train+validation+test and shuffle
    logging.info("Splitting into train, validation and test datasets, and shuffling...")
    train_x, train_y, dev_x, dev_y, test_x, test_y = data_loader.split_shuffle_sample(sig_list, bkg_list, 0.6, 0.2, 0.2)

#       train_x, X_valid, train_y, Y_valid = train_test_split(X, Y, test_size=0.4, random_state=0)
#       dev_x, test_x, dev_y, test_y = train_test_split(X_valid, Y_valid, test_size=0.5, random_state=1)

    ##-------------------
    # Apply RobustScaler (remove outliers, center and scale data)
    transformer=data_loader.get_transformer(train_x)

    # Save transformer
    with open(transformer_data, "wb") as f: pickle.dump(transformer, f)

    #Scale features using the training set transformer
    train_x = data_loader.transform_features(transformer,train_x)
    dev_x = data_loader.transform_features(transformer,dev_x)
    test_x = data_loader.transform_features(transformer,test_x)

    ##---------------------------------
    elapsed_time=time.time()-start_time
    logging.info('Split sample time (minutes) ={}'.format(elapsed_time/60))

    # Save trees
    with open(train_data, "wb") as f: pickle.dump(zip(train_x,train_y), f)
    with open(val_data, "wb") as f: pickle.dump(zip(dev_x,dev_y), f)
  #   if nyu==True:
  #     #We save the weights
  #     with open(test_data, "wb") as f: pickle.dump((zip(test_x,test_y),w), f)
  #   else:
    with open(test_data, "wb") as f: pickle.dump(zip(test_x,test_y), f)

ext: py

reason: import numpy


text:
import numpy as np
import os
import scipy.stats as statss
import scipy
from random import choices
from seqeval.metrics import precision_score, recall_score, f1_score
from collections import Counter
import string
import re
import argparse
import json
import sys
from typing import Any, ClassVar, Dict, List, Optional

'''
QA
'''

def normalize_answer(s):
	"""Lower text and remove punctuation, articles and extra whitespace."""
	def remove_articles(text):
		return re.sub(r'\b(a|an|the)\b', ' ', text)

	def white_space_fix(text):
		return ' '.join(text.split())

	def remove_punc(text):
		exclude = set(string.punctuation)
		return ''.join(ch for ch in text if ch not in exclude)

	def lower(text):
		return text.lower()

	return white_space_fix(remove_articles(remove_punc(lower(s))))

def f1_score_qa_sample_level(prediction:str, ground_truth:str):
	prediction_tokens = normalize_answer(prediction).split()
	ground_truth_tokens = normalize_answer(ground_truth).split()
	common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
	num_same = sum(common.values())
	if num_same == 0:
		return 0
	precision = 1.0 * num_same / len(prediction_tokens)
	recall = 1.0 * num_same / len(ground_truth_tokens)
	f1 = (2 * precision * recall) / (precision + recall)
	return f1

def exact_match_sample_level(prediction:str, ground_truth:str):
	# print("prediction, ground_truth 2: ",prediction, ground_truth)
	# print("normalize_answer: ",normalize_answer(prediction),normalize_answer(ground_truth))
	return (normalize_answer(prediction) == normalize_answer(ground_truth))

def metric_max_over_ground_truths(metric_fn, prediction:str, ground_truths:list):
    # scores_for_ground_truths = []
    # # for ground_truth in ground_truths:
    # score = metric_fn(prediction, ground_truths)
    # scores_for_ground_truths.append(score)
    # return max(scores_for_ground_truths)

    scores_for_ground_truths = []
    for ground_truth in ground_truths:
        # print(prediction, ground_truth)
        score = metric_fn(prediction, ground_truth)
        scores_for_ground_truths.append(score)
    return max(scores_for_ground_truths)

def exact_match_qa(true_answers:List[list], predicted_answer:List[str]):
    exact_match = 0
    total = 0
    # for k1,k2 in zip(true_Anss,pred_Anss):
    #     print(k1, k2)
    for true_ans, pred_ans in zip(true_answers, predicted_answer):
        total += 1
        exact_match1 = metric_max_over_ground_truths(
            exact_match_sample_level, pred_ans, true_ans)

        exact_match+=exact_match1

    exact_match = 100.0 * exact_match / total

    return exact_match

def f1_score_qa(true_answers:List[list], predicted_answer:List[str]):
    f1_dataset_level = 0
    total = 0

    for true_ans, pred_ans in zip(true_answers, predicted_answer):
        total += 1

        f1_sample = metric_max_over_ground_truths(
        	f1_score_qa_sample_level, pred_ans, true_ans)
        f1_dataset_level += f1_sample

    f1_dataset_level = 100.0 * f1_dataset_level / total

    return f1_dataset_level

max_stars_repo_name: abodacs/ExplainaBoard

max_issues_repo_name: abodacs/ExplainaBoard

max_forks_repo_name: abodacs/ExplainaBoard

reason: import numpy,import scipy


text:
from __future__ import print_function
from __future__ import division
from . import _C

import numpy as np
from scipy import stats
from ..strings import xstr
from .xerror import XError
from ..dataframes import DFBuilder
import math
from mlxtend.evaluate import permutation_test
from ..progress_bars import ProgressBar

DEFAULT_TH_PVALUE = .05
NUM_ROUNDS = 1e4
RANDOM_STATE = 0
ALTERNATIVE = 'two-sided' # two-sided less greater
TH_PVALUE_TXT = DEFAULT_TH_PVALUE
SHAPIRO_TH_PVALUE = DEFAULT_TH_PVALUE
N_DECIMALS = _C.N_DECIMALS
PVALUE_CHAR = '$p$'
NULL_CHAR = ''
PVALUE_SYMBOLS = {
	'***':[-np.inf, .001],
	'**':[.001, .01],
	'*':[.01, .05],
	'+':[.05, .1],
	'':[.1, np.inf],
	}
CHECK_SAMPLES = True

###################################################################################################################################################

def _check_xe(xe):
	assert len(xe)>=1 and len(xe.shape)==1

def get_pvalue_symbol(pvalue,
	get_upper_bound=False,
	):
	if pvalue is None:
		return 'x'
	symbols = list(PVALUE_SYMBOLS.keys())
	for s in symbols:
		lower_bound = PVALUE_SYMBOLS[s][0]
		upper_bound = PVALUE_SYMBOLS[s][-1]
		if pvalue>lower_bound and pvalue<=upper_bound:
			if get_upper_bound:
				return upper_bound, s
			else:
				return s

def get_pvalue_symbols():
	symbols = list(PVALUE_SYMBOLS.keys())
	return {f'<={PVALUE_SYMBOLS[s][-1]}':s for s in symbols}

def format_pvalue(pvalue, mean_diff,
	th_pvalue_txt=TH_PVALUE_TXT,
	n_decimals=N_DECIMALS,
	):
	if pvalue is None:
		return NULL_CHAR
	mean_diff_txt = xstr(mean_diff,
		n_decimals=n_decimals,
		)
	txt = f'$\Delta$={mean_diff_txt}{get_pvalue_symbol(pvalue)}'
	if pvalue>th_pvalue_txt:
		txt = f'{txt}; {PVALUE_CHAR}={xstr(pvalue)}'
	return txt

def _normalitytest(x,
	shapiro_th_pvalue=SHAPIRO_TH_PVALUE,
	):
	pvalue = stats.shapiro(x).pvalue
	is_normal = not pvalue<shapiro_th_pvalue
	return is_normal

def _permutation_test(x1, x2,
	alternative=ALTERNATIVE,
	num_rounds=NUM_ROUNDS,
	random_state=RANDOM_STATE,
	):
	pvalue = permutation_test(x1, x2,
		method='approximate',
		num_rounds=int(num_rounds),
		seed=random_state,
		)
	if alternative=='two-sided':
		pvalue = pvalue
	elif alternative=='greater':
		pvalue = pvalue/2
	else:
		raise Exception(f'{alternative}')
	return pvalue

###################################################################################################################################################

def ttest(x1, x2,
	alternative=ALTERNATIVE,
	shapiro_th_pvalue=SHAPIRO_TH_PVALUE,
	):
	pvalue = None
	if type(x1)==XError and type(x2)==XError:
		_check_xe(x1)
		_check_xe(x2)
		pass

	else:
		raise Exception(f'{type(x1)}; {type(x2)}')

	return pvalue

def welchtest(x1, x2,
	alternative=ALTERNATIVE,
	shapiro_th_pvalue=SHAPIRO_TH_PVALUE,
	):
	pvalue = None
	if type(x1)==XError and type(x2)==XError:
		_check_xe(x1)
		_check_xe(x2)
		normal_kwargs = {
			'shapiro_th_pvalue':shapiro_th_pvalue,
			}
		x1 = x1._x
		x2 = x2._x
		both_are_normal = _normalitytest(x1, **normal_kwargs) and _normalitytest(x2, **normal_kwargs)
		if both_are_normal:
			tvalue, pvalue = stats.ttest_ind(x1, x2,
				alternative=alternative,
				equal_var=False, # welch
				)

	else:
		raise Exception(f'{type(x1)}; {type(x2)}')

	return pvalue

def permutationtest(x1, x2,
	alternative=ALTERNATIVE,
	num_rounds=NUM_ROUNDS,
	random_state=RANDOM_STATE,
	):
	pvalue = None
	if type(x1)==XError and type(x2)==XError:
		_check_xe(x1)
		_check_xe(x2)
		x1 = x1._x
		x2 = x2._x
		pvalue = _permutation_test(x1, x2,
			alternative=alternative,
			num_rounds=num_rounds,
			random_state=random_state,
			)

	else:
		raise Exception(f'{type(x1)}; {type(x2)}')

	return pvalue

###################################################################################################################################################

def gridtest_greater(values_dict, test,
	check_samples=CHECK_SAMPLES,
	th_pvalue_txt=TH_PVALUE_TXT,
	n_decimals=N_DECIMALS,
	test_kwargs={}
	):
	pvalue_symbols = get_pvalue_symbols()
	for pvalue_symbol in pvalue_symbols.keys():
		print(f'{pvalue_symbols[pvalue_symbol]} pvalue{pvalue_symbol}')
	df_builder = DFBuilder()
	n = 10
	keys = list(values_dict.keys())
	bar = ProgressBar(len(keys)**2)
	for key1 in keys:
		d = {}
		x1 = values_dict[key1]
		for key2 in keys:
			bar(f'{key1} v/s {key2}')
			x2 = values_dict[key2]
			pvalue, mean_diff = greatertest(x1, x2, test,
				check_samples=check_samples,
				test_kwargs=test_kwargs,
				)
			pvalue_txt = format_pvalue(pvalue, mean_diff,
				th_pvalue_txt=th_pvalue_txt,
				n_decimals=n_decimals,
				)
			d[f'{key2}; #samples={len(x2)}'] = f'{pvalue_txt}'
		df_builder.append(f'{key1}; #samples={len(x1)}', d)
	bar.done()
	return df_builder.get_df()

def greatertest(x1, x2, test,
	check_samples=CHECK_SAMPLES,
	test_kwargs={},
	):
	'''
	check if mean(x1)>mean(x2) with statistical significance
	'''
	if check_samples:
		assert len(x1)==len(x2)
	mean_diff = x1.get_mean()-x2.get_mean()
	if mean_diff<=0:
		return None, None
	pvalue = test(x1, x2,
		alternative='greater',
		**test_kwargs,
		)
	return pvalue, mean_diff

max_stars_repo_name: oscarpimentel/fuzzy-tools

max_issues_repo_name: oscarpimentel/fuzzy-tools

max_forks_repo_name: oscarpimentel/fuzzy-tools

reason: import numpy,from scipy


text:
# -*- coding: UTF-8 -*-

from caid.cad_geometry import cad_geometry
from evaluator import CurveEvaluator, SurfaceEvaluator, VolumeEvaluator
from numpy import array, asarray
from OpenGL.GL import *
from theme import theme as Theme
from global_vars import strtoArray

theme = Theme()

ALPHA                   = theme.alpha
BETA                    = theme.beta
COLOR_DEFAULT_GEOMETRY  = theme.color_viewer('default_geometry')
COLOR_DEFAULT_PATCH     = theme.color_viewer('default_patch')
COLOR_DEFAULT_MESH      = theme.color_viewer('default_mesh')
COLOR_DEFAULT_POINTS    = theme.color_viewer('default_points')
SIZE_GEOMETRY_POINTS    = theme.size_points('geometry')
SIZE_PATCH_POINTS       = theme.size_points('patch')

def pointsToList(points):
        x = points[...,0]
        y = points[...,1]
        z = points[...,2]
        n = asarray(points.shape[:-1]).prod()
        x = x.reshape(n)
        y = y.reshape(n)
        z = z.reshape(n)

        return x,y,z

def controlToList(control):
        x = control[...,0]
        y = control[...,1]
        z = control[...,2]
        w = control[...,3]
        n = asarray(control.shape[:-1]).prod()
        x = x.reshape(n)
        y = y.reshape(n)
        z = z.reshape(n)
        w = w.reshape(n)

        return x,y,z,w

class PatchInfo(object):
    def __init__(self, patch):
        self.patch = patch
        self.show = True
        self.showPoints = False
        self.showMesh = False

        self._MeshColor      = None
        self._NurbsColor     = None
        self._PointsColor    = None

        self.steps          = None

        self.PointsSize     = SIZE_PATCH_POINTS

        self.name = None

        self.load_attributs()

    @property
    def MeshColor(self):
        return self._MeshColor

    @property
    def NurbsColor(self):
        return self._NurbsColor

    @property
    def PointsColor(self):
        return self._PointsColor

    def Show(self):
        self.show = True

    def Hide(self):
        self.show = False

    def ShowPoints(self):
        self.showPoints = True

    def HidePoints(self):
        self.showPoints = False

    def ShowMesh(self):
        self.showMesh = True

    def HideMesh(self):
        self.showMesh = False

    def set_meshColor(self,color):
        self.MeshColor = color

    def set_nurbsColor(self,color):
        self._NurbsColor = color
        self.patch.set_attribut("color-nurbs", str(self.NurbsColor))

    def set_pointsColor(self,color):
        self.PointsColor = color

    def set_pointsSize(self,size):
        self.PointsSize = size

    def set_steps(self, steps):
        self.steps = steps

    def set_name(self, name):
        self.name = name
        self.patch.set_attribut("name", str(name))

    def load_attributs(self):
        patch = self.patch
        # color nurbs
        try:
            self._NurbsColor = strtoArray(patch.get_attribut('color-nurbs'))
        except:
            self._NurbsColor     = None

class geometry(cad_geometry):
    def __new__(typ, *args, **kwargs):
        obj = cad_geometry.__new__(typ)

        obj.MeshColor   = None
        obj.NurbsColor  = None
        obj.PointsColor = None
        obj.PointsSize  = None
        obj.showPoints  = None
        obj.showMesh    = None
        obj._show       = None
        obj._treeItem   = None
        obj.face        = None
        obj.list_patchInfo = []

        return obj

    def __init__(self, geo=None):
        cad_geometry.__init__(self)
        self.list_patchInfo = []
        if geo is not None:
            for i in range(0, geo.npatchs):
                P = geo[i]
                # this is really a clone and not a copy
                patch = P.clone()
                self.append(patch)
                # add attributs from xml description when it exists
                self.list_patchInfo.append(PatchInfo(patch))
            self.set_internal_faces(geo.internal_faces)
            self.set_external_faces(geo.external_faces)
            self.set_connectivity(geo.connectivity)
            self.set_attributs(geo.attributs)

        self.MeshColor      = COLOR_DEFAULT_MESH
        self.NurbsColor     = Theme().color_viewer('default_patch') #COLOR_DEFAULT_GEOMETRY
        self.PointsColor    = COLOR_DEFAULT_POINTS

        self.PointsSize     = SIZE_GEOMETRY_POINTS

        self.showPoints = False
        self.showMesh = False
        self._show = True

        # needed for the tree inspector
        self._treeItem = None

        # needed when extracting faces for the instpector tree
        self.face = None

    def copy(self):
        geo = geometry.__new__(type(self))
        for i in range(0, self.npatchs):
            P = self._list[i]
            geo.append(P.copy())
        geo.set_internal_faces(self.internal_faces)
        geo.set_external_faces(self.external_faces)
        geo.set_connectivity(self.connectivity)

        geo.MeshColor   = self.MeshColor
        geo.NurbsColor  = self.NurbsColor
        geo.PointsColor = self.PointsColor
        geo.PointsSize  = self.PointsSize
        geo.showPoints  = self.showPoints
        geo.showMesh    = self.showMesh
        geo._show       = self._show
        geo._treeItem   = self._treeItem
        geo.face        = self.face
        geo.list_patchInfo = self.list_patchInfo

        return geo

    @property
    def treeItem(self):
        return self._treeItem

    def set_treeItem(self, item):
        self._treeItem = item

    def append(self, nrb):
        cad_geometry.append(self,nrb)
        self.list_patchInfo.append(PatchInfo(nrb))

    def load_attributs(self):
        # color nurbs
        self.NurbsColor = strtoArray(self.get_attribut('color-nurbs'))

        # Hide/Show geometry
        try:
            self._show = self.get_attribut('show')
            if self._show is None:
                self._show = True
        except:
            self._show = False

    def save_attributs(self):
        # color nurbs
        self.set_attribut('color-nurbs', str(self.NurbsColor))
        # Hide/Show geometry
        self.set_attribut('show', str(self.show))

    @property
    def show(self):
        return self._show

    def Show(self):
        self._show = True

    def Hide(self):
        self._show = False

    def ShowPoints(self):
        self.showPoints = True

    def HidePoints(self):
        self.showPoints = False

    def ShowMesh(self):
        self.showMesh = True

    def HideMesh(self):
        self.showMesh = False

    def set_meshColor(self,color):
        self.MeshColor = color

    def set_nurbsColor(self,color):
        self.NurbsColor = color
        self.set_attribut("color-nurbs", str(self.NurbsColor))

    def set_pointsColor(self,color):
        self.PointsColor = color

    def set_pointsSize(self,size):
        self.PointsSize = size

    def remove_patch(self, patch):
        index = self.index(patch)
        patchInfo = self.list_patchInfo[index]
        self.list_patchInfo.remove(patchInfo)
        self.remove(patch)

    def OnDelete(self, event):
        sel = self.listbox.GetSelection()
        if sel != -1:
            self.listbox.Delete(sel)

    def OnClear(self, event):
        self.listbox.Clear()

    def display(self):
        for i in range(0, self.npatchs):
            self.display_patch(i)

    def display_patch(self, i):
        pass

    def DrawControlPoints(self,nrb,PointsColor=None,alpha=ALPHA, blend=False):
        if PointsColor is None:
            PointsColor = self.PointsColor
        PointsColor = list(PointsColor) + [ALPHA]

        x,y,z = pointsToList(nrb.points)
        if blend:
            glEnable(GL_BLEND)
        glDisable(GL_DEPTH_TEST)
        glColor4f(*PointsColor)
        glBlendFunc(GL_SRC_ALPHA,GL_ONE)
        glPointSize(self.PointsSize)
        glBegin(GL_POINTS)
        for _x,_y,_z in zip(x,y,z):
            glVertex(_x,_y,_z)
        glEnd()
        if blend:
            glDisable(GL_BLEND)
        glEnable(GL_DEPTH_TEST)

    def GetEvaluator(self, nrb \
                     , MeshColor=None \
                     , NurbsColor=None \
                     , alpha=ALPHA \
                     , steps=None):
        if MeshColor is None:
            MeshColor = self.MeshColor
        if NurbsColor is None:
            NurbsColor = self.NurbsColor

        if nrb.dim == 1:
            evaluator = CurveEvaluator(nrb\
                                       , MeshColor=MeshColor\
                                       , NurbsColor=NurbsColor\
                                       , alpha=alpha\
                                       , steps=steps)
        if nrb.dim == 2:
            evaluator = SurfaceEvaluator(nrb\
                                       , MeshColor=MeshColor\
                                       , NurbsColor=NurbsColor\
                                       , alpha=alpha\
                                       , steps=steps)
        if nrb.dim == 3:
            evaluator = VolumeEvaluator(nrb\
                                       , MeshColor=MeshColor\
                                       , NurbsColor=NurbsColor\
                                       , alpha=alpha\
                                       , steps=steps)
        return evaluator

    def Draw(self, nrb=None \
             , MeshColor=None \
             , NurbsColor=None \
             , PointsColor=None \
             , alpha=ALPHA \
             , blend=False):

        if NurbsColor is None:
            if self.NurbsColor is None:
                NurbsColor = list(asarray(Theme().color_viewer('default_patch')).copy())
            else:
                NurbsColor = list(asarray(self.NurbsColor).copy())
        if self.show:
            if nrb is not None:
                list_nrb = [nrb]
            else:
                list_nrb = self._list

            for i in range(0,len(list_nrb)):
                nrb = list_nrb[i]
                nrbInfo = self.list_patchInfo[i]
                if nrbInfo.show:
                    _NurbsColor = asarray(NurbsColor).copy()
                    if nrbInfo.NurbsColor is not None:
                        _NurbsColor = asarray(nrbInfo.NurbsColor).copy()
                    NurbsSteps = nrbInfo.steps
                    evaluator = self.GetEvaluator(nrb \
                                                  , MeshColor=MeshColor \
                                                  , NurbsColor=_NurbsColor \
                                                  , alpha=alpha\
                                                  , steps=NurbsSteps)
                    showMesh = self.showMesh or nrbInfo.showMesh
                    evaluator.draw(mesh=showMesh, nurbs=True, blend=blend)
                if self.showPoints or nrbInfo.showPoints:
                    # Draw control points
                    self.DrawControlPoints(nrb \
                                           , PointsColor=PointsColor \
                                           ,alpha=alpha, blend=blend)

#######################################################################
if __name__ == '__main__':
    from caid.cad_geometry import square
    import matplotlib.pyplot as plt
    s1 = square(n=[3,3])
    geo1 = geometry()
    for nrb in s1:
        geo1.append(nrb)
    geo1.plotMesh() ; plt.show()
    geo2 = geometry(s1)
    geo2.plotMesh() ; plt.show()

max_stars_repo_name: ratnania/caid

max_issues_repo_name: pyccel/caid

max_forks_repo_name: pyccel/caid

reason: from numpy


text:
# -*- coding: utf-8 -*-
"""
Created on Sat Feb  7 15:18:06 2015

@author: Nate
"""
import numpy as np

class ParticleFilter(object):
    """Simple implementation of a particle filter with semi-intelligent
    resample."""
    def __init__(self, sys_model, meas_model=None, num_particles=1000,
                 initial_state=[0.0]):
        """Initialize particle filter class with system and measurement models
        as well as optionally specifying the number of particls and the initial
        state.

        Args:
            sys_model: a function that takes an array of particles, Xt and
                inputs Ut and returns Xt+1
            meas_model: a function that takes an array of particles, Xt and a
                measurement, and returns an array of weights corresponding to
                the input particles
            num_particles(optional): an integrer number of particles to be
                projected
            initial_state(optional): an array indicating the initial state of
                the system
        """
        self.M = num_particles # integer number of particles
        self.initialize_filter(initial_state) # set of particles
        #self.w = np.ones(self.M) # weight associated with each particle
        self._sys_model = sys_model # system model function
        self._meas_model = meas_model # measurement model function
        self._resample_test = resample_when_skewed(1.2) # a resample test

    def initialize_filter(self, initial_state=None, noise_array=0.0):
        """Method initializes the member self.X

        Args:
            initial_state: a list indicating the initial system state
        """
        if initial_state is not None:
            # save a copy of the filter's initial state
            self.x0 = initial_state
        else:
            # otherwise, reset the filter
            initial_state = self.x0
        # set the array of particles and weights
        self.X = np.tile(initial_state, [self.M, 1])
        self.w = np.ones(self.M)

    def sample(self, u):
        """Method samples given the system model and the input and advances
        belief state.

        Args:
            u: control input at the current timestep (should match expected
                input in system_model)

        Returns:
            a copy of st state array in case its needed by anyone
        """
        self.X = self._sys_model(self.X, u, self.M)
        # return in case anyone cares
        return self.X

    def assign_weight(self, z):
        """Method assigns an importance factor to each candidate state.

        Args:
            z: measurement at the current timestep, passed directly to the
                measurement model
        """
        self.w = np.multiply(self.w, self._meas_model(self.X, z))

    def resample(self):
        """Method checks the resample criteria (skewedness of distribution) and
        resamples the particles if necessary.

        No Arguments.
        """
        w = self.normalize_weight() # normalize the particle weights
        if self._resample_test(w):
            self.low_var_resample(w)

    def expectation(self):
        """Method returns the expectation of the state.

        Returns:
            E(X): scalar expected value of the distribution
        """
        w = self.normalize_weight()
        return [np.multiply(self.X[:,i],w).sum() for i in range(self.X.shape[1])]

    def normalize_weight(self):
        """Method returns the normalize weight array of the object.

        No Arguments
        """
        if self.w.sum() == 0.0:
            # if all particles have a weight of zero.. there is an issue
            return np.ones(self.M)/self.M
        elif self.w.sum() != 1.0:
            # if the sum of the weights is something other than 1.0, normalize
            return self.w/self.w.sum()
        else:
            # if the sum of the weights is 1.0, nothing need be done.
            return self.w

    def low_var_resample(self, w):
        """Method implements a low variance deterministic sampler as described
        in Thrun pg. 110.

        Args:
            w: associated weights normalized such that they sum to 1.0 (M)
        Returns:
            X: a set of particles representing the belief at t
        """
        # low_variance_sampler from Thrun pg. 110
        X = np.zeros(self.X.shape)
        r = np.random.uniform(0.0, 1.0/self.M)
        c = w[0]
        i = 0
        for m in range(self.M):
            U = r + (m)*(1.0/self.M)
            while U > c:
                i = i+1
                c = c + w[i]
            X[m] = self.X[i]
        # save the new beleif
        self.X = X
        # reset the weights
        self.w = np.ones(self.M)

def resample_when_skewed(threshold):
    '''Resample when the weights in a particle filter are too skewed.
    threshold: A numeric threshold that determines the tolerated skew above
    which particles should be resampled. If this is around 2, resampling
    occurs whenever about 25 % of the particles have about 75 % of the weight.
    Other easily interpreted thresholds are :
        50/50  1.0000 (i.e. resample wih every observation)
        60/40  1.1667
        70/30  1.7619
        80/20  3.2500
        90/10  8.1111
    '''
    # note, lambda functions are not needed...
    return lambda w: (len(w) * (w * w).sum()) > threshold

max_stars_repo_name: aivian/robots

max_issues_repo_name: aivian/robots

max_forks_repo_name: aivian/robots


text:
'''
    Plot predicted data from random airfoils using Graph Neural Networks and Deep Neural Networks
'''
import os, glob, random, json
import os.path as osp
from labellines import labelLines
import matplotlib.pyplot as plt
from pathlib import Path
import pickle
import pandas as pd
import torch
from performance import performance_predict_dnn, performance_predict_gnn
from MultiLayerLinear import MultiLayerLinear
from gnn_model import GnnModel
import numpy as np

Path('predicted_results').mkdir(parents=True, exist_ok=True)

'''
    Compare the checkpoints found here.

    Note:
        I did some rearranging of the files so that they were in separate folders
'''
# These 2 will be compared
dnn_no_cp = glob.glob('checkpoints_dnn_no_cp' + "/**/*.pt.tar", recursive = True)
gnn_no_cp = glob.glob('checkpoints_gnn_no_cp' + "/**/*.pt.tar", recursive = True)
dnn_no_cp.extend(gnn_no_cp)
compare_no_cp = dnn_no_cp
# These 2 will be compared
dnn_cp = glob.glob('checkpoints_dnn_cp' + "/**/*.pt.tar", recursive = True)
gnn_cp = glob.glob('checkpoints_gnn_cp' + "/**/*.pt.tar", recursive = True)
dnn_cp.extend(gnn_cp)
compare_cp = dnn_cp

def plot_airfoil_performance(predicted_results,folder,airfoil_name,Reynolds,Ncrit, plot_cp:bool):
    """Creates and saves an example plot

    Args:
        predicted_results (List[Dict[str,float]]): List of results generated by main code
        filename (str): filename to save the results as
    """
    plt.rcParams['font.size'] = '14'
    fig = plt.figure(figsize=(8,7),dpi=150,num=1,clear=True)
    ax1 = fig.add_subplot(111) # Airfoil
    ax1.set_title(predicted_results['airfoil_name'])
    ax1.plot(predicted_results['xss'], predicted_results['yss'],'.-',linewidth=2)
    ax1.plot(predicted_results['xps'], predicted_results['yps'],'.-',linewidth=2)
    ax1.set_ylabel('y/c')
    ax1.set_xlabel('x/c')
    ax1.set_xlim([0,1])
    ax1.set_ylim([-0.5,0.5])
    plt.savefig(os.path.join(folder,f'{airfoil_name}.png'),dpi=300)

    plt.clf()
    plt.gcf()

    fig = plt.figure(figsize=(18,15),dpi=150,num=2,clear=True)
    ax1 = fig.add_subplot(221) # Cl vs. alpha
    ax1.set_title("Cl vs. angle of attack")
    ax2 = fig.add_subplot(222) # Cd vs. alpha
    ax2.set_title("Cd vs. angle of attack")
    ax3 = fig.add_subplot(223) # Cdp vs. alpha
    ax3.set_title("Cdp vs. angle of attack")
    ax4 = fig.add_subplot(224) # Cm vs. alpha
    ax4.set_title("Cm vs. angle of attack")
    for model_data in predicted_results['models']:
        model_type = model_data['model_info']['type']
        scaler_type = model_data['model_info']['scaler_type']
        if model_type=='dnn':
            layer_size = model_data['model_info']['Layers'][0]
            nlayers = len(model_data['model_info']['Layers'])
            model_name = f'{scaler_type}-MLP-{layer_size}x{nlayers}'
        else:
            hidden_layers = model_data['model_info']['hiddenLayers']
            if hidden_layers:
                layer_size = model_data['model_info']['hiddenLayers'][0]
                nlayers = len(model_data['model_info']['hiddenLayers'])
                model_name = f'{scaler_type}-GNN-{layer_size}x{nlayers}'
            else:
                model_name = f'{scaler_type}-GNN-None'

        df = pd.DataFrame(model_data['predicted_polars']).sort_values(by='alpha')
        ax1.plot(df['alpha'],df['Cl'],'-', label=model_name,linewidth=1.1)
        ax2.plot(df['alpha'],df['Cd'],'-', label=model_name,linewidth=1.1)
        ax3.plot(df['alpha'],df['Cdp'],'-', label=model_name,linewidth=1.1)
        ax4.plot(df['alpha'],df['Cm'],'-', label=model_name,linewidth=1.1)

    df_actual = pd.DataFrame(predicted_results['actual_polars']).sort_values(by='alpha')
    model_name='xfoil'
    ax1.plot(df_actual['alpha'],df_actual['Cl'],'.-', label=model_name,linewidth=1.1)
    ax2.plot(df_actual['alpha'],df_actual['Cd'],'.-', label=model_name,linewidth=1.1)
    ax3.plot(df_actual['alpha'],df_actual['Cdp'],'.-', label=model_name,linewidth=1.1)
    ax4.plot(df_actual['alpha'],df_actual['Cm'],'.-', label=model_name,linewidth=1.1)

    ax1.set_ylabel('Coefficient of Lift (Cl)')
    ax1.set_xlabel('Angle of Attack (alpha)')

    ax2.set_ylabel('Coefficient of Drag (Cd)')
    ax2.set_xlabel('Angle of Attack (alpha)')

    ax3.set_ylabel('Coefficient of Drag zero lift (Cdp)')
    ax3.set_xlabel('Angle of Attack (alpha)')

    ax4.set_ylabel('Coefficient of Moment (Cm)')
    ax4.set_xlabel('Angle of Attack (alpha)')
    labelLines(ax1.get_lines(), zorder=2.5,fontsize=14)
    labelLines(ax2.get_lines(), zorder=2.5,fontsize=14)
    labelLines(ax3.get_lines(), zorder=2.5,fontsize=14)
    labelLines(ax4.get_lines(), zorder=2.5,fontsize=14)

    # ax1.legend(loc="upper left",fontsize=10)
    # ax2.legend(loc="upper left",fontsize=10)
    # ax3.legend(loc="upper left",fontsize=10)
    # ax4.legend(loc="upper left",fontsize=10)
    fig.tight_layout(pad=3.0)
    if plot_cp:
        plt.savefig(os.path.join(folder,f'{airfoil_name}-Re_{Reynolds}-Ncrit_{Ncrit}_cp_models.png'),dpi=300)
    else:
        plt.savefig(os.path.join(folder,f'{airfoil_name}-Re_{Reynolds}-Ncrit_{Ncrit}.png'),dpi=300)

    # Plot random examples of Cp
    if plot_cp:
        alphas = [p['alpha'] for p in predicted_results['models'][0]['predicted_polars']]
        alpha_selected = random.sample(alphas, 4) # Pick 3 random angle of attacks
        alpha_selected.sort()
        fig = plt.figure(figsize=(18,15),dpi=150,num=3,clear=True)
        ax1 = fig.add_subplot(221) # alpha 1
        ax2 = fig.add_subplot(222) # alpha 2
        ax3 = fig.add_subplot(223) # alpha 3
        ax4 = fig.add_subplot(224) # alpha 4
        ax1.set_title(f"{name}: Cp (Re {Reynolds} Ncrit {Ncrit} alpha {alpha_selected[0]})")
        ax2.set_title(f"{name}: Cp (Re {Reynolds} Ncrit {Ncrit} alpha {alpha_selected[1]})")
        ax3.set_title(f"{name}: Cp (Re {Reynolds} Ncrit {Ncrit} alpha {alpha_selected[2]})")
        ax4.set_title(f"{name}: Cp (Re {Reynolds} Ncrit {Ncrit} alpha {alpha_selected[3]})")

        for model_data in predicted_results['models']:
            model_type = model_data['model_info']['type']
            if model_type=='dnn':
                layer_size = model_data['model_info']['Layers'][0]
                nlayers = len(model_data['model_info']['Layers'])
                model_name = f'{scaler_type}-MLP-{layer_size}x{nlayers}'

                polars = pd.DataFrame(model_data['predicted_polars'])
                polar1 = polars[polars['alpha']==alpha_selected[0]]
                polar2 = polars[polars['alpha']==alpha_selected[1]]
                polar3 = polars[polars['alpha']==alpha_selected[2]]
                polar4 = polars[polars['alpha']==alpha_selected[3]]

                x = np.linspace(0,1,len(polar1['Cp_ss'].to_numpy()[0]))
                # Plot the polars
                ax1.plot(x, polar1['Cp_ss'].to_numpy()[0], label=model_name, linewidth=1.1)
                ax1.plot(x, polar1['Cp_ps'].to_numpy()[0], label=model_name, linewidth=1.1)

                ax2.plot(x, polar2['Cp_ss'].to_numpy()[0], label=model_name, linewidth=1.1)
                ax2.plot(x, polar2['Cp_ps'].to_numpy()[0], label=model_name, linewidth=1.1)

                ax3.plot(x, polar3['Cp_ss'].to_numpy()[0], label=model_name, linewidth=1.1)
                ax3.plot(x, polar3['Cp_ps'].to_numpy()[0], label=model_name, linewidth=1.1)

                ax4.plot(x, polar4['Cp_ss'].to_numpy()[0], label=model_name, linewidth=1.1)
                ax4.plot(x, polar4['Cp_ps'].to_numpy()[0], label=model_name, linewidth=1.1)

        model_name='xfoil'
        polar1 = df_actual[df_actual['alpha']==alpha_selected[0]]
        polar2 = df_actual[df_actual['alpha']==alpha_selected[1]]
        polar3 = df_actual[df_actual['alpha']==alpha_selected[2]]
        polar4 = df_actual[df_actual['alpha']==alpha_selected[3]]
        x = np.linspace(0,1,len(polar1['Cp_ss'].to_numpy()[0]))

        ax1.plot(x,polar1['Cp_ss'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)
        ax1.plot(x,polar1['Cp_ps'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)

        ax2.plot(x,polar2['Cp_ss'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)
        ax2.plot(x,polar2['Cp_ps'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)

        ax3.plot(x,polar3['Cp_ss'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)
        ax3.plot(x,polar3['Cp_ps'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)

        ax4.plot(x,polar4['Cp_ss'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)
        ax4.plot(x,polar4['Cp_ps'].to_numpy()[0],'.-', label=model_name,linewidth=1.1)

        labelLines(ax1.get_lines(), zorder=2.5,fontsize=14)
        labelLines(ax2.get_lines(), zorder=2.5,fontsize=14)
        labelLines(ax3.get_lines(), zorder=2.5,fontsize=14)
        labelLines(ax4.get_lines(), zorder=2.5,fontsize=14)

        fig.tight_layout(pad=3.0)
        plt.savefig(os.path.join(folder,f'{airfoil_name}-Re_{Reynolds}-Ncrit_{Ncrit}_Cp.png'),dpi=300)

def get_random_airfoil_data():
    """Retrieves a random airfoil by searching through the JSON

    Returns:
        [type]: [description]
    """
    # Load a random airfoil
    all_json = glob.glob('../generate_xfoil/json/*.json')
    # Plot random data
    filename = all_json[random.randrange(0, len(all_json), 1)]

    with open(filename,'r') as f:
        airfoil = json.load(f)

        name = airfoil['name']
        xss = airfoil['xss']
        yss = airfoil['yss']

        xps = airfoil['xps']
        yps = airfoil['yps']

        polars = pd.DataFrame(airfoil['polars'])
        unique_reynolds = polars['Re'].unique()
        reynolds = unique_reynolds[random.randrange(0, len(unique_reynolds), 1)]

        polars = polars[polars['Re'] == reynolds]
        unique_ncrit = polars['Ncrit'].unique()
        ncrit = unique_ncrit[random.randrange(0, len(unique_ncrit), 1)]
        polars = polars[polars['Ncrit'] == ncrit]

    return name,xss,yss,xps,yps,reynolds,ncrit,polars,airfoil

def compare_models(saved_models,name,xss,yss,xps,yps,scalers):
    """[summary]

    Args:
        saved_models (Dict): Dictionary
        name (str): Name describing the geometry
        xss (np.ndarray): [description]
        yss (np.ndarray): [description]
        xps (np.ndarray): [description]
        yps (np.ndarray): [description]
        scalers ([type]): [description]
    """
    plot_cp = False
    models_to_evaluate = list()
    for filename in saved_models:
        model_data = torch.load(filename)
        train_params = model_data['parameters']
        if "Gnn" in filename:
            # Load the model_settings used for training into memory
            linear_layer = MultiLayerLinear(in_channels=train_params['input_size']*train_params['GnnLayers'][1],out_channels=train_params['output_size'],h_sizes=train_params['hiddenLayers'])
            train_params['type']='gnn'
            if 'minmax' in filename:
                train_params['scaler_type']='minmax'
                train_params['model']=GnnModel(train_params['input_size'],train_params['GnnLayers'],linear_layers=linear_layer)
                train_params['model'].load_state_dict(model_data['state_dict'])
            else:
                train_params['scaler_type']='standard'
                train_params['model']=GnnModel(train_params['input_size'],train_params['GnnLayers'],linear_layers=linear_layer)
                train_params['model'].load_state_dict(model_data['state_dict'])
        else:
            train_params['type']='dnn'
            if 'minmax' in filename:
                train_params['scaler_type']='minmax'
                train_params['model']=MultiLayerLinear(in_channels=train_params['input_size'],out_channels=train_params['output_size'],h_sizes=train_params['Layers'])
                train_params['model'].load_state_dict(model_data['state_dict'])
            else:
                train_params['scaler_type']='standard'
                train_params['model']=MultiLayerLinear(in_channels=train_params['input_size'],out_channels=train_params['output_size'],h_sizes=train_params['Layers'])
                train_params['model'].load_state_dict(model_data['state_dict'])
        train_params['model_name']=str(train_params['model'])
        train_params['use_cp'] = True if train_params['output_size']>4 else False
        plot_cp = train_params['use_cp']
        models_to_evaluate.append(train_params)

    ''' This code below is used to evaluate all the models, all models predict the polar on a random airfoil '''

    if len(models_to_evaluate)>0:
        predicted_results = {'airfoil_name': name, 'xss':xss, 'yss':yss,'xps':xps,'yps':yps,'models':[]}
        for model_data in models_to_evaluate:
            if model_data['scaler_type'] == 'minmax':
                scaler = scalers['min_max']
                scaler_cp = scalers['min_max_cp']
            else:
                scaler = scalers['standard']
                scaler_cp = scalers['standard_cp']

            predicted_results['model'] = {'type':model_data['type'],'scaler_type':model_data['scaler_type']}

            results = list()
            actual_results = list()
            for p in range(len(polars)):
                alpha = polars.iloc[p]['alpha']
                Re = polars.iloc[p]['Re']
                Ncrit = polars.iloc[p]['Ncrit']

                if model_data['type'] == 'gnn':
                    Cl, Cd, Cdp, Cm, Cp_ss, Cp_ps = performance_predict_gnn(model_data['model'],xss,yss,xps,yps,alpha,Re,Ncrit,scaler,scaler_cp,model_data['use_cp'])
                else:
                    Cl, Cd, Cdp, Cm, Cp_ss, Cp_ps = performance_predict_dnn(model_data['model'],xss,yss,xps,yps,alpha,Re,Ncrit,scaler,scaler_cp,model_data['use_cp'])

                results.append({'alpha':alpha,'Re':Re,'Ncrit':Ncrit,
                                'Cl':Cl,
                                'Cd':Cd,
                                'Cm':Cm,
                                'Cdp':Cdp,
                                'Cp_ss':Cp_ss,
                                'Cp_ps':Cp_ps})

                actual_results.append({'alpha':alpha,'Re':Re,'Ncrit':Ncrit,
                                    'Cl':polars.iloc[p]['Cl'],
                                    'Cd':polars.iloc[p]['Cd'],
                                    'Cm':polars.iloc[p]['Cm'],
                                    'Cdp':polars.iloc[p]['Cdp'],
                                    'Cp_ss':polars.iloc[p]['Cp_ss'],
                                    'Cp_ps':polars.iloc[p]['Cp_ps']})
            predicted_results['models'].append({'model_info':model_data, 'predicted_polars':results})
        predicted_results['actual_polars'] = actual_results

        '''
            Create a Plot to compare all models, save to png
        '''
        plot_airfoil_performance(predicted_results,'predicted_results',name,Re,Ncrit,plot_cp)

if __name__ == "__main__":
    # Pick a random json file
    # Make a folder with the same name as this file. This is where we save al the images

    # Load the scalers into memory
    with open(osp.join('../generate_xfoil/','scalers.pickle'),'rb') as f: # min max of individual y and Cp positions
        scalers = pickle.load(f)

    polars= pd.DataFrame()
    while (polars.empty):   # Make sure I get a random airfoil with data
        name,xss,yss,xps,yps,reynolds,ncrit,polars,airfoil = get_random_airfoil_data()

    # Load all the models into memory for later prediction
    compare_models(compare_no_cp,name,xss,yss,xps,yps,scalers)
    compare_models(compare_cp,name,xss,yss,xps,yps,scalers)

max_stars_repo_name: nasa/airfoil-learning

max_issues_repo_name: nasa/airfoil-learning

max_forks_repo_name: nasa/airfoil-learning