Skip to content

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

Open
ChrisW09 wants to merge 1 commit into
mainfrom
perf/tprs-transform-projection
Open

perf(splines): drop the n_train square projector from tprs transform#25
ChrisW09 wants to merge 1 commit into
mainfrom
perf/tprs-transform-projection

Conversation

@ChrisW09

Copy link
Copy Markdown
Collaborator

Addresses the transform-side half of #13. Deliberately does not close it — see below.

Problem

ThinPlateSplineTransformer.transform rebuilt the null-space projector on every call:

P_new = np.eye(Z.shape[0]) - Z @ ZTZ_inv @ Z.T   # n_train x n_train
K_new_proj = K_new @ P_new

Z.shape[0] is the training size, so the cost of transforming a few rows scaled
quadratically with the fit — 250 MB and 42 ms to transform ten rows against an 8k-row fit —
and the pinv was recomputed each time.

Fix

The dense projector is never needed. Since P = I - Z (Z'Z)^-1 Z':

K_new @ P  ==  K_new - (K_new @ Z) (Z'Z)^-1 Z'

(Z'Z)^-1 is now cached at fit time as _ztz_inv_.

Result

Numerically identical — max absolute difference 1.4e-14 against the explicit-projector
formulation over a 300-row fit and 37 query points including out-of-range ones.

n_train transform 10 rows, before after peak alloc before after
2000 6.3 ms 0.41 ms 0.64 MB
8000 42.4 ms 0.72 ms 256.7 MB (at n=4000) 2.56 MB

Scope — why this does not close #13

#13 also covers the fit side, where a dense n x n cdist and a full eigh make 8k rows
take ~40 s and 50k rows get OOM-killed. That cost is inherent to the exact TPS
construction
, so fixing it means either documenting a supported ceiling, raising above a
threshold, or implementing a knot-subsampled low-rank variant (Wood's thin-plate regression
splines). Each is a product decision rather than a bug fix, and adding a size guard would
introduce a new user-visible failure mode.

I have left #13 open for that decision and marked this Refs, not Closes.

Tests

Three added to tests/test_thinplate_transformer.py:

  • test_tprs_transform_matches_explicit_projector — pins numerical equivalence against a
    literal np.eye-based projector, so the algebraic rewrite can't silently drift
  • test_tprs_transform_does_not_allocate_a_train_sized_matrixtracemalloc peak must stay
    under a quarter of a dense n_train² float64 array; this fails on the old code
  • test_tprs_caches_the_pseudo_inverse

Full suite: 483 passed, 9 xfailed. ruff check clean on changed files; pyright unchanged at
its pre-existing 71 errors.

🤖 Generated with Claude Code

``ThinPlateSplineTransformer.transform`` rebuilt the null-space projector
``P = I - Z (Z'Z)^-1 Z'`` in full on every call. ``Z.shape[0]`` is the
training-set size, so transforming a handful of rows against a moderate fit
allocated an n_train x n_train float64 matrix -- 250 MB and 42 ms to transform
ten rows after fitting on 8000 -- and redid the pseudo-inverse each time.

The dense projector is never needed: distributing the product gives
``K_new @ P == K_new - (K_new @ Z) (Z'Z)^-1 Z'``. Cache ``(Z'Z)^-1`` at fit
time and use that form instead.

Numerically identical to the previous formulation (max abs difference 1.4e-14
over a 300-row fit and 37 out-of-range query points).

    n=2000: 6.3 ms -> 0.41 ms,  peak 0.64 MB
    n=8000: 42.4 ms -> 0.72 ms, peak 2.56 MB (was 256.7 MB at n=4000)

This addresses the transform side only. The O(n^2) cdist and O(n^3) eigen-
decomposition in ``fit`` are inherent to the exact TPS construction and are
left for a separate decision.

Refs #13

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf(splines): tprs is O(n^2) memory and OOMs on moderate datasets

1 participant