Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/scipy/sparse/linalg/_eigen/lobpcg/lobpcg.py: 3%

1""" 

2Locally Optimal Block Preconditioned Conjugate Gradient Method (LOBPCG). 

3 

4References 

5---------- 

6.. [1] A. V. Knyazev (2001), 

7 Toward the Optimal Preconditioned Eigensolver: Locally Optimal 

8 Block Preconditioned Conjugate Gradient Method. 

9 SIAM Journal on Scientific Computing 23, no. 2, 

10 pp. 517-541. :doi:`10.1137/S1064827500366124` 

11 

12.. [2] A. V. Knyazev, I. Lashuk, M. E. Argentati, and E. Ovchinnikov (2007), 

13 Block Locally Optimal Preconditioned Eigenvalue Xolvers (BLOPEX) 

14 in hypre and PETSc. :arxiv:`0705.2626` 

15 

16.. [3] A. V. Knyazev's C and MATLAB implementations: 

17 https://github.com/lobpcg/blopex 

18""" 

19 

20import warnings 

21import numpy as np 

22from scipy.linalg import (inv, eigh, cho_factor, cho_solve, 

23 cholesky, LinAlgError) 

24from scipy.sparse.linalg import LinearOperator 

25from scipy.sparse import issparse 

26 

27__all__ = ["lobpcg"] 

28 

29 

30def _report_nonhermitian(M, name): 

31 """ 

32 Report if `M` is not a Hermitian matrix given its type. 

33 """ 

34 from scipy.linalg import norm 

35 

36 md = M - M.T.conj() 

37 nmd = norm(md, 1) 

38 tol = 10 * np.finfo(M.dtype).eps 

39 tol = max(tol, tol * norm(M, 1)) 

40 if nmd > tol: 

41 warnings.warn( 

42 f"Matrix {name} of the type {M.dtype} is not Hermitian: " 

43 f"condition: {nmd} < {tol} fails.", 

44 UserWarning, stacklevel=4 

45 ) 

46 

47def _as2d(ar): 

48 """ 

49 If the input array is 2D return it, if it is 1D, append a dimension, 

50 making it a column vector. 

51 """ 

52 if ar.ndim == 2: 

53 return ar 

54 else: # Assume 1! 

55 aux = np.array(ar, copy=False) 

56 aux.shape = (ar.shape[0], 1) 

57 return aux 

58 

59 

60def _makeMatMat(m): 

61 if m is None: 

62 return None 

63 elif callable(m): 

64 return lambda v: m(v) 

65 else: 

66 return lambda v: m @ v 

67 

68 

69def _matmul_inplace(x, y, verbosityLevel=0): 

70 """Perform 'np.matmul' in-place if possible. 

71 

72 If some sufficient conditions for inplace matmul are met, do so. 

73 Otherwise try inplace update and fall back to overwrite if that fails. 

74 """ 

75 if x.flags["CARRAY"] and x.shape[1] == y.shape[1] and x.dtype == y.dtype: 

76 # conditions where we can guarantee that inplace updates will work; 

77 # i.e. x is not a view/slice, x & y have compatible dtypes, and the 

78 # shape of the result of x @ y matches the shape of x. 

79 np.matmul(x, y, out=x) 

80 else: 

81 # ideally, we'd have an exhaustive list of conditions above when 

82 # inplace updates are possible; since we don't, we opportunistically 

83 # try if it works, and fall back to overwriting if necessary 

84 try: 

85 np.matmul(x, y, out=x) 

86 except Exception: 

87 if verbosityLevel: 

88 warnings.warn( 

89 "Inplace update of x = x @ y failed, " 

90 "x needs to be overwritten.", 

91 UserWarning, stacklevel=3 

92 ) 

93 x = x @ y 

94 return x 

95 

96 

97def _applyConstraints(blockVectorV, factYBY, blockVectorBY, blockVectorY): 

98 """Changes blockVectorV in-place.""" 

99 YBV = blockVectorBY.T.conj() @ blockVectorV 

100 tmp = cho_solve(factYBY, YBV) 

101 blockVectorV -= blockVectorY @ tmp 

102 

103 

104def _b_orthonormalize(B, blockVectorV, blockVectorBV=None, 

105 verbosityLevel=0): 

106 """in-place B-orthonormalize the given block vector using Cholesky.""" 

107 if blockVectorBV is None: 

108 if B is None: 

109 blockVectorBV = blockVectorV 

110 else: 

111 try: 

112 blockVectorBV = B(blockVectorV) 

113 except Exception as e: 

114 if verbosityLevel: 

115 warnings.warn( 

116 f"Secondary MatMul call failed with error\n" 

117 f"{e}\n", 

118 UserWarning, stacklevel=3 

119 ) 

120 return None, None, None 

121 if blockVectorBV.shape != blockVectorV.shape: 

122 raise ValueError( 

123 f"The shape {blockVectorV.shape} " 

124 f"of the orthogonalized matrix not preserved\n" 

125 f"and changed to {blockVectorBV.shape} " 

126 f"after multiplying by the secondary matrix.\n" 

127 ) 

128 

129 VBV = blockVectorV.T.conj() @ blockVectorBV 

130 try: 

131 # VBV is a Cholesky factor from now on... 

132 VBV = cholesky(VBV, overwrite_a=True) 

133 VBV = inv(VBV, overwrite_a=True) 

134 blockVectorV = _matmul_inplace( 

135 blockVectorV, VBV, 

136 verbosityLevel=verbosityLevel 

137 ) 

138 if B is not None: 

139 blockVectorBV = _matmul_inplace( 

140 blockVectorBV, VBV, 

141 verbosityLevel=verbosityLevel 

142 ) 

143 return blockVectorV, blockVectorBV, VBV 

144 except LinAlgError: 

145 if verbosityLevel: 

146 warnings.warn( 

147 "Cholesky has failed.", 

148 UserWarning, stacklevel=3 

149 ) 

150 return None, None, None 

151 

152 

153def _get_indx(_lambda, num, largest): 

154 """Get `num` indices into `_lambda` depending on `largest` option.""" 

155 ii = np.argsort(_lambda) 

156 if largest: 

157 ii = ii[:-num - 1:-1] 

158 else: 

159 ii = ii[:num] 

160 

161 return ii 

162 

163 

164def _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel): 

165 if verbosityLevel: 

166 _report_nonhermitian(gramA, "gramA") 

167 _report_nonhermitian(gramB, "gramB") 

168 

169 

170def lobpcg( 

171 A, 

172 X, 

173 B=None, 

174 M=None, 

175 Y=None, 

176 tol=None, 

177 maxiter=None, 

178 largest=True, 

179 verbosityLevel=0, 

180 retLambdaHistory=False, 

181 retResidualNormsHistory=False, 

182 restartControl=20, 

183): 

184 """Locally Optimal Block Preconditioned Conjugate Gradient Method (LOBPCG). 

185 

186 LOBPCG is a preconditioned eigensolver for large real symmetric and complex 

187 Hermitian definite generalized eigenproblems. 

188 

189 Parameters 

190 ---------- 

191 A : {sparse matrix, ndarray, LinearOperator, callable object} 

192 The Hermitian linear operator of the problem, usually given by a 

193 sparse matrix. Often called the "stiffness matrix". 

194 X : ndarray, float32 or float64 

195 Initial approximation to the ``k`` eigenvectors (non-sparse). 

196 If `A` has ``shape=(n,n)`` then `X` must have ``shape=(n,k)``. 

197 B : {sparse matrix, ndarray, LinearOperator, callable object} 

198 Optional. By default ``B = None``, which is equivalent to identity. 

199 The right hand side operator in a generalized eigenproblem if present. 

200 Often called the "mass matrix". Must be Hermitian positive definite. 

201 M : {sparse matrix, ndarray, LinearOperator, callable object} 

202 Optional. By default ``M = None``, which is equivalent to identity. 

203 Preconditioner aiming to accelerate convergence. 

204 Y : ndarray, float32 or float64, default: None 

205 An ``n-by-sizeY`` ndarray of constraints with ``sizeY < n``. 

206 The iterations will be performed in the ``B``-orthogonal complement 

207 of the column-space of `Y`. `Y` must be full rank if present. 

208 tol : scalar, optional 

209 The default is ``tol=n*sqrt(eps)``. 

210 Solver tolerance for the stopping criterion. 

211 maxiter : int, default: 20 

212 Maximum number of iterations. 

213 largest : bool, default: True 

214 When True, solve for the largest eigenvalues, otherwise the smallest. 

215 verbosityLevel : int, optional 

216 By default ``verbosityLevel=0`` no output. 

217 Controls the solver standard/screen output. 

218 retLambdaHistory : bool, default: False 

219 Whether to return iterative eigenvalue history. 

220 retResidualNormsHistory : bool, default: False 

221 Whether to return iterative history of residual norms. 

222 restartControl : int, optional. 

223 Iterations restart if the residuals jump ``2**restartControl`` times 

224 compared to the smallest recorded in ``retResidualNormsHistory``. 

225 The default is ``restartControl=20``, making the restarts rare for 

226 backward compatibility. 

227 

228 Returns 

229 ------- 

230 lambda : ndarray of the shape ``(k, )``. 

231 Array of ``k`` approximate eigenvalues. 

232 v : ndarray of the same shape as ``X.shape``. 

233 An array of ``k`` approximate eigenvectors. 

234 lambdaHistory : ndarray, optional. 

235 The eigenvalue history, if `retLambdaHistory` is ``True``. 

236 ResidualNormsHistory : ndarray, optional. 

237 The history of residual norms, if `retResidualNormsHistory` 

238 is ``True``. 

239 

240 Notes 

241 ----- 

242 The iterative loop runs ``maxit=maxiter`` (20 if ``maxit=None``) 

243 iterations at most and finishes earler if the tolerance is met. 

244 Breaking backward compatibility with the previous version, LOBPCG 

245 now returns the block of iterative vectors with the best accuracy rather 

246 than the last one iterated, as a cure for possible divergence. 

247 

248 If ``X.dtype == np.float32`` and user-provided operations/multiplications 

249 by `A`, `B`, and `M` all preserve the ``np.float32`` data type, 

250 all the calculations and the output are in ``np.float32``. 

251 

252 The size of the iteration history output equals to the number of the best 

253 (limited by `maxit`) iterations plus 3: initial, final, and postprocessing. 

254 

255 If both `retLambdaHistory` and `retResidualNormsHistory` are ``True``, 

256 the return tuple has the following format 

257 ``(lambda, V, lambda history, residual norms history)``. 

258 

259 In the following ``n`` denotes the matrix size and ``k`` the number 

260 of required eigenvalues (smallest or largest). 

261 

262 The LOBPCG code internally solves eigenproblems of the size ``3k`` on every 

263 iteration by calling the dense eigensolver `eigh`, so if ``k`` is not 

264 small enough compared to ``n``, it makes no sense to call the LOBPCG code. 

265 Moreover, if one calls the LOBPCG algorithm for ``5k > n``, it would likely 

266 break internally, so the code calls the standard function `eigh` instead. 

267 It is not that ``n`` should be large for the LOBPCG to work, but rather the 

268 ratio ``n / k`` should be large. It you call LOBPCG with ``k=1`` 

269 and ``n=10``, it works though ``n`` is small. The method is intended 

270 for extremely large ``n / k``. 

271 

272 The convergence speed depends basically on three factors: 

273 

274 1. Quality of the initial approximations `X` to the seeking eigenvectors. 

275 Randomly distributed around the origin vectors work well if no better 

276 choice is known. 

277 

278 2. Relative separation of the desired eigenvalues from the rest 

279 of the eigenvalues. One can vary ``k`` to improve the separation. 

280 

281 3. Proper preconditioning to shrink the spectral spread. 

282 For example, a rod vibration test problem (under tests 

283 directory) is ill-conditioned for large ``n``, so convergence will be 

284 slow, unless efficient preconditioning is used. For this specific 

285 problem, a good simple preconditioner function would be a linear solve 

286 for `A`, which is easy to code since `A` is tridiagonal. 

287 

288 References 

289 ---------- 

290 .. [1] A. V. Knyazev (2001), 

291 Toward the Optimal Preconditioned Eigensolver: Locally Optimal 

292 Block Preconditioned Conjugate Gradient Method. 

293 SIAM Journal on Scientific Computing 23, no. 2, 

294 pp. 517-541. :doi:`10.1137/S1064827500366124` 

295 

296 .. [2] A. V. Knyazev, I. Lashuk, M. E. Argentati, and E. Ovchinnikov 

297 (2007), Block Locally Optimal Preconditioned Eigenvalue Xolvers 

298 (BLOPEX) in hypre and PETSc. :arxiv:`0705.2626` 

299 

300 .. [3] A. V. Knyazev's C and MATLAB implementations: 

301 https://github.com/lobpcg/blopex 

302 

303 Examples 

304 -------- 

305 Our first example is minimalistic - find the largest eigenvalue of 

306 a diagonal matrix by solving the non-generalized eigenvalue problem 

307 ``A x = lambda x`` without constraints or preconditioning. 

308 

309 >>> import numpy as np 

310 >>> from scipy.sparse import spdiags 

311 >>> from scipy.sparse.linalg import LinearOperator, aslinearoperator 

312 >>> from scipy.sparse.linalg import lobpcg 

313 

314 The square matrix size is 

315 

316 >>> n = 100 

317 

318 and its diagonal entries are 1, ..., 100 defined by 

319 

320 >>> vals = np.arange(1, n + 1).astype(np.int16) 

321 

322 The first mandatory input parameter in this test is 

323 the sparse diagonal matrix `A` 

324 of the eigenvalue problem ``A x = lambda x`` to solve. 

325 

326 >>> A = spdiags(vals, 0, n, n) 

327 >>> A = A.astype(np.int16) 

328 >>> A.toarray() 

329 array([[ 1, 0, 0, ..., 0, 0, 0], 

330 [ 0, 2, 0, ..., 0, 0, 0], 

331 [ 0, 0, 3, ..., 0, 0, 0], 

332 ..., 

333 [ 0, 0, 0, ..., 98, 0, 0], 

334 [ 0, 0, 0, ..., 0, 99, 0], 

335 [ 0, 0, 0, ..., 0, 0, 100]], dtype=int16) 

336 

337 The second mandatory input parameter `X` is a 2D array with the 

338 row dimension determining the number of requested eigenvalues. 

339 `X` is an initial guess for targeted eigenvectors. 

340 `X` must have linearly independent columns. 

341 If no initial approximations available, randomly oriented vectors 

342 commonly work best, e.g., with components normally distributed 

343 around zero or uniformly distributed on the interval [-1 1]. 

344 Setting the initial approximations to dtype ``np.float32`` 

345 forces all iterative values to dtype ``np.float32`` speeding up 

346 the run while still allowing accurate eigenvalue computations. 

347 

348 >>> k = 1 

349 >>> rng = np.random.default_rng() 

350 >>> X = rng.normal(size=(n, k)) 

351 >>> X = X.astype(np.float32) 

352 

353 >>> eigenvalues, _ = lobpcg(A, X, maxiter=60) 

354 >>> eigenvalues 

355 array([100.]) 

356 >>> eigenvalues.dtype 

357 dtype('float32') 

358 

359 `lobpcg` needs only access the matrix product with `A` rather 

360 then the matrix itself. Since the matrix `A` is diagonal in 

361 this example, one can write a function of the matrix product 

362 ``A @ X`` using the diagonal values ``vals`` only, e.g., by 

363 element-wise multiplication with broadcasting in the lambda-function 

364 

365 >>> A_lambda = lambda X: vals[:, np.newaxis] * X 

366 

367 or the regular function 

368 

369 >>> def A_matmat(X): 

370 ... return vals[:, np.newaxis] * X 

371 

372 and use the handle to one of these callables as an input 

373 

374 >>> eigenvalues, _ = lobpcg(A_lambda, X, maxiter=60) 

375 >>> eigenvalues 

376 array([100.]) 

377 >>> eigenvalues, _ = lobpcg(A_matmat, X, maxiter=60) 

378 >>> eigenvalues 

379 array([100.]) 

380 

381 The traditional callable `LinearOperator` is no longer 

382 necessary but still supported as the input to `lobpcg`. 

383 Specifying ``matmat=A_matmat`` explicitely improves performance.  

384 

385 >>> A_lo = LinearOperator((n, n), matvec=A_matmat, matmat=A_matmat, dtype=np.int16) 

386 >>> eigenvalues, _ = lobpcg(A_lo, X, maxiter=80) 

387 >>> eigenvalues 

388 array([100.]) 

389 

390 The least efficient callable option is `aslinearoperator`: 

391 

392 >>> eigenvalues, _ = lobpcg(aslinearoperator(A), X, maxiter=80) 

393 >>> eigenvalues 

394 array([100.]) 

395 

396 We now switch to computing the three smallest eigenvalues specifying 

397 

398 >>> k = 3 

399 >>> X = np.random.default_rng().normal(size=(n, k)) 

400 

401 and ``largest=False`` parameter 

402 

403 >>> eigenvalues, _ = lobpcg(A, X, largest=False, maxiter=80) 

404 >>> print(eigenvalues)  

405 [1. 2. 3.] 

406 

407 The next example illustrates computing 3 smallest eigenvalues of 

408 the same matrix `A` given by the function handle ``A_matmat`` but 

409 with constraints and preconditioning. 

410 

411 Constraints - an optional input parameter is a 2D array comprising 

412 of column vectors that the eigenvectors must be orthogonal to 

413 

414 >>> Y = np.eye(n, 3) 

415 

416 The preconditioner acts as the inverse of `A` in this example, but 

417 in the reduced precision ``np.float32`` even though the initial `X` 

418 and thus all iterates and the output are in full ``np.float64``. 

419 

420 >>> inv_vals = 1./vals 

421 >>> inv_vals = inv_vals.astype(np.float32) 

422 >>> M = lambda X: inv_vals[:, np.newaxis] * X 

423 

424 Let us now solve the eigenvalue problem for the matrix `A` first 

425 without preconditioning requesting 80 iterations 

426 

427 >>> eigenvalues, _ = lobpcg(A_matmat, X, Y=Y, largest=False, maxiter=80) 

428 >>> eigenvalues 

429 array([4., 5., 6.]) 

430 >>> eigenvalues.dtype 

431 dtype('float64') 

432 

433 With preconditioning we need only 20 iterations from the same `X` 

434 

435 >>> eigenvalues, _ = lobpcg(A_matmat, X, Y=Y, M=M, largest=False, maxiter=20) 

436 >>> eigenvalues 

437 array([4., 5., 6.]) 

438 

439 Note that the vectors passed in `Y` are the eigenvectors of the 3 

440 smallest eigenvalues. The results returned above are orthogonal to those. 

441 

442 The primary matrix `A` may be indefinite, e.g., after shifting 

443 ``vals`` by 50 from 1, ..., 100 to -49, ..., 50, we still can compute 

444 the 3 smallest or largest eigenvalues. 

445 

446 >>> vals = vals - 50 

447 >>> X = rng.normal(size=(n, k)) 

448 >>> eigenvalues, _ = lobpcg(A_matmat, X, largest=False, maxiter=99) 

449 >>> eigenvalues 

450 array([-49., -48., -47.]) 

451 >>> eigenvalues, _ = lobpcg(A_matmat, X, largest=True, maxiter=99) 

452 >>> eigenvalues 

453 array([50., 49., 48.]) 

454 

455 """ 

456 blockVectorX = X 

457 bestblockVectorX = blockVectorX 

458 blockVectorY = Y 

459 residualTolerance = tol 

460 if maxiter is None: 

461 maxiter = 20 

462 

463 bestIterationNumber = maxiter 

464 

465 sizeY = 0 

466 if blockVectorY is not None: 

467 if len(blockVectorY.shape) != 2: 

468 warnings.warn( 

469 f"Expected rank-2 array for argument Y, instead got " 

470 f"{len(blockVectorY.shape)}, " 

471 f"so ignore it and use no constraints.", 

472 UserWarning, stacklevel=2 

473 ) 

474 blockVectorY = None 

475 else: 

476 sizeY = blockVectorY.shape[1] 

477 

478 # Block size. 

479 if blockVectorX is None: 

480 raise ValueError("The mandatory initial matrix X cannot be None") 

481 if len(blockVectorX.shape) != 2: 

482 raise ValueError("expected rank-2 array for argument X") 

483 

484 n, sizeX = blockVectorX.shape 

485 

486 # Data type of iterates, determined by X, must be inexact 

487 if not np.issubdtype(blockVectorX.dtype, np.inexact): 

488 warnings.warn( 

489 f"Data type for argument X is {blockVectorX.dtype}, " 

490 f"which is not inexact, so casted to np.float32.", 

491 UserWarning, stacklevel=2 

492 ) 

493 blockVectorX = np.asarray(blockVectorX, dtype=np.float32) 

494 

495 if retLambdaHistory: 

496 lambdaHistory = np.zeros((maxiter + 3, sizeX), 

497 dtype=blockVectorX.dtype) 

498 if retResidualNormsHistory: 

499 residualNormsHistory = np.zeros((maxiter + 3, sizeX), 

500 dtype=blockVectorX.dtype) 

501 

502 if verbosityLevel: 

503 aux = "Solving " 

504 if B is None: 

505 aux += "standard" 

506 else: 

507 aux += "generalized" 

508 aux += " eigenvalue problem with" 

509 if M is None: 

510 aux += "out" 

511 aux += " preconditioning\n\n" 

512 aux += "matrix size %d\n" % n 

513 aux += "block size %d\n\n" % sizeX 

514 if blockVectorY is None: 

515 aux += "No constraints\n\n" 

516 else: 

517 if sizeY > 1: 

518 aux += "%d constraints\n\n" % sizeY 

519 else: 

520 aux += "%d constraint\n\n" % sizeY 

521 print(aux) 

522 

523 if (n - sizeY) < (5 * sizeX): 

524 warnings.warn( 

525 f"The problem size {n} minus the constraints size {sizeY} " 

526 f"is too small relative to the block size {sizeX}. " 

527 f"Using a dense eigensolver instead of LOBPCG iterations." 

528 f"No output of the history of the iterations.", 

529 UserWarning, stacklevel=2 

530 ) 

531 

532 sizeX = min(sizeX, n) 

533 

534 if blockVectorY is not None: 

535 raise NotImplementedError( 

536 "The dense eigensolver does not support constraints." 

537 ) 

538 

539 # Define the closed range of indices of eigenvalues to return. 

540 if largest: 

541 eigvals = (n - sizeX, n - 1) 

542 else: 

543 eigvals = (0, sizeX - 1) 

544 

545 try: 

546 if isinstance(A, LinearOperator): 

547 A = A(np.eye(n, dtype=int)) 

548 elif callable(A): 

549 A = A(np.eye(n, dtype=int)) 

550 if A.shape != (n, n): 

551 raise ValueError( 

552 f"The shape {A.shape} of the primary matrix\n" 

553 f"defined by a callable object is wrong.\n" 

554 ) 

555 elif issparse(A): 

556 A = A.toarray() 

557 else: 

558 A = np.asarray(A) 

559 except Exception as e: 

560 raise Exception( 

561 f"Primary MatMul call failed with error\n" 

562 f"{e}\n") 

563 

564 if B is not None: 

565 try: 

566 if isinstance(B, LinearOperator): 

567 B = B(np.eye(n, dtype=int)) 

568 elif callable(B): 

569 B = B(np.eye(n, dtype=int)) 

570 if B.shape != (n, n): 

571 raise ValueError( 

572 f"The shape {B.shape} of the secondary matrix\n" 

573 f"defined by a callable object is wrong.\n" 

574 ) 

575 elif issparse(B): 

576 B = B.toarray() 

577 else: 

578 B = np.asarray(B) 

579 except Exception as e: 

580 raise Exception( 

581 f"Secondary MatMul call failed with error\n" 

582 f"{e}\n") 

583 

584 try: 

585 vals, vecs = eigh(A, 

586 B, 

587 subset_by_index=eigvals, 

588 check_finite=False) 

589 if largest: 

590 # Reverse order to be compatible with eigs() in 'LM' mode. 

591 vals = vals[::-1] 

592 vecs = vecs[:, ::-1] 

593 

594 return vals, vecs 

595 except Exception as e: 

596 raise Exception( 

597 f"Dense eigensolver failed with error\n" 

598 f"{e}\n" 

599 ) 

600 

601 if (residualTolerance is None) or (residualTolerance <= 0.0): 

602 residualTolerance = np.sqrt(np.finfo(blockVectorX.dtype).eps) * n 

603 

604 A = _makeMatMat(A) 

605 B = _makeMatMat(B) 

606 M = _makeMatMat(M) 

607 

608 # Apply constraints to X. 

609 if blockVectorY is not None: 

610 

611 if B is not None: 

612 blockVectorBY = B(blockVectorY) 

613 if blockVectorBY.shape != blockVectorY.shape: 

614 raise ValueError( 

615 f"The shape {blockVectorY.shape} " 

616 f"of the constraint not preserved\n" 

617 f"and changed to {blockVectorBY.shape} " 

618 f"after multiplying by the secondary matrix.\n" 

619 ) 

620 else: 

621 blockVectorBY = blockVectorY 

622 

623 # gramYBY is a dense array. 

624 gramYBY = blockVectorY.T.conj() @ blockVectorBY 

625 try: 

626 # gramYBY is a Cholesky factor from now on... 

627 gramYBY = cho_factor(gramYBY, overwrite_a=True) 

628 except LinAlgError as e: 

629 raise ValueError("Linearly dependent constraints") from e 

630 

631 _applyConstraints(blockVectorX, gramYBY, blockVectorBY, blockVectorY) 

632 

633 ## 

634 # B-orthonormalize X. 

635 blockVectorX, blockVectorBX, _ = _b_orthonormalize( 

636 B, blockVectorX, verbosityLevel=verbosityLevel) 

637 if blockVectorX is None: 

638 raise ValueError("Linearly dependent initial approximations") 

639 

640 ## 

641 # Compute the initial Ritz vectors: solve the eigenproblem. 

642 blockVectorAX = A(blockVectorX) 

643 if blockVectorAX.shape != blockVectorX.shape: 

644 raise ValueError( 

645 f"The shape {blockVectorX.shape} " 

646 f"of the initial approximations not preserved\n" 

647 f"and changed to {blockVectorAX.shape} " 

648 f"after multiplying by the primary matrix.\n" 

649 ) 

650 

651 gramXAX = blockVectorX.T.conj() @ blockVectorAX 

652 

653 _lambda, eigBlockVector = eigh(gramXAX, check_finite=False) 

654 ii = _get_indx(_lambda, sizeX, largest) 

655 _lambda = _lambda[ii] 

656 if retLambdaHistory: 

657 lambdaHistory[0, :] = _lambda 

658 

659 eigBlockVector = np.asarray(eigBlockVector[:, ii]) 

660 blockVectorX = _matmul_inplace( 

661 blockVectorX, eigBlockVector, 

662 verbosityLevel=verbosityLevel 

663 ) 

664 blockVectorAX = _matmul_inplace( 

665 blockVectorAX, eigBlockVector, 

666 verbosityLevel=verbosityLevel 

667 ) 

668 if B is not None: 

669 blockVectorBX = _matmul_inplace( 

670 blockVectorBX, eigBlockVector, 

671 verbosityLevel=verbosityLevel 

672 ) 

673 

674 ## 

675 # Active index set. 

676 activeMask = np.ones((sizeX,), dtype=bool) 

677 

678 ## 

679 # Main iteration loop. 

680 

681 blockVectorP = None # set during iteration 

682 blockVectorAP = None 

683 blockVectorBP = None 

684 

685 smallestResidualNorm = np.abs(np.finfo(blockVectorX.dtype).max) 

686 

687 iterationNumber = -1 

688 restart = True 

689 forcedRestart = False 

690 explicitGramFlag = False 

691 while iterationNumber < maxiter: 

692 iterationNumber += 1 

693 

694 if B is not None: 

695 aux = blockVectorBX * _lambda[np.newaxis, :] 

696 else: 

697 aux = blockVectorX * _lambda[np.newaxis, :] 

698 

699 blockVectorR = blockVectorAX - aux 

700 

701 aux = np.sum(blockVectorR.conj() * blockVectorR, 0) 

702 residualNorms = np.sqrt(np.abs(aux)) 

703 if retResidualNormsHistory: 

704 residualNormsHistory[iterationNumber, :] = residualNorms 

705 residualNorm = np.sum(np.abs(residualNorms)) / sizeX 

706 

707 if residualNorm < smallestResidualNorm: 

708 smallestResidualNorm = residualNorm 

709 bestIterationNumber = iterationNumber 

710 bestblockVectorX = blockVectorX 

711 elif residualNorm > 2**restartControl * smallestResidualNorm: 

712 forcedRestart = True 

713 blockVectorAX = A(blockVectorX) 

714 if blockVectorAX.shape != blockVectorX.shape: 

715 raise ValueError( 

716 f"The shape {blockVectorX.shape} " 

717 f"of the restarted iterate not preserved\n" 

718 f"and changed to {blockVectorAX.shape} " 

719 f"after multiplying by the primary matrix.\n" 

720 ) 

721 if B is not None: 

722 blockVectorBX = B(blockVectorX) 

723 if blockVectorBX.shape != blockVectorX.shape: 

724 raise ValueError( 

725 f"The shape {blockVectorX.shape} " 

726 f"of the restarted iterate not preserved\n" 

727 f"and changed to {blockVectorBX.shape} " 

728 f"after multiplying by the secondary matrix.\n" 

729 ) 

730 

731 ii = np.where(residualNorms > residualTolerance, True, False) 

732 activeMask = activeMask & ii 

733 currentBlockSize = activeMask.sum() 

734 

735 if verbosityLevel: 

736 print(f"iteration {iterationNumber}") 

737 print(f"current block size: {currentBlockSize}") 

738 print(f"eigenvalue(s):\n{_lambda}") 

739 print(f"residual norm(s):\n{residualNorms}") 

740 

741 if currentBlockSize == 0: 

742 break 

743 

744 activeBlockVectorR = _as2d(blockVectorR[:, activeMask]) 

745 

746 if iterationNumber > 0: 

747 activeBlockVectorP = _as2d(blockVectorP[:, activeMask]) 

748 activeBlockVectorAP = _as2d(blockVectorAP[:, activeMask]) 

749 if B is not None: 

750 activeBlockVectorBP = _as2d(blockVectorBP[:, activeMask]) 

751 

752 if M is not None: 

753 # Apply preconditioner T to the active residuals. 

754 activeBlockVectorR = M(activeBlockVectorR) 

755 

756 ## 

757 # Apply constraints to the preconditioned residuals. 

758 if blockVectorY is not None: 

759 _applyConstraints(activeBlockVectorR, 

760 gramYBY, 

761 blockVectorBY, 

762 blockVectorY) 

763 

764 ## 

765 # B-orthogonalize the preconditioned residuals to X. 

766 if B is not None: 

767 activeBlockVectorR = activeBlockVectorR - ( 

768 blockVectorX @ 

769 (blockVectorBX.T.conj() @ activeBlockVectorR) 

770 ) 

771 else: 

772 activeBlockVectorR = activeBlockVectorR - ( 

773 blockVectorX @ 

774 (blockVectorX.T.conj() @ activeBlockVectorR) 

775 ) 

776 

777 ## 

778 # B-orthonormalize the preconditioned residuals. 

779 aux = _b_orthonormalize( 

780 B, activeBlockVectorR, verbosityLevel=verbosityLevel) 

781 activeBlockVectorR, activeBlockVectorBR, _ = aux 

782 

783 if activeBlockVectorR is None: 

784 warnings.warn( 

785 f"Failed at iteration {iterationNumber} with accuracies " 

786 f"{residualNorms}\n not reaching the requested " 

787 f"tolerance {residualTolerance}.", 

788 UserWarning, stacklevel=2 

789 ) 

790 break 

791 activeBlockVectorAR = A(activeBlockVectorR) 

792 

793 if iterationNumber > 0: 

794 if B is not None: 

795 aux = _b_orthonormalize( 

796 B, activeBlockVectorP, activeBlockVectorBP, 

797 verbosityLevel=verbosityLevel 

798 ) 

799 activeBlockVectorP, activeBlockVectorBP, invR = aux 

800 else: 

801 aux = _b_orthonormalize(B, activeBlockVectorP, 

802 verbosityLevel=verbosityLevel) 

803 activeBlockVectorP, _, invR = aux 

804 # Function _b_orthonormalize returns None if Cholesky fails 

805 if activeBlockVectorP is not None: 

806 activeBlockVectorAP = _matmul_inplace( 

807 activeBlockVectorAP, invR, 

808 verbosityLevel=verbosityLevel 

809 ) 

810 restart = forcedRestart 

811 else: 

812 restart = True 

813 

814 ## 

815 # Perform the Rayleigh Ritz Procedure: 

816 # Compute symmetric Gram matrices: 

817 

818 if activeBlockVectorAR.dtype == "float32": 

819 myeps = 1 

820 else: 

821 myeps = np.sqrt(np.finfo(activeBlockVectorR.dtype).eps) 

822 

823 if residualNorms.max() > myeps and not explicitGramFlag: 

824 explicitGramFlag = False 

825 else: 

826 # Once explicitGramFlag, forever explicitGramFlag. 

827 explicitGramFlag = True 

828 

829 # Shared memory assingments to simplify the code 

830 if B is None: 

831 blockVectorBX = blockVectorX 

832 activeBlockVectorBR = activeBlockVectorR 

833 if not restart: 

834 activeBlockVectorBP = activeBlockVectorP 

835 

836 # Common submatrices: 

837 gramXAR = np.dot(blockVectorX.T.conj(), activeBlockVectorAR) 

838 gramRAR = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorAR) 

839 

840 gramDtype = activeBlockVectorAR.dtype 

841 if explicitGramFlag: 

842 gramRAR = (gramRAR + gramRAR.T.conj()) / 2 

843 gramXAX = np.dot(blockVectorX.T.conj(), blockVectorAX) 

844 gramXAX = (gramXAX + gramXAX.T.conj()) / 2 

845 gramXBX = np.dot(blockVectorX.T.conj(), blockVectorBX) 

846 gramRBR = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorBR) 

847 gramXBR = np.dot(blockVectorX.T.conj(), activeBlockVectorBR) 

848 else: 

849 gramXAX = np.diag(_lambda).astype(gramDtype) 

850 gramXBX = np.eye(sizeX, dtype=gramDtype) 

851 gramRBR = np.eye(currentBlockSize, dtype=gramDtype) 

852 gramXBR = np.zeros((sizeX, currentBlockSize), dtype=gramDtype) 

853 

854 if not restart: 

855 gramXAP = np.dot(blockVectorX.T.conj(), activeBlockVectorAP) 

856 gramRAP = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorAP) 

857 gramPAP = np.dot(activeBlockVectorP.T.conj(), activeBlockVectorAP) 

858 gramXBP = np.dot(blockVectorX.T.conj(), activeBlockVectorBP) 

859 gramRBP = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorBP) 

860 if explicitGramFlag: 

861 gramPAP = (gramPAP + gramPAP.T.conj()) / 2 

862 gramPBP = np.dot(activeBlockVectorP.T.conj(), 

863 activeBlockVectorBP) 

864 else: 

865 gramPBP = np.eye(currentBlockSize, dtype=gramDtype) 

866 

867 gramA = np.block( 

868 [ 

869 [gramXAX, gramXAR, gramXAP], 

870 [gramXAR.T.conj(), gramRAR, gramRAP], 

871 [gramXAP.T.conj(), gramRAP.T.conj(), gramPAP], 

872 ] 

873 ) 

874 gramB = np.block( 

875 [ 

876 [gramXBX, gramXBR, gramXBP], 

877 [gramXBR.T.conj(), gramRBR, gramRBP], 

878 [gramXBP.T.conj(), gramRBP.T.conj(), gramPBP], 

879 ] 

880 ) 

881 

882 _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel) 

883 

884 try: 

885 _lambda, eigBlockVector = eigh(gramA, 

886 gramB, 

887 check_finite=False) 

888 except LinAlgError as e: 

889 # raise ValueError("eigh failed in lobpcg iterations") from e 

890 if verbosityLevel: 

891 warnings.warn( 

892 f"eigh failed at iteration {iterationNumber} \n" 

893 f"with error {e} causing a restart.\n", 

894 UserWarning, stacklevel=2 

895 ) 

896 # try again after dropping the direction vectors P from RR 

897 restart = True 

898 

899 if restart: 

900 gramA = np.block([[gramXAX, gramXAR], [gramXAR.T.conj(), gramRAR]]) 

901 gramB = np.block([[gramXBX, gramXBR], [gramXBR.T.conj(), gramRBR]]) 

902 

903 _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel) 

904 

905 try: 

906 _lambda, eigBlockVector = eigh(gramA, 

907 gramB, 

908 check_finite=False) 

909 except LinAlgError as e: 

910 # raise ValueError("eigh failed in lobpcg iterations") from e 

911 warnings.warn( 

912 f"eigh failed at iteration {iterationNumber} with error\n" 

913 f"{e}\n", 

914 UserWarning, stacklevel=2 

915 ) 

916 break 

917 

918 ii = _get_indx(_lambda, sizeX, largest) 

919 _lambda = _lambda[ii] 

920 eigBlockVector = eigBlockVector[:, ii] 

921 if retLambdaHistory: 

922 lambdaHistory[iterationNumber + 1, :] = _lambda 

923 

924 # Compute Ritz vectors. 

925 if B is not None: 

926 if not restart: 

927 eigBlockVectorX = eigBlockVector[:sizeX] 

928 eigBlockVectorR = eigBlockVector[sizeX: 

929 sizeX + currentBlockSize] 

930 eigBlockVectorP = eigBlockVector[sizeX + currentBlockSize:] 

931 

932 pp = np.dot(activeBlockVectorR, eigBlockVectorR) 

933 pp += np.dot(activeBlockVectorP, eigBlockVectorP) 

934 

935 app = np.dot(activeBlockVectorAR, eigBlockVectorR) 

936 app += np.dot(activeBlockVectorAP, eigBlockVectorP) 

937 

938 bpp = np.dot(activeBlockVectorBR, eigBlockVectorR) 

939 bpp += np.dot(activeBlockVectorBP, eigBlockVectorP) 

940 else: 

941 eigBlockVectorX = eigBlockVector[:sizeX] 

942 eigBlockVectorR = eigBlockVector[sizeX:] 

943 

944 pp = np.dot(activeBlockVectorR, eigBlockVectorR) 

945 app = np.dot(activeBlockVectorAR, eigBlockVectorR) 

946 bpp = np.dot(activeBlockVectorBR, eigBlockVectorR) 

947 

948 blockVectorX = np.dot(blockVectorX, eigBlockVectorX) + pp 

949 blockVectorAX = np.dot(blockVectorAX, eigBlockVectorX) + app 

950 blockVectorBX = np.dot(blockVectorBX, eigBlockVectorX) + bpp 

951 

952 blockVectorP, blockVectorAP, blockVectorBP = pp, app, bpp 

953 

954 else: 

955 if not restart: 

956 eigBlockVectorX = eigBlockVector[:sizeX] 

957 eigBlockVectorR = eigBlockVector[sizeX: