Class AbstractGraphTransform

climart/data_transform/transforms.py:122–171

# -------------------------------------------- Graph specific transforms
class AbstractGraphTransform(AbstractTransform, ABC):
    def __init__(self, exp_type: str):
        super().__init__(exp_type)

    @property
    def n_nodes(self) -> int:
        return self._n_nodes

    def get_edge_idxs(self) -> Tuple[np.ndarray, np.ndarray]:
        # one-way: node i has an edge to node i+1
        one_way_senders = np.arange(self.n_nodes - 1)  # n_lay - 1 = n_lev - 2 edges
        one_way_receivers = one_way_senders + 1
        # one-way: node i+1 has an edge to node i
        other_way_senders = np.arange(1, self.n_nodes)
        other_way_receivers = other_way_senders - 1

        senders = np.concatenate((one_way_senders, other_way_senders))
        receivers = np.concatenate((one_way_receivers, other_way_receivers))
        return senders, receivers

    def get_adj(self, degree_normalized: bool = False, improved: bool = False) -> Tensor:
        """ Adjacency matrix of a line graph with global node and self-loops """
        adj = torch.zeros((self.n_nodes, self.n_nodes))

        if hasattr(self, 'GLOBAL_NODE'):
            adj[:, self.GLOBAL_NODE] = 1
            adj[self.GLOBAL_NODE, :] = 1

        for i in range(1, self.n_nodes):
            adj[i, i - 1:i + 2] = 1
            adj[i - 1:i + 2, i] = 1

        if degree_normalized:
            self.log.info("-> Adjacency matrix is normalized by in-degree")
            return normalize_adjacency_matrix_torch(adj, improved=improved, add_self_loops=True)
        return adj

    def get_level_nodes_mask(self) -> Tensor:
        """
        Indexing an array 'A' with shape (n, self.n_nodes, out-dim)
        leads to A[:, get_level_nodes_mask(), :] have shape (n, len(self.LEVEL_NODES), out-dim)
        :return: A torch.Tensor mask that indexes all the level nodes
        """

        if not hasattr(self, 'LEVEL_NODES'):
            raise ValueError(f"This transform {self} does not have level nodes.")
        else:
            level_mask = torch.zeros(self.n_nodes)
            level_mask[self.LEVEL_NODES] = 1
        return level_mask.bool()


class OnlyLayersAreNodesTransforms(AbstractGraphTransform, ABC):