ENH Add dtype preservation for Isomap (#24714)

rprkh · jjerphan · web-flow · commit 53234c5b6bae · 2022-10-27T10:27:35.000+02:00
Co-authored-by: Julien Jerphanion &lt;git@jjerphan.xyz&gt;
diff --git a/doc/whats_new/v1.2.rst b/doc/whats_new/v1.2.rst
@@ -413,6 +413,9 @@ Changelog
   `eigen_tol="auto"` in version 1.3.
   :pr:`23210` by :user:`Meekail Zain <micky774>`.
 
+- |Enhancement| :class:`manifold.Isomap` now preserves
+  dtype for `np.float32` inputs. :pr:`24714` by :user:`Rahil Parikh <rprkh>`.
+
 :mod:`sklearn.metrics`
 ......................
 
diff --git a/sklearn/manifold/_isomap.py b/sklearn/manifold/_isomap.py
@@ -294,6 +294,11 @@ def _fit_transform(self, X):
 
         self.dist_matrix_ = shortest_path(nbg, method=self.path_method, directed=False)
 
+        if self.nbrs_._fit_X.dtype == np.float32:
+            self.dist_matrix_ = self.dist_matrix_.astype(
+                self.nbrs_._fit_X.dtype, copy=False
+            )
+
         G = self.dist_matrix_**2
         G *= -0.5
 
@@ -412,3 +417,6 @@ def transform(self, X):
         G_X *= -0.5
 
         return self.kernel_pca_.transform(G_X)
+
+    def _more_tags(self):
+        return {"preserves_dtype": [np.float64, np.float32]}
diff --git a/sklearn/manifold/tests/test_isomap.py b/sklearn/manifold/tests/test_isomap.py
@@ -250,6 +250,32 @@ def test_isomap_fit_precomputed_radius_graph():
     assert_allclose(precomputed_result, result)
 
 
+def test_isomap_fitted_attributes_dtype(global_dtype):
+    """Check that the fitted attributes are stored accordingly to the
+    data type of X."""
+    iso = manifold.Isomap(n_neighbors=2)
+
+    X = np.array([[1, 2], [3, 4], [5, 6]], dtype=global_dtype)
+
+    iso.fit(X)
+
+    assert iso.dist_matrix_.dtype == global_dtype
+    assert iso.embedding_.dtype == global_dtype
+
+
+def test_isomap_dtype_equivalence():
+    """Check the equivalence of the results with 32 and 64 bits input."""
+    iso_32 = manifold.Isomap(n_neighbors=2)
+    X_32 = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
+    iso_32.fit(X_32)
+
+    iso_64 = manifold.Isomap(n_neighbors=2)
+    X_64 = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float64)
+    iso_64.fit(X_64)
+
+    assert_allclose(iso_32.dist_matrix_, iso_64.dist_matrix_)
+
+
 def test_isomap_raise_error_when_neighbor_and_radius_both_set():
     # Isomap.fit_transform must raise a ValueError if
     # radius and n_neighbors are provided.
