diff --git a/pretab/transformers/splines/thinplate_spline.py b/pretab/transformers/splines/thinplate_spline.py index 65bfea0..fcd3df4 100644 --- a/pretab/transformers/splines/thinplate_spline.py +++ b/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: diff --git a/tests/test_thinplate_transformer.py b/tests/test_thinplate_transformer.py index 6875635..f824252 100644 --- a/tests/test_thinplate_transformer.py +++ b/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_) + )