Expose sparse matrix normalization (#24)

* extract utility to normalize a sparse matrix
* add test for sparse matrix normalization utility
* remove unused flake8 ignores & fix comments
* update GHA to test on torch 1.13 too

Author: mberr

.github/workflows/tests.yml

@@ -55,7 +55,7 @@ jobs:
       matrix:
         os: [ "ubuntu-latest" ]
         python-version: [ "3.8", "3.9", "3.10" ]
-        torch-version: [ "torch-1.11", "torch-1.12" ]
+        torch-version: [ "torch-1.11", "torch-1.12", "torch-1.13" ]
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v2

setup.cfg

@@ -131,12 +131,10 @@ strictness = short
 #########################
 [flake8]
 ignore =
-    S301 # pickle
-    S403 # pickle
-    S404
-    S603
-    W503 # Line break before binary operator (flake8 is wrong)
-    E203  # whitespace before ':'
+    # Line break before binary operator (flake8 is wrong)
+    W503
+    # whitespace before ':'
+    #E203
 exclude =
     .tox,
     .git,

src/torch_ppr/utils.py

@@ -17,6 +17,7 @@
     "prepare_x0",
     "power_iteration",
     "batched_personalized_page_rank",
+    "sparse_normalize",
 ]
 
 logger = logging.getLogger(__name__)
@@ -172,6 +173,33 @@ def sparse_diagonal(values: torch.Tensor) -> torch.Tensor:
     )
 
 
+def sparse_normalize(matrix: torch.Tensor, dim: int = 0) -> torch.Tensor:
+    """
+    Normalize a sparse matrix to row/column sum of 1.
+
+    :param matrix:
+        the sparse matrix
+    :param dim:
+        the dimension along which to normalize, either 0 for rows or 1 for columns
+
+    :return:
+        the normalized sparse matrix
+    """
+    # calculate row/column sum
+    row_or_column_sum = (
+        torch.sparse.sum(matrix, dim=dim).to_dense().clamp_min(min=torch.finfo(matrix.dtype).eps)
+    )
+    # invert and create diagonal matrix
+    scaling_matrix = sparse_diagonal(values=torch.reciprocal(row_or_column_sum))
+    # multiply matrix
+    if dim == 0:
+        args = (matrix, scaling_matrix)
+    else:
+        args = (scaling_matrix, matrix)
+    # note: we do not pass by keyword due to instable API
+    return torch.sparse.mm(*args)
+
+
 def prepare_page_rank_adjacency(
     adj: Optional[torch.Tensor] = None,
     edge_index: Optional[torch.LongTensor] = None,
@@ -219,12 +247,8 @@ def prepare_page_rank_adjacency(
     if add_identity:
         adj = adj + sparse_diagonal(torch.ones(adj.shape[0], dtype=adj.dtype, device=adj.device))
 
-    # adjacency normalization: normalize to col-sum = 1
-    degree_inv = torch.reciprocal(
-        torch.sparse.sum(adj, dim=0).to_dense().clamp_min(min=torch.finfo(adj.dtype).eps)
-    )
-    degree_inv = sparse_diagonal(values=degree_inv)
-    return torch.sparse.mm(adj, degree_inv)
+    # adjacency normalization: normalize to row-sum = 1
+    return sparse_normalize(matrix=adj, dim=0)
 
 
 def validate_x(x: torch.Tensor, n: Optional[int] = None) -> None:
@@ -259,7 +283,9 @@ def validate_x(x: torch.Tensor, n: Optional[int] = None) -> None:
 
 
 def prepare_x0(
-    x0: Optional[torch.Tensor] = None, indices: Collection[int] = None, n: Optional[int] = None
+    x0: Optional[torch.Tensor] = None,
+    indices: Optional[Collection[int]] = None,
+    n: Optional[int] = None,
 ) -> torch.Tensor:
     """
     Prepare a start value.

tests/test_utils.py

@@ -181,3 +181,19 @@ def test_sparse_diagonal(n: int):
     assert matrix.shape == (n, n)
     assert matrix.is_sparse
     assert torch.allclose(matrix.to_dense(), torch.diag(values))
+
+
+@pytest.mark.parametrize("seed", [21, 42, 63])
+def test_sparse_normalize(seed: int):
+    """Test for sparse matrix normalization."""
+    generator = torch.manual_seed(seed=seed)
+    n_rows, n_cols = torch.randint(10, 20, size=(2,), generator=generator)
+    matrix = torch.rand(size=(n_rows, n_cols), generator=generator)
+    # make sparse
+    matrix[matrix < 0.5] = 0
+    matrix = matrix.to_sparse()
+    # normalize
+    for dim in (0, 1):
+        matrix_norm = utils.sparse_normalize(matrix=matrix, dim=dim)
+        sparse_sum = torch.sparse.sum(matrix_norm, dim=dim)
+        assert torch.allclose(sparse_sum.values(), torch.ones_like(sparse_sum.values()))

tox.ini

@@ -23,7 +23,7 @@ envlist =
     docstr-coverage
     docs-test
     # the actual tests
-    py-torch-{1.11,1.12}
+    py-torch-{1.11,1.12,1.13}
     # always keep coverage-report last
     # coverage-report
 
@@ -39,6 +39,7 @@ setenv =
 deps =
     torch-1.11: torch~=1.11.0
     torch-1.12: torch~=1.12.0
+    torch-1.13: torch~=1.13.0
 extras =
     # See the [options.extras_require] entry in setup.cfg for "tests"
     tests