feat: multivariate derivative estimation for surface / PDE-style discovery

[jkitchin]

Closes #15.

estimate_derivatives() differentiates along a single axis, which covers ODE discovery but not problems whose data is a surface and whose regression needs more than one partial — PDE-style discovery (u_t = F(u, u_x, u_xx, ...)) or transform laws such as time–temperature superposition, where y(x, T) = f(x + s(T)) implies y_T = s'(T)·y_x and both partials must come from one smoothed surface.

What this adds

New module jaxsr.derivatives, exported as SurfaceDerivatives and estimate_partial_derivatives:

est = SurfaceDerivatives(method="tensor_spline")   # or "local_poly", "gp"
est.fit(coords, values, sigma=...)                 # coords (n, d), values (n,)
y, dy = est.derivatives(coords, order=[(1, 0), (0, 1)])
y, dy, dy_se = est.derivatives(coords, order=[(1, 0)], return_std=True)

Gridded data goes in as axes plus an N-D array (fit([x, T], Y_grid)), so the rectangular-grid case needs no meshgrid glue; scattered data works in any dimension. est.coords_ holds the flattened sample locations.

Three smoothers, all returning analytic partials of the fitted smoother rather than finite differences of noisy raw data:

Method	Data	Cost	Uncertainty
"tensor_spline" (default)	gridded or scattered, any d	one penalized least-squares solve	posterior of the penalized fit
"local_poly"	irregular sampling	one weighted fit per query point	sandwich variance of the local fit
"gp"	irregular sampling, n up to a few hundred	cubic in n, capped by max_points	exact posterior, via the Hermite form of the kernel derivatives

Smoothing selection is blind to the downstream symbolic score, by construction. smoothing="auto" uses GCV (spline, local poly) or log marginal likelihood (GP); smoothing="sigma" uses a supplied noise level (the s = n·σ² rule); a float is used verbatim. There is no option to tune against a regression score — a smoother selected that way can manufacture whichever law the regression prefers, and the failure is silent.

The smoothing level actually used is reported, so smoothing-induced bias is visible rather than inferred: smoothing_, smoothing_source_, effective_dof_, residual_std_, noise_std_, and summary(). smoothing_scale re-runs the estimate at a deliberately different level, which turns the ×1/×3 sensitivity table from the issue into a two-line loop.

Validation

On a synthetic TTS surface (20 frequencies × 12 temperatures, Arrhenius shift, true E = 55.85 kJ/mol), the effective activation energy recovered from y_T / y_x:

noise	tensor_spline	local_poly	gp
0%	55.86	55.86	55.85
1%	55.83	55.58	55.98
3%	54.99	55.40	55.76

A heat-equation case (u_t = 0.1·u_xx, 0.2% noise) recovers y = 0.09583*u_xx, drifting to 0.09043 at smoothing_scale=10 — the smoothing bias made visible.

Testing

• 52 new tests in tests/test_derivatives.py (92% coverage of the new module), including recovery of the activation energy from a synthetic shift law and the visibility of deliberate over-smoothing.
• Full suite green: 673 passed, 10 skipped.
• Passes on the NumPy backend (scripts/test_under_numpy.py) — no new JAX API, so the shim needed no changes.
• black --check and ruff check clean; module doctests pass. Slowest new test is 1.3s locally.

Docs

docs/guides/surface-derivatives.md (every code block executed; quoted outputs are real), a skill guide plus a SKILL.md decision-tree entry, docs/api/derivatives.rst, and TOC / README / CHANGELOG entries. estimate_derivatives and SurfaceDerivatives now cross-reference each other.

Not included

No example notebook — the guide covers the same ground without adding an execution-time dependency. Nothing here touches #18's structured library blocks.

🤖 Generated with Claude Code

https://claude.ai/code/session_01S2L62FUGSWfnCgvPYEupj4

---

Generated by Claude Code