uot.py

import numpy as np
import skglm
import celer
import sklearn.linear_model
import scipy.sparse as sp


def get_X_lasso(n, m):
    """ Construction of design matrix H

    :param n: int
        Dimension of source histogram
    :param m: int
        Dimension of target histogram
    :return:
        Design (or dictionary) matrix H
    """

    jHa = np.arange(m * n)
    iHa = np.repeat(np.arange(n), m)
    jHb = np.arange(m * n)
    iHb = np.tile(np.arange(m), n) + n
    j = np.concatenate((jHa, jHb))
    i = np.concatenate((iHa, iHb))
    H = sp.csc_matrix((np.ones(n * m * 2), (i, j)), shape=(n+m, n*m))
    return H

n = 10
X = get_X_lasso(n, n)
y = np.ones(2 * n) / (2 * n)

tol = 1e-14
nitermax = 1000

alpha_max = np.max(np.abs(X.T @ y)) / len(y)
reg_ = alpha_max / 2

print("skglm")
model = skglm.Lasso(
    alpha=reg_, max_iter=nitermax, fit_intercept=False, tol=tol, verbose=True)
model.fit(X, y)
coef_skglm = model.coef_

print("sklearn")
model = sklearn.linear_model.Lasso(
    alpha=reg_, max_iter=10**5, fit_intercept=False, tol=tol)
model.fit(X, y)
coef_sklearn = model.coef_

print("celer")
model = celer.Lasso(
    reg_, max_iter=nitermax, fit_intercept=False, tol=tol, verbose=True)
model.fit(X, y)
coef_celer = model.coef_

def obj(coef_):
    return ((X @ coef_ - y) ** 2).mean() + reg_ * np.abs(coef_).sum()

print("Obj skglm:%e" % obj(coef_skglm))
print("Obj sklearn:%e" % obj(coef_sklearn))
print("Obj celer:%e" % obj(coef_celer))

np.testing.assert_allclose(coef_skglm, coef_celer)
np.testing.assert_allclose(coef_skglm, coef_sklearn)
np.testing.assert_allclose(coef_celer, coef_sklearn)