add batching

regain/covariance/time_graphical_lasso_.py

@@ -59,9 +59,9 @@ def batch_log_likelihood(emp_cov, precision):
 
 def loss(emp_cov, precision, n_samples=None):
     """Loss function for time-varying graphical lasso."""
-    batch_dim = emp_cov.shape[0]
     if n_samples is None:
         # 1 sample for each batch, ie, do not scale.
+        batch_dim = emp_cov.shape[0]
         n_samples = np.ones(batch_dim)
 
     return np.sum(-n_samples * batch_log_likelihood(emp_cov, precision))
@@ -254,15 +254,15 @@ def time_graphical_lasso(
             rnorm=rnorm,
             snorm=snorm,
             e_pri=np.sqrt(K.size + 2 * Z_1.size) * tol
-            + rtol
-            * max(
+                  + 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)),
+                   + 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()

regain/generalized_linear_model/glm_time_ising.py

@@ -46,18 +46,21 @@
 from regain.validation import check_norm_prox
 
 
-def loss_ising(X, K, n_samples=None):
+def loss_ising(X, precision, n_samples=None):
     """Loss function for time-varying ising model."""
     if n_samples is None:
-        n_samples = np.ones(X.shape[0])
-    return sum(-ni * loss(x, k) for x, k, ni in zip(X, K, n_samples))
+        # 1 sample for each batch, ie, do not scale.
+        batch_dim = X.shape[0]
+        n_samples = np.ones(batch_dim)
+
+    return sum(-ni * loss(x, k) for x, k, ni in zip(X, precision, n_samples))
 
 
 def objective(X, K, Z_M, alpha, kernel, psi):
     """Objective function for time-varying ising model."""
 
     obj = loss_ising(X, K)
-    obj += alpha * sum(map(l1_od_norm, K))
+    obj += alpha * np.sum(l1_od_norm(K))
 
     for m in range(1, K.shape[0]):
         # all possible non markovians jumps

regain/generalized_linear_model/glm_time_poisson.py

@@ -31,7 +31,6 @@
 import warnings
 
 import numpy as np
-from six.moves import map, range, zip
 from sklearn.base import BaseEstimator
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.gaussian_process import kernels
@@ -58,7 +57,7 @@ def objective(X, K, Z_M, alpha, kernel, psi):
     """Objective function for time-varying poisson model."""
 
     obj = loss_poisson(X, K)
-    obj += alpha * sum(map(l1_od_norm, K))
+    obj += alpha * np.sum(l1_od_norm(K))
 
     for m in range(1, K.shape[0]):
         # all possible non markovians jumps