ENH: add squared option

Author: William de Vazelhes

metric_learn/base_metric.py

@@ -243,24 +243,35 @@ def transform(self, X):
   def get_metric(self):
     transformer_T = self.transformer_.T.copy()
 
-    def metric_fun(u, v):
+    def metric_fun(u, v, squared=False):
       """This function computes the metric between u and v, according to the
       previously learned metric.
 
       Parameters
       ----------
       u : array-like, shape=(n_features,)
         The first point involved in the distance computation.
+
       v : array-like, shape=(n_features,)
         The second point involved in the distance computation.
+
+      squared : `bool`
+        If True, the function will return the squared metric between u and
+        v, which is faster to compute.
+
       Returns
       -------
       distance: float
         The distance between u and v according to the new metric.
       """
       u = validate_vector(u)
       v = validate_vector(v)
-      return euclidean(u.dot(transformer_T), v.dot(transformer_T))
+      transformed_diff = (u - v).dot(transformer_T)
+      dist = transformed_diff.dot(transformed_diff.T)
+      if not squared:
+        dist = np.sqrt(dist)
+      return dist
+
     return metric_fun
 
   get_metric.__doc__ = BaseMetricLearner.get_metric.__doc__

test/test_mahalanobis_mixin.py

@@ -245,4 +245,24 @@ def test_get_metric_compatible_with_scikit_learn(estimator, build_dataset):
   set_random_state(model)
   model.fit(input_data, labels)
   clustering = DBSCAN(metric=model.get_metric())
-  clustering.fit(X)
+  clustering.fit(X)
+
+
+@pytest.mark.parametrize('estimator, build_dataset', metric_learners,
+                         ids=ids_metric_learners)
+def test_get_squared_metric(estimator, build_dataset):
+  """Test that the squared metric returned is indeed the square of the
+  metric"""
+  input_data, labels, _, X = build_dataset()
+  model = clone(estimator)
+  set_random_state(model)
+  model.fit(input_data, labels)
+  metric = model.get_metric()
+
+  n_features = X.shape[1]
+  for seed in range(10):
+    rng = np.random.RandomState(seed)
+    a, b = (rng.randn(n_features) for _ in range(2))
+    assert_allclose(metric(a, b, squared=True),
+                    metric(a, b, squared=False)**2,
+                    rtol=1e-15)