init

Author: aaronzo

.gitignore

@@ -0,0 +1,73 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# poetry
+poetry.lock
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# ruff
+.ruff_cache

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Aaron Z.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,2 @@
+# graph-diffusers
+Diffusion patterns for graph machine learning

graph_diffusers/__init__.py

@@ -0,0 +1,7 @@
+from graph_diffusers import models, operators, io
+
+__all__ = [
+    "models",
+    "operators",
+    "io",
+]

graph_diffusers/_extras.py

@@ -0,0 +1,4 @@
+from importlib.util import find_spec
+
+TORCH = bool(find_spec("torch"))
+TORCH_SPARSE = bool(find_spec("torch_sparse"))

graph_diffusers/_typing.py

@@ -0,0 +1,23 @@
+from typing import Protocol
+from types import ModuleType
+
+import graphblas.binary
+import graphblas.core.matrix
+import graphblas.core.vector
+import graphblas.io
+import graphblas.select
+import graphblas.unary
+
+
+class _GraphblasModule(Protocol):
+    def __new__(cls) -> "_GraphblasModule":
+        raise NotImplementedError
+
+    class Matrix(graphblas.core.matrix.Matrix, Protocol): ...
+
+    class Vector(graphblas.core.vector.Vector, Protocol): ...
+
+    unary: ModuleType = graphblas.unary
+    binary: ModuleType = graphblas.binary
+    select: ModuleType = graphblas.select
+    io: ModuleType = graphblas.io

graph_diffusers/_utils.py

@@ -0,0 +1,27 @@
+import graphblas as gb
+from typing import TYPE_CHECKING, Optional, Union, Callable, Any
+
+if TYPE_CHECKING:
+    from graph_diffusers._typing import _GraphblasModule as gb
+    from graphblas.core.matrix import MatrixExpression, TransposedMatrix
+
+
+def eval_expr(expr: "MatrixExpression", out: Optional[gb.Matrix] = None, **kw: Any) -> gb.Matrix:
+    if out is None:
+        return expr.new(**kw)
+    out(**kw) << expr
+    return out
+
+
+def transpose_if(
+    condition: bool,
+    matrix: "Union[gb.Matrix, MatrixExpression]",
+) -> "Union[MatrixExpression, MatrixExpression, TransposedMatrix]":
+    return matrix.T if condition else matrix
+
+
+def _recover_triad_count(x: int) -> float:
+    return (1 - x) / 2 if x % 2 else x
+
+
+recover_triad_count: Callable[[gb.Matrix], "MatrixExpression"] = gb.unary.register_anonymous(_recover_triad_count)

graph_diffusers/io/__init__.py

@@ -0,0 +1,8 @@
+from graph_diffusers import _extras
+
+__all__ = []
+
+if _extras.TORCH:
+    from graph_diffusers.io import torch
+
+    __all__.append("torch")

graph_diffusers/io/torch.py

@@ -0,0 +1,88 @@
+import graphblas as gb
+from typing import TYPE_CHECKING, Optional, Union, Hashable
+import functools as ft
+import torch
+from torch_sparse import SparseTensor
+import warnings
+
+if TYPE_CHECKING:
+    from graph_diffusers._typing import _GraphblasModule as gb
+    from numpy.typing import DTypeLike
+
+# TODO check if torch_sparse installed
+
+
+def torch_to_graphblas(
+    edge_index: torch.Tensor,
+    *,
+    num_nodes: Optional[int] = None,
+    weighted: bool = False,
+    dtype: "Optional[DTypeLike]" = None,
+) -> gb.Matrix:
+    if isinstance(edge_index, SparseTensor):
+        return torch_sparse_tensor_to_graphblas(edge_index, weighted=weighted, dtype=dtype)
+    if edge_index.is_sparse_csr:
+        return torch_sparse_csr_to_graphblas(edge_index, weighted=weighted, dtype=dtype)
+    return torch_edge_index_to_graphblas(edge_index, num_nodes=num_nodes, dtype=dtype)
+
+
+def torch_sparse_csr_to_graphblas(
+    adj_t: torch.Tensor, *, weighted: bool = False, dtype: "Optional[DTypeLike]" = None
+) -> gb.Matrix:
+    if not isinstance(dtype, Hashable):
+        warnings.warn(
+            f"Unhashable dtype {dtype} passed when converting from torch to graphblas." "The result will not be cached."
+        )
+        return _torch_edge_index_to_graphblas.__wrapped__(adj_t, weighted=weighted, dtype=dtype)
+    return _torch_sparse_csr_to_graphblas(adj_t, weighted=weighted, dtype=dtype)
+
+
+def torch_sparse_tensor_to_graphblas(
+    adj_t: SparseTensor, *, weighted: bool = False, dtype: "Optional[DTypeLike]" = None
+) -> gb.Matrix:
+    return torch_sparse_csr_to_graphblas(
+        adj_t.to_torch_sparse_csr_tensor(),
+        weighted=weighted,
+        dtype=dtype,
+    )
+
+
+def torch_edge_index_to_graphblas(
+    edge_index: Union[torch.Tensor, SparseTensor],
+    *,
+    num_nodes: Optional[int] = None,
+    dtype: "Optional[DTypeLike]" = None,
+) -> gb.Matrix:
+    if not isinstance(dtype, Hashable):
+        warnings.warn(
+            f"Unhashable dtype {dtype} passed when converting from torch to graphblas." "The result will not be cached."
+        )
+        return _torch_edge_index_to_graphblas.__wrapped__(edge_index, num_nodes=num_nodes, dtype=dtype)
+    return _torch_edge_index_to_graphblas(edge_index, num_nodes=num_nodes, dtype=dtype)
+
+
+@ft.lru_cache(maxsize=1)
+def _torch_sparse_csr_to_graphblas(
+    adj_t: torch.Tensor,
+    weighted: bool,
+    dtype: "Optional[DTypeLike]",
+) -> gb.Matrix:
+    if not adj_t.is_sparse_csr:
+        adj_t = adj_t.to_sparse_csr()
+    return gb.Matrix.from_csr(
+        indptr=adj_t.crow_indices().detach().cpu().numpy(),
+        col_indices=adj_t.col_indices().detach().cpu().numpy(),
+        values=1.0 if not weighted else adj_t.values().detach().cpu().numpy(),
+        nrows=adj_t.shape[0],
+        ncols=adj_t.shape[0],
+        dtype=dtype,
+    )
+
+
+@ft.lru_cache(maxsize=1)
+def _torch_edge_index_to_graphblas(
+    edge_index: torch.Tensor,
+    num_nodes: Optional[int],
+    dtype: "Optional[DTypeLike]",
+) -> gb.Matrix:
+    return gb.Matrix.from_coo(*edge_index, dtype=dtype, nrows=num_nodes, ncols=num_nodes)

graph_diffusers/models.py

@@ -0,0 +1,71 @@
+import numpy as np
+from typing import Union, Literal, TYPE_CHECKING
+import graphblas as gb
+from itertools import chain
+from abc import ABCMeta, abstractmethod
+import graph_diffusers.operators as op
+
+if TYPE_CHECKING:
+    from graph_diffusers._typing import _GraphblasModule as gb
+
+
+class Diffusion(metaclass=ABCMeta):
+    @abstractmethod
+    def num_features(self, in_features: int) -> int: ...
+
+    @abstractmethod
+    def propagate(self, adj: gb.Matrix, X: np.ndarray) -> np.ndarray: ...
+
+
+class SimpleGCNDiffusion(Diffusion):
+    def __init__(self, k: int) -> None:
+        self.k = k
+
+    def num_features(self, in_features: int) -> int:
+        return in_features
+
+    def propagate(self, adj: gb.Matrix, X: np.ndarray) -> np.ndarray:
+        adj_gcn = op.gcn_norm(adj)
+        return op.power(adj_gcn, self.k)(X)
+
+
+class SIGNDiffusion(Diffusion):
+    def __init__(
+        self,
+        s: int,
+        p: int,
+        t: int,
+        s_norm: Union[Literal["gcn"], Literal["rw"]] = "gcn",
+        p_norm: Union[Literal["gcn"], Literal["rw"]] = "rw",
+        t_norm: Union[Literal["gcn"], Literal["rw"]] = "rw",
+    ) -> None:
+        self.s = s
+        self.p = p
+        self.t = t
+        self.s_norm = s_norm
+        self.p_norm = p_norm
+        self.t_norm = t_norm
+        if self.r < 1:
+            raise ValueError
+
+    @property
+    def r(self) -> int:
+        return self.s + self.p + self.t
+
+    def num_features(self, in_features: int) -> int:
+        return self.r * in_features
+
+    def propagate(self, adj: gb.Matrix, X: np.ndarray) -> np.ndarray:
+        ops = []
+        if s := self.s:
+            simple_diffuser = op.simple(op.norm(adj, method=self.s_norm))
+            ops.append(op.diffuse_powers(simple_diffuser, X, s))
+        if p := self.p:
+            ppr_diffuser = op.appnp(op.norm(adj, method=self.p_norm))
+            ops.append(op.diffuse_powers(ppr_diffuser, X, p))
+        if t := self.t:
+            adj_triangle = op.triangle(adj, directed=False)
+            triangle_diffuser = op.simple(op.norm(adj_triangle, method=self.t_norm, copy=False))
+            ops.append(op.diffuse_powers(triangle_diffuser, X, t))
+
+        return np.hstack(tuple(chain(*ops)), dtype=X.dtype)