perf(splines): drop the n_train square projector from tprs transform

pretab/transformers/splines/thinplate_spline.py

@@ -120,6 +120,8 @@ def fit(self, X, y=None):
         K = self._tps_kernel(r)
 
         ZTZ_inv = np.linalg.pinv(Z.T @ Z)
+        # Cached so ``transform`` does not redo the pseudo-inverse on every call.
+        self._ztz_inv_ = ZTZ_inv
         P = np.eye(n) - Z @ ZTZ_inv @ Z.T
         KP = P @ K @ P
 
@@ -149,9 +151,11 @@ def transform(self, X):
         K_new = self._tps_kernel(r_new)
 
         Z = self.Z_
-        ZTZ_inv = np.linalg.pinv(Z.T @ Z)
-        P_new = np.eye(Z.shape[0]) - Z @ ZTZ_inv @ Z.T
-        K_new_proj = K_new @ P_new
+        # ``P = I - Z (Z'Z)^-1 Z'`` is n_train x n_train, so materializing it (and
+        # the identity it is built from) cost O(n_train^2) memory on every call --
+        # 250 MB to transform ten rows against an 8k-row fit. Distributing the
+        # product avoids it entirely: K_new @ P == K_new - (K_new @ Z) (Z'Z)^-1 Z'.
+        K_new_proj = K_new - (K_new @ Z) @ self._ztz_inv_ @ Z.T
 
         out = K_new_proj @ self.basis_
         if self.include_bias:

tests/test_thinplate_transformer.py

@@ -78,3 +78,53 @@ def test_tprs_transform_requires_fit():
         transformer.transform(np.random.rand(5, 1))
     with pytest.raises(NotFittedError):
         transformer.get_penalty_matrix()
+
+
+# --------------------------------------------------------------------------- #
+# ``transform`` must not materialize the n_train x n_train projector.
+#
+# ``P = I - Z (Z'Z)^-1 Z'`` was rebuilt in full on every call, so transforming a
+# handful of rows against a large fit allocated hundreds of MB. Distributing the
+# product gives the same numbers without the square intermediate.
+# --------------------------------------------------------------------------- #
+def test_tprs_transform_matches_explicit_projector():
+    from scipy.spatial.distance import cdist
+
+    X = np.linspace(0, 1, 300).reshape(-1, 1)
+    X_new = np.linspace(-0.2, 1.2, 37).reshape(-1, 1)
+    transformer = ThinPlateSplineTransformer(output_dim=6).fit(X)
+
+    Z = transformer.Z_
+    K_new = transformer._tps_kernel(cdist(X_new, transformer.x_))
+    explicit = (K_new @ (np.eye(Z.shape[0]) - Z @ np.linalg.pinv(Z.T @ Z) @ Z.T)) @ transformer.basis_
+
+    np.testing.assert_allclose(transformer.transform(X_new), explicit, rtol=1e-9, atol=1e-9)
+
+
+def test_tprs_transform_does_not_allocate_a_train_sized_matrix():
+    import tracemalloc
+
+    n_train = 1200
+    transformer = ThinPlateSplineTransformer(output_dim=6).fit(
+        np.linspace(0, 1, n_train).reshape(-1, 1)
+    )
+    dense_projector_bytes = n_train * n_train * 8  # float64 n_train x n_train
+
+    tracemalloc.start()
+    try:
+        transformer.transform(np.linspace(0, 1, 10).reshape(-1, 1))
+        peak = tracemalloc.get_traced_memory()[1]
+    finally:
+        tracemalloc.stop()
+
+    assert peak < dense_projector_bytes / 4
+
+
+def test_tprs_caches_the_pseudo_inverse():
+    X = np.linspace(0, 1, 50).reshape(-1, 1)
+    transformer = ThinPlateSplineTransformer(output_dim=4).fit(X)
+
+    assert hasattr(transformer, "_ztz_inv_")
+    np.testing.assert_allclose(
+        transformer._ztz_inv_, np.linalg.pinv(transformer.Z_.T @ transformer.Z_)
+    )