flow_matrix.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. # Helpers for current-flow betweenness and current-flow closness
  2. # Lazy computations for inverse Laplacian and flow-matrix rows.
  3. import networkx as nx
  4. def flow_matrix_row(G, weight=None, dtype=float, solver="lu"):
  5. # Generate a row of the current-flow matrix
  6. import numpy as np
  7. solvername = {
  8. "full": FullInverseLaplacian,
  9. "lu": SuperLUInverseLaplacian,
  10. "cg": CGInverseLaplacian,
  11. }
  12. n = G.number_of_nodes()
  13. L = nx.laplacian_matrix(G, nodelist=range(n), weight=weight).asformat("csc")
  14. L = L.astype(dtype)
  15. C = solvername[solver](L, dtype=dtype) # initialize solver
  16. w = C.w # w is the Laplacian matrix width
  17. # row-by-row flow matrix
  18. for u, v in sorted(sorted((u, v)) for u, v in G.edges()):
  19. B = np.zeros(w, dtype=dtype)
  20. c = G[u][v].get(weight, 1.0)
  21. B[u % w] = c
  22. B[v % w] = -c
  23. # get only the rows needed in the inverse laplacian
  24. # and multiply to get the flow matrix row
  25. row = B @ C.get_rows(u, v)
  26. yield row, (u, v)
  27. # Class to compute the inverse laplacian only for specified rows
  28. # Allows computation of the current-flow matrix without storing entire
  29. # inverse laplacian matrix
  30. class InverseLaplacian:
  31. def __init__(self, L, width=None, dtype=None):
  32. global np
  33. import numpy as np
  34. (n, n) = L.shape
  35. self.dtype = dtype
  36. self.n = n
  37. if width is None:
  38. self.w = self.width(L)
  39. else:
  40. self.w = width
  41. self.C = np.zeros((self.w, n), dtype=dtype)
  42. self.L1 = L[1:, 1:]
  43. self.init_solver(L)
  44. def init_solver(self, L):
  45. pass
  46. def solve(self, r):
  47. raise nx.NetworkXError("Implement solver")
  48. def solve_inverse(self, r):
  49. raise nx.NetworkXError("Implement solver")
  50. def get_rows(self, r1, r2):
  51. for r in range(r1, r2 + 1):
  52. self.C[r % self.w, 1:] = self.solve_inverse(r)
  53. return self.C
  54. def get_row(self, r):
  55. self.C[r % self.w, 1:] = self.solve_inverse(r)
  56. return self.C[r % self.w]
  57. def width(self, L):
  58. m = 0
  59. for i, row in enumerate(L):
  60. w = 0
  61. x, y = np.nonzero(row)
  62. if len(y) > 0:
  63. v = y - i
  64. w = v.max() - v.min() + 1
  65. m = max(w, m)
  66. return m
  67. class FullInverseLaplacian(InverseLaplacian):
  68. def init_solver(self, L):
  69. self.IL = np.zeros(L.shape, dtype=self.dtype)
  70. self.IL[1:, 1:] = np.linalg.inv(self.L1.todense())
  71. def solve(self, rhs):
  72. s = np.zeros(rhs.shape, dtype=self.dtype)
  73. s = self.IL @ rhs
  74. return s
  75. def solve_inverse(self, r):
  76. return self.IL[r, 1:]
  77. class SuperLUInverseLaplacian(InverseLaplacian):
  78. def init_solver(self, L):
  79. import scipy as sp
  80. import scipy.sparse.linalg # call as sp.sparse.linalg
  81. self.lusolve = sp.sparse.linalg.factorized(self.L1.tocsc())
  82. def solve_inverse(self, r):
  83. rhs = np.zeros(self.n, dtype=self.dtype)
  84. rhs[r] = 1
  85. return self.lusolve(rhs[1:])
  86. def solve(self, rhs):
  87. s = np.zeros(rhs.shape, dtype=self.dtype)
  88. s[1:] = self.lusolve(rhs[1:])
  89. return s
  90. class CGInverseLaplacian(InverseLaplacian):
  91. def init_solver(self, L):
  92. global sp
  93. import scipy as sp
  94. import scipy.sparse.linalg # call as sp.sparse.linalg
  95. ilu = sp.sparse.linalg.spilu(self.L1.tocsc())
  96. n = self.n - 1
  97. self.M = sp.sparse.linalg.LinearOperator(shape=(n, n), matvec=ilu.solve)
  98. def solve(self, rhs):
  99. s = np.zeros(rhs.shape, dtype=self.dtype)
  100. s[1:] = sp.sparse.linalg.cg(self.L1, rhs[1:], M=self.M, atol=0)[0]
  101. return s
  102. def solve_inverse(self, r):
  103. rhs = np.zeros(self.n, self.dtype)
  104. rhs[r] = 1
  105. return sp.sparse.linalg.cg(self.L1, rhs[1:], M=self.M, atol=0)[0]