jaxsr.derivatives

jaxsr.derivatives#

Multivariate derivative estimation for surface and PDE-style discovery.

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.

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 \(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 SurfaceDerivatives.smoothing_ and by SurfaceDerivatives.summary(), so smoothing-induced bias is visible rather than inferred.

class jaxsr.derivatives.SurfaceDerivatives(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)#

Bases: object

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

coords_#

Training coordinates.

Type:

np.ndarray of shape (n_points, n_dims)

values_#

Training values.

Type:

np.ndarray of shape (n_points,)

n_features_in_#

Number of coordinate dimensions.

Type:

int

smoothing_#

Smoothing level actually used (lambda, bandwidth, or noise variance, depending on method).

Type:

float

smoothing_source_#

How smoothing_ was chosen: "gcv", "marginal_likelihood", "sigma", or "fixed".

Type:

str

effective_dof_#

Effective degrees of freedom of the fitted smoother.

Type:

float

residual_std_#

Residual standard deviation of the fit at the training points.

Type:

float

noise_std_#

Noise level used for the reported uncertainties: the supplied sigma (its root-mean-square if per-point) when given, otherwise residual_std_.

Type:

float

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)
fit(coords: ndarray | Sequence[ndarray], values: ndarray, sigma: float | 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:

The fitted estimator (self).

Return type:

SurfaceDerivatives

Raises:

ValueError – If the inputs are inconsistent, too small for the requested smoother, or if smoothing="sigma" was requested without sigma.

predict(coords: ndarray, return_std: bool = False) → ndarray | tuple[ndarray, 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.

derivatives(coords: ndarray, order: Any, return_std: bool = False) → tuple[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.

summary() → 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:

Multi-line summary.

Return type:

str

Raises:

RuntimeError – If the estimator has not been fitted.

jaxsr.derivatives.estimate_partial_derivatives(coords: ndarray | Sequence[ndarray], values: ndarray, order: Any, method: str = 'tensor_spline', sigma: float | ndarray | None = None, query: ndarray | Sequence[ndarray] | None = None, return_std: bool = False, **kwargs: Any) → tuple[ndarray, ...]#

Estimate partial derivatives of a smoothed N-D surface in one call.

Convenience wrapper around 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 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 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)