Rename lr to init_step_size everywhere

examples/empirical_frechet_mean_uncertainty_sn.py

@@ -64,7 +64,7 @@ def empirical_frechet_var_bubble(n_samples, theta, dim, n_expectation=1000):
         data = gs.concatenate([rest_col, last_col], axis=-1)
 
         estimator = FrechetMean(
-            sphere.metric, method="adaptive", max_iter=32, init_point=north_pole
+            sphere.metric, max_iter=32, method="adaptive", init_point=north_pole
         )
         estimator.fit(data)
         current_mean = estimator.estimate_

examples/geodesic_regression_grassmannian.py

@@ -48,9 +48,9 @@ def main():
         metric=METRIC,
         center_X=False,
         method="riemannian",
-        verbose=True,
         max_iter=50,
-        learning_rate=0.1,
+        init_step_size=0.1,
+        verbose=True,
     )
 
     gr.fit(data, target, compute_training_score=True)

examples/geodesic_regression_se2.py

@@ -66,9 +66,9 @@ def main():
         metric=METRIC,
         center_X=False,
         method="riemannian",
-        verbose=True,
         max_iter=100,
-        learning_rate=0.1,
+        init_step_size=0.1,
+        verbose=True,
         initialization="frechet",
     )
     gr.fit(X, y, compute_training_score=True)

geomstats/geometry/product_manifold.py

@@ -1,6 +1,6 @@
 """Product of manifolds.
 
-Lead author: Yann Cabanes.
+Lead author: Nicolas Guigui.
 """
 
 import joblib

geomstats/geometry/sub_riemannian_metric.py

@@ -69,7 +69,7 @@ def frame(self, point):
         """
         raise NotImplementedError("The frame field is not implemented.")
 
-    def cometric_sub_matrix(self, basepoint):
+    def cometric_sub_matrix(self, base_point):
         """Cometric  sub matrix of dimension dist_dim x dist_dim.
 
         Let {X_i}, i = 1, .., dist_dim, be an arbitrary frame for the distribution
@@ -120,7 +120,7 @@ def inner_coproduct(self, cotangent_vec_a, cotangent_vec_b, base_point):
         ----------
         cotangent_vec_a : array-like, shape=[..., dim]
             Cotangent vector at `base_point`.
-        cotangent_vet_b : array-like, shape=[..., dim]
+        cotangent_vec_b : array-like, shape=[..., dim]
             Cotangent vector at `base_point`.
         base_point : array-like, shape=[..., dim]
             Point on the manifold.
@@ -182,12 +182,14 @@ def vector(state):
         return vector
 
     def symp_euler(self, step_size):
-        r"""Compute a function which calculates a step of symplectic euler integration.
+        """Compute a function which calculates a step of symplectic euler integration.
+
+        The output function computes a symplectic euler step of the Hamiltonian system
+        of equations associated with the cometric and obtained by the method
+        :meth:`~sub_remannian_metric.SubRiemannianMetric.symp_grad`.
 
         Parameters
         ----------
-        hamiltonian : callable
-            The hamiltonian function from the tangent bundle to the reals.
         step_size : float
             Step size of the symplectic euler step.
 
@@ -256,7 +258,7 @@ def symp_flow(self, end_time=1.0, n_steps=20):
         return self.iterate(step(step_size), n_steps)
 
     def exp(self, cotangent_vec, base_point, n_steps=20, **kwargs):
-        """Exponential map associated to the cometric. I
+        """Exponential map associated to the cometric.
 
         Exponential map at base_point of cotangent_vec computed by integration
         of the Hamiltonian equation (initial value problem), using the cometric.
@@ -272,9 +274,6 @@ def exp(self, cotangent_vec, base_point, n_steps=20, **kwargs):
         n_steps : int
             Number of discrete time steps to take in the integration.
             Optional, default: N_STEPS.
-        point_type : str, {'vector', 'matrix'}
-            Type of representation used for points.
-            Optional, default: None.
 
         Returns
         -------

geomstats/integrator.py

@@ -34,7 +34,7 @@ def euler_step(force, state, time, dt):
     state : array-like, shape=[2, dim]
         State at time t, corresponds to position and velocity variables at
         time t.
-    time ; float
+    time : float
         Time variable.
     dt : float
         Time-step in the integration.
@@ -61,6 +61,8 @@ def symplectic_euler_step(force, state, time, dt):
         time t.
     force : callable
         Vector field that is being integrated.
+    time : float
+        Time variable.
     dt : float
         Time-step in the integration.
 
@@ -84,6 +86,8 @@ def leapfrog_step(force, state, time, dt):
         time t.
     force : callable
         Vector field that is being integrated.
+    time : float
+        Time variable.
     dt : float
         Time-step in the integration.
 
@@ -107,7 +111,7 @@ def rk2_step(force, state, time, dt):
     state : array-like, shape=[2, dim]
         State at time t, corresponds to position and velocity variables at
         time t.
-    time ; float
+    time : float
         Time variable.
     dt : float
         Time-step in the integration.

geomstats/learning/expectation_maximization.py

@@ -224,7 +224,7 @@ class RiemannianEM(TransformerMixin, ClusterMixin, BaseEstimator):
         Optional, default: 1e-2.
         Convergence tolerance. If the difference of mean distance
         between two steps is lower than tol.
-    lr_mean : float
+    init_step_size : float
         Learning rate in the gradient descent computation of the Frechet means.
         Optional, default: 1.
     max_iter : int
@@ -276,7 +276,7 @@ def __init__(
         n_gaussians=8,
         initialisation_method="random",
         tol=DEFAULT_TOL,
-        lr_mean=1.0,
+        init_step_size=1.0,
         max_iter=100,
         max_iter_mean=100,
         tol_mean=1e-4,
@@ -296,7 +296,7 @@ def __init__(
         self.variances_range = None
         self.normalization_factor_var = None
         self.phi_inv_var = None
-        self.lr_mean = lr_mean
+        self.init_step_size = init_step_size
         self.max_iter = max_iter
         self.max_iter_mean = max_iter_mean
         self.tol_mean = tol_mean
@@ -318,11 +318,11 @@ def update_means(self, data, posterior_probabilities):
 
         mean = FrechetMean(
             metric=self.metric,
-            method=self.mean_method,
-            lr=self.lr_mean,
-            epsilon=self.tol_mean,
             max_iter=self.max_iter_mean,
+            epsilon=self.tol_mean,
             point_type=self.point_type,
+            method=self.mean_method,
+            init_step_size=self.init_step_size,
         )
 
         data_expand = gs.expand_dims(data, 1)
@@ -469,8 +469,8 @@ def fit(self, data):
                 metric=self.metric,
                 n_clusters=self.n_gaussians,
                 init="random",
+                init_step_size=self.init_step_size,
                 mean_method="batch",
-                lr=self.lr_mean,
             )
 
             centroids = kmeans.fit(X=data)

geomstats/learning/exponential_barycenter.py

@@ -15,7 +15,13 @@
 
 
 def _default_gradient_descent(
-    group, points, weights=None, max_iter=32, step=1.0, epsilon=EPSILON, verbose=False
+    group,
+    points,
+    weights=None,
+    max_iter=32,
+    init_step_size=1.0,
+    epsilon=EPSILON,
+    verbose=False,
 ):
     """Compute the (weighted) group exponential barycenter of `points`.
 
@@ -35,7 +41,7 @@ def _default_gradient_descent(
         Tolerance to reach convergence. The exstrinsic norm of the
         gradient is used as criterion.
         Optional, default: 1e-6.
-    step : float
+    init_step_size : float
         Learning rate in the gradient descent.
         Optional, default: 1.
     verbose : bool
@@ -69,7 +75,9 @@ def _default_gradient_descent(
         inv_mean = group.inverse(mean)
         centered_points = group.compose(inv_mean, points)
         logs = group.log(point=centered_points)
-        tangent_mean = step * gs.einsum("n, nk...->k...", weights / sum_weights, logs)
+        tangent_mean = init_step_size * gs.einsum(
+            "n, nk...->k...", weights / sum_weights, logs
+        )
         mean_next = group.compose(mean, group.exp(tangent_vec=tangent_mean))
 
         grad_norm = gs.linalg.norm(tangent_mean)
@@ -103,7 +111,7 @@ class ExponentialBarycenter(BaseEstimator):
         Tolerance to reach convergence. The exstrinsic norm of the
         gradient is used as criterion.
         Optional, default: 1e-6.
-    step : float
+    init_step_size : float
         Learning rate in the gradient descent.
         Optional, default: 1.
     verbose : bool
@@ -120,15 +128,15 @@ def __init__(
         group,
         max_iter=32,
         epsilon=EPSILON,
-        step=1.0,
+        init_step_size=1.0,
         point_type=None,
         verbose=False,
     ):
         self.group = group
         self.max_iter = max_iter
         self.epsilon = epsilon
         self.verbose = verbose
-        self.step = step
+        self.init_step_size = init_step_size
         self.point_type = point_type
         self.estimate_ = None
 
@@ -160,12 +168,12 @@ def fit(self, X, y=None, weights=None):
 
         else:
             mean = _default_gradient_descent(
+                group=self.group,
                 points=X,
                 weights=weights,
-                group=self.group,
                 max_iter=self.max_iter,
+                init_step_size=self.init_step_size,
                 epsilon=self.epsilon,
-                step=self.step,
                 verbose=self.verbose,
             )
         self.estimate_ = mean