umap lower n_neighbors

koaning · Dec 28, 2023 · 77e7a94 · 77e7a94
1 parent 091a166
commit 77e7a94
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tests/test_estimators/test_umap_reconstruction.py b/tests/test_estimators/test_umap_reconstruction.py
@@ -27,7 +27,7 @@
     ),
 )
 def test_estimator_checks(test_fn):
-    outlier_mod = UMAPOutlierDetection(n_components=2, threshold=0.1)
+    outlier_mod = UMAPOutlierDetection(n_components=2, threshold=0.1, n_neighbors=3)
     test_fn(UMAPOutlierDetection.__name__, outlier_mod)
 
 
@@ -38,6 +38,12 @@ def dataset():
 
 
 def test_obvious_usecase(dataset):
-    mod = UMAPOutlierDetection(n_components=2, threshold=7.5, random_state=42, variant="absolute").fit(dataset)
+    mod = UMAPOutlierDetection(
+        n_components=2,
+        threshold=7.5,
+        random_state=42,
+        variant="absolute",
+        n_neighbors=3,
+    ).fit(dataset)
     assert mod.predict([[10] * 10]) == np.array([-1])
     assert mod.predict([[0.01] * 10]) == np.array([1])