benchopt / benchopt/benchmark_lasso

ENH add ADMM solver

Open
#76 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
14
Forks
24
PR merge metrics
No merged PRs in 30d

Description

@josephsalmon here's a basic ADMM solver for the Lasso:
```python
import numpy as np
from numpy.linalg import norm
import matplotlib.pyplot as plt

from scipy import io
from sklearn.preprocessing import MinMaxScaler

from skglm.utils import ST_vec

def primal_lasso(X, y, alpha, w):
"""Compute the primal objective of Lasso.

Parameters
----------
X : array, shape (n_samples, n_features)
Design matrix.

y : array, shape (n_samples,)
Target vector.

alpha : float
Regularization parameter.

w : array, shape (n_features)
Coefficient vector.

Returns
-------
p_obj : float
The primal objective.
"""
r = X @ w - y
return 1/2 * r @ r + alpha * norm(w, ord=1)

def extrapolate(last_K_u, K, extrap_type="LPinf", s=100):
"""Construct an extrapolated point from past iterates.

Parameters
----------
last_K_u : array, shape (n_features, K)
Array of past iterates.

K : int
Number of past iterates to use for extrapolation.

extrap_type : str, (`LP` | `LPinf`), optional
Extrapolation type.

s : int
LP steps (used if extra_type=`LP`).

Returns
-------
w : array, shape (n_features,)
Extrapolated vectors
"""
q = K - 2

V = np.diff(last_K_u, 1, axis=-1)

V_k = V[:, 1:]
v_k = V[:, -1]
V_prev = V[:, :-1]

# Compute coefficient
VtV = V_prev.T @ V_prev

try:
c = np.linalg.solve(VtV, V_prev.T @ v_k)
except np.linalg.LinAlgError:
return v_k
else:
# Iteration matrix
C = np.diag(np.ones(q-1), -1)
C[:, -1] = c

rho = norm(np.linalg.eigvals(C), ord=np.inf)

if extrap_type == "LPinf":
if rho < 1:
tmp = np.eye(q) - C
S = np.linalg.solve(tmp.T, C.T).T
else:
S = 0 * C
else:
if rho < 1:
pC = np.linalg.matrix_power(C, s)
tmp = np.eye(q) - C
S = np.linalg.solve(tmp.T, (C - pC).T, ).T
else:
S = 0 * C

return V_k @ S[:, -1]

def admm(X, y, alpha, gamma=2., max_iter=1000, tol=1e-5, check_gap_freq=50, a=0, K=6,
use_accel=True, verbose=True):
"""Run Alternate Direction Method of multipliers optimization scheme for Lasso.

Parameters
----------
X : array, shape (n_samples, n_features)
Design matrix.

y : array, shape (n_samples,)
Target vector.

alpha : float
Regularization parameter.

gamma : float
Augmented Lagrangian parameter.

max_iter : int
Maximum number of iterations.

tol : float
Tolerance.

check_gap_freq : int
Frequency for checking convergence.

a : float
Inertia parameter.

K : int
Number of past iterates to compute extrapolated point.

use_accel : bool
Use extrapolation.

verbose : bool
Verbosity.

Returns
-------
w : array, shape (n_features,)
Coefficient vector.
"""
n_features = X.shape[1]
residuals = []
iterates = []

# Acceleration variables
last_K_u = np.zeros((n_features, K))

# Optimization variables
w = np.ones(n_features) # Primal iterates
z = np.ones(n_features)
psi = np.ones(n_features) # Dual iterates
u = psi + gamma * w
u_bar = u

v = u - u

# Pre-compute useful quantities
XtX_scaled = X.T @ X / gamma
Xty_scaled = X.T @ y / gamma
L = np.linalg.cholesky(XtX_scaled + np.eye(n_features))
U = L.T

for iter in range(1, max_iter + 1):
u_prev = u.copy()

# Proximal step for datafit
z = np.linalg.solve(U, np.linalg.solve(L, Xty_scaled + u_bar / gamma))
psi = u_bar - gamma * z # Dual update

# Proximal step for pen
w = ST_vec((u_bar - 2 * psi) / gamma, alpha / gamma)
u = psi + gamma * w
iterates.append(w)

# Inertial step
v = u - u_prev
u_bar = u + a * v

last_K_u = np.column_stack((last_K_u[:, 1:], u))

if use_accel and iter % (K + 1) == 0:
e = extrapolate(last_K_u, K)
with np.errstate(divide="ignore"):
# Removes warning for zero division at first iteration
# Parameter safeguard - avoid numerical errors
coeff = np.minimum(1., 1e5 / (iter**1.1 * norm(e)))
u = u + coeff * e
u_bar = u

res = norm(v)
residuals.append(res)

if iter % check_gap_freq == 0:
p_obj = primal_lasso(X, y, alpha, w)
if verbose:
print(f"iter {iter} :: residual {res:.5f} :: obj {p_obj:.4f}")

if res < tol:
break
return w, residuals, iterates

if __name__ == "__main__":
# Matrices can be downloaded at:
# https://github.com/jliang993/A3DMM/tree/master/codes/data
X = io.loadmat('covtype_sample.mat')["h"]
y = io.loadmat('covtype_label.mat')["l"]

y = np.ravel(y)

scaler = MinMaxScaler(feature_range=(-1, 1))
X = X.toarray()
X = scaler.fit_transform(X)

alpha = 1
w, residuals, iterates = admm(X, y, alpha, tol=1e-6, use_accel=True,
max_iter=50_000, check_gap_freq=10)
print("#" * 25)
w_no_accel, residuals_no_acc, iterates_no_acc = admm(X, y, alpha, tol=1e-6,
use_accel=False,
max_iter=50_000,
check_gap_freq=100)

np.testing.assert_allclose(w, w_no_accel, rtol=1e-4)

# Plotting
norms_accel = list(map(lambda wc: np.log(norm(wc - w)), iterates))
norms_no_accel = list(map(lambda wc: np.log(
norm(wc - w_no_accel)), iterates_no_acc))
plt.plot(norms_accel, label="Accelerated")
plt.plot(norms_no_accel, label="No accel")
plt.legend()
plt.title("ADMM - ||x - x^*||")
plt.show()
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.