decagon-pytorch/src/triacontagon/data.py

#
# Copyright (C) Stanislaw Adaszewski, 2020
# License: GPLv3
#


from dataclasses import dataclass
from typing import Callable, \
    Tuple, \
    List
import types
from .util import _nonzero_sum, \
    _diag
import torch


@dataclass
class DecodingMatrices(object):
    global_interaction: torch.Tensor
    local_variation: torch.Tensor


@dataclass
class VertexType(object):
    name: str
    count: int


@dataclass
class EdgeType(object):
    name: str
    vertex_type_row: int
    vertex_type_column: int
    adjacency_matrices: List[torch.Tensor]
    decoder_factory: Callable[[], DecodingMatrices]
    total_connectivity: torch.Tensor


class Data(object):
    vertex_types: List[VertexType]
    edge_types: List[EdgeType]

    def __init__(self, target_value: int = 1) -> None:
        self.vertex_types = []
        self.edge_types = {}
        self.target_value = int(target_value)

    def add_vertex_type(self, name: str, count: int) -> None:
        name = str(name)
        count = int(count)
        if not name:
            raise ValueError('You must provide a non-empty vertex type name')
        if count <= 0:
            raise ValueError('You must provide a positive vertex count')
        self.vertex_types.append(VertexType(name, count))

    def add_edge_type(self, name: str,
        vertex_type_row: int, vertex_type_column: int,
        adjacency_matrices: List[torch.Tensor],
        decoder_factory: Callable[[], DecodingMatrices]) -> None:

        name = str(name)
        vertex_type_row = int(vertex_type_row)
        vertex_type_column = int(vertex_type_column)

        if not isinstance(adjacency_matrices, list):
            raise TypeError('adjacency_matrices must be a list of tensors')

        if not callable(decoder_factory):
            raise TypeError('decoder_factory must be callable')

        if (vertex_type_row, vertex_type_column) in self.edge_types:
            raise KeyError('Edge type for given combination of row and column already exists')

        if vertex_type_row == vertex_type_column and \
            any(torch.any(_diag(adj_mat).to(torch.bool)) \
                for adj_mat in adjacency_matrices):
            raise ValueError('Adjacency matrices for same row/column vertex types must have empty diagonals')

        if any(adj_mat.shape[0] != self.vertex_types[vertex_type_row].count \
            or adj_mat.shape[1] != self.vertex_types[vertex_type_column].count \
                for adj_mat in adjacency_matrices):
            raise ValueError('Adjacency matrices must have as many rows as row vertex type count and as many columns as column vertex type count')

        total_connectivity = _nonzero_sum(adjacency_matrices)

        self.edge_types[vertex_type_row, vertex_type_column] = \
            EdgeType(name, vertex_type_row, vertex_type_column,
                adjacency_matrices, decoder_factory, total_connectivity)