text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 13 17:07:35 2019

@author: arjunbalaji
"""
#caps net idea for oct lumen profiler
#hi
#this iteration of code runs a conv capsule idea

import numpy as np
import torch

###############################################################################
#need this special squash function for caps net

def squash(s, axis = -1, epsilon = 1e-7):
    squared_norm = torch.sum(s*s, dim = axis)

    safe_norm = torch.sqrt(squared_norm + epsilon)

    squash_factor = squared_norm / (1. + squared_norm)

    #the unsqueeze here is very important!
    #for safe norm to be broadcasted appropriately
    unit_vector = torch.div(s, safe_norm.unsqueeze(-1))

    #for squash factor to be broadcasted appropriately
    return torch.mul(squash_factor.unsqueeze(-1), unit_vector)

###############################################################################
class Get_Primary_Caps(torch.nn.Module):
    """This is the primary caps block. It takes in an input image of 1 channel
    and 512x512 pixels and outputs a caps1_n_caps primary caps each which is a
    caps1_n_dims dimensional vector. there is work to be done here so that
    the numbers all make themselves work. at the moment you have to carefully
    check each one to make sure the model runs"""
    def __init__(self,
                 input_channels,
                 caps1_n_maps,
                 caps1_caps_grid_ydim,
                 caps1_caps_grid_xdim,
                 caps1_n_dims):
        super(Get_Primary_Caps, self).__init__()
        self.input_channels = input_channels
        self.caps1_n_maps = caps1_n_maps
        self.caps1_caps_grid_ydim = caps1_caps_grid_ydim
        self.caps1_caps_grid_xdim = caps1_caps_grid_xdim
        self.caps1_n_dims = caps1_n_dims

        self.relu = torch.nn.ReLU()

        #these must be calculated from the input shape and the convs we use!!!
        #this is so important dont get it wrong
        #O = (W - K - 2P)/S +1

        conv1_parameters = {'i': self.input_channels, 'o': 32, 'k': 5, 's': 2, 'p':2}
        self.conv1 = torch.nn.Conv2d(in_channels=conv1_parameters['i'],
                                     out_channels=conv1_parameters['o'],
                                     kernel_size=conv1_parameters['k'],
                                     stride=conv1_parameters['s'],
                                     padding=conv1_parameters['p'])

        conv2_parameters = {'i': 32, 'o': self.caps1_n_dims * self.caps1_n_maps, 'k': 5, 's': 2, 'p':2}
        self.conv2 = torch.nn.Conv2d(in_channels=conv2_parameters['i'],
                                     out_channels=conv2_parameters['o'],
                                     kernel_size=conv2_parameters['k'],
                                     stride=conv2_parameters['s'],
                                     padding=conv2_parameters['p'])

    def output_params(self):
        return {'maps out': self.caps1_n_maps,
                'caps dim': self.caps1_n_dims,
                'h': self.caps1_caps_grid_ydim,
                'w': self.caps1_caps_grid_xdim}

    def forward(self, x):
        #print(x.size())
        x = self.conv1(x)
        #print(x.size())
        x = self.relu(x)
        x = self.conv2(x)
        #print(x.size())
        x = self.relu(x)

        '''
        THIS IS IMPORTANT !!!!!!!!!!!!!!!!!!
        if we just went ;
        x = x.view([-1,
                    self.caps1_n_maps,
                    self.caps1_caps_grid_ydim,
                    self.caps1_caps_grid_xdim,
                    self.caps1_n_dims])
        then pytorch replicates the images and you get this weird ass grid.
        so instead we view it as below, then transpose the axes till we get
        what we want
        '''
        x = x.view([-1,
                    self.caps1_n_maps,
                    self.caps1_n_dims,
                    self.caps1_caps_grid_ydim,
                    self.caps1_caps_grid_xdim])
        #x = torch.transpose(x, 3, 4)
        #x = torch.transpose(x, 2, -1)
        x = x.permute(0, 1, 3, 4, 2)
        #x = squash(x)
        return x

###############################################################################
class Get_Abstract_Caps_Down(torch.nn.Module):
    """This is the abstract caps layer. We take in an input of the capsules
    of the previous layer and then output predictions of abstract capsules."""
    def __init__(self,
                 batch_size,
                 capsin_n_maps,
                 capsin_n_dims,
                 capsout_n_maps,
                 capsout_n_dims,
                 old_h,
                 old_w,
                 y_kernel,
                 x_kernel,
                 stride,
                 padding,
                 across = True):
        super(Get_Abstract_Caps_Down, self).__init__()

        self.batch_size = batch_size
        self.capsin_n_dims = capsin_n_dims
        self.capsin_n_maps = capsin_n_maps
        self.capsout_n_maps = capsout_n_maps
        self.capsout_n_dims = capsout_n_dims

        self.old_h = old_h
        self.old_w = old_w
        self.y_kernel = y_kernel
        self.x_kernel = x_kernel
        self.stride = stride
        self.padding = padding
        self.across = across
        '''
        #this is our weight matrix that goes into our conv layer to get
        #prediction vectors
        self.W = torch.nn.Parameter(torch.Tensor(self.capsout_n_maps * self.capsout_n_dims,
                                                      self.capsin_n_maps,
                                                      self.y_kernel,
                                                      self.x_kernel))
        '''
        #this is a SPECIAL BIAS! it doesnt go into our conv opertaion
        #we replicate it later and use it for the caps
        self.bias = torch.nn.Parameter(torch.Tensor(1,
                                                    self.capsout_n_maps,
                                                    1,
                                                    1,
                                                    self.capsout_n_dims))

        self.reset_parameters()

        #self.predict_vectors = Predict_Vectors_Down(self.W, self.bias)
        self.capsconv2d_down = torch.nn.Conv2d(in_channels = self.capsin_n_maps*self.capsin_n_dims,
                                          out_channels = self.capsout_n_maps*self.capsout_n_dims,
                                          kernel_size = (self.y_kernel, self.x_kernel),
                                          stride = self.stride,
                                          padding = self.padding,
                                          bias = True)
        '''
        if self.across:
            self.capsconv2d_across = torch.nn.Conv2d(in_channels = self.capsout_n_maps*self.capsout_n_dims,
                                                     out_channels = self.capsin_n_maps*self.capsout_n_maps*self.capsout_n_dims,
                                                     kernel_size = 1,
                                                     stride = 1,
                                                     padding = 0,
                                                     bias = False)
        '''

        #calculate final prediction heights and widths for routing
        self.new_hl = (self.old_h - self.y_kernel + 2*self.padding)/self.stride + 1
        #self.new_hl = (self.new_hl - 1 + 0) / 1 + 1

        self.new_wl = (self.old_w - self.x_kernel + 2*self.padding)/self.stride + 1
        #self.new_wl = (self.new_wl - 1 + 0) / 1 + 1

        #init routing algorithm
        self.routing = Agreement_Routing_Down(bias=self.bias,
                                              input_caps_maps = self.capsin_n_maps,
                                              input_caps_dim = self.capsin_n_dims,
                                              output_caps_maps = self.capsout_n_maps,
                                              output_caps_dim = self.capsout_n_dims,
                                              new_hl = self.new_hl,
                                              new_wl = self.new_wl,
                                              num_iterations = 2)

    def infer_shapes(self):
        return {'caps maps': self.capsout_n_maps,
                'caps dims': self.capsout_n_dims,
                'h': self.new_hl,
                'w': self.new_wl}

    def reset_parameters(self):
        stdv = 1. / np.sqrt(self.capsin_n_maps)
        #self.W.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x):
        #print('no we doing abstract caps1')
        #print(self.W.size(), 'weight size')
        #print(self.bias.size(), 'bias size')
        #print(x.size(), 'caps before conv capsule preds' )
        #x = self.predict_vectors(primary_caps)
        #print(self.new_hl)
        batch, input_maps, hold, wold, input_capdims = x.size()

        #these bits are becaause of the view 'replication' issue i was having
        #print(batch, input_maps, hold, wold, input_capdims)
        #x = torch.transpose(x, -1, 2)
        #x = torch.transpose(x, 3, 4)
        x = x.permute(0, 1, 4, 2, 3)
        #print(x.size())

        x = x.contiguous().view([-1,
                    input_maps * input_capdims,
                    hold,
                    wold])

        x = self.capsconv2d_down(x)
        #print(x.size(), 'after down conv')

        #if self.across:
        #    x = self.capsconv2d_across(x)
            #print(x.size(), 'after across conv')

        _, _, hnew, wnew = x.size()

        #view issue
        x = x.view([batch,
                    self.capsout_n_maps,
                    self.capsout_n_dims,
                    hnew,
                    wnew,
                    -1])

        x = x.permute(0, 1, 3, 4, 5, 2)
        #print(x.size())
        '''
        x = x.view([batch,
                    self.capsout_n_maps,
                    hnew,
                    wnew,
                    -1,
                    self.capsout_n_dims])
        '''
        #print(x.size(), 'resizing to normal')

        x = self.routing(x)

        return x

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

class Agreement_Routing_Down(torch.nn.Module):
    """This is the localised agreement routing algorithm. It takes in the total
    prediction vectors from a layer l and computes the routing weights for
    those predictions. It then squashes the prediction vectors using the
    custom squash function."""
    def __init__(self, bias,
                 input_caps_maps,
                 input_caps_dim,
                 output_caps_maps,
                 output_caps_dim,
                 new_hl,
                 new_wl,
                 num_iterations):
        super(Agreement_Routing_Down, self).__init__()

        self.input_caps_maps = input_caps_maps
        self.input_caps_dim = input_caps_dim
        self.output_caps_maps = output_caps_maps
        self.output_caps_dim = output_caps_dim
        self.new_hl = int(new_hl)
        self.new_wl = int(new_wl)
        self.num_iterations = num_iterations
        self.softmax = torch.nn.Softmax(dim = -1)
        #self.bias = bias.repeat((1, 1, self.new_hl, self.new_wl, 1))

        self.b = torch.nn.Parameter(torch.zeros((1,
                                                 self.output_caps_maps,
                                                 self.new_hl,
                                                 self.new_wl,
                                                 self.input_caps_maps))) #  , device = 'cuda:0')

        #self.bias = self.bias.repeat((1, 1, self.new_hl, self.new_wl, 1))

    def forward(self, tensor_of_prediction_vector):
        #print('now we doing routing')
        c = self.softmax(self.b)

        #print(c.size(),'c')
        #print(tensor_of_prediction_vector.size(), 'pred vectors')

        #output_vectors = torch.mul(tensor_of_prediction_vector, c.unsqueeze(-1))
        output_vectors = torch.mul(c.unsqueeze(-1), tensor_of_prediction_vector)

        #print(output_vectors.size(), 'output vectors after *c')
        #print(self.bias.size(), 'repeated bias size')

        output_vectors = output_vectors.sum(dim=-2) #+ self.bias

        #print(output_vectors.size(), 'output vectors after sum')

        #print(output_vectors.size())
        output_vectors = squash(output_vectors, axis = -1)
        #print(output_vectors.size(), 'should be good')
        b_batch = self.b

        #print('routing loop')
        for d in range(self.num_iterations):

            #print('preds', tensor_of_prediction_vector.size())
            #print('out', output_vectors.size())

            b_batch = b_batch + torch.mul(tensor_of_prediction_vector,
                                output_vectors.unsqueeze(-2)).sum(dim = -1)
            '''
            distances = torch.mul(tensor_of_prediction_vector,
                                output_vectors.unsqueeze(-2)).sum(dim = -1)

            self.b = torch.add(self.b, distances)
            '''
            #b_batch = b_batch.sum(dim = -1)

            #b_batch =
            #print('bbatch', b_batch.size())

            c = self.softmax(b_batch)
            #print(c.size())
            output_vectors = torch.mul(tensor_of_prediction_vector,
                                       c.unsqueeze(-1))

            output_vectors = output_vectors.sum(-2)
            output_vectors = squash(output_vectors,
                                    axis = -1)
            #print(c.size())
        #print(tensor_of_prediction_vector.size())
        #print(output_vectors.size())
        self.c = c

        return output_vectors
###############################################################################

class Get_Abstract_Caps_Up(torch.nn.Module):
    """This is the abstract caps layer. We take in an input of the capsules
    of the previous layer and then output predictions of abstract capsules.

        notes:
            . padding must be tuple (x,y)
            . if uptype == 'upsample' used
                > padding is autocalculated, even with user input
                > stride = 1 is set, even with user stride != 1
            .
            . output_padding only required if uptype == 'deconv'
            . output_padding, if used, must be tuple (x,y)
                """
    def __init__(self,
                 batch_size,
                 capsin_n_maps,
                 capsin_n_dims,
                 capsout_n_maps,
                 capsout_n_dims,
                 old_h,
                 old_w,
                 y_kernel,
                 x_kernel,
                 stride,
                 padding, #MUST E A TUPLE CUNNY
                 output_padding, #outpadding must be a TUPLE BIATCH, also i
                 uptype,
                 across = True):
        super(Get_Abstract_Caps_Up, self).__init__()

        self.batch_size = batch_size
        self.capsin_n_dims = capsin_n_dims
        self.capsin_n_maps = capsin_n_maps
        self.capsout_n_maps = capsout_n_maps
        self.capsout_n_dims = capsout_n_dims

        self.old_h = old_h
        self.old_w = old_w
        self.y_kernel = y_kernel
        self.x_kernel = x_kernel
        self.stride = stride
        self.padding = padding
        self.output_padding = output_padding
        self.uptype = uptype
        self.across = across

        #this is a SPECIAL BIAS! it doesnt go into our conv opertaion
        #we replicate it later and use it for the caps
        self.bias = torch.nn.Parameter(torch.Tensor(1,
                                                    self.capsout_n_maps,
                                                    1,
                                                    1,
                                                    self.capsout_n_dims))

        self.reset_parameters()

        pady = self.padding[0]
        padx = self.padding[1]

        if self.uptype == 'deconv':
            '''
            HUGELY IMPORTANT TO GET DECONV TO WORK WE ADD 1 to the end
            '''
            self.new_hl = (self.old_h-1)*self.stride + self.y_kernel - 2*pady + self.output_padding[0]
            self.new_wl = (self.old_w-1)*self.stride + self.x_kernel - 2*padx + self.output_padding[1]

            self.capsconv2d_up = torch.nn.ConvTranspose2d(in_channels = self.capsin_n_maps*self.capsin_n_dims,
                                          out_channels = self.capsout_n_maps*self.capsout_n_dims,
                                          kernel_size = (self.y_kernel, self.x_kernel),
                                          stride = self.stride,
                                          padding = self.padding,
                                          output_padding = self.output_padding,
                                          bias = True)

        elif self.uptype == 'upsample':
            #pady = int((self.y_kernel - 1)/2)
            #padx = int((self.x_kernel - 1)/2)
            #self.padding = (pady, padx)
            self.new_hl = 2*self.old_h
            self.new_wl = 2*self.old_w
            self.capsconv2d_up = torch.nn.Conv2d(in_channels = self.capsin_n_maps*self.capsin_n_dims,
                                          out_channels = self.capsout_n_maps*self.capsout_n_dims,
                                          kernel_size = (self.y_kernel, self.x_kernel),
                                          stride = 1,
                                          padding = self.padding,
                                          bias = True)

        #init routing algorithm
        self.routing = Agreement_Routing_Down(bias=self.bias,
                                              input_caps_maps = self.capsin_n_maps,
                                              input_caps_dim = self.capsin_n_dims,
                                              output_caps_maps = self.capsout_n_maps,
                                              output_caps_dim = self.capsout_n_dims,
                                              new_hl = self.new_hl,
                                              new_wl = self.new_wl,
                                              num_iterations = 2)

    def infer_shapes(self):
        return {'caps maps': self.capsout_n_maps,
                'caps dims': self.capsout_n_dims,
                'h': self.new_hl,
                'w': self.new_wl}

    def reset_parameters(self):
        stdv = 1. / np.sqrt(self.capsin_n_maps)
        #self.W.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)

    def forward(self, x):
        #print('no we doing abstract caps2u')

        #print(self.bias.size(), 'bias size')
        #print(x.size(), 'caps before conv capsule preds' )
        #x = self.predict_vectors(primary_caps)
        #print(self.new_hl)
        batch, input_maps, hold, wold, input_capdims = x.size()

        x = x.permute(0, 1, 4, 2, 3)
        #print(x.size())

        x = x.contiguous().view([-1,
                    input_maps * input_capdims,
                    hold,
                    wold])
        '''
        x = x.view([-1,
                    input_maps * input_capdims,
                    hold,
                    wold])
        '''
        if self.uptype == 'deconv':
            x = self.capsconv2d_up(x)
            #print(x.size(), 'after up deconv')

        elif self.uptype == 'upsample':
            #print(x.size(),'pre up')
            x = torch.nn.functional.upsample(x,
                                             size=[self.new_hl, self.new_wl],
                                             mode='bilinear')
            #print(x.size())
            x = self.capsconv2d_up(x)

        '''
        if self.across:
            x = self.capsconv2d_across(x)
            #print(x.size(), 'after across conv')
        '''

        #just
        _, _, hnew, wnew = x.size()

        if self.new_hl != hnew:
            print('Something funny going on with user defined hnew and actual')

        if self.new_wl != wnew:
            print('Something funny going on with user defined wnew and actual')

        #x = x.permute(0, 1, 3, 4, 5, 2)
        x = x.view([batch,
                    self.capsout_n_maps,
                    self.capsout_n_dims,
                    hnew,
                    wnew,
                    -1])

        x = x.permute(0, 1, 3, 4, 5, 2)

        #print(x.size(), 'resizing to normal')

        x = self.routing(x)
        #print(x.size(), 'out')
        return x

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

class Reconstruction_Layer(torch.nn.Module):
    """TThis is the reconstruction layer for the network to learn how to remake
    the original input image"""
    def __init__(self,
                 batch_size,
                 capsin_n_maps,
                 capsin_n_dims,
                 reconstruct_channels):
        super(Reconstruction_Layer, self).__init__()

        self.batch_size = batch_size
        self.capsin_n_dims = capsin_n_dims
        self.capsin_n_maps = capsin_n_maps
        self.reconstruct_channels = reconstruct_channels
        self.relu = torch.nn.ReLU()
        self.sigmoid = torch.nn.Sigmoid()

        self.conv1_params = {'i':int(self.capsin_n_maps*self.capsin_n_dims),
                             'o':64,
                             'k':1,
                             's':1,
                             'p':0}
        self.conv1 = torch.nn.Conv2d(in_channels = self.conv1_params['i'],
                                     out_channels = self.conv1_params['o'],
                                     kernel_size = self.conv1_params['k'],
                                     stride = self.conv1_params['s'],
                                     padding = self.conv1_params['p'])

        self.conv2_params = {'i':int(self.conv1_params['o']),
                             'o':128,
                             'k':1,
                             's':1,
                             'p':0}
        self.conv2 = torch.nn.Conv2d(in_channels = self.conv2_params['i'],
                                     out_channels = self.conv2_params['o'],
                                     kernel_size = self.conv2_params['k'],
                                     stride = self.conv2_params['s'],
                                     padding = self.conv2_params['p'])

        self.conv3_params = {'i':int(self.conv2_params['o']),
                             'o':self.reconstruct_channels,
                             'k':1,
                             's':1,
                             'p':0}
        self.conv3 = torch.nn.Conv2d(in_channels = self.conv3_params['i'],
                                     out_channels = self.conv3_params['o'],
                                     kernel_size = self.conv3_params['k'],
                                     stride = self.conv3_params['s'],
                                     padding = self.conv3_params['p'])

    def forward(self, x):

        _, _, h, w, _ = x.size()
        #print(x.size())
        x = x.permute(0, 1, 4, 2, 3)
        #print(x.size())
        x = x.contiguous().view([-1,
                    self.capsin_n_maps * self.capsin_n_dims,
                    h,
                    w])
        #print(x.size())
        #x = torch.transpose(x, 3, 4)
        #x = torch.transpose(x, 2, -1)

        x = self.conv1(x)
        x = self.relu(x)

        x = self.conv2(x)
        x = self.relu(x)

        x = self.conv3(x)
        x = self.sigmoid(x)

        return x
###############################################################################

class CapsNet(torch.nn.Module):
    """This is the actual model. It has a down line, a bottom pass and then a
    series of up passes. On the up passes we concatenate prior down passes as
    this improves the networks localisation. it is important the the tensors we
    concatenate are the same size. so we use upsampling. Also be aware of the
    channels, we want a lot of channels (~1000) so the network learns intricate
    features."""
    def __init__(self, opt):
        super(CapsNet, self).__init__()
        self.opt = opt

        if self.opt.transforms:
            caps1_ygrid = self.opt.c_size[0]
            caps1_xgrid = self.opt.c_size[1]
        else:
            caps1_ygrid = self.opt.start_size[0]
            caps1_xgrid = self.opt.start_size[1]

        self.get_prim_caps = Get_Primary_Caps(input_channels=self.opt.inputchannels,
                                              caps1_n_maps = self.opt.primmaps,
                                              caps1_caps_grid_ydim = int(caps1_ygrid / 4),
                                              caps1_caps_grid_xdim = int(caps1_xgrid / 4),
                                              caps1_n_dims = self.opt.primdims)
        prim_params = self.get_prim_caps.output_params()

        self.get_abstract_caps1 = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = prim_params['maps out'],
                                                         capsin_n_dims = prim_params['caps dim'],
                                                         capsout_n_maps = self.opt.maps1,
                                                         capsout_n_dims = self.opt.dims1,
                                                         old_h = prim_params['h'],
                                                         old_w = prim_params['w'],
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=True)
        caps1_params = self.get_abstract_caps1.infer_shapes()

        self.get_abstract_caps1a = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps1_params['caps maps'],
                                                         capsin_n_dims = caps1_params['caps dims'],
                                                         capsout_n_maps = caps1_params['caps maps'],
                                                         capsout_n_dims = caps1_params['caps dims'],
                                                         old_h = caps1_params['h'],
                                                         old_w = caps1_params['w'],
                                                         y_kernel = 6,
                                                         x_kernel = 6,
                                                         stride = 2,
                                                         padding = 2,
                                                         across=True)
        caps1a_params = self.get_abstract_caps1a.infer_shapes()

        self.get_abstract_caps2 = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps1a_params['caps maps'],
                                                         capsin_n_dims = caps1a_params['caps dims'],
                                                         capsout_n_maps = self.opt.maps2,
                                                         capsout_n_dims =  self.opt.dims2,
                                                         old_h = int(caps1a_params['h']),
                                                         old_w = int(caps1a_params['w']),
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=False)
        caps2_params = self.get_abstract_caps2.infer_shapes()

        self.get_abstract_caps2a = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps2_params['caps maps'],
                                                         capsin_n_dims = caps2_params['caps dims'],
                                                         capsout_n_maps = caps2_params['caps maps'],
                                                         capsout_n_dims = caps2_params['caps dims'],
                                                         old_h = caps2_params['h'],
                                                         old_w = caps2_params['w'],
                                                         y_kernel = 6,
                                                         x_kernel = 6,
                                                         stride = 2,
                                                         padding = 2,
                                                         across=True)
        caps2a_params = self.get_abstract_caps2a.infer_shapes()

        self.get_abstract_caps3 = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps2a_params['caps maps'],
                                                         capsin_n_dims = caps2a_params['caps dims'],
                                                         capsout_n_maps = self.opt.maps3,
                                                         capsout_n_dims = self.opt.dims3,
                                                         old_h = int(caps2a_params['h']),
                                                         old_w = int(caps2a_params['w']),
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=False)
        caps3_params = self.get_abstract_caps3.infer_shapes()

        self.get_abstract_caps3a = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps3_params['caps maps'],
                                                         capsin_n_dims = caps3_params['caps dims'],
                                                         capsout_n_maps = caps3_params['caps maps'],
                                                         capsout_n_dims = caps3_params['caps dims'],
                                                         old_h = caps3_params['h'],
                                                         old_w = caps3_params['w'],
                                                         y_kernel = 6,
                                                         x_kernel = 6,
                                                         stride = 2,
                                                         padding = 2,
                                                         across=True)
        caps3a_params = self.get_abstract_caps3a.infer_shapes()

        self.get_abstract_caps_bot = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps3a_params['caps maps'],
                                                         capsin_n_dims = caps3a_params['caps dims'],
                                                         capsout_n_maps = caps3a_params['caps maps'],
                                                         capsout_n_dims = caps3a_params['caps dims'],
                                                         old_h = int(caps3a_params['h']),
                                                         old_w = int(caps3a_params['w']),
                                                         y_kernel = 1,
                                                         x_kernel = 1,
                                                         stride = 1,
                                                         padding = 0,
                                                         across=True)
        capsbot_params = self.get_abstract_caps_bot.infer_shapes()

        if self.opt.uptype == 'upsample':
            upa_outputpadding = (0,0)
        elif self.opt.uptype == 'deconv':
            upa_outputpadding = (1,1)

        self.get_abstract_caps3u = Get_Abstract_Caps_Up(self.opt.batch_size,
                                                         capsin_n_maps = capsbot_params['caps maps'] + caps3_params['caps maps'],
                                                         capsin_n_dims = capsbot_params['caps dims'],
                                                         capsout_n_maps = caps2_params['caps maps'],
                                                         capsout_n_dims = caps2_params['caps dims'],
                                                         old_h = int(capsbot_params['h']),
                                                         old_w = int(capsbot_params['w']),
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 2,
                                                         padding = (2,2),
                                                         output_padding=upa_outputpadding,
                                                         uptype = self.opt.uptype,
                                                         across=True)
        caps3u_params = self.get_abstract_caps3u.infer_shapes()

        self.get_abstract_caps3ua = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps3u_params['caps maps'],
                                                         capsin_n_dims = caps3u_params['caps dims'],
                                                         capsout_n_maps = caps3u_params['caps maps'],
                                                         capsout_n_dims = caps3u_params['caps dims'],
                                                         old_h = caps3u_params['h'],
                                                         old_w = caps3u_params['w'],
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=True)
        caps3ua_params = self.get_abstract_caps3ua.infer_shapes()

        self.get_abstract_caps2u = Get_Abstract_Caps_Up(self.opt.batch_size,
                                                         capsin_n_maps = caps3ua_params['caps maps'] + caps2_params['caps maps'],
                                                         capsin_n_dims = caps3ua_params['caps dims'],
                                                         capsout_n_maps = caps1_params['caps maps'],
                                                         capsout_n_dims = caps1_params['caps dims'],
                                                         old_h = int(caps3ua_params['h']),
                                                         old_w = int(caps3ua_params['w']),
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 2,
                                                         padding = (2,2),
                                                         output_padding=upa_outputpadding,
                                                         uptype = self.opt.uptype,
                                                         across=True)
        caps2u_params = self.get_abstract_caps2u.infer_shapes()

        self.get_abstract_caps2ua = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps2u_params['caps maps'],
                                                         capsin_n_dims = caps2u_params['caps dims'],
                                                         capsout_n_maps = caps2u_params['caps maps'],
                                                         capsout_n_dims = caps2u_params['caps dims'],
                                                         old_h = caps2u_params['h'],
                                                         old_w = caps2u_params['w'],
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=True)
        caps2ua_params = self.get_abstract_caps2ua.infer_shapes()

        self.get_abstract_caps1u = Get_Abstract_Caps_Up(self.opt.batch_size,
                                                         capsin_n_maps = caps2ua_params['caps maps'] + caps1_params['caps maps'],
                                                         capsin_n_dims = caps2ua_params['caps dims'],
                                                         capsout_n_maps = prim_params['maps out'],
                                                         capsout_n_dims = prim_params['caps dim'],
                                                         old_h = int(caps2ua_params['h']),
                                                         old_w = int(caps2ua_params['w']),
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 2,
                                                         padding = (2,2),
                                                         output_padding=upa_outputpadding,
                                                         uptype = self.opt.uptype,
                                                         across=True)
        caps1u_params = self.get_abstract_caps1u.infer_shapes()

        self.get_abstract_caps1ua = Get_Abstract_Caps_Down(self.opt.batch_size,
                                                         capsin_n_maps = caps1u_params['caps maps'],
                                                         capsin_n_dims = caps1u_params['caps dims'],
                                                         capsout_n_maps = caps1u_params['caps maps'],
                                                         capsout_n_dims = caps1u_params['caps dims'],
                                                         old_h = caps1u_params['h'],
                                                         old_w = caps1u_params['w'],
                                                         y_kernel = 5,
                                                         x_kernel = 5,
                                                         stride = 1,
                                                         padding = 2,
                                                         across=True)
        caps1ua_params = self.get_abstract_caps1ua.infer_shapes()

        self.get_abstract_caps_final1 = Get_Abstract_Caps_Up(self.opt.batch_size,
                                                         capsin_n_maps = caps1ua_params['caps maps'] + prim_params['maps out'],
                                                         capsin_n_dims = caps1ua_params['caps dims'],
                                                         capsout_n_maps = self.opt.f1maps,
                                                         capsout_n_dims = self.opt.f1dims,
                                                         old_h = int(caps1ua_params['h']),
                                                         old_w = int(caps1ua_params['w']),
                                                         y_kernel = 7,
                                                         x_kernel = 7,
                                                         stride = 2,
                                                         padding = (3,3),
                                                         output_padding=upa_outputpadding,
                                                         uptype = self.opt.uptype,
                                                         across=True)
        capsfinal1_params = self.get_abstract_caps_final1.infer_shapes()

        self.get_abstract_caps_final2 = Get_Abstract_Caps_Up(self.opt.batch_size,
                                                         capsin_n_maps = self.opt.f1maps,
                                                         capsin_n_dims = self.opt.f1dims,
                                                         capsout_n_maps = self.opt.f2maps,
                                                         capsout_n_dims = self.opt.f2dims,
                                                         old_h = int(capsfinal1_params['h']),
                                                         old_w = int(capsfinal1_params['w']),
                                                         y_kernel = 7,
                                                         x_kernel = 7,
                                                         stride = 2,
                                                         padding = (3,3),
                                                         output_padding=upa_outputpadding,
                                                         uptype = self.opt.uptype,
                                                         across=False)

        self.reconstruct = Reconstruction_Layer(self.opt.batch_size,
                                                capsin_n_maps = self.opt.f2maps,
                                                capsin_n_dims = self.opt.f2dims,
                                                reconstruct_channels=self.opt.reconchannels)

    def forward(self, x):
        x = self.get_prim_caps(x)
        x_prim = x
        #print(x.size(),'0')
        #self.x_prim = x_prim
        #print('##########################FINISHED PRIM#######################')
        x = self.get_abstract_caps1(x)

        #print(x.size(), '1')
        x = self.get_abstract_caps1a(x)
        x_1 = x
        #print(x.size(), '1a')
        #self.x_1 = x_1
        #print('##########################FINISHED 1#######################')

        x = self.get_abstract_caps2(x)
        #print(x.size(), '2')
        x = self.get_abstract_caps2a(x)
        x_2 = x
        #print(x.size(), '2a')
        #self.x_2 = x_2

        x = self.get_abstract_caps3(x)
        #print(x.size(), '3')
        x = self.get_abstract_caps3a(x)
        x_3 = x
        #print(x.size(), '3a')
        #self.x_3 = x_3

        x = self.get_abstract_caps_bot(x)
        #x_bot = x
        #print(x.size(), 'bot')

        #gotta be careful on the way up there are double maps
        x = torch.cat((x, x_3), 1)

        x = self.get_abstract_caps3u(x)
        #print(x.size(), '3u')
        x = self.get_abstract_caps3ua(x)
        #print(x.size(), '3ua')

        x = torch.cat((x, x_2), 1)
        x = self.get_abstract_caps2u(x)
        #print(x.size(), '2u')
        x = self.get_abstract_caps2ua(x)
        #print(x.size(), '2ua')

        x = torch.cat((x, x_1), 1)
        x = self.get_abstract_caps1u(x)
        #print(x.size(), '1u')
        x = self.get_abstract_caps1ua(x)
        #print(x.size(), '1ua')

        x = torch.cat((x, x_prim), 1)
        x = self.get_abstract_caps_final1(x)
        #print(x.size(), 'final1')

        x = self.get_abstract_caps_final2(x)
        #print(x.size(), 'final2')
        #reconstruct = self.reconstruct(x)

        x = safe_norm(x)

        #print(x.size(), 'last')
        return x#torch.argmax(x, 1, keepdim=True) #{'preds': x,
                #'reconstruct': reconstruct}

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

def safe_norm(s, axis=-1, epsilon=1e-7):
    squared_norm = torch.mul(s,s).sum(dim=axis)
    return torch.sqrt(squared_norm + epsilon)

###############################################################################
class Dice_Loss(torch.nn.Module):
    """This is a custom Dice Similarity Coefficient loss function that we use
    to the accuracy of the segmentation. it is defined as ;
    DSC = 2 * (pred /intersect label) / (pred /union label) for the losss we use
    1- DSC so gradient descent leads to better outputs."""
    def __init__(self, weight=None, size_average=False):
        super(Dice_Loss, self).__init__()

    def forward(self, pred, label):
        smooth = 1.              #helps with backprop
        intersection = torch.sum(pred * label)
        union = torch.sum(pred) + torch.sum(label)
        loss = (2. * intersection + smooth) / (union + smooth)
        #return 1-loss because we want to minimise dissimilarity
        return 1 - (loss)

'''
 def printskeet(i):
    f, (ax1i, ax1df, ax1l) = plt.subplots(1,3, sharey=True)
    pyplot.tight_layout()

    #raw image
    ax1i.imshow(allimages[i,:,:],
           aspect = 'equal')
    #pred
    ax1df.imshow(alldf[i,:,:],
           aspect = 'equal')
    #pred after thresholding
    ax1l.imshow(alllabels[i,:,:],
                   aspect = 'equal')
    #original label

    f.colorbar(ax1i.imshow(allimages[i,:,:],
                      aspect = 'equal'))
    pyplot.show()
'''

max_stars_repo_name: arjunsbalaji/oct

max_issues_repo_name: arjunsbalaji/oct

max_forks_repo_name: arjunsbalaji/oct


text:
"""
  This python script takes some real data and creates a number of tests data sets for
  examining rotation.

  The examples start with real data, but the original source of the data is unimportant
  and the test data could be recreated with seismic data from any 3-component broadband
  sensor.
  IU.ANMO.10.BH1 | 2017-08-26T00:00:00.019500Z - 2017-08-26T23:59:59.994500Z | 40.0 Hz, 3456000 samples
  IU.ANMO.10.BH2 | 2017-08-26T00:00:00.019500Z - 2017-08-26T23:59:59.994500Z | 40.0 Hz, 3456000 samples

"""
from obspy.core import *
import obspy.signal.rotate as rotate
import numpy as np
import numpy.testing as npt
import os

def rms(x):
  """ Calculate the RMS for validating test data sets.
  x is a numpy array
  Output
    Returns a scalar
  """
  return np.sqrt(np.mean(np.power(x,2)))

def create_rotation_examples(N,E,angles):
  for angle in angles:
    # We add 180 because the angle is back azimuth.  The rotation does not allow
    # negative angles so we continue on around the other way.
    if angle<180.:
      rangle=angle+180.
    else:
      rangle=angle-180.

    r,t=rotate.rotate_ne_rt(N[0].data,E[0].data,rangle)
    dir="rotation/%03d" % (angle)
    try:
      os.mkdir(dir)
    except:
      # Directory exists do nothing
      pass
    tr_n=N[0].copy()
    tr_n.data=r.astype('int32')
    tr_e=E[0].copy()
    tr_e.data=t.astype('int32')
    tr_n.write("%s/%s" % (dir,tr_n.id),format='MSEED')
    tr_e.write("%s/%s" % (dir,tr_e.id),format='MSEED')

def create_orthoganality_examples(N,E,angles):
  """Create examples where the second horizontal component is rotated out of orthoganality
  by the amount of degrees specified. The angle is the angle difference between the two
  components holding component 1 fixed."""
  for angle in angles:
    # We add 180 because the angle is back azimuth.  The rotation does not allow
    # negative angles so we continue on around the other way.
    rot_angle=180+angle-90.

    if rot_angle<0.:
      r,t=rotate.rotate_ne_rt(N[0].data,E[0].data,(360.+rot_angle))
    elif rot_angle>360.:
      r,t=rotate.rotate_ne_rt(N[0].data,E[0].data,(rot_angle-360.))
    else:
      r,t=rotate.rotate_ne_rt(N[0].data,E[0].data,rot_angle)
    dir="orthogonality/%03d" % (angle)
    try:
      os.mkdir(dir)
    except:
      # Directory exists do nothing
      pass
    tr_e=E[0].copy()
    tr_e.data=t.astype('int32')
    N[0].write("%s/%s" % (dir,N[0].id),format='MSEED')
    tr_e.write("%s/%s" % (dir,tr_e.id),format='MSEED')

def check_rotation_examples(N,E,angles):
  """ Check rotation angles by rotating the partially rotated the rest of the way to 360
  since obspy doesn't do the negative and compare to original values.  Should be close to
  zero"""
  ntests=0
  testsfailed=0
  for angle in angles:
    dir="%03d" % (angle)
    n=read("%s/%s" % (dir,N[0].stats.id))
    e=read("%s/%s" % (dir,N[0].stats.id))
    r,t=rotate.rotate_ne_rt(n[0].data,e[0].data,360.-angle)
    try:
      npt.assert_almost_equal(N[0].data,r)
      ntests+=1
      npt.assert_almost_equal(E[0].data,t)
      ntests+=1
    except:
      testsfailed+=1
  print("Check rotation %d tests passed test %d tests failed" % (ntests,testsfailed))
  return True

if __name__=='__main__':
  # Low Noise Examples for data rotation
  N=read('IU.ANMO.10.BH1')
  E=read('IU.ANMO.10.BH2')
  simple_rotation_angles=np.array([2.,5.,20.,90.,110.,180.,200.,270.,290.,355.,358.])
  create_rotation_examples(N,E,simple_rotation_angles)
  #check_rotation_examples(N,E,simple_rotation_angles)

  internal_angles=([60.,80.,85.,87.,88.,89.,89.5,90.5,91.,92.,93.,95.,100.,120.,270.])
  create_orthoganality_examples(N,E,internal_angles)

  # High Noise Examples
  N=read('IU.FUNA.00.BH1')
  E=read('IU.FUNA.00.BH2')
  create_rotation_examples(N,E,simple_rotation_angles)
  #check_rotation_examples(N,E,simple_rotation_angles)

  create_orthoganality_examples(N,E,internal_angles)

ext: py

max_stars_repo_name: usgs/seedscan-testdata

max_issues_repo_name: usgs/seedscan-testdata

max_forks_repo_name: usgs/seedscan-testdata

reason: import numpy


text:
from sklearn.neighbors import KernelDensity
from scipy.spatial.distance import euclidean
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import MDS
from sklearn import preprocessing
from sklearn.manifold import MDS
def Select_TI_For_LocalArea_plot(Num_Selected_TI,
                                TargetBloc,
                                LineBloc_space,
                                TI_Assignment_Matrix,
                                Block_distances,
                                MDS_coordinates_2D,
                                MDS_coordinates_2D_List,
                                kde_bw = 0.3,
                                plot = True):

    '''
    This function select several TI for a specified local area considering both hard data and global correlations
     with other blocks
     @Parameters
    Num_Selected_TI: Number of TIs assgined for each block, int. Fixed for all local block areas

    TargetBloc: Index location of the target block, int.
                It indicates the index number of target block, same as the designed simulation path values.

    LineBloc_space:Simluation Grid with all block indexes. 2D array. The index values is same as simulation path value.

    TI_Assignment_Matrix: selected TI from Step2 based on hard data only.

    Block_distances: Distance matrix between the Indexed blocks.It is the blockwised distances calculated from Step 3.
                        2D array. [Total_Block_number, Total_Block_number]

    MDS_coordinates_2D: TI distance from MDS
    MDS_coordinates_2D_List: "MDS_coordinates_2D" in list format
    '''

    # the key parameter in kernel density estimator
    KDE_Bandwidth = kde_bw

    # initialize the transition probability vector
    # MDS_coordinates_2D.shape[0] is the number of candidate TI
    TransitionRatio = np.zeros((MDS_coordinates_2D.shape[0]))

    Aggregation_weight_sum = 0.0

    # derive aggregation weights from Block distance
    if np.any(Block_distances != 0):
        Aggregation_weights = 1-Block_distances/Block_distances.max()

    # if the target loation is empty,  Block_distances == 0.
    # Then get the Aggregation_weights by spatial distance between the blocks.
    elif np.all(Block_distances== 0):
        argwhere_TargetBloc = np.argwhere(LineBloc_space==TargetBloc)
        x_centr = (argwhere_TargetBloc[:,1].max() + argwhere_TargetBloc[:,1].min() )//2
        y_centr = (argwhere_TargetBloc[:,0].max() + argwhere_TargetBloc[:,0].min() )//2
        Spatial_dist = []
        for bloc_index in range(len(Block_distances)):
            argwhere_Bloc = np.argwhere(LineBloc_space==bloc_index)
            x_centr_Bloc = (argwhere_Bloc[:,1].max() + argwhere_Bloc[:,1].min() )//2
            y_centr_Bloc = (argwhere_Bloc[:,0].max() + argwhere_Bloc[:,0].min() )//2
            dist = np.sqrt( (x_centr-x_centr_Bloc)**2 + (y_centr-y_centr_Bloc)**2)
            Spatial_dist.append(dist)
        Aggregation_weights = 1-Spatial_dist/max(Spatial_dist)

    # visit each local area in the simulation domain
    for bloc_index in range(len(Block_distances)):
        ### Modified, if the location is empty
        if np.all(Block_distances== 0):
             # obtain the specified training image
            # test 0
            TI_Index = TI_Assignment_Matrix[bloc_index]

            Aggregation_weight = Aggregation_weights[bloc_index]

            Aggregation_weight_sum += Aggregation_weight

            KDE_instance = \
                KernelDensity(kernel='gaussian',bandwidth=KDE_Bandwidth).fit(MDS_coordinates_2D[TI_Index])

            TransitionProbability_oneInstance = np.exp(KDE_instance.score_samples(MDS_coordinates_2D_List))

            TransitionProbability_oneInstance = TransitionProbability_oneInstance / np.sum(TransitionProbability_oneInstance)

            TransitionProbability_oneInstance += 0.00000001

            TransitionRatio_oneInstance = np.log( TransitionProbability_oneInstance / (1-TransitionProbability_oneInstance) )

            # accumulate the transition probability
            TransitionRatio = TransitionRatio + Aggregation_weight * TransitionRatio_oneInstance

        ###! If the location is the block itself and not empty

        elif(np.any(Block_distances != 0) and bloc_index == TargetBloc):

            # obtain the specified training image
            # test ==
            TI_Index =  TI_Assignment_Matrix[bloc_index]

            # Weight = (1+Block_Density/smallest_distance)
#             Aggregation_weight = (1+Block_Density)/min(Block_distances[Block_distances>0])
            Aggregation_weight = 1

            Aggregation_weight_sum += Aggregation_weight

            # calculate transition probability
            KDE_instance = KernelDensity(kernel='gaussian',bandwidth=KDE_Bandwidth).fit(MDS_coordinates_2D[TI_Index])

            TransitionProbability_oneInstance = np.exp(KDE_instance.score_samples(MDS_coordinates_2D_List))

            TransitionProbability_oneInstance = TransitionProbability_oneInstance / np.sum(TransitionProbability_oneInstance)

            TransitionProbability_oneInstance[TransitionProbability_oneInstance<=0] =  0.00000001
            TransitionProbability_oneInstance[TransitionProbability_oneInstance>=1] =  1-0.00000001

            TransitionRatio_oneInstance = np.log( TransitionProbability_oneInstance / (1-TransitionProbability_oneInstance) )

            # accumulate the transition probability
            if plot:
                TransitionRatio = TransitionRatio + Aggregation_weight * TransitionRatio_oneInstance
                plt.scatter(MDS_coordinates_2D[:, 0], MDS_coordinates_2D[:, 1], c=TransitionProbability_oneInstance, cmap='jet', s=36)
                plt.xlabel(str('MDS1'),fontsize='large')
                plt.ylabel(str('MDS2'),fontsize='large')
                plt.tick_params(direction='in',labelsize='large')
                plt.colorbar()
                plt.title(f'The TI selection probability for B{bloc_index}')
                plt.show()

        # Including all other areas with flight lines
        elif np.any(Block_distances != 0):
            # obtain the specified training image
            # test all
            TI_Index = TI_Assignment_Matrix[bloc_index]

            # Modified calculate the importance / weight / influence Use distacne from Step 3

            Aggregation_weight = Aggregation_weights[bloc_index]

            Aggregation_weight = float(Aggregation_weight)
            Aggregation_weight_sum += Aggregation_weight

            # calculate transition probability
            KDE_instance = KernelDensity(kernel='gaussian',bandwidth=KDE_Bandwidth).fit(MDS_coordinates_2D[TI_Index])

            TransitionProbability_oneInstance = np.exp(KDE_instance.score_samples(MDS_coordinates_2D_List))

            TransitionProbability_oneInstance = TransitionProbability_oneInstance / np.sum(TransitionProbability_oneInstance)

#             TransitionProbability_oneInstance += 0.00000001
            TransitionProbability_oneInstance[TransitionProbability_oneInstance<=0] =  0.00000001
            TransitionProbability_oneInstance[TransitionProbability_oneInstance>=1] =  1-0.00000001

            TransitionRatio_oneInstance = np.log( TransitionProbability_oneInstance / (1-TransitionProbability_oneInstance) )

            # accumulate the transition probability
            TransitionRatio = TransitionRatio + Aggregation_weight * TransitionRatio_oneInstance
            if plot:
                plt.scatter(MDS_coordinates_2D[:, 0], MDS_coordinates_2D[:, 1], c=TransitionProbability_oneInstance, cmap='jet', s=36)
                plt.xlabel(str('MDS1'),fontsize='large')
                plt.ylabel(str('MDS2'),fontsize='large')
                plt.tick_params(direction='in',labelsize='large')
                plt.colorbar()
                plt.title(f'The TI selection probability for A{bloc_index}')
                plt.show()

#             plt.scatter(MDS_coordinates_2D[:, 0], MDS_coordinates_2D[:, 1], c=TransitionRatio_oneInstance, cmap='jet', s=36)
#             plt.xlabel(str('MDS1'),fontsize='large')
#             plt.ylabel(str('MDS2'),fontsize='large')
#             plt.tick_params(direction='in',labelsize='large')
#             plt.colorbar()
#             plt.title(f'The transition ratio for area y{index_y}, x{index_x}')
#             plt.show()

    # normalize the transition probability
#     TransitionRatio = TransitionRatio / Aggregation_weight_sum!!!
    Transition_TemporaryVariable = np.exp(TransitionRatio)

    TransitionProbability = Transition_TemporaryVariable / (1+Transition_TemporaryVariable)

#     print(np.sum(TransitionProbability))
    TransitionProbability[np.isnan(TransitionProbability)] = 0

    TransitionProbability = TransitionProbability / np.sum(TransitionProbability)

    pdf = np.copy(TransitionProbability)
    if plot:
        colormap_max = np.max(TransitionProbability)
        plt.scatter(MDS_coordinates_2D[:, 0], MDS_coordinates_2D[:, 1],
                    c=TransitionProbability,
                    vmin=0.0,vmax=colormap_max,
                    cmap='jet', s=36)
        plt.xlabel(str('MDS1'),fontsize='large')
        plt.ylabel(str('MDS2'),fontsize='large')
        plt.tick_params(direction='in',labelsize='large')
        plt.colorbar()
        plt.title('The accumulative transition probability')
        plt.show()

    Candidate_TI_Index = [167]
    while Candidate_TI_Index[0]>150:
        Candidate_TI_Index = np.random.choice(a=MDS_coordinates_2D.shape[0],
                                              size=Num_Selected_TI,
                                              replace=False,
                                              p=TransitionProbability.tolist())

    print(f'the selected TI {Candidate_TI_Index}')

    if plot:
        plt.scatter(MDS_coordinates_2D[:, 0], MDS_coordinates_2D[:, 1], s=36)
        plt.scatter(MDS_coordinates_2D[Candidate_TI_Index, 0], MDS_coordinates_2D[Candidate_TI_Index, 1],
                    color='red', s=49, label='Selected TI')
        plt.xlabel(str('MDS1'),fontsize='large')
        plt.ylabel(str('MDS2'),fontsize='large')
        plt.tick_params(direction='in',labelsize='large')
        plt.legend(scatterpoints=1, loc='upper right', shadow=False,fontsize='large')
        plt.title('The selected TI')
    #     plt.colorbar()
        plt.show()

    ### Plot TI pdf.

        plt.figure(figsize=(10,5))

        argsort = pdf.argsort()[::-1]
    #     argsort = np.delete(argsort, np.argwhere(argsort[:10]>150)[:,0])

        plt.plot(pdf[argsort],'k.', markersize=15)
        plt.ylim(-0.01, pdf.max()*1.2 )

        plt.title('PDF vs TI - A'+str(TargetBloc))
        plt.ylabel('PDF', fontsize = 14)
        plt.xlabel('Training images rank (after transition)', fontsize = 15)
        plt.xticks(fontsize=12), plt.yticks(fontsize=12)
        plt.show()

    return Candidate_TI_Index, pdf, Aggregation_weights, TransitionRatio

reason: import numpy,from scipy


text:
from typing import Any, Dict, List, Optional, Union

from copy import copy
from random import randint
import datetime as dt
import re

from lunchbox.enforce import Enforce
from pandas import DataFrame, DatetimeIndex
from schematics.exceptions import DataError
import cufflinks as cf  # noqa: F401
import lunchbox.tools as lbt
import numpy as np
import pandasql
import pyparsing as pp
import rolling_pin.blob_etl as rpb
import webcolors

from shekels.core.config import ConformAction
import shekels.core.config as cfg
import shekels.enforce.enforce_tools as eft
# ------------------------------------------------------------------------------

COLOR_COERCION_LUT = {
    '#00CC96': '#5F95DE',
    '#0D0887': '#444459',
    '#19D3F3': '#5F95DE',
    '#242424': '#242424',
    '#276419': '#343434',
    '#2A3F5F': '#444459',
    '#343434': '#343434',
    '#444444': '#444444',
    '#46039F': '#444459',
    '#4D9221': '#444444',
    '#636EFA': '#5F95DE',
    '#7201A8': '#5D5D7A',
    '#7FBC41': '#8BD155',
    '#8E0152': '#444459',
    '#9C179E': '#5D5D7A',
    '#A4A4A4': '#A4A4A4',
    '#AB63FA': '#AC92DE',
    '#B6E880': '#A0D17B',
    '#B6ECF3': '#B6ECF3',
    '#B8E186': '#A0D17B',
    '#BD3786': '#F77E70',
    '#C51B7D': '#F77E70',
    '#C8D4E3': '#B6ECF3',
    '#D8576B': '#F77E70',
    '#DE77AE': '#DE958E',
    '#DE958E': '#DE958E',
    '#E5ECF6': '#F4F4F4',
    '#E6F5D0': '#E9EABE',
    '#EBF0F8': '#F4F4F4',
    '#ED7953': '#F77E70',
    '#EF553B': '#F77E70',
    '#F0F921': '#E8EA7E',
    '#F1B6DA': '#C98FDE',
    '#F4F4F4': '#F4F4F4',
    '#F7F7F7': '#F4F4F4',
    '#FB9F3A': '#EB9E58',
    '#FDCA26': '#EB9E58',
    '#FDE0EF': '#F4F4F4',
    '#FECB52': '#EB9E58',
    '#FF6692': '#F77E70',
    '#FF97FF': '#C98FDE',
    '#FFA15A': '#EB9E58',
}

def conform(data, actions=[], columns=[]):
    # type: (DataFrame, List[dict], List[str]) -> DataFrame
    '''
    Conform given mint transaction data.

    Args:
        data (DataFrame): Mint transactions DataFrame.
        actions (list[dict], optional): List of conform actions. Default: [].
        columns (list[str], optional): List of columns. Default: [].

    Raises:
        DataError: If invalid conform action given.
        ValueError: If source column not found in data columns.

    Returns:
        DataFrame: Conformed DataFrame.
    '''
    for action in actions:
        ConformAction(action).validate()

    data.rename(lbt.to_snakecase, axis=1, inplace=True)
    lut = dict(
        account_name='account',
        transaction_type='type'
    )
    data.rename(lambda x: lut.get(x, x), axis=1, inplace=True)
    data.date = DatetimeIndex(data.date)
    data.amount = data.amount.astype(float)
    data.category = data.category \
        .apply(lambda x: re.sub('&', 'and', lbt.to_snakecase(x)))
    data.account = data.account.apply(lbt.to_snakecase)

    for action in actions:
        source = action['source_column']
        if source not in data.columns:
            msg = f'Source column {source} not found in columns. '
            msg += f'Legal columns include: {data.columns.tolist()}.'
            raise ValueError(msg)

        target = action['target_column']
        if target not in data.columns:
            data[target] = None

        for regex, val in action['mapping'].items():
            if action['action'] == 'overwrite':
                mask = data[source] \
                    .apply(lambda x: re.search(regex, x, flags=re.I)).astype(bool)
                data.loc[mask, target] = val
            elif action['action'] == 'substitute':
                data[target] = data[source] \
                    .apply(lambda x: re.sub(regex, val, str(x), flags=re.I))

    if columns != []:
        data = data[columns]
    return data

def filter_data(data, column, comparator, value):
    # type: (DataFrame, str, str, Any) -> DataFrame
    '''
    Filters given data via comparator(column value, value).

    Legal comparators:

        * == ``lambda a, b: a == b``
        * != ``lambda a, b: a != b``
        *  > ``lambda a, b: a > b``
        * >= ``lambda a, b: a >= b``
        *  < ``lambda a, b: a < b``
        * =< ``lambda a, b: a <= b``
        *  ~ ``lambda a, b: bool(re.search(a, b, flags=re.I))``
        * !~ ``lambda a, b: not bool(re.search(a, b, flags=re.I))``

    Args:
        data (DataFrame): DataFrame to be filtered.
        column (str): Column name.
        comparator (str): String representation of comparator.
        value (object): Value to be compared.

    Raises:
        EnforceError: If data is not a DataFrame.
        EnforceError: If column is not a string.
        EnforceError: If column not in data columns.
        EnforceError: If illegal comparator given.
        EnforceError: If comparator is ~ or !~ and value is not a string.

    Returns:
        DataFrame: Filtered data.
    '''
    Enforce(data, 'instance of', DataFrame)
    msg = 'Column must be a str. {a} is not str.'
    Enforce(column, 'instance of', str, message=msg)
    eft.enforce_columns_in_dataframe([column], data)

    lut = {
        '==': lambda a, b: a == b,
        '!=': lambda a, b: a != b,
        '>': lambda a, b: a > b,
        '>=': lambda a, b: a >= b,
        '<': lambda a, b: a < b,
        '<=': lambda a, b: a <= b,
        '~': lambda a, b: bool(re.search(b, a, flags=re.I)),
        '!~': lambda a, b: not bool(re.search(b, a, flags=re.I)),
    }
    msg = 'Illegal comparator. {a} not in [==, !=, >, >=, <, <=, ~, !~].'
    Enforce(comparator, 'in', lut.keys(), message=msg)

    if comparator in ['~', '!~']:
        msg = 'Value must be string if comparator is ~ or !~. {a} is not str.'
        Enforce(value, 'instance of', str, message=msg)
    # --------------------------------------------------------------------------

    op = lut[comparator]
    mask = data[column].apply(lambda x: op(x, value))
    data = data[mask]
    return data

def group_data(data, columns, metric, datetime_column='date'):
    # type: (DataFrame, Union[str, List[str]], str, str) -> DataFrame
    '''
    Groups given data by given columns according to given metric.
    If a legal time interval is given in the columns, then an additional special
    column of that same name is added to the data for grouping.

    Legal metrics:

        *   max ``lambda x: x.max()``
        *  mean ``lambda x: x.mean()``
        *   min ``lambda x: x.min()``
        *   std ``lambda x: x.std()``
        *   sum ``lambda x: x.sum()``
        *   var ``lambda x: x.var()``
        * count ``lambda x: x.count()``

    Legal time intervals:

        * year
        * quarter
        * month
        * two_week
        * week
        * day
        * hour
        * half_hour
        * quarter_hour
        * minute
        * second
        * microsecond

    Args:
        data (DataFrame): DataFrame to be grouped.
        columns (str or list[str]): Columns to group data by.
        metric (str): String representation of metric.
        datetime_column (str, optinal): Datetime column for time grouping.
            Default: date.

    Raises:
        EnforceError: If data is not a DataFrame.
        EnforceError: If columns not in data columns.
        EnforceError: If illegal metric given.
        EnforceError: If time interval in columns and datetime_column not in
            columns.

    Returns:
        DataFrame: Grouped data.
    '''
    # luts
    met_lut = {
        'max': lambda x: x.max(),
        'mean': lambda x: x.mean(),
        'min': lambda x: x.min(),
        'std': lambda x: x.std(),
        'sum': lambda x: x.sum(),
        'var': lambda x: x.var(),
        'count': lambda x: x.count(),
    }

    time_lut = {
        'year': lambda x: dt.datetime(x.year, 1, 1),
        'quarter': lambda x: dt.datetime(
            x.year, int(np.ceil(x.month / 3) * 3 - 2), 1
        ),
        'month': lambda x: dt.datetime(x.year, x.month, 1),
        'two_week': lambda x: dt.datetime(
            x.year, x.month, min(int(np.ceil(x.day / 14) * 14 - 13), 28)
        ),
        'week': lambda x: dt.datetime(
            x.year, x.month, max(1, min([int(x.month / 7) * 7, 28]))
        ),
        'day': lambda x: dt.datetime(x.year, x.month, x.day),
        'hour': lambda x: dt.datetime(x.year, x.month, x.day, x.hour),
        'half_hour': lambda x: dt.datetime(
            x.year, x.month, x.day, x.hour, int(x.minute / 30) * 30
        ),
        'quarter_hour': lambda x: dt.datetime(
            x.year, x.month, x.day, x.hour, int(x.minute / 15) * 15
        ),
        'minute': lambda x: dt.datetime(
            x.year, x.month, x.day, x.hour, x.minute
        ),
        'second': lambda x: dt.datetime(
            x.year, x.month, x.day, x.hour, x.minute, x.second
        ),
        'microsecond': lambda x: dt.datetime(
            x.year, x.month, x.day, x.hour, x.minute, x.second, x.microsecond
        ),
    }
    # --------------------------------------------------------------------------

    # enforcements
    Enforce(data, 'instance of', DataFrame)
    columns_ = columns  # type: Any
    if type(columns_) != list:
        columns_ = [columns_]

    cols = list(filter(lambda x: x not in time_lut.keys(), columns_))
    eft.enforce_columns_in_dataframe(cols, data)

    msg = '{a} is not a legal metric. Legal metrics: {b}.'
    Enforce(metric, 'in', sorted(list(met_lut.keys())), message=msg)

    # time column
    if len(columns_) > len(cols):
        eft.enforce_columns_in_dataframe([datetime_column], data)
        msg = 'Datetime column of type {a}, it must be of type {b}.'
        Enforce(
            data[datetime_column].dtype.type, '==', np.datetime64, message=msg
        )
    # --------------------------------------------------------------------------

    for col in columns_:
        if col in time_lut.keys():
            op = time_lut[col]
            data[col] = data[datetime_column].apply(op)
    agg = met_lut[metric]
    cols = data.columns.tolist()
    grp = data.groupby(columns_, as_index=False)
    output = agg(grp)

    # get first value for columns that cannot be computed by given metric
    diff = set(cols).difference(output.columns.tolist())
    if len(diff) > 0:
        first = grp.first()
        for col in diff:
            output[col] = first[col]
    return output

def pivot_data(data, columns, values=[], index=None):
    # type: (DataFrame, List[str], List[str], Optional[str]) -> DataFrame
    '''
    Pivots a given dataframe via a list of columns.

    Legal time columns:

        * date
        * year
        * quarter
        * month
        * two_week
        * week
        * day
        * hour
        * half_hour
        * quarter_hour
        * minute
        * second
        * microsecond

    Args:
        data (DataFrame): DataFrame to be pivoted.
        columns (list[str]): Columns whose unique values become separate traces
            within a plot.
        values (list[str], optional): Columns whose values become the values
            within each trace of a plot. Default: [].
        index (str, optional): Column whose values become the y axis values of a
            plot. Default: None.

    Raises:
        EnforceError: If data is not a DataFrame.
        EnforceError: If data is of zero length.
        EnforceError: If columns not in data columns.
        EnforceError: If values not in data columns.
        EnforceError: If index not in data columns or legal time columns.

    Returns:
        DataFrame: Pivoted data.
    '''
    time_cols = [
        'date', 'year', 'quarter', 'month', 'two_week', 'week', 'day', 'hour',
        'half_hour', 'quarter_hour', 'minute', 'second', 'microsecond',
    ]

    Enforce(data, 'instance of', DataFrame)
    msg = 'DataFrame must be at least 1 in length. Given length: {a}.'
    Enforce(len(data), '>=', 1, message=msg)
    eft.enforce_columns_in_dataframe(columns, data)
    eft.enforce_columns_in_dataframe(values, data)
    if index is not None:
        msg = '{a} is not in legal column names: {b}.'
        Enforce(index, 'in', data.columns.tolist() + time_cols, message=msg)
    # --------------------------------------------------------------------------

    vals = copy(values)
    if index is not None and index not in values:
        vals.append(index)

    if index in time_cols:
        data[index] = data[index] \
            .apply(lambda x: x + dt.timedelta(microseconds=randint(0, 999999)))

    data = data.pivot(columns=columns, values=vals, index=index)
    data = data[values]
    data.columns = data.columns.droplevel(0)
    return data

def get_figure(
    data,              # type: DataFrame
    filters=[],        # type: List[dict]
    group=None,        # type: Optional[dict]
    pivot=None,        # type: Optional[dict]
    kind='bar',        # type: str
    color_scheme={},   # type: Dict[str, str]
    x_axis=None,       # type: Optional[str]
    y_axis=None,       # type: Optional[str]
    title=None,        # type: Optional[str]
    x_title=None,      # type: Optional[str]
    y_title=None,      # type: Optional[str]
    bins=50,           # type: int
    bar_mode='stack',  # type: str
):
    '''
    Generates a plotly figure dictionary from given data and manipulations.

    Args:
        data (DataFrame): Data.
        filters (list[dict], optional): List of filters for data. Default: [].
        group (dict, optional): Grouping operation. Default: None.
        pivot (dict, optional): Pivot operation. Default: None.
        kind (str, optional): Kind of plot. Default: bar.
        color_scheme (dict[str, str], optional): Color scheme. Default: {}.
        x_axis (str): Column to use as x axis: Default: None.
        y_axis (str): Column to use as y axis: Default: None.
        title (str, optional): Title of plot. Default: None.
        x_title (str, optional): Title of x axis. Default: None.
        y_title (str, optional): Title of y axis. Default: None.
        bins (int, optional): Number of bins if histogram. Default: 50.
        bar_mode (str, optional): How bars in bar graph are presented.
            Default: stack.

    Raises:
        DataError: If any filter in filters is invalid.
        DataError: If group is invalid.
        DataError: If pivot is invalid.

    Returns:
        dict: Plotly Figure as dictionary.
    '''
    data = data.copy()

    # filter
    for f in filters:
        f = cfg.FilterAction(f)
        try:
            f.validate()
        except DataError as e:
            raise DataError({'Invalid filter': e.to_primitive()})

        f = f.to_primitive()
        if len(data) == 0:
            break
        data = filter_data(data, f['column'], f['comparator'], f['value'])

    # group
    if group is not None:
        grp = group  # type: Any
        grp = cfg.GroupAction(grp)
        try:
            grp.validate()
        except DataError as e:
            raise DataError({'Invalid group': e.to_primitive()})
        grp = grp.to_primitive()

        data = group_data(
            data,
            grp['columns'],
            grp['metric'],
            datetime_column=grp['datetime_column'],
        )

    # pivot
    if pivot is not None:
        pvt = pivot  # type: Any
        pvt = cfg.PivotAction(pvt)
        try:
            pvt.validate()
        except DataError as e:
            raise DataError({'Invalid pivot': e.to_primitive()})
        pvt = pvt.to_primitive()

        data = pivot_data(
            data, pvt['columns'], values=pvt['values'], index=pvt['index']
        )

    # create figure
    figure = data.iplot(
        kind=kind, asFigure=True, theme='henanigans', colorscale='henanigans',
        x=x_axis, y=y_axis, title=title, xTitle=x_title, yTitle=y_title,
        barmode=bar_mode, bins=bins
    ).to_dict()  # type: dict
    figure['layout']['title']['font']['color'] = '#F4F4F4'
    figure['layout']['xaxis']['title']['font']['color'] = '#F4F4F4'
    figure['layout']['yaxis']['title']['font']['color'] = '#F4F4F4'
    if color_scheme != {}:
        figure = conform_figure(figure, color_scheme)
    return figure

def parse_rgba(string):
    '''
    Parses rgb and rgba strings into tuples of numbers.

    Example:
        >>>parse_rgba('rgb(255, 0, 0)')
        (255, 0, 0)
        >>>parse_rgba('rgba(255, 0, 0, 0.5)')
        (255, 0, 0, 0.5)
        >>>parse_rgba('foo') is None
        True

    Args:
        string (str): String to be parsed.

    Returns:
        tuple: (red, green, blue) or (red, green, blue, alpha)
    '''
    result = re.search(r'rgba?\((\d+, \d+, \d+(, \d+\.?\d*)?)\)', string)
    if result is None:
        return None

    result = result.group(1)
    result = re.split(', ', result)
    if len(result) == 3:
        result = tuple(map(int, result))
        return result

    result = list(map(int, result[:-1])) + [float(result[-1])]
    result = tuple(result)
    return result

def conform_figure(figure, color_scheme):
    '''
    Conforms given figure to use given color scheme.

    Args:
        figure (dict): Plotly figure.
        color_scheme (dict): Color scheme dictionary.

    Returns:
        dict: Conformed figure.
    '''
    # create hex to hex lut
    lut = {}
    for key, val in cfg.COLOR_SCHEME.items():
        if key in color_scheme:
            lut[val] = color_scheme[key]

    # rgba? to hex --> coerce to standard colors --> coerce with color_scheme
    figure = rpb.BlobETL(figure) \
        .set(
            predicate=lambda k, v: isinstance(v, str) and 'rgb' in v,
            value_setter=lambda k, v: webcolors.rgb_to_hex(parse_rgba(v)[:3]).upper()) \
        .set(
            predicate=lambda k, v: isinstance(v, str),
            value_setter=lambda k, v: COLOR_COERCION_LUT.get(v, v)) \
        .set(
            predicate=lambda k, v: isinstance(v, str),
            value_setter=lambda k, v: lut.get(v, v)) \
        .to_dict()
    return figure

# SQL-PARSING-------------------------------------------------------------------
def get_sql_grammar():
    '''
    Creates a grammar for parsing SQL queries.

    Returns:
       MatchFirst: SQL parser.
    '''
    select = pp.Regex('select', flags=re.I) \
        .setParseAction(lambda s, l, t: 'select') \
        .setResultsName('operator')
    from_ = pp.Suppress(pp.Regex('from', flags=re.I))
    table = (from_ + pp.Regex('[a-z]+', flags=re.I)) \
        .setParseAction(lambda s, l, t: t[0]) \
        .setResultsName('table')
    regex = pp.Regex('~|regex').setParseAction(lambda s, l, t: '~')
    not_regex = pp.Regex('!~|not regex').setParseAction(lambda s, l, t: '!~')
    any_op = pp.Regex('[^ ]*')
    operator = pp.Or([not_regex, regex, any_op]).setResultsName('operator')
    quote = pp.Suppress(pp.Optional("'"))
    value = (quote + pp.Regex('[^\']+', flags=re.I) + quote) \
        .setResultsName('value') \
        .setParseAction(lambda s, l, t: t[0])
    columns = pp.delimitedList(pp.Regex('[^, ]*'), delim=pp.Regex(', *')) \
        .setResultsName('display_columns')
    column = pp.Regex('[a-z]+', flags=re.I).setResultsName('column')
    conditional = column + operator + value
    head = select + columns + table
    grammar = head | conditional
    return grammar

def query_data(data, query):
    '''
    Parses SQL + regex query and applies it to given data.

    Regex operators:

        * ~, regex - Match regular expression
        * !~, not regex - Do not match regular expression

    Args:
        data (DataFrame): DataFrame to be queried.
        query (str): SQL query that may include regex operators.

    Returns:
        DataFrame: Data filtered by query.
    '''
    # split queries by where/and/or
    queries = re.split(' where | and | or ', query, flags=re.I)

    # detect whether any sub query has a regex operator
    has_regex = False
    for q in queries:
        if re.search(' regex | ~ | !~ ', q, flags=re.I):
            has_regex = True
            break

    # if no regex operator is found just submit query to pandasql
    if not has_regex:
        data = pandasql.sqldf(query, {'data': data})

    else:
        grammar = get_sql_grammar()

        # move select statement to end
        if 'select' in queries[0]:
            q = queries.pop(0)
            queries.append(q)

        for q in queries:
            # get column, operator and value
            parse = grammar.parseString(q).asDict()
            op = parse['operator']

            # initial select statement
            if op == 'select':
                data = pandasql.sqldf(q, {'data': data})

            # regex search
            elif op == '~':
                mask = data[parse['column']] \
                    .astype(str) \
                    .apply(lambda x: re.search(parse['value'], x, flags=re.I)) \
                    .astype(bool)
                data = data[mask]

            # regex not search
            elif op == '!~':
                mask = data[parse['column']] \
                    .astype(str) \
                    .apply(lambda x: re.search(parse['value'], x, flags=re.I)) \
                    .astype(bool)
                data = data[~mask]

            # ther SQL query
            else:
                data = pandasql.sqldf('select * from data where ' + q, {'data': data})

            if len(data) == 0:
                break
    return data

def query_dict(data, query):
    # type: (dict, str) -> dict
    '''
    Query a given diction with a given SQL query.

    Args:
        data (dict): Dictionary to be queried.
        query (str): SQL query.

    Returns:
        dict: Queried dictionary.
    '''
    data_ = data  # type: Any
    data_ = rpb.BlobETL(data_) \
        .to_flat_dict() \
        .items()
    data_ = DataFrame(list(data_), columns=['key', 'value'])
    data_ = query_data(data_, query)
    data_ = dict(zip(data_.key.tolist(), data_.value.tolist()))
    data_ = rpb.BlobETL(data_).to_dict()
    return data_

max_stars_repo_name: theNewFlesh/shekels

max_issues_repo_name: theNewFlesh/shekels

max_forks_repo_name: theNewFlesh/shekels