mr2.algorithms.optimizers.admm_linear

Linearized ADMM for \(\min_x f(x) + g(Ax)\).

This routine solves convex composite problems of the form

\(\min_x f(x) + g(Ax)\),

where \(f\) and \(g\) are proximable and \(A\) is linear. It supports single-variable as well as block-variable formulations by operating on tuples of tensors and LinearOperatorMatrix.

The algorithm applies the following updates:

\[\begin{split}x_{k+1} = \mathrm{prox}_{\mu f}\left(x_k - \frac{\mu}{\tau} A^H(Ax_k-z_k+u_k)\right)\\ 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}\]

with scaled dual variable \(u\).

For convergence of the linearized method, mu / tau must be chosen sufficiently small relative to the squared operator norm of \(A\). This is not checked automatically.

Parameters:

f (ProximableFunctionalSeparableSum | ProximableFunctional) – Proximable functional \(f\). Can be a single functional or a ProximableFunctionalSeparableSum.
g (ProximableFunctionalSeparableSum | ProximableFunctional) – Proximable functional \(g\). Can be a single functional or a ProximableFunctionalSeparableSum.
operator (LinearOperator | LinearOperatorMatrix | None) – Linear operator \(A\). If None, an identity operator matrix is used (requiring matching number of components in f and g).
initial_values (Sequence[Tensor] | Tensor) – Initial primal variable(s). Single tensor or tuple of tensors.
tau (float | Tensor) – Positive ADMM penalty parameter for the z-update.
mu (float | Tensor) – Positive linearization/proximal parameter for the x-update. Together with tau, it must satisfy the convergence condition of linearized ADMM.
max_iterations (int, default: 128) – Maximum number of 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.
callback (Callable[[ADMMLinearStatus], bool | None] | None, default: None) – Optional callback called after each iteration with ADMMLinearStatus. 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/mu, mismatched operator/functional dimensions, or incompatible initial variables).