Add float32 compatibility to KMedoids (#120)

* add float32 compatibility

* black

* fix CI and add test on type

* add test on transform

Author: TimotheeMathieu

sklearn_extra/cluster/_k_medoids.py

@@ -94,8 +94,8 @@ class KMedoids(BaseEstimator, ClusterMixin, TransformerMixin):
     >>> kmedoids.predict([[0,0], [4,4]])
     array([0, 1])
     >>> kmedoids.cluster_centers_
-    array([[1, 2],
-           [4, 2]])
+    array([[1., 2.],
+           [4., 2.]])
     >>> kmedoids.inertia_
     8.0
 
@@ -185,7 +185,9 @@ def fit(self, X, y=None):
         random_state_ = check_random_state(self.random_state)
 
         self._check_init_args()
-        X = check_array(X, accept_sparse=["csr", "csc"])
+        X = check_array(
+            X, accept_sparse=["csr", "csc"], dtype=[np.float64, np.float32]
+        )
         if self.n_clusters > X.shape[0]:
             raise ValueError(
                 "The number of medoids (%d) must be less "
@@ -315,7 +317,9 @@ def transform(self, X):
         X_new : {array-like, sparse matrix}, shape=(n_query, n_clusters)
             X transformed in the new space of distances to cluster centers.
         """
-        X = check_array(X, accept_sparse=["csr", "csc"])
+        X = check_array(
+            X, accept_sparse=["csr", "csc"], dtype=[np.float64, np.float32]
+        )
 
         if self.metric == "precomputed":
             check_is_fitted(self, "medoid_indices_")
@@ -345,7 +349,9 @@ def predict(self, X):
         labels : array, shape = (n_query,)
             Index of the cluster each sample belongs to.
         """
-        X = check_array(X, accept_sparse=["csr", "csc"])
+        X = check_array(
+            X, accept_sparse=["csr", "csc"], dtype=[np.float64, np.float32]
+        )
 
         if self.metric == "precomputed":
             check_is_fitted(self, "medoid_indices_")

sklearn_extra/cluster/tests/test_k_medoids.py

@@ -34,15 +34,18 @@
 @pytest.mark.parametrize(
     "init", ["random", "heuristic", "build", "k-medoids++"]
 )
-def test_kmedoid_results(method, init):
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_kmedoid_results(method, init, dtype):
     expected = np.hstack([np.zeros(50), np.ones(50)])
     km = KMedoids(n_clusters=2, init=init, method=method, random_state=rng)
-    km.fit(X_cc)
+    km.fit(X_cc.astype(dtype))
     # This test use data that are not perfectly separable so the
     # accuracy is not 1. Accuracy around 0.85
     assert (np.mean(km.labels_ == expected) > 0.8) or (
         1 - np.mean(km.labels_ == expected) > 0.8
     )
+    assert dtype is np.dtype(km.cluster_centers_.dtype).type
+    assert dtype is np.dtype(km.transform(X_cc.astype(dtype)).dtype).type
 
 
 def test_medoids_invalid_method():