quantax.optimizer.lstsq_shift_cg#

quantax.optimizer.lstsq_shift_cg(ashift: float = 0.0001, rtol: float = 0.01, atol: float = 0.0, maxiter: int = 500, dtype: str | type[Any] | dtype | SupportsDType | None = None) → Callable[[...], Array]#

Obtain the least-square solver for the linear equation \(Ax=b\) using diagonal shift, corresponding to SR. The solution \(x = (A^† A + \epsilon I)^{-1} A^† b\) is obtained by the conjugate gradient method jax.scipy.sparse.linalg.cg, applied to the shifted normal equation \((A^† A + \epsilon I) x = A^† b\) matrix-free, without forming \(A^† A\). This is suitable for overdetermined problems where the number of samples exceeds the number of parameters.

Parameters:

  • ashift – The absolute diagonal shift \(\epsilon\).

  • rtol – The relative tolerance for terminating the CG iteration.

  • atol – The absolute tolerance for terminating the CG iteration.

  • maxiter – The maximum number of CG iterations.

  • dtype – The dtype used internally in the CG iteration. By default (None) the iteration is carried out in double precision (float64 for real inputs, complex128 for complex inputs), regardless of the dtype of A. Running CG on the normal-equation operator \(A^† A\) squares the condition number, so in single precision the iteration easily loses conjugacy and diverges to NaN whenever \(A^† A\) is ill-conditioned. Pass dtype=jnp.float32 to keep single precision for large-scale problems where the memory footprint matters and the operator is well conditioned.

Returns:

A solver function with two arguments A and b and one output x as the solution of \(A x = b\). It also accepts a keyword argument x0 as the initial guess of the CG iteration.