grater_jax.disk_model.interpolated_univariate_spline

interpolated_univariate_spline.py

JAX implementation of 1D spline interpolation.

This module provides a JAX-compatible version of SciPy’s InterpolatedUnivariateSpline with support for linear, quadratic, and cubic splines. It includes evaluation, derivatives, antiderivatives, and definite integrals, while omitting unsupported options such as weights and smoothing.

Functions

array(object[, dtype, copy, order, ndmin, ...])	Convert an object to a JAX array.
concatenate(arrays[, axis, dtype])	Join arrays along an existing axis.
ones(shape[, dtype, device])	Create an array full of ones.
register_pytree_node_class(cls)	Extends the set of types that are considered internal nodes in pytrees.
vmap(fun[, in_axes, out_axes, axis_name, ...])	Vectorizing map.
zeros(shape[, dtype, device])	Create an array full of zeros.

Classes

InterpolatedUnivariateSpline(x, y[, k, ...])

class grater_jax.disk_model.interpolated_univariate_spline.InterpolatedUnivariateSpline(x, y, k=3, endpoints='not-a-knot', coefficients=None)

Bases: object

tree_flatten()

Flatten the spline into JAX-compatible pytree children and auxiliary data.

Returns:

A 2-tuple (children, aux_data) where children is (x, y, coefficients) and aux_data is a dict with 'endpoints' and 'k'.

Return type:

tuple

classmethod tree_unflatten(aux_data, children)

Reconstruct a spline from JAX pytree children and auxiliary data.

Parameters:

• aux_data (dict) – Dictionary with keys 'endpoints' and 'k' as produced by tree_flatten.

• children (tuple) – Tuple (x, y, coefficients) as produced by tree_flatten.

Returns:

Reconstructed spline object.

Return type:

InterpolatedUnivariateSpline

derivative(x, n=1)

Analytic nth derivative of the spline.

The spline has derivatives up to its order k.

antiderivative(xs)

Computes the antiderivative of first order of this spline

integral(a, b)

Compute a definite integral over a piecewise polynomial. :param a: Lower integration bound :type a: float :param b: Upper integration bound :type b: float

Returns:

ig – Definite integral of the piecewise polynomial over [a, b]

Return type:

array_like