diff --git a/seapig/scores/knn.py b/seapig/scores/knn.py index bee843c..c4a2344 100644 --- a/seapig/scores/knn.py +++ b/seapig/scores/knn.py @@ -34,7 +34,7 @@ class KNNScore(EmbeddingScore, FAISSIndexMixin, ABC): ---------- k : int, default 1 Number of nearest neighbors used to compute the distance score. - stat : {'max', 'mean', 'median', 'min'}, default 'max' + stat : {'max', 'mean', 'median', 'min'}, default 'min' Statistic applied to aggregate distances across the k neighbors. pca : `TensorPCA` or None, default None Optional `TensorPCA` object for dimensionality reduction prior to scoring. @@ -57,7 +57,7 @@ class KNNScore(EmbeddingScore, FAISSIndexMixin, ABC): def __init__( self, k: int = 1, - stat: str = "max", + stat: str = "min", pca: TensorPCA | None = None, save_index: bool | Path = False, ) -> None: @@ -110,7 +110,7 @@ def _filter_outliers(self, q: float | bool = False) -> None: if self.index is None: self._setup_index() # temporary index scores, _ = self._distance(self.ref_embeddings, offset=1) - scores = self._stat(-scores) + scores = self._stat(scores) threshold = torch.quantile(scores.float(), q=q) index = scores >= threshold self.ref_embeddings = self.ref_embeddings[index, :] @@ -210,7 +210,7 @@ class EuclideanScore(KNNScore): ---------- k : int, default 1 Number of nearest neighbors to use. - stat : {'max', 'mean', 'median', 'min'}, default 'max' + stat : {'max', 'mean', 'median', 'min'}, default 'min' Statistic to aggregate distances across the k neighbors. pca : `TensorPCA` or None, default None Optional `TensorPCA` object for dimensionality reduction prior to scoring. @@ -241,7 +241,7 @@ class EuclideanScore(KNNScore): def __init__( self, k: int = 1, - stat: str = "max", + stat: str = "min", pca: TensorPCA | None = None, save_index: bool | Path = False, ) -> None: @@ -287,7 +287,7 @@ class CosineScore(KNNScore): ---------- k : int, default 1 Number of nearest neighbors to use. - stat : {'max', 'mean', 'median', 'min'}, default 'max' + stat : {'max', 'mean', 'median', 'min'}, default 'min' Statistic to aggregate distances across the k neighbors. pca : `TensorPCA` or None, default None Optional `TensorPCA` object for dimensionality reduction prior to scoring. @@ -307,7 +307,7 @@ class CosineScore(KNNScore): def __init__( self, k: int = 1, - stat: str = "max", + stat: str = "min", pca: TensorPCA | None = None, save_index: bool | Path = False, ) -> None: @@ -366,7 +366,7 @@ class MahalanobisScore(KNNScore): ---------- k : int, default 1 Number of nearest neighbors to use. - stat : {'max', 'mean', 'median', 'min'}, default 'max' + stat : {'max', 'mean', 'median', 'min'}, default 'min' Statistic to aggregate distances across the k neighbors. pca : `TensorPCA` or None, default None Optional `TensorPCA` object for dimensionality reduction prior to scoring. @@ -387,7 +387,7 @@ class MahalanobisScore(KNNScore): def __init__( self, k: int = 1, - stat: str = "max", + stat: str = "min", pca: TensorPCA | None = None, save_index: bool | Path = False, ) -> None: diff --git a/tests/test_knn.py b/tests/test_knn.py index f9531bf..500500a 100644 --- a/tests/test_knn.py +++ b/tests/test_knn.py @@ -323,3 +323,18 @@ def test_knnstat_all_branches( result = score._stat(dists) expected_val = expected(dists) torch.testing.assert_close(result[0], expected_val) + + +def test_k_returns_diff(): + ref = torch.rand(5000, 256) + query = torch.rand(100, 256) + + s1 = EuclideanScore(k=1) + s2 = EuclideanScore(k=10) + s1.fit(ref) + s2.fit(ref) + + s_1 = s1.score(query) + s_2 = s2.score(query) + + assert not torch.allclose(s_1, s_2)