text:
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import sys

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        # self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 2)

    def forward(self, x):
        x1 = F.relu(F.max_pool2d(self.conv1(x), 2))
        x2 = F.relu(F.max_pool2d(self.conv2(x1), 2))
        x22 = x2.view(-1, 320)
        x3 = F.relu(self.fc1(x22))
        # x = F.dropout(x, training=self.training)
        x4 = self.fc2(x3)
        return F.log_softmax(x4, dim=1), (x1, x2, x3, x4)

class variableWidthNet(nn.Module):
    def __init__(self, hiddenDim):
        super(variableWidthNet, self).__init__()
        self.fc1 = nn.Linear(1, hiddenDim)  # 45  65 ORIGINALLY
        self.fc2 = nn.Linear(hiddenDim, 1)

    def forward(self, x):
        x1 = F.tanh(self.fc1(x))
        x2 = self.fc2(x1)
        return x2

class variableDepthNet(nn.Module):
    def __init__(self, hiddenDims, training):
        super(variableDepthNet, self).__init__()
        self.training = training
        if len(hiddenDims) == 2:
            self.hiddenLayers = 2
            self.fc1 = nn.Linear(1, hiddenDims[0])
            self.fc2 = nn.Linear(hiddenDims[0], hiddenDims[1])
            self.fc3 = nn.Linear(hiddenDims[1], 1)
        if len(hiddenDims) == 3:
            self.hiddenLayers = 3
            self.fc1 = nn.Linear(1, hiddenDims[0])
            self.fc2 = nn.Linear(hiddenDims[0], hiddenDims[1])
            self.fc3 = nn.Linear(hiddenDims[1], hiddenDims[2])
            self.fc4 = nn.Linear(hiddenDims[2], 1)
        if len(hiddenDims) == 4:
            self.hiddenLayers = 4
            self.fc1 = nn.Linear(1, hiddenDims[0])
            self.fc2 = nn.Linear(hiddenDims[0], hiddenDims[1])
            self.fc3 = nn.Linear(hiddenDims[1], hiddenDims[2])
            self.fc4 = nn.Linear(hiddenDims[2], hiddenDims[3])
            self.fc5= nn.Linear(hiddenDims[3], 1)
        if len(hiddenDims) == 5:
            self.hiddenLayers = 5
            self.fc1 = nn.Linear(1, hiddenDims[0])
            self.fc2 = nn.Linear(hiddenDims[0], hiddenDims[1])
            self.fc3 = nn.Linear(hiddenDims[1], hiddenDims[2])
            self.fc4 = nn.Linear(hiddenDims[2], hiddenDims[3])
            self.fc5= nn.Linear(hiddenDims[3], hiddenDims[4])
            self.fc6 = nn.Linear(hiddenDims[4], 1)
        if len(hiddenDims) == 6:
            self.hiddenLayers = 6
            self.fc1 = nn.Linear(1, hiddenDims[0])
            self.fc2 = nn.Linear(hiddenDims[0], hiddenDims[1])
            self.fc3 = nn.Linear(hiddenDims[1], hiddenDims[2])
            self.fc4 = nn.Linear(hiddenDims[2], hiddenDims[3])
            self.fc5= nn.Linear(hiddenDims[3], hiddenDims[4])
            self.fc6 = nn.Linear(hiddenDims[4], hiddenDims[5])
            self.fc7 = nn.Linear(hiddenDims[5], 1)

    def forward(self, x):
        if self.hiddenLayers == 2:
            x1 = F.tanh(self.fc1(x))
            x1 = F.dropout(x1,p=0.25, training=self.training)
            x2 = F.tanh(self.fc2(x1))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x3 = self.fc3(x2)
            return x3
        if self.hiddenLayers == 3:
            x1 = F.tanh(self.fc1(x))
            x1 = F.dropout(x1,p=0.25, training=self.training)
            x2 = F.tanh(self.fc2(x1))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x3 = F.tanh(self.fc3(x2))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x4 = self.fc4(x3)
            return x4
        if self.hiddenLayers == 4:
            x1 = F.tanh(self.fc1(x))
            x1 = F.dropout(x1,p=0.25, training=self.training)
            x2 = F.tanh(self.fc2(x1))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x3 = F.tanh(self.fc3(x2))
            x3 = F.dropout(x3,p=0.25, training=self.training)
            x4 = F.tanh(self.fc4(x3))
            x4 = F.dropout(x4,p=0.25, training=self.training)
            x5 = self.fc5(x4)
            return x5
        if self.hiddenLayers == 5:
            x1 = F.tanh(self.fc1(x))
            x1 = F.dropout(x1,p=0.25, training=self.training)
            x2 = F.tanh(self.fc2(x1))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x3 = F.tanh(self.fc3(x2))
            x3 = F.dropout(x3,p=0.25, training=self.training)
            x4 = F.tanh(self.fc4(x3))
            x4 = F.dropout(x4,p=0.25, training=self.training)
            x5 = F.tanh(self.fc5(x4))
            x5 = F.dropout(x5,p=0.25, training=self.training)
            x6 = self.fc6(x5)
            return x6
        if self.hiddenLayers == 6:
            x1 = F.tanh(self.fc1(x))
            x1 = F.dropout(x1,p=0.25, training=self.training)
            x2 = F.tanh(self.fc2(x1))
            x2 = F.dropout(x2,p=0.25, training=self.training)
            x3 = F.tanh(self.fc3(x2))
            x3 = F.dropout(x3,p=0.25, training=self.training)
            x4 = F.tanh(self.fc4(x3))
            x4 = F.dropout(x4,p=0.25, training=self.training)
            x5 = F.tanh(self.fc5(x4))
            x5 = F.dropout(x5,p=0.25, training=self.training)
            x6 = F.tanh(self.fc6(x5))
            x6 = F.dropout(x6,p=0.25, training=self.training)
            x7 = self.fc7(x6)
            return x7

class variableDepthNet_v2(nn.Module):
     def __init__(self, hiddenDims):
        super(variableDepthNet_v2, self).__init__()
        self.fc1 = nn.Linear(1, hiddenDims[0])

        # self.hiddenLayers = {}

        # for h in range(len(hiddenDims)-1):
        #     self.hiddenLayers[h] = nn.Linear(hiddenDims[h], hiddenDims[h+1]).cpu()

        self.hidden = nn.ModuleList()
        for k in range(len(hiddenDims)-1):
            self.hidden.append(nn.Linear(hiddenDims[k], hiddenDims[k+1]))

        self.fc2 = nn.Linear(hiddenDims[len(hiddenDims)-1], 1)

        # self.hidden = nn.ModuleList()
        # for k in range(len(hiddenDims)-1):
        #     self.hidden.append(nn.Linear(hiddenDims[k], hiddenDims[k+1]))

     def forward(self, x):
        x1 = F.leaky_relu(self.fc1(x))
        for h in range(len(self.hidden)):
            x1 = F.leaky_relu(self.hidden[h](x1))
            # x1 = F.dropout(x1,p=0.15)
        x1 = self.fc2(x1)
        return x1

class variableDepthNet_v3(nn.Module):
     def __init__(self, inSize, outSize,hiddenDims):
        super(variableDepthNet_v3, self).__init__()
        self.fc1 = nn.Linear(inSize, hiddenDims[0])

        # self.hiddenLayers = {}

        # for h in range(len(hiddenDims)-1):
        #     self.hiddenLayers[h] = nn.Linear(hiddenDims[h], hiddenDims[h+1]).cpu()

        self.hidden = nn.ModuleList()
        for k in range(len(hiddenDims)-1):
            self.hidden.append(nn.Linear(hiddenDims[k], hiddenDims[k+1]))

        self.fc2 = nn.Linear(hiddenDims[len(hiddenDims)-1], outSize)

        # self.hidden = nn.ModuleList()
        # for k in range(len(hiddenDims)-1):
        #     self.hidden.append(nn.Linear(hiddenDims[k], hiddenDims[k+1]))

     def forward(self, x):
        x1 = F.leaky_relu(self.fc1(x))
        for h in range(len(self.hidden)):
            x1 = F.leaky_relu(self.hidden[h](x1))
            # x1 = F.dropout(x1,p=0.15)
        x1 = self.fc2(x1)
        return x1

class META_LINEAR_MODEL(nn.Module):
     def __init__(self, inputDim, outputDim):
        super(META_LINEAR_MODEL, self).__init__()
        self.fc1 = nn.Linear(inputDim, 45)  # 45  65 ORIGINALLY
        self.fc2 = nn.Linear(45, 65)
        self.fc3 = nn.Linear(65, 45)
        self.fc4 = nn.Linear(45, outputDim)

     def forward(self, x):
        x1 = F.tanh(self.fc1(x))
        x2 = F.tanh(self.fc2(x1))
        x3 = F.tanh(self.fc3(x2))
        x4 = self.fc4(x3)
        return x4

class META_LINEAR_MODEL_2(nn.Module):
     def __init__(self, inputDim, outputDim):
        super(META_LINEAR_MODEL_2, self).__init__()
        self.fc1 = nn.Linear(inputDim, 55)
        self.fc2 = nn.Linear(55, 100)
        self.fc3 = nn.Linear(100, 55)
        self.fc4 = nn.Linear(55, outputDim)

     def forward(self, x):
        x1 = F.tanh(self.fc1(x))
        x2 = F.tanh(self.fc2(x1))
        x3 = F.tanh(self.fc3(x2))
        x4 = self.fc4(x3)
        return x4

class metaAE(nn.Module):
    def __init__(self, inputDim, outputDim):
        super(metaAE, self).__init__()
        self.fc1 = nn.Linear(inputDim, 15)
        self.fc2 = nn.Linear(15, outputDim)

    def encoder(self, x):
        h1 = self.fc1(x)
        return h1

    def decoder(self, z):
        h2 = self.fc2(z)
        return h2

    def forward(self, x):
        h1 = self.encoder(x)
        h2 = self.decoder(h1)
        return h2

class metaAE3(nn.Module):
    def __init__(self, inputDim, outputDim):
        super(metaAE3, self).__init__()
        self.fc1 = nn.Linear(inputDim, 2000)
        self.fc2 = nn.Linear(2000, 800)
        self.fc3 = nn.Linear(800, 2000)
        self.fc4 = nn.Linear(2000, outputDim)

    def encoder(self, x):
        h1 = F.tanh(self.fc1(x))
        h2 = self.fc2(h1)
        return h2

    def decoder(self, z):
        h3 = F.tanh(self.fc3(z))
        h4 = self.fc4(h3)
        return h4

    def forward(self, x):
        h1 = self.encoder(x)
        h2 = self.decoder(h1)
        return h2

class RegressionAgent(nn.Module):
    def __init__(self):
        super(RegressionAgent, self).__init__()
        self.fc1 = nn.Linear(40, 40)
        self.fc2 = nn.Linear(40, 40)

    def forward(self, x):
        x1 = self.fc1(x)
        x2 = self.fc2(x1)
        return x2

class metaModel_CR(nn.Module):
     def __init__(self, inputDim, outputDim):
        super(RegressionAgent, self).__init__()
        self.fc1 = nn.Linear(inputDim, 40)
        self.fc2 = nn.Linear(40, outputDim)

     def forward(self, x):
        x1 = self.fc1(x)
        x2 = self.fc2(x1)
        return x2

class CAE_91(nn.Module):
    def __init__(self):
        super(CAE, self).__init__()
        self.fc1 = nn.Linear(21432, 2000)
        self.fc2 = nn.Linear(2000,200)
        self.fc3 = nn.Linear(200,2000)
        self.fc4 = nn.Linear(2000, 21432)
        self.sigmoid = nn.Tanh()

    def encoder(self, x):
        h1 = self.sigmoid(self.fc1(x.view(-1, 21432)))
        return self.sigmoid(self.fc2(h1))

    def decoder(self, z):
        h2 = self.sigmoid(self.fc3(z))
        return self.sigmoid(self.fc4(h2))

    def forward(self, x):
        h1 = self.encoder(x)
        h2 = self.decoder(h1)
        return h1, h2

class CAE(nn.Module):
    def __init__(self):
        super(CAE, self).__init__()
        self.fc1 = nn.Linear(21432, 5)

        self.fc2 = nn.Linear(5, 21432)
        self.relu = nn.ELU()
        self.sigmoid = nn.ELU()

    def encoder(self, x):
        h1 = self.fc1(x.view(-1, 21432))
        return h1

    def decoder(self, z):
        h2 = (self.fc2(z))
        return h2

    def forward(self, x):
        h1 = self.encoder(x)
        h2 = self.decoder(h1)
        return h1, h2

class CAE_multiLayer(nn.Module):
    def __init__(self):
        super(CAE_multiLayer, self).__init__()
        self.fc1 = nn.Linear(21432, 4)
        self.fc11 = nn.Linear(4, 3)
        self.fc2 = nn.Linear(3, 4)
        self.fc22 = nn.Linear(4, 21432)
        self.elu = nn.ELU()

    def encoder(self, x):
        h1 = self.elu(self.fc1(x.view(-1, 21432)))
        h2 = self.fc11(h1)
        return h2

    def decoder(self, z):
        h3 = self.elu(self.fc2(z))
        h4 = self.fc22(h3)
        return h4

    def forward(self, x):
        h1 = self.encoder(x)
        h2 = self.decoder(h1)
        return h1, h2

"""
###########################################################################################
    (START) DAN-NET BUILDING BLOCKS
###########################################################################################
"""
#-----------------------------------------------------------------------------
class Single_Synapse(nn.Module):
    """
        A Single_Synapse is an MLP which distills n_synaptic channels into a single value
        traversing the connection.
    """
    def __init__(self, n_channels):
        super(Single_Synapse, self).__init__()
        self.fc1 = nn.Linear(n_channels, 3)
        self.fc2 = nn.Linear(3, 1)
    def forward(self, x):
        x1 = torch.tanh(self.fc1(x))
        x2 = torch.tanh(self.fc2(x1))
        return x2

#-----------------------------------------------------------------------------
class Node_Aggregator(nn.Module):
    """
        A Node_Aggregator is an MLP which distills the filtered signals arriving on num_inbound
        connections into a single outbound activation scalar.
    """
    def __init__(self, num_inbound):
        super(Node_Aggregator, self).__init__()
        self.fc1 = nn.Linear(num_inbound, 3)
        self.fc2 = nn.Linear(3, 1)
    def forward(self, x):
        x1 = torch.tanh(self.fc1(x))
        x2 = self.fc2(x1)
        return x2
#-----------------------------------------------------------------------------
class DAN(nn.Module):
    """
        A DAN is composed of num_inbound Single_Synapses and a single Node_Aggregator
            num_inbound:  number of inbound connections
            n_channels:  number of synaptic channels per connection

    """
    def __init__(self, num_inbound, n_channels, bias_vector_size=None, output_node=None):
        super(DAN, self).__init__()

        # Synapses ###########
        self.synapses = nn.ModuleDict()
        self.output_node = output_node
        self.num_inbound = num_inbound
        for n in range(num_inbound):
            self.synapses[str(n)] = Single_Synapse(n_channels)

        # Node Aggregator ##########
        if bias_vector_size is not None:
            self.node_aggregator = Node_Aggregator(num_inbound+bias_vector_size)
        else:
            self.node_aggregator = Node_Aggregator(num_inbound)

    def forward(self, x, bias_vector=None):
        # pass filtered input through the corresponding synapse
        self.synapse_activations = torch.Tensor(np.zeros(self.num_inbound))
        i = 0

        for synapse in self.synapses:
            self.synapse_activations[i] = self.synapses[synapse](x[i])
            i+=1
        # concatenate with the node location to bias the node & ...
        # pass the synaptic outputs through the node aggregator
        if bias_vector is not None:
            out = self.node_aggregator(torch.cat((self.synapse_activations, bias_vector)))
        else:
            out = self.node_aggregator(self.synapse_activations)
        # if self.output_node == False:
        #     out = torch.tanh(out)
        return out

class Simple_DAN(nn.Module):
    def __init__(self, n_channels):
        super(Simple_DAN, self).__init__()
        self.dendrites = nn.Linear(n_channels, 15)
        self.fc2 = nn.Linear(15, 8)
        self.fc3 = nn.Linear(8, 1)
    def forward(self, x):
        x1 = F.leaky_relu(self.dendrites(x))
        x2 = F.leaky_relu(self.fc2(x1))
        x3 = self.fc3(x2)
        return x3

#-----------------------------------------------------------------------------
class Single_VECs(nn.Module):
    """
        A VECS splits input along n synaptic channels

    """
    def __init__(self, n_channels):
        super(Single_VECs, self).__init__()
        self.fc1 = nn.Linear(1, n_channels)
    def forward(self, x):
        x1 = torch.tanh(self.fc1(x))
        return x1
#-----------------------------------------------------------------------------
class NETWORK_OF_DANS_and_VECS(nn.Module):
    """
        params:
            dims (array, ints): array of layer widths
            n_channels (int):  synaptic channels per connection
            VECS_layers (array): array of ModuleDicts, which hold the VECs (single synapses) for that layer
            DAN_layers (array): array of ModuleDicts which hold the DANs for that layer

        (1) we need to instantiate a VECS from every presynaptic node to every postsynaptic node.
        (2) we will therefore have a layer of VECS
        (3) followed by a layer of DANs
    """
    def __init__(self, dims, n_channels):
        super(NETWORK_OF_DANS_and_VECS, self).__init__()
        self.dims = dims
        self.n_channels = n_channels
        self.VECS_layers = nn.ModuleList() #nn.ModuleDict() # container (dictionary) for the layers of VECS
        self.DAN_topologies = []

        # # node bias ------------------------------------
        self.layer_codes = np.eye(len(dims))[range(len(dims))]
        self.node_codes = []
        for layer in range(len(dims)):
            layer_code = self.layer_codes[layer]
            for node_idx in range(dims[layer]):
                node_code = np.eye(np.max(dims))[node_idx]
                self.node_codes.append(node_code)
        self.layer_codes = torch.Tensor(self.layer_codes)
        self.node_codes = torch.Tensor(self.node_codes)
        self.bias_vector_size = len(torch.cat((self.layer_codes[0], self.node_codes[0])))

        # construct fully-connected layers of VECS -------------------------------------------
        for layer_idx in range(len(dims)-1):   #for each layer, instantiate a layer of VECs
            VECs = nn.ModuleDict()
            for pre_node in range(dims[layer_idx]):
                for post_node in range(dims[layer_idx+1]):
                    VECs["pre_"+str(pre_node)+"_post_"+str(post_node)] = Single_VECs(n_channels)
            self.VECS_layers.append(VECs)

        # construct layers of DANs ----------------------------------------------------
        self.DAN_layers = nn.ModuleList()
        for layer_idx, dim in enumerate(dims):
            layer_DANs = nn.ModuleDict()
            if layer_idx > 0:
                for node_idx in range(dim):
                    torch.manual_seed(9)   # initialize all DANs to same parameters
                    if node_idx == 0:
                        self.DAN_topologies.append([dims[layer_idx-1], n_channels, self.bias_vector_size])
                    if layer_idx == len(self.dims)-1:
                        output_node = True
                    else:
                        output_node = False
                    layer_DANs["node_"+str(node_idx)] = DAN(num_inbound=dims[layer_idx-1], n_channels=n_channels, bias_vector_size=self.bias_vector_size, output_node=output_node)
                self.DAN_layers.append(layer_DANs)

    def forward(self, x):
        self.layer_activations = {}
        # input layer is a special case, no VECs filtering necessary
        for node in range(self.dims[0]):
            self.layer_activations["layer_0_node_"+str(node)] = x[node].view(-1,1)

        # get the filtered input for each synapse on each DAN in the next layer --------------------
        for layer_idx in range(len(self.dims)-1):
            inbound_activation_vectors = torch.Tensor(np.zeros((self.dims[layer_idx+1], self.dims[layer_idx], self.n_channels))) # of size num_inbound_connections * n_channels

            for node in range(self.dims[layer_idx+1]):  # for each post-synaptic node
                # for each synapse on that post-synaptic node, equals number of ...
                # ... inbound connections == number of nodes in pre layer
                self.layer_activations["layer_"+str(layer_idx+1)+"_node_"+str(node)] = []
                for synapse in range(self.dims[layer_idx]):
                    node_key = "node_"+str(node)
                    synapse_key = str(synapse)
                    VECs_key = "pre_"+str(synapse_key)+"_post_"+str(node)
                    #get the filtered activation for that synapse, will be of size n_channels
                    filtered_act = self.VECS_layers[layer_idx][VECs_key](self.layer_activations["layer_"+str(layer_idx)+"_node_"+str(synapse)]) #clone not needed
                    inbound_activation_vectors[node][synapse] = filtered_act #clone not needed
                """
                    Once we have our filtered activations for each node, we need to pass them through the correct DAN for that layer, and
                    collect the output. which will itself be filtered by VECS in the next layer.
                """
                # node bias ------------------------------------
                bias_vector = torch.cat((self.layer_codes[layer_idx+1], self.node_codes[node]))
                # pass the synapse activations through the node
                node_out = self.DAN_layers[layer_idx]["node_"+str(node)](inbound_activation_vectors[node].clone(), bias_vector=bias_vector) #CLONE NEEDED
                self.layer_activations["layer_"+str(layer_idx+1)+"_node_"+str(node)] = node_out.view(-1,1) #clone not needed
        net_out = torch.Tensor()
        for node in range(self.dims[-1]):
            net_out = torch.cat((self.layer_activations["layer_"+str(len(self.dims)-1)+"_node_"+str(node)], net_out)) #clone not needed
        return net_out

    def avg_DAN_grads(self):
        for layer_idx, dan_layer in enumerate(self.DAN_layers):
            mu_DAN = DAN(num_inbound=self.DAN_topologies[layer_idx][0], n_channels=self.n_channels, bias_vector_size=self.DAN_topologies[layer_idx][2])
            for name, p in mu_DAN.named_parameters():
                p.data.fill_(0)
            for node in dan_layer:
                for (node_paramName, node_param), (muDAN_paramName, mu_param) in zip(dan_layer[node].named_parameters(),mu_DAN.named_parameters()):
                    mu_param.data = mu_param.data + node_param.grad
            for name, p in mu_DAN.named_parameters():
                p.data = p.data / len(dan_layer)
            for node in dan_layer:
                for (node_paramName, node_param), (muDAN_paramName, mu_param) in zip(dan_layer[node].named_parameters(),mu_DAN.named_parameters()):

                    # print("pre", node_paramName, node_param.grad)
                    node_param.grad.data = mu_param

    def avg_DANs(self):
        for layer_idx, dan_layer in enumerate(self.DAN_layers):
            mu_DAN = DAN(num_inbound=self.DAN_topologies[layer_idx][0], n_channels=self.n_channels, bias_vector_size=self.DAN_topologies[layer_idx][2])
            for name, p in mu_DAN.named_parameters():
                p.data.fill_(0)
            for node in dan_layer:
                for (node_paramName, node_param), (muDAN_paramName, mu_param) in zip(dan_layer[node].named_parameters(),mu_DAN.named_parameters()):
                    mu_param.data = mu_param.data + node_param
            for name, p in mu_DAN.named_parameters():
                p.data = p.data / len(dan_layer)
            for node in dan_layer:
                for (node_paramName, node_param), (muDAN_paramName, mu_param) in zip(dan_layer[node].named_parameters(),mu_DAN.named_parameters()):

                    # print("pre", node_paramName, node_param.grad)
                    node_param.data = mu_param

                    # print("post", node_paramName, node_param.grad)
                    # print("mu_param",  mu_param)

            #     print("^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^")
            # print("############################################################")
        # sys.exit()
#-----------------------------------------------------------------------------
class NETWORK_OF_DANS(nn.Module):
    """
        params:
            dims (array of ints): array of layer widths.  Must be multiples of n_channels
            n_channels (int):  denotes input dimensions for DANs
            DAN (MLP):  Simple_DAN, distibuted genome
            VECS_layers (dict): ModuleDicts which hold the adapt parameters for that layer

        (1) we need to instantiate a fully-connected layer from the presynaptic layer (layer_idx) to the
            post-synaptic layer (layer_idx+1)
        (2) we will therefore have a layer of VECS (theta)
        (3) followed by a layer of DANs (which themselves are multi-layer MLPs)
            - in practice, each layer has but a single DAN, which is distributed throughout the layer
    """
    def __init__(self, dims, n_channels, genome=None):
        super(NETWORK_OF_DANS, self).__init__()
        self.dims = dims
        self.n_channels = n_channels
        self.VECS_layers = nn.ModuleDict() #nn.ModuleDict() # container (dictionary) for the layers of VECS
        self.Skip_layers = nn.ModuleDict()
        #self.DAN_topologies = []
        #self.DAN_layers = nn.ModuleDict()#Simple_DAN(self.n_channels)
        self.layer_bias_codes = torch.Tensor(np.eye(len(self.dims)-1)).cpu()
        if genome is not None:
            self.DAN = genome
        else:
            self.DAN = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))
        self.dan_acts = torch.Tensor().cpu()

        self.DAN_out = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))

        self.Skip_layers["0_2"] = nn.Linear(1, self.dims[2]*self.n_channels)
        self.Skip_layers["1_3"] = nn.Linear(self.dims[1], self.n_channels)

        # construct fully-connected layers of VECS -------------------------------------------
        for l in range(len(self.dims)-1):
            self.VECS_layers["layer_"+str(l)] = nn.Linear(self.dims[l], self.dims[l+1]*n_channels)
            #if we want sparse VECS
            # if l > 0:
            #     i = 0
            #     for p in self.VECS_layers["layer_"+str(l)].parameters():
            #         if i == 0:
            #             mask = np.zeros(tuple(p.t().shape))
            #             for idx, row in enumerate(mask):
            #                 mask[idx][idx::self.n_channels] = 1
            #             mask = mask.transpose()

            #             p.data = p * torch.Tensor(mask)
            #         i+=1

        #construct a DAN for each layer -------------------------------------------
        # for l in range(len(self.dims)-1):
        #     self.DAN_layers["layer_"+str(l)] = Simple_DAN(n_channels)

    # if we want sparse VECS
    def maskVECS(self):
        for l in range(len(self.dims)-1):
            if l > 0:
                i = 0
                for p in self.VECS_layers.Standard_layers["layer_"+str(l)].parameters():
                    if i == 0:
                        mask = np.zeros(tuple(p.t().shape))
                        for idx, row in enumerate(mask):
                            mask[idx][idx::self.n_channels] = 1
                        mask = mask.transpose()
                        p.data = p * torch.Tensor(mask)
                    i+=1

    def forward(self, x):
        hidden_acts = x
        #node_idx = 0
        layer_idx = 0
        self.skip_acts_0_2 = torch.Tensor().cpu()
        self.skip_acts_1_3 = torch.Tensor().cpu()
        for layer_key in self.VECS_layers:
            #print("hidden acts pre", hidden_acts)
            # print("-----------------------")
            # for p in self.VECS_layers[layer_key].parameters():

            #     print(p)
            #     print(p.shape)
            # print("--------------------------")

            hidden_acts_copy = hidden_acts.clone().detach()
            hidden_acts = self.VECS_layers[layer_key](hidden_acts)

            if layer_idx == 1:
                self.skip_acts_0_2 = self.Skip_layers["0_2"](x)
                self.skip_acts_1_3 = self.Skip_layers["1_3"](hidden_acts_copy)
                hidden_acts = hidden_acts.add_(self.skip_acts_0_2)
            if layer_idx == 2:
                hidden_acts = hidden_acts.add_(self.skip_acts_1_3)

            #hidden_acts = torch.tanh(hidden_acts)

            #slice the hidden activations for the layer into sub-vectors to be passed through the DAN
            self.dan_acts = torch.Tensor().cpu()
            node_idx_ext = 0
            for z in [hidden_acts[0][node_idx:node_idx + self.n_channels] for node_idx in range(0, len(hidden_acts[0]), self.n_channels)]:
                #concat the output of each DAN in the layer to form the activation vector for the layer of DANs
                #print("node_idx", node_idx)
                if layer_idx == len(self.dims)-1:
                    #self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key](z)))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, self.DAN(z)))
                else:
                    #self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key](z))))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN(z))))
                node_idx_ext += 1
            #sys.exit()
            #assign as hidden_acts to be passed through the next layer of VECS
            hidden_acts = self.dan_acts.view(1, len(self.dan_acts))
            layer_idx +=1

        return hidden_acts

###########################################################################################################
class NETWORK_OF_DANS_batch(nn.Module):
    """
        params:
            dims (array of ints): array of layer widths.  Must be multiples of n_channels
            n_channels (int):  denotes input dimensions for DANs
            DAN (MLP):  Simple_DAN, distibuted genome
            VECS_layers (dict): ModuleDicts which hold the adapt parameters for that layer

        (1) we need to instantiate a fully-connected layer from the presynaptic layer (layer_idx) to the
            post-synaptic layer (layer_idx+1)
        (2) we will therefore have a layer of VECS (theta)
        (3) followed by a layer of DANs (which themselves are multi-layer MLPs)
            - in practice, each layer has but a single DAN, which is distributed throughout the layer
    """
    def __init__(self, dims, n_channels, batch_size=1,genome=None,):
        super(NETWORK_OF_DANS_batch, self).__init__()
        self.dims = dims
        self.batch_size = batch_size
        self.num_DAN_layers = len(dims)-1
        self.n_channels = n_channels
        self.VECS_layers = nn.ModuleDict() #nn.ModuleDict() # container (dictionary) for the layers of VECS
        self.Skip_layers = nn.ModuleDict()
        #self.DAN_topologies = []

        #self.DAN_layers = nn.ModuleDict()#Simple_DAN(self.n_channels)
        self.layer_bias_codes = torch.Tensor(np.eye(len(self.dims)-1)).cpu()
        self.bias_codes = torch.zeros((self.num_DAN_layers,self.batch_size,self.num_DAN_layers))
        for l_idx in range(len(self.layer_bias_codes)):
            for y in range(self.batch_size):
                bias_zeros = torch.zeros((self.batch_size, len(self.layer_bias_codes[0])))
                for z_idx in range(len(bias_zeros)):
                    bias_zeros[z_idx] = self.layer_bias_codes[l_idx]

                self.bias_codes[l_idx] = bias_zeros
        self.layer_bias_codes = self.bias_codes

        if genome is not None:
            self.DAN = genome
        else:
            self.DAN = Simple_DAN(n_channels+len(self.layer_bias_codes[0][0]))
        self.dan_acts = torch.Tensor().cpu()

        self.Skip_layers["0_2"] = nn.Linear(self.dims[0], self.dims[2]*self.n_channels)
        self.Skip_layers["1_3"] = nn.Linear(self.dims[1], self.dims[3]*self.n_channels)

        # construct fully-connected layers of VECS -------------------------------------------
        for l in range(len(self.dims)-1):
            self.VECS_layers["layer_"+str(l)] = nn.Linear(self.dims[l], self.dims[l+1]*n_channels)

            # #for unique
            # for node_idx in range(self.dims[l+1]):
            #     self.DAN_layers["layer_"+str(l)+"_node_"+str(node_idx)] = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))

    def forward(self, x):
        hidden_acts = x
        #node_idx = 0
        layer_idx = 0
        self.skip_acts_0_2 = torch.Tensor().cpu()
        self.skip_acts_1_3 = torch.Tensor().cpu()
        for layer_key in self.VECS_layers:

            hidden_acts_copy = hidden_acts.clone().detach()
            hidden_acts = self.VECS_layers[layer_key](hidden_acts)

            if layer_idx == 1:

                self.skip_acts_0_2 = self.Skip_layers["0_2"](x)
                self.skip_acts_1_3 = self.Skip_layers["1_3"](hidden_acts_copy)
                hidden_acts = hidden_acts.add_(self.skip_acts_0_2)
            if layer_idx == 2:
                hidden_acts = hidden_acts.add_(self.skip_acts_1_3)

            #hidden_acts = torch.tanh(hidden_acts)

            #slice the hidden activations for the layer into sub-vectors to be passed through the DAN
            self.dan_acts = torch.Tensor().cpu()
            node_idx_ext = 0

            #print("hidden_acts", hidden_acts.size())
            for z in [hidden_acts[:,node_idx:node_idx + self.n_channels] for node_idx in range(0, len(hidden_acts[0]), self.n_channels)]:
                #concat the output of each DAN in the layer to form the activation vector for the layer of DANs

                if layer_idx == len(self.dims)-1:
                    #self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key](z)))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]), dim=1)
                    self.dan_acts = torch.cat((self.dan_acts, self.DAN(z)), dim=1)
                else:
                    #self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key](z))))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]), dim=1)
                    self.dan_acts = torch.cat((self.dan_acts, F.leaky_relu(self.DAN(z))), dim=1)

                node_idx_ext += 1

            #assign as hidden_acts to be passed through the next layer of VECS

            hidden_acts = self.dan_acts.view(self.dan_acts.size()[0], self.dan_acts.size()[1])
            layer_idx +=1

        return hidden_acts
##########  End DANs batch #########################################################################

##########################################################################################################
class NETWORK_OF_DANS_perLayerGenomes(nn.Module):
    """
        params:
            dims (array of ints): array of layer widths.  Must be multiples of n_channels
            n_channels (int):  denotes input dimensions for DANs
            DAN (MLP):  Simple_DAN, distibuted genome
            VECS_layers (dict): ModuleDicts which hold the adapt parameters for that layer

        (1) we need to instantiate a fully-connected layer from the presynaptic layer (layer_idx) to the
            post-synaptic layer (layer_idx+1)
        (2) we will therefore have a layer of VECS (theta)
        (3) followed by a layer of DANs (which themselves are multi-layer MLPs)
            - in practice, each layer has but a single DAN, which is distributed throughout the layer
    """
    def __init__(self, dims, n_channels, genome=None):
        super(NETWORK_OF_DANS_perLayerGenomes, self).__init__()
        self.dims = dims
        self.n_channels = n_channels
        self.VECS_layers = nn.ModuleDict() #nn.ModuleDict() # container (dictionary) for the layers of VECS
        self.Skip_layers = nn.ModuleDict()
        #self.DAN_topologies = []
        self.DAN_layers = nn.ModuleDict()#Simple_DAN(self.n_channels)
        self.layer_bias_codes = torch.Tensor(np.eye(len(self.dims)-1)).cpu()
        # if genome is not None:
        #     self.DAN = genome
        # else:
        #     self.DAN = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))
        # self.dan_acts = torch.Tensor().cpu()

        # self.DAN_out = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))

        self.Skip_layers["0_2"] = nn.Linear(1, self.dims[2]*self.n_channels)
        self.Skip_layers["1_3"] = nn.Linear(self.dims[1], self.n_channels)

        # construct fully-connected layers of VECS -------------------------------------------
        for l in range(len(self.dims)-1):
            self.VECS_layers["layer_"+str(l)] = nn.Linear(self.dims[l], self.dims[l+1]*n_channels)
            self.DAN_layers["layer_"+str(l)] = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))
            #if we want sparse VECS
            # if l > 0:
            #     i = 0
            #     for p in self.VECS_layers["layer_"+str(l)].parameters():
            #         if i == 0:
            #             mask = np.zeros(tuple(p.t().shape))
            #             for idx, row in enumerate(mask):
            #                 mask[idx][idx::self.n_channels] = 1
            #             mask = mask.transpose()

            #             p.data = p * torch.Tensor(mask)
            #         i+=1

        #construct a DAN for each layer -------------------------------------------
        # for l in range(len(self.dims)-1):
        #     self.DAN_layers["layer_"+str(l)] = Simple_DAN(n_channels)

    # if we want sparse VECS
    def maskVECS(self):
        for l in range(len(self.dims)-1):
            if l > 0:
                i = 0
                for p in self.VECS_layers.Standard_layers["layer_"+str(l)].parameters():
                    if i == 0:
                        mask = np.zeros(tuple(p.t().shape))
                        for idx, row in enumerate(mask):
                            mask[idx][idx::self.n_channels] = 1
                        mask = mask.transpose()
                        p.data = p * torch.Tensor(mask)
                    i+=1

    def forward(self, x):
        hidden_acts = x
        #node_idx = 0
        layer_idx = 0
        self.skip_acts_0_2 = torch.Tensor().cpu()
        self.skip_acts_1_3 = torch.Tensor().cpu()
        for layer_key in self.VECS_layers:
            #print("hidden acts pre", hidden_acts)
            # print("-----------------------")
            # for p in self.VECS_layers[layer_key].parameters():

            #     print(p)
            #     print(p.shape)
            # print("--------------------------")

            hidden_acts_copy = hidden_acts.clone().detach()
            hidden_acts = self.VECS_layers[layer_key](hidden_acts)

            if layer_idx == 1:
                self.skip_acts_0_2 = self.Skip_layers["0_2"](x)
                self.skip_acts_1_3 = self.Skip_layers["1_3"](hidden_acts_copy)
                hidden_acts = hidden_acts.add_(self.skip_acts_0_2)
            if layer_idx == 2:
                hidden_acts = hidden_acts.add_(self.skip_acts_1_3)

            #hidden_acts = torch.tanh(hidden_acts)

            #slice the hidden activations for the layer into sub-vectors to be passed through the DAN
            self.dan_acts = torch.Tensor().cpu()
            node_idx_ext = 0
            for z in [hidden_acts[0][node_idx:node_idx + self.n_channels] for node_idx in range(0, len(hidden_acts[0]), self.n_channels)]:
                #concat the output of each DAN in the layer to form the activation vector for the layer of DANs
                #print("node_idx", node_idx)
                if layer_idx == len(self.dims)-1:
                    #self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key](z)))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key](z)))
                else:
                    #self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key](z))))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key](z))))
                node_idx_ext += 1
            #sys.exit()
            #assign as hidden_acts to be passed through the next layer of VECS
            hidden_acts = self.dan_acts.view(1, len(self.dan_acts))
            layer_idx +=1

        return hidden_acts

class NETWORK_OF_DANS_uniqueNeurons(nn.Module):
    """
        params:
            dims (array of ints): array of layer widths.  Must be multiples of n_channels
            n_channels (int):  denotes input dimensions for DANs
            DAN (MLP):  Simple_DAN, distibuted genome
            VECS_layers (dict): ModuleDicts which hold the adapt parameters for that layer

        (1) we need to instantiate a fully-connected layer from the presynaptic layer (layer_idx) to the
            post-synaptic layer (layer_idx+1)
        (2) we will therefore have a layer of VECS (theta)
        (3) followed by a layer of DANs (which themselves are multi-layer MLPs)
            - in practice, each layer has but a single DAN, which is distributed throughout the layer
    """
    def __init__(self, dims, n_channels, genome=None):
        super(NETWORK_OF_DANS_uniqueNeurons, self).__init__()
        self.dims = dims
        self.n_channels = n_channels
        self.VECS_layers = nn.ModuleDict() #nn.ModuleDict() # container (dictionary) for the layers of VECS
        self.Skip_layers = nn.ModuleDict()
        #self.DAN_topologies = []
        self.DAN_layers = nn.ModuleDict()#Simple_DAN(self.n_channels)
        self.layer_bias_codes = torch.Tensor(np.eye(len(self.dims)-1)).cpu()

        self.Skip_layers["0_2"] = nn.Linear(1, self.dims[2]*self.n_channels)
        self.Skip_layers["1_3"] = nn.Linear(self.dims[1], self.n_channels)

        # construct fully-connected layers of VECS -------------------------------------------
        for l in range(len(self.dims)-1):
            self.VECS_layers["layer_"+str(l)] = nn.Linear(self.dims[l], self.dims[l+1]*n_channels)
            for node_idx in range(self.dims[l+1]):
                self.DAN_layers["layer_"+str(l)+"_node_"+str(node_idx)] = Simple_DAN(n_channels+len(self.layer_bias_codes[0]))

    def forward(self, x):
        hidden_acts = x
        #node_idx = 0
        layer_idx = 0
        self.skip_acts_0_2 = torch.Tensor().cpu()
        self.skip_acts_1_3 = torch.Tensor().cpu()
        for layer_key in self.VECS_layers:

            hidden_acts_copy = hidden_acts.clone().detach()
            hidden_acts = self.VECS_layers[layer_key](hidden_acts)

            if layer_idx == 1:
                self.skip_acts_0_2 = self.Skip_layers["0_2"](x)
                self.skip_acts_1_3 = self.Skip_layers["1_3"](hidden_acts_copy)
                hidden_acts = hidden_acts.add_(self.skip_acts_0_2)
            if layer_idx == 2:
                hidden_acts = hidden_acts.add_(self.skip_acts_1_3)

            #hidden_acts = torch.tanh(hidden_acts)

            #slice the hidden activations for the layer into sub-vectors to be passed through the DAN
            self.dan_acts = torch.Tensor().cpu()
            node_idx_ext = 0
            for z in [hidden_acts[0][node_idx:node_idx + self.n_channels] for node_idx in range(0, len(hidden_acts[0]), self.n_channels)]:
                #concat the output of each DAN in the layer to form the activation vector for the layer of DANs
                #print("node_idx", node_idx)
                if layer_idx == len(self.dims)-1:
                    #self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key](z)))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, self.DAN_layers[layer_key+"_node_"+str(node_idx_ext)](z)))
                else:
                    #self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key](z))))
                    z = torch.cat((z,self.layer_bias_codes[layer_idx]))
                    self.dan_acts = torch.cat((self.dan_acts, torch.tanh(self.DAN_layers[layer_key+"_node_"+str(node_idx_ext)](z))))
                node_idx_ext += 1

            #assign as hidden_acts to be passed through the next layer of VECS
            hidden_acts = self.dan_acts.view(1, len(self.dan_acts))
            layer_idx +=1

        return hidden_acts


text:
#!/usr/bin/env python
# coding: utf-8

# # Preferential Bayesian Optimization: Multinomial Predictive Entropy Search

# In[ ]:

import numpy as np
import gpflow
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt
import sys
import os
import pickle

from gpflow.utilities import set_trainable, print_summary
gpflow.config.set_default_summary_fmt("notebook")

sys.path.append(os.path.split(os.path.split(os.path.split(os.getcwd())[0])[0])[0]) # Move 3 levels up directory to import project files as module
import importlib
PBO = importlib.import_module("Top-k-Ranking-Bayesian-Optimization")

# In[ ]:

gpu_to_use = 0

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    # Restrict TensorFlow to only use the first GPU
    try:
        for gpu in gpus:
              tf.config.experimental.set_memory_growth(gpu, True)
        tf.config.experimental.set_visible_devices(gpus[gpu_to_use], 'GPU')
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
    except RuntimeError as e:
        # Visible devices must be set before GPUs have been initialized
        print(e)

# In[ ]:

features = pickle.load( open( "sushi_features.p", "rb" ) )

# In[ ]:

fvals = pickle.load( open( "fvals.p", "rb" ) )

# In[ ]:

# construct dict
feat_to_fval_dict = {}
for i in range(len(features)):
    key = features[i].data.tobytes()
    feat_to_fval_dict[key] = fvals[i]

# In[ ]:

objective = lambda x: PBO.objectives.sushi(x, feat_to_fval_dict)
objective_low = np.min(features)
objective_high = np.max(features)
objective_name = "SUSHI"
acquisition_name = "MPES"
experiment_name = acquisition_name + "_" + objective_name

# In[ ]:

num_runs = 10
num_evals = 35
num_choices = 2
input_dims = 6
num_maximizers = 20
num_maximizers_init = 50
num_fourier_features = 1000
num_init_prefs = 10

# In[ ]:

results_dir = os.getcwd() + '/results/' + experiment_name + '/'

try:
    # Create target Directory
    os.makedirs(results_dir)
    print("Directory " , results_dir ,  " created ")
except FileExistsError:
    print("Directory " , results_dir ,  " already exists")

# In[ ]:

def get_noisy_observation(X, objective):
    f = PBO.objectives.objective_get_f_neg(X, objective)
    return PBO.observation_model.gen_observation_from_f(X, f, 1)

# In[ ]:

def train_and_visualize(X, y, title, lengthscale_init=None, signal_variance_init=None):

    # Train model with data
    result = PBO.models.learning_fullgp.train_model_fullcov(
                        X, y,
                        obj_low=objective_low,
                        obj_high=objective_high,
                        lengthscale_init=lengthscale_init,
                        signal_variance_init=signal_variance_init,
                        indifference_threshold=0.,
                        n_sample=1000,
                        deterministic=True, # only sample f values once, not re-sampling
                        num_steps=3000)

    q_mu = result['q_mu']
    q_sqrt = result['q_sqrt']
    u = result['u']
    inputs = result['inputs']
    k = result['kernel']

    likelihood = gpflow.likelihoods.Gaussian()
    model = PBO.models.learning.init_SVGP_fullcov(q_mu, q_sqrt, u, k, likelihood)
    u_mean = q_mu.numpy()
    inducing_vars = u.numpy()

    return model, inputs, u_mean, inducing_vars

# Generate rank dictionary and immediate regret dictionary.

# In[ ]:

fval_idx_tuples = pickle.load(open("fval_idx_tuples.p", "rb"))

# In[ ]:

rank_dict = {}

for i in range(len(fval_idx_tuples)):
    rank_dict[features[fval_idx_tuples[i][1]].data.tobytes()] = i + 1

# This function is our main metric for the performance of the acquisition function.

# In[ ]:

def get_max_sushi(model, features, rank_dict):
    """
    :param model: gpflow model
    :param features: sushi features
    :param rank_dict: dictionary from sushi idx to place in ranking
    :return: tuple (index of max sushi, rank)
    """
    f_preds = model.predict_f(features)[0]
    max_idx = np.argmax(f_preds)

    return (max_idx, rank_dict[features[max_idx].data.tobytes()])

# Store the results in these arrays:

# In[ ]:

num_data_at_end = int(num_init_prefs + num_evals)
X_results = np.zeros([num_runs, num_data_at_end, num_choices, input_dims])
y_results = np.zeros([num_runs, num_data_at_end, 1, input_dims])
immediate_regret = np.zeros([num_runs, num_evals], np.int32)

# Create the initial values for each run:

# In[ ]:

np.random.seed(0)
random_indices = np.random.choice(features.shape[0], [num_runs, num_init_prefs, num_choices])
init_vals = np.take(features, random_indices, axis=0)

# The following loops carry out the Bayesian optimization algorithm over a number of runs, with a fixed number of evaluations per run.

# In[ ]:

for run in range(num_runs):  # CHECK IF STARTING RUN IS CORRECT
    print("Beginning run %s" % (run))

    X = init_vals[run]
    y = get_noisy_observation(X, objective)

    model, inputs, u_mean, inducing_vars = train_and_visualize(X, y, "Run_{}:_Initial_model".format(run))

    for evaluation in range(num_evals):
        print("Beginning evaluation %s" % (evaluation))

        success = False
        fail_count = 0
        while not success:
            # TODO: THIS ONLY WORKS FOR TOP-1 OF 2, CHANGE TO APPROPRIATE QUERY SAMPLING FOR HIGHER NUMBER OF CHOICES
            samples = PBO.models.learning_fullgp.construct_input_pairs(inputs, features)

            # Sample maximizers
            print("Evaluation %s: Sampling maximizers" % (evaluation))
            maximizers = PBO.fourier_features.sample_maximizers(X=inducing_vars,
                                                                count=num_maximizers,
                                                                n_init=num_maximizers_init,
                                                                D=num_fourier_features,
                                                                model=model,
                                                                min_val=objective_low,
                                                                max_val=objective_high)
            print(maximizers)

            # Calculate PES value I for each possible next query
            print("Evaluation %s: Calculating I" % (evaluation))
            I_vals = PBO.acquisitions.pes.I_batch(samples, maximizers, model)

            # Select query that maximizes I
            next_idx = np.argmax(I_vals)
            next_query = samples[next_idx]
            print("Evaluation %s: Next query is %s with I value of %s" % (evaluation, next_query, I_vals[next_idx]))

            X_temp = np.concatenate([X, [next_query]])
            # Evaluate objective function
            y_temp = np.concatenate([y, get_noisy_observation(np.expand_dims(next_query, axis=0), objective)], axis=0)

            try:
                print("Evaluation %s: Training model" % (evaluation))
                model, inputs, u_mean, inducing_vars = train_and_visualize(X_temp, y_temp,
                                                                           "Run_{}_Evaluation_{}".format(run, evaluation))
                success = True

            except ValueError as err:
                print(err)
                print("Retrying sampling random inputs")
                fail_count += 1

            if fail_count >= 3:
                print("Retry limit exceeded")
                raise ValueError("Failed")

        X = X_temp
        y = y_temp

        # Save model
        pickle.dump((X, y, inputs,
                     model.kernel.variance,
                     model.kernel.lengthscale,
                     model.likelihood.variance,
                     inducing_vars,
                     model.q_mu,
                     model.q_sqrt,
                     maximizers),
                    open(results_dir + "Model_Run_{}_Evaluation_{}.p".format(run, evaluation), "wb"))

        (max_idx, rank) = get_max_sushi(model, features, rank_dict)
        immediate_regret[run, evaluation] = rank - 1

        print("Maximizing sushi has index {} and rank {}".format(max_idx, rank))

    X_results[run] = X
    y_results[run] = y
    print("Run {} immediate regret: ".format(run))
    print(immediate_regret[run])

# In[ ]:

pickle.dump((X_results, y_results, immediate_regret), open(results_dir + "res.p", "wb"))

# In[ ]:

ir = immediate_regret
mean = np.mean(ir, axis=0)
std_dev = np.std(ir, axis=0)
std_err = std_dev / np.sqrt(ir.shape[0])

# In[ ]:

print("Mean immediate regret at each evaluation averaged across all runs:")
print(mean)

# In[ ]:

print("Standard error of immediate regret at each evaluation averaged across all runs:")
print(std_err)

# In[ ]:

with open(results_dir + acquisition_name + "_" + objective_name + "_" + "mean_sem" + ".txt", "w") as text_file:
    print("Mean immediate regret at each evaluation averaged across all runs:", file=text_file)
    print(mean, file=text_file)
    print("Standard error of immediate regret at each evaluation averaged across all runs:", file=text_file)
    print(std_err, file=text_file)

# In[ ]:

pickle.dump((mean, std_err), open(results_dir + acquisition_name + "_" + objective_name + "_" + "mean_sem.p", "wb"))

ext: py

max_stars_repo_name: sebtsh/PBO

max_issues_repo_name: sebtsh/PBO

max_forks_repo_name: sebtsh/PBO

reason: import numpy


text:
import numpy as np

if __name__ == '__main__':
    # path = './drugs/Fentanyl500ugResults/FentanylScenarioResults.txt'
    path = './drugs/Succinylcholine5mgResults/SuccinylcholineScenarioResults.txt'
    data = np.loadtxt(path, delimiter='\t', skiprows=1)
    selected = [0, 2, 8, 9, 14, 15]
    selected_data = data[:,selected]
    np.savetxt("./drugs/SuccinylcholineScenarioResults.csv", selected_data,
            fmt='%.2f', delimiter=',', header="time,hr,sp,dp,rr,spo2",comments='')

max_stars_repo_name: ScreenBasedSimulator/BioGears-Paper

max_issues_repo_name: ScreenBasedSimulator/BioGears-Paper

max_forks_repo_name: ScreenBasedSimulator/BioGears-Paper


text:
# -*- coding: utf-8 -*-
#  Copyright (c) 2021. Distributed under the terms of the MIT License.

from scipy.constants import physical_constants as pc
au_to_angstrom = pc["atomic unit of length"][0] * 10 ** 10  # = 0.5291772109

max_stars_repo_name: kumagai-group/vise

max_issues_repo_name: kumagai-group/vise

max_forks_repo_name: kumagai-group/vise

reason: from scipy


text:
# Preprocess data and plot

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import seaborn as sns
import json_tricks as json
from collections import OrderedDict
import os
import sys

sns.set(style="darkgrid")

# AVERAGE_CHILDREN_BY_AGE = pd.read_excel("Births and Fertility.xlsx", sheet_name="T1")
AVERAGE_CHILDREN_BY_EDUCATION = pd.read_excel("Births and Fertility.xlsx", sheet_name="T2")
BIRTHS_AND_FERTILITY = pd.read_excel("Births and Fertility.xlsx", sheet_name="T3")
# AVERAGE_CHILDREN = pd.read_excel("Births and Fertility.xlsx", sheet_name="T1")
MARRIAGE_RATE_BY_AGE = pd.read_excel("Marriage.xlsx", sheet_name="T2")
# MARRIAGE_RATE_BY_AGE = pd.read_excel("Marriage.xlsx", sheet_name="T2")

AVERAGE_CHILDREN_BY_EDUCATION = AVERAGE_CHILDREN_BY_EDUCATION.iloc[4:12, 1:]
BIRTHS_AND_FERTILITY = BIRTHS_AND_FERTILITY.iloc[4:21, 1:]
MALE_MARRIAGE_RATE_BY_AGE = MARRIAGE_RATE_BY_AGE.iloc[5:17, 1:]
FEMALE_MARRIAGE_RATE_BY_AGE = MARRIAGE_RATE_BY_AGE.iloc[17:30, 1:]

x=AVERAGE_CHILDREN_BY_EDUCATION.iloc[0,:]
# print(x)

AVERAGE_CHILDREN_BY_EDUCATION = AVERAGE_CHILDREN_BY_EDUCATION.astype(np.float64)
print(AVERAGE_CHILDREN_BY_EDUCATION.dtypes)

plot1 = sns.lineplot(data=AVERAGE_CHILDREN_BY_EDUCATION, x=AVERAGE_CHILDREN_BY_EDUCATION.iloc[0],
                     y=AVERAGE_CHILDREN_BY_EDUCATION.iloc[1],)
plot1.set_title('AVERAGE_CHILDREN_BY_EDUCATION')
plot1.set_ylabel('avaerage birth rate')
plot1.set_xlabel('year')

fig1 = plot1.get_figure()
fig1.savefig("output_1.png")

# print(BIRTHS_AND_FERTILITY)

max_issues_repo_name: c-liangyu/ShiroDarumas

max_forks_repo_name: c-liangyu/ShiroDarumas


text:
#!/usr/bin/env python
# coding: utf-8

import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from mtcnn.mtcnn import MTCNN
import sys,os
detector = MTCNN()

def extract_face(image):
    try:
        res = detector.detect_faces(image)
    except:
        return
    x, y, x1, y1 = res[0]['box']

    height = y+y1
    width = x+x1
    return x, y, height, width

def create_folder():
    name = input("Enter the name of the person:")
    path = os.getcwd()+"\data_set\\train"
    path = path+"\\"+name
    try:
        os.mkdir(path)
    except:
        print("Directory available")
    return name, path

def capture_faces():
    name, path =create_folder()
    cap = cv2.VideoCapture(0)
    res, frame = cap.read()
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    cap.set(cv2.CAP_PROP_FPS,25)
    itr = 1
    while cap.isOpened():
        print(itr)
        x, y, height, width = extract_face(frame)
        #height = int(height +(height*0.15))
        #width = int(width + (width*0.05))
        #x=x+(x*0.30)
        #img_out = frame[y:height, x:width]
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        #img_out = cv2.resize(img_out, (200,200))
        plt.imsave(path+"\\"+"Madhan_"+str(itr)+".jpg", frame)
        img = cv2.rectangle(frame, (x,y),(width, height), (0,0,255), 2)

        img = cv2.putText(img, "Image #: "+str(x)+","+str(y+1),(50,50),cv2.FONT_HERSHEY_SIMPLEX , 0.5, (0,0,255), 2 )
        if itr==15 :
            break
        cv2.imshow("windows", img)
        itr=itr+1

        res, frame =cap.read()
        if cv2.waitKey(200) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    return

capture_faces()

max_stars_repo_name: gcmadhan/Face-ID-from-webcam-or-Video

max_issues_repo_name: gcmadhan/Face-ID-from-webcam-or-Video

max_forks_repo_name: gcmadhan/Face-ID-from-webcam-or-Video


text:
import lmfit
import numpy
import scipy.optimize
import numbers
import copy

import suspect.basis

def complex_to_real(complex_fid):
    """This function takes a complex FID and constructs a real version.

    Standard optimization routines as used in lmfit require real data.
    Concatenates the imaginary part to the complex part.
    The imaginary part is also reversed to keep the maxima at each end of the FID
    and avoid discontinuities in the center.

    Parameters
    ----------
    complex_fid : MRSData instance
        The complex FID to be converted to real.

    Returns
    -------
    real_fid : MRSData instance
        The real FID, which has twice as many points as the input.

    """
    np = complex_fid.shape[0]
    real_fid = numpy.zeros(np * 2)
    real_fid[:np] = complex_fid.real
    real_fid[np:] = complex_fid.imag[::-1]
    return real_fid

def real_to_complex(real_fid):
    """This function takes a real FID and converts it back into a true complex form.

    Standard optimization routines as used in lmfit require real data.

    Parameters
    ----------
    real_fid : MRSData instance
        The real FID to be converted to complex, generated from the optimization routine.

    Returns
    -------
    complex_fid : MRSData instance
        The complex version of the FID.

    """
    np = int(real_fid.shape[0] / 2)
    complex_fid = numpy.zeros(np, 'complex')
    complex_fid[:] = real_fid[:np]

    # the imaginary part of the FID has been reversed, have to flip it back
    imag_fid = real_fid[np:]
    complex_fid += 1j * imag_fid[::-1]
    return complex_fid

def fit(fid, model, baseline_points=16):
    """Fit fid with model parameters.

    Parameters
    ----------
    fid : MRSData instance
        FID to be fit
    model : dict
        Dictionary of model fit parameters
    baseline_points : int
        The number of points at the start of the FID to ignore

    Returns
    -------
    dict
        Dictionary containing ["model": optimized model, "fit": fitting data, "err": dictionary of standard errors]

    """
    # Get list of metabolite names.
    def get_metabolites(model_input):
        metabolites = []
        for fid_property_name, fid_property_value in model_input.items():
            if type(fid_property_value) is dict:
                metabolites.append(fid_property_name)
        return metabolites

    # Get standard errors from lmfit MinimizerResult object.
    def get_errors(result):
        errors = {}
        for name, param in result.params.items():
            errors[name] = param.stderr
        return errors

    def make_basis(params, time_axis):
        """This function generates a basis set.

        Parameters
        ----------
        params: lmfit Parameters object
             Contains fitting parameters
        time_axis : arange ndarray
            The time axis in ms

        Returns
        -------
        basis_matrix : matrix
            A 2D array containing the generated basis set.

        """
        basis_matrix = numpy.matrix(numpy.zeros((len(metabolite_name_list), len(time_axis) * 2)))
        for i, metabolite_name in enumerate(metabolite_name_list):
            gaussian = suspect.basis.gaussian(time_axis,
                                              params["{}_frequency".format(metabolite_name)],
                                              params["{}_phase".format(metabolite_name)].value,
                                              params["{}_fwhm".format(metabolite_name)])
            real_gaussian = complex_to_real(gaussian)
            basis_matrix[i, :] = real_gaussian
        return basis_matrix

    # unphase the FID to make it pure absorptive
    def unphase(data, params):

        unphased_data = suspect.adjust_phase(data, -params['phase0'], -params['phase1'])
        real_unphased_data = complex_to_real(unphased_data)

        return real_unphased_data

    def do_fit(params, time_axis, real_unphased_data):
        """This function performs the fitting.

        Parameters
        ----------
        params: lmfit Parameters object
            Contains fitting parameters.
        time_axis : arange ndarray
            The time axis in ms
        real_unphased_data : MRSData instance
            Unphased FID

        Returns
        -------
        list
            Contains fitted data points and amplitudes of each singlet.

        """
        basis = make_basis(params, time_axis)

        weights = scipy.optimize.nnls(basis[:, baseline_points:-baseline_points].T,
                                      real_unphased_data[baseline_points:-baseline_points])[0]

        fitted_data = numpy.array(numpy.dot(weights, basis)).squeeze()
        return fitted_data, weights

    def residual(params, time_axis, data):
        """Calculates the residual to be minimized by the least squares means method.

        Parameters
        ----------
        params : lmfit Parameters object
            Contains fitting parameters
        time_axis : array
            The time axis in ms
        data : MRSData instance
            FID to be fitted

        Returns
        -------
        array
            Residual values

        """
        real_unphased_data = unphase(data, params)
        fitted_data, weights = do_fit(params, time_axis, real_unphased_data)
        res = fitted_data - real_unphased_data

        return res[baseline_points:-baseline_points]

    def fit_data(data, initial_params):
        """
        This function takes an FID and a set of parameters contained in an lmfit Parameters object,
        and fits the data using the least squares means method.

        Parameters
        ----------
        data : MRSData instance
            FID to be fitted.
        initial_params : lmfit Parameters object
            Contains initial fitting parameters.

        Returns
        -------
        fitting_weights : list
            Tuple of weights
        fitted_data : MRSData
            Fitted data
        fitting_result : lmfit MinimizerResult instance

        """
        fitting_result = lmfit.minimize(residual,
                                        initial_params,
                                        args=(data.time_axis(), data),
                                        xtol=5e-3)

        real_fitted_data, fitting_weights = do_fit(fitting_result.params, data.time_axis(), unphase(data, fitting_result.params))
        fitted_data = data.inherit(real_to_complex(real_fitted_data))

        return fitting_weights, fitted_data, fitting_result

    # Convert lmfit parameters to model format
    def parameters_to_model(parameters_obj, param_weights):
        new_model = {}
        for param_name, param in parameters_obj.items():
            name = param_name.split("_")
            name1 = name[0]
            if len(name) == 1:  # i.e. phase
                new_model[name1] = param.value
            else:
                name2 = name[1]
                if name1 not in new_model:
                    new_model[name1] = {name2: param.value}
                else:
                    new_model[name1][name2] = param.value

        for i, metabolite_name in enumerate(metabolite_name_list):
            new_model[metabolite_name]["amplitude"] = param_weights[i]

        return new_model

    # Convert initial model to lmfit parameters.
    def model_to_parameters(model_dict):
        lmfit_parameters = lmfit.Parameters()
        params = []
        ordered_params = []
        # Calculate dependencies/references for each parameter.
        depend_dict = calculate_dependencies(model_dict)

        model_dict_copy = copy.deepcopy(model_dict)
        params.append(("phase0", model_dict_copy.pop("phase0")))
        params.append(("phase1", model_dict_copy.pop("phase1")))

        # Construct lmfit Parameter input for each parameter.
        for peak_name, peak_properties in model_dict_copy.items():
                # Fix phase value to 0 by default.
                if "phase" not in peak_properties:
                    params.append(("{0}_{1}".format(peak_name, "phase"), None, None, None, None, "0"))
                for property_name, property_value in peak_properties.items():
                    # Initialize lmfit parameter arguments.
                    name = "{0}_{1}".format(peak_name, property_name)
                    value = None
                    vary = True
                    lmfit_min = None
                    lmfit_max = None
                    expr = None
                    if isinstance(property_value, numbers.Number):
                        value = property_value
                    elif isinstance(property_value, str):
                        expr = property_value
                    elif isinstance(property_value, dict):
                        if "value" in property_value:
                            value = property_value["value"]
                        if "min" in property_value:
                            lmfit_min = property_value["min"]
                        if "max" in property_value:
                            lmfit_max = property_value["max"]
                    # Add parameter object with defined parameters.
                    params.append((name, value, vary, lmfit_min, lmfit_max, expr))  # (lmfit Parameter input format)

        # Order parameters based on dependencies.
        in_oparams = []
        while len(params) > 0:
            front = params.pop(0)
            name = front[0]
            # If no dependencies, add parameter to list and mark parameter as added.
            if name not in depend_dict or depend_dict[name] is None:
                ordered_params.append(front)
                in_oparams.append(name)
            else:
                dependencies_present = True
                for dependency in depend_dict[name]:
                    # If dependency not yet added, mark parameter to move to back of queue.
                    if dependency not in in_oparams:
                        dependencies_present = False
                # If all dependencies present, add parameter to list and mark parameter as added.
                if dependencies_present:
                    ordered_params.append(front)
                    in_oparams.append(name)
                # If dependencies missing, move parameter to back of queue.
                else:
                    params.append(front)

        # Convert all parameters to lmfit Parameter objects.
        lmfit_parameters.add_many(*ordered_params)

        return lmfit_parameters

    # Check if all model input types are correct.
    def check_errors(check_model):
        # Allowed keys in the model.
        allowed_keys = ["min", "max", "value", "phase", "amplitude"]

        # Scan model.
        for model_property, model_values in check_model.items():
            if not isinstance(model_values, (numbers.Number, dict)):
                raise TypeError("Value of {0} must be a number (for phases), or a dictionary.".format(model_property))
            elif type(model_values) is dict:  # i.e. type(value) is not int
                for peak_property, peak_value in model_values.items():
                    if not isinstance(peak_value, (numbers.Number, dict, str)):
                        raise TypeError("Value of {0}_{1} must be a value, an expression, or a dictionary."
                                        .format(model_property, peak_property))
                    if type(peak_value) is dict:
                        for width_param in peak_value:
                            # Dictionary must have 'value' key.
                            if "value" not in peak_value:
                                raise KeyError("Dictionary {0}_{1} is missing 'value' key."
                                               .format(model_property, peak_property))
                            # Dictionary can only have 'min,' 'max,' and 'value'.
                            if width_param not in allowed_keys:
                                raise KeyError("In {0}_{1}, '{2}' is not an allowed key."
                                               .format(model_property, peak_property, width_param))

    # Calculate references to determine order for Parameters.
    def calculate_dependencies(unordered_model):
        dependencies = {}  # (name, [dependencies])

        # Compile dictionary of effective names.
        for model_property, model_values in unordered_model.items():
            if type(model_values) is dict:  # i.e. pcr, not phase
                for peak_property in model_values:
                    dependencies["{0}_{1}".format(model_property, peak_property)] = None

        # Find dependencies for each effective name.
        for model_property, model_values in unordered_model.items():
            if type(model_values) is dict:  # i.e. not phase
                for peak_property, peak_value in model_values.items():
                    if type(peak_value) is str:
                        lmfit_name = "{0}_{1}".format(model_property, peak_property)
                        dependencies[lmfit_name] = []
                        for depend in dependencies:
                            if depend in peak_value:
                                dependencies[lmfit_name].append(depend)

        # Check for circular dependencies.
        for name, dependents in dependencies.items():
            if type(dependents) is list:
                for dependent in dependents:
                    if dependencies[dependent] is not None and name in dependencies[dependent]:
                        raise ReferenceError("{0} and {1} reference each other, creating a circular reference."
                                             .format(name, dependent))
        return dependencies

    # Do singlet fitting
    # Minimize and fit 31P data.

    check_errors(model)  # Check for errors in model formatting.

    metabolite_name_list = get_metabolites(model)  # Set list of metabolite names.

    parameters = model_to_parameters(model)  # Convert model to lmfit Parameters object.

    fitted_weights, fitted_data, fitted_results = fit_data(fid, parameters)  # Fit data.

    final_model = parameters_to_model(fitted_results.params, fitted_weights)  # Convert fit parameters to model format.

    stderr = get_errors(fitted_results)  # Get stderr values for each parameter.

    return_dict = {"model": final_model, "fit": fitted_data, "errors": stderr}  # Compile output into a dictionary.
    return return_dict

reason: import numpy,import scipy


text:
"""Toolkit for dataset AFW

Useage: https://github.com/yinguobing/facial-landmark-dataset/issues/4
"""

import cv2
import numpy as np

from fmd.mark_dataset.dataset import MarkDataset
from fmd.mark_dataset.util import FileListGenerator

class AFW(MarkDataset):
    # To use this class, there are two functions need to be overridden.

    def populate_dataset(self, image_dir):
        """Populate the AFW dataset with essential data.

        Args:
            image_dir: the direcotry of the dataset images.
        """
        # As required by the abstract method, we need to override this function.
        # 1. populate the image file list.
        lg = FileListGenerator()
        self.image_files = lg.generate_list(image_dir)

        # 2. Populate the mark file list. Note the order should be same with the
        # image file list. Since the AFW dataset had the mark file named after
        # the image file but with different extention name `pts`. We will make
        # use of this.
        self.mark_files = [img_path.split(
            ".")[-2] + ".pts" for img_path in self.image_files]

        # 3 Set the key marks indices. Here key marks are: left eye left corner,
        #  left eye right corner, right eye left corner, right eye right corner,
        #  mouse left corner, mouse right corner. For AFW the indices are 36,
        # 39, 42, 45, 48, 54. Most of the time you need to do this manually.
        # Refer to the mark dataset for details.
        self.key_marks_indices = [36, 39, 42, 45, 48, 54]

        # Even optional, it is highly recommended to update the meta data.
        self.meta.update({"authors": "Imperial College London",
                          "year": 2013,
                          "num_marks": 68,
                          "num_samples": len(self.image_files)
                          })

    def get_marks_from_file(self, mark_file):
        """This function should read the mark file and return the marks as a
        numpy array in form of [[x, y, z], [x, y, z]]."""
        marks = []
        with open(mark_file) as fid:
            for line in fid:
                if "version" in line or "points" in line or "{" in line or "}" in line:
                    continue
                else:
                    loc_x, loc_y = line.strip().split(sep=" ")
                    marks.append([float(loc_x), float(loc_y), 0.0])
        marks = np.array(marks, dtype=np.float)
        assert marks.shape[1] == 3, "Marks should be 3D, check z axis values."
        return marks

max_stars_repo_name: yinguobing/facial-landmark-dataset

max_issues_repo_name: UBCDingXin/facial-landmark-dataset

max_forks_repo_name: UBCDingXin/facial-landmark-dataset


text:
import numpy as np
import torch

from allrank.data.dataset_loading import PADDED_Y_VALUE
from allrank.models.model_utils import get_torch_device

def ndcg(y_pred, y_true, ats=None, gain_function=lambda x: torch.pow(2, x) - 1):
    # y_true and y_pred have dimensions [listing, document_score]
    # returns a tensor with ndcg values at positions specified by 'ats' with dimensions [listing, dcg_at]

    idcg = dcg(y_true, y_true, ats, gain_function)
    ndcg_ = dcg(y_pred, y_true, ats, gain_function) / idcg
    idcg_mask = idcg == 0
    ndcg_[idcg_mask] = 0.  # if idcg == 0 , set ndcg to 0

    assert (ndcg_ < 0.0).sum() >= 0, "every ndcg should be non-negative"

    return ndcg_

def __apply_mask_and_get_true_sorted_by_preds(y_pred, y_true):
    mask = y_true == PADDED_Y_VALUE

    y_pred[mask] = float('-inf')
    y_true[mask] = 0.0

    _, indices = y_pred.sort(descending=True, dim=-1)
    return torch.gather(y_true, dim=1, index=indices)

def dcg(y_pred, y_true, ats=None, gain_function=lambda x: torch.pow(2, x) - 1):
    # y_true and y_pred have dimensions [listing, document_score]
    # returns a tensor with ndcg values at positions specified by 'ats' with dimensions [listing, dcg_at]
    y_true = y_true.clone()
    y_pred = y_pred.clone()

    actual_length = y_true.shape[1]

    if ats is None:
        ats = [actual_length]
    ats = [min(at, actual_length) for at in ats]

    true_sorted_by_preds = __apply_mask_and_get_true_sorted_by_preds(y_pred, y_true)

    dev = get_torch_device()

    discounts = (torch.tensor(1) / torch.log2(torch.arange(true_sorted_by_preds.shape[1], dtype=torch.float) + 2.0)).to(
        device=dev)

    gains = gain_function(true_sorted_by_preds)

    discounted_gains = (gains * discounts)[:, :np.max(ats)]

    cum_dcg = torch.cumsum(discounted_gains, dim=1)

    ats_tensor = torch.tensor(ats, dtype=torch.long) - torch.tensor(1)

    dcg = cum_dcg[:, ats_tensor]

    return dcg

def mrr(y_pred, y_true, ats=None):
    # y_true and y_pred have dimensions [listing, document_score]
    # returns a tensor with mrr values at positions specified by 'ats' with dimensions [listing, mrr_at]
    y_true = y_true.clone()
    y_pred = y_pred.clone()

    if ats is None:
        ats = [y_true.shape[1]]

    true_sorted_by_preds = __apply_mask_and_get_true_sorted_by_preds(y_pred, y_true)

    values, indices = torch.max(true_sorted_by_preds, dim=1)
    indices = indices.type_as(values).unsqueeze(dim=0).t().expand(len(y_true), len(ats))

    dev = get_torch_device()

    ats_rep = torch.tensor(data=ats, device=dev, dtype=torch.float32).expand(len(y_true), len(ats))

    within_at_mask = (indices < ats_rep).type(torch.float32)

    result = torch.tensor(1.0) / (indices + torch.tensor(1.0))

    zero_sum_mask = torch.sum(values) == 0.0
    result[zero_sum_mask] = 0.0

    result = result * within_at_mask

    return result


text:
from polyhedral_analysis.coordination_polyhedron import CoordinationPolyhedron
from polyhedral_analysis.atom import Atom
from typing import List, Tuple
import numpy as np # type: ignore

# Functions for analysing octahedra

VertexPairs = Tuple[Tuple[Atom, Atom],
                    Tuple[Atom, Atom],
                    Tuple[Atom, Atom]]

def check_octahedra(polyhedron: CoordinationPolyhedron) -> None:
    """
    Check whether a polyhedron is considered an octahedron.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The :obj:`CoordinationPolyhedron`
            to be tested.

    Returns:
        None

    Raises:
        ValueError: If the "best fit" geometry (lowest continuous symmetry
            measure) is not an octahedron, a ValueError is raised.

    """
    if not polyhedron.best_fit_geometry['geometry'] is 'Octahedron':
        raise ValueError( 'This polyhedron is not recognised a an octahedron' )

def opposite_vertex_pairs(polyhedron: CoordinationPolyhedron,
                          check: bool = True) -> VertexPairs:
    """For an octahedral polyhedron, find the pairs of vertices opposite each other.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this
            polyhedron is an octahedron. Default is `True`.

    Returns:
        (tuple): 3 pairs of vertex atoms.

    """
    if check:
        check_octahedra(polyhedron)
    assert {len(n) for n in polyhedron.edge_graph.values()} == {4}, "Edge graph does not describe an octahedron."
    vertex_pairs = []
    seen_indices: List[int] = []
    for v1 in polyhedron.vertices:
        if v1.index in seen_indices:
            continue
        v1_neighbours = polyhedron.vertices_by_indices(v1.neighbours[polyhedron.index])
        v2 = next(v for v in polyhedron.vertices if v not in [v1, *v1_neighbours])
        vertex_pairs.append((v1, v2))
        seen_indices.extend([v1.index, v2.index])
    return (vertex_pairs[0], vertex_pairs[1], vertex_pairs[2])

def opposite_vertex_distances(polyhedron: CoordinationPolyhedron,
                              check: bool = True) -> Tuple[float, float, float]:
    """For an octahedral polyhedron, return the distances between pairs of cis- vertices.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this polyhedron
            is an octahedron. Default is `True`.

    Returns:
        (tuple): a length-3 tuple of distances.

    """
    distances = tuple(vp[0].distance(vp[1])
                      for vp in opposite_vertex_pairs(polyhedron, check=check))
    return (distances[0], distances[1], distances[2])

def isomer_is_trans(polyhedron: CoordinationPolyhedron,
             check: bool = True) -> bool:
    """For an octahedral polyhedron with 2+4 coordination, return True for trans coordination,
    and False for cis coordination.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this polyhedron
            is an octahedron. Default is `True`.

    Return:
        (bool): `True` for trans, `False` for cis.

    """
    coord = list(polyhedron.vertex_count.values())
    if sorted(coord) != [2, 4]:
        raise ValueError(f'cis/trans not defined for {polyhedron.vertex_count.values()} coordination.')
    to_return = bool(np.all([vp[0].label == vp[1].label
                     for vp in opposite_vertex_pairs(polyhedron, check=check)]))
    return to_return

def isomer_is_cis(polyhedron: CoordinationPolyhedron,
                  check: bool = True) -> bool:
    """For an octahedral polyhedron with 2+4 coordination, return True for cis coordination,
    and False for trans coordination.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this polyhedron
            is an octahedron. Default is `True`.

    Return:
        (bool): `True` for cis, `False` for trans.

    """
    return not isomer_is_trans(polyhedron, check=check)

def isomer_is_fac(polyhedron: CoordinationPolyhedron,
                  check: bool = True) -> bool:
    """For an octahedral polyhedron with 3+3 coordination, return True for fac coordination,
    and False for mer coordination.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this polyhedron
            is an octahedron. Default is `True`.

    Return:
        (bool): `True` for fac, `False` for mer.

    """
    coord = list(polyhedron.vertex_count.values())
    if sorted(coord) != [3, 3]:
        raise ValueError(f'fac/mer not defined for {polyhedron.vertex_count.values()} coordination.')
    to_return = bool(np.all([vp[0].label != vp[1].label
                     for vp in opposite_vertex_pairs(polyhedron, check=check)]))
    return to_return

def isomer_is_mer(polyhedron: CoordinationPolyhedron,
                  check: bool = True) -> bool:
    """For an octahedral polyhedron with 3+3 coordination, return True for mer coordination,
    and False for fac coordination.

    Args:
        polyhedron (:obj:`CoordinationPolyhedron`): The polyhedron to be analysed.
        check: (Optional, :obj:`bool`): Optional flag to set whether to check that this polyhedron
            is an octahedron. Default is `True`.

    Return:
        (bool): `True` for fac, `False` for mer.

    """
    return not isomer_is_fac(polyhedron, check=check)

max_stars_repo_name: bjmorgan/polyhedral-analysis

max_issues_repo_name: bjmorgan/polyhedral-analysis

max_forks_repo_name: bjmorgan/polyhedral-analysis