geomstats · luisfpereira · May 16, 2022 · May 2, 2022 · May 4, 2022 · May 4, 2022
diff --git a/geomstats/geometry/discrete_curves.py b/geomstats/geometry/discrete_curves.py
@@ -423,10 +423,11 @@ def __init__(self, ambient_manifold, metric=None, translation_invariant=True):
                 )
         else:
             self.ambient_metric = metric
+        self.default_point_type = "matrix"
         self.l2_metric = L2CurvesMetric(ambient_manifold, metric)
         self.translation_invariant = translation_invariant
 
-    def srv_transform(self, curve):
+    def srv_transform(self, curve, tol=gs.atol):
         """Square Root Velocity Transform (SRVT).
 
         Compute the square root velocity representation of a curve. The
@@ -440,11 +441,23 @@ def srv_transform(self, curve):
         curve : array-like, shape=[..., n_sampling_points, ambient_dim]
             Discrete curve.
 
+        tol : float
+            Tolerance value to decide duplicity of two consecutive sample
+            points on a given Discrete Curve.
+
         Returns
         -------
         srv : array-like, shape=[..., n_sampling_points - 1, ambient_dim]
             Square-root velocity representation of a discrete curve.
         """
+        if gs.any(
+            self.ambient_metric.norm(curve[..., 1:, :] - curve[..., :-1, :]) < tol
+        ):
+            raise AssertionError(
+                "The square root velocity framework "
+                "is only defined for discrete curves "
+                "with distinct consecutive sample points."
+            )
         curve_ndim = gs.ndim(curve)
         curve = gs.to_ndarray(curve, to_ndim=3)
         n_curves, n_sampling_points, n_coords = curve.shape

diff --git a/geomstats/learning/frechet_mean.py b/geomstats/learning/frechet_mean.py
@@ -11,6 +11,7 @@
 import geomstats.backend as gs
 import geomstats.errors as error
 import geomstats.vectorization
+from geomstats.geometry.discrete_curves import SRVMetric
 from geomstats.geometry.hypersphere import Hypersphere
 
 EPSILON = 1e-4
@@ -138,7 +139,13 @@ def _default_gradient_descent(
         sq_dists_between_iterates.append(sq_dist)
 
         var_is_0 = gs.isclose(var, 0.0)
-        sq_dist_is_small = gs.less_equal(sq_dist, epsilon * metric.dim)
+
+        metric_dim = metric.dim
+        if isinstance(metric, SRVMetric):
+            metric_dim = tangent_mean.shape[-2] * tangent_mean.shape[-1]
+
+        sq_dist_is_small = gs.less_equal(sq_dist, epsilon * metric_dim)
+
         condition = ~gs.logical_or(var_is_0, sq_dist_is_small)
         if not (condition or iteration == 0):
             break

diff --git a/geomstats/learning/riemannian_mean_shift.py b/geomstats/learning/riemannian_mean_shift.py
@@ -190,3 +190,16 @@ def predict(self, points):
         indices = self.metric.closest_neighbor_index(points, self.centers)
 
         return gs.take(self.centers, indices, axis=0)
+
+    def predict_labels(self, points):
+        """Predict the closest cluster label each point in 'points' belongs to.
+
+        Parameters
+        ----------
+        points : array-like, shape=[..., n_features]
+            Clusters of points.
+        """
+        if self.centers is None:
+            raise Exception("Not fitted")
+
+        return self.metric.closest_neighbor_index(points, self.centers)