Coverage for dragonfly_sim/utils/pod_utils.py: 99%

1""" 

2Helpers for core/pod.py: PETSc dense-matrix plumbing for the distributed 

3snapshot SVD, and the compact-support cubic weight function of the 

4local_weighted basis. Also the distributed thin QR (tsqr) of the reduced 

5adjoint in core/reduced_order_model.py. 

6 

7Storage convention (see core/pod.py:SnapshotMatrix): 

8 

9 - Q (the basis), distributed MATDENSE with the row layout of the state 

10 vector; its column count (the rank) is the same on every rank. 

11 - R, Sigma, VT, small PETSc objects on PETSc.COMM_SELF, one identical copy 

12 per rank. Their size scales with the number of snapshots, not with the 

13 mesh, so there is nothing to gain from distributing them. 

14 

15Rank-sized quantities such as Q^T W x are formed as a local dense product on 

16each rank's row block followed by one Allreduce, which keeps them available 

17as replicated numpy arrays. 

18""" 

19import numpy as np 

20from mpi4py import MPI 

21from petsc4py import PETSc 

22 

23 

24# ---------------------------------------------------------------------- 

25# PETSc dense plumbing 

26# ---------------------------------------------------------------------- 

27def petsc_append_dense_column(comm, old_mat, new_col_local): 

28 """ 

29 A new distributed MATDENSE with one column more than `old_mat` (None for 

30 a fresh single-column matrix): its data, then `new_col_local` (this 

31 rank's slice of the new column) as the last column. 

32 """ 

33 m_local = new_col_local.shape[0] 

34 M_global = comm.allreduce(m_local, op=MPI.SUM) 

35 r_old = 0 if old_mat is None else old_mat.getSize()[1] 

36 mat = PETSc.Mat().createDense(size=((m_local, M_global), (PETSc.DECIDE, r_old + 1)), comm=comm) 

37 mat.setUp() 

38 arr = mat.getDenseArray() 

39 if old_mat is not None: 

40 arr[:, :r_old] = old_mat.getDenseArray() 

41 arr[:, r_old] = new_col_local 

42 mat.assemble() 

43 return mat 

44 

45 

46def petsc_dense_from_local_block(comm, local_block, M_global=None): 

47 """A distributed MATDENSE whose local row block is `local_block` 

48 (shape (m_local, ncols)).""" 

49 m_local, ncols = local_block.shape 

50 if M_global is None: 

51 M_global = comm.allreduce(m_local, op=MPI.SUM) 

52 mat = PETSc.Mat().createDense(size=((m_local, M_global), (PETSc.DECIDE, ncols)), comm=comm) 

53 mat.setUp() 

54 mat.getDenseArray()[:, :] = local_block 

55 mat.assemble() 

56 return mat 

57 

58 

59def petsc_small_mat_from_numpy(arr): 

60 """A small numpy array as a PETSc.COMM_SELF MATDENSE (the caller builds 

61 `arr` identically on every rank).""" 

62 mat = PETSc.Mat().createDense(size=arr.shape, comm=PETSc.COMM_SELF) 

63 mat.setUp() 

64 if arr.size > 0: 

65 mat.getDenseArray()[:, :] = arr 

66 mat.assemble() 

67 return mat 

68 

69 

70def petsc_small_vec_from_numpy(arr): 

71 """A small 1D numpy array as a PETSc.COMM_SELF Vec (replicated).""" 

72 vec = PETSc.Vec().createSeq(arr.shape[0], comm=PETSc.COMM_SELF) 

73 if arr.size > 0: 

74 vec.setArray(arr) 

75 vec.assemble() 

76 return vec 

77 

78 

79def petsc_small_mat_to_numpy(mat): 

80 if mat is None: 

81 return np.zeros((0, 0)) 

82 return mat.getDenseArray().copy() 

83 

84 

85def replicated_transpose_product(comm, V, x_local): 

86 """V^T x as a numpy array on every rank, for a distributed MATDENSE V 

87 and this rank's owned entries x_local of a vector in V's row layout.""" 

88 return comm.allreduce(V.getDenseArray().T @ x_local, op=MPI.SUM) 

89 

90 

91def tsqr(comm, X_local): 

92 """ 

93 Thin QR, X = Q R, of a tall matrix X (N x k) distributed by rows, 

94 X_local being this rank's rows (TSQR): each rank factors its block, the 

95 k x k factors are gathered in rank order and factored once more. 

96 

97 Returns (Q_local, R): this rank's rows of Q (Q^T Q = I) and the k x k 

98 upper-triangular R, identical on every rank (every rank factors the same 

99 gathered stack). A rank with fewer than k rows (or none) pads its factor 

100 with zero rows. 

101 """ 

102 n_local, k = X_local.shape 

103 Q1 = np.zeros((n_local, k)) 

104 R1 = np.zeros((k, k)) 

105 if n_local > 0: 

106 q, r = np.linalg.qr(X_local, mode='reduced') 

107 m = r.shape[0] 

108 Q1[:, :m] = q 

109 R1[:m, :] = r 

110 Q2, R = np.linalg.qr(np.vstack(comm.allgather(R1)), mode='reduced') 

111 rank = comm.Get_rank() 

112 return Q1 @ Q2[rank * k:(rank + 1) * k, :], R 

113 

114 

115class DiagonalWeight: 

116 """ 

117 A diagonal inner-product weight W = diag(w) with the `.mult(x, y)` that 

118 core/pod.py:SnapshotMatrix calls, w a Vec in the state layout. 

119 """ 

120 

121 def __init__(self, w): 

122 self.w = w 

123 

124 def mult(self, x, y): 

125 y.pointwiseMult(self.w, x) 

126 

127 def destroy(self): 

128 self.w.destroy() 

129 

130 

131# ---------------------------------------------------------------------- 

132# Cubic weight function of the local_weighted basis 

133# ---------------------------------------------------------------------- 

134def cubic_dist_hat(parameter_distnorm, min_dist, max_dist, n_nonzero_weights, 

135 tail_eps_scale=1e-3, tie_rtol=1e-9): 

136 """ 

137 The cubic weight's support radius dist_hat such that exactly the 

138 k = min(n_nonzero_weights, n_snapshots) snapshots nearest the query point 

139 get a nonzero weight and the (k+1)-th nearest gets zero. 

140 

141 k >= n_snapshots: there is no (k+1)-th neighbour to cut at, so dist_hat 

142 is pushed a small relative distance beyond max_dist (tail_eps_scale keeps 

143 the farthest snapshot's weight above the 1e-10 pruning threshold in 

144 PODBasis.construct_basis). 

145 

146 k or more snapshots tied at min_dist (e.g. a symmetric pair of design 

147 perturbations) would put dist_hat at min_dist; k is walked forward past 

148 the tie first. Ties exactly at the (k+1)-th boundary need no handling: 

149 both snapshots land on dist_hat and both get zero. 

150 """ 

151 n_snapshots = parameter_distnorm.shape[0] 

152 k = max(0, min(int(n_nonzero_weights), n_snapshots)) 

153 

154 if k < n_snapshots: 

155 sorted_dist = np.sort(parameter_distnorm) 

156 tol = tie_rtol * (max_dist - min_dist) 

157 while k < n_snapshots and (sorted_dist[k] - min_dist) <= tol: 

158 k += 1 

159 if k < n_snapshots: 

160 return sorted_dist[k] 

161 

162 return max_dist + tail_eps_scale * (max_dist - min_dist) 

163 

164 

165def cubic_dist_hat_fixed(min_dist, max_dist, c): 

166 """ 

167 The cubic weight's support radius from a user-set c in [0, 1): 

168 

169 dist_hat = c min_dist + (1 - c) max_dist. 

170 

171 c = 0 puts dist_hat at the farthest snapshot (only it gets zero weight); 

172 larger c narrows the support towards the nearest snapshot. 

173 """ 

174 if not 0.0 <= c < 1.0: 

175 raise ValueError("the cubic cutoff c must lie in [0, 1), got {!r}".format(c)) 

176 return c * min_dist + (1.0 - c) * max_dist 

177 

178 

179def basis_size_floor(accepted_distnorm, rb_size, spread, margin_scale=1e-3): 

180 """ 

181 The smallest cubic support radius that leaves the rb_size snapshots 

182 nearest the query point (among `accepted_distnorm`, the snapshots that 

183 carry a basis column) with a nonzero weight. 

184 

185 The cubic weight vanishes AT its support radius, so the radius has to lie 

186 strictly beyond d_k, the k-th smallest distance (k = min(rb_size, number 

187 of snapshots)): it is the next larger distance (that snapshot gets zero 

188 weight), but at least d_k + margin_scale * spread, with `spread` the 

189 spread of all distances (max - min). The margin keeps the k-th weight 

190 clear of the 1e-10 pruning threshold when the next distance is barely 

191 larger, and places the radius beyond d_k when no larger distance exists. 

192 Snapshots tied with d_k all stay inside. 

193 """ 

194 d = np.sort(np.asarray(accepted_distnorm, dtype=float)) 

195 k = max(1, min(int(rb_size), d.shape[0])) 

196 d_k = d[k - 1] 

197 floor = d_k + margin_scale * spread 

198 larger = d[d > d_k] 

199 if larger.size: 

200 floor = max(floor, larger[0]) 

201 return floor 

202 

203 

204def cubic_weights(parameter_distnorm, min_dist, dist_hat): 

205 """ 

206 The cubic Hermite weight with w(min_dist) = 1, w(dist_hat) = 0 and zero 

207 slope at both ends, zero beyond dist_hat: 

208 

209 w = 2 t^3 - 3 t^2 + 1, t = (d - min_dist) / (dist_hat - min_dist). 

210 

211 This closed form is the exact solution of the 4x4 monomial interpolation 

212 system, without that system's ill-conditioning at large absolute 

213 distances. The zeroing beyond dist_hat is required: the polynomial turns 

214 back up for t > 1. 

215 """ 

216 t = (parameter_distnorm - min_dist) / (dist_hat - min_dist) 

217 snapshot_weights = 2.0 * np.power(t, 3) - 3.0 * np.power(t, 2) + 1.0 

218 snapshot_weights[parameter_distnorm > dist_hat] = 0. 

219 return snapshot_weights