decagon-pytorch/src/triacontagon/model.py

from .data import Data, \
    EdgeType
import torch
from dataclasses import dataclass
from .weights import init_glorot
import types
from typing import List, \
    Dict, \
    Callable
from .util import _sparse_coo_tensor


@dataclass
class TrainingBatch(object):
    vertex_type_row: int
    vertex_type_column: int
    relation_type_index: int
    edges: torch.Tensor


class Model(torch.nn.Module):
    def __init__(self, data: Data, layer_dimensions: List[int],
        keep_prob: float,
        conv_activation: Callable[[torch.Tensor], torch.Tensor],
        dec_activation: Callable[[torch.Tensor], torch.Tensor],
        **kwargs) -> None:
        super().__init__(**kwargs)

        if not isinstance(data, Data):
            raise TypeError('data must be an instance of Data')

        if not isinstance(conv_activation, types.FunctionType):
            raise TypeError('conv_activation must be a function')

        if not isinstance(dec_activation, types.FunctionType):
            raise TypeError('dec_activation must be a function')

        self.data = data
        self.layer_dimensions = list(layer_dimensions)
        self.keep_prob = float(keep_prob)
        self.conv_activation = conv_activation
        self.dec_activation = dec_activation

        self.conv_weights = None
        self.dec_weights = None
        self.build()


    def build(self) -> None:
        self.conv_weights = torch.nn.ParameterDict()
        for i in range(len(self.layer_dimensions) - 1):
            in_dimension = self.layer_dimensions[i]
            out_dimension = self.layer_dimensions[i + 1]

            for _, et in self.data.edge_types.items():
                weight = init_glorot(in_dimension, out_dimension)
                self.conv_weights[et.vertex_type_row, et.vertex_type_column, i] = \
                    torch.nn.Parameter(weight)

        self.dec_weights = torch.nn.ParameterDict()
        for _, et in self.data.edge_types.items():
            global_interaction, local_variation = \
                et.decoder_factory(self.layer_dimensions[-1],
                    len(et.adjacency_matrices))
            self.dec_weights[et.vertex_type_row, et.vertex_type_column] = \
                torch.nn.ParameterList([
                    torch.nn.Parameter(global_interaction),
                    torch.nn.Parameter(local_variation)
                ])

    def limit_adjacency_matrix_to_rows(self, adjacency_matrix: torch.Tensor,
        rows: torch.Tensor) -> torch.Tensor:

        adj_mat = adjacency_matrix.coalesce()
        adj_mat = torch.index_select(adj_mat, 0, rows)

        adj_mat = adj_mat.coalesce()
        indices = adj_mat.indices()
        indices[0] = rows

        adj_mat = _sparse_coo_tensor(indices, adj_mat.values(), adjacency_matrix.shape)

    def temporary_adjacency_matrix(self, adjacency_matrix: torch.Tensor,
        batch: TrainingBatch, total_connectivity: torch.Tensor) -> torch.Tensor:

        col = batch.vertex_type_column
        rows = batch.edges[:, 0]

        columns = batch.edges[:, 1].sum(dim=0).flatten()
        columns = torch.nonzero(columns)

        for i in range(len(self.layer_dimensions) - 1):

            columns =


    def temporary_adjacency_matrices(self, batch: TrainingBatch) ->
        Dict[Tuple[int, int], List[List[torch.Tensor]]]:

        col = batch.vertex_type_column
        batch.edges[:, 1]

        res = {}

        for _, et in self.data.edge_types.items():
            sum_nonzero = _nonzero_sum(et.adjacency_matrices)
            res[et.vertex_type_row, et.vertex_type_column] = \
                [ self.temporary_adjacency_matrix(adj_mat, batch,
                    et.total_connectivity) \
                        for adj_mat in et.adjacency_matrices ]

        return res

    def forward(self, initial_repr: List[torch.Tensor],
        batch: TrainingBatch) -> torch.Tensor:

        if not isinstance(initial_repr, list):
            raise TypeError('initial_repr must be a list')

        if len(initial_repr) != len(self.data.vertex_types):
            raise ValueError('initial_repr must contain representations for all vertex types')

        if not isinstance(batch, TrainingBatch):
            raise TypeError('batch must be an instance of TrainingBatch')

        adj_matrices = self.temporary_adjacency_matrices(batch)

        row_vertices = initial_repr[batch.vertex_type_row]
        column_vertices = initial_repr[batch.vertex_type_column]