Additional test on Iris inverse transformation
@@ -5,6 +5,7 @@ from scipy import sparse
import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics import adjusted_rand_score
+from sklearn.neighbors import KDTree
# ===================================================
@@ -129,3 +130,23 @@ def test_umap_clusterability_on_supervised_iris(supervised_iris_model, iris):
embedding = supervised_iris_model.embedding_
clusters = KMeans(3).fit_predict(embedding)
assert_greater_equal(adjusted_rand_score(clusters, iris.target), 0.95)
+
+
+# UMAP Inverse transform on Iris
+# ------------------------------
+def test_umap_inverse_transform_on_iris(iris, iris_model):
+ highd_tree = KDTree(iris.data)
+ fitter = iris_model
+ lowd_tree = KDTree(fitter.embedding_)
+ for i in range(1, 150, 20):
+ query_point = fitter.embedding_[i]
+ near_points = lowd_tree.query([query_point], k=5, return_distance=False)
+ centroid = np.mean(np.squeeze(fitter.embedding_[near_points]), axis=0)
+ highd_centroid = fitter.inverse_transform([centroid])
+ highd_near_points = highd_tree.query(
+ highd_centroid, k=10, return_distance=False
+ )
+ assert_greater_equal(
+ np.intersect1d(near_points, highd_near_points[0]).shape[0], 3
+ )
+
