Enable out-of-sample transform with precomputed_affinity

.gitignore

@@ -55,6 +55,9 @@ Python/doc/build/
 .ipynb_checkpoints/
 Python/tutorial/cache
 
+# Virtual environments
+.venv/
+
 # Mac
 .DS_Store
 

Python/phate/phate.py

@@ -10,6 +10,7 @@
 import graphtools
 from sklearn.base import BaseEstimator
 from sklearn.exceptions import NotFittedError
+from sklearn.preprocessing import normalize
 from scipy import sparse
 import warnings
 import tasklogger
@@ -941,7 +942,12 @@ def transform(self, X=None, t_max=100, plot_optimal_t=False, ax=None):
             Accepted data types: `numpy.ndarray`,
             `scipy.sparse.spmatrix`, `pd.DataFrame`, `anndata.AnnData`. If
             `knn_dist` is 'precomputed', `data` should be a n_samples x
-            n_samples distance or affinity matrix
+            n_samples distance or affinity matrix. Exception: if `knn_dist`
+            is 'precomputed_affinity', `X` may instead be a
+            n_query x n_samples query-to-train affinity matrix, in which
+            case PHATE performs an out-of-sample extension by row-
+            normalizing `X` (aggregated by landmark first, if the graph was
+            fit with landmarks) and applying it to the fitted embedding.
 
         t_max : int, optional, default: 100
             maximum t to test if `t` is set to 'auto'
@@ -974,6 +980,12 @@ def transform(self, X=None, t_max=100, plot_optimal_t=False, ax=None):
                 RuntimeWarning,
             )
             if (
+                self.knn_dist == "precomputed_affinity"
+                and isinstance(self.graph, graphtools.graphs.TraditionalGraph)
+                and self.graph.precomputed == "affinity"
+            ):
+                return self._transform_precomputed_affinity(X)
+            elif (
                 isinstance(self.graph, graphtools.graphs.TraditionalGraph)
                 and self.graph.precomputed is not None
             ):
@@ -987,27 +999,72 @@ def transform(self, X=None, t_max=100, plot_optimal_t=False, ax=None):
                 transitions = self.graph.extend_to_data(X)
                 return self.graph.interpolate(self.embedding, transitions)
         else:
-            diff_potential = self._calculate_potential(
-                t_max=t_max, plot_optimal_t=plot_optimal_t, ax=ax
-            )
-            if self.embedding is None:
-                with _logger.log_task("{} MDS".format(self.mds)):
-                    self.embedding = mds.embed_MDS(
-                        diff_potential,
-                        ndim=self.n_components,
-                        how=self.mds,
-                        solver=self.mds_solver,
-                        distance_metric=self.mds_dist,
-                        n_jobs=self.n_jobs,
-                        seed=self.random_state,
-                        verbose=max(self.verbose - 1, 0),
-                    )
+            self._ensure_embedded(t_max=t_max, plot_optimal_t=plot_optimal_t, ax=ax)
             if isinstance(self.graph, graphtools.graphs.LandmarkGraph):
                 _logger.log_debug("Extending to original data...")
                 return self.graph.interpolate(self.embedding)
             else:
                 return self.embedding
 
+    def _ensure_embedded(self, t_max=100, plot_optimal_t=False, ax=None):
+        """Ensures `self.embedding` holds the MDS embedding of `self.diff_op`
+
+        For a landmark graph, this is the embedding of the landmarks; for a
+        non-landmark graph, this is the embedding of the fitted data.
+        """
+        if self.embedding is None:
+            diff_potential = self._calculate_potential(
+                t_max=t_max, plot_optimal_t=plot_optimal_t, ax=ax
+            )
+            with _logger.log_task("{} MDS".format(self.mds)):
+                self.embedding = mds.embed_MDS(
+                    diff_potential,
+                    ndim=self.n_components,
+                    how=self.mds,
+                    solver=self.mds_solver,
+                    distance_metric=self.mds_dist,
+                    n_jobs=self.n_jobs,
+                    seed=self.random_state,
+                    verbose=max(self.verbose - 1, 0),
+                )
+        return self.embedding
+
+    def _transform_precomputed_affinity(self, X):
+        """Out-of-sample extension for a precomputed affinity graph
+
+        Given a query-to-train affinity matrix `X`, computes the PHATE
+        embedding of the query points as a transition-weighted combination
+        of the training (or landmark) embedding:
+
+        - No landmarks: row-normalize `X` to get the query-train transition
+          matrix, then apply it to the training embedding.
+        - Landmarks: aggregate the columns of `X` by landmark cluster
+          assignment, row-normalize to get the query-landmark transition
+          matrix, then apply it to the landmark embedding.
+        """
+        self._ensure_embedded()
+        if isinstance(self.graph, graphtools.graphs.LandmarkGraph):
+            clusters = self.graph.clusters
+            landmarks = np.unique(clusters)
+            if sparse.issparse(X):
+                pnm = sparse.hstack(
+                    [
+                        sparse.csr_matrix(X[:, clusters == i].sum(axis=1))
+                        for i in landmarks
+                    ]
+                )
+            else:
+                pnm = np.array(
+                    [np.sum(X[:, clusters == i], axis=1) for i in landmarks]
+                ).T
+            transitions = normalize(pnm, norm="l1", axis=1)
+        else:
+            transitions = normalize(X, norm="l1", axis=1)
+        embedding = transitions.dot(self.embedding)
+        if sparse.issparse(embedding):
+            embedding = embedding.toarray()
+        return np.asarray(embedding)
+
     def fit_transform(self, X, **kwargs):
         """Computes the diffusion operator and the position of the cells in the
         embedding space

Python/test/test_phate.py

@@ -14,7 +14,9 @@
 import phate
 import graphtools
 import pytest
+from scipy import sparse
 from scipy.spatial.distance import pdist, squareform
+from sklearn.preprocessing import normalize
 
 # Optional dependencies
 try:
@@ -423,6 +425,122 @@ def test_phate_precomputed_affinity():
     print("✓ Test 13 PASSED\n")
 
 
+def test_phate_precomputed_affinity_out_of_sample_transform():
+    """Out-of-sample transform with precomputed_affinity (no landmarks, issue #181)"""
+    print("=" * 70)
+    print("TEST 13b: Out-of-sample transform with precomputed affinity, no landmarks")
+    print("=" * 70)
+
+    data, _ = create_test_data()
+    train_data, test_data = data[:250], data[250:]
+
+    # Build a real affinity structure the way an external proximity model would:
+    # a kernel among training points, and a kernel from query points to
+    # training points.
+    G = graphtools.Graph(
+        train_data, knn=5, decay=40, distance="euclidean", random_state=42, verbose=False
+    )
+    K_train = G.kernel
+    K_test_train = G.build_kernel_to_data(test_data)
+
+    phate_op = phate.PHATE(
+        knn_dist="precomputed_affinity",
+        n_landmark=None,
+        t=10,
+        random_state=42,
+        verbose=False,
+    )
+    Z_train = phate_op.fit_transform(K_train)
+
+    with pytest.warns(RuntimeWarning, match="Pre-fit PHATE"):
+        Z_test = phate_op.transform(K_test_train)
+
+    assert Z_test.shape == (test_data.shape[0], 2)
+    assert np.all(np.isfinite(Z_test))
+
+    # Expected: row-normalize the query-train affinity, then apply to Z_train
+    P_query_train = normalize(K_test_train, norm="l1", axis=1)
+    expected = np.asarray(P_query_train.dot(Z_train))
+    assert np.allclose(Z_test, expected), "Out-of-sample embedding does not match P @ Z_train"
+
+    print("✓ Test 13b PASSED\n")
+
+
+def test_phate_precomputed_affinity_out_of_sample_transform_landmark():
+    """Out-of-sample transform with precomputed_affinity and landmarks (issue #181)"""
+    print("=" * 70)
+    print("TEST 13c: Out-of-sample transform with precomputed affinity, landmarks")
+    print("=" * 70)
+
+    data, _ = create_test_data()
+    train_data, test_data = data[:250], data[250:]
+
+    G = graphtools.Graph(
+        train_data, knn=5, decay=40, distance="euclidean", random_state=42, verbose=False
+    )
+    K_train = G.kernel
+    K_test_train = G.build_kernel_to_data(test_data)
+
+    phate_op = phate.PHATE(
+        knn_dist="precomputed_affinity",
+        n_landmark=10,
+        t=10,
+        random_state=42,
+        verbose=False,
+    )
+    phate_op.fit_transform(K_train)
+    assert isinstance(phate_op.graph, graphtools.graphs.LandmarkGraph)
+    Z_landmark = phate_op.embedding
+
+    with pytest.warns(RuntimeWarning, match="Pre-fit PHATE"):
+        Z_test = phate_op.transform(K_test_train)
+
+    assert Z_test.shape == (test_data.shape[0], 2)
+    assert np.all(np.isfinite(Z_test))
+
+    # Expected: aggregate query-train affinities by landmark cluster,
+    # row-normalize, then apply to the landmark embedding
+    clusters = phate_op.graph.clusters
+    landmarks = np.unique(clusters)
+    pnm = sparse.hstack(
+        [
+            sparse.csr_matrix(K_test_train[:, clusters == i].sum(axis=1))
+            for i in landmarks
+        ]
+    )
+    P_query_landmark = normalize(pnm, norm="l1", axis=1)
+    expected = np.asarray(P_query_landmark.dot(Z_landmark))
+    assert np.allclose(
+        Z_test, expected
+    ), "Out-of-sample landmark embedding does not match P @ Z_landmark"
+
+    print("✓ Test 13c PASSED\n")
+
+
+def test_phate_precomputed_distance_out_of_sample_transform_raises():
+    """Out-of-sample transform with precomputed_distance should still raise"""
+    print("=" * 70)
+    print("TEST 13d: Out-of-sample transform with precomputed distance still raises")
+    print("=" * 70)
+
+    data, _ = create_test_data()
+    train_data, test_data = data[:250], data[250:]
+
+    D_train = squareform(pdist(train_data, "euclidean"))
+
+    phate_op = phate.PHATE(
+        knn=5, t=10, knn_dist="precomputed_distance", verbose=False, random_state=42
+    )
+    phate_op.fit_transform(D_train)
+
+    D_test_train = np.zeros((test_data.shape[0], train_data.shape[0]))
+    with pytest.warns(RuntimeWarning, match="Pre-fit PHATE"):
+        with pytest.raises(ValueError, match="Cannot transform additional data"):
+            phate_op.transform(D_test_train)
+
+    print("✓ Test 13d PASSED\n")
+
+
 #####################################################
 # Input type tests
 #####################################################