From eb6fed709e8569f4b289e988cc50cb99ce65c7b8 Mon Sep 17 00:00:00 2001
From: William de Vazelhes <william.de-vazelhes@inria.fr>
Date: Wed, 22 May 2019 10:31:40 +0200
Subject: [PATCH] Add checks for bounds argument

---
 metric_learn/itml.py      | 18 +++++++++++-------
 test/metric_learn_test.py | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/metric_learn/itml.py b/metric_learn/itml.py
index 6cb34313..e3ff515a 100644
--- a/metric_learn/itml.py
+++ b/metric_learn/itml.py
@@ -69,9 +69,13 @@ def _fit(self, pairs, y, bounds=None):
       X = np.vstack({tuple(row) for row in pairs.reshape(-1, pairs.shape[2])})
       self.bounds_ = np.percentile(pairwise_distances(X), (5, 95))
     else:
-      assert len(bounds) == 2
+      bounds = check_array(bounds, allow_nd=False, ensure_min_samples=0,
+                           ensure_2d=False)
+      bounds = bounds.ravel()
+      if bounds.size != 2:
+        raise ValueError("`bounds` should be an array-like of two elements.")
       self.bounds_ = bounds
-    self.bounds_[self.bounds_==0] = 1e-9
+    self.bounds_[self.bounds_ == 0] = 1e-9
     # init metric
     if self.A0 is None:
       A = np.identity(pairs.shape[2])
@@ -134,7 +138,7 @@ class ITML(_BaseITML, _PairsClassifierMixin):
 
   Attributes
   ----------
-  bounds_ : array-like, shape=(2,)
+  bounds_ : `numpy.ndarray`, shape=(2,)
       Bounds on similarity, aside slack variables, s.t.
       ``d(a, b) < bounds_[0]`` for all given pairs of similar points ``a``
       and ``b``, and ``d(c, d) > bounds_[1]`` for all given pairs of
@@ -171,7 +175,7 @@ def fit(self, pairs, y, bounds=None, calibration_params=None):
         preprocessor.
     y: array-like, of shape (n_constraints,)
         Labels of constraints. Should be -1 for dissimilar pair, 1 for similar.
-    bounds : `list` of two numbers
+    bounds : array-like of two numbers
         Bounds on similarity, aside slack variables, s.t.
         ``d(a, b) < bounds_[0]`` for all given pairs of similar points ``a``
         and ``b``, and ``d(c, d) > bounds_[1]`` for all given pairs of
@@ -192,7 +196,7 @@ def fit(self, pairs, y, bounds=None, calibration_params=None):
     calibration_params = (calibration_params if calibration_params is not
                           None else dict())
     self._validate_calibration_params(**calibration_params)
-    self._fit(pairs, y)
+    self._fit(pairs, y, bounds=bounds)
     self.calibrate_threshold(pairs, y, **calibration_params)
     return self
 
@@ -202,7 +206,7 @@ class ITML_Supervised(_BaseITML, TransformerMixin):
 
   Attributes
   ----------
-  bounds_ : array-like, shape=(2,)
+  bounds_ : `numpy.ndarray`, shape=(2,)
       Bounds on similarity, aside slack variables, s.t.
       ``d(a, b) < bounds_[0]`` for all given pairs of similar points ``a``
       and ``b``, and ``d(c, d) > bounds_[1]`` for all given pairs of
@@ -275,7 +279,7 @@ def fit(self, X, y, random_state=np.random, bounds=None):
     random_state : numpy.random.RandomState, optional
         If provided, controls random number generation.
 
-    bounds : `list` of two numbers
+    bounds : array-like of two numbers
         Bounds on similarity, aside slack variables, s.t.
         ``d(a, b) < bounds_[0]`` for all given pairs of similar points ``a``
         and ``b``, and ``d(c, d) > bounds_[1]`` for all given pairs of
diff --git a/test/metric_learn_test.py b/test/metric_learn_test.py
index a785d60d..c3efd9a3 100644
--- a/test/metric_learn_test.py
+++ b/test/metric_learn_test.py
@@ -18,7 +18,7 @@
   HAS_SKGGM = True
 from metric_learn import (LMNN, NCA, LFDA, Covariance, MLKR, MMC,
                           LSML_Supervised, ITML_Supervised, SDML_Supervised,
-                          RCA_Supervised, MMC_Supervised, SDML)
+                          RCA_Supervised, MMC_Supervised, SDML, ITML)
 # Import this specially for testing.
 from metric_learn.constraints import wrap_pairs
 from metric_learn.lmnn import python_LMNN
@@ -109,6 +109,43 @@ def test_deprecation_bounds(self):
     assert_warns_message(DeprecationWarning, msg, itml_supervised.fit, X, y)
 
 
+@pytest.mark.parametrize('bounds', [None, (20., 100.), [20., 100.],
+                                    np.array([20., 100.]),
+                                    np.array([[20., 100.]]),
+                                    np.array([[20], [100]])])
+def test_bounds_parameters_valid(bounds):
+  """Asserts that we can provide any array-like of two elements as bounds,
+  and that the attribute bound_ is a numpy array"""
+
+  pairs = np.array([[[-10., 0.], [10., 0.]], [[0., 50.], [0., -60]]])
+  y_pairs = [1, -1]
+  itml = ITML()
+  itml.fit(pairs, y_pairs, bounds=bounds)
+
+  X = np.array([[0, 0], [0, 1], [2, 0], [2, 1]])
+  y = np.array([1, 0, 1, 0])
+  itml_supervised = ITML_Supervised()
+  itml_supervised.fit(X, y, bounds=bounds)
+
+
+@pytest.mark.parametrize('bounds', ['weird', ['weird1', 'weird2'],
+                                    np.array([1, 2, 3])])
+def test_bounds_parameters_invalid(bounds):
+  """Assert that if a non array-like is put for bounds, or an array-like
+  of length different than 2, an error is returned"""
+  pairs = np.array([[[-10., 0.], [10., 0.]], [[0., 50.], [0., -60]]])
+  y_pairs = [1, -1]
+  itml = ITML()
+  with pytest.raises(Exception):
+    itml.fit(pairs, y_pairs, bounds=bounds)
+
+  X = np.array([[0, 0], [0, 1], [2, 0], [2, 1]])
+  y = np.array([1, 0, 1, 0])
+  itml_supervised = ITML_Supervised()
+  with pytest.raises(Exception):
+    itml_supervised.fit(X, y, bounds=bounds)
+
+
 class TestLMNN(MetricTestCase):
   def test_iris(self):
     # Test both impls, if available.