mr2.algorithms.optimizers.admm_l2

mr2.algorithms.optimizers.admm_l2(g: ProximableFunctionalSeparableSum | ProximableFunctional, op: LinearOperator | LinearOperatorMatrix, b: Sequence[Tensor] | Tensor, a: LinearOperator | LinearOperatorMatrix, initial_values: Sequence[Tensor] | Tensor, *, tau: float | Tensor, max_iterations: int = 128, tolerance: float = 0.0, initial_z: Sequence[Tensor] | Tensor | None = None, initial_u: Sequence[Tensor] | Tensor | None = None, cg_max_iterations: int = 128, cg_tolerance: float = 1e-6, callback: Callable[[ADMML2Status], bool | None] | None = None) → tuple[Tensor, ...][source]

ADMM for \(\min_x \frac{1}{2}\|Op\,x-b\|_2^2 + g(Ax)\).

This routine targets inverse problems with quadratic data fidelity and proximable regularization in transformed coordinates:

\(\min_x \frac{1}{2}\|Op\,x-b\|_2^2 + g(Ax)\).

Introducing \(z = Ax\), ADMM uses:

\[\begin{split}x_{k+1} = \arg\min_x \frac{1}{2}\|Op\,x-b\|_2^2 + \frac{1}{2\tau}\|Ax-z_k+u_k\|_2^2\\ z_{k+1} = \mathrm{prox}_{\tau g}(Ax_{k+1}+u_k)\\ u_{k+1} = u_k + Ax_{k+1} - z_{k+1}\end{split}\]

The x-update is solved with cg on the normal equations:

\[(Op^H Op + \tau^{-1} A^H A)x = Op^H b + \tau^{-1}A^H(z-u).\]
Parameters:

  • g (ProximableFunctionalSeparableSum | ProximableFunctional) – Proximable regularizer. Can be a single functional or a ProximableFunctionalSeparableSum.

  • op (LinearOperator | LinearOperatorMatrix) – Forward operator \(Op\) in the data term.

  • b (Sequence[Tensor] | Tensor) – Data tensor(s). Single tensor or tuple of tensors, matching rows of op.

  • a (LinearOperator | LinearOperatorMatrix) – Linear operator \(A\) used for splitting in \(g(Ax)\).

  • initial_values (Sequence[Tensor] | Tensor) – Initial primal variable(s). Single tensor or tuple of tensors.

  • tau (float | Tensor) – Positive ADMM penalty parameter.

  • max_iterations (int, default: 128) – Maximum number of ADMM iterations.

  • tolerance (float, default: 0.0) – Relative stopping tolerance on the primal update \(\|x_{k+1}-x_k\|_2 / \|x_{k+1}\|_2\). If zero, no tolerance-based early stopping is applied.

  • initial_z (Sequence[Tensor] | Tensor | None, default: None) – Optional initial split variable(s). If None, initialized as \(z_0 = A x_0\).

  • initial_u (Sequence[Tensor] | Tensor | None, default: None) – Optional initial scaled dual variable(s). If None, initialized to zero with shape matching z.

  • cg_max_iterations (int, default: 128) – Maximum CG iterations used inside each ADMM x-update.

  • cg_tolerance (float, default: 1e-6) – Residual tolerance used by CG for each x-update.

  • callback (Callable[[ADMML2Status], bool | None] | None, default: None) – Optional callback called after each iteration with ADMML2Status. If it returns False, iterations stop early.

Returns:

Tuple of tensors representing the final primal variable(s).

Raises:

ValueError – If parameters are inconsistent (non-positive tau, incompatible operator dimensions, mismatched number of regularizers, or invalid initial values).