in admm.py [0:0]
def adm(f, K, D, mu, beta, n_iter):
"""
ADMM iterations for an l_2 group lasso.
Runs n_iter iterations of the alternating direction method of multipliers
(ADMM) with coupling beta to minimize over x the following objective:
mu |Kx - f|_2^2 / 2 + sum_i sqrt((Dx)_{2i}^2 + (Dx)_{2i+1}^2)
where |Kx - f|_2^2 denotes the square of the Euclidean norm of Kx-f.
Parameters
----------
f : ndarray
dependent variables in the least-squares regression
K : ndarray
independent variables (design matrix) in the least-squares regression
D : ndarray
matrix in the regularization term sqrt((Dx)_{2i}^2 + (Dx)_{2i+1}^2)
mu : float
regularization parameter
beta : float
coupling parameter for the ADMM iterations
n_iter : int
number of iterations to conduct
Returns
-------
ndarray
argmin[ mu |Kx - f|_2^2 / 2 + sum_i sqrt((Dx)_{2i}^2 + (Dx)_{2i+1}^2) ]
where |Kx - f|_2^2 denotes the square of the Euclidean norm of Kx-f
float
objective value at the end of the ADMM iterations
"""
assert type(f) == type(K)
assert type(f) == type(D)
n = K.shape[1]
# Compute the upper triangular factor U in the Cholesky decomposition of
# D^* D + mu / beta * K^* K.
U = scipy.linalg.lu_factor(D.conj().T @ D + (mu / beta) * (K.conj().T @ K))
# Apply (mu / beta) * K^* to the input f, obtaining Ktf.
Ktf = (mu / beta) * (K.conj().T @ f)
# Initialize the primal (x) and dual (la) solution vectors to zeros.
x = np.zeros(n)
la = np.zeros(2 * n)
# Conduct iterations of alternating minimization.
for i in range(n_iter):
# Apply shrinkage via formula (2.7) from Tao-Yang, dividing both
# arguments of the "max" operator in formula (2.7) by the denominator
# of the rightmost factor in formula (2.7).
a = (D @ x + la / beta).reshape(2, n)
b = scipy.linalg.norm(a, axis=0, keepdims=True)
if i > 0:
y = (a * np.maximum(1 - 1 / (beta * b), 0)).reshape(2 * n)
else:
y = np.zeros(2 * n)
# Solve formula (2.8) from Tao-Yang using the Cholesky factor U
# (that is, perform a backward solve followed by a forward solve).
c = D.conj().T @ (y - la / beta) + Ktf
x = scipy.linalg.lu_solve(U, c)
# Update the Lagrange multipliers via formula (2.9) from Tao-Yang.
la = la - beta * (y - D @ x)
# Calculate the loss in formula (1.4) from Tao-Yang...
loss = scipy.linalg.norm((D @ x).reshape(2, n), axis=0).sum()
# ... adding in the term for the fidelity of the reconstruction.
loss += (mu / 2) * scipy.linalg.norm(K @ x - f)**2
# Discard the imaginary part of the primal solution,
# returning only the real part and the loss.
return x.real, loss