Graph operators

src/beignet/__init__.py

@@ -358,6 +358,16 @@
     trim_probabilists_hermite_polynomial_coefficients,
 )
 from .special import error_erf, error_erfc
+from ._validate_graph_matrix import validate_graph_matrix
+from ._tensor_to_masked_graph_matrix import tensor_to_masked_graph_matrix
+from ._tensor_to_graph_matrix import tensor_to_graph_matrix
+from ._masked_tensor_to_graph_matrix import masked_tensor_to_graph_matrix
+from ._graph_matrix_to_tensor import graph_matrix_to_tensor
+from ._graph_matrix_to_masked_tensor import graph_matrix_to_masked_tensor
+from ._predecessor_matrix_to_distance_matrix import predecessor_matrix_to_distance_matrix
+from ._reconstruct_path import reconstruct_path
+
+
 
 __all__ = [
     "add_chebyshev_polynomial",
@@ -453,6 +463,8 @@
     "gauss_legendre_quadrature",
     "gauss_physicists_hermite_polynomial_quadrature",
     "gauss_probabilists_hermite_polynomial_quadrature",
+    "graph_matrix_to_masked_tensor",
+    "graph_matrix_to_tensor",
     "integrate_chebyshev_polynomial",
     "integrate_laguerre_polynomial",
     "integrate_legendre_polynomial",
@@ -497,6 +509,7 @@
     "linear_physicists_hermite_polynomial",
     "linear_polynomial",
     "linear_probabilists_hermite_polynomial",
+    "masked_tensor_to_graph_matrix",
     "multiply_chebyshev_polynomial",
     "multiply_chebyshev_polynomial_by_x",
     "multiply_laguerre_polynomial",
@@ -538,6 +551,7 @@
     "polynomial_vandermonde_3d",
     "polynomial_x",
     "polynomial_zero",
+    "predecessor_matrix_to_distance_matrix",
     "probabilists_hermite_polynomial_companion",
     "probabilists_hermite_polynomial_domain",
     "probabilists_hermite_polynomial_from_roots",
@@ -562,6 +576,7 @@
     "random_quaternion",
     "random_rotation_matrix",
     "random_rotation_vector",
+    "reconstruct_path",
     "rotation_matrix_identity",
     "rotation_matrix_magnitude",
     "rotation_matrix_mean",
@@ -580,11 +595,14 @@
     "subtract_physicists_hermite_polynomial",
     "subtract_polynomial",
     "subtract_probabilists_hermite_polynomial",
+    "tensor_to_graph_matrix",
+    "tensor_to_masked_graph_matrix",
     "translation_identity",
     "trim_chebyshev_polynomial_coefficients",
     "trim_laguerre_polynomial_coefficients",
     "trim_legendre_polynomial_coefficients",
     "trim_physicists_hermite_polynomial_coefficients",
     "trim_polynomial_coefficients",
     "trim_probabilists_hermite_polynomial_coefficients",
+    "validate_graph_matrix",
 ]

src/beignet/_floyd_warshall.py

@@ -0,0 +1,61 @@
+import torch
+from torch import Tensor
+
+
+def floyd_warshall(
+    input: Tensor,
+    directed: bool = True,
+    unweighted: bool = False,
+):
+    r"""
+    ...
+
+    Parameters
+    ----------
+    input : Tensor, shape=(..., N, N)
+        ...
+
+    directed : bool
+        If `False`, symmetrizes `input`. Default, `True`.
+
+    unweighted : bool
+        If `True`, distance of non-zero connections is 1. Default, `False`.
+
+    Returns
+    -------
+    output : Tensor, shape=(..., N, N)
+        ...
+    """
+    output = input.clone()
+
+    if not directed:
+        output = 0.5 * (output + output.transpose(-1, -2))
+
+    if unweighted:
+        output = torch.where(
+            output != 0,
+            torch.ones_like(output),
+            torch.zeros_like(output),
+        )
+
+    n = output.shape[-1]
+
+    eye = torch.eye(n, device=output.device, dtype=output.dtype)
+
+    eye = torch.expand_copy(eye, output.shape)
+
+    output[((output == 0) & (~eye.bool()))] = torch.inf
+
+    output = torch.where(
+        eye.to(dtype=torch.bool),
+        torch.zeros_like(output),
+        output,
+    )
+
+    for k in range(n):
+        a = torch.unsqueeze(output[..., :, k], dim=-1)
+        b = torch.unsqueeze(output[..., k, :], dim=-2)
+
+        output = torch.minimum(output, a + b)
+
+    return output

src/beignet/_graph_matrix_to_masked_tensor.py

@@ -0,0 +1,11 @@
+import numpy
+from numpy.ma import MaskedArray
+from scipy.sparse import csr_matrix
+
+from ._graph_matrix_to_tensor import graph_matrix_to_tensor
+
+
+def graph_matrix_to_masked_tensor(input: csr_matrix) -> MaskedArray:
+    output = graph_matrix_to_tensor(input, numpy.nan)
+
+    return numpy.ma.masked_invalid(output)

src/beignet/_graph_matrix_to_tensor.py

@@ -0,0 +1,54 @@
+import numpy
+import scipy
+import scipy.sparse
+from scipy.sparse import csr_matrix
+
+def _populate_graph(data, indices, indptr, graph, null_value):
+    N = graph.shape[0]
+    null_flag = numpy.ones((N, N), dtype=bool, order='C')
+
+    # Calculate the number of non-zero entries per row
+    row_counts = indptr[1:] - indptr[:-1]
+
+    # Generate row indices for all non-zero entries
+    rows = numpy.repeat(numpy.arange(N), row_counts)
+
+    # Update null_flag to mark positions that have edges
+    null_flag[rows, indices] = False
+
+    # Update the graph with the minimum values for each edge
+    graph[rows, indices] = numpy.minimum(data, graph[rows, indices])
+
+    # Assign null_value to positions with no edges
+    graph[null_flag] = null_value
+
+def graph_matrix_to_tensor(
+        input: csr_matrix,
+        null_value: float = 0,
+) -> numpy.ndarray:
+    if not scipy.sparse.issparse(input):
+        raise ValueError
+
+    if input.format not in {"lil", "csc", "csr"}:
+        raise ValueError
+
+    input = input.tocsr()
+
+    n = input.shape[0]
+
+    if input.shape[1] != n:
+        raise ValueError
+
+    data = numpy.asarray(input.data, dtype=numpy.float64, order="C")
+
+    indices = numpy.asarray(input.indices, dtype=numpy.int32, order="C")
+
+    indptr = numpy.asarray(input.indptr, dtype=numpy.int32, order="C")
+
+    output = numpy.empty(input.shape, dtype=numpy.float64)
+
+    output.fill(numpy.inf)
+
+    _populate_graph(data, indices, indptr, output, null_value)
+
+    return output

src/beignet/_masked_tensor_to_graph_matrix.py

@@ -0,0 +1,33 @@
+import numpy
+from numpy.ma import MaskedArray
+from scipy.sparse import csr_matrix
+
+
+def masked_tensor_to_graph_matrix(input: MaskedArray) -> csr_matrix:
+    input = numpy.ma.asarray(input)
+
+    if input.ndim != 2:
+        raise ValueError
+
+    n = input.shape[0]
+
+    if input.shape[1] != n:
+        raise ValueError
+
+    compressed = input.compressed()
+
+    mask = ~input.mask
+
+    compressed = numpy.asarray(compressed, dtype=numpy.int32, order="c")
+
+    idx_grid = numpy.empty((n, n), dtype=numpy.int32)
+
+    idx_grid[:] = numpy.arange(n, dtype=numpy.int32)
+
+    indices = numpy.asarray(idx_grid[mask], dtype=numpy.int32, order="c")
+
+    indptr = numpy.zeros(n + 1, dtype=numpy.int32)
+
+    indptr[1:] = mask.sum(1).cumsum()
+
+    return csr_matrix((compressed, indices, indptr), (n, n))

src/beignet/_predecessor_matrix_to_distance_matrix.py

@@ -0,0 +1,79 @@
+import numpy
+import torch
+from scipy.sparse import csr_matrix
+
+from ._validate_graph_matrix import validate_graph_matrix
+
+NULL_IDX = -9999
+
+def _predecessor_matrix_to_distance_matrix(
+        input: numpy.ndarray,
+        predecessor_matrix: numpy.ndarray,
+        distance_matrix: numpy.ndarray,
+        directed: bool,
+        null_value: float,
+):
+    n = input.shape[0]
+
+    # symmetrize matrix, if necessary
+    if not directed:
+        input[input == 0] = numpy.inf
+
+        for i in range(n):
+            for j in range(i + 1, n):
+                if input[j, i] <= input[i, j]:
+                    input[i, j] = input[j, i]
+                else:
+                    input[j, i] = input[i, j]
+
+    for i in range(n):
+        for j in range(n):
+            null_path = True
+
+            k2 = j
+
+            while k2 != i:
+                k1 = predecessor_matrix[i, k2]
+
+                if k1 == NULL_IDX:
+                    break
+
+                distance_matrix[i, j] += input[k1, k2]
+
+                null_path = False
+
+                k2 = k1
+
+            if null_path and i != j:
+                distance_matrix[i, j] = null_value
+
+def predecessor_matrix_to_distance_matrix(
+        input: numpy.ndarray | csr_matrix,
+        predecessor_matrix: numpy.ndarray,
+        directed: bool = True,
+        null_value: float = numpy.inf,
+) -> numpy.ndarray:
+    input = validate_graph_matrix(
+        input,
+        directed,
+        dtype=torch.float64,
+        csr_output=False,
+        copy_if_dense=not directed,
+    )
+
+    predecessor_matrix = numpy.asarray(predecessor_matrix)
+
+    if predecessor_matrix.shape != input.shape:
+        raise ValueError
+
+    distance_matrix = numpy.zeros(input.shape, dtype=numpy.float64)
+
+    _predecessor_matrix_to_distance_matrix(
+        input,
+        predecessor_matrix,
+        distance_matrix,
+        directed,
+        null_value,
+    )
+
+    return distance_matrix

src/beignet/_reconstruct_path.py

@@ -0,0 +1,47 @@
+import numpy
+import scipy
+import scipy.sparse
+from scipy.sparse import csr_matrix
+
+from ._validate_graph_matrix import validate_graph_matrix
+
+
+def reconstruct_path(
+        input: numpy.ndarray | csr_matrix,
+        predecessors: numpy.ndarray,
+        directed: bool = True,
+) -> csr_matrix:
+    input = validate_graph_matrix(input, directed, dense_output=False)
+
+    n = input.shape[0]
+
+    nnull = (predecessors < 0).sum()
+
+    indices = numpy.argsort(predecessors)[nnull:].astype(numpy.int32)
+
+    pind = predecessors[indices]
+
+    indptr = pind.searchsorted(numpy.arange(n + 1)).astype(numpy.int32)
+
+    data = input[pind, indices]
+
+    if scipy.sparse.issparse(data):
+        data = data.todense()
+
+    data = data.getA1()
+
+    if not directed:
+        data2 = input[indices, pind]
+
+        if scipy.sparse.issparse(data2):
+            data2 = data2.todense()
+
+        data2 = data2.getA1()
+
+        data[data == 0] = numpy.inf
+
+        data2[data2 == 0] = numpy.inf
+
+        data = numpy.minimum(data, data2)
+
+    return csr_matrix((data, indices, indptr), shape=(n, n))

src/beignet/_tensor_to_graph_matrix.py

@@ -0,0 +1,21 @@
+import numpy
+from scipy.sparse import csr_matrix
+
+from ._tensor_to_masked_graph_matrix import tensor_to_masked_graph_matrix
+from ._masked_tensor_to_graph_matrix import masked_tensor_to_graph_matrix
+
+
+def tensor_to_graph_matrix(
+        input: numpy.ndarray,
+        null_value: float = 0.0,
+        nan_is_null_edge: bool = True,
+        infinity_is_null_edge: bool = True,
+) -> csr_matrix:
+    output = tensor_to_masked_graph_matrix(
+        input,
+        null_value,
+        nan_is_null_edge,
+        infinity_is_null_edge,
+    )
+
+    return masked_tensor_to_graph_matrix(output)