From 98c78a02186298dcf20652f10cd197862d26ae43 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 12 Aug 2026 21:13:08 +0000 Subject: [PATCH] feat: multivariate derivative estimation for surface/PDE-style discovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 derivative — 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. Adds jaxsr.derivatives with SurfaceDerivatives, which fits a smoother to scattered or gridded N-D data and returns analytic partial derivatives of that smoother — never finite differences of noisy raw data — with standard errors: - "tensor_spline" (default): penalized tensor-product B-splines, any dimension, gridded or scattered, penalty by GCV. - "local_poly": local polynomial regression for irregular sampling. - "gp": anisotropic squared-exponential Gaussian process, giving the derivative posterior directly (Hermite form of the kernel derivatives). The smoothing hyperparameter is selectable only by criteria blind to the downstream symbolic score — GCV, log marginal likelihood, or a supplied noise level (smoothing="sigma", the s = n·sigma^2 rule). A smoother tuned against the regression that consumes it can manufacture whichever law the regression prefers, and the failure is silent. Smoothing flattens derivatives and biases any coefficient read off them, so the level actually used is reported: smoothing_, smoothing_source_, effective_dof_, residual_std_, noise_std_ and summary(). smoothing_scale re-runs the estimate at a deliberately different level to expose that bias. Also adds estimate_partial_derivatives() as a one-call wrapper, 52 tests (including recovery of an activation energy from a synthetic shift law), a user guide, a skill guide, and an API reference page. Closes #15 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01S2L62FUGSWfnCgvPYEupj4 --- .claude/skills/jaxsr/SKILL.md | 7 + .../jaxsr/guides/surface-derivatives.md | 184 ++ CHANGELOG.md | 18 + README.md | 1 + docs/_toc.yml | 1 + docs/api/derivatives.rst | 7 + docs/api/index.rst | 1 + docs/guides/surface-derivatives.md | 257 +++ src/jaxsr/__init__.py | 6 + src/jaxsr/derivatives.py | 1535 +++++++++++++++++ src/jaxsr/dynamics.py | 6 + src/jaxsr/skill/SKILL.md | 7 + src/jaxsr/skill/guides/surface-derivatives.md | 184 ++ tests/test_derivatives.py | 564 ++++++ 14 files changed, 2778 insertions(+) create mode 100644 .claude/skills/jaxsr/guides/surface-derivatives.md create mode 100644 docs/api/derivatives.rst create mode 100644 docs/guides/surface-derivatives.md create mode 100644 src/jaxsr/derivatives.py create mode 100644 src/jaxsr/skill/guides/surface-derivatives.md create mode 100644 tests/test_derivatives.py diff --git a/.claude/skills/jaxsr/SKILL.md b/.claude/skills/jaxsr/SKILL.md index 5b1c908..eb3b58d 100644 --- a/.claude/skills/jaxsr/SKILL.md +++ b/.claude/skills/jaxsr/SKILL.md @@ -386,6 +386,13 @@ See `guides/rsm.md` for RSM designs, canonical analysis, and optimization. See `guides/active-learning.md` for acquisition functions and adaptive sampling. +### "My data is a surface and I need partial derivatives" + +See `guides/surface-derivatives.md` for `SurfaceDerivatives`: several partials of one +smoothed N-D surface (`y_x` and `y_T`, or `u_t = F(u, u_x, u_xx, ...)`), with the +smoothing level chosen by GCV/marginal likelihood and reported. For a single time axis +(`dX/dt`), use `estimate_derivatives` / `discover_dynamics` instead. + ### "One expression isn't enough / the signal is a sum of many effects" See `guides/additive.md` for boosting-style additive symbolic regression diff --git a/.claude/skills/jaxsr/guides/surface-derivatives.md b/.claude/skills/jaxsr/guides/surface-derivatives.md new file mode 100644 index 0000000..a4bf5b1 --- /dev/null +++ b/.claude/skills/jaxsr/guides/surface-derivatives.md @@ -0,0 +1,184 @@ +# Multivariate Derivative Estimation (`SurfaceDerivatives`) + +Estimate **partial derivatives of a surface** — several partials from one smoothed +fit — for problems where the regression target or the basis library contains +derivatives. + +## When to use this + +| Situation | Use | +|-----------|-----| +| One state trajectory over time, need `dX/dt` | `estimate_derivatives(X, t, ...)` (see `jaxsr.dynamics`) | +| Whole ODE system from time series | `discover_dynamics(X, t, ...)` | +| Data is a surface over 2+ coordinates, need `y_x` **and** `y_T` | `SurfaceDerivatives` | +| PDE-style discovery: `u_t = F(u, u_x, u_xx, ...)` | `SurfaceDerivatives` | +| Transform/shift laws, e.g. `y(x, T) = f(x + s(T))` ⟹ `y_T = s'(T)·y_x` | `SurfaceDerivatives` | + +`estimate_derivatives` differentiates along a **single** axis (`X` is +`(n_times, n_states)`, `t` is 1-D). It cannot give you two partials of one surface. + +## API + +The snippets in this section assume `coords` (an `(n, d)` array), `values` (`(n,)`) and +`sigma` come from the user's data; the worked examples further down are self-contained. + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +est = SurfaceDerivatives(method="tensor_spline") # or "local_poly", "gp" +est.fit(coords, values, sigma=0.01) # coords (n, d), values (n,) + +y, dy = est.derivatives(coords, order=[(1, 0), (0, 1)]) +# y -> (n,) smoothed surface +# dy -> (n, 2) column 0 = d/dx0, column 1 = d/dx1 + +y, dy, dy_se = est.derivatives(coords, order=[(1, 0), (0, 1)], return_std=True) +print(est.summary()) # method, smoothing level, effective dof, residual std +``` + +Gridded data can be passed as axes plus an N-D array — no meshgrid needed: + +```python +est = SurfaceDerivatives().fit([x_axis, T_axis], Y_grid) # Y_grid (len(x), len(T)) +y, dy = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) # coords_ is the flat (n, 2) grid +``` + +One-call convenience wrapper: + +```python +from jaxsr import estimate_partial_derivatives + +y, dy = estimate_partial_derivatives(coords, values, order=[(1, 0), (0, 1)], + method="tensor_spline", sigma=0.01) +``` + +**Signature notes** (common mistakes): + +| API | Wrong | Right | +|-----|-------|-------| +| `derivatives()` | `dy = est.derivatives(...)` | returns a **tuple** `(y, dy)`, or `(y, dy, std)` with `return_std=True` | +| `order` | `order=1`, `order="x"` | a tuple per dimension: `(1, 0)`, or a list of them | +| single order | expecting shape `(n,)` | a single tuple still returns `(n, 1)` | +| query points | grid axes | `derivatives()` takes an `(n_query, d)` array — use `est.coords_` for the sample locations | +| `sigma` | a variance | a **standard deviation**, scalar or per point | + +## Choosing a method + +| Method | Data | Cost | Derivative uncertainty | Notes | +|--------|------|------|------------------------|-------| +| `"tensor_spline"` (default) | gridded or scattered, any `d` | fast | from the penalized-fit posterior | Best default. Penalty chosen by GCV. | +| `"local_poly"` | scattered, irregular | moderate (per-point fits) | sandwich variance of the local fit | Good on uneven sampling; degree bounds the total order. | +| `"gp"` | scattered, irregular, small `n` | `O(n³)`, capped by `max_points=800` | exact posterior, grows away from data | Best uncertainty; use when `n` is a few hundred. | + +Derivative orders are always analytic partials of the fitted smoother — never finite +differences of noisy raw data. + +## Choosing the smoothing level + +**The smoothing hyperparameter must never be tuned against the downstream symbolic +score.** If the smoother is selected by which law the regression likes best, it can +manufacture that law; the fit looks excellent and the failure is silent. + +| `smoothing=` | Meaning | +|--------------|---------| +| `"auto"` (default) | GCV (spline, local poly) or marginal likelihood (GP) | +| `"sigma"` | requires `sigma` at `fit()`; matches residual scatter to the known noise (the `s = n·σ²` rule) | +| float | use it verbatim: penalty `λ` (spline), bandwidth (local poly), noise variance (GP) | + +`smoothing_scale=3.0` multiplies whatever was selected — the cheapest way to check how +much a discovered coefficient depends on the derivative stage. + +## Reporting the smoothing actually used + +Smoothing biases derivatives toward zero, and any coefficient read off them inherits +that bias. Make it visible: + +```python +print(est.summary()) +est.smoothing_ # λ / bandwidth / noise variance actually used +est.smoothing_source_ # "gcv", "marginal_likelihood", "sigma", or "fixed" +est.effective_dof_ # effective degrees of freedom of the smoother +est.residual_std_ # residual scatter of the fit +``` + +Sensitivity check — rerun the whole pipeline at several smoothing scales and report the +spread, not just one number: + +```python +for scale in (1.0, 3.0, 10.0): + est = SurfaceDerivatives(smoothing_scale=scale).fit(coords, values, sigma=sigma) + ... # refit the symbolic stage, record the coefficient +``` + +## Example: PDE-style discovery + +```python +import numpy as np +from jaxsr import BasisLibrary, SurfaceDerivatives, SymbolicRegressor + +# Heat equation data: u_t = 0.1 * u_xx +x = np.linspace(0, 2 * np.pi, 40) +t = np.linspace(0, 1.0, 30) +xx, tt = np.meshgrid(x, t, indexing="ij") +u = np.exp(-0.1 * tt) * np.sin(xx) + 0.5 * np.exp(-0.9 * tt) * np.sin(3 * xx) +noise = 0.002 +u_obs = u + np.random.default_rng(0).normal(0, noise, u.shape) + +est = SurfaceDerivatives().fit([x, t], u_obs, sigma=noise) +values, d = est.derivatives(est.coords_, order=[(1, 0), (2, 0), (0, 1)]) +u_x, u_xx, u_t = d[:, 0], d[:, 1], d[:, 2] + +library = ( + BasisLibrary(n_features=3, feature_names=["u", "u_x", "u_xx"]) + .add_linear() + .add_interactions(max_order=2) +) +model = SymbolicRegressor(basis_library=library, max_terms=1).fit( + np.column_stack([values, u_x, u_xx]), u_t +) +print(model.expression_) # y = 0.0958*u_xx (true 0.1; the gap is smoothing bias) +``` + +## Example: shift law (`y_T = s'(T)·y_x`) + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +# Synthetic time-temperature superposition surface (Arrhenius shift, E = 55.85 kJ/mol) +gas_r, energy, t_ref = 8.314e-3, 55.85, 350.0 +x_axis = np.linspace(-2.0, 4.0, 20) # log frequency +T_axis = np.linspace(320.0, 400.0, 12) # temperature, K +xx, TT = np.meshgrid(x_axis, T_axis, indexing="ij") +shift = -(energy / (gas_r * np.log(10.0))) * (1.0 / TT - 1.0 / t_ref) +sigma = 0.01 +Y_grid = 2.0 + 1.5 * np.tanh(0.8 * (xx + shift - 1.0)) +Y_grid = Y_grid + np.random.default_rng(0).normal(0, sigma, Y_grid.shape) + +est = SurfaceDerivatives(smoothing="sigma").fit([x_axis, T_axis], Y_grid, sigma=sigma) +_, d = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) +y_x, y_T = d[:, 0], d[:, 1] + +keep = np.abs(y_x) > 0.15 * np.abs(y_x).max() # the ratio is ill-posed where y_x ≈ 0 +s_prime = y_T[keep] / y_x[keep] +T_keep = est.coords_[keep, 1] + +print(np.median(s_prime * gas_r * np.log(10.0) * T_keep**2)) # ≈ 55.9 kJ/mol +``` + +Then regress `s_prime` against `T_keep` with a `SymbolicRegressor` rather than assuming +the Arrhenius form. + +## Pitfalls + +- **Boundaries.** Every smoother is weakest at the edge of the data. Drop a margin + before reading derivatives, or expect the largest errors there. +- **Dividing partials.** Ratios like `y_T / y_x` blow up where the denominator crosses + zero. Mask small denominators, as above. +- **High orders.** Each extra order costs accuracy. `degree` must be at least the + highest order requested (`tensor_spline`: per dimension; `local_poly`: total). +- **GP size.** `method="gp"` is `O(n³)`; above `max_points=800` it raises rather than + hanging. Subsample or switch to `"tensor_spline"`. +- **Basis resolution.** The spline defaults to at most 12 basis functions per + dimension because GCV under-smooths derivatives; pass `n_basis=` for finer structure. diff --git a/CHANGELOG.md b/CHANGELOG.md index 6fb8650..d814f69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,24 @@ for details. ## [Unreleased] +### Added +- **Multivariate derivative estimation** (`jaxsr.derivatives`) — analytic partial + derivatives of a smoothed N-D surface, for problems whose regression needs more + than one partial (PDE-style discovery `u_t = F(u, u_x, u_xx, ...)`, or transform + laws such as `y(x, T) = f(x + s(T))` where `y_T = s'(T)·y_x`): + - `SurfaceDerivatives` — fits a smoother to scattered or gridded data and returns + requested mixed partials with standard errors. Three smoothers: `"tensor_spline"` + (penalized tensor-product B-splines, the default), `"local_poly"` (local + polynomial regression), and `"gp"` (Gaussian process with derivative posterior). + - `estimate_partial_derivatives` — one-call convenience wrapper. + - Smoothing is selected only by criteria blind to the downstream symbolic score + (GCV, log marginal likelihood, or a supplied noise level), and the level actually + used is reported via `smoothing_`, `smoothing_source_`, `effective_dof_`, + `residual_std_`, and `summary()`. `smoothing_scale` re-runs the estimate at a + deliberately different smoothing level to expose smoothing-induced bias. +- Documentation for multivariate derivative estimation: user guide + (`docs/guides/surface-derivatives.md`), API reference page, and skill guide. + ## [0.3.0] - 2026-07-02 ### Added diff --git a/README.md b/README.md index c510c0c..d9b33b2 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ JAXSR is a fully open-source symbolic regression library built on JAX that disco - **JAX-Accelerated**: JIT compilation and GPU support for fast computation - **Symbolic Classification**: Discover interpretable logistic models for binary and multiclass problems via IRLS + sparse selection - **Additive Symbolic Regression**: Boosting-style ensembles of small symbolic expressions — fit residuals stagewise for many simple, interpretable terms (`jaxsr.additive`) +- **Derivative Estimation**: Time derivatives for ODE discovery (`estimate_derivatives`, `discover_dynamics`), and analytic partial derivatives of a smoothed N-D surface for PDE-style or shift-law discovery (`SurfaceDerivatives`) - **Scikit-learn Compatible**: Full estimator protocol (`get_params`/`set_params`/`clone`) — works with `cross_val_score`, `GridSearchCV`, `Pipeline` - **Symbolic Export**: Export to SymPy, LaTeX, or pure Python/NumPy functions diff --git a/docs/_toc.yml b/docs/_toc.yml index 90ba8ff..d3ca09d 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -15,6 +15,7 @@ parts: - file: guides/claude_code_skills - file: guides/performance - file: guides/sklearn-integration + - file: guides/surface-derivatives - caption: Examples chapters: diff --git a/docs/api/derivatives.rst b/docs/api/derivatives.rst new file mode 100644 index 0000000..646e771 --- /dev/null +++ b/docs/api/derivatives.rst @@ -0,0 +1,7 @@ +jaxsr.derivatives +================= + +.. automodule:: jaxsr.derivatives + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/index.rst b/docs/api/index.rst index ba6a6f7..8690731 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -18,3 +18,4 @@ Full API documentation generated from source docstrings. plotting classifier additive + derivatives diff --git a/docs/guides/surface-derivatives.md b/docs/guides/surface-derivatives.md new file mode 100644 index 0000000..8626ae8 --- /dev/null +++ b/docs/guides/surface-derivatives.md @@ -0,0 +1,257 @@ +# Multivariate Derivative Estimation + +Some discovery problems need derivatives *of a surface* rather than of a trajectory. +Two examples: + +- **PDE-style discovery.** The target is `u_t` and the candidate library contains + `u`, `u_x`, `u_xx`, ... — every one of them a partial derivative of the same + measured field. +- **Transform / shift laws.** Time–temperature superposition rests on the identity + + ``` + y(x, T) = f(x + s(T)) => y_T = s'(T) * y_x + ``` + + with `x = log(omega)` and `y = log(G)`. Both partials have to come from **one** + smoothed surface over `(x, T)` before the symbolic stage can see `s'(T)` at all. + +`jaxsr.dynamics.estimate_derivatives` differentiates along a single axis: `X` is +`(n_times, n_states)` and `t` is a 1-D vector. That covers ODE discovery, but not +either case above. `SurfaceDerivatives` fills the gap. + +## Quick start + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +# A surface sampled on a rectangular grid +x = np.linspace(0.0, 2.0, 25) +T = np.linspace(-1.0, 1.0, 20) +xx, TT = np.meshgrid(x, T, indexing="ij") +Y = np.sin(2 * xx) * np.exp(0.5 * TT) + +sigma = 0.01 +Y_obs = Y + np.random.default_rng(0).normal(0, sigma, Y.shape) + +est = SurfaceDerivatives(method="tensor_spline").fit([x, T], Y_obs, sigma=sigma) + +y, dy = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) +y_x, y_T = dy[:, 0], dy[:, 1] + +print(est.summary()) +``` + +``` +SurfaceDerivatives +======================================== +method : tensor_spline +data : 500 points, 2 dimensions +penalty lambda : 0.158489 (chosen by gcv) +effective dof : 97.89 +residual std : 0.00988053 +noise std used : 0.01 +basis per dim : [12, 12] (degree 3) +``` + +Gridded data is passed as a list of axis arrays plus an N-D value array; scattered data +as an `(n_points, n_dims)` coordinate array plus a flat value array. Either way, +`est.coords_` holds the flattened sample locations, which is usually what you want to +evaluate at. + +The one-call form: + +```python +from jaxsr import estimate_partial_derivatives + +y, dy = estimate_partial_derivatives( + [x, T], Y_obs, order=[(1, 0), (0, 1)], method="tensor_spline", sigma=sigma +) +``` + +## Reading the API + +`derivatives()` returns a tuple. The first element is the smoothed surface; the second +is a column per requested order, in the order you asked for them (continuing from the +quick start above): + +```python +y, dy = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) # dy.shape == (n, 2) +y, dy = est.derivatives(est.coords_, order=(2, 0)) # dy.shape == (n, 1) +y, dy, dy_se = est.derivatives(est.coords_, order=[(1, 1)], return_std=True) +``` + +Each order is a tuple with one entry per coordinate: `(1, 0)` is the first partial with +respect to dimension 0, `(0, 1)` the first partial with respect to dimension 1, +`(2, 0)` the second partial in dimension 0, and `(1, 1)` the mixed partial. `sigma` is a +**standard deviation** (scalar or per point), not a variance. + +`predict()` is the order-zero case, with optional standard errors: + +```python +mean, std = est.predict(est.coords_, return_std=True) +``` + +## Choosing a smoother + +```python +SurfaceDerivatives(method="tensor_spline") # default +SurfaceDerivatives(method="local_poly") +SurfaceDerivatives(method="gp") +``` + +| Method | Best for | Cost | Uncertainty | +|--------|----------|------|-------------| +| `"tensor_spline"` | the default; gridded or scattered data in any dimension | fast — one penalized least-squares solve | posterior of the penalized fit | +| `"local_poly"` | irregular sampling, local structure | moderate — one weighted fit per query point | sandwich variance of the local fit | +| `"gp"` | irregular sampling, honest uncertainty, `n` up to a few hundred | cubic in `n`, capped by `max_points` | exact posterior; grows away from the data | + +All three return **analytic** partials of the fitted smoother. None of them finite-differences +the raw data, which is what makes second derivatives usable at all under noise. + +`degree` bounds what you can ask for: for `"tensor_spline"` no single dimension may be +differentiated more than `degree` times; for `"local_poly"` the *total* order may not +exceed `degree`. The default `degree=3` covers `u_xx` and mixed second partials. + +## Choosing the smoothing level + +This is the part that decides whether the downstream symbolic result is trustworthy. + +**Never select the smoothing hyperparameter by the downstream symbolic score.** A +smoother tuned against the regression that consumes it can manufacture whichever law the +regression prefers — the fit looks excellent and the failure is silent. JAXSR therefore +only offers criteria that are blind to the symbolic stage: + +| `smoothing=` | How the level is chosen | +|--------------|-------------------------| +| `"auto"` (default) | generalized cross-validation for `"tensor_spline"` and `"local_poly"`; log marginal likelihood for `"gp"` | +| `"sigma"` | from the noise you supply to `fit(..., sigma=...)`: the level whose residual sum of squares matches `n_points * sigma**2` | +| a float | used verbatim — the penalty `λ` for `"tensor_spline"`, the bandwidth for `"local_poly"`, the noise variance for `"gp"` | + +Replicates are the cleanest source of `sigma`: pool the within-replicate variance and +pass its square root. + +## Reporting and diagnosing smoothing bias + +Smoothing flattens derivatives, and a flatter derivative reports a smaller coefficient. +The bias tracks the noise level and the smoothing level, and is roughly flat in the +amount of data — so more data does not remove it. It is predictable and calibratable, +but only if the smoothing level is visible (continuing from the quick start above): + +```python +est.smoothing_ # λ, bandwidth, or noise variance actually used +est.smoothing_source_ # "gcv", "marginal_likelihood", "sigma", or "fixed" +est.effective_dof_ # effective degrees of freedom of the smoother +est.residual_std_ # residual scatter at the sample points +est.noise_std_ # noise level backing the reported uncertainties +print(est.summary()) # all of the above, formatted +``` + +`smoothing_scale` multiplies the selected level, which turns "how much does my answer +depend on the derivative stage?" into a three-line experiment: + +```python +for scale in (1.0, 3.0, 10.0): + est = SurfaceDerivatives(smoothing_scale=scale).fit([x, T], Y_obs, sigma=sigma) + _, dy = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) + ... # rerun the symbolic stage, record the coefficient +``` + +Report the spread across scales alongside the point estimate. A coefficient that moves +by 10% between `×1` and `×10` is telling you where its error bar really comes from. + +## Worked example: PDE-style discovery + +Recovering the heat equation `u_t = 0.1 * u_xx` from a noisy field: + +```python +import numpy as np +from jaxsr import BasisLibrary, SurfaceDerivatives, SymbolicRegressor + +x = np.linspace(0, 2 * np.pi, 40) +t = np.linspace(0, 1.0, 30) +xx, tt = np.meshgrid(x, t, indexing="ij") +u = np.exp(-0.1 * tt) * np.sin(xx) + 0.5 * np.exp(-0.9 * tt) * np.sin(3 * xx) + +noise = 0.002 +u_obs = u + np.random.default_rng(0).normal(0, noise, u.shape) + +# One smoothed surface -> every partial the library needs +est = SurfaceDerivatives().fit([x, t], u_obs, sigma=noise) +values, d = est.derivatives(est.coords_, order=[(1, 0), (2, 0), (0, 1)]) +u_x, u_xx, u_t = d[:, 0], d[:, 1], d[:, 2] + +library = ( + BasisLibrary(n_features=3, feature_names=["u", "u_x", "u_xx"]) + .add_linear() + .add_interactions(max_order=2) +) +model = SymbolicRegressor(basis_library=library, max_terms=1).fit( + np.column_stack([values, u_x, u_xx]), u_t +) +print(model.expression_) +``` + +``` +y = 0.09583*u_xx +``` + +The right term, with a coefficient about 4% low. Raising `smoothing_scale` to `10.0` +moves it to `0.09043` — the derivative stage, not the symbolic stage, is what sets that +digit. + +## Worked example: a shift law + +For `y(x, T) = f(x + s(T))`, the slope ratio *is* the derivative of the shift law. Here +the shift is Arrhenius, so `s'(T) = E / (R ln(10) T**2)` and the activation energy `E` +falls out of the ratio: + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +# Synthetic time-temperature superposition data +gas_r, energy, t_ref = 8.314e-3, 55.85, 350.0 # kJ/mol/K, kJ/mol, K +x_axis = np.linspace(-2.0, 4.0, 20) # log frequency +T_axis = np.linspace(320.0, 400.0, 12) # temperature, K +xx, TT = np.meshgrid(x_axis, T_axis, indexing="ij") +shift = -(energy / (gas_r * np.log(10.0))) * (1.0 / TT - 1.0 / t_ref) + +sigma = 0.01 +Y_grid = 2.0 + 1.5 * np.tanh(0.8 * (xx + shift - 1.0)) +Y_grid = Y_grid + np.random.default_rng(0).normal(0, sigma, Y_grid.shape) + +# Both partials from one smoothed surface +est = SurfaceDerivatives(smoothing="sigma").fit([x_axis, T_axis], Y_grid, sigma=sigma) +_, d = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) +y_x, y_T = d[:, 0], d[:, 1] + +# The ratio is ill-posed wherever the master curve is flat +keep = np.abs(y_x) > 0.15 * np.abs(y_x).max() +s_prime = y_T[keep] / y_x[keep] +T_keep = est.coords_[keep, 1] + +E_eff = np.median(s_prime * gas_r * np.log(10.0) * T_keep**2) +print(f"E_eff = {E_eff:.2f} kJ/mol") # E_eff = 55.93 kJ/mol (true 55.85) +``` + +`s_prime` versus `T_keep` is also an ordinary symbolic regression problem in its own +right — fit it with a `SymbolicRegressor` over a library containing `1/T**2` instead of +assuming the Arrhenius form, and let the selection decide. + +## Pitfalls + +- **Boundaries.** Every smoother is weakest at the edge of its data. Drop a margin + before reading derivatives, or expect the largest errors there. +- **Ratios of partials.** Mask points where the denominator is near zero, as above. +- **High orders.** Accuracy degrades with each order. Second partials need a well-sampled + surface; third and higher rarely survive realistic noise. +- **GP size.** `method="gp"` is cubic in the number of points and raises above + `max_points=800` rather than hanging. Subsample, or use `"tensor_spline"`. +- **Basis resolution.** The spline uses at most 12 basis functions per dimension by + default, because GCV picks the penalty that is best for the *fit* and tends to + under-smooth derivatives. Pass `n_basis=` when the surface genuinely has finer + structure. +- **Irregular sampling.** Rectangular grids are the best-tested case. For strongly + irregular sampling prefer `"gp"` (or `"local_poly"`), and check the reported + uncertainty rather than assuming it. diff --git a/src/jaxsr/__init__.py b/src/jaxsr/__init__.py index 4442aba..4aadf38 100644 --- a/src/jaxsr/__init__.py +++ b/src/jaxsr/__init__.py @@ -39,6 +39,9 @@ # Classification from .classifier import SymbolicClassifier, fit_symbolic_classification from .constraints import Constraint, Constraints, ConstraintType, build_constraint_scorer + +# Multivariate derivative estimation +from .derivatives import SurfaceDerivatives, estimate_partial_derivatives from .dynamics import DynamicsResult, discover_dynamics, estimate_derivatives # Metrics @@ -201,6 +204,9 @@ def _get_plotting(): "DynamicsResult", "discover_dynamics", "estimate_derivatives", + # Multivariate derivative estimation + "SurfaceDerivatives", + "estimate_partial_derivatives", # Metrics "ModelComparison", "compare_models", diff --git a/src/jaxsr/derivatives.py b/src/jaxsr/derivatives.py new file mode 100644 index 0000000..1847c64 --- /dev/null +++ b/src/jaxsr/derivatives.py @@ -0,0 +1,1535 @@ +""" +Multivariate derivative estimation for surface and PDE-style discovery. + +:func:`jaxsr.dynamics.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 derivative -- PDE-style discovery +(``u_t = F(u, u_x, u_xx, ...)``) or transform discovery 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. + +:class:`SurfaceDerivatives` fits a smoother to scattered or gridded N-D data and +returns *analytic* partial derivatives of that smoother, never finite differences of +noisy raw data. Three smoothers are available: + +``"tensor_spline"`` + Penalized tensor-product B-splines (P-splines). Fast, works in any dimension for + scattered or gridded data, smoothing chosen by GCV or from a known noise level. + +``"local_poly"`` + Local polynomial (LOESS-style) regression. Robust to irregular sampling; the + derivative of order *m* is read off the local polynomial coefficient. + +``"gp"`` + Gaussian process with an anisotropic squared-exponential kernel. Gives derivative + uncertainty directly and handles irregular sampling, at :math:`O(n^3)` cost. + +The smoothing hyperparameter is always selected *without reference to any downstream +symbolic score* -- by GCV, by marginal likelihood, or from a supplied noise level. +Tuning a smoother against the regression that consumes it can manufacture whichever +law the regression prefers, and the failure is silent. The level actually used is +reported in :attr:`SurfaceDerivatives.smoothing_` and by +:meth:`SurfaceDerivatives.summary`, so smoothing-induced bias is visible rather than +inferred. +""" + +from __future__ import annotations + +import itertools +import math +from collections.abc import Sequence +from typing import Any + +import numpy as np + +__all__ = [ + "SurfaceDerivatives", + "estimate_partial_derivatives", +] + +_VALID_METHODS = ("tensor_spline", "local_poly", "gp") +_JITTER = 1e-10 + + +# --------------------------------------------------------------------------- +# Input normalization helpers +# --------------------------------------------------------------------------- + + +def _normalize_coords( + coords: np.ndarray | Sequence[np.ndarray], + values: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """ + Normalize scattered or gridded inputs to flat ``(n, d)`` / ``(n,)`` arrays. + + Parameters + ---------- + coords : np.ndarray of shape (n_points, n_dims), or sequence of 1-D arrays + Sample locations. A sequence of ``d`` 1-D axis arrays is interpreted as a + rectangular grid, in which case *values* must have shape + ``(len(axis_0), ..., len(axis_{d-1}))``. + values : np.ndarray + Observed values, shape ``(n_points,)`` for scattered input or the grid shape + for gridded input. + + Returns + ------- + coords : np.ndarray of shape (n_points, n_dims) + Flattened coordinates. + values : np.ndarray of shape (n_points,) + Flattened values. + + Raises + ------ + ValueError + If shapes are inconsistent or the inputs are not finite. + """ + values = np.asarray(values, dtype=np.float64) + + # Gridded form: a sequence of d axis arrays paired with a d-dimensional values array. + is_grid = ( + isinstance(coords, (list, tuple)) + and len(coords) > 0 + and values.ndim == len(coords) + and all(np.ndim(axis) == 1 for axis in coords) + ) + + if is_grid: + axes = [np.asarray(a, dtype=np.float64).ravel() for a in coords] + shape = tuple(len(a) for a in axes) + if values.shape != shape: + raise ValueError( + f"Gridded input: values shape {values.shape} does not match the grid " + f"shape implied by the axes {shape}" + ) + mesh = np.meshgrid(*axes, indexing="ij") + coords_arr = np.column_stack([m.ravel() for m in mesh]) + values = values.ravel() + else: + coords_arr = np.asarray(coords, dtype=np.float64) + if coords_arr.ndim == 1: + coords_arr = coords_arr.reshape(-1, 1) + if coords_arr.ndim != 2: + raise ValueError(f"coords must be 1-D or 2-D, got {coords_arr.ndim}-D") + values = values.ravel() + + if coords_arr.shape[0] != values.shape[0]: + raise ValueError( + f"Number of coordinate rows ({coords_arr.shape[0]}) must match the number " + f"of values ({values.shape[0]})" + ) + if not np.all(np.isfinite(coords_arr)): + raise ValueError("coords contains non-finite values") + if not np.all(np.isfinite(values)): + raise ValueError("values contains non-finite entries") + + return coords_arr, values + + +def _normalize_query(coords: np.ndarray, n_dims: int) -> np.ndarray: + """ + Normalize query coordinates to a ``(n_query, n_dims)`` array. + + Parameters + ---------- + coords : np.ndarray of shape (n_query, n_dims) + Query locations. A 1-D array is read as a column of points when the surface is + 1-D, and as a single point otherwise. + n_dims : int + Expected number of dimensions. + + Returns + ------- + np.ndarray of shape (n_query, n_dims) + Query coordinates. + + Raises + ------ + ValueError + If the coordinate dimension does not match *n_dims* or entries are not finite. + """ + query = np.asarray(coords, dtype=np.float64) + if query.ndim == 1: + if n_dims == 1: + query = query.reshape(-1, 1) + elif query.shape[0] == n_dims: + query = query.reshape(1, n_dims) + + if query.ndim != 2 or query.shape[1] != n_dims: + raise ValueError( + f"Query coordinates must have shape (n_query, {n_dims}), got {query.shape}" + ) + if not np.all(np.isfinite(query)): + raise ValueError("Query coordinates contain non-finite values") + return query + + +def _normalize_orders(order: Any, n_dims: int) -> np.ndarray: + """ + Normalize a derivative-order specification to a ``(n_orders, n_dims)`` int array. + + Parameters + ---------- + order : sequence + Either a single order tuple such as ``(1, 0)`` or a sequence of them such as + ``[(1, 0), (0, 1)]``. Each tuple gives the differentiation order per dimension. + n_dims : int + Number of coordinate dimensions. + + Returns + ------- + np.ndarray of shape (n_orders, n_dims) + Non-negative integer derivative orders. + + Raises + ------ + ValueError + If the specification is malformed, has the wrong length, or is negative. + """ + if order is None: + raise ValueError("order must be provided, e.g. order=[(1, 0), (0, 1)]") + + seq = list(order) if isinstance(order, (list, tuple, np.ndarray)) else [order] + if len(seq) == 0: + raise ValueError("order must contain at least one derivative order") + + if all(np.ndim(item) == 0 for item in seq): + orders = [seq] + else: + orders = [list(item) for item in seq] + + arr = np.asarray(orders) + if arr.ndim != 2 or arr.shape[1] != n_dims: + raise ValueError( + f"Each derivative order must have {n_dims} entries (one per dimension); " + f"got {arr.shape}" + ) + if not np.issubdtype(arr.dtype, np.integer): + if np.any(arr != np.round(arr)): + raise ValueError("Derivative orders must be integers") + arr = np.round(arr).astype(int) + if np.any(arr < 0): + raise ValueError("Derivative orders must be non-negative") + return arr.astype(int) + + +def _normalize_sigma(sigma: float | np.ndarray | None, n_points: int) -> np.ndarray | None: + """ + Normalize a noise specification to a per-point standard-deviation array. + + Parameters + ---------- + sigma : float, np.ndarray of shape (n_points,), or None + Measurement noise standard deviation, scalar or per point. + n_points : int + Number of data points. + + Returns + ------- + np.ndarray of shape (n_points,) or None + Per-point standard deviations, or ``None`` if *sigma* was ``None``. + + Raises + ------ + ValueError + If *sigma* is non-positive, non-finite, or the wrong length. + """ + if sigma is None: + return None + arr = np.asarray(sigma, dtype=np.float64) + if arr.ndim == 0: + arr = np.full(n_points, float(arr)) + else: + arr = arr.ravel() + if arr.shape[0] != n_points: + raise ValueError(f"sigma must be scalar or have {n_points} entries, got {arr.shape[0]}") + if not np.all(np.isfinite(arr)) or np.any(arr <= 0): + raise ValueError("sigma must be finite and strictly positive") + return arr + + +# --------------------------------------------------------------------------- +# B-spline helpers (tensor_spline) +# --------------------------------------------------------------------------- + + +def _knot_vector(lo: float, hi: float, n_basis: int, degree: int) -> np.ndarray: + """ + Build a clamped, uniformly spaced knot vector. + + Parameters + ---------- + lo, hi : float + Lower and upper bounds of the data along this dimension. + n_basis : int + Number of B-spline basis functions. + degree : int + Spline degree. + + Returns + ------- + np.ndarray + Knot vector of length ``n_basis + degree + 1``. + + Raises + ------ + ValueError + If *n_basis* is smaller than ``degree + 1``. + """ + if n_basis < degree + 1: + raise ValueError(f"n_basis ({n_basis}) must be at least degree + 1 ({degree + 1})") + if hi <= lo: + hi = lo + 1.0 + n_interior = n_basis - degree - 1 + interior = np.linspace(lo, hi, n_interior + 2)[1:-1] + return np.concatenate([np.full(degree + 1, lo), interior, np.full(degree + 1, hi)]) + + +def _bspline_basis(x: np.ndarray, knots: np.ndarray, degree: int, nu: int = 0) -> np.ndarray: + """ + Evaluate all B-spline basis functions (or their derivatives) at *x*. + + Parameters + ---------- + x : np.ndarray of shape (n_points,) + Evaluation points. + knots : np.ndarray + Knot vector. + degree : int + Spline degree. + nu : int + Derivative order to evaluate. + + Returns + ------- + np.ndarray of shape (n_points, n_basis) + Basis (or basis-derivative) values. + """ + from scipy.interpolate import BSpline + + n_basis = len(knots) - degree - 1 + out = np.zeros((x.shape[0], n_basis)) + if nu > degree: + return out + coef = np.zeros(n_basis) + for j in range(n_basis): + coef[:] = 0.0 + coef[j] = 1.0 + out[:, j] = BSpline(knots, coef, degree, extrapolate=True)(x, nu) + return out + + +def _row_tensor(bases: list[np.ndarray]) -> np.ndarray: + """ + Row-wise tensor (Khatri-Rao) product of per-dimension basis matrices. + + Parameters + ---------- + bases : list of np.ndarray + One ``(n_points, m_i)`` matrix per dimension. + + Returns + ------- + np.ndarray of shape (n_points, prod(m_i)) + Tensor-product design matrix, with dimension 0 varying slowest. + """ + out = bases[0] + for basis in bases[1:]: + out = (out[:, :, None] * basis[:, None, :]).reshape(basis.shape[0], -1) + return out + + +def _difference_penalty(sizes: Sequence[int], penalty_order: int) -> np.ndarray: + """ + Build the additive tensor-product difference penalty matrix. + + Parameters + ---------- + sizes : sequence of int + Number of basis functions per dimension. + penalty_order : int + Order of the difference penalty applied along each dimension. + + Returns + ------- + np.ndarray of shape (prod(sizes), prod(sizes)) + Sum over dimensions of ``I (x) ... (x) D.T @ D (x) ... (x) I``. + """ + total = int(np.prod(sizes)) + penalty = np.zeros((total, total)) + for axis, size in enumerate(sizes): + order = min(penalty_order, size - 1) + if order <= 0: + block = np.eye(size) + else: + diff = np.diff(np.eye(size), n=order, axis=0) + block = diff.T @ diff + term = np.array([[1.0]]) + for other, other_size in enumerate(sizes): + term = np.kron(term, block if other == axis else np.eye(other_size)) + penalty += term + return penalty + + +# --------------------------------------------------------------------------- +# Gaussian process helpers +# --------------------------------------------------------------------------- + + +def _hermite_e(n: int, u: np.ndarray) -> np.ndarray: + """ + Evaluate the probabilists' Hermite polynomial ``He_n``. + + Parameters + ---------- + n : int + Polynomial order. + u : np.ndarray + Evaluation points. + + Returns + ------- + np.ndarray + ``He_n(u)``, same shape as *u*. + """ + coef = np.zeros(n + 1) + coef[n] = 1.0 + return np.polynomial.hermite_e.hermeval(u, coef) + + +def _hermite_e_at_zero(n: int) -> float: + """ + Evaluate ``He_n(0)``. + + Parameters + ---------- + n : int + Polynomial order. + + Returns + ------- + float + ``He_n(0)``: zero for odd *n*, ``(-1)^(n/2) * (n-1)!!`` otherwise. + """ + if n % 2 == 1: + return 0.0 + return float(_hermite_e(n, np.zeros(1))[0]) + + +# --------------------------------------------------------------------------- +# Main estimator +# --------------------------------------------------------------------------- + + +class SurfaceDerivatives: + """ + Smoothed N-D surface with analytic partial derivatives. + + Fits a smoother to scattered or gridded data over ``n_dims`` coordinates and + evaluates arbitrary mixed partial derivatives of that smoother analytically, + together with their standard errors. + + Parameters + ---------- + method : str + Smoother to fit. One of ``"tensor_spline"`` (penalized tensor-product + B-splines), ``"local_poly"`` (local polynomial regression), or ``"gp"`` + (Gaussian process with an anisotropic squared-exponential kernel). + degree : int + Spline degree for ``"tensor_spline"``, or the local polynomial degree for + ``"local_poly"``. Must be at least the highest total derivative order + requested. Ignored by ``"gp"``. + n_basis : int or sequence of int, optional + Number of B-spline basis functions per dimension (``"tensor_spline"`` only). + Defaults to a value derived from the number of distinct coordinates per + dimension, capped so the design matrix stays well determined. + penalty_order : int + Order of the difference penalty on the spline coefficients + (``"tensor_spline"`` only). ``2`` penalizes curvature. + smoothing : float or str + How much to smooth, and how that level is chosen: + + - ``"auto"`` (default): GCV for ``"tensor_spline"`` and ``"local_poly"``, + marginal likelihood for ``"gp"``. + - ``"sigma"``: requires *sigma* at :meth:`fit`; chooses the smoothing level + whose weighted residual sum of squares equals ``n_points`` (equivalently, + unweighted residual sum of squares equal to ``n_points * sigma**2``). + - float: use this value directly -- the ridge parameter ``lambda`` for + ``"tensor_spline"``, the bandwidth for ``"local_poly"`` (in standardized + coordinate units), or the noise variance for ``"gp"``. + + Never selected against a downstream regression score. + smoothing_scale : float + Multiplier applied to the selected smoothing level. Useful for deliberately + over- or under-smoothing to expose the sensitivity of a discovered law to the + derivative stage. + length_scale : float or sequence of float, optional + Fixed kernel length scales for ``"gp"`` in standardized coordinate units. If + ``None`` they are learned by maximizing the log marginal likelihood. + max_basis : int + Guard on the total number of tensor-product spline basis functions. + max_points : int + Guard on the number of training points for ``"gp"``, whose cost is cubic. + random_state : int, optional + Seed for the subsampling used during ``"local_poly"`` bandwidth selection. + + Attributes + ---------- + coords_ : np.ndarray of shape (n_points, n_dims) + Training coordinates. + values_ : np.ndarray of shape (n_points,) + Training values. + n_features_in_ : int + Number of coordinate dimensions. + smoothing_ : float + Smoothing level actually used (``lambda``, bandwidth, or noise variance, + depending on *method*). + smoothing_source_ : str + How :attr:`smoothing_` was chosen: ``"gcv"``, ``"marginal_likelihood"``, + ``"sigma"``, or ``"fixed"``. + effective_dof_ : float + Effective degrees of freedom of the fitted smoother. + residual_std_ : float + Residual standard deviation of the fit at the training points. + noise_std_ : float + Noise level used for the reported uncertainties: the supplied *sigma* (its + root-mean-square if per-point) when given, otherwise :attr:`residual_std_`. + + Raises + ------ + ValueError + If *method*, *degree*, *penalty_order*, *smoothing*, or *smoothing_scale* is + invalid. + + See Also + -------- + jaxsr.dynamics.estimate_derivatives : Derivatives along a single axis, for + time-series / ODE discovery. + + Examples + -------- + >>> import numpy as np + >>> from jaxsr import SurfaceDerivatives + >>> x = np.linspace(0, 1, 25) + >>> T = np.linspace(0, 2, 15) + >>> XX, TT = np.meshgrid(x, T, indexing="ij") + >>> Z = np.sin(XX) * TT**2 + >>> est = SurfaceDerivatives(method="tensor_spline").fit([x, T], Z) + >>> coords = np.column_stack([XX.ravel(), TT.ravel()]) + >>> z, dz = est.derivatives(coords, order=[(1, 0), (0, 1)]) + >>> dz.shape + (375, 2) + """ + + def __init__( + self, + method: str = "tensor_spline", + *, + degree: int = 3, + n_basis: int | Sequence[int] | None = None, + penalty_order: int = 2, + smoothing: float | str = "auto", + smoothing_scale: float = 1.0, + length_scale: float | Sequence[float] | None = None, + max_basis: int = 512, + max_points: int = 800, + random_state: int | None = None, + ) -> None: + if method not in _VALID_METHODS: + raise ValueError(f"Unknown method {method!r}. Choose from {list(_VALID_METHODS)}") + if not isinstance(degree, (int, np.integer)) or degree < 1: + raise ValueError(f"degree must be a positive integer, got {degree!r}") + if not isinstance(penalty_order, (int, np.integer)) or penalty_order < 0: + raise ValueError(f"penalty_order must be a non-negative integer, got {penalty_order!r}") + if isinstance(smoothing, str): + if smoothing not in ("auto", "sigma"): + raise ValueError( + f"smoothing must be 'auto', 'sigma', or a float, got {smoothing!r}" + ) + else: + smoothing = float(smoothing) + if smoothing <= 0: + raise ValueError("Numeric smoothing must be strictly positive") + if smoothing_scale <= 0: + raise ValueError("smoothing_scale must be strictly positive") + if max_basis < 1: + raise ValueError("max_basis must be a positive integer") + if max_points < 1: + raise ValueError("max_points must be a positive integer") + + self.method = method + self.degree = int(degree) + self.n_basis = n_basis + self.penalty_order = int(penalty_order) + self.smoothing = smoothing + self.smoothing_scale = float(smoothing_scale) + self.length_scale = length_scale + self.max_basis = int(max_basis) + self.max_points = int(max_points) + self.random_state = random_state + + self._is_fitted = False + + # -- public API --------------------------------------------------------- + + def fit( + self, + coords: np.ndarray | Sequence[np.ndarray], + values: np.ndarray, + sigma: float | np.ndarray | None = None, + ) -> SurfaceDerivatives: + """ + Fit the smoother to surface data. + + Parameters + ---------- + coords : np.ndarray of shape (n_points, n_dims), or sequence of 1-D arrays + Sample locations. A sequence of ``n_dims`` 1-D axis arrays is treated as a + rectangular grid, in which case *values* must have the grid shape. + values : np.ndarray + Observed values, shape ``(n_points,)`` for scattered coordinates or the + grid shape for gridded coordinates. + sigma : float or np.ndarray of shape (n_points,), optional + Known measurement noise standard deviation, scalar or per point. Supplying + it enables ``smoothing="sigma"`` and makes the reported derivative + uncertainty reflect the measurement noise rather than the residual scatter. + + Returns + ------- + SurfaceDerivatives + The fitted estimator (``self``). + + Raises + ------ + ValueError + If the inputs are inconsistent, too small for the requested smoother, or + if ``smoothing="sigma"`` was requested without *sigma*. + """ + coords_arr, values_arr = _normalize_coords(coords, values) + n_points, n_dims = coords_arr.shape + + if n_points < 4: + raise ValueError(f"Need at least 4 data points, got {n_points}") + + sigma_arr = _normalize_sigma(sigma, n_points) + if self.smoothing == "sigma" and sigma_arr is None: + raise ValueError("smoothing='sigma' requires the sigma argument to fit()") + + self.coords_ = coords_arr + self.values_ = values_arr + self.n_features_in_ = n_dims + self._sigma = sigma_arr + self._weights = np.ones(n_points) if sigma_arr is None else 1.0 / sigma_arr**2 + + if self.method == "tensor_spline": + self._fit_tensor_spline() + elif self.method == "local_poly": + self._fit_local_poly() + else: + self._fit_gp() + + residuals = self.values_ - self._fitted_values + dof = max(n_points - self.effective_dof_, 1.0) + self.residual_std_ = float(np.sqrt(np.sum(residuals**2) / dof)) + if sigma_arr is None: + self.noise_std_ = self.residual_std_ + else: + self.noise_std_ = float(np.sqrt(np.mean(sigma_arr**2))) + + self._is_fitted = True + return self + + def predict( + self, + coords: np.ndarray, + return_std: bool = False, + ) -> np.ndarray | tuple[np.ndarray, np.ndarray]: + """ + Evaluate the smoothed surface. + + Parameters + ---------- + coords : np.ndarray of shape (n_query, n_dims) + Query locations. + return_std : bool + If ``True``, also return the standard error of the smoothed value. + + Returns + ------- + values : np.ndarray of shape (n_query,) + Smoothed values. + std : np.ndarray of shape (n_query,) + Standard errors. Only returned when *return_std* is ``True``. + + Raises + ------ + RuntimeError + If the estimator has not been fitted. + """ + if not self._is_fitted: + raise RuntimeError("SurfaceDerivatives must be fitted before calling predict()") + zero = tuple([0] * self.n_features_in_) + out = self.derivatives(coords, order=[zero], return_std=return_std) + if return_std: + values, _, std = out + return values, std[:, 0] + return out[0] + + def derivatives( + self, + coords: np.ndarray, + order: Any, + return_std: bool = False, + ) -> tuple[np.ndarray, ...]: + """ + Evaluate analytic partial derivatives of the fitted smoother. + + Parameters + ---------- + coords : np.ndarray of shape (n_query, n_dims) + Query locations. Use ``estimator.coords_`` to evaluate at the sample + locations, including for data supplied in gridded form. + order : tuple of int, or sequence of tuples + Derivative orders per dimension. ``(1, 0)`` is a single first partial with + respect to dimension 0; ``[(1, 0), (0, 1)]`` requests both first partials. + return_std : bool + If ``True``, also return the standard error of each partial derivative. + + Returns + ------- + values : np.ndarray of shape (n_query,) + The smoothed surface at the query points. + partials : np.ndarray of shape (n_query, n_orders) + One column per entry of *order*, in the given order. A single order tuple + still yields a column vector of shape ``(n_query, 1)``. + std : np.ndarray of shape (n_query, n_orders) + Standard errors of *partials*. Only returned when *return_std* is ``True``. + + Raises + ------ + RuntimeError + If the estimator has not been fitted. + ValueError + If *order* is malformed or exceeds what the smoother can differentiate. + """ + if not self._is_fitted: + raise RuntimeError("SurfaceDerivatives must be fitted before calling derivatives()") + + query = _normalize_query(coords, self.n_features_in_) + orders = _normalize_orders(order, self.n_features_in_) + self._validate_orders(orders) + + zero = np.zeros((1, self.n_features_in_), dtype=int) + all_orders = np.vstack([zero, orders]) + + if self.method == "tensor_spline": + vals, stds = self._eval_tensor_spline(query, all_orders, return_std) + elif self.method == "local_poly": + vals, stds = self._eval_local_poly(query, all_orders, return_std) + else: + vals, stds = self._eval_gp(query, all_orders, return_std) + + values = vals[:, 0] + partials = vals[:, 1:] + if return_std: + return values, partials, stds[:, 1:] + return values, partials + + def summary(self) -> str: + """ + Return a human-readable description of the fitted smoother. + + Reports the smoothing level actually used and how it was chosen, so that + smoothing-induced bias in downstream results is visible rather than inferred. + + Returns + ------- + str + Multi-line summary. + + Raises + ------ + RuntimeError + If the estimator has not been fitted. + """ + if not self._is_fitted: + raise RuntimeError("SurfaceDerivatives must be fitted before calling summary()") + + label = { + "tensor_spline": "penalty lambda", + "local_poly": "bandwidth", + "gp": "noise variance", + }[self.method] + + lines = [ + "SurfaceDerivatives", + "=" * 40, + f"method : {self.method}", + f"data : {self.coords_.shape[0]} points, " + f"{self.n_features_in_} dimensions", + f"{label:<18}: {self.smoothing_:.6g} (chosen by {self.smoothing_source_})", + f"effective dof : {self.effective_dof_:.2f}", + f"residual std : {self.residual_std_:.6g}", + f"noise std used : {self.noise_std_:.6g}", + ] + if self.method == "gp": + scales = ", ".join(f"{s:.4g}" for s in self._gp_length_scale) + lines.append(f"length scales : [{scales}] (standardized units)") + elif self.method == "tensor_spline": + lines.append(f"basis per dim : {list(self._spline_sizes)} (degree {self.degree})") + return "\n".join(lines) + + # -- validation --------------------------------------------------------- + + def _validate_orders(self, orders: np.ndarray) -> None: + """ + Check that requested derivative orders are supported by the smoother. + + Parameters + ---------- + orders : np.ndarray of shape (n_orders, n_dims) + Requested derivative orders. + + Raises + ------ + ValueError + If an order exceeds what the fitted smoother can differentiate. + """ + if self.method == "tensor_spline": + worst = int(orders.max(initial=0)) + if worst > self.degree: + raise ValueError( + f"Derivative order {worst} exceeds the spline degree {self.degree}; " + f"refit with degree >= {worst}" + ) + elif self.method == "local_poly": + worst = int(orders.sum(axis=1).max(initial=0)) + if worst > self.degree: + raise ValueError( + f"Total derivative order {worst} exceeds the local polynomial degree " + f"{self.degree}; refit with degree >= {worst}" + ) + + # -- tensor-product P-splines ------------------------------------------ + + def _default_basis_sizes(self) -> list[int]: + """ + Choose a per-dimension basis size that stays well determined. + + Returns + ------- + list of int + Number of B-spline basis functions per dimension. + + Raises + ------ + ValueError + If an explicitly supplied *n_basis* is invalid or too large. + """ + n_points, n_dims = self.coords_.shape + floor = self.degree + 1 + + if self.n_basis is not None: + if np.ndim(self.n_basis) == 0: + sizes = [int(self.n_basis)] * n_dims + else: + sizes = [int(v) for v in self.n_basis] # type: ignore[union-attr] + if len(sizes) != n_dims: + raise ValueError(f"n_basis must be a scalar or have {n_dims} entries") + if any(s < floor for s in sizes): + raise ValueError(f"Each n_basis entry must be at least degree + 1 = {floor}") + total = int(np.prod(sizes)) + if total > self.max_basis: + raise ValueError( + f"n_basis={sizes} needs {total} tensor-product basis functions, " + f"above max_basis={self.max_basis}" + ) + if total > n_points: + raise ValueError( + f"n_basis={sizes} needs {total} basis functions but only {n_points} " + "data points are available" + ) + return sizes + + # Cap the per-dimension resolution: GCV picks the penalty that is best for the + # *fit*, which tends to under-smooth derivatives, so a moderate basis is a + # useful second brake. Pass n_basis explicitly for finer structure. + cap = max(floor, 12) + sizes = [] + for dim in range(n_dims): + n_unique = len(np.unique(self.coords_[:, dim])) + sizes.append(int(np.clip(n_unique, floor, cap))) + + budget = min(self.max_basis, max(n_points // 2, floor**n_dims)) + while int(np.prod(sizes)) > budget and any(s > floor for s in sizes): + largest = int(np.argmax(sizes)) + if sizes[largest] <= floor: + break + sizes[largest] -= 1 + if int(np.prod(sizes)) > self.max_basis: + raise ValueError( + f"Minimal tensor-product basis needs {int(np.prod(sizes))} functions, above " + f"max_basis={self.max_basis}; reduce degree or raise max_basis" + ) + return sizes + + def _spline_design(self, coords: np.ndarray, order: np.ndarray) -> np.ndarray: + """ + Build the tensor-product design matrix for one derivative order. + + Parameters + ---------- + coords : np.ndarray of shape (n_points, n_dims) + Evaluation points. + order : np.ndarray of shape (n_dims,) + Derivative order per dimension. + + Returns + ------- + np.ndarray of shape (n_points, n_basis_total) + Design matrix. + """ + bases = [ + _bspline_basis(coords[:, dim], self._spline_knots[dim], self.degree, int(order[dim])) + for dim in range(self.n_features_in_) + ] + return _row_tensor(bases) + + def _fit_tensor_spline(self) -> None: + """ + Fit penalized tensor-product B-splines and select the penalty weight. + + Raises + ------ + ValueError + If the basis specification is inconsistent with the data. + """ + from scipy.linalg import cho_factor, cho_solve + + n_points = self.coords_.shape[0] + sizes = self._default_basis_sizes() + self._spline_sizes = sizes + self._spline_knots = [ + _knot_vector( + float(self.coords_[:, dim].min()), + float(self.coords_[:, dim].max()), + sizes[dim], + self.degree, + ) + for dim in range(self.n_features_in_) + ] + + zero_order = np.zeros(self.n_features_in_, dtype=int) + design = self._spline_design(self.coords_, zero_order) + sqrt_w = np.sqrt(self._weights) + design_w = design * sqrt_w[:, None] + values_w = self.values_ * sqrt_w + + gram = design_w.T @ design_w + rhs = design_w.T @ values_w + penalty = _difference_penalty(sizes, self.penalty_order) + + # Scale the penalty so that lambda is dimensionless and O(1) grids work. + scale = np.trace(gram) / max(np.trace(penalty), _JITTER) + penalty = penalty * scale + ridge = _JITTER * np.trace(gram) / gram.shape[0] * np.eye(gram.shape[0]) + + def solve_for(lam: float) -> tuple[np.ndarray, float, float]: + matrix = gram + lam * penalty + ridge + factor = cho_factor(matrix, lower=True) + coef = cho_solve(factor, rhs) + resid = values_w - design_w @ coef + rss = float(resid @ resid) + edf = float(np.trace(cho_solve(factor, gram))) + return coef, rss, edf + + lam_grid = np.geomspace(1e-8, 1e8, 21) + results = [solve_for(float(lam)) for lam in lam_grid] + rss_grid = np.array([r[1] for r in results]) + edf_grid = np.array([r[2] for r in results]) + + if isinstance(self.smoothing, float): + lam = self.smoothing + source = "fixed" + elif self.smoothing == "sigma": + lam = float(_interp_log(lam_grid, rss_grid, float(n_points))) + source = "sigma" + else: + denom = np.maximum(n_points - edf_grid, 1e-6) + gcv = n_points * rss_grid / denom**2 + gcv[edf_grid >= n_points] = np.inf + lam = float(lam_grid[int(np.argmin(gcv))]) + source = "gcv" + + lam *= self.smoothing_scale + coef, _, edf = solve_for(lam) + + self.smoothing_ = lam + self.smoothing_source_ = source + self.effective_dof_ = edf + self._spline_coef = coef + self._fitted_values = design @ coef + + matrix = gram + lam * penalty + ridge + factor = cho_factor(matrix, lower=True) + self._spline_cov = cho_solve(factor, np.eye(matrix.shape[0])) + + def _eval_tensor_spline( + self, query: np.ndarray, orders: np.ndarray, return_std: bool + ) -> tuple[np.ndarray, np.ndarray]: + """ + Evaluate spline partials at the query points. + + Parameters + ---------- + query : np.ndarray of shape (n_query, n_dims) + Query points. + orders : np.ndarray of shape (n_orders, n_dims) + Derivative orders. + return_std : bool + Whether to compute standard errors. + + Returns + ------- + values : np.ndarray of shape (n_query, n_orders) + Evaluated partials. + std : np.ndarray of shape (n_query, n_orders) + Standard errors, left as zeros when *return_std* is ``False``. + """ + scale = 1.0 if self._sigma is not None else self.residual_std_**2 + values = np.empty((query.shape[0], orders.shape[0])) + stds = np.zeros_like(values) + + for k, order in enumerate(orders): + design = self._spline_design(query, order) + values[:, k] = design @ self._spline_coef + if return_std: + var = np.einsum("ij,jk,ik->i", design, self._spline_cov, design) + stds[:, k] = np.sqrt(np.maximum(scale * var, 0.0)) + return values, stds + + # -- local polynomial regression --------------------------------------- + + def _poly_exponents(self) -> np.ndarray: + """ + Enumerate monomial exponents up to the local polynomial degree. + + Returns + ------- + np.ndarray of shape (n_terms, n_dims) + Exponent vectors with total degree at most :attr:`degree`. + """ + dims = self.n_features_in_ + combos = [ + exps + for exps in itertools.product(range(self.degree + 1), repeat=dims) + if sum(exps) <= self.degree + ] + combos.sort(key=lambda e: (sum(e), e)) + return np.asarray(combos, dtype=int) + + def _local_solve( + self, point: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """ + Fit the weighted local polynomial around one query point. + + Parameters + ---------- + point : np.ndarray of shape (n_dims,) + Query point in standardized coordinates. + + Returns + ------- + coef : np.ndarray of shape (n_terms,) + Local polynomial coefficients in centered, standardized coordinates. + var_diag : np.ndarray of shape (n_terms,) + Sandwich variance of each coefficient, in units of the noise variance. + value_weights : np.ndarray of shape (n_neighbors,) + Linear smoother weights producing the fitted value at *point*. + neighbors : np.ndarray of shape (n_neighbors,) + Indices of the training points used, aligned with *value_weights*. + """ + idx = np.asarray(self._tree.query_ball_point(point, self._bandwidth), dtype=int) + if idx.size < self._min_points: + _, idx = self._tree.query(point, k=self._min_points) + idx = np.atleast_1d(np.asarray(idx, dtype=int)) + + local = self._coords_z[idx] - point + radius = float(np.max(np.linalg.norm(local, axis=1))) + radius = max(radius, _JITTER) + dist = np.linalg.norm(local, axis=1) / radius + kernel = np.clip(1.0 - np.clip(dist, 0.0, 1.0) ** 3, 0.0, None) ** 3 + kernel = np.maximum(kernel, 1e-6) + weights = kernel * self._weights[idx] + + design = np.prod(local[:, None, :] ** self._exponents[None, :, :], axis=2) + gram = design.T @ (weights[:, None] * design) + gram_pinv = np.linalg.pinv(gram) + smoother = gram_pinv @ (design.T * weights) + coef = smoother @ self.values_[idx] + + noise = 1.0 / self._weights[idx] if self._sigma is not None else np.ones(idx.size) + var_diag = np.einsum("ij,j,ij->i", smoother, noise, smoother) + return coef, var_diag, smoother[0], idx + + def _local_pass(self, points: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ + Evaluate the local fit at a set of training points. + + Parameters + ---------- + points : np.ndarray of shape (n_eval,) + Indices into the training set. + + Returns + ------- + fitted : np.ndarray of shape (n_eval,) + Fitted values. + hat : np.ndarray of shape (n_eval,) + Diagonal of the smoother matrix at those points. + """ + fitted = np.empty(points.shape[0]) + hat = np.empty(points.shape[0]) + for j, i in enumerate(points): + coef, _, value_weights, neighbors = self._local_solve(self._coords_z[i]) + fitted[j] = coef[0] + match = np.flatnonzero(neighbors == i) + hat[j] = float(value_weights[match[0]]) if match.size else 0.0 + return fitted, hat + + def _fit_local_poly(self) -> None: + """ + Fit local polynomial regression and select the bandwidth. + + Raises + ------ + ValueError + If there are too few points to support the requested local degree. + """ + from scipy.spatial import cKDTree + + n_points = self.coords_.shape[0] + self._center = self.coords_.mean(axis=0) + spread = self.coords_.std(axis=0) + spread[spread <= 0] = 1.0 + self._scale = spread + self._coords_z = (self.coords_ - self._center) / self._scale + + self._exponents = self._poly_exponents() + n_terms = self._exponents.shape[0] + if n_points < n_terms + 1: + raise ValueError( + f"Local polynomial of degree {self.degree} in {self.n_features_in_} dimensions " + f"needs more than {n_terms} points, got {n_points}" + ) + self._min_points = min(2 * n_terms, n_points) + self._tree = cKDTree(self._coords_z) + + knn_dist, _ = self._tree.query(self._coords_z, k=self._min_points) + knn_dist = np.atleast_2d(knn_dist)[:, -1] + h_min = float(np.percentile(knn_dist, 75)) + diameter = float(np.linalg.norm(self._coords_z.max(axis=0) - self._coords_z.min(axis=0))) + h_max = max(diameter, 2 * h_min) + h_min = max(h_min, _JITTER) + + rng = np.random.default_rng(self.random_state) + n_probe = min(120, n_points) + probe = rng.choice(n_points, size=n_probe, replace=False) + + if isinstance(self.smoothing, float): + bandwidth = self.smoothing + source = "fixed" + else: + grid = np.geomspace(h_min, h_max, 10) + rss_grid = np.empty(grid.shape[0]) + score_grid = np.empty(grid.shape[0]) + for i, h in enumerate(grid): + self._bandwidth = float(h) + fitted, hat = self._local_pass(probe) + resid = self.values_[probe] - fitted + weighted = float(np.sum(self._weights[probe] * resid**2)) * n_points / n_probe + rss_grid[i] = weighted + denom = max(1.0 - float(np.mean(hat)), 1e-6) + score_grid[i] = float(np.mean(resid**2)) / denom**2 + if self.smoothing == "sigma": + bandwidth = float(_interp_log(grid, rss_grid, float(n_points))) + source = "sigma" + else: + bandwidth = float(grid[int(np.argmin(score_grid))]) + source = "gcv" + + self._bandwidth = bandwidth * self.smoothing_scale + self.smoothing_ = self._bandwidth + self.smoothing_source_ = source + + fitted, hat = self._local_pass(np.arange(n_points)) + self._fitted_values = fitted + self.effective_dof_ = float(np.sum(hat)) + + def _eval_local_poly( + self, query: np.ndarray, orders: np.ndarray, return_std: bool + ) -> tuple[np.ndarray, np.ndarray]: + """ + Evaluate local polynomial partials at the query points. + + Parameters + ---------- + query : np.ndarray of shape (n_query, n_dims) + Query points. + orders : np.ndarray of shape (n_orders, n_dims) + Derivative orders. + return_std : bool + Whether to compute standard errors. + + Returns + ------- + values : np.ndarray of shape (n_query, n_orders) + Evaluated partials. + std : np.ndarray of shape (n_query, n_orders) + Standard errors, left as zeros when *return_std* is ``False``. + """ + query_z = (query - self._center) / self._scale + n_orders = orders.shape[0] + + # Map each requested order to its monomial column and derivative factor. + columns = np.empty(n_orders, dtype=int) + factors = np.empty(n_orders) + for k, order in enumerate(orders): + match = np.flatnonzero(np.all(self._exponents == order[None, :], axis=1)) + columns[k] = int(match[0]) + factors[k] = np.prod([math.factorial(int(o)) for o in order]) / np.prod( + self._scale ** order.astype(float) + ) + + scale = 1.0 if self._sigma is not None else self.residual_std_**2 + values = np.empty((query.shape[0], n_orders)) + stds = np.zeros_like(values) + + for i in range(query.shape[0]): + coef, var_diag, _, _ = self._local_solve(query_z[i]) + values[i] = coef[columns] * factors + if return_std: + stds[i] = np.sqrt(np.maximum(scale * var_diag[columns] * factors**2, 0.0)) + return values, stds + + # -- Gaussian process --------------------------------------------------- + + def _gp_kernel(self, a: np.ndarray, b: np.ndarray, sf2: float, ls: np.ndarray) -> np.ndarray: + """ + Anisotropic squared-exponential kernel matrix. + + Parameters + ---------- + a : np.ndarray of shape (n_a, n_dims) + First set of standardized coordinates. + b : np.ndarray of shape (n_b, n_dims) + Second set of standardized coordinates. + sf2 : float + Signal variance. + ls : np.ndarray of shape (n_dims,) + Length scales. + + Returns + ------- + np.ndarray of shape (n_a, n_b) + Kernel matrix. + """ + diff = (a[:, None, :] - b[None, :, :]) / ls[None, None, :] + return sf2 * np.exp(-0.5 * np.sum(diff**2, axis=2)) + + def _fit_gp(self) -> None: + """ + Fit a Gaussian process and select its hyperparameters. + + Raises + ------ + ValueError + If the training set is larger than :attr:`max_points`. + """ + from scipy.linalg import cho_factor, cho_solve + from scipy.optimize import minimize + + n_points, n_dims = self.coords_.shape + if n_points > self.max_points: + raise ValueError( + f"method='gp' has cubic cost and is capped at max_points={self.max_points}; " + f"got {n_points} points. Subsample, raise max_points, or use " + "method='tensor_spline'." + ) + + self._center = self.coords_.mean(axis=0) + spread = self.coords_.std(axis=0) + spread[spread <= 0] = 1.0 + self._scale = spread + self._coords_z = (self.coords_ - self._center) / self._scale + + self._y_mean = float(self.values_.mean()) + y = self.values_ - self._y_mean + y_var = max(float(np.var(y)), _JITTER) + + if self.length_scale is None: + init_ls = np.ones(n_dims) + fixed_ls = None + else: + fixed_ls = np.asarray(self.length_scale, dtype=np.float64).ravel() + if fixed_ls.size == 1: + fixed_ls = np.full(n_dims, float(fixed_ls[0])) + if fixed_ls.size != n_dims: + raise ValueError(f"length_scale must be a scalar or have {n_dims} entries") + if np.any(fixed_ls <= 0): + raise ValueError("length_scale entries must be strictly positive") + init_ls = fixed_ls + + if isinstance(self.smoothing, float): + noise_var = self.smoothing + noise_source = "fixed" + elif self._sigma is not None: + # Covers both smoothing='sigma' and smoothing='auto' with a known noise level. + noise_var = float(np.mean(self._sigma**2)) + noise_source = "sigma" + else: + noise_var = 0.01 * y_var + noise_source = "marginal_likelihood" + + if self._sigma is None: + noise_shape = np.ones(n_points) + else: + noise_shape = self._sigma**2 / float(np.mean(self._sigma**2)) + + fit_noise = noise_source == "marginal_likelihood" + + def unpack(theta: np.ndarray) -> tuple[float, np.ndarray, float]: + sf2 = math.exp(theta[0]) + ls = fixed_ls if fixed_ls is not None else np.exp(theta[1 : 1 + n_dims]) + nv = math.exp(theta[-1]) if fit_noise else noise_var + return sf2, ls, nv + + def negative_log_marginal(theta: np.ndarray) -> float: + sf2, ls, nv = unpack(theta) + kernel = self._gp_kernel(self._coords_z, self._coords_z, sf2, ls) + kernel[np.diag_indices(n_points)] += nv * noise_shape + _JITTER * sf2 + try: + factor = cho_factor(kernel, lower=True) + except np.linalg.LinAlgError: + return 1e12 + alpha = cho_solve(factor, y) + log_det = 2.0 * float(np.sum(np.log(np.diag(factor[0])))) + return 0.5 * float(y @ alpha) + 0.5 * log_det + + theta0 = [math.log(y_var)] + bounds = [(math.log(y_var) - 8, math.log(y_var) + 8)] + if fixed_ls is None: + theta0 += list(np.log(init_ls)) + bounds += [(math.log(0.02), math.log(20.0))] * n_dims + if fit_noise: + theta0.append(math.log(noise_var)) + bounds.append((math.log(1e-8 * y_var), math.log(y_var))) + + opt = minimize( + negative_log_marginal, + np.asarray(theta0), + method="L-BFGS-B", + bounds=bounds, + options={"maxiter": 200}, + ) + sf2, ls, noise_var = unpack(opt.x) + noise_var *= self.smoothing_scale + + kernel = self._gp_kernel(self._coords_z, self._coords_z, sf2, ls) + kernel_noisy = kernel.copy() + kernel_noisy[np.diag_indices(n_points)] += noise_var * noise_shape + _JITTER * sf2 + factor = cho_factor(kernel_noisy, lower=True) + self._gp_alpha = cho_solve(factor, y) + self._gp_kinv = cho_solve(factor, np.eye(n_points)) + self._gp_signal_var = sf2 + self._gp_length_scale = ls + self._gp_noise_var = noise_var + + self._fitted_values = kernel @ self._gp_alpha + self._y_mean + self.effective_dof_ = float(np.trace(kernel @ self._gp_kinv)) + self.smoothing_ = float(noise_var) + self.smoothing_source_ = noise_source + + def _gp_cross_covariance(self, query_z: np.ndarray, order: np.ndarray) -> np.ndarray: + """ + Derivative cross-covariance between a partial at the query points and the data. + + Parameters + ---------- + query_z : np.ndarray of shape (n_query, n_dims) + Standardized query points. + order : np.ndarray of shape (n_dims,) + Derivative order per dimension. + + Returns + ------- + np.ndarray of shape (n_query, n_points) + The requested partial derivative of the kernel with respect to the query + coordinates, in original (unstandardized) units. + """ + ls = self._gp_length_scale + diff = (query_z[:, None, :] - self._coords_z[None, :, :]) / ls[None, None, :] + cov = self._gp_signal_var * np.exp(-0.5 * np.sum(diff**2, axis=2)) + for dim in range(self.n_features_in_): + m = int(order[dim]) + if m == 0: + continue + step = ls[dim] * self._scale[dim] + cov = cov * ((-1.0) ** m) * _hermite_e(m, diff[:, :, dim]) / step**m + return cov + + def _gp_prior_variance(self, order: np.ndarray) -> float: + """ + Prior variance of a partial derivative of the GP. + + Parameters + ---------- + order : np.ndarray of shape (n_dims,) + Derivative order per dimension. + + Returns + ------- + float + ``Var(d^order f(x))`` under the prior. + """ + var = self._gp_signal_var * (-1.0) ** int(np.sum(order)) + for dim in range(self.n_features_in_): + m = int(order[dim]) + if m == 0: + continue + step = self._gp_length_scale[dim] * self._scale[dim] + var *= _hermite_e_at_zero(2 * m) / step ** (2 * m) + return float(var) + + def _eval_gp( + self, query: np.ndarray, orders: np.ndarray, return_std: bool + ) -> tuple[np.ndarray, np.ndarray]: + """ + Evaluate GP posterior partials at the query points. + + Parameters + ---------- + query : np.ndarray of shape (n_query, n_dims) + Query points. + orders : np.ndarray of shape (n_orders, n_dims) + Derivative orders. + return_std : bool + Whether to compute posterior standard deviations. + + Returns + ------- + values : np.ndarray of shape (n_query, n_orders) + Posterior mean partials. + std : np.ndarray of shape (n_query, n_orders) + Posterior standard deviations, left as zeros when *return_std* is ``False``. + """ + query_z = (query - self._center) / self._scale + values = np.empty((query.shape[0], orders.shape[0])) + stds = np.zeros_like(values) + + for k, order in enumerate(orders): + cross = self._gp_cross_covariance(query_z, order) + values[:, k] = cross @ self._gp_alpha + if int(np.sum(order)) == 0: + values[:, k] += self._y_mean + if return_std: + prior = self._gp_prior_variance(order) + explained = np.einsum("ij,jk,ik->i", cross, self._gp_kinv, cross) + stds[:, k] = np.sqrt(np.maximum(prior - explained, 0.0)) + return values, stds + + +def _interp_log(grid: np.ndarray, response: np.ndarray, target: float) -> float: + """ + Invert a monotonically increasing response curve on a logarithmic grid. + + Parameters + ---------- + grid : np.ndarray + Strictly positive, increasing smoothing levels. + response : np.ndarray + Response (residual sum of squares) at each grid point; assumed increasing. + target : float + Desired response value. + + Returns + ------- + float + The grid value whose response matches *target*, clipped to the grid range. + """ + monotone = np.maximum.accumulate(response) + if target <= monotone[0]: + return float(grid[0]) + if target >= monotone[-1]: + return float(grid[-1]) + return float(np.exp(np.interp(target, monotone, np.log(grid)))) + + +def estimate_partial_derivatives( + coords: np.ndarray | Sequence[np.ndarray], + values: np.ndarray, + order: Any, + method: str = "tensor_spline", + sigma: float | np.ndarray | None = None, + query: np.ndarray | Sequence[np.ndarray] | None = None, + return_std: bool = False, + **kwargs: Any, +) -> tuple[np.ndarray, ...]: + """ + Estimate partial derivatives of a smoothed N-D surface in one call. + + Convenience wrapper around :class:`SurfaceDerivatives` for the common case of + fitting a smoother and evaluating partials at the sample locations. + + Parameters + ---------- + coords : np.ndarray of shape (n_points, n_dims), or sequence of 1-D arrays + Sample locations, scattered or a rectangular grid (see + :meth:`SurfaceDerivatives.fit`). + values : np.ndarray + Observed values, flat or grid-shaped to match *coords*. + order : tuple of int, or sequence of tuples + Derivative orders per dimension, e.g. ``[(1, 0), (0, 1)]``. + method : str + Smoother to use: ``"tensor_spline"``, ``"local_poly"``, or ``"gp"``. + sigma : float or np.ndarray of shape (n_points,), optional + Known measurement noise standard deviation. + query : np.ndarray of shape (n_query, n_dims), optional + Where to evaluate. Defaults to the (flattened) sample locations. + return_std : bool + If ``True``, also return the standard error of each partial derivative. + **kwargs + Extra keyword arguments forwarded to :class:`SurfaceDerivatives`. + + Returns + ------- + values : np.ndarray of shape (n_query,) + The smoothed surface at the query points. + partials : np.ndarray of shape (n_query, n_orders) + One column per requested derivative order. + std : np.ndarray of shape (n_query, n_orders) + Standard errors. Only returned when *return_std* is ``True``. + + Raises + ------ + ValueError + If the inputs or the derivative orders are invalid. + + Examples + -------- + >>> import numpy as np + >>> from jaxsr import estimate_partial_derivatives + >>> x = np.linspace(0, 1, 20) + >>> t = np.linspace(0, 1, 20) + >>> XX, TT = np.meshgrid(x, t, indexing="ij") + >>> Z = XX**2 + 3 * TT + >>> _, dz = estimate_partial_derivatives([x, t], Z, order=[(1, 0), (0, 1)]) + >>> dz.shape + (400, 2) + """ + estimator = SurfaceDerivatives(method=method, **kwargs) + estimator.fit(coords, values, sigma=sigma) + target = estimator.coords_ if query is None else query + return estimator.derivatives(target, order=order, return_std=return_std) diff --git a/src/jaxsr/dynamics.py b/src/jaxsr/dynamics.py index 19115fc..a233618 100644 --- a/src/jaxsr/dynamics.py +++ b/src/jaxsr/dynamics.py @@ -58,6 +58,12 @@ def estimate_derivatives( ValueError If *method* is unknown, shapes are inconsistent, *t* is not monotonically increasing, or savgol is used with non-uniform spacing. + + See Also + -------- + jaxsr.derivatives.SurfaceDerivatives : Partial derivatives of an N-D surface, + for problems needing more than one partial (PDE-style discovery, shift + laws). This function differentiates along a single axis only. """ X = np.asarray(X, dtype=np.float64) t = np.asarray(t, dtype=np.float64).ravel() diff --git a/src/jaxsr/skill/SKILL.md b/src/jaxsr/skill/SKILL.md index 5b1c908..eb3b58d 100644 --- a/src/jaxsr/skill/SKILL.md +++ b/src/jaxsr/skill/SKILL.md @@ -386,6 +386,13 @@ See `guides/rsm.md` for RSM designs, canonical analysis, and optimization. See `guides/active-learning.md` for acquisition functions and adaptive sampling. +### "My data is a surface and I need partial derivatives" + +See `guides/surface-derivatives.md` for `SurfaceDerivatives`: several partials of one +smoothed N-D surface (`y_x` and `y_T`, or `u_t = F(u, u_x, u_xx, ...)`), with the +smoothing level chosen by GCV/marginal likelihood and reported. For a single time axis +(`dX/dt`), use `estimate_derivatives` / `discover_dynamics` instead. + ### "One expression isn't enough / the signal is a sum of many effects" See `guides/additive.md` for boosting-style additive symbolic regression diff --git a/src/jaxsr/skill/guides/surface-derivatives.md b/src/jaxsr/skill/guides/surface-derivatives.md new file mode 100644 index 0000000..a4bf5b1 --- /dev/null +++ b/src/jaxsr/skill/guides/surface-derivatives.md @@ -0,0 +1,184 @@ +# Multivariate Derivative Estimation (`SurfaceDerivatives`) + +Estimate **partial derivatives of a surface** — several partials from one smoothed +fit — for problems where the regression target or the basis library contains +derivatives. + +## When to use this + +| Situation | Use | +|-----------|-----| +| One state trajectory over time, need `dX/dt` | `estimate_derivatives(X, t, ...)` (see `jaxsr.dynamics`) | +| Whole ODE system from time series | `discover_dynamics(X, t, ...)` | +| Data is a surface over 2+ coordinates, need `y_x` **and** `y_T` | `SurfaceDerivatives` | +| PDE-style discovery: `u_t = F(u, u_x, u_xx, ...)` | `SurfaceDerivatives` | +| Transform/shift laws, e.g. `y(x, T) = f(x + s(T))` ⟹ `y_T = s'(T)·y_x` | `SurfaceDerivatives` | + +`estimate_derivatives` differentiates along a **single** axis (`X` is +`(n_times, n_states)`, `t` is 1-D). It cannot give you two partials of one surface. + +## API + +The snippets in this section assume `coords` (an `(n, d)` array), `values` (`(n,)`) and +`sigma` come from the user's data; the worked examples further down are self-contained. + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +est = SurfaceDerivatives(method="tensor_spline") # or "local_poly", "gp" +est.fit(coords, values, sigma=0.01) # coords (n, d), values (n,) + +y, dy = est.derivatives(coords, order=[(1, 0), (0, 1)]) +# y -> (n,) smoothed surface +# dy -> (n, 2) column 0 = d/dx0, column 1 = d/dx1 + +y, dy, dy_se = est.derivatives(coords, order=[(1, 0), (0, 1)], return_std=True) +print(est.summary()) # method, smoothing level, effective dof, residual std +``` + +Gridded data can be passed as axes plus an N-D array — no meshgrid needed: + +```python +est = SurfaceDerivatives().fit([x_axis, T_axis], Y_grid) # Y_grid (len(x), len(T)) +y, dy = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) # coords_ is the flat (n, 2) grid +``` + +One-call convenience wrapper: + +```python +from jaxsr import estimate_partial_derivatives + +y, dy = estimate_partial_derivatives(coords, values, order=[(1, 0), (0, 1)], + method="tensor_spline", sigma=0.01) +``` + +**Signature notes** (common mistakes): + +| API | Wrong | Right | +|-----|-------|-------| +| `derivatives()` | `dy = est.derivatives(...)` | returns a **tuple** `(y, dy)`, or `(y, dy, std)` with `return_std=True` | +| `order` | `order=1`, `order="x"` | a tuple per dimension: `(1, 0)`, or a list of them | +| single order | expecting shape `(n,)` | a single tuple still returns `(n, 1)` | +| query points | grid axes | `derivatives()` takes an `(n_query, d)` array — use `est.coords_` for the sample locations | +| `sigma` | a variance | a **standard deviation**, scalar or per point | + +## Choosing a method + +| Method | Data | Cost | Derivative uncertainty | Notes | +|--------|------|------|------------------------|-------| +| `"tensor_spline"` (default) | gridded or scattered, any `d` | fast | from the penalized-fit posterior | Best default. Penalty chosen by GCV. | +| `"local_poly"` | scattered, irregular | moderate (per-point fits) | sandwich variance of the local fit | Good on uneven sampling; degree bounds the total order. | +| `"gp"` | scattered, irregular, small `n` | `O(n³)`, capped by `max_points=800` | exact posterior, grows away from data | Best uncertainty; use when `n` is a few hundred. | + +Derivative orders are always analytic partials of the fitted smoother — never finite +differences of noisy raw data. + +## Choosing the smoothing level + +**The smoothing hyperparameter must never be tuned against the downstream symbolic +score.** If the smoother is selected by which law the regression likes best, it can +manufacture that law; the fit looks excellent and the failure is silent. + +| `smoothing=` | Meaning | +|--------------|---------| +| `"auto"` (default) | GCV (spline, local poly) or marginal likelihood (GP) | +| `"sigma"` | requires `sigma` at `fit()`; matches residual scatter to the known noise (the `s = n·σ²` rule) | +| float | use it verbatim: penalty `λ` (spline), bandwidth (local poly), noise variance (GP) | + +`smoothing_scale=3.0` multiplies whatever was selected — the cheapest way to check how +much a discovered coefficient depends on the derivative stage. + +## Reporting the smoothing actually used + +Smoothing biases derivatives toward zero, and any coefficient read off them inherits +that bias. Make it visible: + +```python +print(est.summary()) +est.smoothing_ # λ / bandwidth / noise variance actually used +est.smoothing_source_ # "gcv", "marginal_likelihood", "sigma", or "fixed" +est.effective_dof_ # effective degrees of freedom of the smoother +est.residual_std_ # residual scatter of the fit +``` + +Sensitivity check — rerun the whole pipeline at several smoothing scales and report the +spread, not just one number: + +```python +for scale in (1.0, 3.0, 10.0): + est = SurfaceDerivatives(smoothing_scale=scale).fit(coords, values, sigma=sigma) + ... # refit the symbolic stage, record the coefficient +``` + +## Example: PDE-style discovery + +```python +import numpy as np +from jaxsr import BasisLibrary, SurfaceDerivatives, SymbolicRegressor + +# Heat equation data: u_t = 0.1 * u_xx +x = np.linspace(0, 2 * np.pi, 40) +t = np.linspace(0, 1.0, 30) +xx, tt = np.meshgrid(x, t, indexing="ij") +u = np.exp(-0.1 * tt) * np.sin(xx) + 0.5 * np.exp(-0.9 * tt) * np.sin(3 * xx) +noise = 0.002 +u_obs = u + np.random.default_rng(0).normal(0, noise, u.shape) + +est = SurfaceDerivatives().fit([x, t], u_obs, sigma=noise) +values, d = est.derivatives(est.coords_, order=[(1, 0), (2, 0), (0, 1)]) +u_x, u_xx, u_t = d[:, 0], d[:, 1], d[:, 2] + +library = ( + BasisLibrary(n_features=3, feature_names=["u", "u_x", "u_xx"]) + .add_linear() + .add_interactions(max_order=2) +) +model = SymbolicRegressor(basis_library=library, max_terms=1).fit( + np.column_stack([values, u_x, u_xx]), u_t +) +print(model.expression_) # y = 0.0958*u_xx (true 0.1; the gap is smoothing bias) +``` + +## Example: shift law (`y_T = s'(T)·y_x`) + +```python +import numpy as np +from jaxsr import SurfaceDerivatives + +# Synthetic time-temperature superposition surface (Arrhenius shift, E = 55.85 kJ/mol) +gas_r, energy, t_ref = 8.314e-3, 55.85, 350.0 +x_axis = np.linspace(-2.0, 4.0, 20) # log frequency +T_axis = np.linspace(320.0, 400.0, 12) # temperature, K +xx, TT = np.meshgrid(x_axis, T_axis, indexing="ij") +shift = -(energy / (gas_r * np.log(10.0))) * (1.0 / TT - 1.0 / t_ref) +sigma = 0.01 +Y_grid = 2.0 + 1.5 * np.tanh(0.8 * (xx + shift - 1.0)) +Y_grid = Y_grid + np.random.default_rng(0).normal(0, sigma, Y_grid.shape) + +est = SurfaceDerivatives(smoothing="sigma").fit([x_axis, T_axis], Y_grid, sigma=sigma) +_, d = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) +y_x, y_T = d[:, 0], d[:, 1] + +keep = np.abs(y_x) > 0.15 * np.abs(y_x).max() # the ratio is ill-posed where y_x ≈ 0 +s_prime = y_T[keep] / y_x[keep] +T_keep = est.coords_[keep, 1] + +print(np.median(s_prime * gas_r * np.log(10.0) * T_keep**2)) # ≈ 55.9 kJ/mol +``` + +Then regress `s_prime` against `T_keep` with a `SymbolicRegressor` rather than assuming +the Arrhenius form. + +## Pitfalls + +- **Boundaries.** Every smoother is weakest at the edge of the data. Drop a margin + before reading derivatives, or expect the largest errors there. +- **Dividing partials.** Ratios like `y_T / y_x` blow up where the denominator crosses + zero. Mask small denominators, as above. +- **High orders.** Each extra order costs accuracy. `degree` must be at least the + highest order requested (`tensor_spline`: per dimension; `local_poly`: total). +- **GP size.** `method="gp"` is `O(n³)`; above `max_points=800` it raises rather than + hanging. Subsample or switch to `"tensor_spline"`. +- **Basis resolution.** The spline defaults to at most 12 basis functions per + dimension because GCV under-smooths derivatives; pass `n_basis=` for finer structure. diff --git a/tests/test_derivatives.py b/tests/test_derivatives.py new file mode 100644 index 0000000..d13b5a8 --- /dev/null +++ b/tests/test_derivatives.py @@ -0,0 +1,564 @@ +"""Tests for multivariate (N-D) derivative estimation.""" + +import numpy as np +import pytest + +from jaxsr.derivatives import SurfaceDerivatives, estimate_partial_derivatives + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def smooth_surface(): + """f(x, y) = sin(2x) * exp(y/2) on a rectangular grid, with exact partials.""" + x = np.linspace(0.0, 2.0, 25) + y = np.linspace(-1.0, 1.0, 20) + xx, yy = np.meshgrid(x, y, indexing="ij") + values = np.sin(2 * xx) * np.exp(0.5 * yy) + coords = np.column_stack([xx.ravel(), yy.ravel()]) + exact = { + (1, 0): (2 * np.cos(2 * xx) * np.exp(0.5 * yy)).ravel(), + (0, 1): (0.5 * np.sin(2 * xx) * np.exp(0.5 * yy)).ravel(), + (2, 0): (-4 * np.sin(2 * xx) * np.exp(0.5 * yy)).ravel(), + } + return x, y, values, coords, exact + + +@pytest.fixture() +def bilinear_surface(): + """f(x, y) = 3x + 2y + x*y: exactly representable, so partials are exact.""" + x = np.linspace(0.0, 1.0, 12) + y = np.linspace(0.0, 2.0, 11) + xx, yy = np.meshgrid(x, y, indexing="ij") + values = 3 * xx + 2 * yy + xx * yy + coords = np.column_stack([xx.ravel(), yy.ravel()]) + return x, y, values, coords + + +def interior_mask(coords, frac=0.12): + """Mask selecting points away from the boundary, where smoothers are weakest.""" + mask = np.ones(coords.shape[0], dtype=bool) + for dim in range(coords.shape[1]): + lo, hi = coords[:, dim].min(), coords[:, dim].max() + pad = frac * (hi - lo) + mask &= (coords[:, dim] >= lo + pad) & (coords[:, dim] <= hi - pad) + return mask + + +# =========================================================================== +# TestInputHandling +# =========================================================================== + + +class TestInputHandling: + """Tests for coordinate, value, order, and sigma normalization.""" + + def test_grid_and_scattered_forms_agree(self, bilinear_surface): + """Gridded axes + N-D values give the same fit as flat (n, d) coordinates.""" + x, y, values, coords = bilinear_surface + grid_fit = SurfaceDerivatives().fit([x, y], values) + flat_fit = SurfaceDerivatives().fit(coords, values.ravel()) + _, d_grid = grid_fit.derivatives(coords, order=[(1, 0)]) + _, d_flat = flat_fit.derivatives(coords, order=[(1, 0)]) + np.testing.assert_allclose(d_grid, d_flat, atol=1e-8) + np.testing.assert_allclose(grid_fit.coords_, coords) + + def test_grid_shape_mismatch_raises(self, bilinear_surface): + """Grid axes inconsistent with the values array raise ValueError.""" + x, y, values, _ = bilinear_surface + with pytest.raises(ValueError, match="does not match the grid"): + SurfaceDerivatives().fit([x, y[:-1]], values) + + def test_length_mismatch_raises(self): + """Mismatched coordinate and value counts raise ValueError.""" + coords = np.random.default_rng(0).uniform(size=(20, 2)) + with pytest.raises(ValueError, match="must match the number"): + SurfaceDerivatives().fit(coords, np.zeros(19)) + + def test_non_finite_values_raise(self): + """Non-finite values raise ValueError.""" + coords = np.random.default_rng(0).uniform(size=(20, 2)) + values = np.zeros(20) + values[3] = np.nan + with pytest.raises(ValueError, match="non-finite"): + SurfaceDerivatives().fit(coords, values) + + def test_too_few_points_raise(self): + """Fewer than four points raise ValueError.""" + coords = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.5]]) + with pytest.raises(ValueError, match="at least 4 data points"): + SurfaceDerivatives().fit(coords, np.zeros(3)) + + def test_single_order_tuple_gives_one_column(self, bilinear_surface): + """A single order tuple yields a (n_query, 1) partials array.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives().fit([x, y], values) + vals, partials = est.derivatives(coords, order=(1, 0)) + assert vals.shape == (coords.shape[0],) + assert partials.shape == (coords.shape[0], 1) + + def test_wrong_order_length_raises(self, bilinear_surface): + """Derivative orders must have one entry per dimension.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives().fit([x, y], values) + with pytest.raises(ValueError, match="one per dimension"): + est.derivatives(coords, order=[(1, 0, 0)]) + + def test_negative_order_raises(self, bilinear_surface): + """Negative derivative orders raise ValueError.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives().fit([x, y], values) + with pytest.raises(ValueError, match="non-negative"): + est.derivatives(coords, order=[(-1, 0)]) + + def test_order_above_spline_degree_raises(self, bilinear_surface): + """Requesting a derivative above the spline degree raises ValueError.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives(degree=2).fit([x, y], values) + with pytest.raises(ValueError, match="exceeds the spline degree"): + est.derivatives(coords, order=[(3, 0)]) + + def test_order_above_local_poly_degree_raises(self, bilinear_surface): + """Local polynomial degree bounds the total derivative order.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives(method="local_poly", degree=2).fit([x, y], values) + with pytest.raises(ValueError, match="exceeds the local polynomial degree"): + est.derivatives(coords, order=[(2, 1)]) + + def test_query_dimension_mismatch_raises(self, bilinear_surface): + """Query coordinates must have the fitted number of dimensions.""" + x, y, values, _ = bilinear_surface + est = SurfaceDerivatives().fit([x, y], values) + with pytest.raises(ValueError, match=r"shape \(n_query, 2\)"): + est.derivatives(np.zeros((5, 3)), order=[(1, 0)]) + + def test_unfitted_raises(self): + """Calling derivatives(), predict(), or summary() before fit() raises.""" + est = SurfaceDerivatives() + with pytest.raises(RuntimeError, match="must be fitted"): + est.derivatives(np.zeros((2, 2)), order=[(1, 0)]) + with pytest.raises(RuntimeError, match="must be fitted"): + est.predict(np.zeros((2, 2))) + with pytest.raises(RuntimeError, match="must be fitted"): + est.summary() + + def test_invalid_constructor_arguments(self): + """Invalid constructor arguments raise ValueError.""" + with pytest.raises(ValueError, match="Unknown method"): + SurfaceDerivatives(method="bogus") + with pytest.raises(ValueError, match="degree must be"): + SurfaceDerivatives(degree=0) + with pytest.raises(ValueError, match="smoothing must be"): + SurfaceDerivatives(smoothing="cv") + with pytest.raises(ValueError, match="smoothing_scale"): + SurfaceDerivatives(smoothing_scale=0.0) + with pytest.raises(ValueError, match="Numeric smoothing"): + SurfaceDerivatives(smoothing=-1.0) + + def test_sigma_validation(self, bilinear_surface): + """sigma must be positive and either scalar or per point.""" + x, y, values, coords = bilinear_surface + with pytest.raises(ValueError, match="strictly positive"): + SurfaceDerivatives().fit(coords, values.ravel(), sigma=0.0) + with pytest.raises(ValueError, match="scalar or have"): + SurfaceDerivatives().fit(coords, values.ravel(), sigma=np.ones(3)) + + def test_smoothing_sigma_requires_sigma(self, bilinear_surface): + """smoothing='sigma' without a supplied sigma raises ValueError.""" + x, y, values, _ = bilinear_surface + with pytest.raises(ValueError, match="requires the sigma"): + SurfaceDerivatives(smoothing="sigma").fit([x, y], values) + + +# =========================================================================== +# TestTensorSpline +# =========================================================================== + + +class TestTensorSpline: + """Tests for the penalized tensor-product B-spline smoother.""" + + def test_exact_on_bilinear(self, bilinear_surface): + """A bilinear surface is in the spline space, so partials are essentially exact.""" + x, y, values, coords = bilinear_surface + est = SurfaceDerivatives(smoothing=1e-8).fit([x, y], values) + _, partials = est.derivatives(coords, order=[(1, 0), (0, 1), (1, 1)]) + np.testing.assert_allclose(partials[:, 0], 3 + coords[:, 1], atol=1e-6) + np.testing.assert_allclose(partials[:, 1], 2 + coords[:, 0], atol=1e-6) + np.testing.assert_allclose(partials[:, 2], 1.0, atol=1e-6) + + def test_first_partials_on_noisy_surface(self, smooth_surface): + """Both first partials are recovered from one noisy surface.""" + x, y, values, coords, exact = smooth_surface + rng = np.random.default_rng(0) + noisy = values + rng.normal(0, 0.01, values.shape) + est = SurfaceDerivatives().fit([x, y], noisy, sigma=0.01) + _, partials = est.derivatives(coords, order=[(1, 0), (0, 1)]) + mask = interior_mask(coords) + assert np.abs(partials[mask, 0] - exact[(1, 0)][mask]).max() < 0.15 + assert np.abs(partials[mask, 1] - exact[(0, 1)][mask]).max() < 0.15 + + def test_second_partial(self, smooth_surface): + """Second partials track the analytic result on clean data.""" + x, y, values, coords, exact = smooth_surface + est = SurfaceDerivatives().fit([x, y], values) + _, partials = est.derivatives(coords, order=[(2, 0)]) + mask = interior_mask(coords) + rel = np.abs(partials[mask, 0] - exact[(2, 0)][mask]).max() / np.abs(exact[(2, 0)]).max() + assert rel < 0.1 + + def test_predict_matches_zero_order(self, smooth_surface): + """predict() equals the order-zero output of derivatives().""" + x, y, values, coords, _ = smooth_surface + est = SurfaceDerivatives().fit([x, y], values) + vals, _ = est.derivatives(coords, order=[(1, 0)]) + np.testing.assert_allclose(est.predict(coords), vals) + + def test_predict_return_std(self, smooth_surface): + """predict(return_std=True) returns positive, correctly shaped standard errors.""" + x, y, values, coords, _ = smooth_surface + rng = np.random.default_rng(1) + est = SurfaceDerivatives().fit([x, y], values + rng.normal(0, 0.02, values.shape)) + mean, std = est.predict(coords, return_std=True) + assert mean.shape == std.shape == (coords.shape[0],) + assert np.all(std > 0) + + def test_std_grows_with_noise(self, smooth_surface): + """Derivative standard errors increase with the measurement noise.""" + x, y, values, coords, _ = smooth_surface + rng = np.random.default_rng(2) + quiet = SurfaceDerivatives().fit([x, y], values + rng.normal(0, 0.005, values.shape)) + loud = SurfaceDerivatives().fit([x, y], values + rng.normal(0, 0.05, values.shape)) + _, _, std_quiet = quiet.derivatives(coords, order=[(1, 0)], return_std=True) + _, _, std_loud = loud.derivatives(coords, order=[(1, 0)], return_std=True) + assert std_loud.mean() > std_quiet.mean() + + def test_gcv_selects_smoothing(self, smooth_surface): + """With smoothing='auto' the penalty is chosen by GCV and reported.""" + x, y, values, _, _ = smooth_surface + rng = np.random.default_rng(3) + est = SurfaceDerivatives().fit([x, y], values + rng.normal(0, 0.02, values.shape)) + assert est.smoothing_source_ == "gcv" + assert est.smoothing_ > 0 + assert 0 < est.effective_dof_ < values.size + + def test_sigma_criterion_matches_residuals(self, smooth_surface): + """smoothing='sigma' picks a penalty whose residual scatter matches sigma.""" + x, y, values, _, _ = smooth_surface + rng = np.random.default_rng(4) + sigma = 0.02 + est = SurfaceDerivatives(smoothing="sigma").fit( + [x, y], values + rng.normal(0, sigma, values.shape), sigma=sigma + ) + assert est.smoothing_source_ == "sigma" + assert 0.5 * sigma < est.residual_std_ < 2.0 * sigma + + def test_smoothing_scale_increases_smoothing(self, smooth_surface): + """smoothing_scale multiplies the selected penalty and lowers the effective dof.""" + x, y, values, _, _ = smooth_surface + rng = np.random.default_rng(5) + noisy = values + rng.normal(0, 0.02, values.shape) + base = SurfaceDerivatives().fit([x, y], noisy) + scaled = SurfaceDerivatives(smoothing_scale=5.0).fit([x, y], noisy) + assert scaled.smoothing_ == pytest.approx(5.0 * base.smoothing_) + assert scaled.effective_dof_ < base.effective_dof_ + + def test_fixed_smoothing_is_reported(self, bilinear_surface): + """A numeric smoothing value is used verbatim and reported as fixed.""" + x, y, values, _ = bilinear_surface + est = SurfaceDerivatives(smoothing=0.5).fit([x, y], values) + assert est.smoothing_ == pytest.approx(0.5) + assert est.smoothing_source_ == "fixed" + + def test_scattered_data(self, smooth_surface): + """Scattered (non-gridded) samples are supported.""" + _, _, _, _, _ = smooth_surface + rng = np.random.default_rng(6) + coords = rng.uniform([0, -1], [2, 1], size=(500, 2)) + values = np.sin(2 * coords[:, 0]) * np.exp(0.5 * coords[:, 1]) + est = SurfaceDerivatives().fit(coords, values) + _, partials = est.derivatives(coords, order=[(1, 0)]) + exact = 2 * np.cos(2 * coords[:, 0]) * np.exp(0.5 * coords[:, 1]) + mask = interior_mask(coords, frac=0.15) + assert np.abs(partials[mask, 0] - exact[mask]).max() < 0.3 + + def test_three_dimensional(self): + """Three-dimensional surfaces are supported.""" + axis = np.linspace(0.0, 1.0, 8) + aa, bb, cc = np.meshgrid(axis, axis, axis, indexing="ij") + values = aa**2 + 2 * bb * cc + est = SurfaceDerivatives(smoothing=1e-8).fit([axis, axis, axis], values) + _, partials = est.derivatives(est.coords_, order=[(1, 0, 0), (0, 1, 1)]) + np.testing.assert_allclose(partials[:, 0], 2 * est.coords_[:, 0], atol=1e-5) + np.testing.assert_allclose(partials[:, 1], 2.0, atol=1e-5) + + def test_one_dimensional(self): + """One-dimensional data reduces to a smoothing spline derivative.""" + t = np.linspace(0.0, 2 * np.pi, 200).reshape(-1, 1) + values = np.sin(t.ravel()) + est = SurfaceDerivatives().fit(t, values) + _, partials = est.derivatives(t, order=(1,)) + mask = interior_mask(t) + assert np.abs(partials[mask, 0] - np.cos(t.ravel())[mask]).max() < 0.05 + + def test_explicit_n_basis(self, smooth_surface): + """n_basis controls the per-dimension basis size.""" + x, y, values, coords, _ = smooth_surface + est = SurfaceDerivatives(n_basis=[8, 6]).fit([x, y], values) + assert est._spline_sizes == [8, 6] + assert "basis per dim : [8, 6]" in est.summary() + + def test_n_basis_too_large_raises(self, bilinear_surface): + """A basis larger than max_basis raises ValueError.""" + x, y, values, _ = bilinear_surface + with pytest.raises(ValueError, match="above max_basis"): + SurfaceDerivatives(n_basis=30, max_basis=100).fit([x, y], values) + + def test_n_basis_below_degree_raises(self, bilinear_surface): + """Fewer basis functions than degree + 1 raises ValueError.""" + x, y, values, _ = bilinear_surface + with pytest.raises(ValueError, match="at least degree"): + SurfaceDerivatives(degree=3, n_basis=3).fit([x, y], values) + + def test_summary_reports_smoothing(self, smooth_surface): + """summary() reports the method, smoothing level, and how it was chosen.""" + x, y, values, _, _ = smooth_surface + est = SurfaceDerivatives().fit([x, y], values) + text = est.summary() + assert "tensor_spline" in text + assert "penalty lambda" in text + assert "effective dof" in text + assert "chosen by" in text + + +# =========================================================================== +# TestLocalPoly +# =========================================================================== + + +class TestLocalPoly: + """Tests for the local polynomial smoother.""" + + def test_first_partials(self, smooth_surface): + """Both first partials are recovered on the interior.""" + x, y, values, coords, exact = smooth_surface + est = SurfaceDerivatives(method="local_poly").fit([x, y], values) + _, partials = est.derivatives(coords, order=[(1, 0), (0, 1)]) + mask = interior_mask(coords) + assert np.abs(partials[mask, 0] - exact[(1, 0)][mask]).max() < 0.15 + assert np.abs(partials[mask, 1] - exact[(0, 1)][mask]).max() < 0.15 + + def test_exact_on_quadratic(self): + """A quadratic is reproduced exactly by a degree-2 local fit.""" + rng = np.random.default_rng(7) + coords = rng.uniform(-1, 1, size=(200, 2)) + values = 1.0 + 2 * coords[:, 0] - 3 * coords[:, 1] + 0.5 * coords[:, 0] * coords[:, 1] + est = SurfaceDerivatives(method="local_poly", degree=2).fit(coords, values) + _, partials = est.derivatives(coords, order=[(1, 0), (0, 1), (1, 1)]) + np.testing.assert_allclose(partials[:, 0], 2 + 0.5 * coords[:, 1], atol=1e-6) + np.testing.assert_allclose(partials[:, 1], -3 + 0.5 * coords[:, 0], atol=1e-6) + np.testing.assert_allclose(partials[:, 2], 0.5, atol=1e-6) + + def test_std_is_reported(self, smooth_surface): + """Standard errors are positive and shaped like the partials.""" + x, y, values, coords, _ = smooth_surface + rng = np.random.default_rng(8) + est = SurfaceDerivatives(method="local_poly").fit( + [x, y], values + rng.normal(0, 0.02, values.shape) + ) + _, partials, std = est.derivatives(coords[:50], order=[(1, 0), (0, 1)], return_std=True) + assert std.shape == partials.shape + assert np.all(std > 0) + + def test_bandwidth_selected_by_gcv(self, smooth_surface): + """The bandwidth is chosen by GCV and reported.""" + x, y, values, _, _ = smooth_surface + rng = np.random.default_rng(9) + est = SurfaceDerivatives(method="local_poly").fit( + [x, y], values + rng.normal(0, 0.02, values.shape) + ) + assert est.smoothing_source_ == "gcv" + assert est.smoothing_ > 0 + assert "bandwidth" in est.summary() + + def test_fixed_bandwidth(self, smooth_surface): + """A numeric smoothing value is used as the bandwidth verbatim.""" + x, y, values, _, _ = smooth_surface + est = SurfaceDerivatives(method="local_poly", smoothing=0.7).fit([x, y], values) + assert est.smoothing_ == pytest.approx(0.7) + assert est.smoothing_source_ == "fixed" + + def test_too_few_points_for_degree(self): + """A local degree needing more terms than there are points raises ValueError.""" + rng = np.random.default_rng(10) + coords = rng.uniform(size=(8, 2)) + with pytest.raises(ValueError, match="needs more than"): + SurfaceDerivatives(method="local_poly", degree=3).fit(coords, np.zeros(8)) + + +# =========================================================================== +# TestGaussianProcess +# =========================================================================== + + +class TestGaussianProcess: + """Tests for the Gaussian process smoother.""" + + def test_first_partials(self): + """Both first partials are recovered from a noisy GP fit.""" + rng = np.random.default_rng(11) + x = np.linspace(0.0, 2.0, 16) + y = np.linspace(-1.0, 1.0, 14) + xx, yy = np.meshgrid(x, y, indexing="ij") + values = np.sin(2 * xx) * np.exp(0.5 * yy) + est = SurfaceDerivatives(method="gp").fit( + [x, y], values + rng.normal(0, 0.01, values.shape), sigma=0.01 + ) + _, partials = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) + exact_x = (2 * np.cos(2 * xx) * np.exp(0.5 * yy)).ravel() + exact_y = (0.5 * np.sin(2 * xx) * np.exp(0.5 * yy)).ravel() + mask = interior_mask(est.coords_) + assert np.abs(partials[mask, 0] - exact_x[mask]).max() < 0.2 + assert np.abs(partials[mask, 1] - exact_y[mask]).max() < 0.2 + + def test_derivative_uncertainty(self): + """The GP reports positive derivative uncertainty that grows away from data.""" + rng = np.random.default_rng(12) + coords = rng.uniform(-1, 1, size=(120, 2)) + values = np.sin(coords[:, 0]) + coords[:, 1] ** 2 + est = SurfaceDerivatives(method="gp").fit(coords, values, sigma=0.01) + _, _, std_in = est.derivatives(coords, order=[(1, 0)], return_std=True) + _, _, std_out = est.derivatives(np.array([[4.0, 4.0]]), order=[(1, 0)], return_std=True) + assert np.all(std_in > 0) + assert std_out[0, 0] > std_in.mean() + + def test_fixed_length_scale(self): + """A supplied length_scale is used instead of being learned.""" + rng = np.random.default_rng(13) + coords = rng.uniform(-1, 1, size=(80, 2)) + values = np.sin(coords[:, 0]) + coords[:, 1] + est = SurfaceDerivatives(method="gp", length_scale=[1.0, 2.0]).fit(coords, values) + np.testing.assert_allclose(est._gp_length_scale, [1.0, 2.0]) + assert "length scales" in est.summary() + + def test_bad_length_scale_raises(self): + """Invalid length scales raise ValueError.""" + rng = np.random.default_rng(14) + coords = rng.uniform(-1, 1, size=(40, 2)) + values = coords[:, 0] + with pytest.raises(ValueError, match="length_scale must be"): + SurfaceDerivatives(method="gp", length_scale=[1.0, 2.0, 3.0]).fit(coords, values) + with pytest.raises(ValueError, match="strictly positive"): + SurfaceDerivatives(method="gp", length_scale=-1.0).fit(coords, values) + + def test_max_points_guard(self): + """Exceeding max_points raises a ValueError that names the cap.""" + rng = np.random.default_rng(15) + coords = rng.uniform(size=(60, 2)) + with pytest.raises(ValueError, match="max_points"): + SurfaceDerivatives(method="gp", max_points=50).fit(coords, np.zeros(60)) + + def test_marginal_likelihood_without_sigma(self): + """Without sigma the noise level is learned by marginal likelihood.""" + rng = np.random.default_rng(16) + coords = rng.uniform(-1, 1, size=(100, 2)) + values = np.sin(coords[:, 0]) + rng.normal(0, 0.05, 100) + est = SurfaceDerivatives(method="gp").fit(coords, values) + assert est.smoothing_source_ == "marginal_likelihood" + assert est.smoothing_ > 0 + + +# =========================================================================== +# TestEstimatePartialDerivatives +# =========================================================================== + + +class TestEstimatePartialDerivatives: + """Tests for the convenience wrapper.""" + + def test_returns_values_and_partials(self, bilinear_surface): + """The wrapper evaluates at the sample locations by default.""" + x, y, values, coords = bilinear_surface + vals, partials = estimate_partial_derivatives( + [x, y], values, order=[(1, 0), (0, 1)], smoothing=1e-8 + ) + assert vals.shape == (coords.shape[0],) + assert partials.shape == (coords.shape[0], 2) + np.testing.assert_allclose(partials[:, 0], 3 + coords[:, 1], atol=1e-6) + + def test_query_and_return_std(self, bilinear_surface): + """Explicit query points and standard errors are supported.""" + x, y, values, _ = bilinear_surface + query = np.array([[0.5, 1.0], [0.25, 0.5]]) + vals, partials, std = estimate_partial_derivatives( + [x, y], values, order=[(0, 1)], query=query, return_std=True + ) + assert vals.shape == (2,) + assert partials.shape == std.shape == (2, 1) + + def test_method_passthrough(self, bilinear_surface): + """The method and extra keyword arguments reach the estimator.""" + x, y, values, coords = bilinear_surface + _, partials = estimate_partial_derivatives( + coords, values.ravel(), order=[(1, 0)], method="local_poly", degree=2 + ) + np.testing.assert_allclose(partials[:, 0], 3 + coords[:, 1], atol=1e-6) + + +# =========================================================================== +# TestShiftLawRecovery +# =========================================================================== + + +class TestShiftLawRecovery: + """The motivating application: a time-temperature superposition shift law. + + ``y(x, T) = f(x + s(T))`` implies ``y_T = s'(T) * y_x``, so both partials must + come from one smoothed surface. + """ + + @staticmethod + def _surface(noise, seed): + gas_r = 8.314e-3 # kJ/mol/K + energy = 55.85 # kJ/mol + t_ref = 350.0 + x = np.linspace(-2.0, 4.0, 20) + temps = np.linspace(320.0, 400.0, 12) + xx, tt = np.meshgrid(x, temps, indexing="ij") + shift = -(energy / (gas_r * np.log(10.0))) * (1.0 / tt - 1.0 / t_ref) + values = 2.0 + 1.5 * np.tanh(0.8 * (xx + shift - 1.0)) + if noise: + values = values + np.random.default_rng(seed).normal(0, noise, values.shape) + return x, temps, values, energy, gas_r + + @staticmethod + def _effective_energy(est, gas_r): + _, partials = est.derivatives(est.coords_, order=[(1, 0), (0, 1)]) + y_x, y_t = partials[:, 0], partials[:, 1] + keep = np.abs(y_x) > 0.15 * np.abs(y_x).max() + temps = est.coords_[keep, 1] + return float(np.median(y_t[keep] / y_x[keep] * gas_r * np.log(10.0) * temps**2)) + + @pytest.mark.parametrize("method", ["tensor_spline", "local_poly", "gp"]) + def test_clean_surface(self, method): + """All smoothers recover the activation energy from noise-free data.""" + x, temps, values, energy, gas_r = self._surface(0.0, 0) + est = SurfaceDerivatives(method=method).fit([x, temps], values) + assert self._effective_energy(est, gas_r) == pytest.approx(energy, rel=0.02) + + def test_noisy_surface(self): + """The recovered activation energy stays close under 1% noise.""" + x, temps, values, energy, gas_r = self._surface(0.01, 17) + est = SurfaceDerivatives(smoothing="sigma").fit([x, temps], values, sigma=0.01) + assert self._effective_energy(est, gas_r) == pytest.approx(energy, rel=0.05) + + def test_oversmoothing_is_visible(self): + """Deliberate over-smoothing lowers the effective dof and is reported.""" + x, temps, values, _, _ = self._surface(0.01, 18) + base = SurfaceDerivatives().fit([x, temps], values, sigma=0.01) + rough = SurfaceDerivatives(smoothing_scale=20.0).fit([x, temps], values, sigma=0.01) + assert rough.effective_dof_ < base.effective_dof_ + assert rough.smoothing_ > base.smoothing_ + assert f"{rough.smoothing_:.6g}" in rough.summary()