lobpcg.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982
  1. """
  2. Locally Optimal Block Preconditioned Conjugate Gradient Method (LOBPCG).
  3. References
  4. ----------
  5. .. [1] A. V. Knyazev (2001),
  6. Toward the Optimal Preconditioned Eigensolver: Locally Optimal
  7. Block Preconditioned Conjugate Gradient Method.
  8. SIAM Journal on Scientific Computing 23, no. 2,
  9. pp. 517-541. :doi:`10.1137/S1064827500366124`
  10. .. [2] A. V. Knyazev, I. Lashuk, M. E. Argentati, and E. Ovchinnikov (2007),
  11. Block Locally Optimal Preconditioned Eigenvalue Xolvers (BLOPEX)
  12. in hypre and PETSc. :arxiv:`0705.2626`
  13. .. [3] A. V. Knyazev's C and MATLAB implementations:
  14. https://github.com/lobpcg/blopex
  15. """
  16. import warnings
  17. import numpy as np
  18. from scipy.linalg import (inv, eigh, cho_factor, cho_solve,
  19. cholesky, LinAlgError)
  20. from scipy.sparse.linalg import LinearOperator
  21. from scipy.sparse import isspmatrix
  22. from numpy import block as bmat
  23. __all__ = ["lobpcg"]
  24. def _report_nonhermitian(M, name):
  25. """
  26. Report if `M` is not a Hermitian matrix given its type.
  27. """
  28. from scipy.linalg import norm
  29. md = M - M.T.conj()
  30. nmd = norm(md, 1)
  31. tol = 10 * np.finfo(M.dtype).eps
  32. tol = max(tol, tol * norm(M, 1))
  33. if nmd > tol:
  34. warnings.warn(
  35. f"Matrix {name} of the type {M.dtype} is not Hermitian: "
  36. f"condition: {nmd} < {tol} fails.",
  37. UserWarning, stacklevel=4
  38. )
  39. def _as2d(ar):
  40. """
  41. If the input array is 2D return it, if it is 1D, append a dimension,
  42. making it a column vector.
  43. """
  44. if ar.ndim == 2:
  45. return ar
  46. else: # Assume 1!
  47. aux = np.array(ar, copy=False)
  48. aux.shape = (ar.shape[0], 1)
  49. return aux
  50. def _makeMatMat(m):
  51. if m is None:
  52. return None
  53. elif callable(m):
  54. return lambda v: m(v)
  55. else:
  56. return lambda v: m @ v
  57. def _applyConstraints(blockVectorV, factYBY, blockVectorBY, blockVectorY):
  58. """Changes blockVectorV in place."""
  59. YBV = np.dot(blockVectorBY.T.conj(), blockVectorV)
  60. tmp = cho_solve(factYBY, YBV)
  61. blockVectorV -= np.dot(blockVectorY, tmp)
  62. def _b_orthonormalize(B, blockVectorV, blockVectorBV=None,
  63. verbosityLevel=0):
  64. """in-place B-orthonormalize the given block vector using Cholesky."""
  65. normalization = blockVectorV.max(axis=0) + np.finfo(blockVectorV.dtype).eps
  66. blockVectorV = blockVectorV / normalization
  67. if blockVectorBV is None:
  68. if B is not None:
  69. try:
  70. blockVectorBV = B(blockVectorV)
  71. except Exception as e:
  72. if verbosityLevel:
  73. warnings.warn(
  74. f"Secondary MatMul call failed with error\n"
  75. f"{e}\n",
  76. UserWarning, stacklevel=3
  77. )
  78. return None, None, None, normalization
  79. if blockVectorBV.shape != blockVectorV.shape:
  80. raise ValueError(
  81. f"The shape {blockVectorV.shape} "
  82. f"of the orthogonalized matrix not preserved\n"
  83. f"and changed to {blockVectorBV.shape} "
  84. f"after multiplying by the secondary matrix.\n"
  85. )
  86. else:
  87. blockVectorBV = blockVectorV # Shared data!!!
  88. else:
  89. blockVectorBV = blockVectorBV / normalization
  90. VBV = blockVectorV.T.conj() @ blockVectorBV
  91. try:
  92. # VBV is a Cholesky factor from now on...
  93. VBV = cholesky(VBV, overwrite_a=True)
  94. VBV = inv(VBV, overwrite_a=True)
  95. blockVectorV = blockVectorV @ VBV
  96. # blockVectorV = (cho_solve((VBV.T, True), blockVectorV.T)).T
  97. if B is not None:
  98. blockVectorBV = blockVectorBV @ VBV
  99. # blockVectorBV = (cho_solve((VBV.T, True), blockVectorBV.T)).T
  100. return blockVectorV, blockVectorBV, VBV, normalization
  101. except LinAlgError:
  102. if verbosityLevel:
  103. warnings.warn(
  104. "Cholesky has failed.",
  105. UserWarning, stacklevel=3
  106. )
  107. return None, None, None, normalization
  108. def _get_indx(_lambda, num, largest):
  109. """Get `num` indices into `_lambda` depending on `largest` option."""
  110. ii = np.argsort(_lambda)
  111. if largest:
  112. ii = ii[:-num - 1:-1]
  113. else:
  114. ii = ii[:num]
  115. return ii
  116. def _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel):
  117. if verbosityLevel:
  118. _report_nonhermitian(gramA, "gramA")
  119. _report_nonhermitian(gramB, "gramB")
  120. def lobpcg(
  121. A,
  122. X,
  123. B=None,
  124. M=None,
  125. Y=None,
  126. tol=None,
  127. maxiter=None,
  128. largest=True,
  129. verbosityLevel=0,
  130. retLambdaHistory=False,
  131. retResidualNormsHistory=False,
  132. restartControl=20,
  133. ):
  134. """Locally Optimal Block Preconditioned Conjugate Gradient Method (LOBPCG).
  135. LOBPCG is a preconditioned eigensolver for large symmetric positive
  136. definite (SPD) generalized eigenproblems.
  137. Parameters
  138. ----------
  139. A : {sparse matrix, dense matrix, LinearOperator, callable object}
  140. The symmetric linear operator of the problem, usually a
  141. sparse matrix. Often called the "stiffness matrix".
  142. X : ndarray, float32 or float64
  143. Initial approximation to the ``k`` eigenvectors (non-sparse). If `A`
  144. has ``shape=(n,n)`` then `X` should have shape ``shape=(n,k)``.
  145. B : {dense matrix, sparse matrix, LinearOperator, callable object}
  146. Optional.
  147. The right hand side operator in a generalized eigenproblem.
  148. By default, ``B = Identity``. Often called the "mass matrix".
  149. M : {dense matrix, sparse matrix, LinearOperator, callable object}
  150. Optional.
  151. Preconditioner to `A`; by default ``M = Identity``.
  152. `M` should approximate the inverse of `A`.
  153. Y : ndarray, float32 or float64, optional.
  154. An n-by-sizeY matrix of constraints (non-sparse), sizeY < n.
  155. The iterations will be performed in the B-orthogonal complement
  156. of the column-space of Y. Y must be full rank.
  157. tol : scalar, optional.
  158. Solver tolerance (stopping criterion).
  159. The default is ``tol=n*sqrt(eps)``.
  160. maxiter : int, optional.
  161. Maximum number of iterations. The default is ``maxiter=20``.
  162. largest : bool, optional.
  163. When True, solve for the largest eigenvalues, otherwise the smallest.
  164. verbosityLevel : int, optional
  165. Controls solver output. The default is ``verbosityLevel=0``.
  166. retLambdaHistory : bool, optional.
  167. Whether to return eigenvalue history. Default is False.
  168. retResidualNormsHistory : bool, optional.
  169. Whether to return history of residual norms. Default is False.
  170. restartControl : int, optional.
  171. Iterations restart if the residuals jump up 2**restartControl times
  172. compared to the smallest ones recorded in retResidualNormsHistory.
  173. The default is ``restartControl=20``, making the restarts rare for
  174. backward compatibility.
  175. Returns
  176. -------
  177. w : ndarray
  178. Array of ``k`` eigenvalues.
  179. v : ndarray
  180. An array of ``k`` eigenvectors. `v` has the same shape as `X`.
  181. lambdas : ndarray, optional
  182. The eigenvalue history, if `retLambdaHistory` is True.
  183. rnorms : ndarray, optional
  184. The history of residual norms, if `retResidualNormsHistory` is True.
  185. Notes
  186. -----
  187. The iterative loop in lobpcg runs maxit=maxiter (or 20 if maxit=None)
  188. iterations at most and finishes earler if the tolerance is met.
  189. Breaking backward compatibility with the previous version, lobpcg
  190. now returns the block of iterative vectors with the best accuracy rather
  191. than the last one iterated, as a cure for possible divergence.
  192. The size of the iteration history output equals to the number of the best
  193. (limited by maxit) iterations plus 3 (initial, final, and postprocessing).
  194. If both ``retLambdaHistory`` and ``retResidualNormsHistory`` are True,
  195. the return tuple has the following format
  196. ``(lambda, V, lambda history, residual norms history)``.
  197. In the following ``n`` denotes the matrix size and ``k`` the number
  198. of required eigenvalues (smallest or largest).
  199. The LOBPCG code internally solves eigenproblems of the size ``3k`` on every
  200. iteration by calling the dense eigensolver `eigh`, so if ``k`` is not
  201. small enough compared to ``n``, it makes no sense to call the LOBPCG code.
  202. Moreover, if one calls the LOBPCG algorithm for ``5k > n``, it would likely
  203. break internally, so the code calls the standard function `eigh` instead.
  204. It is not that ``n`` should be large for the LOBPCG to work, but rather the
  205. ratio ``n / k`` should be large. It you call LOBPCG with ``k=1``
  206. and ``n=10``, it works though ``n`` is small. The method is intended
  207. for extremely large ``n / k``.
  208. The convergence speed depends basically on two factors:
  209. 1. Relative separation of the seeking eigenvalues from the rest
  210. of the eigenvalues. One can vary ``k`` to improve the absolute
  211. separation and use proper preconditioning to shrink the spectral spread.
  212. For example, a rod vibration test problem (under tests
  213. directory) is ill-conditioned for large ``n``, so convergence will be
  214. slow, unless efficient preconditioning is used. For this specific
  215. problem, a good simple preconditioner function would be a linear solve
  216. for `A`, which is easy to code since `A` is tridiagonal.
  217. 2. Quality of the initial approximations `X` to the seeking eigenvectors.
  218. Randomly distributed around the origin vectors work well if no better
  219. choice is known.
  220. References
  221. ----------
  222. .. [1] A. V. Knyazev (2001),
  223. Toward the Optimal Preconditioned Eigensolver: Locally Optimal
  224. Block Preconditioned Conjugate Gradient Method.
  225. SIAM Journal on Scientific Computing 23, no. 2,
  226. pp. 517-541. :doi:`10.1137/S1064827500366124`
  227. .. [2] A. V. Knyazev, I. Lashuk, M. E. Argentati, and E. Ovchinnikov
  228. (2007), Block Locally Optimal Preconditioned Eigenvalue Xolvers
  229. (BLOPEX) in hypre and PETSc. :arxiv:`0705.2626`
  230. .. [3] A. V. Knyazev's C and MATLAB implementations:
  231. https://github.com/lobpcg/blopex
  232. Examples
  233. --------
  234. Solve ``A x = lambda x`` with constraints and preconditioning.
  235. >>> import numpy as np
  236. >>> from scipy.sparse import spdiags, issparse
  237. >>> from scipy.sparse.linalg import lobpcg, LinearOperator
  238. The square matrix size:
  239. >>> n = 100
  240. >>> vals = np.arange(1, n + 1)
  241. The first mandatory input parameter, in this test
  242. a sparse 2D array representing the square matrix
  243. of the eigenvalue problem to solve:
  244. >>> A = spdiags(vals, 0, n, n)
  245. >>> A.toarray()
  246. array([[ 1, 0, 0, ..., 0, 0, 0],
  247. [ 0, 2, 0, ..., 0, 0, 0],
  248. [ 0, 0, 3, ..., 0, 0, 0],
  249. ...,
  250. [ 0, 0, 0, ..., 98, 0, 0],
  251. [ 0, 0, 0, ..., 0, 99, 0],
  252. [ 0, 0, 0, ..., 0, 0, 100]])
  253. Initial guess for eigenvectors, should have linearly independent
  254. columns. The second mandatory input parameter, a 2D array with the
  255. row dimension determining the number of requested eigenvalues.
  256. If no initial approximations available, randomly oriented vectors
  257. commonly work best, e.g., with components normally disrtibuted
  258. around zero or uniformly distributed on the interval [-1 1].
  259. >>> rng = np.random.default_rng()
  260. >>> X = rng.normal(size=(n, 3))
  261. Constraints - an optional input parameter is a 2D array comprising
  262. of column vectors that the eigenvectors must be orthogonal to:
  263. >>> Y = np.eye(n, 3)
  264. Preconditioner in the inverse of A in this example:
  265. >>> invA = spdiags([1./vals], 0, n, n)
  266. The preconditiner must be defined by a function:
  267. >>> def precond( x ):
  268. ... return invA @ x
  269. The argument x of the preconditioner function is a matrix inside `lobpcg`,
  270. thus the use of matrix-matrix product ``@``.
  271. The preconditioner function is passed to lobpcg as a `LinearOperator`:
  272. >>> M = LinearOperator(matvec=precond, matmat=precond,
  273. ... shape=(n, n), dtype=np.float64)
  274. Let us now solve the eigenvalue problem for the matrix A:
  275. >>> eigenvalues, _ = lobpcg(A, X, Y=Y, M=M, largest=False)
  276. >>> eigenvalues
  277. array([4., 5., 6.])
  278. Note that the vectors passed in Y are the eigenvectors of the 3 smallest
  279. eigenvalues. The results returned are orthogonal to those.
  280. """
  281. blockVectorX = X
  282. bestblockVectorX = blockVectorX
  283. blockVectorY = Y
  284. residualTolerance = tol
  285. if maxiter is None:
  286. maxiter = 20
  287. bestIterationNumber = maxiter
  288. sizeY = 0
  289. if blockVectorY is not None:
  290. if len(blockVectorY.shape) != 2:
  291. warnings.warn(
  292. f"Expected rank-2 array for argument Y, instead got "
  293. f"{len(blockVectorY.shape)}, "
  294. f"so ignore it and use no constraints.",
  295. UserWarning, stacklevel=2
  296. )
  297. blockVectorY = None
  298. else:
  299. sizeY = blockVectorY.shape[1]
  300. # Block size.
  301. if blockVectorX is None:
  302. raise ValueError("The mandatory initial matrix X cannot be None")
  303. if len(blockVectorX.shape) != 2:
  304. raise ValueError("expected rank-2 array for argument X")
  305. n, sizeX = blockVectorX.shape
  306. # Data type of iterates, determined by X, must be inexact
  307. if not np.issubdtype(blockVectorX.dtype, np.inexact):
  308. warnings.warn(
  309. f"Data type for argument X is {blockVectorX.dtype}, "
  310. f"which is not inexact, so casted to np.float32.",
  311. UserWarning, stacklevel=2
  312. )
  313. blockVectorX = np.asarray(blockVectorX, dtype=np.float32)
  314. if retLambdaHistory:
  315. lambdaHistory = np.zeros((maxiter + 3, sizeX),
  316. dtype=blockVectorX.dtype)
  317. if retResidualNormsHistory:
  318. residualNormsHistory = np.zeros((maxiter + 3, sizeX),
  319. dtype=blockVectorX.dtype)
  320. if verbosityLevel:
  321. aux = "Solving "
  322. if B is None:
  323. aux += "standard"
  324. else:
  325. aux += "generalized"
  326. aux += " eigenvalue problem with"
  327. if M is None:
  328. aux += "out"
  329. aux += " preconditioning\n\n"
  330. aux += "matrix size %d\n" % n
  331. aux += "block size %d\n\n" % sizeX
  332. if blockVectorY is None:
  333. aux += "No constraints\n\n"
  334. else:
  335. if sizeY > 1:
  336. aux += "%d constraints\n\n" % sizeY
  337. else:
  338. aux += "%d constraint\n\n" % sizeY
  339. print(aux)
  340. if (n - sizeY) < (5 * sizeX):
  341. warnings.warn(
  342. f"The problem size {n} minus the constraints size {sizeY} "
  343. f"is too small relative to the block size {sizeX}. "
  344. f"Using a dense eigensolver instead of LOBPCG iterations."
  345. f"No output of the history of the iterations.",
  346. UserWarning, stacklevel=2
  347. )
  348. sizeX = min(sizeX, n)
  349. if blockVectorY is not None:
  350. raise NotImplementedError(
  351. "The dense eigensolver does not support constraints."
  352. )
  353. # Define the closed range of indices of eigenvalues to return.
  354. if largest:
  355. eigvals = (n - sizeX, n - 1)
  356. else:
  357. eigvals = (0, sizeX - 1)
  358. try:
  359. if isinstance(A, LinearOperator):
  360. A = A(np.eye(n, dtype=int))
  361. elif callable(A):
  362. A = A(np.eye(n, dtype=int))
  363. if A.shape != (n, n):
  364. raise ValueError(
  365. f"The shape {A.shape} of the primary matrix\n"
  366. f"defined by a callable object is wrong.\n"
  367. )
  368. elif isspmatrix(A):
  369. A = A.toarray()
  370. else:
  371. A = np.asarray(A)
  372. except Exception as e:
  373. raise Exception(
  374. f"Primary MatMul call failed with error\n"
  375. f"{e}\n")
  376. if B is not None:
  377. try:
  378. if isinstance(B, LinearOperator):
  379. B = B(np.eye(n, dtype=int))
  380. elif callable(B):
  381. B = B(np.eye(n, dtype=int))
  382. if B.shape != (n, n):
  383. raise ValueError(
  384. f"The shape {B.shape} of the secondary matrix\n"
  385. f"defined by a callable object is wrong.\n"
  386. )
  387. elif isspmatrix(B):
  388. B = B.toarray()
  389. else:
  390. B = np.asarray(B)
  391. except Exception as e:
  392. raise Exception(
  393. f"Secondary MatMul call failed with error\n"
  394. f"{e}\n")
  395. try:
  396. vals, vecs = eigh(A,
  397. B,
  398. subset_by_index=eigvals,
  399. check_finite=False)
  400. if largest:
  401. # Reverse order to be compatible with eigs() in 'LM' mode.
  402. vals = vals[::-1]
  403. vecs = vecs[:, ::-1]
  404. return vals, vecs
  405. except Exception as e:
  406. raise Exception(
  407. f"Dense eigensolver failed with error\n"
  408. f"{e}\n"
  409. )
  410. if (residualTolerance is None) or (residualTolerance <= 0.0):
  411. residualTolerance = np.sqrt(np.finfo(blockVectorX.dtype).eps) * n
  412. A = _makeMatMat(A)
  413. B = _makeMatMat(B)
  414. M = _makeMatMat(M)
  415. # Apply constraints to X.
  416. if blockVectorY is not None:
  417. if B is not None:
  418. blockVectorBY = B(blockVectorY)
  419. if blockVectorBY.shape != blockVectorY.shape:
  420. raise ValueError(
  421. f"The shape {blockVectorY.shape} "
  422. f"of the constraint not preserved\n"
  423. f"and changed to {blockVectorBY.shape} "
  424. f"after multiplying by the secondary matrix.\n"
  425. )
  426. else:
  427. blockVectorBY = blockVectorY
  428. # gramYBY is a dense array.
  429. gramYBY = np.dot(blockVectorY.T.conj(), blockVectorBY)
  430. try:
  431. # gramYBY is a Cholesky factor from now on...
  432. gramYBY = cho_factor(gramYBY)
  433. except LinAlgError as e:
  434. raise ValueError("Linearly dependent constraints") from e
  435. _applyConstraints(blockVectorX, gramYBY, blockVectorBY, blockVectorY)
  436. ##
  437. # B-orthonormalize X.
  438. blockVectorX, blockVectorBX, _, _ = _b_orthonormalize(
  439. B, blockVectorX, verbosityLevel=verbosityLevel)
  440. if blockVectorX is None:
  441. raise ValueError("Linearly dependent initial approximations")
  442. ##
  443. # Compute the initial Ritz vectors: solve the eigenproblem.
  444. blockVectorAX = A(blockVectorX)
  445. if blockVectorAX.shape != blockVectorX.shape:
  446. raise ValueError(
  447. f"The shape {blockVectorX.shape} "
  448. f"of the initial approximations not preserved\n"
  449. f"and changed to {blockVectorAX.shape} "
  450. f"after multiplying by the primary matrix.\n"
  451. )
  452. gramXAX = np.dot(blockVectorX.T.conj(), blockVectorAX)
  453. _lambda, eigBlockVector = eigh(gramXAX, check_finite=False)
  454. ii = _get_indx(_lambda, sizeX, largest)
  455. _lambda = _lambda[ii]
  456. if retLambdaHistory:
  457. lambdaHistory[0, :] = _lambda
  458. eigBlockVector = np.asarray(eigBlockVector[:, ii])
  459. blockVectorX = np.dot(blockVectorX, eigBlockVector)
  460. blockVectorAX = np.dot(blockVectorAX, eigBlockVector)
  461. if B is not None:
  462. blockVectorBX = np.dot(blockVectorBX, eigBlockVector)
  463. ##
  464. # Active index set.
  465. activeMask = np.ones((sizeX,), dtype=bool)
  466. ##
  467. # Main iteration loop.
  468. blockVectorP = None # set during iteration
  469. blockVectorAP = None
  470. blockVectorBP = None
  471. smallestResidualNorm = np.abs(np.finfo(blockVectorX.dtype).max)
  472. iterationNumber = -1
  473. restart = True
  474. forcedRestart = False
  475. explicitGramFlag = False
  476. while iterationNumber < maxiter:
  477. iterationNumber += 1
  478. if B is not None:
  479. aux = blockVectorBX * _lambda[np.newaxis, :]
  480. else:
  481. aux = blockVectorX * _lambda[np.newaxis, :]
  482. blockVectorR = blockVectorAX - aux
  483. aux = np.sum(blockVectorR.conj() * blockVectorR, 0)
  484. residualNorms = np.sqrt(np.abs(aux))
  485. if retResidualNormsHistory:
  486. residualNormsHistory[iterationNumber, :] = residualNorms
  487. residualNorm = np.sum(np.abs(residualNorms)) / sizeX
  488. if residualNorm < smallestResidualNorm:
  489. smallestResidualNorm = residualNorm
  490. bestIterationNumber = iterationNumber
  491. bestblockVectorX = blockVectorX
  492. elif residualNorm > 2**restartControl * smallestResidualNorm:
  493. forcedRestart = True
  494. blockVectorAX = A(blockVectorX)
  495. if blockVectorAX.shape != blockVectorX.shape:
  496. raise ValueError(
  497. f"The shape {blockVectorX.shape} "
  498. f"of the restarted iterate not preserved\n"
  499. f"and changed to {blockVectorAX.shape} "
  500. f"after multiplying by the primary matrix.\n"
  501. )
  502. if B is not None:
  503. blockVectorBX = B(blockVectorX)
  504. if blockVectorBX.shape != blockVectorX.shape:
  505. raise ValueError(
  506. f"The shape {blockVectorX.shape} "
  507. f"of the restarted iterate not preserved\n"
  508. f"and changed to {blockVectorBX.shape} "
  509. f"after multiplying by the secondary matrix.\n"
  510. )
  511. ii = np.where(residualNorms > residualTolerance, True, False)
  512. activeMask = activeMask & ii
  513. currentBlockSize = activeMask.sum()
  514. if verbosityLevel:
  515. print(f"iteration {iterationNumber}")
  516. print(f"current block size: {currentBlockSize}")
  517. print(f"eigenvalue(s):\n{_lambda}")
  518. print(f"residual norm(s):\n{residualNorms}")
  519. if currentBlockSize == 0:
  520. break
  521. activeBlockVectorR = _as2d(blockVectorR[:, activeMask])
  522. if iterationNumber > 0:
  523. activeBlockVectorP = _as2d(blockVectorP[:, activeMask])
  524. activeBlockVectorAP = _as2d(blockVectorAP[:, activeMask])
  525. if B is not None:
  526. activeBlockVectorBP = _as2d(blockVectorBP[:, activeMask])
  527. if M is not None:
  528. # Apply preconditioner T to the active residuals.
  529. activeBlockVectorR = M(activeBlockVectorR)
  530. ##
  531. # Apply constraints to the preconditioned residuals.
  532. if blockVectorY is not None:
  533. _applyConstraints(activeBlockVectorR,
  534. gramYBY,
  535. blockVectorBY,
  536. blockVectorY)
  537. ##
  538. # B-orthogonalize the preconditioned residuals to X.
  539. if B is not None:
  540. activeBlockVectorR = activeBlockVectorR - (
  541. blockVectorX @
  542. (blockVectorBX.T.conj() @ activeBlockVectorR)
  543. )
  544. else:
  545. activeBlockVectorR = activeBlockVectorR - (
  546. blockVectorX @
  547. (blockVectorX.T.conj() @ activeBlockVectorR)
  548. )
  549. ##
  550. # B-orthonormalize the preconditioned residuals.
  551. aux = _b_orthonormalize(
  552. B, activeBlockVectorR, verbosityLevel=verbosityLevel)
  553. activeBlockVectorR, activeBlockVectorBR, _, _ = aux
  554. if activeBlockVectorR is None:
  555. warnings.warn(
  556. f"Failed at iteration {iterationNumber} with accuracies "
  557. f"{residualNorms}\n not reaching the requested "
  558. f"tolerance {residualTolerance}.",
  559. UserWarning, stacklevel=2
  560. )
  561. break
  562. activeBlockVectorAR = A(activeBlockVectorR)
  563. if iterationNumber > 0:
  564. if B is not None:
  565. aux = _b_orthonormalize(
  566. B, activeBlockVectorP, activeBlockVectorBP,
  567. verbosityLevel=verbosityLevel
  568. )
  569. activeBlockVectorP, activeBlockVectorBP, invR, normal = aux
  570. else:
  571. aux = _b_orthonormalize(B, activeBlockVectorP,
  572. verbosityLevel=verbosityLevel)
  573. activeBlockVectorP, _, invR, normal = aux
  574. # Function _b_orthonormalize returns None if Cholesky fails
  575. if activeBlockVectorP is not None:
  576. activeBlockVectorAP = activeBlockVectorAP / normal
  577. activeBlockVectorAP = np.dot(activeBlockVectorAP, invR)
  578. restart = forcedRestart
  579. else:
  580. restart = True
  581. ##
  582. # Perform the Rayleigh Ritz Procedure:
  583. # Compute symmetric Gram matrices:
  584. if activeBlockVectorAR.dtype == "float32":
  585. myeps = 1
  586. else:
  587. myeps = np.sqrt(np.finfo(activeBlockVectorR.dtype).eps)
  588. if residualNorms.max() > myeps and not explicitGramFlag:
  589. explicitGramFlag = False
  590. else:
  591. # Once explicitGramFlag, forever explicitGramFlag.
  592. explicitGramFlag = True
  593. # Shared memory assingments to simplify the code
  594. if B is None:
  595. blockVectorBX = blockVectorX
  596. activeBlockVectorBR = activeBlockVectorR
  597. if not restart:
  598. activeBlockVectorBP = activeBlockVectorP
  599. # Common submatrices:
  600. gramXAR = np.dot(blockVectorX.T.conj(), activeBlockVectorAR)
  601. gramRAR = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorAR)
  602. gramDtype = activeBlockVectorAR.dtype
  603. if explicitGramFlag:
  604. gramRAR = (gramRAR + gramRAR.T.conj()) / 2
  605. gramXAX = np.dot(blockVectorX.T.conj(), blockVectorAX)
  606. gramXAX = (gramXAX + gramXAX.T.conj()) / 2
  607. gramXBX = np.dot(blockVectorX.T.conj(), blockVectorBX)
  608. gramRBR = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorBR)
  609. gramXBR = np.dot(blockVectorX.T.conj(), activeBlockVectorBR)
  610. else:
  611. gramXAX = np.diag(_lambda).astype(gramDtype)
  612. gramXBX = np.eye(sizeX, dtype=gramDtype)
  613. gramRBR = np.eye(currentBlockSize, dtype=gramDtype)
  614. gramXBR = np.zeros((sizeX, currentBlockSize), dtype=gramDtype)
  615. if not restart:
  616. gramXAP = np.dot(blockVectorX.T.conj(), activeBlockVectorAP)
  617. gramRAP = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorAP)
  618. gramPAP = np.dot(activeBlockVectorP.T.conj(), activeBlockVectorAP)
  619. gramXBP = np.dot(blockVectorX.T.conj(), activeBlockVectorBP)
  620. gramRBP = np.dot(activeBlockVectorR.T.conj(), activeBlockVectorBP)
  621. if explicitGramFlag:
  622. gramPAP = (gramPAP + gramPAP.T.conj()) / 2
  623. gramPBP = np.dot(activeBlockVectorP.T.conj(),
  624. activeBlockVectorBP)
  625. else:
  626. gramPBP = np.eye(currentBlockSize, dtype=gramDtype)
  627. gramA = bmat(
  628. [
  629. [gramXAX, gramXAR, gramXAP],
  630. [gramXAR.T.conj(), gramRAR, gramRAP],
  631. [gramXAP.T.conj(), gramRAP.T.conj(), gramPAP],
  632. ]
  633. )
  634. gramB = bmat(
  635. [
  636. [gramXBX, gramXBR, gramXBP],
  637. [gramXBR.T.conj(), gramRBR, gramRBP],
  638. [gramXBP.T.conj(), gramRBP.T.conj(), gramPBP],
  639. ]
  640. )
  641. _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel)
  642. try:
  643. _lambda, eigBlockVector = eigh(gramA,
  644. gramB,
  645. check_finite=False)
  646. except LinAlgError as e:
  647. # raise ValueError("eigh failed in lobpcg iterations") from e
  648. if verbosityLevel:
  649. warnings.warn(
  650. f"eigh failed at iteration {iterationNumber} \n"
  651. f"with error {e} causing a restart.\n",
  652. UserWarning, stacklevel=2
  653. )
  654. # try again after dropping the direction vectors P from RR
  655. restart = True
  656. if restart:
  657. gramA = bmat([[gramXAX, gramXAR], [gramXAR.T.conj(), gramRAR]])
  658. gramB = bmat([[gramXBX, gramXBR], [gramXBR.T.conj(), gramRBR]])
  659. _handle_gramA_gramB_verbosity(gramA, gramB, verbosityLevel)
  660. try:
  661. _lambda, eigBlockVector = eigh(gramA,
  662. gramB,
  663. check_finite=False)
  664. except LinAlgError as e:
  665. # raise ValueError("eigh failed in lobpcg iterations") from e
  666. warnings.warn(
  667. f"eigh failed at iteration {iterationNumber} with error\n"
  668. f"{e}\n",
  669. UserWarning, stacklevel=2
  670. )
  671. break
  672. ii = _get_indx(_lambda, sizeX, largest)
  673. _lambda = _lambda[ii]
  674. eigBlockVector = eigBlockVector[:, ii]
  675. if retLambdaHistory:
  676. lambdaHistory[iterationNumber + 1, :] = _lambda
  677. # Compute Ritz vectors.
  678. if B is not None:
  679. if not restart:
  680. eigBlockVectorX = eigBlockVector[:sizeX]
  681. eigBlockVectorR = eigBlockVector[sizeX:
  682. sizeX + currentBlockSize]
  683. eigBlockVectorP = eigBlockVector[sizeX + currentBlockSize:]
  684. pp = np.dot(activeBlockVectorR, eigBlockVectorR)
  685. pp += np.dot(activeBlockVectorP, eigBlockVectorP)
  686. app = np.dot(activeBlockVectorAR, eigBlockVectorR)
  687. app += np.dot(activeBlockVectorAP, eigBlockVectorP)
  688. bpp = np.dot(activeBlockVectorBR, eigBlockVectorR)
  689. bpp += np.dot(activeBlockVectorBP, eigBlockVectorP)
  690. else:
  691. eigBlockVectorX = eigBlockVector[:sizeX]
  692. eigBlockVectorR = eigBlockVector[sizeX:]
  693. pp = np.dot(activeBlockVectorR, eigBlockVectorR)
  694. app = np.dot(activeBlockVectorAR, eigBlockVectorR)
  695. bpp = np.dot(activeBlockVectorBR, eigBlockVectorR)
  696. blockVectorX = np.dot(blockVectorX, eigBlockVectorX) + pp
  697. blockVectorAX = np.dot(blockVectorAX, eigBlockVectorX) + app
  698. blockVectorBX = np.dot(blockVectorBX, eigBlockVectorX) + bpp
  699. blockVectorP, blockVectorAP, blockVectorBP = pp, app, bpp
  700. else:
  701. if not restart:
  702. eigBlockVectorX = eigBlockVector[:sizeX]
  703. eigBlockVectorR = eigBlockVector[sizeX:
  704. sizeX + currentBlockSize]
  705. eigBlockVectorP = eigBlockVector[sizeX + currentBlockSize:]
  706. pp = np.dot(activeBlockVectorR, eigBlockVectorR)
  707. pp += np.dot(activeBlockVectorP, eigBlockVectorP)
  708. app = np.dot(activeBlockVectorAR, eigBlockVectorR)
  709. app += np.dot(activeBlockVectorAP, eigBlockVectorP)
  710. else:
  711. eigBlockVectorX = eigBlockVector[:sizeX]
  712. eigBlockVectorR = eigBlockVector[sizeX:]
  713. pp = np.dot(activeBlockVectorR, eigBlockVectorR)
  714. app = np.dot(activeBlockVectorAR, eigBlockVectorR)
  715. blockVectorX = np.dot(blockVectorX, eigBlockVectorX) + pp
  716. blockVectorAX = np.dot(blockVectorAX, eigBlockVectorX) + app
  717. blockVectorP, blockVectorAP = pp, app
  718. if B is not None:
  719. aux = blockVectorBX * _lambda[np.newaxis, :]
  720. else:
  721. aux = blockVectorX * _lambda[np.newaxis, :]
  722. blockVectorR = blockVectorAX - aux
  723. aux = np.sum(blockVectorR.conj() * blockVectorR, 0)
  724. residualNorms = np.sqrt(np.abs(aux))
  725. # Use old lambda in case of early loop exit.
  726. if retLambdaHistory:
  727. lambdaHistory[iterationNumber + 1, :] = _lambda
  728. if retResidualNormsHistory:
  729. residualNormsHistory[iterationNumber + 1, :] = residualNorms
  730. residualNorm = np.sum(np.abs(residualNorms)) / sizeX
  731. if residualNorm < smallestResidualNorm:
  732. smallestResidualNorm = residualNorm
  733. bestIterationNumber = iterationNumber + 1
  734. bestblockVectorX = blockVectorX
  735. if np.max(np.abs(residualNorms)) > residualTolerance:
  736. warnings.warn(
  737. f"Exited at iteration {iterationNumber} with accuracies \n"
  738. f"{residualNorms}\n"
  739. f"not reaching the requested tolerance {residualTolerance}.\n"
  740. f"Use iteration {bestIterationNumber} instead with accuracy \n"
  741. f"{smallestResidualNorm}.\n",
  742. UserWarning, stacklevel=2
  743. )
  744. if verbosityLevel:
  745. print(f"Final iterative eigenvalue(s):\n{_lambda}")
  746. print(f"Final iterative residual norm(s):\n{residualNorms}")
  747. blockVectorX = bestblockVectorX
  748. # Making eigenvectors "exactly" satisfy the blockVectorY constrains
  749. if blockVectorY is not None:
  750. _applyConstraints(blockVectorX,
  751. gramYBY,
  752. blockVectorBY,
  753. blockVectorY)
  754. # Making eigenvectors "exactly" othonormalized by final "exact" RR
  755. blockVectorAX = A(blockVectorX)
  756. if blockVectorAX.shape != blockVectorX.shape:
  757. raise ValueError(
  758. f"The shape {blockVectorX.shape} "
  759. f"of the postprocessing iterate not preserved\n"
  760. f"and changed to {blockVectorAX.shape} "
  761. f"after multiplying by the primary matrix.\n"
  762. )
  763. gramXAX = np.dot(blockVectorX.T.conj(), blockVectorAX)
  764. blockVectorBX = blockVectorX
  765. if B is not None:
  766. blockVectorBX = B(blockVectorX)
  767. if blockVectorBX.shape != blockVectorX.shape:
  768. raise ValueError(
  769. f"The shape {blockVectorX.shape} "
  770. f"of the postprocessing iterate not preserved\n"
  771. f"and changed to {blockVectorBX.shape} "
  772. f"after multiplying by the secondary matrix.\n"
  773. )
  774. gramXBX = np.dot(blockVectorX.T.conj(), blockVectorBX)
  775. _handle_gramA_gramB_verbosity(gramXAX, gramXBX, verbosityLevel)
  776. gramXAX = (gramXAX + gramXAX.T.conj()) / 2
  777. gramXBX = (gramXBX + gramXBX.T.conj()) / 2
  778. try:
  779. _lambda, eigBlockVector = eigh(gramXAX,
  780. gramXBX,
  781. check_finite=False)
  782. except LinAlgError as e:
  783. raise ValueError("eigh has failed in lobpcg postprocessing") from e
  784. ii = _get_indx(_lambda, sizeX, largest)
  785. _lambda = _lambda[ii]
  786. eigBlockVector = np.asarray(eigBlockVector[:, ii])
  787. blockVectorX = np.dot(blockVectorX, eigBlockVector)
  788. blockVectorAX = np.dot(blockVectorAX, eigBlockVector)
  789. if B is not None:
  790. blockVectorBX = np.dot(blockVectorBX, eigBlockVector)
  791. aux = blockVectorBX * _lambda[np.newaxis, :]
  792. else:
  793. aux = blockVectorX * _lambda[np.newaxis, :]
  794. blockVectorR = blockVectorAX - aux
  795. aux = np.sum(blockVectorR.conj() * blockVectorR, 0)
  796. residualNorms = np.sqrt(np.abs(aux))
  797. if retLambdaHistory:
  798. lambdaHistory[bestIterationNumber + 1, :] = _lambda
  799. if retResidualNormsHistory:
  800. residualNormsHistory[bestIterationNumber + 1, :] = residualNorms
  801. if retLambdaHistory:
  802. lambdaHistory = lambdaHistory[
  803. : bestIterationNumber + 2, :]
  804. if retResidualNormsHistory:
  805. residualNormsHistory = residualNormsHistory[
  806. : bestIterationNumber + 2, :]
  807. if np.max(np.abs(residualNorms)) > residualTolerance:
  808. warnings.warn(
  809. f"Exited postprocessing with accuracies \n"
  810. f"{residualNorms}\n"
  811. f"not reaching the requested tolerance {residualTolerance}.",
  812. UserWarning, stacklevel=2
  813. )
  814. if verbosityLevel:
  815. print(f"Final postprocessing eigenvalue(s):\n{_lambda}")
  816. print(f"Final residual norm(s):\n{residualNorms}")
  817. if retLambdaHistory:
  818. lambdaHistory = np.vsplit(lambdaHistory, np.shape(lambdaHistory)[0])
  819. lambdaHistory = [np.squeeze(i) for i in lambdaHistory]
  820. if retResidualNormsHistory:
  821. residualNormsHistory = np.vsplit(residualNormsHistory,
  822. np.shape(residualNormsHistory)[0])
  823. residualNormsHistory = [np.squeeze(i) for i in residualNormsHistory]
  824. if retLambdaHistory:
  825. if retResidualNormsHistory:
  826. return _lambda, blockVectorX, lambdaHistory, residualNormsHistory
  827. else:
  828. return _lambda, blockVectorX, lambdaHistory
  829. else:
  830. if retResidualNormsHistory:
  831. return _lambda, blockVectorX, residualNormsHistory
  832. else:
  833. return _lambda, blockVectorX