jax.experimental.sparse.linalg.lobpcg_standard

jax.experimental.sparse.linalg.lobpcg_standard#

jax.experimental.sparse.linalg.lobpcg_standard(A, X, m=100, tol=None)[source]#

使用 LOBPCG 例程计算前 k 个标准特征值。

LOBPCG [1] 代表局部最优块预处理共轭梯度法。该方法能够以加速器友好的方式找到前 k 个特征向量。

这个初始实验版本有一些注意事项。

  • 只支持标准特征值问题 A U = lambda U，不支持一般特征值问题。

  • 梯度代码不可用。

  • f64 只能在 jnp.linalg.eigh 支持该类型的情况下工作。

  • 尚未支持查找最小的特征向量。因此，我们目前不支持预处理，而这主要用于这种情况。

该实现基于 [2] 和 [3]；但是，我们在几个方面偏离了这些来源以提高鲁棒性或简化实现。

  • 尽管迭代成本增加，我们始终为块搜索方向维护一个正交基。

  • 我们更改了收敛标准；请参阅 tol 参数。

  • 故意没有实现软锁定 [4]；它依赖于选择适当的特定于问题的容差，以防止收敛附近的近零残差的灾难性抵消导致爆炸。相反，实现的方法倾向于截断迭代基础。

[1]: http://ccm.ucdenver.edu/reports/rep149.pdf [2]: https://arxiv.org/abs/1704.07458 [3]: https://arxiv.org/abs/0705.2626 [4]: DOI 10.13140/RG.2.2.11794.48327

参数:

  • A (jax.Array | Callable[[jax.Array], jax.Array]) – 表示方阵的 (n, n) 数组，或其操作的调用函数。

  • X (jax.Array) – 表示用于查找前 k 个特征向量的初始搜索方向的 (n, k) 数组。 该数组不必是正交的，但必须是数值线性无关的 (X 将被正交化)。 注意，必须有 0 < k * 5 < n。

  • m (int) – 最大迭代次数；LOBPCG 只探索 (Krylov 基 {X, A X, A^2 X, …, A^m X} 的) 子空间。

  • tol (jax.Array | float | None | None) – 浮点收敛容差；当特征对 (lambda, v) 的残差 L2 范数 r = |A v - lambda v| 低于 tol * 10 * n * (lambda + |A v|) 时，该特征对就收敛了，这大致估计了理想特征向量最坏情况下的浮点误差。 如果所有 k 个特征向量都满足容差比较，则 LOBPCG 会提前退出。 如果保留为 None，则设置为 A.dtype 的浮点 epsilon。

返回:

theta, U, i，其中 theta 是一个 (k,) 的特征值数组，U 是一个 (n, k) 的特征向量数组，i 是执行的迭代次数。

引发:

ValueError – 如果 A,X 数据类型或 n 维度不匹配，或者 k 太大 (仅支持 k * 5 < n)，或者 k == 0。

内容