google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Utilities for data loading and preprocessing."""
17

18
from typing import Any
19
from typing import Callable
20
import numpy as np
21
import scipy.io
22
import scipy.sparse as sps
23

24

25
def preprocess_adjacency(
26
    adjacency,
27
    *,
28
    convert_to_csr = True,
29
    convert_to_unweighted = True,
30
    remove_self_loops = True,
31
    remove_isolated_nodes = True,
32
):
33
  """Pre-processes input adjacency matrix.
34

35
  Args:
36
    adjacency: Input adjacency matrix.
37
    convert_to_csr: Whether to convert the input matrix to the CSR format with
38
      fast matrix multiplications.
39
    convert_to_unweighted: Whether to discard input weights.
40
    remove_self_loops: Whether to remove self-loops from the graph.
41
    remove_isolated_nodes: Whether to remove isolated nodes from the graph.
42

43
  Returns:
44
    Clean adjacency matrix.
45
  """
46
  if adjacency.ndim != 2:
47
    raise ValueError(
48
        f'Adjacency matrix should be a 2D tensor, got {adjacency.ndim}'
49
    )
50
  if adjacency.shape[0] != adjacency.shape[1]:
51
    raise ValueError(
52
        f'Adjacency matrix should be square, got {adjacency.shape}'
53
    )
54
  if convert_to_csr:
55
    adjacency = adjacency.tocsr()
56
  if convert_to_unweighted:
57
    adjacency.data = np.ones_like(adjacency.data)
58
  if remove_self_loops:
59
    adjacency = adjacency - sps.diags(adjacency.diagonal())
60
  if remove_isolated_nodes:
61
    nonzero_rows = (adjacency.sum(0) != 0).A1
62
    adjacency = adjacency[nonzero_rows, :][:, nonzero_rows]
63
  return adjacency
64

65

66
def load_matfile(
67
    filepath,
68
    matfile_variable_name = 'network',
69
    convert_to_unweighted = True,
70
    open_fn = open,
71
):
72
  with open_fn(filepath, 'rb') as inf:
73
    data = scipy.io.loadmat(inf)
74
    adjacency = data[matfile_variable_name].tocsr()
75
  if convert_to_unweighted:
76
    adjacency.data = np.ones_like(adjacency.data)
77
  return adjacency
78