Coverage for jetgp/full_degp_sparse/sparse_cholesky.py: 48%
296 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-04-09 18:57 -0500
1"""
2Sparse Cholesky utilities using the paper's geometric sparsity criterion.
4Sparsity pattern:
5 dist(x_P(i), x_P(j)) <= rho * l(j)
7where P is the MMD (maximin) ordering and l(j) is the fill-distance at step j.
9Key property: the sparsity pattern depends only on the training points and rho,
10NOT on the kernel hyperparameters. It is therefore computed once at model
11initialisation and reused across all NLML evaluations during optimisation.
13References
14----------
15Schäfer et al. (2021), "Sparse Cholesky factorization by Kullback-Leibler
16minimization", SIAM Journal on Scientific Computing.
17"""
19import numpy as np
20from scipy.linalg import lu_factor, lu_solve, cho_factor, cho_solve, solve_triangular
21from scipy.spatial.distance import cdist
22from line_profiler import profile
25# =============================================================================
26# Maximin (MMD) ordering
27# =============================================================================
29def mmd_ordering(X):
30 """
31 Compute the maximin ordering of training points.
33 The first point is fixed at index 0. Each subsequent point is chosen as
34 the one farthest from all already-selected points (maximum of minimum
35 distances). Also returns the fill-distances l[j], where l[j] is the
36 minimum distance from X[P[j]] to all previously selected points.
38 Parameters
39 ----------
40 X : ndarray of shape (N, d)
41 Training input points.
43 Returns
44 -------
45 P : ndarray of int, shape (N,)
46 Permutation indices (MMD ordering). P[0] is always 0.
47 l : ndarray of float, shape (N,)
48 Fill-distances. l[0] = 0 by convention; l[j] > 0 for j >= 1.
49 """
50 N = len(X)
51 P = np.zeros(N, dtype=int)
52 l = np.zeros(N)
53 rem = list(range(1, N))
54 for q in range(1, N):
55 dists = cdist(X[rem], X[P[:q]]).min(axis=1)
56 best = int(np.argmax(dists))
57 P[q] = rem[best]
58 l[q] = dists[best]
59 rem.pop(best)
60 return P, l
63# =============================================================================
64# Sparsity pattern (paper's geometric criterion)
65# =============================================================================
67def build_sparsity_pattern(X_ord, l, rho):
68 """
69 Build the sparsity pattern using the paper's criterion.
71 For column j, the neighbour set S[j] contains j itself plus all i < j
72 such that dist(X_ord[i], X_ord[j]) <= rho * l[j].
74 Because the criterion depends only on (X_ord, l, rho) and NOT on the
75 kernel hyperparameters, this can be computed once and reused.
77 Parameters
78 ----------
79 X_ord : ndarray of shape (N, d)
80 Training points reordered by MMD (i.e. X[P]).
81 l : ndarray of float, shape (N,)
82 Fill-distances from mmd_ordering.
83 rho : float
84 Sparsity radius multiplier. Larger rho → denser pattern.
86 Returns
87 -------
88 S : dict
89 S[j] = sorted list of row indices (neighbours) for column j,
90 always including j.
91 """
92 N = len(X_ord)
93 S = {}
94 for j in range(N):
95 nb = [j]
96 if l[j] > 0:
97 for i in range(j):
98 if np.linalg.norm(X_ord[i] - X_ord[j]) <= rho * l[j]:
99 nb.append(i)
100 S[j] = sorted(nb)
101 return S
104# =============================================================================
105# Sparse inverse-Cholesky factor U (column-by-column)
106# =============================================================================
108def build_U(K_ord, S, N, block_size=0, out=None):
109 """
110 Build the sparse inverse-Cholesky factor U column by column.
112 For each column j, we solve the small local system K_sub * v = e_j
113 where K_sub = K_ord[np.ix_(S[j], S[j])] and e_j is the unit vector
114 for the diagonal entry. The column of U is then v / sqrt(v[jj]).
116 When block_size > 0, consecutive columns are grouped into blocks.
117 Within each block the sparsity patterns are nested (S[j] ⊂ S[j+1]),
118 so the last column's pattern is the union. We factor that union
119 K_sub once via Cholesky and batch-solve all RHS vectors at once.
121 Parameters
122 ----------
123 K_ord : ndarray of shape (N, N)
124 Kernel matrix in MMD order.
125 S : dict
126 Sparsity pattern from build_sparsity_pattern.
127 N : int
128 Number of training points.
129 block_size : int
130 If > 0, batch columns in groups of this size (typically n_dim + 1).
131 Set to 0 to use the original column-by-column path.
132 out : ndarray of shape (N, N), optional
133 Pre-allocated output buffer. Zeroed before use.
135 Returns
136 -------
137 U : ndarray of shape (N, N)
138 Sparse upper-triangular inverse-Cholesky factor.
139 """
140 if out is not None:
141 U = out
142 U[:] = 0.0
143 else:
144 U = np.zeros((N, N))
146 if block_size > 0:
147 for start in range(0, N, block_size):
148 end = min(start + block_size, N)
149 cols = list(range(start, end))
150 n_cols = len(cols)
152 # Union pattern = last column's pattern (nested property)
153 nb = S[end - 1] if isinstance(S[end - 1], np.ndarray) else np.asarray(S[end - 1])
154 m = len(nb)
156 K_sub = K_ord[np.ix_(nb, nb)]
157 try:
158 L, low = cho_factor(K_sub, lower=True)
159 use_cho = True
160 except np.linalg.LinAlgError:
161 L, low = lu_factor(K_sub)
162 use_cho = False
164 # Build RHS matrix: each column k has a 1 at the position of cols[k] in nb
165 positions = np.searchsorted(nb, cols)
166 E = np.zeros((m, n_cols))
167 E[positions, np.arange(n_cols)] = 1.0
169 if use_cho:
170 V = cho_solve((L, low), E)
171 else:
172 V = lu_solve((L, low), E)
174 # Normalise each column: v /= sqrt(v[diag_pos])
175 diag_vals = V[positions, np.arange(n_cols)]
176 np.maximum(diag_vals, 1e-30, out=diag_vals)
177 np.sqrt(diag_vals, out=diag_vals)
178 V /= diag_vals
180 U[np.ix_(nb, cols)] = V
181 else:
182 for j in range(N):
183 nb = S[j] if isinstance(S[j], np.ndarray) else np.asarray(S[j])
184 K_sub = K_ord[np.ix_(nb, nb)]
185 dp = int(np.searchsorted(nb, j))
186 e = np.zeros(len(nb))
187 e[dp] = 1.0
188 v = np.linalg.solve(K_sub, e)
189 u_jj = np.sqrt(max(v[dp], 1e-30))
190 U[nb, j] = v / u_jj
192 return U
195def build_U_from_phi(phi_exp_3d, S, N, block_size,
196 k_type, k_phys, deriv_lookup, sign_lookup,
197 P_full, sigma_n_sq, sigma_data_diag,
198 out=None):
199 """
200 Build sparse U directly from phi_exp_3d, skipping full K construction.
202 Instead of building the full K matrix and extracting K_sub via fancy
203 indexing, this assembles each K_sub on the fly from the kernel's
204 intermediate phi_exp_3d representation.
206 Parameters
207 ----------
208 phi_exp_3d : ndarray of shape (n_derivs, n_rows_func, n_cols_func)
209 S : dict
210 Sparsity pattern (in MMD-ordered indices).
211 N : int
212 Total number of rows/columns.
213 block_size : int
214 Batch columns in groups of this size.
215 k_type, k_phys : int64 arrays of shape (N_total,)
216 Maps from original K index to (derivative type, physical point).
217 deriv_lookup : int64 array of shape (n_types, n_types)
218 sign_lookup : float64 array of shape (n_types,)
219 P_full : int64 array of shape (N_total,)
220 MMD permutation: P_full[mmd_idx] = original K index.
221 sigma_n_sq : float
222 Noise variance.
223 sigma_data_diag : float64 array of shape (N_total,)
224 Diagonal of sigma_data**2 in MMD order.
225 out : ndarray, optional
226 Pre-allocated output buffer of shape (N, N).
228 Returns
229 -------
230 U : ndarray of shape (N, N)
231 """
232 from jetgp.full_degp_sparse.optimizer import _extract_K_sub
234 if out is not None:
235 U = out
236 # No need to zero: the sparsity pattern S is fixed, so the loop
237 # overwrites exactly the same entries every call. Non-pattern
238 # entries stay at zero from the initial np.zeros allocation.
239 else:
240 U = np.zeros((N, N))
242 for start in range(0, N, block_size):
243 end = min(start + block_size, N)
245 # Union pattern = last column's pattern (nested property)
246 nb_union = S[end - 1] if isinstance(S[end - 1], np.ndarray) else np.asarray(S[end - 1])
247 m_union = len(nb_union)
249 # Map MMD-ordered union neighbourhood to original K indices
250 orig_nb_union = P_full[nb_union]
251 nb_type_union = k_type[orig_nb_union]
252 nb_phys_union = k_phys[orig_nb_union]
253 sd_diag_union = sigma_data_diag[nb_union]
255 # Assemble K_sub for the union neighbourhood once
256 K_sub_union = np.empty((m_union, m_union))
257 _extract_K_sub(phi_exp_3d, nb_type_union, nb_phys_union,
258 deriv_lookup, sign_lookup,
259 sigma_n_sq, sd_diag_union, m_union, K_sub_union)
261 # Batch solve: factor K_sub_union once, solve all columns at once.
262 # Within each block, earlier columns condition on the union
263 # neighbourhood (slightly larger than their own S[j]), which gives
264 # a tighter Vecchia approximation at minimal accuracy cost.
265 n_cols = end - start
266 positions = np.searchsorted(nb_union, np.arange(start, end))
267 E = np.zeros((m_union, n_cols))
268 E[positions, np.arange(n_cols)] = 1.0
269 try:
270 L_u, low_u = cho_factor(K_sub_union, lower=True)
271 V = cho_solve((L_u, low_u), E)
272 except np.linalg.LinAlgError:
273 V = np.linalg.solve(K_sub_union, E)
274 diag_vals = V[positions, np.arange(n_cols)]
275 np.maximum(diag_vals, 1e-30, out=diag_vals)
276 np.sqrt(diag_vals, out=diag_vals)
277 V /= diag_vals
278 U[nb_union, start:end] = V
280 return U
283def build_U_from_phi_flat(phi_exp_3d, block_maps, N, sigma_n_sq, out=None):
284 """
285 Build sparse U directly from phi_exp_3d using precomputed flat indices.
287 Same result as build_U_from_phi but uses a single ravel()[flat_idx]
288 gather per block instead of the numba _extract_K_sub loop.
290 Parameters
291 ----------
292 phi_exp_3d : ndarray of shape (n_derivs, n_rows_func, n_cols_func)
293 block_maps : list of dict
294 Precomputed per-block: 'nb', 'flat_idx', 'sign_mat', 'sd_diag',
295 'positions'.
296 N : int
297 sigma_n_sq : float
298 out : ndarray, optional
300 Returns
301 -------
302 U : ndarray of shape (N, N)
303 """
304 if out is not None:
305 U = out
306 else:
307 U = np.zeros((N, N))
309 phi_flat = phi_exp_3d.ravel()
311 for bm in block_maps:
312 nb = bm['nb']
313 m = len(nb)
314 start = bm['start']
316 K_sub = phi_flat[bm['flat_idx']] * bm['sign_mat']
317 diag_idx = np.arange(m)
318 K_sub[diag_idx, diag_idx] += sigma_n_sq + bm['sd_diag']
320 positions = bm['positions']
321 n_cols = len(positions)
322 E = np.zeros((m, n_cols))
323 E[positions, np.arange(n_cols)] = 1.0
324 try:
325 L_u, low_u = cho_factor(K_sub, lower=True)
326 V = cho_solve((L_u, low_u), E)
327 except np.linalg.LinAlgError:
328 V = np.linalg.solve(K_sub, E)
329 diag_vals = V[positions, np.arange(n_cols)]
330 np.maximum(diag_vals, 1e-30, out=diag_vals)
331 np.sqrt(diag_vals, out=diag_vals)
332 V /= diag_vals
333 U[nb, start:start + n_cols] = V
335 return U
338# =============================================================================
339# Supernodes
340# =============================================================================
342def build_supernodes(X_ord, l, S, lam=1.5):
343 """
344 Aggregate adjacent columns into supernodes.
346 Columns j-1 and j are merged into the same supernode if
347 dist(X_ord[P(j-1)], X_ord[P(j)]) <= lam * l[j].
349 The children of a supernode are the union of all neighbour sets of its
350 parent columns. This allows a single LU factorisation to be shared across
351 all parent columns in the supernode.
353 Parameters
354 ----------
355 X_ord : ndarray of shape (N, d)
356 l : ndarray of float, shape (N,)
357 S : dict
358 Sparsity pattern.
359 lam : float
360 Merging threshold (should be >= 1).
362 Returns
363 -------
364 supernodes : list of dict
365 Each dict has keys 'parents' (list of column indices) and
366 'children' (sorted list of row indices = union of S[p]).
367 """
368 N = len(X_ord)
369 supernodes = []
370 current_parents = [0]
371 for j in range(1, N):
372 prev = current_parents[-1]
373 dist_prev = np.linalg.norm(X_ord[prev] - X_ord[j])
374 if l[j] > 0 and dist_prev <= lam * l[j]:
375 current_parents.append(j)
376 else:
377 children = set()
378 for p in current_parents:
379 children.update(S[p])
380 supernodes.append({
381 'parents': list(current_parents),
382 'children': sorted(children)
383 })
384 current_parents = [j]
385 children = set()
386 for p in current_parents:
387 children.update(S[p])
388 supernodes.append({
389 'parents': list(current_parents),
390 'children': sorted(children)
391 })
392 return supernodes
395def build_U_supernodes(K_ord, supernodes, N):
396 """
397 Build sparse U using supernode structure.
399 Each supernode factorises K_sub (children x children) once via LU and
400 solves for all parent columns. LU is used instead of Cholesky because
401 derivative-enhanced kernel submatrices can be poorly conditioned and
402 Cholesky may fail on them even when the full K is PD.
404 Parameters
405 ----------
406 K_ord : ndarray of shape (N, N)
407 supernodes : list of dict
408 Output of build_supernodes or expand_supernodes_to_blocks.
409 N : int
411 Returns
412 -------
413 U : ndarray of shape (N, N)
414 n_factorizations : int
415 Number of LU factorisations performed.
417 Raises
418 ------
419 np.linalg.LinAlgError
420 If a supernode's K_sub is singular.
421 """
422 U = np.zeros((N, N))
423 n_factorizations = 0
424 for sn in supernodes:
425 # Use pre-computed numpy array and position lookup if available
426 ch = sn.get('children_arr')
427 if ch is None:
428 ch = np.asarray(sn['children'])
429 ch_pos = sn.get('ch_pos')
430 if ch_pos is None:
431 ch_pos = {c: i for i, c in enumerate(sn['children'])}
433 m = len(ch)
434 K_sub = K_ord[np.ix_(ch, ch)]
435 try:
436 L, low = cho_factor(K_sub, lower=True)
437 use_cho = True
438 except np.linalg.LinAlgError:
439 L, low = lu_factor(K_sub)
440 if np.any(np.abs(np.diag(L)) < 1e-30):
441 raise np.linalg.LinAlgError("Singular submatrix in supernode")
442 use_cho = False
443 n_factorizations += 1
445 # Solve all parent columns in one batch
446 parents = sn['parents']
447 n_parents = len(parents)
448 parent_positions = sn.get('parent_positions')
449 if parent_positions is None:
450 parent_positions = np.array([ch_pos[p] for p in parents])
452 E = np.zeros((m, n_parents))
453 E[parent_positions, np.arange(n_parents)] = 1.0
454 if use_cho:
455 V = cho_solve((L, low), E)
456 else:
457 V = lu_solve((L, low), E) # (m, n_parents)
459 # Vectorised normalisation: extract diagonal entries, clamp, sqrt
460 diag_vals = V[parent_positions, np.arange(n_parents)]
461 np.maximum(diag_vals, 1e-30, out=diag_vals)
462 np.sqrt(diag_vals, out=diag_vals)
463 V /= diag_vals # broadcast: (m, n_parents) / (n_parents,)
465 # Place all columns at once via fancy indexing
466 U[np.ix_(ch, parents)] = V
467 return U, n_factorizations
471def build_U_supernodes_from_phi(phi_exp_3d, supernodes, N, sigma_n_sq):
472 """
473 Build sparse U using supernode structure directly from phi_exp_3d.
475 Like build_U_supernodes but assembles each supernode's K_sub on the fly
476 via a precomputed flat index into phi_exp_3d, skipping full K construction.
478 Requires that each supernode dict has precomputed keys:
479 'phi_flat_idx', 'phi_sign_mat', 'phi_sd_diag'
480 (set up once by the optimizer's _ensure_phi_index_maps).
482 Parameters
483 ----------
484 phi_exp_3d : ndarray of shape (n_derivs, n_rows_func, n_cols_func)
485 supernodes : list of dict
486 N : int
487 sigma_n_sq : float
489 Returns
490 -------
491 U : ndarray of shape (N, N)
492 n_factorizations : int
493 """
494 phi_flat = phi_exp_3d.ravel()
496 U = np.zeros((N, N))
497 n_factorizations = 0
498 for sn in supernodes:
499 ch = sn.get('children_arr')
500 if ch is None:
501 ch = np.asarray(sn['children'])
503 m = len(ch)
505 # Single vectorised gather + elementwise multiply
506 K_sub = phi_flat[sn['phi_flat_idx']] * sn['phi_sign_mat']
507 diag_idx = np.arange(m)
508 K_sub[diag_idx, diag_idx] += sigma_n_sq + sn['phi_sd_diag']
510 try:
511 L, low = cho_factor(K_sub, lower=True)
512 use_cho = True
513 except np.linalg.LinAlgError:
514 L, low = lu_factor(K_sub)
515 if np.any(np.abs(np.diag(L)) < 1e-30):
516 raise np.linalg.LinAlgError("Singular submatrix in supernode")
517 use_cho = False
518 n_factorizations += 1
520 parents = sn['parents']
521 n_parents = len(parents)
522 parent_positions = sn.get('parent_positions')
523 if parent_positions is None:
524 ch_pos = sn.get('ch_pos')
525 if ch_pos is None:
526 ch_pos = {c: i for i, c in enumerate(sn['children'])}
527 parent_positions = np.array([ch_pos[p] for p in parents])
529 E = np.zeros((m, n_parents))
530 E[parent_positions, np.arange(n_parents)] = 1.0
531 if use_cho:
532 V = cho_solve((L, low), E)
533 else:
534 V = lu_solve((L, low), E)
536 diag_vals = V[parent_positions, np.arange(n_parents)]
537 np.maximum(diag_vals, 1e-30, out=diag_vals)
538 np.sqrt(diag_vals, out=diag_vals)
539 V /= diag_vals
541 U[np.ix_(ch, parents)] = V
542 return U, n_factorizations
545def build_deriv_supernodes(X_ord, l, S_phys, lam=1.5):
546 """
547 Expand physical supernodes to derivative index pairs.
549 Each physical index p maps to derivative indices [2p, 2p+1].
551 Parameters
552 ----------
553 X_ord : ndarray of shape (N, d)
554 l : ndarray of float, shape (N,)
555 S_phys : dict
556 Physical sparsity pattern.
557 lam : float
559 Returns
560 -------
561 der_sns : list of dict
562 Supernodes with 'parents' and 'children' in derivative index space.
563 """
564 phys_sns = build_supernodes(X_ord, l, S_phys, lam)
565 der_sns = []
566 for sn in phys_sns:
567 der_parents = []
568 for p in sn['parents']:
569 der_parents.extend([2 * p, 2 * p + 1])
570 der_children = []
571 for c in sn['children']:
572 der_children.extend([2 * c, 2 * c + 1])
573 der_sns.append({
574 'parents': der_parents,
575 'children': sorted(set(der_children))
576 })
577 return der_sns
580# =============================================================================
581# NLML via sparse U
582# =============================================================================
585def nlml_from_U(U, f):
586 """
587 Compute the negative log marginal likelihood from the sparse U factor.
589 NLML = 0.5 * ||U.T f||^2 - sum(log|diag(U)|) + 0.5 * N * log(2pi)
591 This is equivalent to the standard NLML when U.T @ U = K^{-1} exactly.
592 For sparse U it is an approximation.
594 Parameters
595 ----------
596 U : ndarray of shape (N, N)
597 Sparse inverse-Cholesky factor (upper triangular in construction,
598 stored as a full matrix with zeros outside the sparsity pattern).
599 f : ndarray of shape (N,)
600 Training targets (in MMD order).
602 Returns
603 -------
604 float
605 Approximate NLML value.
606 """
607 N = len(f)
608 Ut_f = U.T @ f
609 log_det_term = -np.sum(np.log(np.abs(np.diag(U)) + 1e-300))
610 return 0.5 * np.dot(Ut_f, Ut_f) + log_det_term + 0.5 * N * np.log(2 * np.pi)
613# =============================================================================
614# Expansion utilities: physical ordering → full K-matrix ordering
615# =============================================================================
617def expand_mmd_permutation(P_phys, N, derivative_locations):
618 """
619 Expand the physical MMD ordering (size N) to a full K-matrix permutation.
621 JetGP DEGP stores K in a block layout:
622 rows 0..N-1 : function values f(x_0), ..., f(x_{N-1})
623 rows N..N+|dl[0]|-1 : 1st deriv at points derivative_locations[0]
624 rows N+|dl[0]|.. : 2nd deriv at points derivative_locations[1]
625 ...
627 We build an interleaved ordering: for each physical point P_phys[q] (in
628 MMD order), we first emit its function-value row, then each derivative row
629 (if P_phys[q] appears in that derivative block). This groups each point
630 with all its observations before moving to the next MMD point.
632 Parameters
633 ----------
634 P_phys : ndarray of int, shape (N,)
635 Physical MMD permutation from mmd_ordering.
636 N : int
637 Number of physical training points.
638 derivative_locations : list of list of int
639 derivative_locations[k] = list of physical indices with derivative k.
641 Returns
642 -------
643 P_full : ndarray of int, shape (N_total,)
644 Permutation of {0, ..., N_total-1} for use with K[np.ix_(P_full, P_full)].
645 phys_to_rows : list of list of int
646 phys_to_rows[q] = all K-row indices belonging to physical point P_phys[q],
647 in the order they appear in P_full. Needed to build S_full.
648 """
649 # Build mapping: physical index p → list of K row indices
650 p_to_k_rows = {p: [p] for p in range(N)} # function-value block first
652 cum = N
653 for k, dl in enumerate(derivative_locations):
654 for pos, p in enumerate(dl):
655 p_to_k_rows[p].append(cum + pos)
656 cum += len(dl)
658 P_full = []
659 phys_to_rows = []
660 for p in P_phys:
661 rows = p_to_k_rows[p]
662 phys_to_rows.append(rows)
663 P_full.extend(rows)
665 return np.array(P_full, dtype=int), phys_to_rows
668def expand_sparsity_to_blocks(S_phys, phys_to_rows):
669 """
670 Expand a physical sparsity pattern to the full K-matrix index space.
672 S_phys[q] gives the physical MMD-order neighbours of column q (indices ≤ q).
673 S_full[j] gives the K_ord_full row indices (in P_full-indexed space) that
674 are neighbours of column j (indices ≤ j).
676 Because the interleaved ordering places all rows of physical point q at
677 P_full positions [n_rows_before_q .. n_rows_before_q + n_blocks_q - 1],
678 column j (= P_full position of some row of physical point q) has the same
679 physical neighbours as q.
681 Parameters
682 ----------
683 S_phys : dict
684 S_phys[q] = sorted list of physical MMD indices, 0 ≤ i ≤ q, for q in range(N).
685 phys_to_rows : list of list of int
686 phys_to_rows[q] = P_full positions belonging to physical point q
687 (output of expand_mmd_permutation, already in P_full index space
688 since they are consecutive starting from the offset for q).
690 Returns
691 -------
692 S_full : dict
693 S_full[j] = sorted list of P_full-space indices ≤ j that are neighbours
694 of column j.
695 """
696 # Build a flat lookup: P_full position j → physical MMD index q
697 j_to_q = {}
698 P_full_offset = 0
699 for q, rows in enumerate(phys_to_rows):
700 for _ in rows:
701 j_to_q[P_full_offset] = q
702 P_full_offset += 1
704 N_total = P_full_offset
706 # For each physical point q, compute the starting P_full index
707 q_start = {}
708 offset = 0
709 for q, rows in enumerate(phys_to_rows):
710 q_start[q] = offset
711 offset += len(rows)
713 S_full = {}
714 for j in range(N_total):
715 q_j = j_to_q[j]
716 nb = []
717 for q_nb in S_phys[q_j]: # physical neighbours (≤ q_j)
718 start = q_start[q_nb]
719 for b in range(len(phys_to_rows[q_nb])):
720 row = start + b
721 if row <= j:
722 nb.append(row)
723 S_full[j] = sorted(nb) if nb else [j]
725 return S_full
728def expand_supernodes_to_blocks(supernodes, phys_to_rows):
729 """
730 Expand physical-level supernodes to cover all derivative rows.
732 Each physical parent p maps to all P_full positions belonging to p.
733 Children are expanded similarly.
735 Parameters
736 ----------
737 supernodes : list of dict
738 Physical supernodes from build_supernodes.
739 phys_to_rows : list of list of int
740 phys_to_rows[q] = P_full positions for physical point q.
742 Returns
743 -------
744 full_supernodes : list of dict
745 Supernodes in P_full-indexed space.
746 """
747 # Compute starting P_full offset for each physical point q
748 q_start = {}
749 offset = 0
750 for q, rows in enumerate(phys_to_rows):
751 q_start[q] = offset
752 offset += len(rows)
754 full_supernodes = []
755 for sn in supernodes:
756 full_parents = []
757 for p in sn['parents']:
758 start = q_start[p]
759 full_parents.extend(range(start, start + len(phys_to_rows[p])))
761 full_children = []
762 for c in sn['children']:
763 start = q_start[c]
764 full_children.extend(range(start, start + len(phys_to_rows[c])))
766 full_supernodes.append({
767 'parents': full_parents,
768 'children': sorted(set(full_children)),
769 })
771 return full_supernodes
775def alpha_from_U(U, f):
776 """
777 Compute K^{-1} f using the sparse U factor.
779 Since U.T @ U ≈ K^{-1}, we have K^{-1} f ≈ U @ (U.T @ f).
781 Parameters
782 ----------
783 U : ndarray of shape (N, N)
784 f : ndarray of shape (N,)
786 Returns
787 -------
788 ndarray of shape (N,)
789 """
790 return U @ (U.T @ f)