Merge pull request #217 from dkobak/precomputed_affinity

Implement precomputed affinity class

Author: pavlin-policar

openTSNE/affinity.py

@@ -1264,3 +1264,26 @@ def to_new(self, data, k_neighbors=None, return_distances=False):
             return P, neighbors, distances
 
         return P
+
+
+class PrecomputedAffinities(Affinities):
+    """Use a precomputed affinity matrix.
+
+    Parameters
+    ----------
+    affinities: scipy.sparse.csr_matrix, np.ndarray
+        A N x N matrix containing the affinities.
+    normalize: bool
+        Normalize the affinity matrix to sum to 1. Default is True.
+
+    """
+
+    def __init__(self, affinities, normalize=True):
+        if not isinstance(affinities, sp.csr_matrix):
+            affinities = sp.csr_matrix(affinities)
+        if normalize:
+            affinities /= np.sum(affinities)
+        self.P = affinities
+
+    def to_new(self, data, return_distances=False):
+        raise RuntimeError("Precomputed affinity matrices cannot be queried.")

tests/test_affinities.py

@@ -340,3 +340,26 @@ def test_to_new(self):
         for method_name, cls in AFFINITY_CLASSES:
             aff = cls(knn_index=knn_index)
             aff.to_new(self.iris)
+
+class TestPrecomputedAffinity(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.iris = datasets.load_iris().data
+
+    def test_precomputed_affinity_matches(self):
+        aff1 = PerplexityBasedNN(self.iris)
+        aff2 = affinity.PrecomputedAffinities(aff1.P)
+        np.testing.assert_allclose(
+            aff1.P.toarray(),
+            aff2.P.toarray(),
+            err_msg="Precomputed affinities are not passed in correctly.",
+        )
+
+    def test_precomputed_affinity_normalization(self):
+        aff1 = PerplexityBasedNN(self.iris)
+        aff2 = affinity.PrecomputedAffinities(aff1.P * 2)
+        np.testing.assert_allclose(
+            aff1.P.toarray(),
+            aff2.P.toarray(),
+            err_msg="Precomputed affinities are not normalized to sum to 1.",
+        )

tests/test_different_usages.py

@@ -109,6 +109,15 @@ def test_affinity(self):
             new_embedding.optimize(10, learning_rate=0.1, inplace=True)
             self.eval_embedding(new_embedding, self.y, f"transform::{aff.__class__.__name__}")
 
+    def test_precomputed_affinity(self):
+        # Setup precomputed affinity
+        aff = affinity.Uniform(self.x).P
+        aff *= 10  # rescale
+        precomputed_affinity = affinity.PrecomputedAffinities(aff, normalize=True)
+
+        embedding = TSNE().fit(affinities=precomputed_affinity, initialization="spectral")
+        self.eval_embedding(embedding, self.y, aff.__class__.__name__)
+
 
 class TestUsageWithCustomAffinityAndCustomNeighbors(TestUsage):
     def test_affinity_with_queryable_knn_index(self):