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

.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

.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.

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

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
 

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:

docs/api/derivatives.rst

@@ -0,0 +1,7 @@
+jaxsr.derivatives
+=================
+
+.. automodule:: jaxsr.derivatives
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api/index.rst

@@ -18,3 +18,4 @@ Full API documentation generated from source docstrings.
    plotting
    classifier
    additive
+   derivatives