Repository: fdtomasi/regain

Issue:
Renaming of GraphLasso into GraphicalLasso
As `sklearn` v0.20, we need to import `GraphicalLasso` instead of `GraphLasso` to avoid warnings.

At the same time, it does make sense to rename also our classes.

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 66cc7dd..7b22ee2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -102,3 +102,5 @@ ENV/
 
 *~
 .magicindex.json
+
+\.vscode/
diff --git a/regain/__init__.py b/regain/__init__.py
index 124e462..7fd229a 100644
--- a/regain/__init__.py
+++ b/regain/__init__.py
@@ -1,1 +1,1 @@
-__version__ = '0.1.7'
+__version__ = '0.2.0'
diff --git a/regain/bayesian/gwishart_inference.py b/regain/bayesian/gwishart_inference.py
index 0f8c2d8..86c9b0a 100755
--- a/regain/bayesian/gwishart_inference.py
+++ b/regain/bayesian/gwishart_inference.py
@@ -16,16 +16,19 @@ from sklearn.linear_model import LassoLars
 from sklearn.utils import Bunch, check_array
 from sklearn.utils.extmath import fast_logdet
 
-from regain.covariance.graph_lasso_ import GraphLasso, graph_lasso
+from regain.covariance.graphical_lasso_ import GraphicalLasso, graphical_lasso
 
 
 def mk_all_ugs(n_dim):
     """Utility for generating all possible """
     nedges = int(comb(n_dim, 2))
-    m = 2 ** nedges
+    m = 2**nedges
 
-    ind = np.array([list(binary_repr(x, width=len(binary_repr(m-1))))
-                    for x in range(m)]).astype(int)
+    ind = np.array(
+        [
+            list(binary_repr(x, width=len(binary_repr(m - 1))))
+            for x in range(m)
+        ]).astype(int)
     ord = np.argsort(ind.sum(axis=1))
     ind = ind[ord]
 
@@ -42,8 +45,9 @@ def mk_all_ugs(n_dim):
 
 def markov_blankets(graphs, boolean=True, tol=1e-8, unique=False):
     """For each variable, list each of its markov blankets in graphs."""
-    m_blankets = [np.array([G[i] for G in graphs]) for i in
-                  range(graphs[0].shape[0])]
+    m_blankets = [
+        np.array([G[i] for G in graphs]) for i in range(graphs[0].shape[0])
+    ]
     for i, mb in enumerate(m_blankets):
         mb[:, i] = 0
         mb[np.abs(mb) < tol] = 0
@@ -71,25 +75,30 @@ def score_blankets(blankets, X, alphas=(0.01, 0.5, 1)):
                 X_mb = np.zeros((X.shape[0], 1))
 
             y_mb = X[:, i]
-            score = np.sum([LassoLars(alpha=alpha).fit(
-                X_mb, y_mb).score(X_mb, y_mb) for alpha in alphas])
+            score = np.sum(
+                [
+                    LassoLars(alpha=alpha).fit(X_mb, y_mb).score(X_mb, y_mb)
+                    for alpha in alphas
+                ])
 
             scores.append(score)
         scores_all.append(scores)
 
     ss = np.exp(np.array(scores_all))
-    normalized_scores = ss / np.repeat(ss.sum(axis=1)[:, None],
-                                       ss.shape[1], axis=1)
+    normalized_scores = ss / np.repeat(
+        ss.sum(axis=1)[:, None], ss.shape[1], axis=1)
     return normalized_scores
 
 
 def get_graphs(blankets, scores, n_dim, n_resampling=200):
-    idx = np.array([np.random.choice(
-        scores.shape[1], p=scores[i], size=n_resampling)
-        for i in range(n_dim)])
-
-    graphs_ = np.array([blankets[i][idx[i]] for i in
-                        range(n_dim)]).transpose(1, 0, 2)
+    idx = np.array(
+        [
+            np.random.choice(scores.shape[1], p=scores[i], size=n_resampling)
+            for i in range(n_dim)
+        ])
+
+    graphs_ = np.array([blankets[i][idx[i]] for i in range(n_dim)]).transpose(
+        1, 0, 2)
     # symmetrise with AND operator -> product
     graphs = np.array([g * g.T for g in graphs_])
     return graphs
@@ -102,12 +111,12 @@ def covsel(x, p, nonZero, C):
     Based on sparse GGM estimation code by Mark Schmidt.
     """
     X = np.zeros((p, p))
-    X[nonZero] = x   # fill the diagonal and upper triangle
+    X[nonZero] = x  # fill the diagonal and upper triangle
     X += np.triu(X, 1).T  # fill the lower triangle
 
     # Fast Way to compute -logdet(X) + tr(X*C)
     # f = -2*sum(log(diag(R))) + sum(sum(C.*X)) + (lambda/2)*sum(X(:).^2);
-    f = - fast_logdet(X) + np.sum(C * X)
+    f = -fast_logdet(X) + np.sum(C * X)
 
     if f < np.inf:
         g = C - linalg.pinvh(X)
@@ -159,7 +168,7 @@ def GWishartFit(X, G, GWprior, mode='covsel'):
         # convert G to alpha
         alpha = np.zeros_like(G, dtype=float)
         alpha[~(G + G.T)] = np.inf
-        precision = graph_lasso(emp_cov=C, alpha=alpha)[0]
+        precision = graphical_lasso(emp_cov=C, alpha=alpha)[0]
 
     return precision, S
 
@@ -239,11 +248,11 @@ def compute_score(X, G, P, S, GWprior=None, score_method='bic'):
 
             tmp2 = A[0, :].dot(B[:, 0]) + A[1, :].dot(B[:, 1])
 
-            diagH[e1] = -(dn-2) * tmp2 / 2
+            diagH[e1] = -(dn - 2) * tmp2 / 2
             # diagH(e1) = -(dn-2) * trace(M1(nz2,nz1)*M2(nz1,nz2))/2;
 
         logdetHdiag = sum(np.log(-diagH))
-        lognormconst = dof * np.log(2*np.pi)/2 + logh - logdetHdiag / 2.
+        lognormconst = dof * np.log(2 * np.pi) / 2 + logh - logdetHdiag / 2.
         score = lognormconst - GWprior.lognormconst - \
             n_samples * n_dim * np.log(2 * np.pi) / 2
         GWpost.lognormconst = lognormconst
@@ -266,11 +275,11 @@ def compute_score(X, G, P, S, GWprior=None, score_method='bic'):
                 # tmp2 = A[0, :].dot(B[:, 0]) + A[1, :].dot(B[:, 1])
                 # tmp2 = np.trace(A.dot(B))
                 tmp2 = (A * B.T).sum()
-                H[e2, e1] = H[e1, e2] = -(dn-2) * tmp2 / 2.
+                H[e2, e1] = H[e1, e2] = -(dn - 2) * tmp2 / 2.
 
         # neg Hessian will be posdef
-        logdetH = 2*sum(np.log(np.diag(linalg.cholesky(-H))))
-        lognormconst = dof * np.log(2*np.pi)/2 + logh - logdetH / 2.
+        logdetH = 2 * sum(np.log(np.diag(linalg.cholesky(-H))))
+        lognormconst = dof * np.log(2 * np.pi) / 2 + logh - logdetH / 2.
         score = lognormconst - GWprior.lognormconst - \
             n_samples * n_dim * np.log(2 * np.pi) / 2
         GWpost.lognormconst = lognormconst
@@ -283,7 +292,7 @@ def GWishartScore(X, G, d0=3, S0=None, score_method='bic', mode='covsel'):
     # %Initialize GW prior structure
     n_samples, n_dim = X.shape
     if S0 is None:
-        S0 = (d0 - 2) * np.diag(np.diag(np.cov(X,bias=True)))
+        S0 = (d0 - 2) * np.diag(np.diag(np.cov(X, bias=True)))
 
     GWprior = Bunch(d0=d0, S0=S0, lognormconst=0, lognormconstDiag=0)
 
@@ -292,8 +301,9 @@ def GWishartScore(X, G, d0=3, S0=None, score_method='bic', mode='covsel'):
         noData = np.zeros((0, n_dim))
 
         P0, S_noData = GWishartFit(noData, G, GWprior)
-        GWtemp = compute_score(noData, G, P0, S_noData, GWprior=GWprior,
-                               score_method=score_method)
+        GWtemp = compute_score(
+            noData, G, P0, S_noData, GWprior=GWprior,
+            score_method=score_method)
         GWprior.lognormconst = GWtemp.lognormconst
 
     # Compute the map precision matrix P
@@ -304,7 +314,7 @@ def GWishartScore(X, G, d0=3, S0=None, score_method='bic', mode='covsel'):
     return GWpost
 
 
-def bayesian_graph_lasso(
+def bayesian_graphical_lasso(
         X, tol=1e-8, alphas=None, n_resampling=200, mode='gl', top_n=1,
         scoring='diaglaplace', assume_centered=False, max_iter=100):
     """Bayesian graphical lasso."""
@@ -314,17 +324,16 @@ def bayesian_graph_lasso(
         alphas = np.logspace(-2, 0, 20)
 
     # get a series of Markov blankets for vaiours alphas
-    mdl = GraphLasso(assume_centered=assume_centered, tol=tol,
-                     max_iter=max_iter, verbose=False)
-    precisions = [
-        mdl.set_params(alpha=a).fit(X).precision_
-        for a in alphas]
+    mdl = GraphicalLasso(
+        assume_centered=assume_centered, tol=tol, max_iter=max_iter,
+        verbose=False)
+    precisions = [mdl.set_params(alpha=a).fit(X).precision_ for a in alphas]
     mblankets = markov_blankets(precisions, tol=tol, unique=True)
 
     normalized_scores = score_blankets(mblankets, X=X, alphas=[0.01, 0.5, 1])
 
-    graphs = get_graphs(mblankets, normalized_scores, n_dim=n_dim,
-                        n_resampling=n_resampling)
+    graphs = get_graphs(
+        mblankets, normalized_scores, n_dim=n_dim, n_resampling=n_resampling)
 
     nonzeros_all = [np.triu(g, 1) + np.eye(n_dim, dtype=bool) for g in graphs]
 
@@ -335,14 +344,16 @@ def bayesian_graph_lasso(
     # Find non-zero elements of upper triangle of G
     # make sure diagonal is non-zero
     # G = nonzeros_all[1] # probably can discard if all zeros?
-    res = [GWishartScore(X, G, d0=d0, S0=S0, mode=mode, score_method=scoring)
-           for G in nonzeros_all]
+    res = [
+        GWishartScore(X, G, d0=d0, S0=S0, mode=mode, score_method=scoring)
+        for G in nonzeros_all
+    ]
 
     top_n = [x.P for x in sorted(res, key=lambda x: x.score)[::-1][:top_n]]
     return np.mean(top_n, axis=0)
 
 
-class BayesianGraphLasso(GraphLasso):
+class BayesianGraphicalLasso(GraphicalLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -400,10 +411,11 @@ class BayesianGraphLasso(GraphLasso):
 
     """
 
-    def __init__(self, alpha=0.01, max_iter=100, alphas=None, n_resampling=200,
-                 tol=1e-4, verbose=False, assume_centered=False, mode='gl',
-                 scoring='diaglaplace', top_n=1):
-        super(GraphLasso, self).__init__(
+    def __init__(
+            self, alpha=0.01, max_iter=100, alphas=None, n_resampling=200,
+            tol=1e-4, verbose=False, assume_centered=False, mode='gl',
+            scoring='diaglaplace', top_n=1):
+        super(GraphicalLasso, self).__init__(
             alpha=alpha, tol=tol, max_iter=max_iter, verbose=verbose,
             assume_centered=assume_centered, mode=mode)
         self.alphas = alphas
@@ -422,13 +434,13 @@ class BayesianGraphLasso(GraphLasso):
 
         """
         # Covariance does not make sense for a single feature
-        X = check_array(X, ensure_min_features=2, ensure_min_samples=2,
-                        estimator=self)
+        X = check_array(
+            X, ensure_min_features=2, ensure_min_samples=2, estimator=self)
 
         if self.alphas is None:
             self.alphas = np.logspace(-2, 0, 20)
 
-        self.precision_ = bayesian_graph_lasso(
+        self.precision_ = bayesian_graphical_lasso(
             X, tol=self.tol, alphas=self.alphas,
             n_resampling=self.n_resampling, mode=self.mode,
             scoring=self.scoring, assume_centered=self.assume_centered,
diff --git a/regain/bayesian/wishart_process_.py b/regain/bayesian/wishart_process_.py
index 285af40..fae372d 100644
--- a/regain/bayesian/wishart_process_.py
+++ b/regain/bayesian/wishart_process_.py
@@ -11,8 +11,9 @@ from sklearn.utils.validation import check_X_y
 
 from regain.bayesian import stats
 from regain.bayesian.gaussian_process_ import sample as sample_gp
-from regain.bayesian.sampling import elliptical_slice, sample_hyper_kernel, sample_ell, GWP_construct
-from regain.covariance.time_graph_lasso_ import TimeGraphLasso
+from regain.bayesian.sampling import (GWP_construct, elliptical_slice,
+                                      sample_ell, sample_hyper_kernel)
+from regain.covariance.time_graphical_lasso_ import TimeGraphicalLasso
 
 
 def fit(
@@ -151,12 +152,11 @@ def kernel(X, Y=None, var=None, inverse_width=None, normalised=False):
     return k
 
 
-def periodic_kernel(X, Y=None, inverse_width=None):
-    k = kernels.ExpSineSquared(X, Y=Y, length_scale=inverse_width)
-    return k
+def periodic_kernel(X, Y=None, inverse_width=1):
+    return kernels.ExpSineSquared(length_scale=inverse_width)(X, Y=Y)
 
 
-class WishartProcess(TimeGraphLasso):
+class WishartProcess(TimeGraphicalLasso):
     def __init__(
             self, theta=100, var_prop=1, mu_prior=1, var_prior=10,
             var_Lprop=10, mu_Lprior=1, var_Lprior=1, n_iter=500, burn_in=None,
diff --git a/regain/covariance/__init__.py b/regain/covariance/__init__.py
index 4b5f685..8b0f557 100644
--- a/regain/covariance/__init__.py
+++ b/regain/covariance/__init__.py
@@ -1,4 +1,4 @@
-from .graph_lasso_ import GraphLasso
-from .latent_graph_lasso_ import LatentGraphLasso
-from .time_graph_lasso_ import TimeGraphLasso
-from .latent_time_graph_lasso_ import LatentTimeGraphLasso
+from .graphical_lasso_ import GraphicalLasso
+from .latent_graphical_lasso_ import LatentGraphicalLasso
+from .time_graphical_lasso_ import TimeGraphicalLasso
+from .latent_time_graphical_lasso_ import LatentTimeGraphicalLasso
diff --git a/regain/covariance/__old/latent_time_graph_lasso_.py b/regain/covariance/__old/latent_time_graph_lasso_.py
deleted file mode 100644
index 9322356..0000000
--- a/regain/covariance/__old/latent_time_graph_lasso_.py
+++ /dev/null
@@ -1,463 +0,0 @@
-"""Graphical latent variable models selection over time via ADMM."""
-from __future__ import division
-
-import numpy as np
-import warnings
-
-from functools import partial
-from six.moves import range
-from sklearn.covariance import empirical_covariance
-from sklearn.covariance import EmpiricalCovariance
-from sklearn.utils.extmath import squared_norm, fast_logdet
-from sklearn.utils.validation import check_array
-
-from regain.admm.time_graph_lasso_ import log_likelihood, log_likelihood_trace
-from regain.norm import l1_od_norm
-from regain.prox import soft_thresholding_sign
-from regain.prox import prox_logdet
-from regain.prox import prox_trace_indicator
-from regain.utils import convergence, error_norm_time
-from regain.validation import check_norm_prox
-
-
-def objective(S, R, Z_0, Z_1, Z_2, W_0, W_1, W_2,
-              alpha, tau, beta, eta, psi, phi):
-    """Objective function for time-varying graphical lasso."""
-    obj = np.sum(- log_likelihood(s, r) for s, r in zip(S, R))
-    obj += alpha * np.sum(map(l1_od_norm, Z_0))
-    obj += tau * np.sum(map(partial(np.linalg.norm, ord='nuc'), W_0))
-    obj += beta * np.sum(map(psi, Z_2 - Z_1))
-    obj += eta * np.sum(map(phi, W_2 - W_1))
-    return obj
-
-# def _objective(S, K, L, alpha, tau, beta, eta, psi, phi):
-#     obj = np.sum(- log_likelihood(s, r) for s, r in zip(S, K-L))
-#     obj += alpha * np.sum(map(l1_od_norm, K))
-#     obj += tau * np.sum(map(partial(np.linalg.norm, ord='nuc'), L))
-#     obj += beta * np.sum(map(psi, K[2:] - K[:-1]))
-#     obj += eta * np.sum(map(phi, L[2:] - L[:-1]))
-#     return obj
-
-def latent_time_graph_lasso(
-        emp_cov, alpha=1, tau=1, rho=1, beta=1., eta=1., max_iter=1000,
-        verbose=False, psi='laplacian', phi='laplacian', mode=None,
-        tol=1e-4, rtol=1e-2, assume_centered=False,
-        return_history=False, return_n_iter=True):
-    r"""Time-varying latent variable graphical lasso solver.
-
-    Solves the following problem via ADMM:
-        min sum_{i=1}^T -n_i log_likelihood(K_i-L_i) + alpha ||K_i||_{od,1}
-            + tau ||L_i||_*
-            + beta sum_{i=2}^T Psi(K_i - K_{i-1})
-            + eta sum_{i=2}^T Phi(L_i - L_{i-1})
-
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
-
-    Parameters
-    ----------
-    data_list : list of 2-dimensional matrices.
-        Input matrices.
-    alpha, tau : float, optional
-        Regularisation parameters.
-    rho : float, optional
-        Augmented Lagrangian parameter.
-    max_iter : int, optional
-        Maximum number of iterations.
-    tol : float, optional
-        Absolute tolerance for convergence.
-    rtol : float, optional
-        Relative tolerance for convergence.
-    return_history : bool, optional
-        Return the history of computed values.
-
-    Returns
-    -------
-    K, L : numpy.array, 3-dimensional (T x d x d)
-        Solution to the problem for each time t=1...T .
-    history : list
-        If return_history, then also a structure that contains the
-        objective value, the primal and dual residual norms, and tolerances
-        for the primal and dual residual norms at each iteration.
-    """
-    psi, prox_psi = check_norm_prox(psi)
-    phi, prox_phi = check_norm_prox(phi)
-
-    # S = np.array(map(empirical_covariance, data_list))
-    # n_samples = np.array([s for s in [1.]])
-
-    K = np.zeros_like(emp_cov)
-    Z_0 = np.zeros_like(K)
-    Z_1 = np.zeros_like(K)[:-1]
-    Z_2 = np.zeros_like(K)[1:]
-    W_0 = np.zeros_like(K)
-    W_1 = np.zeros_like(K)[:-1]
-    W_2 = np.zeros_like(K)[1:]
-    X_0 = np.zeros_like(K)
-    X_1 = np.zeros_like(K)[:-1]
-    X_2 = np.zeros_like(K)[1:]
-
-    Z_consensus = np.zeros_like(K)
-    # Z_consensus_old = np.zeros_like(K)
-    W_consensus = np.zeros_like(K)
-    # W_consensus_old = np.zeros_like(K)
-    R_old = np.zeros_like(K)
-
-    # divisor for consensus variables, accounting for two less matrices
-    divisor = np.full(K.shape[0], 3, dtype=float)
-    divisor[0] -= 1
-    divisor[-1] -= 1
-
-    checks = []
-    for iteration_ in range(max_iter):
-        # update R
-        A = Z_0 - W_0 - X_0
-        A[:-1] += Z_1 - W_1 - X_1
-        A[1:] += Z_2 - W_2 - X_2
-        A /= divisor[:, None, None]
-
-        # A += np.array(map(np.transpose, A))
-        # A /= 2.
-
-        # A *= - rho / n_samples[:, None, None]
-        A *= - rho
-        A += emp_cov
-
-        R = np.array([prox_logdet(a, lamda=1. / rho) for a in A])
-
-        # update Z_0
-        # Zold = Z
-        # X_hat = alpha * X + (1 - alpha) * Zold
-        soft_thresholding = partial(soft_thresholding_sign, lamda=alpha / rho)
-        Z_0 = np.array(map(soft_thresholding, R + W_0 + X_0))
-
-        # update Z_1, Z_2
-        # prox_l = partial(prox_laplacian, beta=2. * beta / rho)
-        # prox_e = np.array(map(prox_l, K[1:] - K[:-1] + U_2 - U_1))
-        if beta != 0:
-            A_1 = R[:-1] + W_1 + X_1
-            # A_1 = Z_0[:-1].copy()
-            A_2 = R[1:] + W_2 + X_2
-            # A_2 = Z_0[1:].copy()
-            prox_e = prox_psi(A_2 - A_1, lamda=2. * beta / rho)
-            Z_1 = .5 * (A_1 + A_2 - prox_e)
-            Z_2 = .5 * (A_1 + A_2 + prox_e)
-        else:
-            Z_1 = Z_0[:-1].copy()
-            Z_2 = Z_0[1:].copy()
-
-        # update W_0
-        A = Z_0 - R - X_0
-        W_0 = np.array(map(partial(prox_trace_indicator, lamda=tau / rho), A))
-
-        # update W_1, W_2
-        if eta != 0:
-            A_1 = Z_1 - R[:-1] - X_1
-            # A_1 = W_0[:-1].copy()
-            A_2 = Z_2 - R[1:] - X_2
-            # A_2 = W_0[1:].copy()
-            prox_e = prox_phi(A_2 - A_1, lamda=2. * eta / rho)
-            W_1 = .5 * (A_1 + A_2 - prox_e)
-            W_2 = .5 * (A_1 + A_2 + prox_e)
-        else:
-            W_1 = W_0[:-1].copy()
-            W_2 = W_0[1:].copy()
-
-        # update residuals
-        X_0 += R - Z_0 + W_0
-        X_1 += R[:-1] - Z_1 + W_1
-        X_2 += R[1:] - Z_2 + W_2
-
-        # diagnostics, reporting, termination checks
-        X_consensus = X_0.copy()
-        X_consensus[:-1] += X_1
-        X_consensus[1:] += X_2
-        X_consensus /= divisor[:, None, None]
-
-        Z_consensus = Z_0.copy()
-        Z_consensus[:-1] += Z_1
-        Z_consensus[1:] += Z_2
-        Z_consensus /= divisor[:, None, None]
-
-        W_consensus = W_0.copy()
-        W_consensus[:-1] += W_1
-        W_consensus[1:] += W_2
-        W_consensus /= divisor[:, None, None]
-
-        check = convergence(
-            obj=objective(emp_cov, R, Z_0, Z_1, Z_2, W_0, W_1, W_2,
-                          alpha, tau, beta, eta, psi, phi),
-            rnorm=np.linalg.norm(R - Z_consensus + W_consensus),
-            snorm=np.linalg.norm(rho * (R - R_old)),
-            e_pri=np.sqrt(np.prod(K.shape)) * tol + rtol * max(
-                np.linalg.norm(R),
-                np.sqrt(squared_norm(Z_consensus) - squared_norm(W_consensus))),
-            e_dual=np.sqrt(np.prod(K.shape)) * tol + rtol * np.linalg.norm(
-                rho * X_consensus)
-        )
-        R_old = R.copy()
-
-        if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check)
-
-        checks.append(check)
-        if check.rnorm <= check.e_pri and check.snorm <= check.e_dual:
-            break
-
-        # if iteration_ % 10 == 0:
-        #     rho = rho * 0.8
-    else:
-        warnings.warn("Objective did not converge.")
-
-    # return_list = [Z_consensus, W_consensus, emp_cov]
-    return_list = [Z_consensus, W_0, W_1, W_2, emp_cov]
-    if return_history:
-        return_list.append(checks)
-    if return_n_iter:
-        return_list.append(iteration_)
-    return return_list
-
-
-class LatentTimeGraphLasso(EmpiricalCovariance):
-    """Sparse inverse covariance estimation with an l1-penalized estimator.
-
-    Read more in the :ref:`User Guide <sparse_inverse_covariance>`.
-
-    Parameters
-    ----------
-    alpha : positive float, default 0.01
-        The regularization parameter: the higher alpha, the more
-        regularization, the sparser the inverse covariance.
-
-    mode : {'cd', 'lars'}, default 'cd'
-        The Lasso solver to use: coordinate descent or LARS. Use LARS for
-        very sparse underlying graphs, where p > n. Elsewhere prefer cd
-        which is more numerically stable.
-
-    tol : positive float, default 1e-4
-        The tolerance to declare convergence: if the dual gap goes below
-        this value, iterations are stopped.
-
-    enet_tol : positive float, optional
-        The tolerance for the elastic net solver used to calculate the descent
-        direction. This parameter controls the accuracy of the search direction
-        for a given column update, not of the overall parameter estimate. Only
-        used for mode='cd'.
-
-    max_iter : integer, default 100
-        The maximum number of iterations.
-
-    verbose : boolean, default False
-        If verbose is True, the objective function and dual gap are
-        plotted at each iteration.
-
-    assume_centered : boolean, default False
-        If True, data are not centered before computation.
-        Useful when working with data whose mean is almost, but not exactly
-        zero.
-        If False, data are centered before computation.
-
-    Attributes
-    ----------
-    covariance_ : array-like, shape (n_features, n_features)
-        Estimated covariance matrix
-
-    precision_ : array-like, shape (n_features, n_features)
-        Estimated pseudo inverse matrix.
-
-    n_iter_ : int
-        Number of iterations run.
-
-    See Also
-    --------
-    graph_lasso, GraphLassoCV
-
-    """
-
-    def __init__(self, alpha=1., tau=1., beta=1., eta=1., mode='cd', rho=1.,
-                 bypass_transpose=True, tol=1e-4, rtol=1e-4,
-                 psi='laplacian', phi='laplacian', max_iter=100,
-                 verbose=False, assume_centered=False):
-        super(LatentTimeGraphLasso, self).__init__(assume_centered=assume_centered)
-        self.alpha = alpha
-        self.tau = tau
-        self.beta = beta
-        self.eta = eta
-        self.rho = rho
-        self.mode = mode
-        self.tol = tol
-        self.rtol = rtol
-        self.psi = psi
-        self.phi = phi
-        self.max_iter = max_iter
-        self.verbose = verbose
-        # for splitting purposes, data may come transposed, with time in the
-        # last index. Set bypass_transpose=True if X comes with time in the
-        # first dimension already
-        self.bypass_transpose = bypass_transpose
-
-    def fit(self, X, y=None):
-        """Fit the GraphLasso model to X.
-
-        Parameters
-        ----------
-        X : ndarray, shape (n_time, n_samples, n_features), or
-                (n_samples, n_features, n_time)
-            Data from which to compute the covariance estimate.
-            If shape is (n_samples, n_features, n_time), then set
-            `bypass_transpose = False`.
-        y : (ignored)
-        """
-        if not self.bypass_transpose:
-            X = X.transpose(2, 0, 1)  # put time as first dimension
-        # Covariance does not make sense for a single feature
-        # X = check_array(X, allow_nd=True, estimator=self)
-        # if X.ndim != 3:
-        #     raise ValueError("Found array with dim %d. %s expected <= 2."
-        #                      % (X.ndim, self.__class__.__name__))
-        X = np.array([check_array(x, ensure_min_features=2,
-                      ensure_min_samples=2, estimator=self) for x in X])
-        if self.assume_centered:
-            self.location_ = np.zeros((X.shape[0], 1, X.shape[2]))
-        else:
-            self.location_ = X.mean(1).reshape(X.shape[0], 1, X.shape[2])
-        self.emp_cov = np.array([empirical_covariance(
-            x, assume_centered=self.assume_centered) for x in X])
-        self.precision_, self.latent_, self.covariance_, self.n_iter_ = \
-            latent_time_graph_lasso(
-                self.emp_cov, alpha=self.alpha, tau=self.tau, rho=self.rho,
-                beta=self.beta, eta=self.eta, mode=self.mode,
-                tol=self.tol, rtol=self.rtol, psi=self.psi, phi=self.phi,
-                max_iter=self.max_iter, verbose=self.verbose,
-                return_n_iter=True, return_history=False)
-        return self
-
-    def score(self, X_test, y=None):
-        """Computes the log-likelihood of a Gaussian data set with
-        `self.covariance_` as an estimator of its covariance matrix.
-
-        Parameters
-        ----------
-        X_test : array-like, shape = [n_samples, n_features]
-            Test data of which we compute the likelihood, where n_samples is
-            the number of samples and n_features is the number of features.
-            X_test is assumed to be drawn from the same distribution than
-            the data used in fit (including centering).
-
-        y : not used, present for API consistence purpose.
-
-        Returns
-        -------
-        res : float
-            The likelihood of the data set with `self.covariance_` as an
-            estimator of its covariance matrix.
-
-        """
-        if not self.bypass_transpose:
-            X_test = X_test.transpose(2, 0, 1)  # put time as first dimension
-        # compute empirical covariance of the test set
-        test_cov = np.array([empirical_covariance(
-            x, assume_centered=True) for x in X_test - self.location_])
-
-        # ALLA SKLEARN
-        res = np.sum([log_likelihood(S, K-L) for S, K,L in zip(
-            test_cov, self.precision_, self.latent_)])
-        # -----------------------------------------------------------
-
-        # ALLA TIBSHIRANI
-        # score_f = lambda x, y : - fast_logdet(np.linalg.inv(x))\
-        #         - np.trace(y.dot(x))
-        # #estimated_cov = np.array(map(np.linalg.inv, self.precision_ - self.latent_))
-        #
-        #
-        # res = 0
-        # for train, test in zip(self.precision_ - self.latent_ ,test_cov):
-        #     res += score_f(train, test)
-        # res /= self.precision_.shape[0]
-        # -------------------------------------------------------------------
-
-        # ALLA  MATLAB1
-        # ranks = [np.linalg.matrix_rank(L) for L in self.latent_]
-        # # print(ranks)
-        # scores_ranks = np.square(ranks-np.sqrt(L.shape[0]))
-        # #scores_ranks[np.array(ranks)==0] = 1e10
-        # res = np.mean([log_likelihood_trace(S, K-L) for S, K,L in zip(
-        #      test_cov, self.precision_, self.latent_)])
-        # -----------------------------------------------------------------
-
-        # ALLA MATLAB2
-        # try:
-        #     chols = [2*np.sum(np.log(np.diag(np.linalg.cholesky(K))))
-        #              for K in self.precision_ - self.latent_]
-        #     res = np.sum([c - t.ravel(order='F').T.dot(K.ravel(order="F"))
-        #                   for c, t, K in
-        #               zip(chols, test_cov, self.precision_ - self.latent_)])
-        # except:
-        #     res = -np.inf
-        # print(self.alpha, self.tau, res )#-  np.sum(scores_ranks))
-        return res  # -  np.sum(scores_ranks)
-
-    def error_norm(self, comp_cov, norm='frobenius', scaling=True,
-                   squared=True):
-        """Compute the Mean Squared Error between two covariance estimators.
-        (In the sense of the Frobenius norm).
-
-        Parameters
-        ----------
-        comp_cov : array-like, shape = [n_features, n_features]
-            The covariance to compare with.
-
-        norm : str
-            The type of norm used to compute the error. Available error types:
-            - 'frobenius' (default): sqrt(tr(A^t.A))
-            - 'spectral': sqrt(max(eigenvalues(A^t.A))
-            where A is the error ``(comp_cov - self.covariance_)``.
-
-        scaling : bool
-            If True (default), the squared error norm is divided by n_features.
-            If False, the squared error norm is not rescaled.
-
-        squared : bool
-            Whether to compute the squared error norm or the error norm.
-            If True (default), the squared error norm is returned.
-            If False, the error norm is returned.
-
-        Returns
-        -------
-        The Mean Squared Error (in the sense of the Frobenius norm) between
-        `self` and `comp_cov` covariance estimators.
-
-        """
-        return error_norm_time(self.covariance_, comp_cov, norm=norm,
-                               scaling=scaling, squared=squared)
-
-    def mahalanobis(self, observations):
-        """Computes the squared Mahalanobis distances of given observations.
-
-        Parameters
-        ----------
-        observations : array-like, shape = [n_observations, n_features]
-            The observations, the Mahalanobis distances of the which we
-            compute. Observations are assumed to be drawn from the same
-            distribution than the data used in fit.
-
-        Returns
-        -------
-        mahalanobis_distance : array, shape = [n_observations,]
-            Squared Mahalanobis distances of the observations.
-
-        """
-        if not self.bypass_transpose:
-            # put time as first dimension
-            observations = observations.transpose(2, 0, 1)
-        precision = self.get_precision()
-        # compute mahalanobis distances
-        sum_ = 0.
-        for obs, loc in zip(observations, self.location_):
-            centered_obs = observations - self.location_
-            sum_ += np.sum(
-                np.dot(centered_obs, precision) * centered_obs, 1)
-
-        mahalanobis_dist = sum_ / len(observations)
-        return mahalanobis_dist
diff --git a/regain/covariance/__old/latent_time_graph_lasso_v2_.py b/regain/covariance/__old/latent_time_graph_lasso_v2_.py
deleted file mode 100644
index da9b5dc..0000000
--- a/regain/covariance/__old/latent_time_graph_lasso_v2_.py
+++ /dev/null
@@ -1,457 +0,0 @@
-"""Graphical latent variable models selection over time via ADMM."""
-from __future__ import division
-
-import numpy as np
-import warnings
-
-from functools import partial
-from six.moves import range
-from sklearn.covariance import empirical_covariance
-from sklearn.covariance import EmpiricalCovariance
-from sklearn.utils.extmath import squared_norm
-
-from regain.admm.latent_time_graph_lasso_ import objective
-from regain.norm import l1_od_norm, l1_norm
-from regain.prox import soft_thresholding_od, soft_thresholding_sign
-from regain.prox import blockwise_soft_thresholding, prox_linf
-from regain.prox import prox_logdet, prox_laplacian
-from regain.prox import prox_trace_indicator
-from regain.utils import convergence
-from regain.validation import check_norm_prox
-
-
-def time_latent_graph_lasso(
-        emp_cov, alpha=1., tau=1., rho=1., beta=1., eta=1., max_iter=1000,
-        verbose=False, psi='laplacian', phi='laplacian', assume_centered=False,
-        tol=1e-4, rtol=1e-2, return_history=False, return_n_iter=True,
-        mode=None):
-    r"""Time-varying latent variable graphical lasso solver.
-
-    Solves the following problem via ADMM:
-        min sum_{i=1}^T -n_i log_likelihood(K_i-L_i) + alpha ||K_i||_{od,1}
-            + tau ||L_i||_*
-            + beta sum_{i=2}^T Psi(K_i - K_{i-1})
-            + eta sum_{i=2}^T Phi(L_i - L_{i-1})
-
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
-
-    Parameters
-    ----------
-    data_list : list of 2-dimensional matrices.
-        Input matrices.
-    alpha, tau : float, optional
-        Regularisation parameters.
-    rho : float, optional
-        Augmented Lagrangian parameter.
-    max_iter : int, optional
-        Maximum number of iterations.
-    tol : float, optional
-        Absolute tolerance for convergence.
-    rtol : float, optional
-        Relative tolerance for convergence.
-    return_history : bool, optional
-        Return the history of computed values.
-
-    Returns
-    -------
-    K, L : numpy.array, 3-dimensional (T x d x d)
-        Solution to the problem for each time t=1...T .
-    history : list
-        If return_history, then also a structure that contains the
-        objective value, the primal and dual residual norms, and tolerances
-        for the primal and dual residual norms at each iteration.
-
-    """
-    psi, prox_psi = check_norm_prox(psi)
-    phi, prox_phi = check_norm_prox(phi)
-
-    # emp_cov = np.array([empirical_covariance(
-    #     x, assume_centered=assume_centered) for x in data_list])
-
-    n_samples = np.array([s for s in [1.]])
-
-    K = np.zeros_like(emp_cov)
-    L = np.zeros_like(emp_cov)
-    X = np.zeros_like(emp_cov)
-    Z_0 = np.zeros_like(emp_cov)
-    Z_1 = np.zeros_like(emp_cov)[:-1]
-    Z_2 = np.zeros_like(emp_cov)[1:]
-    W_0 = np.zeros_like(emp_cov)
-    W_1 = np.zeros_like(emp_cov)[:-1]
-    W_2 = np.zeros_like(emp_cov)[1:]
-    U_0 = np.zeros_like(emp_cov)
-    U_1 = np.zeros_like(emp_cov)[:-1]
-    U_2 = np.zeros_like(emp_cov)[1:]
-    Y_0 = np.zeros_like(emp_cov)
-    Y_1 = np.zeros_like(emp_cov)[:-1]
-    Y_2 = np.zeros_like(emp_cov)[1:]
-
-    U_consensus = np.zeros_like(emp_cov)
-    Y_consensus = np.zeros_like(emp_cov)
-    Z_consensus = np.zeros_like(emp_cov)
-    Z_consensus_old = np.zeros_like(emp_cov)
-    W_consensus = np.zeros_like(emp_cov)
-    W_consensus_old = np.zeros_like(emp_cov)
-    R_old = np.zeros_like(emp_cov)
-
-    # divisor for consensus variables, accounting for two less matrices
-    divisor = np.zeros(emp_cov.shape[0]) + 3
-    divisor[0] -= 1
-    divisor[-1] -= 1
-    # eta = np.divide(n_samples, divisor * rho)
-
-    checks = []
-    for iteration_ in range(max_iter):
-        # update R
-        A = K - L - X
-        # A += np.array(map(np.transpose, A))
-        # A /= 2.
-        A *= - rho / n_samples[:, None, None]
-        A += emp_cov
-        R = np.array(map(prox_logdet, A, n_samples / rho))
-
-        # update K, L
-        K = L + R + X + Z_0 - U_0
-        K[:-1] += Z_1 - U_1
-        K[1:] += Z_2 - U_2
-        K /= divisor[:, None, None] + 1
-
-        L = K - R - X + W_0 - Y_0
-        L[:-1] += W_1 - Y_1
-        L[1:] += W_2 - Y_2
-        L /= divisor[:, None, None] + 1
-
-        # update Z_0
-        # Zold = Z
-        # X_hat = alpha * X + (1 - alpha) * Zold
-        soft_thresholding = partial(soft_thresholding_sign, lamda=alpha / rho)
-        Z_0 = np.array(map(soft_thresholding, K + U_0))
-
-        # update Z_1, Z_2
-        A_1 = K[:-1] + U_1
-        A_2 = K[1:] + U_2
-        prox_e = prox_psi(A_2 - A_1, lamda=2. * beta / rho)
-        Z_1 = .5 * (A_1 + A_2 - prox_e)
-        Z_2 = .5 * (A_1 + A_2 + prox_e)
-
-        # update W_0
-        A = L + Y_0
-        W_0 = np.array(map(partial(prox_trace_indicator, lamda=tau / rho), A))
-
-        # update W_1, W_2
-        A_1 = L[:-1] + Y_1
-        A_2 = L[1:] + Y_2
-        prox_e = prox_phi(A_2 - A_1, lamda=2. * eta / rho)
-        W_1 = .5 * (A_1 + A_2 - prox_e)
-        W_2 = .5 * (A_1 + A_2 + prox_e)
-
-        # update residuals
-        X += R - K + L
-
-        U_0 += (K - Z_0)
-        U_1 += (K[:-1] - Z_1)
-        U_2 += (K[1:] - Z_2)
-
-        Y_0 += (L - W_0)
-        Y_1 += (L[:-1] - W_1)
-        Y_2 += (L[1:] - W_2)
-
-        # diagnostics, reporting, termination checks
-        Z_consensus = Z_0.copy()
-        Z_consensus[:-1] += Z_1
-        Z_consensus[1:] += Z_2
-        Z_consensus /= divisor[:, None, None]
-
-        U_consensus = U_0.copy()
-        U_consensus[:-1] += U_1
-        U_consensus[1:] += U_2
-        U_consensus /= divisor[:, None, None]
-
-        W_consensus = W_0.copy()
-        W_consensus[:-1] += W_1
-        W_consensus[1:] += W_2
-        W_consensus /= divisor[:, None, None]
-
-        Y_consensus = Y_0.copy()
-        Y_consensus[:-1] += Y_1
-        Y_consensus[1:] += Y_2
-        Y_consensus /= divisor[:, None, None]
-
-        check = convergence(
-            obj=objective(n_samples, emp_cov, R, Z_0, Z_1, Z_2, W_0, W_1, W_2,
-                          alpha, tau, beta, eta, psi, phi),
-            rnorm=np.sqrt(squared_norm(K - Z_consensus) +
-                          squared_norm(L - W_consensus) +
-                          squared_norm(K - L - R)),
-            snorm=np.sqrt(squared_norm(rho * (Z_consensus - Z_consensus_old)) +
-                          squared_norm(rho * (W_consensus - W_consensus_old)) +
-                          squared_norm(rho * (R - R_old))),
-            e_pri=np.sqrt(np.prod(K.shape) * 3) * tol + rtol * max(
-                np.sqrt(squared_norm(K) + squared_norm(L) + squared_norm(K - L)),
-                np.sqrt(squared_norm(Z_consensus) + squared_norm(W_consensus) + squared_norm(R))),
-            e_dual=np.sqrt(np.prod(K.shape) * 3) * tol +
-                rtol * np.sqrt(squared_norm(rho * (U_consensus)) +
-                               squared_norm(rho * (Y_consensus)) +
-                               squared_norm(rho * (X)))
-        )
-        Z_consensus_old = Z_consensus.copy()
-        W_consensus_old = W_consensus.copy()
-        R_old = R.copy()
-
-        if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check)
-
-        checks.append(check)
-        if check.rnorm <= check.e_pri and check.snorm <= check.e_dual:
-            break
-    else:
-        warnings.warn("Objective did not converge.")
-
-    return_list = [K, L, emp_cov]
-    if return_history:
-        return_list.append(checks)
-    if return_n_iter:
-        return_list.append(iteration_)
-    return return_list
-
-
-def time_latent_graph_lasso_v2(
-        emp_cov, alpha=1., tau=1., rho=1., beta=1., eta=1., max_iter=1000,
-        verbose=False, psi='laplacian', phi='laplacian', assume_centered=False,
-        tol=1e-4, rtol=1e-2, return_history=False, return_n_iter=True,
-        mode=None):
-    r"""Time-varying latent variable graphical lasso solver.
-
-    Solves the following problem via ADMM:
-        min sum_{i=1}^T -n_i log_likelihood(K_i-L_i) + alpha ||K_i||_{od,1}
-            + tau ||L_i||_*
-            + beta sum_{i=2}^T Psi(K_i - K_{i-1})
-            + eta sum_{i=2}^T Phi(L_i - L_{i-1})
-
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
-
-    Parameters
-    ----------
-    data_list : list of 2-dimensional matrices.
-        Input matrices.
-    alpha, tau : float, optional
-        Regularisation parameters.
-    rho : float, optional
-        Augmented Lagrangian parameter.
-    max_iter : int, optional
-        Maximum number of iterations.
-    tol : float, optional
-        Absolute tolerance for convergence.
-    rtol : float, optional
-        Relative tolerance for convergence.
-    return_history : bool, optional
-        Return the history of computed values.
-
-    Returns
-    -------
-    K, L : numpy.array, 3-dimensional (T x d x d)
-        Solution to the problem for each time t=1...T .
-    history : list
-        If return_history, then also a structure that contains the
-        objective value, the primal and dual residual norms, and tolerances
-        for the primal and dual residual norms at each iteration.
-    """
-    if psi == 'laplacian':
-        prox_psi = prox_laplacian
-        psi = squared_norm
-    elif psi == 'l1':
-        prox_psi = soft_thresholding_sign
-        psi = l1_norm
-    elif psi == 'l2':
-        prox_psi = blockwise_soft_thresholding
-        psi = np.linalg.norm
-    elif psi == 'linf':
-        prox_psi = prox_linf
-        psi = partial(np.linalg.norm, ord=np.inf)
-    else:
-        raise ValueError("Value of `psi` not understood.")
-    if phi == 'laplacian':
-        prox_phi = prox_laplacian
-        phi = squared_norm
-    elif phi == 'l1':
-        prox_phi = soft_thresholding_sign
-        phi = l1_norm
-    elif phi == 'l2':
-        prox_phi = blockwise_soft_thresholding
-        phi = np.linalg.norm
-    elif phi == 'linf':
-        prox_phi = prox_linf
-        phi = partial(np.linalg.norm, ord=np.inf)
-    else:
-        raise ValueError("Value of `phi` not understood.")
-
-    # emp_cov = np.array([empirical_covariance(
-    #     x, assume_centered=assume_centered) for x in data_list])
-
-    n_samples = np.array([s for s in [1.]])
-
-    K = np.zeros_like(emp_cov)
-    L = np.zeros_like(emp_cov)
-    X = np.zeros_like(emp_cov)
-    Z_0 = np.zeros_like(K)
-    Z_1 = np.zeros_like(K)[:-1]
-    Z_2 = np.zeros_like(K)[1:]
-    X_0 = np.zeros_like(K)
-    X_1 = np.zeros_like(K)[:-1]
-    X_2 = np.zeros_like(K)[1:]
-    W_0 = np.zeros_like(K)
-    W_1 = np.zeros_like(K)[:-1]
-    W_2 = np.zeros_like(K)[1:]
-    U_0 = np.zeros_like(emp_cov)
-    U_1 = np.zeros_like(emp_cov)[:-1]
-    U_2 = np.zeros_like(emp_cov)[1:]
-    Y_0 = np.zeros_like(emp_cov)
-    Y_1 = np.zeros_like(emp_cov)[:-1]
-    Y_2 = np.zeros_like(emp_cov)[1:]
-
-    U_consensus = np.zeros_like(emp_cov)
-    Y_consensus = np.zeros_like(emp_cov)
-    Z_consensus = np.zeros_like(emp_cov)
-    Z_consensus_old = np.zeros_like(emp_cov)
-    W_consensus = np.zeros_like(emp_cov)
-    W_consensus_old = np.zeros_like(emp_cov)
-    R_old = np.zeros_like(emp_cov)
-
-    # divisor for consensus variables, accounting for two less matrices
-    divisor = np.zeros(emp_cov.shape[0]) + 3
-    divisor[0] -= 1
-    divisor[-1] -= 1
-    # eta = np.divide(n_samples, divisor * rho)
-
-    checks = []
-    for iteration_ in range(max_iter):
-        # update R
-        A = K - L - X
-        A += Z_0 - W_0 - X_0
-        A[:-1] += Z_1 - W_1 - X_1
-        A[1:] += Z_2 - W_2 - X_2
-        A /= divisor[:, None, None] + 1
-
-        A += np.array(map(np.transpose, A))
-        A /= 2.
-        A *= - rho  # / n_samples[:, None, None]
-        A += emp_cov
-        R = np.array(map(prox_logdet, A, n_samples / (rho)))
-
-        # update K, L
-        K = L + R + X + Z_0 - U_0
-        K[:-1] += Z_1 - U_1
-        K[1:] += Z_2 - U_2
-        K /= divisor[:, None, None] + 1
-
-        L = K - R - X + W_0 - Y_0
-        L[:-1] += W_1 - Y_1
-        L[1:] += W_2 - Y_2
-        L /= divisor[:, None, None] + 1
-
-        # update Z_0
-        soft_thresholding = partial(soft_thresholding_sign, lamda=alpha / rho)
-        Z_0 = np.array(map(soft_thresholding, (K + U_0 + R + W_0 + X_0) / 2.))
-
-        # update Z_1, Z_2
-        A_1 = (K[:-1] + U_1 + R[:-1] + W_1 + X_1) / 2.
-        A_2 = (K[1:] + U_2 + R[1:] + W_2 + X_2) / 2.
-        prox_e = prox_psi(A_2 - A_1, lamda=2. * beta / rho)
-        Z_1 = .5 * (A_2 + A_1 - prox_e)
-        Z_2 = .5 * (A_2 + A_1 + prox_e)
-
-        # update W_0
-        A = (L + Y_0 - R + Z_0 - X_0) / 2.
-        W_0 = np.array(map(partial(prox_trace_indicator, lamda=tau / rho), A))
-
-        # update W_1, W_2
-        A_1 = (L[:-1] + Y_1 - R[:-1] + Z_1 - X_1) / 2.
-        A_2 = (L[1:] + Y_2 - R[1:] + Z_2 - X_2) / 2.
-        prox_e = prox_phi(A_2 - A_1, lamda=2. * eta / rho)
-        W_1 = .5 * (A_2 + A_1 - prox_e)
-        W_2 = .5 * (A_2 + A_1 + prox_e)
-
-        # update residuals
-        X += R - K + L
-        X_0 += (R - Z_0 + W_0)
-        X_1 += (R[:-1] - Z_1 + W_1)
-        X_2 += (R[1:] - Z_2 + W_2)
-
-        U_0 += (K - Z_0)
-        U_1 += (K[:-1] - Z_1)
-        U_2 += (K[1:] - Z_2)
-
-        Y_0 += (L - W_0)
-        Y_1 += (L[:-1] - W_1)
-        Y_2 += (L[1:] - W_2)
-
-        # diagnostics, reporting, termination checks
-        Z_consensus = Z_0.copy()
-        Z_consensus[:-1] += Z_1
-        Z_consensus[1:] += Z_2
-        Z_consensus /= divisor[:, None, None]
-
-        U_consensus = U_0.copy()
-        U_consensus[:-1] += U_1
-        U_consensus[1:] += U_2
-        U_consensus /= divisor[:, None, None]
-
-        W_consensus = W_0.copy()
-        W_consensus[:-1] += W_1
-        W_consensus[1:] += W_2
-        W_consensus /= divisor[:, None, None]
-
-        X_consensus = X_0.copy()
-        X_consensus[:-1] += X_1
-        X_consensus[1:] += X_2
-        X_consensus /= divisor[:, None, None]
-
-        Y_consensus = Y_0.copy()
-        Y_consensus[:-1] += Y_1
-        Y_consensus[1:] += Y_2
-        Y_consensus /= divisor[:, None, None]
-
-        check = convergence(
-            obj=objective(n_samples, emp_cov, R, Z_0, Z_1, Z_2, W_0, W_1, W_2,
-                          alpha, tau, beta, eta, psi, phi),
-            rnorm=np.sqrt(squared_norm(K - Z_consensus) +
-                          squared_norm(L - W_consensus) +
-                          squared_norm(Z_consensus - W_consensus - R) +
-                          squared_norm(K - L - R)),
-            snorm=np.sqrt(squared_norm(rho * (Z_consensus - Z_consensus_old)) +
-                          squared_norm(rho * (W_consensus - W_consensus_old)) +
-                          squared_norm(rho * (Z_consensus - W_consensus - (Z_consensus_old - W_consensus_old))) +
-                          squared_norm(rho * (R - R_old))),
-            e_pri=np.sqrt(np.prod(K.shape) * 4) * tol + rtol * max(
-                np.sqrt(squared_norm(K) + squared_norm(L) + squared_norm(K - L)
-                        + squared_norm(Z_consensus - W_consensus)),
-                np.sqrt(squared_norm(Z_consensus) + squared_norm(W_consensus) + squared_norm(R) + squared_norm(R))),
-            e_dual=np.sqrt(np.prod(K.shape) * 4) * tol +
-                rtol * np.sqrt(squared_norm(rho * (U_consensus)) +
-                               squared_norm(rho * (Y_consensus)) +
-                               squared_norm(rho * (X_consensus)) +
-                               squared_norm(rho * (X)))
-        )
-        Z_consensus_old = Z_consensus.copy()
-        W_consensus_old = W_consensus.copy()
-        R_old = R.copy()
-
-        if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check)
-
-        checks.append(check)
-        if (check.rnorm <= check.e_pri and check.snorm <= check.e_dual):
-            break
-    else:
-        warnings.warn("Objective did not converge.")
-
-    return_list = [K, L, emp_cov]
-    if return_history:
-        return_list.append(checks)
-    if return_n_iter:
-        return_list.append(iteration_)
-    return return_list
diff --git a/regain/covariance/__old/time_graph_lasso_.py b/regain/covariance/__old/time_graph_lasso_.py
deleted file mode 100644
index c40a18f..0000000
--- a/regain/covariance/__old/time_graph_lasso_.py
+++ /dev/null
@@ -1,204 +0,0 @@
-"""Sparse inverse covariance selection over time via ADMM.
-
-More information can be found in the paper linked at:
-[cite Hallac]
-"""
-from __future__ import division
-
-import numpy as np
-import warnings
-
-from functools import partial
-from six.moves import range
-from sklearn.covariance import empirical_covariance
-from sklearn.utils.extmath import fast_logdet, squared_norm
-
-from regain.norm import l1_od_norm, l1_norm
-from regain.prox import prox_logdet
-from regain.prox import soft_thresholding_od, soft_thresholding_sign
-from regain.utils import convergence
-from regain.validation import check_norm_prox
-
-
-def log_likelihood(emp_cov, precision):
-    """Gaussian log-likelihood without constant term."""
-    return fast_logdet(precision) - np.sum(emp_cov * precision)
-
-
-def log_likelihood_trace(emp_cov, precision):
-    """Gaussian log-likelihood without constant term."""
-    # return fast_logdet(precision) - np.trace(emp_cov * precision)
-    return np.trace(emp_cov.dot(precision))
-
-
-def objective(S, K, Z_0, Z_1, Z_2, alpha, beta, psi):
-    """Objective function for time-varying graphical lasso."""
-    obj = np.sum(- log_likelihood(emp_cov, precision)
-                 for emp_cov, precision in zip(S, K))
-    obj += alpha * np.sum(map(l1_od_norm, Z_0))
-    obj += beta * np.sum(map(psi, Z_2 - Z_1))
-    return obj
-
-
-def time_graph_lasso(
-        emp_cov, alpha=1, rho=1, beta=1, max_iter=1000,
-        verbose=False, psi='laplacian',
-        tol=1e-4, rtol=1e-2, return_history=False, return_n_iter=True):
-    """Time-varying graphical lasso solver.
-
-    Solves the following problem via ADMM:
-        minimize  trace(S*X) - log det X + lambda*||X||_1
-
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
-
-    Parameters
-    ----------
-    data_list : list of 2-dimensional matrices.
-        Input matrices.
-    lamda : float, optional
-        Regularisation parameter.
-    rho : float, optional
-        Augmented Lagrangian parameter.
-    alpha : float, optional
-        Over-relaxation parameter (typically between 1.0 and 1.8).
-    max_iter : int, optional
-        Maximum number of iterations.
-    tol : float, optional
-        Absolute tolerance for convergence.
-    rtol : float, optional
-        Relative tolerance for convergence.
-    return_history : bool, optional
-        Return the history of computed values.
-
-    Returns
-    -------
-    X : numpy.array, 2-dimensional
-        Solution to the problem.
-    history : list
-        If return_history, then also a structure that contains the
-        objective value, the primal and dual residual norms, and tolerances
-        for the primal and dual residual norms at each iteration.
-
-    """
-    psi, prox_psi = check_norm_prox(psi)
-
-    # S = np.array(map(empirical_covariance, data_list))
-    # n_samples = np.array([1. for data in data_list])
-
-    # K = np.zeros_like(emp_cov)
-    Z_0 = np.zeros_like(emp_cov)
-    Z_1 = np.zeros_like(emp_cov)[:-1]
-    Z_2 = np.zeros_like(emp_cov)[1:]
-    U_0 = np.zeros_like(emp_cov)
-    U_1 = np.zeros_like(emp_cov)[:-1]
-    U_2 = np.zeros_like(emp_cov)[1:]
-
-    U_consensus = np.zeros_like(emp_cov)
-    Z_consensus = np.zeros_like(emp_cov)
-    Z_consensus_old = np.zeros_like(emp_cov)
-    Z_0_old = np.zeros_like(Z_0)
-    Z_1_old = np.zeros_like(Z_1)
-    Z_2_old = np.zeros_like(Z_2)
-
-    # divisor for consensus variables, accounting for two less matrices
-    divisor = np.full(emp_cov.shape[0], 3, dtype=float)
-    divisor[0] -= 1
-    divisor[-1] -= 1
-
-    checks = []
-    for iteration_ in range(max_iter):
-        # x-update
-        # A = Z_consensus - U_consensus
-        A = Z_0 - U_0
-        A[:-1] += Z_1 - U_1
-        A[1:] += Z_2 - U_2
-        A /= divisor[:, None, None]
-
-        # A += np.array(map(np.transpose, A))
-        # A /= 2.
-
-        # A *= - rho / n_samples[:, None, None]
-        A *= - rho
-        A += emp_cov
-
-        K = np.array([prox_logdet(a, lamda=1. / rho) for a in A])
-
-        # z-update with relaxation
-        # Zold = Z
-        # X_hat = alpha * X + (1 - alpha) * Zold
-        soft_thresholding = partial(soft_thresholding_sign, lamda=alpha / rho)
-        Z_0 = np.array(map(soft_thresholding, K + U_0))
-
-        # other Zs
-        if beta != 0:
-            A_1 = K[:-1] + U_1
-            A_2 = K[1:] + U_2
-            prox_e = prox_psi(A_2 - A_1, lamda=2. * beta / rho)
-            Z_1 = .5 * (A_1 + A_2 - prox_e)
-            Z_2 = .5 * (A_1 + A_2 + prox_e)
-        else:
-            Z_1 = Z_0[:-1].copy()
-            Z_2 = Z_0[1:].copy()
-
-        # update residuals
-        U_0 += K - Z_0
-        U_1 += K[:-1] - Z_1
-        U_2 += K[1:] - Z_2
-
-        # diagnostics, reporting, termination checks
-        Z_consensus = Z_0.copy()
-        Z_consensus[:-1] += Z_1
-        Z_consensus[1:] += Z_2
-        Z_consensus /= divisor[:, None, None]
-
-        U_consensus = U_0.copy()
-        U_consensus[:-1] += U_1
-        U_consensus[1:] += U_2
-        U_consensus /= divisor[:, None, None]
-
-        # check = convergence(
-        #     obj=objective(emp_cov, K, Z_0, Z_1, Z_2, alpha, beta, psi),
-        #     rnorm=np.linalg.norm(K - Z_consensus),
-        #     snorm=np.linalg.norm(rho * (Z_consensus - Z_consensus_old)),
-        #     e_pri=np.sqrt(np.prod(K.shape)) * tol + rtol * max(
-        #         np.linalg.norm(K), np.linalg.norm(Z_consensus)),
-        #     e_dual=np.sqrt(np.prod(K.shape)) * tol + rtol * np.linalg.norm(
-        #         rho * U_consensus)
-        # )
-        check = convergence(
-            obj=objective(emp_cov, K, Z_0, Z_1, Z_2, alpha, beta, psi),
-            rnorm=np.sqrt(squared_norm(K - Z_0) +
-                          squared_norm(K[:-1] - Z_1) +
-                          squared_norm(K[1:] - Z_2)),
-            snorm=np.sqrt(squared_norm(rho * (Z_0 - Z_0_old)) +
-                          squared_norm(rho * (Z_1 - Z_1_old)) +
-                          squared_norm(rho * (Z_2 - Z_2_old))),
-            e_pri=np.sqrt(np.prod(K.shape[1:]) * (3 * K.shape[0] - 2)) * tol + rtol * max(
-                np.sqrt(squared_norm(Z_0) + squared_norm(Z_1) + squared_norm(Z_2)), np.sqrt(squared_norm(K) + squared_norm(K[:-1]) + squared_norm(K[1:]))),
-            e_dual=np.sqrt(np.prod(K.shape[1:]) * (3 * K.shape[0] - 2)) * tol + rtol * np.sqrt(
-                squared_norm(rho * U_0) + squared_norm(rho * U_1) +
-                squared_norm(rho * U_2))
-        )
-
-        Z_consensus_old = Z_consensus.copy()
-        Z_0_old = Z_0.copy()
-        Z_1_old = Z_1.copy()
-        Z_2_old = Z_2.copy()
-
-        if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check)
-
-        checks.append(check)
-        if check.rnorm <= check.e_pri and check.snorm <= check.e_dual:
-            break
-    else:
-        warnings.warn("Objective did not converge.")
-
-    return_list = [K, emp_cov]
-    if return_history:
-        return_list.append(checks)
-    if return_n_iter:
-        return_list.append(iteration_)
-    return return_list
diff --git a/regain/covariance/graph_lasso_.py b/regain/covariance/graphical_lasso_.py
similarity index 83%
rename from regain/covariance/graph_lasso_.py
rename to regain/covariance/graphical_lasso_.py
index 2b73e9f..4f758b4 100644
--- a/regain/covariance/graph_lasso_.py
+++ b/regain/covariance/graphical_lasso_.py
@@ -8,6 +8,7 @@ from __future__ import division
 import warnings
 
 import numpy as np
+from scipy import linalg
 from six.moves import range
 from sklearn.covariance import empirical_covariance
 from sklearn.utils.extmath import fast_logdet
@@ -32,21 +33,21 @@ def logl(emp_cov, precision):
 
 
 def objective(emp_cov, x, z, alpha):
-    """Graph lasso objective."""
+    """Graphical lasso objective."""
     return -logl(emp_cov, x) + alpha * l1_od_norm(z)
 
 
-def graph_lasso(
+def graphical_lasso(
         emp_cov, alpha=0.01, rho=1, over_relax=1, max_iter=100, verbose=False,
         tol=1e-4, rtol=1e-4, return_history=False, return_n_iter=True,
-        update_rho_options=None, compute_objective=True, mode='admm'):
-    """Graph lasso solver via ADMM.
+        update_rho_options=None, compute_objective=True, init='empirical'):
+    r"""Graphical lasso solver via ADMM.
 
-    Solves the following problem via ADMM:
+    Solves the following problem:
         minimize  trace(S*X) - log det X + alpha ||X||_{od,1}
 
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
+    where S = (1/n) X^T \times X is the empirical covariance of the data
+    matrix X (training observations by features).
 
     Parameters
     ----------
@@ -70,6 +71,14 @@ def graph_lasso(
         Return the number of iteration before convergence.
     verbose : bool, default False
         Print info at each iteration.
+    update_rho_options : dict, optional
+        Arguments for the rho update.
+        See regain.update_rules.update_rho function for more information.
+    compute_objective : bool, default True
+        Choose to compute the objective value.
+    init : ('empirical', 'zero')
+        Choose how to initialize the precision matrix, with the inverse
+        empirical covariance or zero matrix.
 
     Returns
     -------
@@ -86,13 +95,15 @@ def graph_lasso(
 
     """
     _, n_features = emp_cov.shape
-    covariance_ = emp_cov.copy()
-    covariance_ *= 0.95
-    covariance_.flat[::n_features + 1] = emp_cov.flat[::n_features + 1]
-    from scipy import linalg
-    Z = linalg.pinvh(covariance_)
 
-    # Z = np.zeros_like(emp_cov)
+    if init == 'empirical':
+        covariance_ = emp_cov.copy()
+        covariance_ *= 0.95
+        covariance_.flat[::n_features + 1] = emp_cov.flat[::n_features + 1]
+        Z = linalg.pinvh(covariance_)
+    else:
+        Z = np.zeros_like(emp_cov)
+
     U = np.zeros_like(emp_cov)
     Z_old = np.zeros_like(Z)
 
@@ -116,9 +127,8 @@ def graph_lasso(
         rnorm = np.linalg.norm(K - Z, 'fro')
         snorm = rho * np.linalg.norm(Z - Z_old, 'fro')
         check = convergence(
-            obj=obj, rnorm=rnorm, snorm=snorm,
-            e_pri=np.sqrt(K.size) * tol + rtol * max(
-                np.linalg.norm(K, 'fro'), np.linalg.norm(Z, 'fro')),
+            obj=obj, rnorm=rnorm, snorm=snorm, e_pri=np.sqrt(K.size) * tol +
+            rtol * max(np.linalg.norm(K, 'fro'), np.linalg.norm(Z, 'fro')),
             e_dual=np.sqrt(K.size) * tol + rtol * rho * np.linalg.norm(U))
 
         Z_old = Z.copy()
@@ -148,7 +158,7 @@ def graph_lasso(
     return return_list
 
 
-class GraphLasso(GraphLasso):
+class GraphicalLasso(GraphLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -210,7 +220,7 @@ class GraphLasso(GraphLasso):
             self, alpha=0.01, rho=1., over_relax=1., max_iter=100, mode='admm',
             tol=1e-4, rtol=1e-4, verbose=False, assume_centered=False,
             update_rho_options=None, compute_objective=True):
-        super(GraphLasso, self).__init__(
+        super(GraphicalLasso, self).__init__(
             alpha=alpha, tol=tol, max_iter=max_iter, verbose=verbose,
             assume_centered=assume_centered, mode=mode)
         self.rho = rho
@@ -220,7 +230,7 @@ class GraphLasso(GraphLasso):
         self.compute_objective = compute_objective
 
     def _fit(self, emp_cov):
-        """Fit the GraphLasso model to X.
+        """Fit the GraphicalLasso model to X.
 
         Parameters
         ----------
@@ -228,16 +238,16 @@ class GraphLasso(GraphLasso):
             Empirical covariance of data.
 
         """
-        self.precision_, self.covariance_, self.n_iter_ = graph_lasso(
+        self.precision_, self.covariance_, self.n_iter_ = graphical_lasso(
             emp_cov, alpha=self.alpha, tol=self.tol, rtol=self.rtol,
             max_iter=self.max_iter, over_relax=self.over_relax, rho=self.rho,
             verbose=self.verbose, return_n_iter=True, return_history=False,
-            mode=self.mode, update_rho_options=self.update_rho_options,
+            update_rho_options=self.update_rho_options,
             compute_objective=self.compute_objective)
         return self
 
     def fit(self, X, y=None):
-        """Fit the GraphLasso model to X.
+        """Fit the GraphicalLasso model to X.
 
         Parameters
         ----------
diff --git a/regain/covariance/latent_graph_lasso_.py b/regain/covariance/latent_graphical_lasso_.py
similarity index 73%
rename from regain/covariance/latent_graph_lasso_.py
rename to regain/covariance/latent_graphical_lasso_.py
index b81260a..a2469fd 100644
--- a/regain/covariance/latent_graph_lasso_.py
+++ b/regain/covariance/latent_graphical_lasso_.py
@@ -7,8 +7,8 @@ import numpy as np
 from scipy import linalg
 from six.moves import range
 
-from regain.covariance.graph_lasso_ import GraphLasso, logl
-from regain.norm import l1_od_norm
+from regain.covariance.graphical_lasso_ import GraphicalLasso
+from regain.covariance.graphical_lasso_ import objective as obj_gl
 from regain.prox import prox_logdet, prox_trace_indicator, soft_thresholding
 from regain.update_rules import update_rho
 from regain.utils import convergence
@@ -16,24 +16,22 @@ from regain.utils import convergence
 
 def objective(emp_cov, R, K, L, alpha, tau):
     """Objective function for latent graphical lasso."""
-    obj = - logl(emp_cov, R)
-    obj += alpha * l1_od_norm(K)
+    obj = obj_gl(emp_cov, R, K, alpha=alpha)
     obj += tau * np.linalg.norm(L, ord='nuc')
     return obj
 
 
-def latent_graph_lasso(
-        emp_cov, alpha=1., tau=1., rho=1., max_iter=100,
-        verbose=False, tol=1e-4, rtol=1e-2, return_history=False,
-        return_n_iter=True, mode=None,
-        update_rho_options=None, compute_objective=True):
-    r"""Latent variable graphical lasso solver.
+def latent_graphical_lasso(
+        emp_cov, alpha=1., tau=1., rho=1., max_iter=100, verbose=False,
+        tol=1e-4, rtol=1e-2, return_history=False, return_n_iter=True,
+        update_rho_options=None, compute_objective=True, init='empirical'):
+    r"""Latent variable graphical lasso solver via ADMM.
 
-    Solves the following problem via ADMM:
+    Solves the following problem:
         min - log_likelihood(S, K-L) + alpha ||K||_{od,1} + tau ||L_i||_*
 
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
+    where S = (1/n) X^T \times X is the empirical covariance of the data
+    matrix X (training observations by features).
 
     Parameters
     ----------
@@ -55,6 +53,14 @@ def latent_graph_lasso(
         Return the number of iteration before convergence.
     verbose : bool, default False
         Print info at each iteration.
+    update_rho_options : dict, optional
+        Arguments for the rho update.
+        See regain.update_rules.update_rho function for more information.
+    compute_objective : bool, default True
+        Choose to compute the objective value.
+    init : ('empirical', 'zero')
+        Choose how to initialize the precision matrix, with the inverse
+        empirical covariance or zero matrix.
 
     Returns
     -------
@@ -70,13 +76,16 @@ def latent_graph_lasso(
         for the primal and dual residual norms at each iteration.
 
     """
-    # _, n_features = emp_cov.shape
-    # covariance_ = emp_cov.copy()
-    # covariance_ *= 0.95
-    # covariance_.flat[::n_features + 1] = emp_cov.flat[::n_features + 1]
-    # K = linalg.pinvh(covariance_)
+    _, n_features = emp_cov.shape
+
+    if init == 'empirical':
+        covariance_ = emp_cov.copy()
+        covariance_ *= 0.95
+        covariance_.flat[::n_features + 1] = emp_cov.flat[::n_features + 1]
+        K = linalg.pinvh(covariance_)
+    else:
+        K = np.zeros_like(emp_cov)
 
-    K = np.zeros_like(emp_cov)
     L = np.zeros_like(emp_cov)
     U = np.zeros_like(emp_cov)
     R_old = np.zeros_like(emp_cov)
@@ -106,24 +115,24 @@ def latent_graph_lasso(
         rnorm = np.linalg.norm(R - K + L)
         snorm = rho * np.linalg.norm(R - R_old)
         check = convergence(
-            obj=obj, rnorm=rnorm, snorm=snorm,
-            e_pri=np.sqrt(R.size) * tol + rtol * max(
-                np.linalg.norm(R), np.linalg.norm(K - L)),
-            e_dual=np.sqrt(R.size) * tol + rtol * rho * np.linalg.norm(U)
-        )
+            obj=obj, rnorm=rnorm, snorm=snorm, e_pri=np.sqrt(R.size) * tol +
+            rtol * max(np.linalg.norm(R), np.linalg.norm(K - L)),
+            e_dual=np.sqrt(R.size) * tol + rtol * rho * np.linalg.norm(U))
         R_old = R.copy()
 
         if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
+            print(
+                "obj: %.4f, rnorm: %.4f, snorm: %.4f,"
+                "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
 
         checks.append(check)
         if check.rnorm <= check.e_pri and check.snorm <= check.e_dual:
             break
         if check.obj == np.inf:
             break
-        rho_new = update_rho(rho, rnorm, snorm, iteration=iteration_,
-                             **(update_rho_options or {}))
+        rho_new = update_rho(
+            rho, rnorm, snorm, iteration=iteration_,
+            **(update_rho_options or {}))
         # scaled dual variables should be also rescaled
         U *= rho / rho_new
         rho = rho_new
@@ -139,7 +148,7 @@ def latent_graph_lasso(
     return return_list
 
 
-class LatentGraphLasso(GraphLasso):
+class LatentGraphicalLasso(GraphicalLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -204,13 +213,13 @@ class LatentGraphLasso(GraphLasso):
 
     """
 
-    def __init__(self, alpha=0.01, tau=1., rho=1., tol=1e-4, rtol=1e-4,
-                 max_iter=100, verbose=False, assume_centered=False,
-                 mode='admm', update_rho_options=None, compute_objective=True):
-        super(LatentGraphLasso, self).__init__(
-            alpha=alpha, rho=rho,
-            tol=tol, rtol=rtol, max_iter=max_iter, verbose=verbose,
-            assume_centered=assume_centered, mode=mode,
+    def __init__(
+            self, alpha=0.01, tau=1., rho=1., tol=1e-4, rtol=1e-4,
+            max_iter=100, verbose=False, assume_centered=False, mode='admm',
+            update_rho_options=None, compute_objective=True):
+        super(LatentGraphicalLasso, self).__init__(
+            alpha=alpha, rho=rho, tol=tol, rtol=rtol, max_iter=max_iter,
+            verbose=verbose, assume_centered=assume_centered, mode=mode,
             update_rho_options=update_rho_options,
             compute_objective=compute_objective)
         self.tau = tau
@@ -228,7 +237,7 @@ class LatentGraphLasso(GraphLasso):
         return self.precision_ - self.latent_
 
     def _fit(self, emp_cov):
-        """Fit the LatentGraphLasso model to X.
+        """Fit the LatentGraphicalLasso model to X.
 
         Parameters
         ----------
@@ -237,9 +246,9 @@ class LatentGraphLasso(GraphLasso):
 
         """
         self.precision_, self.latent_, self.covariance_, self.n_iter_ = \
-            latent_graph_lasso(
+            latent_graphical_lasso(
                 emp_cov, alpha=self.alpha, tau=self.tau, rho=self.rho,
-                mode=self.mode, tol=self.tol, rtol=self.rtol,
+                tol=self.tol, rtol=self.rtol,
                 max_iter=self.max_iter, verbose=self.verbose,
                 return_n_iter=True, return_history=False,
                 update_rho_options=self.update_rho_options,
diff --git a/regain/covariance/latent_time_graph_lasso_.py b/regain/covariance/latent_time_graphical_lasso_.py
similarity index 84%
rename from regain/covariance/latent_time_graph_lasso_.py
rename to regain/covariance/latent_time_graphical_lasso_.py
index d0ff7e9..7e815c7 100644
--- a/regain/covariance/latent_time_graph_lasso_.py
+++ b/regain/covariance/latent_time_graphical_lasso_.py
@@ -9,8 +9,8 @@ from scipy import linalg
 from six.moves import map, range, zip
 from sklearn.utils.extmath import squared_norm
 
-from regain.covariance.time_graph_lasso_ import TimeGraphLasso
-from regain.covariance.time_graph_lasso_ import objective as obj_tgl
+from regain.covariance.time_graphical_lasso_ import TimeGraphicalLasso
+from regain.covariance.time_graphical_lasso_ import objective as obj_tgl
 from regain.prox import prox_logdet, prox_trace_indicator, soft_thresholding
 from regain.update_rules import update_rho
 from regain.utils import convergence
@@ -23,7 +23,11 @@ def objective(
     """Objective function for latent variable time-varying graphical lasso."""
     # obj = sum(- n * logl(s, r) for s, r, n in zip(S, R, n_samples))
     obj = obj_tgl(n_samples, S, R, Z_0, Z_1, Z_2, alpha, beta, psi)
-    obj += tau * sum(map(partial(np.linalg.norm, ord='nuc'), W_0))
+
+    if isinstance(tau, np.ndarray):
+        obj += sum(np.linalg.norm(t * w, ord='nuc') for t, w in zip(tau, W_0))
+    else:
+        obj += tau * sum(map(partial(np.linalg.norm, ord='nuc'), W_0))
 
     if isinstance(eta, np.ndarray):
         obj += sum(b[0][0] * m for b, m in zip(eta, map(phi, W_2 - W_1)))
@@ -32,21 +36,22 @@ def objective(
     return obj
 
 
-def latent_time_graph_lasso(
+def latent_time_graphical_lasso(
         emp_cov, alpha=0.01, tau=1., rho=1., beta=1., eta=1., max_iter=100,
-        verbose=False, psi='laplacian', phi='laplacian', mode='admm', tol=1e-4,
-        rtol=1e-4, assume_centered=False, n_samples=None, return_history=False,
-        return_n_iter=True, update_rho_options=None, compute_objective=True):
-    r"""Time-varying latent variable graphical lasso solver.
+        n_samples=None, verbose=False, psi='laplacian', phi='laplacian',
+        mode='admm', tol=1e-4, rtol=1e-4, return_history=False,
+        return_n_iter=True, update_rho_options=None, compute_objective=True,
+        init='empirical'):
+    r"""Latent variable time-varying graphical lasso solver.
 
     Solves the following problem via ADMM:
-        min sum_{i=1}^T -n_i log_likelihood(K_i-L_i) + alpha ||K_i||_{od,1}
-            + tau ||L_i||_*
-            + beta sum_{i=2}^T Psi(K_i - K_{i-1})
-            + eta sum_{i=2}^T Phi(L_i - L_{i-1})
+      min sum_{i=1}^T -n_i log_likelihood(S_i, K_i-L_i) + alpha ||K_i||_{od,1}
+          + tau ||L_i||_*
+          + beta sum_{i=2}^T Psi(K_i - K_{i-1})
+          + eta sum_{i=2}^T Phi(L_i - L_{i-1})
 
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
+    where S_i = (1/n_i) X_i^T \times X_i is the empirical covariance of data
+    matrix X (training observations by features).
 
     Parameters
     ----------
@@ -58,12 +63,26 @@ def latent_time_graph_lasso(
         Augmented Lagrangian parameter.
     max_iter : int, optional
         Maximum number of iterations.
+    n_samples : ndarray
+        Number of samples available for each time point.
     tol : float, optional
         Absolute tolerance for convergence.
     rtol : float, optional
         Relative tolerance for convergence.
     return_history : bool, optional
         Return the history of computed values.
+    return_n_iter : bool, optional
+        Return the number of iteration before convergence.
+    verbose : bool, default False
+        Print info at each iteration.
+    update_rho_options : dict, optional
+        Arguments for the rho update.
+        See regain.update_rules.update_rho function for more information.
+    compute_objective : bool, default True
+        Choose to compute the objective value.
+    init : ('empirical', 'zero')
+        Choose how to initialize the precision matrix, with the inverse
+        empirical covariance or zero matrix.
 
     Returns
     -------
@@ -78,7 +97,17 @@ def latent_time_graph_lasso(
     psi, prox_psi, psi_node_penalty = check_norm_prox(psi)
     phi, prox_phi, phi_node_penalty = check_norm_prox(phi)
 
-    Z_0 = np.zeros_like(emp_cov)
+    if init == 'empirical':
+        n_times, _, n_features = emp_cov.shape
+        covariance_ = emp_cov.copy()
+        covariance_ *= 0.95
+        Z_0 = np.empty_like(emp_cov)
+        for i, (c, e) in enumerate(zip(covariance_, emp_cov)):
+            c.flat[::n_features + 1] = e.flat[::n_features + 1]
+            Z_0[i] = linalg.pinvh(c)
+    else:
+        Z_0 = np.zeros_like(emp_cov)
+
     Z_1 = np.zeros_like(Z_0)[:-1]
     Z_2 = np.zeros_like(Z_0)[1:]
     W_0 = np.zeros_like(Z_0)
@@ -240,7 +269,7 @@ def latent_time_graph_lasso(
     return return_list
 
 
-class LatentTimeGraphLasso(TimeGraphLasso):
+class LatentTimeGraphicalLasso(TimeGraphicalLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -314,11 +343,8 @@ class LatentTimeGraphLasso(TimeGraphLasso):
     covariance_ : array-like, shape (n_times, n_features, n_features)
         Estimated covariance matrix
 
-    precision_ : array-like, shape (n_times, n_features, n_features)
-        Estimated pseudo inverse matrix.
-
-    latent_ : array-like, shape (n_times, n_features, n_features)
-        Estimated latent variable matrix.
+    precision_, latent_ : array-like, shape (n_times, n_features, n_features)
+        Estimated precision and latent variables matrix.
 
     n_iter_ : int
         Number of iterations run.
@@ -331,7 +357,7 @@ class LatentTimeGraphLasso(TimeGraphLasso):
             phi='laplacian', max_iter=100, verbose=False,
             assume_centered=False, update_rho_options=None,
             compute_objective=True):
-        super(LatentTimeGraphLasso, self).__init__(
+        super(LatentTimeGraphicalLasso, self).__init__(
             alpha=alpha, beta=beta, mode=mode, rho=rho, tol=tol, rtol=rtol,
             psi=psi, max_iter=max_iter, verbose=verbose,
             time_on_axis=time_on_axis, assume_centered=assume_centered,
@@ -354,7 +380,7 @@ class LatentTimeGraphLasso(TimeGraphLasso):
         return self.precision_ - self.latent_
 
     def _fit(self, emp_cov, n_samples):
-        """Fit the LatentTimeGraphLasso model to X.
+        """Fit the LatentTimeGraphicalLasso model to X.
 
         Parameters
         ----------
@@ -363,7 +389,7 @@ class LatentTimeGraphLasso(TimeGraphLasso):
 
         """
         self.precision_, self.latent_, self.covariance_, self.n_iter_ = \
-            latent_time_graph_lasso(
+            latent_time_graphical_lasso(
                 emp_cov, n_samples=n_samples,
                 alpha=self.alpha, tau=self.tau, rho=self.rho,
                 beta=self.beta, eta=self.eta, mode=self.mode,
diff --git a/regain/covariance/time_graph_lasso_.py b/regain/covariance/time_graphical_lasso_.py
similarity index 70%
rename from regain/covariance/time_graph_lasso_.py
rename to regain/covariance/time_graphical_lasso_.py
index 4fbbba7..f821668 100644
--- a/regain/covariance/time_graph_lasso_.py
+++ b/regain/covariance/time_graphical_lasso_.py
@@ -8,27 +8,26 @@ from __future__ import division
 import warnings
 
 import numpy as np
-import scipy.sparse as sp
 from scipy import linalg
 from six.moves import map, range, zip
 from sklearn.covariance import empirical_covariance, log_likelihood
 from sklearn.utils.extmath import squared_norm
-from sklearn.utils.validation import check_array
 
-from regain.covariance.graph_lasso_ import GraphLasso, logl
+from regain.covariance.graphical_lasso_ import GraphicalLasso, logl
 from regain.norm import l1_od_norm
 from regain.prox import prox_logdet, soft_thresholding
 from regain.update_rules import update_rho
 from regain.utils import convergence, error_norm_time
-from regain.validation import check_array_dimensions, check_norm_prox
+from regain.validation import check_input, check_norm_prox
 
 
 def loss(S, K, n_samples=None):
     """Loss function for time-varying graphical lasso."""
     if n_samples is None:
         n_samples = np.ones(S.shape[0])
-    return sum(-ni * logl(emp_cov, precision)
-               for emp_cov, precision, ni in zip(S, K, n_samples))
+    return sum(
+        -ni * logl(emp_cov, precision)
+        for emp_cov, precision, ni in zip(S, K, n_samples))
 
 
 def objective(n_samples, S, K, Z_0, Z_1, Z_2, alpha, beta, psi):
@@ -39,7 +38,7 @@ def objective(n_samples, S, K, Z_0, Z_1, Z_2, alpha, beta, psi):
         obj += sum(l1_od_norm(a * z) for a, z in zip(alpha, Z_0))
     else:
         obj += alpha * sum(map(l1_od_norm, Z_0))
-    
+
     if isinstance(beta, np.ndarray):
         obj += sum(b[0][0] * m for b, m in zip(beta, map(psi, Z_2 - Z_1)))
     else:
@@ -47,20 +46,20 @@ def objective(n_samples, S, K, Z_0, Z_1, Z_2, alpha, beta, psi):
     return obj
 
 
-def time_graph_lasso(
+def time_graphical_lasso(
         emp_cov, alpha=0.01, rho=1, beta=1, max_iter=100, n_samples=None,
         verbose=False, psi='laplacian', tol=1e-4, rtol=1e-4,
         return_history=False, return_n_iter=True, mode='admm',
-        update_rho_options=None, compute_objective=True,
-        stop_at=None, stop_when=1e-4):
+        update_rho_options=None, compute_objective=True, stop_at=None,
+        stop_when=1e-4, init='empirical'):
     """Time-varying graphical lasso solver.
 
     Solves the following problem via ADMM:
-        minimize  trace(S*X) - log det X + lambda*||X||_1
+        min sum_{i=1}^T -n_i log_likelihood(S_i, K_i) + alpha*||K_i||_{od,1}
             + beta sum_{i=2}^T Psi(K_i - K_{i-1})
 
-    where S is the empirical covariance of the data
-    matrix D (training observations by features).
+    where S_i = (1/n_i) X_i^T \times X_i is the empirical covariance of data
+    matrix X (training observations by features).
 
     Parameters
     ----------
@@ -72,17 +71,31 @@ def time_graph_lasso(
         Augmented Lagrangian parameter.
     max_iter : int, optional
         Maximum number of iterations.
+    n_samples : ndarray
+        Number of samples available for each time point.
     tol : float, optional
         Absolute tolerance for convergence.
     rtol : float, optional
         Relative tolerance for convergence.
     return_history : bool, optional
         Return the history of computed values.
+    return_n_iter : bool, optional
+        Return the number of iteration before convergence.
+    verbose : bool, default False
+        Print info at each iteration.
+    update_rho_options : dict, optional
+        Arguments for the rho update.
+        See regain.update_rules.update_rho function for more information.
+    compute_objective : bool, default True
+        Choose to compute the objective value.
+    init : ('empirical', 'zero')
+        Choose how to initialize the precision matrix, with the inverse
+        empirical covariance or zero matrix.
 
     Returns
     -------
-    X : numpy.array, 2-dimensional
-        Solution to the problem.
+    K : numpy.array, 3-dimensional (T x d x d)
+        Solution to the problem for each time t=1...T .
     history : list
         If return_history, then also a structure that contains the
         objective value, the primal and dual residual norms, and tolerances
@@ -91,18 +104,20 @@ def time_graph_lasso(
     """
     psi, prox_psi, psi_node_penalty = check_norm_prox(psi)
 
-    n_times, _, n_features = emp_cov.shape
-    covariance_ = emp_cov.copy()
-    covariance_ *= 0.95
-    K = np.empty_like(emp_cov)
-    for i, (c, e) in enumerate(zip(covariance_, emp_cov)):
-        c.flat[::n_features + 1] = e.flat[::n_features + 1]
-        K[i] = linalg.pinvh(c)
-
-    # K = np.zeros_like(emp_cov)
-    Z_0 = K.copy() # np.zeros_like(emp_cov)
-    Z_1 = K.copy()[:-1] # np.zeros_like(emp_cov)[:-1]
-    Z_2 = K.copy()[1:] # np.zeros_like(emp_cov)[1:]
+    if init == 'empirical':
+        n_times, _, n_features = emp_cov.shape
+        covariance_ = emp_cov.copy()
+        covariance_ *= 0.95
+        K = np.empty_like(emp_cov)
+        for i, (c, e) in enumerate(zip(covariance_, emp_cov)):
+            c.flat[::n_features + 1] = e.flat[::n_features + 1]
+            K[i] = linalg.pinvh(c)
+    else:
+        K = np.zeros_like(emp_cov)
+
+    Z_0 = K.copy()  # np.zeros_like(emp_cov)
+    Z_1 = K.copy()[:-1]  # np.zeros_like(emp_cov)[:-1]
+    Z_2 = K.copy()[1:]  # np.zeros_like(emp_cov)[1:]
     # Z_0 = K.copy()
     # Z_1 = K.copy()[:-1]
     # Z_2 = K.copy()[1:]
@@ -123,8 +138,11 @@ def time_graph_lasso(
     if n_samples is None:
         n_samples = np.ones(emp_cov.shape[0])
 
-    checks = [convergence(obj=objective(
-        n_samples, emp_cov, Z_0, K, Z_1, Z_2, alpha, beta, psi))]
+    checks = [
+        convergence(
+            obj=objective(
+                n_samples, emp_cov, Z_0, K, Z_1, Z_2, alpha, beta, psi))
+    ]
     for iteration_ in range(max_iter):
         # update K
         A = Z_0 - U_0
@@ -136,11 +154,14 @@ def time_graph_lasso(
         A += A.transpose(0, 2, 1)
         A /= 2.
 
-        A *= - rho * divisor[:, None, None] / n_samples[:, None, None]
+        A *= -rho * divisor[:, None, None] / n_samples[:, None, None]
         A += emp_cov
 
-        K = np.array([prox_logdet(a, lamda=ni / (rho * div))
-                      for a, div, ni in zip(A, divisor, n_samples)])
+        K = np.array(
+            [
+                prox_logdet(a, lamda=ni / (rho * div))
+                for a, div, ni in zip(A, divisor, n_samples)
+            ])
 
         # update Z_0
         A = K + U_0
@@ -156,9 +177,9 @@ def time_graph_lasso(
             Z_1 = .5 * (A_1 + A_2 - prox_e)
             Z_2 = .5 * (A_1 + A_2 + prox_e)
         else:
-            Z_1, Z_2 = prox_psi(np.concatenate((A_1, A_2), axis=1),
-                                lamda=.5 * beta / rho,
-                                rho=rho, tol=tol, rtol=rtol, max_iter=max_iter)
+            Z_1, Z_2 = prox_psi(
+                np.concatenate((A_1, A_2), axis=1), lamda=.5 * beta / rho,
+                rho=rho, tol=tol, rtol=rtol, max_iter=max_iter)
 
         # update residuals
         U_0 += K - Z_0
@@ -166,25 +187,30 @@ def time_graph_lasso(
         U_2 += K[1:] - Z_2
 
         # diagnostics, reporting, termination checks
-        rnorm = np.sqrt(squared_norm(K - Z_0) +
-                        squared_norm(K[:-1] - Z_1) +
-                        squared_norm(K[1:] - Z_2))
+        rnorm = np.sqrt(
+            squared_norm(K - Z_0) + squared_norm(K[:-1] - Z_1) +
+            squared_norm(K[1:] - Z_2))
 
-        snorm = rho * np.sqrt(squared_norm(Z_0 - Z_0_old) +
-                              squared_norm(Z_1 - Z_1_old) +
-                              squared_norm(Z_2 - Z_2_old))
+        snorm = rho * np.sqrt(
+            squared_norm(Z_0 - Z_0_old) + squared_norm(Z_1 - Z_1_old) +
+            squared_norm(Z_2 - Z_2_old))
 
         obj = objective(
             n_samples, emp_cov, Z_0, K, Z_1, Z_2, alpha, beta, psi) \
             if compute_objective else np.nan
 
         check = convergence(
-            obj=obj, rnorm=rnorm, snorm=snorm,
+            obj=obj,
+            rnorm=rnorm,
+            snorm=snorm,
             e_pri=np.sqrt(K.size + 2 * Z_1.size) * tol + rtol * max(
-                np.sqrt(squared_norm(Z_0) + squared_norm(Z_1) + squared_norm(Z_2)),
-                np.sqrt(squared_norm(K) + squared_norm(K[:-1]) + squared_norm(K[1:]))),
-            e_dual=np.sqrt(K.size + 2 * Z_1.size) * tol + rtol * rho * np.sqrt(
-                squared_norm(U_0) + squared_norm(U_1) + squared_norm(U_2)),
+                np.sqrt(
+                    squared_norm(Z_0) + squared_norm(Z_1) + squared_norm(Z_2)),
+                np.sqrt(
+                    squared_norm(K) + squared_norm(K[:-1]) +
+                    squared_norm(K[1:]))),
+            e_dual=np.sqrt(K.size + 2 * Z_1.size) * tol + rtol * rho *
+            np.sqrt(squared_norm(U_0) + squared_norm(U_1) + squared_norm(U_2)),
             # precision=Z_0.copy()
         )
         Z_0_old = Z_0.copy()
@@ -192,8 +218,9 @@ def time_graph_lasso(
         Z_2_old = Z_2.copy()
 
         if verbose:
-            print("obj: %.4f, rnorm: %.4f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
+            print(
+                "obj: %.4f, rnorm: %.4f, snorm: %.4f,"
+                "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
 
         checks.append(check)
         if stop_at is not None:
@@ -203,8 +230,9 @@ def time_graph_lasso(
         if check.rnorm <= check.e_pri and check.snorm <= check.e_dual:
             break
 
-        rho_new = update_rho(rho, rnorm, snorm, iteration=iteration_,
-                             **(update_rho_options or {}))
+        rho_new = update_rho(
+            rho, rnorm, snorm, iteration=iteration_,
+            **(update_rho_options or {}))
         # scaled dual variables should be also rescaled
         U_0 *= rho / rho_new
         U_1 *= rho / rho_new
@@ -222,7 +250,7 @@ def time_graph_lasso(
     return return_list
 
 
-class TimeGraphLasso(GraphLasso):
+class TimeGraphicalLasso(GraphicalLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -285,7 +313,7 @@ class TimeGraphLasso(GraphLasso):
         Estimated covariance matrix
 
     precision_ : array-like, shape (n_times, n_features, n_features)
-        Estimated pseudo inverse matrix.
+        Estimated precision matrix.
 
     n_iter_ : int
         Number of iterations run.
@@ -299,7 +327,7 @@ class TimeGraphLasso(GraphLasso):
             return_history=False, update_rho_options=None,
             compute_objective=True, stop_at=None, stop_when=1e-4,
             suppress_warn_list=False):
-        super(TimeGraphLasso, self).__init__(
+        super(TimeGraphicalLasso, self).__init__(
             alpha=alpha, rho=rho, tol=tol, rtol=rtol, max_iter=max_iter,
             verbose=verbose, assume_centered=assume_centered, mode=mode,
             update_rho_options=update_rho_options,
@@ -324,7 +352,7 @@ class TimeGraphLasso(GraphLasso):
         return self.get_precision()
 
     def _fit(self, emp_cov, n_samples):
-        """Fit the TimeGraphLasso model to X.
+        """Fit the TimeGraphicalLasso model to X.
 
         Parameters
         ----------
@@ -332,15 +360,14 @@ class TimeGraphLasso(GraphLasso):
             Empirical covariance of data.
 
         """
-        out = time_graph_lasso(
-                emp_cov, alpha=self.alpha, rho=self.rho,
-                beta=self.beta, mode=self.mode, n_samples=n_samples,
-                tol=self.tol, rtol=self.rtol, psi=self.psi,
-                max_iter=self.max_iter, verbose=self.verbose,
-                return_n_iter=True, return_history=self.return_history,
-                update_rho_options=self.update_rho_options,
-                compute_objective=self.compute_objective,
-                stop_at=self.stop_at, stop_when=self.stop_when)
+        out = time_graphical_lasso(
+            emp_cov, alpha=self.alpha, rho=self.rho, beta=self.beta,
+            mode=self.mode, n_samples=n_samples, tol=self.tol, rtol=self.rtol,
+            psi=self.psi, max_iter=self.max_iter, verbose=self.verbose,
+            return_n_iter=True, return_history=self.return_history,
+            update_rho_options=self.update_rho_options,
+            compute_objective=self.compute_objective, stop_at=self.stop_at,
+            stop_when=self.stop_when)
         if self.return_history:
             self.precision_, self.covariance_, self.history_, self.n_iter_ = out
         else:
@@ -348,7 +375,7 @@ class TimeGraphLasso(GraphLasso):
         return self
 
     def fit(self, X, y=None):
-        """Fit the TimeGraphLasso model to X.
+        """Fit the TimeGraphicalLasso model to X.
 
         Parameters
         ----------
@@ -360,36 +387,21 @@ class TimeGraphLasso(GraphLasso):
         y : (ignored)
 
         """
-        if sp.issparse(X):
-            raise TypeError("sparse matrices not supported.")
-
-        X = check_array_dimensions(
-            X, n_dimensions=3, time_on_axis=self.time_on_axis,
-            suppress_warn_list=self.suppress_warn_list)
-
         is_list = isinstance(X, list)
-        # Covariance does not make sense for a single feature
-        X = np.array([check_array(x, ensure_min_features=2,
-                      ensure_min_samples=2, estimator=self) for x in X])
-        if is_list:
-            n_times = len(X)
-            if np.unique([x.shape[1] for x in X]).size != 1:
-                raise ValueError("Input data cannot have different number "
-                                 "of variables.")
-            n_dimensions = X[0].shape[1]
-        else:
-            n_times = X.shape[0]
-            n_dimensions = X.shape[2]
-
-        n_samples = np.array([x.shape[0] for x in X])
+        X, n_samples, n_dimensions, n_times = check_input(
+            X, y=None, time_on_axis=self.time_on_axis,
+            suppress_warn_list=self.suppress_warn_list, estimator=self)
 
         if self.assume_centered:
             self.location_ = np.zeros((n_times, 1, n_dimensions))
         else:
             mean = np.array([x.mean(0) for x in X]) if is_list else X.mean(1)
             self.location_ = mean.reshape(n_times, 1, n_dimensions)
-        emp_cov = np.array([empirical_covariance(
-            x, assume_centered=self.assume_centered) for x in X])
+        emp_cov = np.array(
+            [
+                empirical_covariance(x, assume_centered=self.assume_centered)
+                for x in X
+            ])
 
         return self._fit(emp_cov, n_samples)
 
@@ -409,38 +421,34 @@ class TimeGraphLasso(GraphLasso):
 
         Returns
         -------
-        res : float
+        logp : float
             The likelihood of the data set with `self.covariance_` as an
             estimator of its covariance matrix.
 
         """
-        if sp.issparse(X_test):
-            raise TypeError("sparse matrices not supported.")
-
-        X_test = check_array_dimensions(
-            X_test, n_dimensions=3, time_on_axis=self.time_on_axis)
-
-        # Covariance does not make sense for a single feature
-        X_test = np.array([
-            check_array(x, ensure_min_features=2,
-                        ensure_min_samples=2, estimator=self) for x in X_test])
+        X_test, n_samples, _, _ = check_input(
+            X_test, y=None, time_on_axis=self.time_on_axis,
+            suppress_warn_list=self.suppress_warn_list, estimator=self)
 
         # compute empirical covariance of the test set
-        test_cov = np.array([empirical_covariance(
-            x, assume_centered=True) for x in X_test - self.location_])
+        test_cov = np.array(
+            [
+                empirical_covariance(x, assume_centered=True)
+                for x in X_test - self.location_
+            ])
 
-        n_samples = np.array([x.shape[0] for x in X_test])
-        res = sum(n * log_likelihood(S, K) for S, K, n in zip(
-            test_cov, self.get_observed_precision(), n_samples))
+        logp = sum(
+            n * log_likelihood(S, K) for S, K, n in zip(
+                test_cov, self.get_observed_precision(), n_samples))
 
         # ALLA  MATLAB1
         # ranks = [np.linalg.matrix_rank(L) for L in self.latent_]
         # scores_ranks = np.square(ranks-np.sqrt(L.shape[1]))
 
-        return res  # - np.sum(scores_ranks)
+        return logp  # - np.sum(scores_ranks)
 
-    def error_norm(self, comp_cov, norm='frobenius', scaling=True,
-                   squared=True):
+    def error_norm(
+            self, comp_cov, norm='frobenius', scaling=True, squared=True):
         """Compute the Mean Squared Error between two covariance estimators.
         (In the sense of the Frobenius norm).
 
@@ -470,8 +478,9 @@ class TimeGraphLasso(GraphLasso):
         `self` and `comp_cov` covariance estimators.
 
         """
-        return error_norm_time(self.covariance_, comp_cov, norm=norm,
-                               scaling=scaling, squared=squared)
+        return error_norm_time(
+            self.covariance_, comp_cov, norm=norm, scaling=scaling,
+            squared=squared)
 
     def mahalanobis(self, observations):
         """Computes the squared Mahalanobis distances of given observations.
@@ -497,8 +506,7 @@ class TimeGraphLasso(GraphLasso):
         sum_ = 0.
         for obs, loc in zip(observations, self.location_):
             centered_obs = observations - self.location_
-            sum_ += np.sum(
-                np.dot(centered_obs, precision) * centered_obs, 1)
+            sum_ += np.sum(np.dot(centered_obs, precision) * centered_obs, 1)
 
         mahalanobis_dist = sum_ / len(observations)
         return mahalanobis_dist
diff --git a/regain/datasets.py b/regain/datasets.py
index 1b0bff1..26b28b5 100644
--- a/regain/datasets.py
+++ b/regain/datasets.py
@@ -10,7 +10,8 @@ from scipy.spatial.distance import squareform
 from sklearn.datasets.base import Bunch
 from sklearn.utils import deprecated
 
-from regain.utils import is_pos_def, is_pos_semidef, positive_definite
+from regain.utils import (
+    ensure_posdef, is_pos_def, is_pos_semidef, positive_definite)
 
 
 def normalize_matrix(x):
@@ -21,11 +22,11 @@ def normalize_matrix(x):
     x *= d.T
 
 
-def make_dataset(n_samples=100, n_dim_obs=100, n_dim_lat=10, T=10,
-                 mode=None, time_on_axis='first',
-                 update_ell='l2', update_theta='l2',
-                 normalize_starting_matrices=False, degree=2, epsilon=1e-2,
-                 keep_sparsity=False, proportional=False, **kwargs):
+def make_dataset(
+        n_samples=100, n_dim_obs=100, n_dim_lat=10, T=10, mode=None,
+        time_on_axis='first', update_ell='l2', update_theta='l2',
+        normalize_starting_matrices=False, degree=2, epsilon=1e-2,
+        keep_sparsity=False, proportional=False, **kwargs):
     """Generate a synthetic dataset.
 
     Parameters
@@ -69,22 +70,29 @@ def make_dataset(n_samples=100, n_dim_obs=100, n_dim_lat=10, T=10,
         fixedl1=make_l1,
         yuan=generate_dataset_yuan,
         l1l2=generate_dataset_l1l2,
-        norm=make_l2l2_norm, l1l1=generate_dataset_l1l1,
+        norm=make_l2l2_norm,
+        l1l1=generate_dataset_l1l1,
 
         # the previous are deprecated
         sin=make_sin,
-        fixed_sparsity=make_fixed_sparsity, sincos=make_sin_cos,
-        fede=make_fede, sklearn=make_sparse_low_rank,
-        ma=make_ma_xue_zou, mak=make_ma_xue_zou_rand_k)
+        fixed_sparsity=make_fixed_sparsity,
+        sincos=make_sin_cos,
+        gp=make_exp_sine_squared,
+        fede=make_fede,
+        sklearn=make_sparse_low_rank,
+        ma=make_ma_xue_zou,
+        mak=make_ma_xue_zou_rand_k)
 
     if mode is not None:
         # mode overrides other parameters, for back compatibility
         func = mode if callable(mode) else modes.get(mode, None)
         if func is None:
-            raise ValueError("Unknown mode %s. "
-                             "Choices are: %s" % (mode, modes.keys()))
-        kwargs.update(degree=degree, epsilon=epsilon,
-                      keep_sparsity=keep_sparsity, proportional=proportional)
+            raise ValueError(
+                "Unknown mode %s. "
+                "Choices are: %s" % (mode, modes.keys()))
+        kwargs.update(
+            degree=degree, epsilon=epsilon, keep_sparsity=keep_sparsity,
+            proportional=proportional)
     else:
         func = partial(
             make_covariance, update_ell=update_ell, update_theta=update_theta,
@@ -96,27 +104,31 @@ def make_dataset(n_samples=100, n_dim_obs=100, n_dim_lat=10, T=10,
     sigmas = map(np.linalg.inv, thetas_obs)
     # map(normalize_matrix, sigmas)  # in place
 
-    data = np.array([np.random.multivariate_normal(
-        np.zeros(n_dim_obs), sigma, size=n_samples) for sigma in sigmas])
+    data = np.array(
+        [
+            np.random.multivariate_normal(
+                np.zeros(n_dim_obs), sigma, size=n_samples) for sigma in sigmas
+        ])
 
     if time_on_axis == "last":
         data = data.transpose(1, 2, 0)
-    return Bunch(data=data, thetas=np.array(thetas),
-                 thetas_observed=np.array(thetas_obs),
-                 ells=np.array(ells))
+    return Bunch(
+        data=data, thetas=np.array(thetas),
+        thetas_observed=np.array(thetas_obs), ells=np.array(ells))
 
 
 def make_ell(n_dim_obs=100, n_dim_lat=10):
     """Doc."""
     K_HO = np.zeros((n_dim_lat, n_dim_obs))
     for i in range(n_dim_lat):
-        percentage = int(n_dim_obs * 0.8)
+        percentage = int(n_dim_obs * 0.99)
         indices = np.random.randint(0, high=n_dim_obs, size=percentage)
         K_HO[i, indices] = np.random.rand(percentage) * 0.12
 
-    K_HO /= np.sum(K_HO, axis=1)[:, None] / 2
+    K_HO /= np.sum(K_HO, axis=1)[:, None] / 2.
     L = K_HO.T.dot(K_HO)
-    assert(is_pos_semidef(L))
+    print("{}%".format(np.nonzero(L)[0].size / L.size))
+    assert (is_pos_semidef(L))
     assert np.linalg.matrix_rank(L) == n_dim_lat
     # from sklearn.datasets import make_low_rank_matrix
     # L = make_low_rank_matrix(n_dim_obs, n_dim_obs, effective_rank=n_dim_lat)
@@ -132,9 +144,11 @@ def make_starting(n_dim_obs=100, n_dim_lat=10, degree=2, normalize=False):
     if normalize:
         theta = np.zeros((n_dim_obs, n_dim_obs))
         for i in range(n_dim_obs):
-            possible_idx = list(set(range(n_dim_obs)) - (
-                set(np.nonzero(theta[i, :])[0]) |
-                set(np.where(np.count_nonzero(theta, axis=1) > degree)[0])))
+            possible_idx = list(
+                set(range(n_dim_obs)) - (
+                    set(np.nonzero(theta[i, :])[0]) | set(
+                        np.where(np.count_nonzero(theta, axis=1) > degree)[0]))
+            )
             if not possible_idx:
                 continue
             n_choice = degree - (np.count_nonzero(theta[i, :]) - 1)
@@ -142,14 +156,16 @@ def make_starting(n_dim_obs=100, n_dim_lat=10, degree=2, normalize=False):
                 indices = np.random.choice(possible_idx, n_choice)
             theta[i, indices] = theta[indices, i] = 1. / degree
 
-        theta.flat[::n_dim_obs+1] = np.sum(theta, axis=1) + 0.002
+        theta.flat[::n_dim_obs + 1] = np.sum(theta, axis=1) + 0.002
 
     else:
         theta = np.eye(n_dim_obs)
         for i in range(n_dim_obs):
-            possible_idx = list(set(range(n_dim_obs)) - (
-                set(np.nonzero(theta[i, :])[0]) |
-                set(np.where(np.count_nonzero(theta, axis=1) > degree)[0])))
+            possible_idx = list(
+                set(range(n_dim_obs)) - (
+                    set(np.nonzero(theta[i, :])[0]) | set(
+                        np.where(np.count_nonzero(theta, axis=1) > degree)[0]))
+            )
             if not possible_idx:
                 continue
             n_choice = degree - (np.count_nonzero(theta[i, :]) - 1)
@@ -157,9 +173,9 @@ def make_starting(n_dim_obs=100, n_dim_lat=10, degree=2, normalize=False):
                 indices = np.random.choice(possible_idx, n_choice)
                 theta[i, indices] = theta[indices, i] = .5 / degree
 
-    assert(is_pos_def(theta))
+    assert (is_pos_def(theta))
     theta_observed = theta - L
-    assert(is_pos_def(theta_observed))
+    assert (is_pos_def(theta_observed))
     return theta, theta_observed, L, K_HO
 
 
@@ -190,9 +206,11 @@ def update_theta_l1(theta_init, no, n_dim_obs):
         rows = np.random.randint(0, n_dim_obs, no)
         cols = np.random.randint(0, n_dim_obs, no)
     for r, c in zip(rows, cols):
-        theta[r, c] = np.random.choice([0.12, 0, 0]) if theta[r, c] == 0 else .06  # np.random.rand(1) * .35
+        theta[r, c] = np.random.choice([
+            0.12, 0, 0
+        ]) if theta[r, c] == 0 else .06  # np.random.rand(1) * .35
         theta[c, r] = theta[r, c]
-    assert(is_pos_def(theta))
+    assert (is_pos_def(theta))
     return theta
 
 
@@ -208,8 +226,8 @@ def update_ell_l2(K_HO_old, epsilon, n_dim_obs):
 
 
 def make_covariance(
-    n_dim_obs=100, n_dim_lat=10, T=10, update_ell='l1', update_theta='l1',
-    normalize_starting_matrices=True, degree=2, epsilon=1e-2,
+        n_dim_obs=100, n_dim_lat=10, T=10, update_ell='l1', update_theta='l1',
+        normalize_starting_matrices=True, degree=2, epsilon=1e-2,
         keep_sparsity=False, proportional=False):
     no = int(np.ceil(n_dim_obs / 20)) if proportional else 1
 
@@ -225,16 +243,20 @@ def make_covariance(
     for i in range(1, T):
         if update_theta == 'l1':
             if keep_sparsity:
-                warnings.warn("keep_sparsity is specified but is not "
-                              "implemented with l1.")
+                warnings.warn(
+                    "keep_sparsity is specified but is not "
+                    "implemented with l1.")
             theta = update_theta_l1(thetas[-1], no, n_dim_obs)
         elif update_theta == 'l2':
-            theta = update_theta_l2(thetas[-1], n_dim_obs, degree, epsilon,
-                                    keep_sparsity=keep_sparsity, indices=idx)
+            theta = update_theta_l2(
+                thetas[-1], n_dim_obs, degree, epsilon,
+                keep_sparsity=keep_sparsity, indices=idx)
         else:
             raise ValueError(update_theta)
 
-        if update_ell == 'l1':
+        if update_ell == 'fixed' or n_dim_lat < 1:
+            pass  # L is fixed or empty
+        elif update_ell == 'l1':
             K_HO = K_HOs[-1].copy()
             picks = np.random.permutation(K_HO.size)[:no]
             K_HO = K_HO.ravel()
@@ -250,14 +272,12 @@ def make_covariance(
             r = np.random.randint(0, n_dim_lat, 1)
             K_HO[r, c] = 0.12 if K_HO[r, c] == 0 else 0
             L = K_HO.T.dot(K_HO)
-        elif update_ell == 'fixed':
-            pass  # L is fixed
         else:
             raise ValueError(update_ell)
 
         assert np.linalg.matrix_rank(L) == n_dim_lat
-        assert(is_pos_semidef(L))
-        assert(is_pos_def(theta - L))
+        assert (is_pos_semidef(L))
+        assert (is_pos_def(theta - L))
 
         thetas.append(theta)
         thetas_obs.append(theta - L)
@@ -268,8 +288,7 @@ def make_covariance(
 
 
 @deprecated("it will be removed in v0.2.0. Use `make_covariance` instead")
-def generate_dataset_l1l1(
-        n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
+def generate_dataset_l1l1(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
     """Generate matrices according to a l1-l1 model."""
     return make_covariance(
         n_dim_lat=n_dim_lat, n_dim_obs=n_dim_obs, T=T, update_ell='l1',
@@ -294,8 +313,7 @@ def generate_dataset_yuan(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
 
 
 @deprecated("it will be removed in v0.2.0. Use `make_covariance` instead")
-def make_l2l2(
-        n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
+def make_l2l2(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
     """Generate dataset with evolving L."""
     return make_covariance(
         n_dim_lat=n_dim_lat, n_dim_obs=n_dim_obs, T=T, update_ell='l2',
@@ -336,7 +354,7 @@ def make_fixed_sparsity(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
 
     theta = np.abs(theta)
     theta_observed = theta - L
-    nonzeros = np.nonzero(theta-np.diag(theta))
+    nonzeros = np.nonzero(theta - np.diag(theta))
 
     thetas = [theta]
     thetas_obs = [theta_observed]
@@ -352,7 +370,7 @@ def make_fixed_sparsity(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
         theta[theta < epsilon] = 0
         theta.flat[::n_dim_obs + 1] = np.sum(np.abs(theta), axis=1) \
                                     + np.sum(np.abs(L), axis=1) + .1
-        nonzeros = np.nonzero(theta-np.diag(theta))
+        nonzeros = np.nonzero(theta - np.diag(theta))
 
         theta_observed = theta - L
         assert is_pos_def(theta_observed)
@@ -362,11 +380,12 @@ def make_fixed_sparsity(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
     return thetas, thetas_obs, np.array([L] * T)
 
 
-def make_sin(n_dim_obs, n_dim_lat, T, shape='smooth', closeness=1,
-             normalize=False, **kwargs):
+def make_sin(
+        n_dim_obs, n_dim_lat, T, shape='smooth', closeness=1, normalize=False,
+        **kwargs):
     upper_idx = np.triu_indices(n_dim_obs, 1)
     n_interactions = len(upper_idx[0])
-    x = np.tile(np.linspace(0, (T-1.) / closeness, T), (n_interactions, 1))
+    x = np.tile(np.linspace(0, (T - 1.) / closeness, T), (n_interactions, 1))
     phase = np.random.rand(n_interactions, 1)
     freq = np.random.rand(n_interactions, 1) - .50
     A = (np.random.rand(n_interactions, 1) + 1) / 2.
@@ -379,12 +398,16 @@ def make_sin(n_dim_obs, n_dim_lat, T, shape='smooth', closeness=1,
     # threshold
     y = np.maximum(y, 0)
 
-    Y = np.array([squareform(y[:, j]) + np.diag(np.sum(squareform(y[:, j]), axis=1))
-                  for j in range(y.shape[1])])
+    Y = np.array(
+        [
+            squareform(y[:, j]) + np.diag(np.sum(squareform(y[:, j]), axis=1))
+            for j in range(y.shape[1])
+        ])
 
     if normalize:
         map(normalize_matrix, Y)  # in place
-    assert positive_definite(Y)
+    # assert positive_definite(Y)
+    ensure_posdef(Y)
 
     return Y, Y, np.zeros_like(Y)
 
@@ -410,25 +433,61 @@ def make_sin_cos(n_dim_obs=100, n_dim_lat=10, T=10, **kwargs):
 
     thetas = np.array([np.eye(n_dim_obs) for i in range(T)])
 
-    x = np.linspace(0, 2*np.pi, T)
+    x = np.linspace(0, 2 * np.pi, T)
     for i in range(T):
         for r in range(n_dim_obs):
             for c in range(n_dim_obs):
                 if r == c:
                     continue
                 if clip[r, c]:
-                    thetas[i, r, c] = np.sin((x[i]+phase[r, c])/T**2)
+                    thetas[i, r, c] = np.sin((x[i] + phase[r, c]) / T**2)
                 else:
-                    thetas[i, r, c] = np.sin((x[i]+phase[r, c]))
+                    thetas[i, r, c] = np.sin((x[i] + phase[r, c]))
         thetas[i][clip == 1] = np.clip(thetas[i][clip == 1], 0, 1)
         thetas[i][np.abs(thetas[i]) < eps] = 0
 
-        assert(is_pos_def(thetas[i]))
+        assert (is_pos_def(thetas[i]))
         theta_observed = thetas[i] - L
-        assert(is_pos_def(theta_observed))
+        assert (is_pos_def(theta_observed))
         thetas_obs = [theta_observed]
 
-    return thetas, thetas_obs, np.array([L]*T)
+    return thetas, thetas_obs, np.array([L] * T)
+
+
+def make_exp_sine_squared(n_dim_obs=5, n_dim_lat=0, T=1, **kwargs):
+    from regain.bayesian.gaussian_process_ import sample as samplegp
+    from scipy.spatial.distance import squareform
+    from sklearn.gaussian_process import kernels
+
+    L, K_HO = make_ell(n_dim_obs, n_dim_lat)
+
+    periodicity = kwargs.get('periodicity', np.pi)
+    length_scale = kwargs.get('length_scale', 2)
+    epsilon = kwargs.get('epsilon', 0.5)
+    sparse = kwargs.get('sparse', True)
+    temporal_kernel = kernels.ExpSineSquared(
+        periodicity=periodicity, length_scale=length_scale)(
+            np.arange(T)[:, None])
+
+    u = samplegp(temporal_kernel, p=n_dim_obs * (n_dim_obs - 1) // 2)[0]
+    K, K_obs = [], []
+    for uu in u.T:
+        theta = squareform(uu)
+
+        if sparse:
+            # sparsify
+            theta[np.abs(theta) < epsilon] = 0
+
+        theta += np.diag(np.sum(np.abs(theta), axis=1) + 0.01)
+        K.append(theta)
+
+        assert (is_pos_def(theta))
+        theta_observed = theta - L
+        assert (is_pos_def(theta_observed))
+        K_obs.append(theta_observed)
+
+    thetas = np.array(K)
+    return thetas, np.array(K_obs), np.array([L] * T)
 
 
 def make_fede(
@@ -483,8 +542,8 @@ def make_sparse_low_rank(
 
     for i in range(1, T):
         K = K + make_sparse_spd_matrix(n_dim_obs)
-        L = L + make_low_rank_matrix(n_dim_obs, n_dim_obs,
-                                     effective_rank=n_dim_lat)
+        L = L + make_low_rank_matrix(
+            n_dim_obs, n_dim_obs, effective_rank=n_dim_lat)
 
         # assert is_pos_def(K - L)
         # assert is_pos_semidef(L)
@@ -503,8 +562,8 @@ def make_ma_xue_zou(n_dim_obs=12, n_latent=3, T=1, epsilon=1e-3, sparsity=0.1):
     po = n_dim_obs
     ph = p - n_dim_obs
     W = np.zeros((p, p))
-    non_zeros = int(round(p*p*sparsity))
-    picks = np.random.permutation(p*p)[:non_zeros]
+    non_zeros = int(round(p * p * sparsity))
+    picks = np.random.permutation(p * p)[:non_zeros]
     W = W.ravel(order='F')
     W[picks] = np.random.randn(non_zeros)
     W = np.reshape(W, (p, p), order="F")
@@ -541,7 +600,6 @@ def make_ma_xue_zou_rand_k(
     # p = n_dim_obs + n_latent  # int(n_dim_obs * 0.05)
     p = n_dim_obs + int(n_dim_obs * 0.05)
     po = n_dim_obs
-    ph = p - n_dim_obs
     nnzr = int(sparsity * (np.triu_indices(p, 1)[0].size))
 
     # Generate A, the original inverse covariance, with random sparsity pattern...
diff --git a/regain/forward_backward/time_graph_lasso_.py b/regain/forward_backward/time_graphical_lasso_.py
similarity index 80%
rename from regain/forward_backward/time_graph_lasso_.py
rename to regain/forward_backward/time_graphical_lasso_.py
index 48609cd..3fada0f 100644
--- a/regain/forward_backward/time_graph_lasso_.py
+++ b/regain/forward_backward/time_graphical_lasso_.py
@@ -9,8 +9,9 @@ from scipy import linalg
 from six.moves import map, range, zip
 from sklearn.covariance import empirical_covariance
 
-from regain.covariance.graph_lasso_ import logl
-from regain.covariance.time_graph_lasso_ import TimeGraphLasso
+from regain.covariance.graphical_lasso_ import logl
+from regain.covariance.time_graphical_lasso_ import TimeGraphicalLasso
+from regain.covariance.time_graphical_lasso_ import loss as loss_tgl
 from regain.norm import l1_od_norm, vector_p_norm
 from regain.prox import prox_FL
 from regain.utils import convergence, positive_definite
@@ -18,11 +19,7 @@ from regain.utils import convergence, positive_definite
 
 def loss(S, K, n_samples=None, vareps=0):
     """Loss function for time-varying graphical lasso."""
-    if n_samples is None:
-        n_samples = np.ones(S.shape[0])
-    loss_ = sum(-ni * logl(emp_cov, precision)
-                for emp_cov, precision, ni in zip(S, K, n_samples))
-    # loss_ += vareps / 2. * squared_norm(K)
+    loss_ = loss_tgl(S, K, n_samples=n_samples)
     loss_ += vareps / 2. * _scalar_product(K, K)
     return loss_
 
@@ -43,10 +40,16 @@ def grad_loss(x, emp_cov, n_samples, x_inv=None, vareps=0):
 def penalty(precision, alpha, beta, psi):
     """Penalty for time-varying graphical lasso."""
     if isinstance(alpha, np.ndarray):
-        obj = sum(a[0][0] * m for a, m in zip(alpha, map(l1_od_norm, precision)))
+        obj = sum(
+            a[0][0] * m for a, m in zip(alpha, map(l1_od_norm, precision)))
     else:
         obj = alpha * sum(map(l1_od_norm, precision))
-    obj += beta * psi(precision[1:] - precision[:-1])
+    if isinstance(beta, np.ndarray):
+        obj += sum(
+            b[0][0] * m
+            for b, m in zip(beta, map(psi, precision[1:] - precision[:-1])))
+    else:
+        obj += beta * psi(precision[1:] - precision[:-1])
     return obj
 
 
@@ -71,9 +74,9 @@ def _scalar_product(x, y):
     # return x.ravel().dot(y.ravel())
 
 
-def choose_gamma(gamma, x, emp_cov, n_samples, beta, alpha, lamda, grad,
-                 delta=1e-4, eps=0.5, max_iter=1000, p=1, x_inv=None,
-                 vareps=1e-5, choose='gamma'):
+def choose_gamma(
+        gamma, x, emp_cov, n_samples, beta, alpha, lamda, grad, delta=1e-4,
+        eps=0.5, max_iter=1000, p=1, x_inv=None, vareps=1e-5, choose='gamma'):
     """Choose gamma for backtracking.
 
     References
@@ -105,10 +108,10 @@ def choose_gamma(gamma, x, emp_cov, n_samples, beta, alpha, lamda, grad,
     return gamma, prox
 
 
-def choose_lamda(lamda, x, emp_cov, n_samples, beta, alpha, gamma, delta=1e-4,
-                 eps=0.5, max_iter=1000, criterion='b', p=1, x_inv=None,
-                 grad=None, prox=None, min_eigen_x=None,
-                 vareps=1e-5):
+def choose_lamda(
+        lamda, x, emp_cov, n_samples, beta, alpha, gamma, delta=1e-4, eps=0.5,
+        max_iter=1000, criterion='b', p=1, x_inv=None, grad=None, prox=None,
+        min_eigen_x=None, vareps=1e-5):
     """Choose lambda for backtracking.
 
     References
@@ -162,7 +165,8 @@ def choose_lamda(lamda, x, emp_cov, n_samples, beta, alpha, gamma, delta=1e-4,
                 n_samples, emp_cov, x1, alpha, beta, psi, vareps=vareps) \
                     - objective_x
             # if positive_definite(x1) and obj_diff <= lamda * tolerance:
-            cond = True # lamda > 0 if min_eigen_y >= 0 else lamda < min_eigen_x / (min_eigen_x - min_eigen_y)
+            cond = True
+            # COND IS lamda > 0 if min_eigen_y >= 0 else lamda < min_eigen_x / (min_eigen_x - min_eigen_y)
             if cond and obj_diff <= lamda * tolerance:
                 break
         else:
@@ -172,8 +176,8 @@ def choose_lamda(lamda, x, emp_cov, n_samples, beta, alpha, gamma, delta=1e-4,
 
 
 def fista_step(Y, Y_diff, t):
-    t_next = (1. + np.sqrt(1.0 + 4.0 * t*t)) / 2.
-    return Y + ((t - 1.0)/t_next) * Y_diff, t_next
+    t_next = (1. + np.sqrt(1.0 + 4.0 * t * t)) / 2.
+    return Y + ((t - 1.0) / t_next) * Y_diff, t_next
 
 
 def upper_diag_3d(x):
@@ -229,8 +233,8 @@ def time_graph_lasso(
     """
     available_choose = ('gamma', 'lamda', 'fixed', 'both')
     if choose not in available_choose:
-        raise ValueError("`choose` parameter must be one of %s." %
-                         available_choose)
+        raise ValueError(
+            "`choose` parameter must be one of %s." % available_choose)
 
     n_times, _, n_features = emp_cov.shape
     covariance_ = emp_cov.copy()
@@ -246,9 +250,8 @@ def time_graph_lasso(
     # assert positive_definite(K)
 
     obj_partial = partial(
-        objective, n_samples=n_samples, emp_cov=emp_cov,
-        alpha=alpha, beta=beta, psi=partial(vector_p_norm, p=time_norm),
-        vareps=vareps)
+        objective, n_samples=n_samples, emp_cov=emp_cov, alpha=alpha,
+        beta=beta, psi=partial(vector_p_norm, p=time_norm), vareps=vareps)
 
     max_residual = -np.inf
     n_linesearch = 0
@@ -274,23 +277,34 @@ def time_graph_lasso(
         if choose in ['gamma', 'both']:
             gamma, y = choose_gamma(
                 gamma / eps if iteration_ > 0 else gamma, K, emp_cov,
-                n_samples=n_samples,
-                beta=beta, alpha=alpha, lamda=lamda, grad=grad,
-                delta=delta, eps=eps, max_iter=200, p=time_norm, x_inv=x_inv,
-                vareps=vareps, choose=choose)
+                n_samples=n_samples, beta=beta, alpha=alpha, lamda=lamda,
+                grad=grad, delta=delta, eps=eps, max_iter=200, p=time_norm,
+                x_inv=x_inv, vareps=vareps, choose=choose)
         # print(gamma)
 
         x_hat = K - gamma * grad
         if choose not in ['gamma', 'both']:
-            y = prox_FL(x_hat, beta * gamma, alpha * gamma, p=time_norm, symmetric=True)
+            y = prox_FL(
+                x_hat, beta * gamma, alpha * gamma, p=time_norm,
+                symmetric=True)
 
         if choose in ['lamda', 'both']:
             lamda, n_ls = choose_lamda(
                 min(lamda / eps if iteration_ > 0 else lamda, 1),
-                K, emp_cov, n_samples=n_samples,
-                beta=beta, alpha=alpha, gamma=gamma, delta=delta, eps=eps,
-                criterion=lamda_criterion, max_iter=200, p=time_norm,
-                x_inv=x_inv, grad=grad, prox=y,
+                K,
+                emp_cov,
+                n_samples=n_samples,
+                beta=beta,
+                alpha=alpha,
+                gamma=gamma,
+                delta=delta,
+                eps=eps,
+                criterion=lamda_criterion,
+                max_iter=200,
+                p=time_norm,
+                x_inv=x_inv,
+                grad=grad,
+                prox=y,
                 # min_eigen_x=np.min(eigens),
                 vareps=vareps)
             n_linesearch += n_ls
@@ -309,11 +323,12 @@ def time_graph_lasso(
                 np.linalg.norm(upper_diag_3d(k_previous))),
             e_dual=tol,
             # precision=K.copy()
-            )
+        )
 
         if verbose and iteration_ % (50 if verbose < 2 else 1) == 0:
-            print("obj: %.4f, rnorm: %.7f, snorm: %.4f,"
-                  "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
+            print(
+                "obj: %.4f, rnorm: %.7f, snorm: %.4f,"
+                "eps_pri: %.4f, eps_dual: %.4f" % check[:5])
 
         if return_history:
             checks.append(check)
@@ -333,19 +348,20 @@ def time_graph_lasso(
 
                 if iteration_ == 0:
                     normalizer = res_norm + 1e-6
-                max_residual = max(np.linalg.norm(grad),
-                                   np.linalg.norm(subgrad)) + 1e-6
+                max_residual = max(
+                    np.linalg.norm(grad), np.linalg.norm(subgrad)) + 1e-6
             else:
                 res_norm = np.linalg.norm(K - k_previous) / gamma
                 max_residual = max(max_residual, res_norm)
-                normalizer = max(np.linalg.norm(grad),
-                                 np.linalg.norm(subgrad)) + 1e-6
+                normalizer = max(
+                    np.linalg.norm(grad), np.linalg.norm(subgrad)) + 1e-6
 
             r_rel = res_norm / max_residual
             r_norm = res_norm / normalizer
 
-            if not debug and (r_rel <= tol or r_norm <= tol) and iteration_ > 0: # or (
-                    # check.rnorm <= check.e_pri and iteration_ > 0):
+            if not debug and (r_rel <= tol
+                              or r_norm <= tol) and iteration_ > 0:  # or (
+                # check.rnorm <= check.e_pri and iteration_ > 0):
                 break
             # pass
         # if check.rnorm <= check.e_pri and iteration_ > 0:
@@ -367,7 +383,7 @@ def time_graph_lasso(
     return return_list
 
 
-class TimeGraphLassoForwardBackward(TimeGraphLasso):
+class TimeGraphicalLassoForwardBackward(TimeGraphicalLasso):
     """Sparse inverse covariance estimation with an l1-penalized estimator.
 
     Parameters
@@ -437,16 +453,16 @@ class TimeGraphLassoForwardBackward(TimeGraphLasso):
 
     """
 
-    def __init__(self, alpha=0.01, beta=1., time_on_axis='first', tol=1e-4,
-                 max_iter=100, verbose=False, assume_centered=False,
-                 compute_objective=True, eps=0.5, choose='gamma', lamda=1,
-                 delta=1e-4, gamma=1., lamda_criterion='b', time_norm=1,
-                 return_history=False, debug=False,
-                 return_n_linesearch=False,
-                 vareps=1e-5, stop_at=None, stop_when=1e-4):
-        super(TimeGraphLassoForwardBackward, self).__init__(
-            alpha=alpha, tol=tol, max_iter=max_iter,
-            verbose=verbose, assume_centered=assume_centered,
+    def __init__(
+            self, alpha=0.01, beta=1., time_on_axis='first', tol=1e-4,
+            max_iter=100, verbose=False, assume_centered=False,
+            compute_objective=True, eps=0.5, choose='gamma', lamda=1,
+            delta=1e-4, gamma=1., lamda_criterion='b', time_norm=1,
+            return_history=False, debug=False, return_n_linesearch=False,
+            vareps=1e-5, stop_at=None, stop_when=1e-4):
+        super(TimeGraphicalLassoForwardBackward, self).__init__(
+            alpha=alpha, tol=tol, max_iter=max_iter, verbose=verbose,
+            assume_centered=assume_centered,
             compute_objective=compute_objective, beta=beta,
             time_on_axis=time_on_axis)
         self.delta = delta
@@ -480,12 +496,11 @@ class TimeGraphLassoForwardBackward(TimeGraphLasso):
             emp_cov, n_samples=n_samples, alpha=self.alpha, beta=self.beta,
             tol=self.tol, max_iter=self.max_iter, verbose=self.verbose,
             return_n_iter=True, return_history=self.return_history,
-            compute_objective=self.compute_objective,
-            time_norm=self.time_norm, lamda_criterion=self.lamda_criterion,
-            gamma=self.gamma, delta=self.delta, eps=self.eps,
-            choose=self.choose, lamda=self.lamda, debug=self.debug,
-            return_n_linesearch=self.return_n_linesearch,
-            vareps=self.vareps,
+            compute_objective=self.compute_objective, time_norm=self.time_norm,
+            lamda_criterion=self.lamda_criterion, gamma=self.gamma,
+            delta=self.delta, eps=self.eps, choose=self.choose,
+            lamda=self.lamda, debug=self.debug,
+            return_n_linesearch=self.return_n_linesearch, vareps=self.vareps,
             stop_at=self.stop_at, stop_when=self.stop_when)
 
         if self.return_history:
@@ -506,8 +521,11 @@ class TimeGraphLassoForwardBackward(TimeGraphLasso):
         if is_covariance:
             emp_cov = X
         else:
-            emp_cov = np.array([empirical_covariance(
-                x, assume_centered=self.assume_centered) for x in X])
+            emp_cov = np.array(
+                [
+                    empirical_covariance(
+                        x, assume_centered=self.assume_centered) for x in X
+                ])
         return max(alpha_max(e) for e in emp_cov)
 
 
diff --git a/regain/forward_backward/time_graph_lasso_laplacian.py b/regain/forward_backward/time_graphical_lasso_laplacian.py
similarity index 100%
rename from regain/forward_backward/time_graph_lasso_laplacian.py
rename to regain/forward_backward/time_graphical_lasso_laplacian.py
diff --git a/regain/covariance/__old/__init__.py b/regain/linear_model/__old/__init__.py
similarity index 100%
rename from regain/covariance/__old/__init__.py
rename to regain/linear_model/__old/__init__.py
diff --git a/regain/covariance/__old/group_lasso_overlap_old_.py b/regain/linear_model/__old/group_lasso_overlap_old_.py
similarity index 100%
rename from regain/covariance/__old/group_lasso_overlap_old_.py
rename to regain/linear_model/__old/group_lasso_overlap_old_.py
diff --git a/regain/model_selection/bayesian_optimization.py b/regain/model_selection/bayesian_optimization.py
index 079f560..934c6f9 100644
--- a/regain/model_selection/bayesian_optimization.py
+++ b/regain/model_selection/bayesian_optimization.py
@@ -1,4 +1,9 @@
-"""docstring."""
+"""Deprecated.
+
+This module to port GPyOpt in a `sklearn-usable` way has been deprecated
+in favor of scikit-optimize, which natively offer a Bayesian replacement
+for the GridSearchCV.
+"""
 import numpy as np
 from functools import partial
 from sklearn.base import is_classifier, clone
@@ -6,6 +11,7 @@ from sklearn.metrics.scorer import _check_multimetric_scoring
 from sklearn.model_selection import cross_val_score
 from sklearn.model_selection._split import check_cv
 from sklearn.model_selection._search import BaseSearchCV
+from sklearn.utils import deprecated
 
 import GPyOpt
 from GPyOpt.core.task.objective import SingleObjective
@@ -15,6 +21,7 @@ from GPyOpt.util.arguments_manager import ArgumentsManager
 from GPyOpt.optimization.acquisition_optimizer import AcquisitionOptimizer
 
 
+@deprecated()
 class BayesianOptimization(GPyOpt.methods.BayesianOptimization, BaseSearchCV):
     """Wrapper for BayesianOptimization object from GPyOpt package.
 
@@ -22,16 +29,16 @@ class BayesianOptimization(GPyOpt.methods.BayesianOptimization, BaseSearchCV):
     For more information, see GPyOpt package.
     """
 
-    def __init__(self, estimator, domain=None, constraints=None,
-                 cost_withGradients=None, model_type='GP', X=None, Y=None,
-                 initial_design_numdata=5, initial_design_type='random',
-                 acquisition_type='EI', normalize_Y=True, exact_feval=False,
-                 acquisition_optimizer_type='lbfgs', model_update_interval=1,
-                 evaluator_type='sequential', batch_size=1, num_cores=1,
-                 verbosity_model=False, verbosity=False, de_duplication=False,
-                 max_iter=50, refit=True, cv=None, scoring=None, n_jobs=1,
-                 verbose=False,
-                 **kwargs):
+    def __init__(
+            self, estimator, domain=None, constraints=None,
+            cost_withGradients=None, model_type='GP', X=None, Y=None,
+            initial_design_numdata=5, initial_design_type='random',
+            acquisition_type='EI', normalize_Y=True, exact_feval=False,
+            acquisition_optimizer_type='lbfgs', model_update_interval=1,
+            evaluator_type='sequential', batch_size=1, num_cores=1,
+            verbosity_model=False, verbosity=False, de_duplication=False,
+            max_iter=50, refit=True, cv=None, scoring=None, n_jobs=1,
+            verbose=False, **kwargs):
         """Initialise the estimator."""
         # super(BayesianOptimization, self).__init__(
         #     f=f, domain=domain, constraints=constraints,
@@ -84,8 +91,8 @@ class BayesianOptimization(GPyOpt.methods.BayesianOptimization, BaseSearchCV):
         self.exact_feval = exact_feval
         self.normalize_Y = normalize_Y
 
-        if 'model' in self.kwargs and isinstance(
-                kwargs['model'], GPyOpt.models.base.BOModel):
+        if 'model' in self.kwargs and isinstance(kwargs['model'],
+                                                 GPyOpt.models.base.BOModel):
             self.model = kwargs['model']
             self.model_type = 'User defined model used.'
             if self.verbose:
@@ -171,16 +178,16 @@ class BayesianOptimization(GPyOpt.methods.BayesianOptimization, BaseSearchCV):
         self.run_optimization(max_iter=self.max_iter, verbosity=self.verbosity)
 
         self.best_index_ = self.Y.argmin()
-        self.best_params_ = dict(zip(self.param_names,
-                                     10 ** self.X[self.best_index_]))
+        self.best_params_ = dict(
+            zip(self.param_names, 10**self.X[self.best_index_]))
         self.best_score_ = self.Y[self.Y.argmin()]
 
         # Store the only scorer not as a dict for single metric evaluation
         self.scorer_ = scorers if self.multimetric_ else scorers['score']
 
         if self.refit:
-            self.best_estimator_ = clone(self.estimator).set_params(
-                **self.best_params_)
+            self.best_estimator_ = clone(
+                self.estimator).set_params(**self.best_params_)
             if y is not None:
                 self.best_estimator_.fit(X, y, **fit_params)
             else:
@@ -199,7 +206,10 @@ class BayesianOptimization(GPyOpt.methods.BayesianOptimization, BaseSearchCV):
 
 
 def _fit_score(x, score_function, mdl=None, param_names=None):
-    x = 10 ** np.atleast_2d(x)
-    return - np.array([np.mean(score_function(mdl.set_params(
-        **dict(zip(param_names, pars)))))
-                       for pars in x])[:, None]
+    x = 10**np.atleast_2d(x)
+    return -np.array(
+        [
+            np.mean(
+                score_function(mdl.set_params(**dict(zip(param_names, pars)))))
+            for pars in x
+        ])[:, None]
diff --git a/regain/plotting/graph.py b/regain/plotting/graph.py
index b792f35..1fa5c9f 100644
--- a/regain/plotting/graph.py
+++ b/regain/plotting/graph.py
@@ -7,7 +7,6 @@ def plot_graph_with_latent_variables(
         adjacency_matrix, n_observed, title="", n_latents=0, labels=None,
         ax=None, node_size=100, font_size=5):
     """Plot graph with latent variables."""
-    import networkx as nx
     # plt.figure(figsize=(15,10))
     if ax is None:
         f, ax = plt.subplots(1, 1)
diff --git a/regain/prox.py b/regain/prox.py
index 50015a7..7a9fe43 100644
--- a/regain/prox.py
+++ b/regain/prox.py
@@ -63,8 +63,14 @@ def blockwise_soft_thresholding(a, lamda):
     """Proximal operator for l2 norm."""
     if a.ndim > 2:
         out = np.empty_like(a, dtype=float)
+        if not isinstance(lamda, collections.Iterable):
+            lamda = np.repeat(lamda, a.shape[0])
+        else:
+            lamda = lamda.ravel()
+            assert lamda.shape[0] == a.shape[0]
+
         for t in range(a.shape[0]):
-            out[t] = _blockwise_soft_thresholding_2d(a[t], lamda)
+            out[t] = _blockwise_soft_thresholding_2d(a[t], lamda[t])
     else:
         out = _blockwise_soft_thresholding_2d(a, lamda)
     return out
@@ -75,15 +81,25 @@ def blockwise_soft_thresholding_symmetric(a, lamda):
     col_norms = np.linalg.norm(a, axis=1)
     ones_vect = np.ones(a.shape[1])
 
-    if a.ndim == 2:
-        return np.dot(
+    if a.ndim > 2:
+        out = np.empty_like(a, dtype=float)
+        if not isinstance(lamda, collections.Iterable):
+            lamda = np.repeat(lamda, a.shape[0])
+        else:
+            lamda = lamda.ravel()
+            assert lamda.shape[0] == a.shape[0]
+
+        out = np.empty_like(a, dtype=float)
+        for t, (x, c_norm) in enumerate(zip(a, col_norms)):
+            out[t] = np.dot(
+                x,
+                np.diag((ones_vect - lamda[t] / c_norm) * (c_norm > lamda[t])))
+
+    else:
+        out = np.dot(
             a, np.diag((ones_vect - lamda / col_norms) * (col_norms > lamda)))
 
-    output = np.empty_like(a, dtype=float)
-    for i, (x, c_norm) in enumerate(zip(a, col_norms)):
-        output[i] = np.dot(
-            x, np.diag((ones_vect - lamda / c_norm) * (c_norm > lamda)))
-    return output
+    return out
 
 
 # %% Perform prox operator:   min_x (1/2t)||x-w||^2 subject to |x|<=radius
diff --git a/regain/validation.py b/regain/validation.py
index 4afbdeb..09d82c8 100644
--- a/regain/validation.py
+++ b/regain/validation.py
@@ -3,7 +3,9 @@ import warnings
 from functools import partial
 
 import numpy as np
+import scipy.sparse as sp
 from sklearn.utils.extmath import squared_norm
+from sklearn.utils.validation import check_array
 
 from regain.norm import l1_norm, node_penalty
 from regain.prox import (
@@ -58,3 +60,45 @@ def check_array_dimensions(
         return X.transpose(2, 0, 1)  # put time as first dimension
 
     return X
+
+
+def check_input_3d(
+        X, time_on_axis='first', suppress_warn_list=False, estimator=None):
+    """Old API. X is a 3d matrix."""
+    X = check_array_dimensions(
+        X, n_dimensions=3, time_on_axis=time_on_axis,
+        suppress_warn_list=suppress_warn_list)
+
+    is_list = isinstance(X, list)
+    # Covariance does not make sense for a single feature
+    X = np.array(
+        [
+            check_array(
+                x, ensure_min_features=2, ensure_min_samples=2,
+                estimator=estimator) for x in X
+        ])
+    if is_list:
+        n_times = len(X)
+        if np.unique([x.shape[1] for x in X]).size != 1:
+            raise ValueError(
+                "Input data cannot have different number "
+                "of variables.")
+        n_dimensions = X[0].shape[1]
+    else:
+        n_times = X.shape[0]
+        n_dimensions = X.shape[2]
+
+    n_samples = np.array([x.shape[0] for x in X])
+    return X, n_samples, n_dimensions, n_times
+
+
+def check_input(X, y=None, **kwargs):
+    """Validate data and labels."""
+    if sp.issparse(X):
+        raise TypeError("sparse matrices not supported.")
+
+    if y is None:
+        # needs X.ndim == 3 or X is list
+        return check_input_3d(X, **kwargs)
+    else:
+        raise ValueError("y cannot be not None with this class.")


Test patch:
diff --git a/regain/tests/test_graph_lasso.py b/regain/tests/test_graphical_lasso.py
similarity index 68%
rename from regain/tests/test_graph_lasso.py
rename to regain/tests/test_graphical_lasso.py
index b28389f..a5d0370 100644
--- a/regain/tests/test_graph_lasso.py
+++ b/regain/tests/test_graphical_lasso.py
@@ -1,4 +1,4 @@
-"""Test LatentTimeGraphLasso."""
+"""Test LatentTimeGraphicalLasso."""
 import numpy as np
 from numpy.testing import assert_array_almost_equal
 
@@ -9,14 +9,14 @@ except ImportError:
     # sklean < 0.20
     from sklearn.covariance import GraphLasso as GL
 
-from regain.covariance.graph_lasso_ import GraphLasso
+from regain.covariance.graphical_lasso_ import GraphicalLasso
 
 
 def test_gl():
-    """Check GraphLasso vs sklearn-graph lasso."""
+    """Check GraphicalLasso vs sklearn-graphical lasso."""
     np.random.seed(2)
     X = np.random.multivariate_normal(np.zeros(3), np.eye(3), size=100)
     p1 = GL().fit(X).precision_
-    p2 = GraphLasso().fit(X).precision_
+    p2 = GraphicalLasso().fit(X).precision_
 
     assert_array_almost_equal(p1, p2, 1)
diff --git a/regain/tests/test_latent_graph_lasso.py b/regain/tests/test_latent_graphical_lasso.py
similarity index 62%
rename from regain/tests/test_latent_graph_lasso.py
rename to regain/tests/test_latent_graphical_lasso.py
index dd1221a..4c5e2f8 100644
--- a/regain/tests/test_latent_graph_lasso.py
+++ b/regain/tests/test_latent_graphical_lasso.py
@@ -1,19 +1,19 @@
-"""Test LatentGraphLasso."""
+"""Test LatentGraphicalLasso."""
 import numpy as np
 import warnings
 
 from numpy.testing import assert_array_equal
 
-from regain.covariance.latent_graph_lasso_ import LatentGraphLasso
+from regain.covariance.latent_graphical_lasso_ import LatentGraphicalLasso
 
 
 def test_lgl_zero():
-    """Check that LatentGraphLasso can handle zero data."""
+    """Check that LatentGraphicalLasso can handle zero data."""
     a = np.zeros((3, 3))
 
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        mdl = LatentGraphLasso(max_iter=1, assume_centered=True).fit(a)
+        mdl = LatentGraphicalLasso(max_iter=1, assume_centered=True).fit(a)
 
     mdl.precision_.flat[::4] = 0
 
diff --git a/regain/tests/test_latent_time_graph_lasso.py b/regain/tests/test_latent_time_graphical_lasso.py
similarity index 75%
rename from regain/tests/test_latent_time_graph_lasso.py
rename to regain/tests/test_latent_time_graphical_lasso.py
index 75c9d80..519b146 100644
--- a/regain/tests/test_latent_time_graph_lasso.py
+++ b/regain/tests/test_latent_time_graphical_lasso.py
@@ -1,19 +1,19 @@
-"""Test LatentTimeGraphLasso."""
+"""Test LatentTimeGraphicalLasso."""
 import numpy as np
 import warnings
 
 from numpy.testing import assert_array_equal
 
-from regain.covariance.latent_time_graph_lasso_ import LatentTimeGraphLasso
+from regain.covariance.latent_time_graphical_lasso_ import LatentTimeGraphicalLasso
 
 
 def test_ltgl_zero():
-    """Check that LatentTimeGraphLasso can handle zero data."""
+    """Check that LatentTimeGraphicalLasso can handle zero data."""
     a = np.zeros((3, 3, 3))
 
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        mdl = LatentTimeGraphLasso(
+        mdl = LatentTimeGraphicalLasso(
             time_on_axis='last', max_iter=1, assume_centered=True).fit(a)
 
     for p in mdl.precision_:
@@ -27,12 +27,12 @@ def test_ltgl_zero():
 
 
 def test_ltgl_prox_l1():
-    """Check that LatentTimeGraphLasso can handle zero data."""
+    """Check that LatentTimeGraphicalLasso can handle zero data."""
     a = np.zeros((3, 3, 3))
 
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        mdl = LatentTimeGraphLasso(
+        mdl = LatentTimeGraphicalLasso(
             time_on_axis='first', psi='l2', phi='node', max_iter=1,
             assume_centered=True).fit(a)
 
diff --git a/regain/tests/test_time_graph_lasso.py b/regain/tests/test_time_graphical_lasso.py
similarity index 70%
rename from regain/tests/test_time_graph_lasso.py
rename to regain/tests/test_time_graphical_lasso.py
index a8a4127..9420daf 100644
--- a/regain/tests/test_time_graph_lasso.py
+++ b/regain/tests/test_time_graphical_lasso.py
@@ -1,17 +1,17 @@
-"""Test LatentTimeGraphLasso."""
+"""Test LatentTimeGraphicalLasso."""
 import numpy as np
 import warnings
 from numpy.testing import assert_array_equal
 
-from regain.covariance.time_graph_lasso_ import TimeGraphLasso
+from regain.covariance.time_graphical_lasso_ import TimeGraphicalLasso
 
 
 def test_ltgl_zero():
-    """Check that LatentTimeGraphLasso can handle zero data."""
+    """Check that LatentTimeGraphicalLasso can handle zero data."""
     a = np.zeros((3, 3, 3))
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        mdl = TimeGraphLasso(
+        mdl = TimeGraphicalLasso(
             time_on_axis='last', max_iter=1, assume_centered=True).fit(a)
 
     for p in mdl.precision_: