-
Notifications
You must be signed in to change notification settings - Fork 1
/
sparse.py
514 lines (424 loc) · 18 KB
/
sparse.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""JAX-compatible sparse positive definite symmetric matrices.
This library provides utilities for dealing with sparse positive definite
symmetric matrices with fixed sparsity structure. The main class is
SPDSparsityStructure, which delegates to sksparse to determine a sparse Cholesky
factorization and implements versions of all the methods we need for SPD
matrices (e.g. sparse cholesly factorization, triangular solves, diagonal
updates, log determinant) which operate on the underlying data vector of a
sparse SPD matrix with given sparsity structure.
SPDMatrix is a JAX-compatible class consisting of a (static)
SPDSparsityStructure together with the (dynamic) data vector of entries.
Important caveat: the sparse Cholesky factorization is implemented by simply
looping over the non-zero entries of the Cholesky triangle. So even though
methods using the SPDMatrix class are jittable, vmappable, differentiable, and
can run on accelerators, they're going to be slow. Replacing this loop with a
jax.lax.scan would be swell, but the naive way to do so loses the benefit of
sparsity.
"""
# We follow the naming convensions in sksparse.
# pylint:disable=invalid-name
from __future__ import annotations
import dataclasses
import functools
from typing import Any, Callable, Optional, Tuple
import jax
import jax.numpy as jnp
import numpy as np
import scipy
import scipy.sparse
import sksparse
import sksparse.cholmod
EdgeList = list[Tuple[Any, Any]]
UnaryOperation = Callable[[jnp.ndarray], jnp.ndarray]
BinaryOperation = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
NOT_POPULATED = -1
def data_index(mat: scipy.sparse.csc_matrix, i: int, j: int) -> int:
"""Index of mat[i, j] in mat.data, or NOT_POPULATED if mat[i, j] not set."""
assert mat.has_sorted_indices, 'Indices are not sorted.'
row_indices = mat.indices[mat.indptr[j] : mat.indptr[j + 1]]
k = np.searchsorted(row_indices, i)
if k < len(row_indices) and row_indices[k] == i:
return mat.indptr[j] + k
else:
return NOT_POPULATED
def row_col_index(mat: scipy.sparse.csc_matrix, k: int) -> Tuple[int, int]:
"""Returns (i, j) so that mat[i, j] = mat.data[k]."""
i = mat.indices[k]
j = np.searchsorted(mat.indptr, k, side='right') - 1
return i, j
def get_entries(data: jnp.ndarray, indices: jnp.ndarray) -> jnp.ndarray:
"""data[indices], except an index of NOT_POPULATED produces 0.0."""
return jnp.concatenate([data, jnp.array([0.0])])[indices]
def sparse_cholesky_fn(
Q: scipy.sparse.csc_matrix, factor: Optional[sksparse.cholmod.Factor] = None
):
"""Returns sparse Cholesky factorization function.
The returned method can be used for matrices with the same sparsity structure
as Q.
Args:
Q: a sparse symmetric positive definite matrix.
factor: a sksparse.cholmod.Factor representing the sparse Cholesky
factorization of Q.
Returns:
A function which, given a specification of the entries Q.data, returns the
values L.data of the Cholesky factorization.
"""
if factor is None:
factor = sksparse.cholmod.cholesky(Q)
L = factor.L()
P = factor.P()
# Determine the compressed sparse *row* representation of the Cholesky
# triangle so that we can efficiently compute dot products of rows as we fill
# in the columns of L.
L_row = scipy.sparse.csr_matrix(L)
# row_overlaps[(i, j)] is a pair of index arrays (ind1, ind2) recording the
# overlap of row i with row j. It has the property that
# L_row.data[ind1] @ L_row.data[ind2] + L[i, j] * L[j, j]
# computes the dot product of row i and row j.
# We compute the Cholesky factorization by setting this equal to Q[i, j] and
# solving for L[j, i]. We only populate the row overlaps that will be required
# to compute the Cholesky factorization.
row_overlaps = {}
for j in range(Q.shape[1]):
row1_start, row1_end = L_row.indptr[j], L_row.indptr[j + 1]
columns1 = L_row.indices[row1_start:row1_end]
for i in L.indices[L.indptr[j] : L.indptr[j + 1]]:
row2_start, row2_end = L_row.indptr[i], L_row.indptr[i + 1] - 1
columns2 = L_row.indices[row2_start:row2_end]
_, indices1, indices2 = np.intersect1d(
columns1, columns2, assume_unique=True, return_indices=True
)
row_overlaps[(i, j)] = (row1_start + indices1, row2_start + indices2)
# Q_indices[k] is the index of the element of Q.data which is in the same row
# and column as L.data[k] is within L.
Q_indices = []
for k in range(len(L.data)):
i, j = row_col_index(L, k)
Q_indices.append(data_index(Q, P[i], P[j]))
Q_indices = jnp.array(Q_indices)
def cholesky_factor(Q_data: jnp.ndarray) -> jnp.ndarray:
Q_entries = get_entries(Q_data, Q_indices)
L_data = jnp.zeros_like(L.data)
L_row_data = jnp.zeros_like(L.data)
# num_populated[i] tracks how many entries of row i been populated for
# purposes of determining which entry of L_row_data to set.
num_populated = [0 for _ in range(Q.shape[0])]
for j in range(Q.shape[1]):
# Compute the entries of L in column j.
col_start, col_end = L.indptr[j], L.indptr[j + 1]
row_start, row_end = L_row.indptr[j], L_row.indptr[j + 1]
diag_entry = jnp.sqrt(
Q_entries[col_start]
- L_row_data[row_start : row_end - 1]
@ L_row_data[row_start : row_end - 1]
)
col_entries = [diag_entry]
row_indices = [row_end - 1] # The diagonal entry is the last in its row.
for k in range(col_start + 1, col_end):
i = L.indices[k]
# Compute the entry in row i, column j
ind1, ind2 = row_overlaps[(i, j)]
row_dot_other_row = L_row_data[ind1] @ L_row_data[ind2]
entry = (Q_entries[k] - row_dot_other_row) / diag_entry
col_entries.append(entry)
other_row_start = L_row.indptr[i]
row_indices.append(other_row_start + num_populated[i])
num_populated[i] += 1
L_data = L_data.at[col_start:col_end].set(col_entries)
L_row_data = L_row_data.at[np.array(row_indices)].set(col_entries)
return L_data, L_row_data # pytype: disable=bad-return-type # jax-ndarray
return cholesky_factor
def is_lower_triangular(L: scipy.sparse.csc_matrix) -> bool:
# All indices in the i-th column must be at least i.
return all(
(L.indices[L.indptr[i] : L.indptr[i + 1]] >= i).all()
for i in range(L.shape[1])
)
def is_upper_triangular(U: scipy.sparse.csc_matrix) -> bool:
# All indices in the i-th column must be at most i.
return all(
(U.indices[U.indptr[i] : U.indptr[i + 1]] <= i).all()
for i in range(U.shape[1])
)
def tril_solver(
L: scipy.sparse.csc_matrix, L_T: Optional[scipy.sparse.csc_matrix] = None
) -> BinaryOperation:
"""Returns a method that solves sparse lower triangular linear systems."""
assert isinstance(L, scipy.sparse.csc_matrix)
assert (
L.shape[0] == L.shape[1]
), f'Only square (and full rank) matrices are supported. {L.shape=}'
assert is_lower_triangular(L), 'L must be lower triangular.'
assert all(
[L.indices[L.indptr[i]] == i for i in range(L.shape[0])]
), 'L must have populated diagonal entries'
if L_T is None:
L_T = scipy.sparse.csc_matrix(L.T)
else:
assert isinstance(L_T, scipy.sparse.csc_matrix)
def solve(L_T_data, b):
x = jnp.zeros(L.shape[0])
for i in range(L.shape[0]):
row_start, row_end = L_T.indptr[i], L_T.indptr[i + 1]
indices = L_T.indices[row_start : row_end - 1]
entry = (
b[i] - L_T_data[row_start : row_end - 1] @ x[indices]
) / L_T_data[row_end - 1]
x = x.at[i].set(entry)
return x
return solve
def triu_solver(
U: scipy.sparse.csc_matrix, U_T: Optional[scipy.sparse.csc_matrix] = None
) -> BinaryOperation:
"""Returns a method that solves sparse upper triangular linear systems."""
assert isinstance(U, scipy.sparse.csc_matrix)
assert (
U.shape[0] == U.shape[1]
), f'Only square (and full rank) matrices are supported. {U.shape=}'
assert is_upper_triangular(U), 'U must be upper triangular.'
assert all(
[U.indices[U.indptr[i + 1] - 1] == i for i in range(U.shape[0])]
), 'U must have populated diagonal entries'
if U_T is None:
U_T = scipy.sparse.csc_matrix(U.T)
else:
assert isinstance(U_T, scipy.sparse.csc_matrix)
def solve(U_T_data, b):
x = jnp.zeros(U.shape[0])
for i in reversed(range(U.shape[0])):
row_start, row_end = U_T.indptr[i], U_T.indptr[i + 1]
indices = U_T.indices[row_start + 1 : row_end]
entry = (
b[i] - U_T_data[row_start + 1 : row_end] @ x[indices]
) / U_T_data[row_start]
x = x.at[i].set(entry)
return x
return solve
def add_fn(
Q: scipy.sparse.csc_matrix,
) -> Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]:
"""Returns a method that adds a diagonal matrix diag(x) to Q."""
assert isinstance(Q, scipy.sparse.csc_matrix)
assert (
Q.shape[0] == Q.shape[1]
), f'Only square (and full rank) matrices are supported. {Q.shape=}'
Q_diag_inds = jnp.array([data_index(Q, i, i) for i in range(Q.shape[0])])
assert -1 not in Q_diag_inds, 'Diagonal entries of Q must be populated.'
def Q_add(Q_data, x):
Q_data = Q_data.at[Q_diag_inds].set(Q_data[Q_diag_inds] + x)
return Q_data
return Q_add
def multiply_fn(
Q: scipy.sparse.csc_matrix, Q_T: Optional[scipy.sparse.csc_matrix] = None
) -> BinaryOperation:
"""Returns a method that computes Q @ x, where x is a dense vector."""
assert isinstance(Q, scipy.sparse.csc_matrix)
assert (
Q.shape[0] == Q.shape[1]
), f'Only square (and full rank) matrices are supported. {Q.shape=}'
if Q_T is None:
Q_T = scipy.sparse.csc_matrix(Q.T)
def Q_multiply(Q_T_data, x):
y = []
for i in range(Q.shape[0]):
row_start, row_end = Q_T.indptr[i], Q_T.indptr[i + 1]
indices = Q_T.indices[row_start:row_end]
y.append(Q_T_data[row_start:row_end] @ x[indices])
return jnp.array(y)
return Q_multiply
def logdet_fn(L: scipy.sparse.csc_matrix) -> UnaryOperation:
"""Returns a method that computes the log determinant of L."""
assert isinstance(L, scipy.sparse.csc_matrix)
assert is_lower_triangular(L) or is_upper_triangular(
L
), 'L must be triangular'
assert (
L.shape[0] == L.shape[1]
), f'Only square (and full rank) matrices are supported. {L.shape=}'
L_diag_inds = jnp.array([data_index(L, i, i) for i in range(L.shape[0])])
def L_log_det(L_data):
L_diag = L_data[L_diag_inds]
return jnp.log(L_diag).sum()
return L_log_det
def node_names_and_integer_edges(edges: EdgeList) -> tuple[list[Any], EdgeList]:
node_names = sorted(set.union(*[set(e) for e in edges]))
int_edges = [(node_names.index(i), node_names.index(j)) for i, j in edges]
return node_names, int_edges
def adjacency_matrix(edges: EdgeList) -> scipy.sparse.csc_matrix:
nodes, edges = node_names_and_integer_edges(edges)
row_indices, col_indices = zip(*edges)
dim = len(nodes)
directed_adjacency_matrix = scipy.sparse.csc_matrix(
(np.ones_like(row_indices), (row_indices, col_indices)), shape=(dim, dim)
)
return directed_adjacency_matrix + directed_adjacency_matrix.T
def edge_list(mat: scipy.sparse.spmatrix) -> EdgeList:
"""An edge list matching the sparsity structure of `mat`."""
mat = mat.tocoo()
return list(set(tuple(sorted((r, c))) for r, c in zip(mat.row, mat.col)))
class SPDSparsityStructure:
"""Deals with data vectors related to a given sparsity structure."""
edges: EdgeList
# Runtime asserts can be turned off by setting this member to False. This is
# required to avoid ConcretizationTypeError when you apply jax.vmap.
do_asserts: bool
# Don't set these at initialization; they will be overwritten. These are just
# here to make pytype happy.
node_names: Optional[list[Any]] = None
int_edges: Optional[EdgeList] = None
dim: Optional[int] = None
Q: Optional[scipy.sparse.csc_matrix] = None
data_size: Optional[int] = None
L: Optional[scipy.sparse.csc_matrix] = None
L_T: Optional[scipy.sparse.csc_matrix] = None
P: Optional[np.ndarray] = None
Pinv: Optional[np.ndarray] = None
cholesky_factor: Optional[UnaryOperation] = None
solve_L: Optional[BinaryOperation] = None
solve_LT: Optional[BinaryOperation] = None
multiply_Q: Optional[BinaryOperation] = None
add_diag_Q: Optional[BinaryOperation] = None
logdet_L: Optional[UnaryOperation] = None
def __init__(self, edges: EdgeList):
self.edges = edges
self.do_asserts = True
self.node_names, self.int_edges = node_names_and_integer_edges(self.edges)
self.dim = len(self.node_names)
# Note: the diagonal is multiplied by self.dim to ensure positive
# definiteness so that cholesky factorization doesn't crash.
self.Q = adjacency_matrix(self.int_edges) + self.dim * scipy.sparse.eye(
self.dim
)
self.data_size = len(self.Q.data)
Q_coo = self.Q.tocoo()
self.transpose_indices = np.array(
[data_index(self.Q, row, col) for row, col in zip(Q_coo.col, Q_coo.row)]
)
factor = sksparse.cholmod.cholesky(self.Q)
self.L = factor.L()
self.L_T = scipy.sparse.csc_matrix(self.L.T)
self.P = factor.P()
self.Pinv = np.argsort(self.P)
self.cholesky_factor = sparse_cholesky_fn(self.Q, factor=factor)
self.solve_L = tril_solver(L=self.L, L_T=self.L_T)
self.solve_LT = triu_solver(U=self.L_T, U_T=self.L)
self.multiply_Q = multiply_fn(self.Q, Q_T=self.Q)
self.add_diag_Q = add_fn(self.Q)
self.logdet_L = logdet_fn(self.L)
# WARNING: this method is likely to disappear or change behavior! I'm not
# sure what the best way to handle this is. The sparsity structure gets passed
# around a lot (e.g. it's inherited by all SPDMatrices output by a given
# GMRF), so it's worth making its methods fast.
def jit_me(self) -> None:
for method in [
'cholesky_factor',
'solve_L',
'solve_LT',
'multiply_Q',
'add_diag_Q',
'logdet_L',
]:
setattr(self, method, jax.jit(getattr(self, method)))
def __repr__(self) -> str:
return f'SPDSparsityStructure(edges={self.edges})'
def transpose(self, data: jnp.ndarray) -> jnp.ndarray:
return data[self.transpose_indices]
def data_from_scipy(self, mat: scipy.sparse.spmatrix) -> jnp.ndarray:
"""Data vector for `mat`, aligned with Q.data."""
data = np.zeros_like(self.Q.data)
mat = mat.tocoo()
for value, row, col in zip(mat.data, mat.row, mat.col):
idx = data_index(self.Q, row, col)
if idx == NOT_POPULATED:
raise ValueError(
f'Entry ({row}, {col}) is populated in mat, but edge '
'is not in the graph.'
)
data[idx] = value
return jnp.array(data)
def data_to_scipy(self, data: jnp.ndarray) -> scipy.sparse.csc_matrix:
if self.do_asserts:
assert (
len(data) == self.data_size
), f'Wrong data size ({len(data)} vs {self.data_size})'
indptr = self.Q.indptr
indices = self.Q.indices
dim = self.dim
return scipy.sparse.csc_matrix((data, indices, indptr), shape=(dim, dim))
def to_spd_matrix(self, data: jnp.ndarray) -> SPDMatrix:
spd_matrix = SPDMatrix(self, data)
if self.do_asserts:
assert (
len(data) == self.data_size
), f'Wrong data size ({len(data)} vs {self.data_size})'
assert jnp.all(
self.transpose(data) == data
), 'Data vector does not correspond to a symmetric matrix.'
assert not jnp.isnan(
spd_matrix.logdet
), 'Data vector does not correspond to a positive definite matrix.'
return spd_matrix
@jax.tree_util.register_pytree_node_class
@dataclasses.dataclass
class SPDMatrix:
"""A sparse positive definite matrix."""
sparsity_structure: SPDSparsityStructure
data: jnp.ndarray
def __matmul__(self, vec: jnp.ndarray) -> jnp.ndarray:
return self.sparsity_structure.multiply_Q(self.data, vec)
def add_diag(self, vec: jnp.ndarray) -> SPDMatrix:
return SPDMatrix(
self.sparsity_structure,
self.sparsity_structure.add_diag_Q(self.data, vec),
)
def to_scipy(self) -> scipy.sparse.csc_matrix:
return self.sparsity_structure.data_to_scipy(self.data)
@classmethod
def from_scipy(cls, mat: scipy.sparse.spmatrix) -> SPDMatrix:
mat = mat.tocoo()
edges = edge_list(mat)
sparsity_structure = SPDSparsityStructure(edges)
data = sparsity_structure.data_from_scipy(mat)
return SPDMatrix(sparsity_structure, data)
def solve_LT(self, epsilon: jnp.ndarray) -> jnp.ndarray:
"""Solves L.T @ x = epsilon, where Q = L @ L.T."""
L_data, _ = self._cholesky_factorization
x = self.sparsity_structure.solve_LT(L_data, epsilon)
return x[self.sparsity_structure.Pinv]
def solve(self, b: jnp.ndarray) -> jnp.ndarray:
"""Solves for Q @ x = b."""
L_data, LT_data = self._cholesky_factorization
v = self.sparsity_structure.solve_L(LT_data, b[self.sparsity_structure.P])
x = self.sparsity_structure.solve_LT(L_data, v)
return x[self.sparsity_structure.Pinv]
@functools.cached_property
def _cholesky_factorization(self) -> Tuple[jnp.ndarray, jnp.ndarray]:
return self.sparsity_structure.cholesky_factor(self.data)
@functools.cached_property
def logdet(self) -> jnp.ndarray:
L_data, _ = self._cholesky_factorization
return 2.0 * self.sparsity_structure.logdet_L(L_data)
@property
def shape(self) -> Tuple[int, int]:
return (self.sparsity_structure.dim, self.sparsity_structure.dim)
# Methods below for JAX pytree compatibility.
def tree_flatten(self):
dynamic_values = (self.data,)
static_values = (self.sparsity_structure,)
return dynamic_values, static_values
@classmethod
def tree_unflatten(cls, static_values, dynamic_values):
return cls(*static_values, *dynamic_values)