text:
from turtle import forward
from black import out
from more_itertools import last
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
import torch
import torch.nn as nn
import dgl
from dgl.nn.pytorch import RelGraphConv, GatedGraphConv
from torch.nn.utils.rnn import pad_sequence
from torch.nn.utils import spectral_norm
import dgl.data
import torch
import torch.nn as nn
import torch.nn.functional as F

device = None

def dropout_inference(m: nn.Dropout):
    if isinstance(m, nn.Dropout):
        m.eval()

def init_weights_ortho(m):
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight.data)

class Linear_3(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super(Linear_3, self).__init__()
        self.Dense1 = nn.Linear(in_dim, hidden_dim)
        self.Dense2 = nn.Linear(hidden_dim, hidden_dim)
        self.Dense3 = nn.Linear(hidden_dim, out_dim)

    def forward(self, inputs):
        out = torch.tanh(self.Dense1(inputs))
        out = torch.tanh(self.Dense2(out))
        out = self.Dense3(out)
        return out

class LinearBlock(nn.Module):
    def __init__(
        self,
        in_dim,
        hidden_dim,
        out_dim,
        num_layers,
        drop_out,
        activation,
        normalization=None,
    ):
        super(LinearBlock, self).__init__()
        self.block = nn.ModuleList([])
        for i, _ in enumerate(range(num_layers)):
            block = []
            # transformation
            if i == 0:
                block.append(nn.Linear(in_dim, hidden_dim))

            else:
                if normalization == "SpectralNorm":
                    block.append(spectral_norm(nn.Linear(hidden_dim, hidden_dim)))
                    print("231")
                else:
                    block.append(nn.Linear(hidden_dim, hidden_dim))

            # normalization
            if normalization == "LayerNorm":
                block.append(nn.LayerNorm(hidden_dim))

            # activations
            if activation == "ReLU":
                block.append(nn.ReLU())
            elif activation == "ELU":
                block.append(nn.ELU())
            elif activation == "LeakyReLU":
                block.append(nn.LeakyReLU())

            # dropout
            block.append(nn.Dropout(drop_out))
            self.block.extend(block)

        self.block.append(nn.Linear(hidden_dim, out_dim))

    def forward(self, x):
        for layer in self.block:
            x = layer(x)

        return x

class CriticSqueeze(nn.Module):

    """
    Critic class for advantage estimation
    """

    def __init__(self, in_dim, hidden_dim):
        super(CriticSqueeze, self).__init__()
        self.GGN = dgl.nn.GatedGraphConv(in_dim, hidden_dim, 5, 4)
        self.Dense = Linear_3(hidden_dim + 1 + in_dim, hidden_dim, 1)

    def forward(self, graph, last_action_node, last_node):
        # print(graph.adj())
        h = F.relu(
            self.GGN(graph, graph.ndata["atomic"], graph.edata["type"].squeeze())
        )
        with graph.local_scope():
            graph.ndata["h"] = h
            hg = dgl.mean_nodes(graph, "h")
        cat = torch.cat((hg, last_action_node, last_node), dim=1)
        out = self.Dense(cat)
        return out

class EdgePredictor(nn.Module):
    def __init__(
        self, in_dim: int, hidden_dim: int, num_layers=3, dropout=0.0, activation="ReLU"
    ):
        """torch module for edge prediction. Paramaterizes edge probability by concating
        final node feat with each other node and running that through a mlp

        Args:
            in_dim (int): input dimension
            hidden_dim (int): hidden dimension
            num_layers (int, optional): number of layers. Defaults to 3.
            dropout (float, optional): dropout level. Defaults to 0.0.
            activation (str, optional): what activation to use. Defaults to 'ReLU'.
        """
        super(EdgePredictor, self).__init__()
        self.NodeEmbed = LinearBlock(
            in_dim, hidden_dim, hidden_dim, 1, dropout, activation
        )
        self.Edges = LinearBlock(
            hidden_dim * 2, hidden_dim, 2, num_layers, dropout, activation
        )

    def forward(self, graphs, last_node_batch, h):
        """
        graph is actually a batch of graphs where len(graph.batch_num_nodes()) = len(last_node_batch)
        returns just a list of edges so its unsplit up rn, bbut
        """
        edges_per_graph = []

        node_embed_batch = self.NodeEmbed(last_node_batch)
        batch_node_stacks = []

        """looping over graphs"""
        for i in range(last_node_batch.shape[0]):
            batch_node_stacks.append(
                node_embed_batch[i].repeat(graphs.batch_num_nodes()[i], 1)
            )

        batch_node_stacks = torch.cat(batch_node_stacks, dim=0)

        stack = torch.cat((h, batch_node_stacks), dim=1)
        edges = self.Edges(stack)
        with graphs.local_scope():
            graphs.ndata["bond_pred"] = edges
            graphs = dgl.unbatch(graphs)
            for graph in graphs:
                edges_per_graph.append(graph.ndata["bond_pred"])

        return pad_sequence(
            edges_per_graph, batch_first=True, padding_value=-10000
        ).flatten(1, -1)

class Batch_Edge(nn.Module):
    def __init__(self, in_dim, hidden_dim):
        super(Batch_Edge, self).__init__()
        # need more theres no non linearity
        self.NodeEmbed = nn.Linear(in_dim, hidden_dim)
        self.Edges = Linear_3(hidden_dim * 2, hidden_dim * 2, 2)

    def forward(self, graphs, last_node_batch, h):
        """
        graph is actually a batch of graphs where len(graph.batch_num_nodes()) = len(last_node_batch)
        returns just a list of edges so its unsplit up rn, bbut
        """
        edges_per_graph = []

        node_embed_batch = self.NodeEmbed(last_node_batch)
        # holds number of graphs of each node embedding stacked to match node number per graph
        batch_node_stacks = []
        # so its first dimension is the batch size, and the second 'dimension' is the number of nodes per grah

        """looping over graphs"""
        for i in range(last_node_batch.shape[0]):
            batch_node_stacks.append(
                node_embed_batch[i].repeat(graphs.batch_num_nodes()[i], 1)
            )

        batch_node_stacks = torch.cat(batch_node_stacks, dim=0)

        stack = torch.cat((h, batch_node_stacks), dim=1)
        edges = self.Edges(stack)
        with graphs.local_scope():
            graphs.ndata["bond_pred"] = edges
            graphs = dgl.unbatch(graphs)
            for graph in graphs:
                edges_per_graph.append(graph.ndata["bond_pred"])

        return pad_sequence(
            edges_per_graph, batch_first=True, padding_value=-10000
        ).flatten(1, -1)

class BaseLine(nn.Module):
    """Crappy base line to check improvements against"""

    def __init__(self, in_dim, hidden_dim, num_nodes):
        super(BaseLine, self).__init__()
        self.in_dim = in_dim
        self.num_nodes = num_nodes
        self.GGN = dgl.nn.GatedGraphConv(in_dim, hidden_dim, 5, 4)
        self.AddNode = Linear_3(hidden_dim, hidden_dim, num_nodes)
        self.BatchEdge = Batch_Edge(in_dim, hidden_dim)

    def forward(self, graph, last_action_node, last_node, mask=False, softmax=True):
        # out = []
        h = F.relu(
            self.GGN(graph, graph.ndata["atomic"], graph.edata["type"].squeeze())
        )
        with graph.local_scope():
            graph.ndata["h"] = h
            hg = dgl.mean_nodes(graph, "h")

        addNode = self.AddNode(hg)

        if mask:
            node_mask = torch.unsqueeze(
                torch.cat((torch.zeros(1), torch.ones(self.num_nodes - 1))), dim=0
            ).to(device)
            mask = last_action_node * (node_mask * -100000)
            addNode += mask

        edges = self.BatchEdge(graph, last_node, h)
        out = torch.cat((addNode, edges), dim=1)
        if softmax:
            return torch.softmax(out, dim=1)
        else:
            return out

class Actor(nn.Module):
    def __init__(
        self,
        in_dim,
        hidden_dim,
        out_dim,
        drop_out,
        graph_activation,
        dense_activation,
        model,
        graph_num_layers,
        dense_num_layers,
        norm=None,
        regularizer=None,
        ortho_init=True,
    ):
        super(Actor, self).__init__()
        assert graph_num_layers > 1

        """GRAPH LAYERS"""
        if model == "RelationalGraphConv":
            init_embed = RelGraphConv(
                in_dim,
                hidden_dim,
                3,
                regularizer,
                activation=graph_activation,
                dropout=drop_out,
            )
            self.GraphConv = nn.ModuleList(
                [init_embed]
                + [
                    RelGraphConv(
                        hidden_dim,
                        hidden_dim,
                        3,
                        regularizer,
                        activation=graph_activation,
                        dropout=drop_out,
                    )
                    for _ in range(graph_num_layers)
                ]
            )

        elif model == "GatedGraphConv":
            self.GraphConv = nn.ModuleList(
                [GatedGraphConv(in_dim, hidden_dim, graph_num_layers, 4)]
            )

        else:
            print("Bad Graph Name")
            raise

        """DENSE LAYER"""
        self.dense = LinearBlock(
            hidden_dim + 1,
            hidden_dim * 2,
            out_dim,
            dense_num_layers,
            0,
            dense_activation,
            normalization=norm,
        )

        """EDGE LAYER"""
        self.BatchEdge = EdgePredictor(
            in_dim, hidden_dim, dropout=drop_out, activation=dense_activation
        )

        if ortho_init:
            self.apply(init_weights_ortho)

        self.out_dim = out_dim

    def forward(self, graph, last_action_node, last_node, mask=False, softmax=True):
        node_feats = graph.ndata["atomic"]

        for _, layer in enumerate(self.GraphConv):
            node_feats = layer(graph, node_feats, graph.edata["type"].squeeze())
        with graph.local_scope():
            graph.ndata["node_feats"] = node_feats
            hg = dgl.sum_nodes(graph, "node_feats")

        hg = torch.cat((hg, last_action_node), dim=1)
        addNode = self.dense(hg)
        if mask:
            node_mask = torch.unsqueeze(
                torch.cat((torch.zeros(1), torch.ones(self.out_dim - 1))), dim=0
            ).to(device)
            node_mask = node_mask.cuda()
            print(last_action_node.device,node_mask.device)
            mask = last_action_node * (node_mask * -100000)
            print(mask)
            addNode += mask
            print(addNode,'addnode')

        edges = self.BatchEdge(graph, last_node, node_feats)
        out = torch.cat((addNode, edges), dim=1)
        if softmax:
            return torch.softmax(out, dim=1)
        else:
            return out

class Critic(nn.Module):
    def __init__(
        self,
        in_dim,
        hidden_dim,
        model,
        graph_num_layers,
        dense_num_layers,
        graph_activation,
        dense_activation,
        dropout,
        norm=None,
        regularizer=None,
        ortho_init=True,
    ):
        super(Critic, self).__init__()

        """GRAPH LAYERS"""
        if model == "RelationalGraphConv":
            init_embed = RelGraphConv(
                in_dim,
                hidden_dim,
                3,
                regularizer,
                activation=graph_activation,
                dropout=dropout,
            )
            self.GraphConv = nn.ModuleList(
                [init_embed]
                + [
                    RelGraphConv(
                        hidden_dim,
                        hidden_dim,
                        3,
                        regularizer,
                        activation=graph_activation,
                        dropout=dropout,
                    )
                    for _ in range(graph_num_layers)
                ]
            )

        elif model == "GatedGraphConv":
            self.GraphConv = nn.ModuleList(
                [GatedGraphConv(in_dim, hidden_dim, graph_num_layers, 4)]
            )

        else:
            print("Bad Graph Name")
            raise

        """DENSE LAYER"""
        self.dense = LinearBlock(
            hidden_dim + 1 + in_dim,
            hidden_dim * 2,
            1,
            dense_num_layers,
            0,
            dense_activation,
            normalization=norm,
        )

        # '''EDGE LAYER'''
        # self.BatchEdge = EdgePredictor(
        #     in_dim, hidden_dim, dropout=drop_out, activation=dense_activation)

        if ortho_init:
            self.apply(init_weights_ortho)

    def forward(self, graph, last_action_node, last_node, masking=False, softmax=True):
        node_feats = graph.ndata["atomic"]

        for _, layer in enumerate(self.GraphConv):
            node_feats = layer(graph, node_feats, graph.edata["type"].squeeze())
        with graph.local_scope():
            graph.ndata["node_feats"] = node_feats
            hg = dgl.sum_nodes(graph, "node_feats")

        hg = torch.cat((hg, last_action_node, last_node), dim=1)
        out = self.dense(hg)
        return out

max_stars_repo_name: tsteternlieb/DrugDesignThesis

max_issues_repo_name: tsteternlieb/DrugDesignThesis

max_forks_repo_name: tsteternlieb/DrugDesignThesis


text:
"""
    Test functions in the space of solutions of the
    Euler Lagrange equations of

    \int_{-1}^{1} (2/tau) \alpha dq/ds + (2/tau)^5 (1-\alpha) d^3 q / ds^3 dt
"""
import unittest
from scipy.sparse import csc_matrix
import numpy as np
from gsplines.interpolator.gspline import cSplineCalc
from gsplines.basis.basis1010 import cBasis1010
from gsplines.basis.basis0010 import cBasis0010

class cMyTest(unittest.TestCase):
    def __init__(self, *args, **kwargs):
        super(cMyTest, self).__init__(*args, **kwargs)
        import sys
        np.set_printoptions(
            linewidth=5000000,
            formatter={'float': '{:+10.3e}'.format},
            threshold=sys.maxsize)
        pass

    def testInversion(self):
        import time
        print('Testinf inversion of matrix')
        dim = 6  # np.random.randint(2, 6)
        N = 50  # np.random.randint(3, 120)
        a = np.random.rand()
        splcalc = cSplineCalc(dim, N, cBasis1010(a))
        for i in range(50):
            tauv = np.random.rand(N)
            A1 = splcalc.eval_A(tauv)
#            A0 = self.eval_A(tauv, dim, N, cBasis1010(a))
#            e = np.max(np.abs(A1 - A0.todense()))
#            #            print(A0)
#            #            print('----------------------------------------')
#            #            print(A1)
#            #            print(dim, N)
#            assert e < 1.0e-8

        splcalc.printPerformace()
        pass

    def testcontinuity(self):
        print('Test continuity constraints with plot')
        for i in range(3):
            dim = np.random.randint(2, 3)
            N = np.random.randint(3, 10)
            a = np.random.rand()
            wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
            tauv = 0.5 + np.random.rand(N) * 3.0
            tis = [np.sum(tauv[0:i]) for i in range(0, N + 1)]
            T = np.sum(tauv)
            splcalc = cSplineCalc(dim, N, cBasis1010(a))
            spln = splcalc.getSpline(tauv, wp)
            from matplotlib import pyplot as plt

            t = np.arange(0, T, 0.005)
            q_list = [spln.deriv(i)(t) for i in range(0, 6)]

            fig, axs = plt.subplots(6, dim)

            for i in range(0, 6):
                for j in range(0, dim):
                    axs[i, j].plot(t, q_list[i][:, j])
                    axs[i, j].grid()
                    for ti in tis:
                        axs[i, j].axvline(x=ti, color='b', linestyle='--')

            plt.show()

    def test_eval_b(self):
        import time
        print('Test evaluation of b vector')
        for i in range(20):
            dim = np.random.randint(1, 8)
            N = np.random.randint(3, 200)
            a = np.random.rand()
            wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
            dwp0 = np.zeros((dim, ))
            ddwp0 = np.zeros((dim, ))
            dwpT = np.zeros((dim, ))
            ddwpT = np.zeros((dim, ))
            splcalc = cSplineCalc(dim, N, cBasis1010(a))
            b1 = splcalc.eval_b(wp)
            b2 = self.eval_b(wp, N, dim, dwp0, ddwp0, dwpT, ddwpT)

            e = np.max(np.abs(b1 - b2))

            assert e < 1.0e-8

    def non_zero_diagonals_A(self):
        for i in range(0, 1):
            dim = np.random.randint(1, 8)
            N = np.random.randint(3, 200)
            a = np.random.rand()
            splcalc = cSplineCalc(dim, N, cBasis1010, a)

            tauv = 10.0 * np.random.rand(N) + 0.5
            splcalc.eval_A(tauv)
            A = splcalc.Aeval

            upper_diags = 0
            flag = 0
            for i in range(A.shape[0]):
                if np.max(np.abs(np.diagonal(A, i))) > 1.0e-10:
                    assert flag != 1, 'Matrix is Not Banded!!!'
                    upper_diags += 1
                else:
                    flag = 1

            lower_diags = 0
            flag = 0
            for i in range(A.shape[0]):
                if np.max(np.abs(np.diagonal(A, -i))) > 1.0e-10:
                    assert flag != 1, 'Matrix is Not Banded!!!'
                    lower_diags += 1
                else:
                    flag = 1

#            print('upper diagonas = {:d}  lower diagonals = {:d}'.format(
#                upper_diags, lower_diags))

            assert 4 * dim + 4 == max(upper_diags, lower_diags)

#        wp = (np.random.rand(N + 1, dim) - 0.5) * 2.0 * np.pi
#        b = splcalc.eval_b(wp)
#
#        wp = np.random.rand(N + 1, dim)
#
#        spline = splcalc.solve(wp, tauv)
#
#        tis = [np.sum(tauv[0:i]) for i in range(0, N + 1)]
#
#        t = np.arange(0, tis[-1], 0.1)
#
#        plt.plot(t, spline(t)[:, 0])
#
#        plt.show()

    def test_derivative_b(self):
        '''
            Here we rest the correctness of the numerical output of the basis
            class comparing it with its analitical form optained using sympy
        '''
        print('Test derivative of b w.r.t. waypoint components')
        np.random.seed()
        dim = np.random.randint(1, 8)
        N = np.random.randint(2, 60)
        a = np.random.rand()
        wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
        splcalc = cSplineCalc(dim, N, cBasis1010(a))

        dwp = 0.0005
        for i in range(N + 1):
            for j in range(dim):
                wpidx = i
                i = j

                wp_aux = wp.copy()
                wp_aux[wpidx, i] += -dwp
                b1 = splcalc.eval_b(wp_aux).copy()
                wp_aux[wpidx, i] += 2 * dwp
                b2 = splcalc.eval_b(wp_aux).copy()

                dbdwpij_num = 0.5 * (b2 - b1) / dwp

                dbdwpij = splcalc.eval_dbdwpij(wpidx, i)

                e = np.max(np.abs(dbdwpij_num - dbdwpij))

                if e > 1.0e-8:
                    print('Erroe in db_dwpij:')
                    print('implementation:')
                    print(dbdwpij)
                    print('(b1-b2)/dwp:')
                    print(dbdwpij_num)
                    print('component', i)
                    print('waypoint ', wpidx)
                    print('dimension ', dim)
                    print('number of intervals ', N)
                    raise AssertionError('Error of {:14.7e}'.format(e))

    def test_derivative_y(self):
        ''' Compare the numerical derivate of y w.r.t tau with the nominal one
        '''
        for i in range(40):
            np.random.seed()
            dim = np.random.randint(1, 3)
            N = np.random.randint(2, 6)
            a = np.random.rand()
            wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
            tauv = 1.0 + np.random.rand(N) * 2.0
            splcalc = cSplineCalc(dim, N, cBasis1010(a))
            dydtauNom, y = splcalc.eval_dydtau(tauv, wp)

            y = y.copy()

            dtau = 1.0e-8
            err = 0.0
            errp = 0.0
            for iinter in range(0, N):
                tauv_aux = tauv.copy()
                tauv_aux[iinter] += -2 * dtau
                y0 = splcalc.eval_y(tauv_aux, wp).copy() * (1.0 / 12.0)
                tauv_aux[iinter] += dtau
                y1 = splcalc.eval_y(tauv_aux, wp).copy() * (-2.0 / 3.0)
                tauv_aux[iinter] += 2 * dtau
                y2 = splcalc.eval_y(tauv_aux, wp).copy() * (2.0 / 3.0)
                tauv_aux[iinter] += dtau
                y3 = splcalc.eval_y(tauv_aux, wp).copy() * (-1.0 / 12.0)
                dydtauiTest = (y0 + y1 + y2 + y3) / dtau
                ev = np.abs(dydtauiTest - dydtauNom[:, iinter])
                e = np.max(ev)
                eidx = np.argmax(ev)

                ep = e / abs(dydtauiTest[eidx])

                if e > err:
                    err = e
                if ep > errp:
                    errp = ep

                assert ep < 5.0e-2, '''
                error on dydtau = {:10.7e}
                value of dydtau = {:10.7e}
                relative error  = {:10.7e}
                '''.format(e, dydtauiTest[eidx], ep)

    def test_derivative_wp(self):
        ''' Compare the numerical derivate of y w.r.t waypoints with the nominal one
        '''
        for _ in range(4):
            np.random.seed()
            dim = np.random.randint(1, 8)
            N = np.random.randint(2, 20)
            a = np.random.rand()
            wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
            tauv = 0.5 + np.random.rand(N) * 2.0
            splcalc = cSplineCalc(dim, N, cBasis1010(a))
            _, y = splcalc.eval_dydtau(tauv, wp)
            y = y.copy()

            err = 0.0
            errp = 0.0

            err = 0.0
            errp = 0.0
            dwp = 1.0e-5

            wpidx = [(i, j) for i in range(N + 1) for j in range(dim)]
            dydwpNom = np.zeros((y.shape[0], len(wpidx)))
            dydwpNom, _ = splcalc.eval_dydu(tauv, wp, wpidx, dydwpNom)

            for k, (i, j) in enumerate(wpidx):
                wp_aux = wp.copy()
                wpidx = i
                wpcom = j

                wp_aux[wpidx, wpcom] += -3 * dwp
                y0 = splcalc.eval_y(tauv, wp_aux).copy() * (-1.0 / 60.0)
                wp_aux[wpidx, wpcom] += dwp
                y1 = splcalc.eval_y(tauv, wp_aux).copy() * (3.0 / 20.0)
                wp_aux[wpidx, wpcom] += dwp
                y2 = splcalc.eval_y(tauv, wp_aux).copy() * (-3.0 / 4.0)
                wp_aux[wpidx, wpcom] += 2 * dwp
                y3 = splcalc.eval_y(tauv, wp_aux).copy() * (3.0 / 4.0)
                wp_aux[wpidx, wpcom] += dwp
                y4 = splcalc.eval_y(tauv, wp_aux).copy() * (-3.0 / 20.0)
                wp_aux[wpidx, wpcom] += dwp
                y5 = splcalc.eval_y(tauv, wp_aux).copy() * (1.0 / 60.0)

                dydwpTest = (y0 + y1 + y2 + y3 + y4 + y5) / dwp

                ev = np.abs(dydwpNom[:, k] - dydwpTest)
                e = np.max(ev)
                eidx = np.argmax(ev)
#                print('{:14.7e} {:14.7e} {:14.7e}'.format(
#                    e, dydwpNom[eidx, k], dydwpTest[eidx]))

                ep = e / dydwpTest[eidx]

                if e > err:
                    err = e
                if ep > errp:
                    errp = ep

                if e > 1.0e-4:
                    assert ep < 1.0e-8, '''
                    Relative Error   = {:10.3e}
                    Absolute Error   = {:10.3e}
                    '''.format(ep, e)

#            print('Maximum Error for dy dwp           = {:14.7e}'.format(err))
#            print('Maximum Relative Error for dy dwp  = {:14.7e}'.format(errp))

#                assert e < 5.0e-2, 'error = {:14.7e}'.format(e)

    def test_derivative_y_2(self):
        ''' Second test for the derivative of y wr.t. tau.
        First test the identity A*dydtau + dAdtau y = 0A
        Second test the identity of above using basis
        '''
        for i in range(40):
            np.random.seed()
            dim = np.random.randint(1, 3)
            N = np.random.randint(2, 6)
            a = np.random.rand()
            wp = (np.random.rand(N + 1, dim) - 0.5) * 2 * np.pi
            tauv = 0.5 + np.random.rand(N) * 2.0
            splcalc = cSplineCalc(dim, N, cBasis0010())
            basis = splcalc.basis_
            dydtauNom, y = splcalc.eval_dydtau(tauv, wp)
            A = splcalc.eval_A(tauv).todense()
            y = y.copy()

            # A*dydtau + dAdtau y = 0
            for i, _ in enumerate(tauv):
                v0 = A.dot(dydtauNom[:, i])
                v1 = splcalc.eval_dAdtiy(tauv, i, y).todense()
                res = v0 + v1.T
                e = np.abs(res)

                assert np.max(e) < 1.0e-10, '''
                e = {:14.7e}
                '''.format(e)

            for i, taui in enumerate(tauv):

                B = basis.evalDerivOnWindow(-1, taui, 0)
                dB_dtau = basis.evalDerivWrtTauOnWindow(-1, taui, 0)

                for idim in range(dim):
                    i0 = i * 6 * dim + 6 * idim
                    i1 = i0 + 6
                    e = dydtauNom[i0:i1, i].dot(B) + y[i0:i1].dot(dB_dtau)

                    assert np.abs(e) < 1.0e-10, '''
                    error computing dydtau^\\top B + y^\\top dB_dtau
                    e            = {:14.7e}
                    index of tau = {:d}
                    index of q   = {:d}
                    i0           = {:d}
                    i1           = {:d}
                    '''.format(e, i, idim, i0, i1)

    def eval_A(self, tau, _dim, _N, _basis):
        """
        Alternative way to fill the Matrix A
        WARNING:  This work ok for _N<120
        """

        A = np.zeros(2 * (6 * _dim * _N, ))

        dim = _dim
        nzv = np.zeros((((_N - 1) * dim + (3) * dim) * 4 * 3 + (3 - 1) *
                        (_N - 1) * dim * 8 * 3, ))
        idxv = np.zeros(
            (((_N - 1) * dim + (3) * dim) * 4 * 3 + (3 - 1) *
             (_N - 1) * dim * 8 * 3, ),
            dtype=np.int16)
        ptrv = np.zeros((2 * 3 * _N * dim + 1, ), dtype=np.int16)

        basis = _basis

        Pl = [basis.evalDerivOnWindow(-1.0, tau[0], i) for i in range(0, 5)]
        Pr = [basis.evalDerivOnWindow(1.0, tau[0], i) for i in range(0, 5)]
        # Fill the content for the derivatives at boundaries
        Cpl = -np.vstack(Pl[1:5])
        Cpr = np.vstack(Pr[1:5])

        App1 = np.vstack([Pl[:3], Pr[0]])

        nnz = 0
        nptr = 0
        # --------------------------------------------
        for j in range(0, 6 * dim):
            ptrv[nptr] = nnz
            i0 = (j // 6) * 4
            for i in range(i0, i0 + 4):
                A[i, j] = App1[(i - i0), j % 6]
                nzv[nnz] = App1[(i - i0), j % 6]
                idxv[nnz] = i
                nnz += 1
            i0 += 4 * dim
            for i in range(i0, i0 + 4):
                A[i, j] = Cpr[(i - i0), j % 6]
                nzv[nnz] = Cpr[(i - i0), j % 6]
                idxv[nnz] = i
                nnz += 1
            nptr += 1
        # --------------------------------------------
        for iinter in range(1, _N - 1):
            j0 = iinter * 6 * dim
            i0 = 4 * dim + 6 * dim * (iinter - 1)

            Pl = [
                basis.evalDerivOnWindow(-1.0, tau[iinter], i)
                for i in range(0, 5)
            ]
            Pr = [
                basis.evalDerivOnWindow(1.0, tau[iinter], i)
                for i in range(0, 5)
            ]
            Cpl = -np.vstack(Pl[1:5])
            Cpr = np.vstack(Pr[1:5])
            Appi = np.vstack([Pl[0], Pr[0]])
            # --------------------------------------------
            for j in range(j0, j0 + 6 * dim):
                ptrv[nptr] = nnz
                i1 = i0 + ((j - j0) // 6) * 4
                for i in range(i1, i1 + 4):
                    A[i, j] = Cpl[(i - i1), j % 6]
                    nzv[nnz] = Cpl[(i - i1), j % 6]
                    idxv[nnz] = i
                    nnz += 1
                i1 += (dim - (j - j0) // 6) * 4 + ((j - j0) // 6) * 2
                for i in range(i1, i1 + 2):
                    A[i, j] = Appi[i - i1, j % 6]
                    nzv[nnz] = Appi[i - i1, j % 6]
                    idxv[nnz] = i
                    nnz += 1
                i1 += (dim - (j - j0) // 6) * 2 + ((j - j0) // 6) * 4
                for i in range(i1, i1 + 4):
                    A[i, j] = Cpr[(i - i1), j % 6]
                    nzv[nnz] = Cpr[(i - i1), j % 6]
                    idxv[nnz] = i
                    nnz += 1
                nptr += 1
        # -------------------------------------
        #     Last column of the matrix A
        # ------------------------------------
        i0 = 4 * dim + 6 * dim * (_N - 2)
        j0 = (_N - 1) * 6 * dim
        Pl = [basis.evalDerivOnWindow(-1.0, tau[-1], i) for i in range(0, 5)]
        Pr = [basis.evalDerivOnWindow(1.0, tau[-1], i) for i in range(0, 5)]
        Cpl = -np.vstack(Pl[1:5])
        AppN = np.vstack([Pl[0], Pr[1:3], Pr[0]])
        # --------------------------------------------
        for j in range(j0, j0 + 6 * dim):
            ptrv[nptr] = nnz
            i1 = i0 + ((j - j0) // 6) * 4
            for i in range(i1, i1 + 4):
                A[i, j] = Cpl[(i - i1), j % 6]
                nzv[nnz] = Cpl[(i - i1), j % 6]
                idxv[nnz] = i
                nnz += 1
            i1 += (dim - (j - j0) // 6) * 4 + ((j - j0) // 6) * 4
            for i in range(i1, i1 + 4):
                A[i, j] = AppN[i - i1, j % 6]
                nzv[nnz] = AppN[(i - i1), j % 6]
                idxv[nnz] = i
                nnz += 1
            nptr += 1

        ptrv[nptr] = nnz

        return csc_matrix((nzv, idxv, ptrv), shape=2 * (2 * 3 * _N * dim, ))

    def eval_b(self, _wp, _N, _dim, _dwp0, _dwpT, _ddw0, _ddwT):
        assert _wp.shape[0] == _N + 1 and _wp.shape[1] == _dim
        dim = _dim
        b = np.zeros((2 * _N * 3 * _dim, ))

        for i in range(dim):
            b[4 * i:4 * i + 4] = (_wp[0][i], _dwp0[i], _ddw0[i], _wp[1][i])

        i0 = 8 * dim

        for idxwp, _ in enumerate(_wp[1:-2]):
            idxwp += 1
            for i in range(dim):
                b[i0 + 2 * i:i0 + 2 * i + 2] = (_wp[idxwp][i],
                                                _wp[idxwp + 1][i])
            i0 += 6 * dim

        for i in range(dim):
            b[i0 + 4 * i:i0 + 4 * i + 4] = (_wp[-2][i], _dwpT[i], _ddwT[i],
                                            _wp[-1][i])
        return b

    def test_derivative_A(self):
        for i in range(5):
            dim = 6  # np.random.randint(2, 6)
            N = 20  # np.random.randint(3, 120)
            a = np.random.rand()
            splcalc = cSplineCalc(dim, N, cBasis1010(a))

            tauv0 = 0.5 + np.random.rand(N)
            dtau = 0.0001
            z = np.ones((splcalc.linsys_shape_, ))
            for iinter, taui in enumerate(tauv0):
                tauv1 = tauv0.copy()
                tauv1[iinter] += -dtau
                A0 = splcalc.eval_A(tauv1).copy()
                tauv1[iinter] += 2 * dtau
                A1 = splcalc.eval_A(tauv1)

                dAdtaui = 0.5 * (A1 - A0) / dtau

                a = np.max(np.abs(dAdtaui))

                testVal = dAdtaui.dot(z)
                nomVal = np.ravel(
                    splcalc.eval_dAdtiy(tauv0, iinter, z).todense())

                ev = np.abs(testVal - nomVal)
                testValMax = np.max(testVal)
                nomValMax = np.max(nomVal)
                e = np.max(ev)
                ep = e / (nomValMax)
                assert ep < 5.0e-6, '''
                error              = {:10.3e}
                max nominal value  = {:10.3e}
                max test value     = {:10.3e}
                tau compoenet      = {:d}
                interation         = {:d}
                '''.format(e, nomValMax, testValMax, iinter, i)

def main():
    unittest.main()

if __name__ == '__main__':
    main()

ext: py

max_stars_repo_name: rafaelrojasmiliani/gsplines

max_issues_repo_name: rafaelrojasmiliani/gsplines

max_forks_repo_name: rafaelrojasmiliani/gsplines

reason: import numpy,from scipy


text:
#Create a 3x3x3 array with random values

import numpy as np

vector = np.random.uniform((3,3,3))
print(vector)

'''
Output

[1.61309016 1.09376516 2.9416863 ]

'''

reason: import numpy


text:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from mypath import Path
from networks.backbone import build_backbone
from networks.CoarseNet import CoarseNet
from networks.FineNet import FineNet

affine_par = True

class PSPModule(nn.Module):
    """
    Pyramid Scene Parsing module
    """

    def __init__(
        self, in_features=2048, out_features=512, sizes=(1, 2, 3, 6), n_classes=1
    ):
        super(PSPModule, self).__init__()
        self.stages = []
        self.stages = nn.ModuleList(
            [self._make_stage_1(in_features, size) for size in sizes]
        )
        self.bottleneck = self._make_stage_2(
            in_features * (len(sizes) // 4 + 1), out_features
        )
        self.relu = nn.ReLU()

    def _make_stage_1(self, in_features, size):
        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
        conv = nn.Conv2d(in_features, in_features // 4, kernel_size=1, bias=False)
        bn = nn.BatchNorm2d(in_features // 4, affine=affine_par)
        relu = nn.ReLU(inplace=True)
        return nn.Sequential(prior, conv, bn, relu)

    def _make_stage_2(self, in_features, out_features):
        conv = nn.Conv2d(in_features, out_features, kernel_size=1, bias=False)
        bn = nn.BatchNorm2d(out_features, affine=affine_par)
        relu = nn.ReLU(inplace=True)

        return nn.Sequential(conv, bn, relu)

    def forward(self, feats):
        h, w = feats.size(2), feats.size(3)
        priors = [
            F.upsample(
                input=stage(feats), size=(h, w), mode="bilinear", align_corners=True
            )
            for stage in self.stages
        ]
        priors.append(feats)
        bottle = self.relu(self.bottleneck(torch.cat(priors, 1)))
        return bottle

class SegmentationNetwork(nn.Module):
    def __init__(
        self,
        backbone="resnet",
        output_stride=16,
        num_classes=21,
        nInputChannels=3,
        sync_bn=True,
        freeze_bn=False,
    ):
        super(SegmentationNetwork, self).__init__()
        output_shape = 128
        channel_settings = [512, 1024, 512, 256]
        self.Coarse_net = CoarseNet(channel_settings, output_shape, num_classes)
        self.Fine_net = FineNet(channel_settings[-1], output_shape, num_classes)
        BatchNorm = nn.BatchNorm2d
        self.backbone = build_backbone(
            backbone, output_stride, BatchNorm, nInputChannels, pretrained=False
        )
        self.psp4 = PSPModule(
            in_features=2048, out_features=512, sizes=(1, 2, 3, 6), n_classes=256
        )
        self.upsample = nn.Upsample(
            size=(512, 512), mode="bilinear", align_corners=True
        )
        if freeze_bn:
            self.freeze_bn()

    def forward(self, input):
        (
            low_level_feat_4,
            low_level_feat_3,
            low_level_feat_2,
            low_level_feat_1,
        ) = self.backbone(input)
        low_level_feat_4 = self.psp4(low_level_feat_4)
        res_out = [
            low_level_feat_4,
            low_level_feat_3,
            low_level_feat_2,
            low_level_feat_1,
        ]
        coarse_fms, coarse_outs = self.Coarse_net(res_out)
        fine_out = self.Fine_net(coarse_fms)
        coarse_outs[0] = self.upsample(coarse_outs[0])
        coarse_outs[1] = self.upsample(coarse_outs[1])
        coarse_outs[2] = self.upsample(coarse_outs[2])
        coarse_outs[3] = self.upsample(coarse_outs[3])
        fine_out = self.upsample(fine_out)
        return coarse_outs[0], coarse_outs[1], coarse_outs[2], coarse_outs[3], fine_out

    def inference(self, input, IOG_points):
        """Same as calling mainnetwork.forward
        but adds a dummy None output to match the output of the refinement net
        """
        corse_out_0, corse_out_1, corse_out_2, corse_out_3, fine_out = self.forward(
            input
        )
        return None, corse_out_0, corse_out_1, corse_out_2, corse_out_3, fine_out

    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()

    def get_1x_lr_params(self):
        modules = [self.backbone]
        for i in range(len(modules)):
            for m in modules[i].named_modules():
                if isinstance(m[1], nn.Conv2d) or isinstance(m[1], nn.BatchNorm2d):
                    for p in m[1].parameters():
                        if p.requires_grad:
                            yield p

    def get_10x_lr_params(self):
        modules = [self.Coarse_net, self.Fine_net, self.psp4, self.upsample]
        for i in range(len(modules)):
            for m in modules[i].named_modules():
                if isinstance(m[1], nn.Conv2d) or isinstance(m[1], nn.BatchNorm2d):
                    for p in m[1].parameters():
                        if p.requires_grad:
                            yield p

def Network(
    nInputChannels=5,
    num_classes=1,
    backbone="resnet101",
    output_stride=16,
    sync_bn=None,
    freeze_bn=False,
    pretrained=False,
):
    model = SegmentationNetwork(
        nInputChannels=nInputChannels,
        num_classes=num_classes,
        backbone=backbone,
        output_stride=output_stride,
        sync_bn=sync_bn,
        freeze_bn=freeze_bn,
    )
    if pretrained:
        load_pth_name = Path.models_dir()
        pretrain_dict = torch.load(
            load_pth_name, map_location=lambda storage, loc: storage
        )
        conv1_weight_new = np.zeros((64, 5, 7, 7))
        conv1_weight_new[:, :3, :, :] = pretrain_dict["conv1.weight"].cpu().data
        pretrain_dict["conv1.weight"] = torch.from_numpy(conv1_weight_new)
        state_dict = model.state_dict()
        model_dict = state_dict
        for k, v in pretrain_dict.items():
            kk = "backbone." + k
            if kk in state_dict:
                model_dict[kk] = v
        state_dict.update(model_dict)
        model.load_state_dict(state_dict)
    return model

max_stars_repo_name: Suveksha/labelflow

max_issues_repo_name: Suveksha/labelflow

max_forks_repo_name: Suveksha/labelflow


text:
import torch
import numpy as np

class Buffer():
    def __init__(self, max_size, s_dim, a_dim):
        self.max_size = max_size
        self.s_dim = s_dim
        self.a_dim = a_dim
        self.s_buffer = np.zeros((max_size, s_dim), dtype=np.float32)
        self.a_buffer = np.zeros((max_size, a_dim), dtype=np.float32)
        self.s_p_buffer = np.zeros((max_size, s_dim), dtype=np.float32)
        self.r_buffer = np.zeros((max_size), dtype=np.float32)
        self.done_buffer = np.zeros((max_size), dtype=np.float32)
        self.size = 0
        self.next = 0

    def add_tuple(self, s, a, r, s_p, done):
        insert_idx = self.next % self.max_size
        self.s_buffer[insert_idx,:] = s
        self.a_buffer[insert_idx,:] = a
        self.s_p_buffer[insert_idx,:] = s_p
        self.r_buffer[insert_idx] = r
        self.done_buffer[insert_idx] = done
        self.next += 1
        self.size = max(self.size, insert_idx)

    def sample(self, batch_size):
        batch_size = min(batch_size, self.size)
        idxs = torch.randint(low=0, high=self.size, size=(batch_size,))
        s = self.s_buffer[idxs,:]
        a = self.a_buffer[idxs,:]
        s_p = self.s_p_buffer[idxs,:]
        r = self.r_buffer[idxs]
        done = self.done_buffer[idxs]

        return torch.tensor(s, requires_grad=False), \
               torch.tensor(a, requires_grad=False), \
               torch.tensor(r, requires_grad=False), \
               torch.tensor(s_p, requires_grad=False), \
               torch.tensor(done, requires_grad=False)

max_stars_repo_name: tskTNK/LunarLandingSimulation

max_issues_repo_name: tskTNK/LunarLandingSimulation

max_forks_repo_name: tskTNK/LunarLandingSimulation


text:
import numpy as np
import time
from meta_mb.logger import logger
import gym
from gym import error, spaces
from meta_mb.meta_envs.base import MetaEnv

from blue_interface.blue_interface import BlueInterface

class ArmReacherEnv(MetaEnv, BlueInterface, gym.utils.EzPickle):
    def __init__(self, match_joints=False, fixed=False, side='right', ip='127.0.0.1', port=9090):
        self.goal = np.array([0.1, 0.21, -0.55])
        self.joint_goal_pos = np.zeros(7) # joint position control vs torque control
        self.match_joints = match_joints
        self.fixed = fixed
        max_torques = np.array([5, 5, 4, 4, 3, 2, 2]) #control 7 different joints not including the grippers
        self.frame_skip = 1
        self.dt = 0.2
        super(ArmReacherEnv, self).__init__(side, ip, port)
        self.init_qpos = self.get_joint_positions()
        self._prev_qpos = self.init_qpos.copy()
        self.act_dim = len(max_torques)
        self.obs_dim = len(self._get_obs())
        self._low = -max_torques
        self._high = max_torques
        self.positions = {}
        self.actions = {}
        gym.utils.EzPickle.__init__(self)

    def _get_obs(self):
        return np.concatenate([
            np.concatenate((self.get_joint_positions(), self.goal)),
            self.get_joint_velocities(),
            self.tip_position,
            self.vec_gripper_to_goal,
            ]).reshape(-1)

    def step(self, act):
        self._prev_qpos = self.get_joint_positions()
        self._prev_qvel = self.get_joint_velocities()
        if len(act) == 1:
            act = act[0]
        self.do_simulation(act, self.frame_skip)

        #self.correction() #FIXME
        if not self.match_joints:
            vec = self.vec_gripper_to_goal
            reward_dist = -np.linalg.norm(vec)
        else:
            reward = self.joint_match()

        reward_ctrl = -np.square(act/(2 * self._high)).sum()
        reward = reward_dist + 0.5 * 0.1 * reward_ctrl
        ob = self._get_obs()
        done = False

        if self.actions is not None:
            action_num = len(self.actions)
            self.actions.update({action_num : act})

        if self.positions is not None:
            if len(self.positions) == 0:
                self.positions = dict({0 : np.vstack((self._prev_qpos, self._prev_qvel))})
            else:
                arr = np.vstack((self.get_joint_positions(), self.get_joint_velocities()))
                self.positions.update({len(self.positions) : arr})

        return ob, reward, done, dict(reward_dist=reward_dist, reward_ctrl=reward_ctrl)

    def viewer_setup(self):
        self.viewer.cam.trackbodyid = 0

    def do_simulation(self, action,frame_skip):
        action = np.clip(action, self._low, self._high)
        assert frame_skip > 0
        for _ in range(frame_skip):
            time.sleep(self.dt)
            self.set_joint_torques(action)

    def reward(self, obs, act, obs_next):
        assert obs.ndim == act.ndim == obs_next.ndim
        if obs.ndim == 2:
            assert obs.shape == obs_next.shape and act.shape[0] == obs.shape[0]
            reward_ctrl = -0.5 * 0.1 * np.sum(np.square(act / (2 * self._high)), axis=1)
            if self.match_joints:
                reward_dist = -np.linalg.norm(self.joint_goal())
            else:
                reward_dist = -np.linalg.norm(obs_next[:, -3:], axis=1)
            reward = reward_dist + reward_ctrl
            return np.clip(reward, -1e2, 1e2)

        elif obs.ndim == 1:
            return self.reward(obs[None], act[None], obs_next[None])[0]

        else:
            raise NotImplementedError

    def reset(self):
        self.set_joint_positions(np.zeros(7), duration=5)
        while True:
            #self.goal = np.append(np.random.uniform(low=-0.75, hig-0.25), np.random.uniform(low=-0.2, high=0.2, size=3))
            if self.fixed and self.match_joints:
                self.joint_goal_pos = np.array([-np.pi/4.0, -0.78539816339, np.pi/4.0, 0, 0, 0, 0])
            elif not self.fixed and self.match_joints:
                self.joint_goal_pos = self.joint_goal()
            elif self.fixed:
                self.goal = np.array([-0.44, -0.1, 0.21])
            else:
                self.goal = np.random.uniform(low=[-0.75, -0.25, 0.25], high=[-0.25, 0.25, 0.5])
            if np.linalg.norm(self.goal) < 2:
                break
        return self._get_obs()

    def joint_goal(self):
        base_roll = np.random.uniform(low=-(np.pi/4.0), high=np.pi/4.0)
        right_shoulder_lift = np.random.uniform(low=-2.2944, high=-np.pi/4.0)
        right_shoulder_roll = np.random.uniform(low=-(np.pi/4.0), high=np.pi/4.0)
        right_elbow_lift = np.random.uniform(low=0, high=0)
        right_elbow_roll = np.random.uniform(low=0, high=0)
        right_wrist_lift = np.random.uniform(low=0, high=0)
        right_wrist_roll = np.random.uniform(low=0, high=0)
        return np.array([
            base_roll,
            right_shoulder_lift,
            right_shoulder_roll,
            right_elbow_lift,
            right_elbow_roll,
            right_wrist_lift,
            right_wrist_roll
        ])

    def joint_match(self):
        real_joint_positions = self.get_joint_positions()

        return -np.linalg.norm(real_joint_positions, self.joint_goal_pos)

    def correction(self):
        try:
            joints = self.get_joint_positions()
            joint_velocities = self.get_joint_velocities()
            shoulder_lift_joint = joints[1]
            if shoulder_lift_joint > -0.01:
                correction = -(0.5 * abs(shoulder_lift_joint)) + shoulder_lift_joint
            joints[1] = correction
            joint_velocities[1] = joint_velocities[1] * 0.9
            print("correction!")
            self.set_joint_positions(joints)
            self._joint_velocities = joint_velocities
        except:
            pass

    @property
    def tip_position(self):
        pose = self.get_cartesian_pose()
        return pose['position']

    @property
    def vec_gripper_to_goal(self):
        gripper_pos = self.tip_position
        vec_gripper_to_goal = self.goal - gripper_pos
        return vec_gripper_to_goal

    @property
    def vec_arm_to_goal_pos(self):
        arm_pos = self.get_joint_positions()
        vec_arm_to_goal = self.goal_pos - arm_pos
        return vec_arm_to_goal

    @property
    def action_space(self):
        return spaces.Box(low=self._low, high=self._high, dtype=np.float32)

    @property
    def observation_space(self):
        low = np.ones(self.obs_dim) * -1e6
        high = np.ones(self.obs_dim) * 1e6
        return spaces.Box(low=low, high=high, dtype=np.float32)

    def log_diagnostics(self, paths, prefix=''):
        dist = [-path["env_infos"]['reward_dist'] for path in paths]
        final_dist = [-path["env_infos"]['reward_dist'][-1] for path in paths]
        ctrl_cost = [-path["env_infos"]['reward_ctrl'] for path in paths]

        logger.logkv(prefix + 'AvgDistance', np.mean(dist))
        logger.logkv(prefix + 'AvgFinalDistance', np.mean(final_dist))
        logger.logkv(prefix + 'AvgCtrlCost', np.mean(ctrl_cost))

if __name__ == "__main__":
    env = ArmReacherEnv(match_joints=True)
    while True:
        env.reset()
        for _ in range(100):
            env.step(env.action_space.sample())
            env.render()

max_stars_repo_name: iclavera/meta-mb

max_issues_repo_name: iclavera/meta-mb

max_forks_repo_name: iclavera/meta-mb


text:
import numpy as np
import cv2
import math

def get_center(bbox):
    center = [int((bbox[0]+bbox[2])*0.5),int((bbox[1]+bbox[3])*0.5)] # x,y
    return center

def get_center_distance(center_0,center_1):
    distance = int((abs(center_0[0]-center_1[0])**2 + abs(center_0[1]-center_1[1])**2) ** 0.5)
    return distance

def get_v(bbox_mes_0,bbox_mes_1,v_max,v_norm,v_bias):
    if bbox_mes_0[5] != 'person':
        center_0 = get_center(bbox_mes_0)
        # print(bbox_mes_0[0][0])
        center_1 = get_center(bbox_mes_1)
        v = get_center_distance(center_0 ,center_1)
        v = round(math.exp(v/v_norm)**2*2 +  math.exp(v_bias),2)
        if v >= v_max:
            v = 'v='+str(v)+'--> Too-Fast'
    else:
        v = ' '
    return str(v)
# # example
# a = [[746, 453, 998, 637, 1]]
# b = [[752, 443, 998, 622, 1]]
# num = get_v(a,b,2)
# print(num)

def xyxy2xywh(location):
    x = (location[0]+location[2])*0.5
    y = (location[1] + location[3]) * 0.5
    w = location[2] - location[0]
    h = location[3] - location[1]
    return [x,y,w,h]

def estimateSpeed(location2, location1,v_max,v_norm,v_bias):
    location1 = xyxy2xywh(location1)
    location2 = xyxy2xywh(location2)
    d_pixels = math.sqrt(math.pow(location2[0] - location1[0], 2) + math.pow(location2[1] - location1[1], 2))
    # carWidht = 2
    # ppm = location2[2] / carWidht
    ppm = 10
    d_meters = d_pixels / ppm
    speed = d_meters * 10 * 3.6
    speed = (speed/v_norm)+  math.exp( v_bias )
    return speed


text:
import keras
from keras.models import Sequential
from keras.layers import Dense, Activation, Dropout
from keras.layers import LSTM
from keras.optimizers import RMSprop
import numpy as np
import os

import datetime
import pickle

import argparse
import sys

import p3b2 as bmk
import candle

def initialize_parameters():

    # Build benchmark object
    p3b2Bmk = bmk.BenchmarkP3B2(bmk.file_path, 'p3b2_default_model.txt', 'keras',
    prog='p3b2_baseline', desc='Multi-task (DNN) for data extraction from clinical reports - Pilot 3 Benchmark 1')

    # Initialize parameters
    gParameters = candle.initialize_parameters(p3b2Bmk)
    #bmk.logger.info('Params: {}'.format(gParameters))

    return gParameters

class LossHistory( keras.callbacks.Callback ):
    def on_train_begin( self, logs= {} ):
        self.losses = []

    def on_batch_end( self, batch, logs= {} ):
        self.losses.append( logs.get( 'loss' ) )

def sample( preds, temperature= 1.0 ):
    # helper function to sample an index from a probability array
    preds = np.asarray( preds ).astype( 'float64' )
    preds = np.log( preds ) / temperature
    exp_preds = np.exp( preds )
    preds = exp_preds / np.sum( exp_preds )
    probas = np.random.multinomial( 1, preds, 1 )
    return np.argmax( probas )

def run(gParameters):

    origin = gParameters['data_url']
    train_data = gParameters['train_data']
    data_loc = candle.fetch_file(origin+train_data, untar=True, md5_hash=None, subdir='Pilot3')

    print( 'Data downloaded and stored at: ' + data_loc )
    data_path = os.path.dirname(data_loc)
    print( data_path )

    kerasDefaults = candle.keras_default_config()

    rnn_size = gParameters['rnn_size']
    n_layers = gParameters['n_layers']
    learning_rate = gParameters['learning_rate']
    dropout = gParameters['drop']
    recurrent_dropout = gParameters['recurrent_dropout']
    n_epochs = gParameters['epochs']
    data_train = data_path+'/data.pkl'
    verbose = gParameters['verbose']
    savedir = gParameters['output_dir']
    do_sample = gParameters['do_sample']
    temperature = gParameters['temperature']
    primetext = gParameters['primetext']
    length = gParameters['length']

    # load data from pickle
    f = open( data_train, 'rb' )

    if ( sys.version_info > ( 3, 0 ) ):
        classes = pickle.load( f, encoding= 'latin1' )
        chars = pickle.load( f, encoding= 'latin1' )
        char_indices = pickle.load( f, encoding= 'latin1' )
        indices_char = pickle.load( f, encoding= 'latin1' )

        maxlen = pickle.load( f, encoding= 'latin1' )
        step = pickle.load( f, encoding= 'latin1' )

        X_ind = pickle.load( f, encoding= 'latin1' )
        y_ind = pickle.load( f, encoding= 'latin1' )
    else:
        classes = pickle.load( f )
        chars = pickle.load( f )
        char_indices = pickle.load( f )
        indices_char = pickle.load( f )

        maxlen = pickle.load( f )
        step = pickle.load( f )

        X_ind = pickle.load( f )
        y_ind = pickle.load( f )

    f.close()

    [ s1, s2 ] = X_ind.shape
    print( X_ind.shape )
    print( y_ind.shape )
    print( maxlen )
    print( len( chars ) )

    X = np.zeros( ( s1, s2, len( chars ) ), dtype=np.bool )
    y = np.zeros( ( s1, len( chars ) ), dtype=np.bool )

    for i in range( s1 ):
        for t in range( s2 ):
            X[ i, t, X_ind[ i, t ] ] = 1
        y[ i, y_ind[ i ] ] = 1

    # build the model: a single LSTM
    if verbose:
        print( 'Build model...' )

    model = Sequential()

    # for rnn_size in rnn_sizes:
    for k in range( n_layers ):
        if k < n_layers - 1:
            ret_seq = True
        else:
            ret_seq = False

        if k == 0:
            model.add( LSTM( rnn_size, input_shape= ( maxlen, len( chars ) ), return_sequences= ret_seq,
                                 dropout= dropout, recurrent_dropout= recurrent_dropout ) )
        else:
            model.add( LSTM( rnn_size, dropout= dropout, recurrent_dropout= recurrent_dropout, return_sequences= ret_seq ) )

    model.add( Dense( len( chars ) ) )
    model.add( Activation( gParameters['activation'] ) )

    optimizer = candle.build_optimizer(gParameters['optimizer'],
                                            gParameters['learning_rate'],
                                            kerasDefaults)

    model.compile( loss= gParameters['loss'], optimizer= optimizer )

    if verbose:
        model.summary()

    for iteration in range( 1, n_epochs + 1 ):
        if verbose:
            print()
            print('-' * 50)
            print('Iteration', iteration)

        history = LossHistory()
        model.fit( X, y, batch_size= 100, epochs= 1, callbacks= [ history ] )

        loss = history.losses[ -1 ]
        if verbose:
            print( loss )

        dirname = savedir
        if len( dirname ) > 0 and not dirname.endswith( '/' ):
            dirname = dirname + '/'

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

        # serialize model to JSON
        model_json = model.to_json()
        with open( dirname + "/model_" + str( iteration ) + "_" + "{:f}".format( loss ) + ".json", "w" ) as json_file:
            json_file.write( model_json )

        # serialize weights to HDF5
        model.save_weights( dirname + "/model_" + str( iteration ) + "_" + "{:f}".format( loss ) + ".h5" )

        if verbose:
            print( "Checkpoint saved." )

        if do_sample:
            outtext = open( dirname + "/example_" + str( iteration ) + "_" + "{:f}".format( loss ) + ".txt", "w" , encoding= 'utf-8' )

            diversity = temperature

            outtext.write('----- diversity:' + str( diversity ) + "\n" )

            generated = ''
            seedstr = primetext

            outtext.write('----- Generating with seed: "' + seedstr + '"' + "\n" )

            sentence = " " * maxlen

            # class_index = 0
            generated += sentence
            outtext.write( generated )

            for c in seedstr:
                sentence = sentence[1:] + c
                x = np.zeros( ( 1, maxlen, len( chars ) ) )
                for t, char in enumerate(sentence):
                    x[ 0, t, char_indices[ char ] ] = 1.

                preds = model.predict( x, verbose= verbose )[ 0 ]
                next_index = sample( preds, diversity )
                next_char = indices_char[ next_index ]

                generated += c

                outtext.write( c )

            for i in range( length ):
                x = np.zeros( ( 1, maxlen, len( chars ) ) )
                for t, char in enumerate( sentence ):
                    x[ 0, t, char_indices[ char ] ] = 1.

                preds = model.predict( x, verbose= verbose )[ 0 ]
                next_index = sample( preds, diversity )
                next_char = indices_char[ next_index ]

                generated += next_char
                sentence = sentence[ 1 : ] + next_char

            if (sys.version_info > (3, 0)):
                outtext.write( generated + '\n' )
            else:
                outtext.write( generated.decode('utf-8').encode('utf-8') + '\n' )

            outtext.close()

if __name__  == "__main__":

    gParameters = initialize_parameters()
    run(gParameters)

max_stars_repo_name: mdorier/Benchmarks

max_issues_repo_name: mdorier/Benchmarks

max_forks_repo_name: mdorier/Benchmarks


text:
import os
import numpy as np
import pandas as pd
from sklearn.metrics import hamming_loss, auc, f1_score, precision_recall_curve, roc_auc_score, roc_curve, average_precision_score

    # https://machinelearningmastery.com/roc-curves-and-precision-recall-curves-for-classification-in-python/

    # https://www.kaggle.com/kmkarakaya/multi-label-model-evaluation

    # # Minz metrics
    # roc_aucs  = roc_auc_score(y_true, y_pred_matrix, average = 'macro')
    # pr_aucs = average_precision_score(y_true, y_pred_matrix, average = 'macro')

    # np.average([average_precision_score(y_true[it,:], y_pred_matrix[it,:]) for it in range(y_true.shape[0])])
    # np.average([pr_score(y_true[it, :], y_pred_matrix[it, :]) for it in range(y_true.shape[0])])

    # np.average([auc_score(y_true[:, it], y_pred_matrix[:, it]) for it in range(y_true.shape[1])])
    # np.average([pr_score(y_true[:, it], y_pred_matrix[:, it]) for it in range(y_true.shape[1])])

"""Hamming Loss metric"""
def hamming_loss_eval(y_true, y_pred):
    er = hamming_loss(y_true = y_true, y_pred = y_pred)
    return er

"""Classification accuracy with area under ROC"""
def custom_roc_auc_score(y_true, y_pred, average = 'macro'):
    results = list()
    for it_class in range(y_true.shape[1]):
        if len(np.unique(y_true[:, it_class])) == 1:
            results.append(0.5)
        else:
            results.append(roc_auc_score(y_true[:, it_class], y_pred[:, it_class]))
    return np.average(results)

"""Area under ROC using the Precision/Recall curve"""
def custom_average_precision_score(y_true, y_pred, average = 'macro'):
    results = list()
    for it_class in range(y_true.shape[1]):
        if len(np.unique(y_true[:, it_class])) == 1:
            results.append(0.5)
        else:
            results.append(average_precision_score(y_true[:, it_class], y_pred[:, it_class]))
    return np.average(results)

"""Function for computing the different metrics"""
def compute_metrics(config, y_true, y_prob_matrix):

    data = {'cls' : config.shallow_model,
            'arch' : config.model_type
    }

    # Hamming loss metric
    # data['hl'] = hamming_loss_eval(y_true = y_true, y_pred = y_pred)
    data['hl'] = '-'

    # ROC-AUC:
    data['roc_auc'] = custom_roc_auc_score(y_true = y_true, y_pred = y_prob_matrix, average = 'macro')

    # PR-AUC:
    data['pr_auc'] = custom_average_precision_score(y_true = y_true, y_pred = y_prob_matrix, average = 'macro')

    return data

if __name__ == '__main__':

    y_true = [
        [0, 1, 1, 0],
        [1, 0, 1, 1],
        [0, 0, 0, 1],
        [0, 1, 0, 1]
    ]

    y_pred = [
        [0, 1, 1, 0],
        [1, 1, 0, 1],
        [1, 1, 0, 0],
        [1, 1, 1, 1]
    ]

    print("Evaluation protocol")

max_stars_repo_name: jose-jvmas/sota-music-tagging-models

max_issues_repo_name: jose-jvmas/sota-music-tagging-models

max_forks_repo_name: jose-jvmas/sota-music-tagging-models


text:
import config
import dataset
import torch
import time
import transformers
import numpy as np
from tqdm import tqdm
from torch.utils.data import DataLoader

class BertForFakeNewsDetection:

    def __init__(self):
        self.model = config.MODEL

    def get_features(self,train=True):

        fakenews_dataset = dataset.FakeNewsDataset(train=train)
        n_samples = fakenews_dataset.n_samples
        dataloader = DataLoader(dataset=fakenews_dataset, batch_size=config.BATCH_SIZE, shuffle=False, num_workers=1)
        dataiter = iter(dataloader)
        if train == True:
            labels = []
        features = []

        with torch.no_grad():
            for ind, d in tqdm(enumerate(dataloader),total=len(dataloader)):
                if train == True:
                    label = np.array(d['labels'])
                    labels.extend(label)
                input_ids = d['input_ids']
                attention_mask = d['attention_mask']
                segment_ids = d['segment_ids']
                outputs = self.model(input_ids,token_type_ids = segment_ids, attention_mask=attention_mask)
                last_hidden_state = outputs[1]
                features.append(last_hidden_state[-1][:,0,:].numpy())
                if (ind+1)%10000 == 0:
                    np.save(f"Saved_Model/data_{ind+1}.npy",np.array(features))
        features = np.array(features).reshape(n_samples,768)

        if train == True:
            labels = np.array(labels)
            return features, labels
        else:
            return features


text:
import numpy as np
import torch

from stog.models.model import Model
from stog.utils.nn import masked_log_softmax
from stog.modules.attention import BiaffineAttention
from stog.modules.linear import BiLinear
from stog.metrics import AttachmentScores
from stog.algorithms.maximum_spanning_tree import decode_mst_with_coreference, decode_mst

class DeepBiaffineGraphDecoder(torch.nn.Module):

    def __init__(
            self,
            decode_algorithm,
            head_sentinel,
            edge_node_h_linear,
            edge_node_m_linear,
            edge_label_h_linear,
            edge_label_m_linear,
            encode_dropout,
            biaffine_attention,
            edge_label_bilinear
    ):
        super(DeepBiaffineGraphDecoder, self).__init__()
        self.decode_algorithm = decode_algorithm
        self.head_sentinel = head_sentinel
        self.edge_node_h_linear = edge_node_h_linear
        self.edge_node_m_linear = edge_node_m_linear
        self.edge_label_h_linear = edge_label_h_linear
        self.edge_label_m_linear = edge_label_m_linear
        self.encode_dropout = encode_dropout
        self.biaffine_attention = biaffine_attention
        self.edge_label_bilinear = edge_label_bilinear

        self.metrics = AttachmentScores()

        self.minus_inf = -1e8

    def forward(self, memory_bank, edge_heads, edge_labels, corefs, mask):
        num_nodes = mask.sum().item()

        memory_bank, edge_heads, edge_labels, corefs, mask = self._add_head_sentinel(
            memory_bank, edge_heads, edge_labels, corefs, mask)

        (edge_node_h, edge_node_m), (edge_label_h, edge_label_m) = self.encode(memory_bank)

        edge_node_scores = self._get_edge_node_scores(edge_node_h, edge_node_m, mask)

        edge_node_nll, edge_label_nll = self.get_loss(
            edge_label_h, edge_label_m, edge_node_scores, edge_heads, edge_labels, mask)

        pred_edge_heads, pred_edge_labels = self.decode(
            edge_label_h, edge_label_m, edge_node_scores, corefs, mask)

        self.metrics(
            pred_edge_heads, pred_edge_labels, edge_heads[:, 1:], edge_labels[:, 1:], mask[:, 1:],
            edge_node_nll.item(), edge_label_nll.item()
        )

        return dict(
            edge_heads=pred_edge_heads,
            edge_labels=pred_edge_labels,
            loss=(edge_node_nll + edge_label_nll) / num_nodes,
            total_loss=edge_node_nll + edge_label_nll,
            num_nodes=torch.tensor(float(num_nodes)).type_as(memory_bank)
        )

    def encode(self, memory_bank):
        """
        Map contextual representation into specific space (w/ lower dimensionality).

        :param input: [batch, length, hidden_size]
        :return:
            edge_node: a tuple of (head, modifier) hidden state with size [batch, length, edge_hidden_size]
            edge_label: a tuple of (head, modifier) hidden state with size [batch, length, label_hidden_size]
        """

        # Output: [batch, length, edge_hidden_size]
        edge_node_h = torch.nn.functional.elu(self.edge_node_h_linear(memory_bank))
        edge_node_m = torch.nn.functional.elu(self.edge_node_m_linear(memory_bank))

        # Output: [batch, length, label_hidden_size]
        edge_label_h = torch.nn.functional.elu(self.edge_label_h_linear(memory_bank))
        edge_label_m = torch.nn.functional.elu(self.edge_label_m_linear(memory_bank))

        # Apply dropout to certain node?
        # [batch, length * 2, hidden_size]
        edge_node = torch.cat([edge_node_h, edge_node_m], dim=1)
        edge_label = torch.cat([edge_label_h, edge_label_m], dim=1)
        edge_node = self.encode_dropout(edge_node.transpose(1, 2)).transpose(1, 2)
        edge_label = self.encode_dropout(edge_label.transpose(1, 2)).transpose(1, 2)

        edge_node_h, edge_node_m = edge_node.chunk(2, 1)
        edge_label_h, edge_label_m = edge_label.chunk(2, 1)

        return (edge_node_h, edge_node_m), (edge_label_h, edge_label_m)

    def get_loss(self, edge_label_h, edge_label_m, edge_node_scores, edge_heads, edge_labels, mask):
        """
        :param edge_label_h: [batch, length, hidden_size]
        :param edge_label_m: [batch, length, hidden_size]
        :param edge_node_scores:  [batch, length, length]
        :param edge_heads:  [batch, length]
        :param edge_labels:  [batch, length]
        :param mask: [batch, length]
        :return:  [batch, length - 1]
        """
        batch_size, max_len, _ = edge_node_scores.size()

        edge_node_log_likelihood = masked_log_softmax(
            edge_node_scores, mask.unsqueeze(2) + mask.unsqueeze(1), dim=1)

        edge_label_scores = self._get_edge_label_scores(edge_label_h, edge_label_m, edge_heads)
        edge_label_log_likelihood = torch.nn.functional.log_softmax(edge_label_scores, dim=2)

        # Create indexing matrix for batch: [batch, 1]
        batch_index = torch.arange(0, batch_size).view(batch_size, 1).type_as(edge_heads)
        # Create indexing matrix for modifier: [batch, modifier_length]
        modifier_index = torch.arange(0, max_len).view(1, max_len).expand(batch_size, max_len).type_as(edge_heads)
        # Index the log likelihood of gold edges.
        _edge_node_log_likelihood = edge_node_log_likelihood[
            batch_index, edge_heads.data, modifier_index]
        _edge_label_log_likelihood = edge_label_log_likelihood[
            batch_index, modifier_index, edge_labels.data]

        # Exclude the dummy root.
        # Output [batch, length - 1]
        gold_edge_node_nll = - _edge_node_log_likelihood[:, 1:].sum()
        gold_edge_label_nll = - _edge_label_log_likelihood[:, 1:].sum()

        return gold_edge_node_nll, gold_edge_label_nll

    def decode(self, edge_label_h, edge_label_m, edge_node_scores, corefs, mask):
        if self.decode_algorithm == 'mst':
            return self.mst_decode(edge_label_h, edge_label_m, edge_node_scores, corefs, mask)
        else:
            return self.greedy_decode(edge_label_h, edge_label_m, edge_node_scores, mask)

    def greedy_decode(self, edge_label_h, edge_label_m, edge_node_scores, mask):
        # out_arc shape [batch, length, length]
        edge_node_scores = edge_node_scores.data
        max_len = edge_node_scores.size(1)

        # Set diagonal elements to -inf
        edge_node_scores = edge_node_scores + torch.diag(edge_node_scores.new(max_len).fill_(-np.inf))

        # Set invalid positions to -inf
        minus_mask = (1 - mask.float()) * self.minus_inf
        edge_node_scores = edge_node_scores + minus_mask.unsqueeze(2) + minus_mask.unsqueeze(1)

        # Compute naive predictions.
        # prediction shape = [batch, length]
        _, edge_heads = edge_node_scores.max(dim=1)

        # Based on predicted heads, compute the edge label scores.
        # [batch, length, num_labels]
        edge_label_scores = self._get_edge_label_scores(edge_label_h, edge_label_m, edge_heads)
        _, edge_labels = edge_label_scores.max(dim=2)

        return edge_heads[:, 1:], edge_labels[:, 1:]

    def mst_decode(self, edge_label_h, edge_label_m, edge_node_scores, corefs, mask):
        batch_size, max_length, edge_label_hidden_size = edge_label_h.size()
        lengths = mask.data.sum(dim=1).long().cpu().numpy()

        expanded_shape = [batch_size, max_length, max_length, edge_label_hidden_size]
        edge_label_h = edge_label_h.unsqueeze(2).expand(*expanded_shape).contiguous()
        edge_label_m = edge_label_m.unsqueeze(1).expand(*expanded_shape).contiguous()
        # [batch, max_head_length, max_modifier_length, num_labels]
        edge_label_scores = self.edge_label_bilinear(edge_label_h, edge_label_m)
        edge_label_scores = torch.nn.functional.log_softmax(edge_label_scores, dim=3).permute(0, 3, 1, 2)

        # Set invalid positions to -inf
        minus_mask = (1 - mask.float()) * self.minus_inf
        edge_node_scores = edge_node_scores + minus_mask.unsqueeze(2) + minus_mask.unsqueeze(1)
        # [batch, max_head_length, max_modifier_length]
        edge_node_scores = torch.nn.functional.log_softmax(edge_node_scores, dim=1)

        # [batch, num_labels, max_head_length, max_modifier_length]
        batch_energy = torch.exp(edge_node_scores.unsqueeze(1) + edge_label_scores)

        edge_heads, edge_labels = self._run_mst_decoding(batch_energy, lengths, corefs)
        return edge_heads[:, 1:], edge_labels[:, 1:]

    @staticmethod
    def _run_mst_decoding(batch_energy, lengths, corefs=None):
        edge_heads = []
        edge_labels = []
        for i, (energy, length) in enumerate(zip(batch_energy.detach().cpu(), lengths)):
            # energy: [num_labels, max_head_length, max_modifier_length]
            # scores | label_ids : [max_head_length, max_modifier_length]
            scores, label_ids = energy.max(dim=0)
            # Although we need to include the root node so that the MST includes it,
            # we do not want the dummy root node to be the head of more than one nodes,
            # since there should be only one root in a sentence.
            # Here, we enforce this by setting the scores for all word -> ROOT edges
            # edges to be 0.
            # TODO: set it to -1 seems better?
            scores[0, :] = 0
            # Decode the heads. Because we modify the scores to prevent
            # adding in word -> ROOT edges, we need to find the labels ourselves.
            if corefs is not None:
                coref = corefs[i].detach().cpu().tolist()[:length]
                instance_heads, _ = decode_mst_with_coreference(
                    scores.numpy(), coref, length, has_labels=False)
            else:
                instance_heads, _ = decode_mst(scores.numpy(), length, has_labels=False)

            # Find the labels which correspond to the edges in the max spanning tree.
            instance_head_labels = []
            for child, parent in enumerate(instance_heads):
                instance_head_labels.append(label_ids[parent, child].item())
            # We don't care what the head or tag is for the root token, but by default it's
            # not necessarily the same in the batched vs unbatched case, which is annoying.
            # Here we'll just set them to zero.
            instance_heads[0] = 0
            instance_head_labels[0] = 0
            edge_heads.append(instance_heads)
            edge_labels.append(instance_head_labels)
        return torch.from_numpy(np.stack(edge_heads)), torch.from_numpy(np.stack(edge_labels))

    @classmethod
    def from_params(cls, vocab, params):
        decode_algorithm = params['decode_algorithm']
        input_size = params['input_size']
        edge_node_hidden_size = params['edge_node_hidden_size']
        edge_label_hidden_size = params['edge_label_hidden_size']
        dropout = params['dropout']

        head_sentinel = torch.nn.Parameter(torch.randn([1, 1, input_size]))

        # Transform representations into a space for edge node heads, edge node modifiers
        edge_node_h_linear = torch.nn.Linear(input_size, edge_node_hidden_size)
        edge_node_m_linear = torch.nn.Linear(input_size, edge_node_hidden_size)

        # Transform representations into a space for edge label heads, edge label modifiers
        edge_label_h_linear = torch.nn.Linear(input_size, edge_label_hidden_size)
        edge_label_m_linear = torch.nn.Linear(input_size, edge_label_hidden_size)

        encode_dropout = torch.nn.Dropout2d(p=dropout)

        biaffine_attention = BiaffineAttention(edge_node_hidden_size, edge_node_hidden_size)

        num_labels = vocab.get_vocab_size("head_tags")
        edge_label_bilinear = BiLinear(edge_label_hidden_size, edge_label_hidden_size, num_labels)

        return cls(
            decode_algorithm=decode_algorithm,
            head_sentinel=head_sentinel,
            edge_node_h_linear=edge_node_h_linear,
            edge_node_m_linear=edge_node_m_linear,
            edge_label_h_linear=edge_label_h_linear,
            edge_label_m_linear=edge_label_m_linear,
            encode_dropout=encode_dropout,
            biaffine_attention=biaffine_attention,
            edge_label_bilinear=edge_label_bilinear
        )

    def _add_head_sentinel(self, memory_bank, edge_heads, edge_labels, corefs, mask):
        """
        Add a dummy ROOT at the beginning of each node sequence.
        :param memory_bank: [batch, length, hidden_size]
        :param edge_head: None or [batch, length]
        :param edge_labels: None or [batch, length]
        :param corefs: None or [batch, length]
        :param mask: [batch, length]
        """
        batch_size, _, hidden_size = memory_bank.size()
        head_sentinel = self.head_sentinel.expand([batch_size, 1, hidden_size])
        memory_bank = torch.cat([head_sentinel, memory_bank], 1)
        if edge_heads is not None:
            edge_heads = torch.cat([edge_heads.new_zeros(batch_size, 1), edge_heads], 1)
        if edge_labels is not None:
            edge_labels = torch.cat([edge_labels.new_zeros(batch_size, 1), edge_labels], 1)
        if corefs is not None:
            corefs = torch.cat([corefs.new_zeros(batch_size, 1), corefs], 1)
        mask = torch.cat([mask.new_ones(batch_size, 1), mask], 1)
        return memory_bank, edge_heads, edge_labels, corefs, mask

    def _get_edge_node_scores(self, edge_node_h, edge_node_m, mask):
        edge_node_scores = self.biaffine_attention(edge_node_h, edge_node_m, mask_d=mask, mask_e=mask).squeeze(1)
        return edge_node_scores

    def _get_edge_label_scores(self, edge_label_h, edge_label_m, edge_heads):
        """
        Compute the edge label scores.
        :param edge_label_h: [batch, length, edge_label_hidden_size]
        :param edge_label_m: [batch, length, edge_label_hidden_size]
        :param heads: [batch, length] -- element at [i, j] means the head index of node_j at batch_i.
        :return: [batch, length, num_labels]
        """
        batch_size = edge_label_h.size(0)
        # Create indexing matrix for batch: [batch, 1]
        batch_index = torch.arange(0, batch_size).view(batch_size, 1).type_as(edge_heads.data).long()

        # Select the heads' representations based on the gold/predicted heads.
        # [batch, length, edge_label_hidden_size]
        edge_label_h = edge_label_h[batch_index, edge_heads.data].contiguous()
        edge_label_m = edge_label_m.contiguous()

        # [batch, length, num_labels]
        edge_label_scores = self.edge_label_bilinear(edge_label_h, edge_label_m)

        return edge_label_scores

max_stars_repo_name: sfillwo/stog

max_issues_repo_name: sfillwo/stog

max_forks_repo_name: sfillwo/stog