jax_sbgeom.jax_utils.numerical module

interpolate_fractions(s, nsurf)

Interpolate fractions for uniform sampling

Parameters:

• s (jnp.ndarray) – Normalized parameter(s) between 0 and 1

• nsurf (int) – Number of samples

Returns:

• i0 (jnp.ndarray) – Lower indices

• i1 (jnp.ndarray) – Upper indices

• ds (jnp.ndarray) – Fraction between i0 and i1

interpolate_array(x_interp, s)

Interpolate array for uniform sampling

Parameters:

• x_interp (jnp.ndarray [1D]) – Array to interpolate

• s (jnp.ndarray [1D]) – Normalized parameter(s) between 0 and 1

Returns:

Interpolated array

Return type:

jnp.ndarray

interpolate_fractions_modulo(s, nsurf)

Interpolate fractions for uniform sampling with modulo wrapping I.e., s=1 maps to index 0 again.

Parameters:

• s (jnp.ndarray) – Normalized parameter(s) between 0 and 1

• nsurf (int) – Number of samples

Returns:

• i0 (jnp.ndarray) – Lower indices

• i1 (jnp.ndarray) – Upper indices

• ds (jnp.ndarray) – Fraction between i0 and i1

robust_grid_index_modulo(s, nsurf)

Robustly select the (modulo-wrapped) grid index at or below s*nsurf, for direct (non-blended) lookups.

Unlike interpolate_fractions_modulo, this returns only an index and is meant for callers that pick a single grid point rather than interpolate between two – so it is worth guarding against floating-point noise pushing s*nsurf just below an intended integer boundary (e.g. 57.999999999999996 instead of 58.0). This happens whenever s is deliberately constructed to land exactly on a grid point (e.g. via jnp.linspace(0, 1, nsurf, endpoint=False)), and floor() would otherwise silently select the wrong (previous) index. The tolerance scales with the dtype’s precision and nsurf, since the absolute error in s*nsurf grows with the magnitude of the result.

Parameters:

• s (jnp.ndarray) – Normalized parameter(s) between 0 and 1

• nsurf (int) – Number of samples

Returns:

Grid index/indices (not modulo-wrapped; wrap with % nsurf as needed)

Return type:

jnp.ndarray

interpolate_array_modulo(x_interp, s)

Interpolate array for uniform sampling with modulo wrapping I.e., s=1 maps to index 0 again. :param x_interp: Array to interpolate :type x_interp: jnp.ndarray [1D] :param s: Normalized parameter(s) between 0 and 1 :type s: jnp.ndarray [1D]

Returns:

Interpolated array

Return type:

jnp.ndarray

interpolate_array_modulo_broadcasted(x_interp, s)

Interpolate array for uniform sampling with modulo wrapping I.e., s=1 maps to index 0 again. This version supports broadcasting of s to higher dimensions. :param x_interp: Array to interpolate :type x_interp: jnp.ndarray [s.shape, interpolation_dimension, :] :param s: Normalized parameter(s) between 0 and 1 :type s: jnp.ndarray [s.shape]

Returns:

jnp.ndarray[s.shape, – Interpolated array

Return type:

]

reverse_except_begin(array: Array)

bilinear_interp(norm_array_0, norm_array_1, interpolate_array)

Bilinear interpolation for uniform sampling in 2D It is assumed that interpolate_array is defined on a uniform grid normalised to 1 in both dimensions.

Parameters:

• norm_array_0 (jnp.ndarray [shape]) – Normalized parameter(s) between 0 and 1 in first dimension

• norm_array_1 (jnp.ndarray [shape]) – Normalized parameter(s) between 0 and 1 in second dimension

• interpolate_array (jnp.ndarray [N0, N1]) – Array to interpolate

Returns:

Interpolated array

Return type:

jnp.ndarray [shape]

interp1d_jax(x, y, x_new)

Simple 1D linear interpolation function in JAX

Extrapolates flat with the boundary values.

Parameters:

• x (jnp.ndarray [N]) – x-coordinates of the data points

• y (jnp.ndarray [N]) – y-coordinates of the data points

• x_new (jnp.ndarray [M]) – x-coordinates where to interpolate

Returns:

Interpolated y-coordinates at x_new

Return type:

jnp.ndarray [M]

pchip_interpolation(x, y, x_new)

Piecewise cubic Hermite interpolating polynomial (PCHIP) interpolation in JAX

Convenience function: simply calls _pchip_derivatives and _pchip_evaluation If you need derivatives or evaluation, you can jax.grad and vectorize the _pchip_evalution function. If you need multiple calls with the same x,y data, compute the derivatives once with _pchip_derivatives and pass them to _pchip_evaluation. If you need to do this for many different x,y datasets, consider vmap’ing the _pchip_derivatives function.

Parameters:

• x (jnp.ndarray [N]) – x-coordinates of the data points

• y (jnp.ndarray [N]) – y-coordinates of the data points

• x_new (jnp.ndarray [M]) – x-coordinates where to interpolate

Returns:

Interpolated y-coordinates at x_new

Return type:

jnp.ndarray [M]

cumulative_trapezoid_uniform_periodic(y: Array, dx: float, initial: float = 0.0)

Cumulative trapezoidal integration of y with respect to x, assuming uniform spacing and periodicity.

The y is to be sampled at uniform intervals in x, with spacing dx and not including the endpoint. i.e., jnp.linspace(0, period, n_samples, endpoint=False)

Parameters:

• y (Array) – Values to integrate

• dx (float) – Spacing between samples in x

• initial (float) – Initial value of the integral

Returns:

Cumulative integral of y with respect to x

Return type:

jnp.ndarray [N+1]

resample_uniform_periodic_linear(non_uniform_values: Array, n_points_desired: int)

Resample non-uniform values to uniform values using linear interpolation. The input values are assumed periodic and sampled on a uniform grid.

Parameters:

• non_uniform_values (Array) – Non-uniform values to resample

• n_points_desired (int) – Number of desired uniformly sampled points

Returns:

Resampled uniform values

Return type:

jnp.ndarray [n_points_desired]

resample_uniform_periodic_pchip(non_uniform_values: Array, n_points_desired: int)

Resample non-uniform values to uniform values using PCHIP interpolation. The input values are assumed periodic and sampled on a uniform grid.

Parameters:

• non_uniform_values (Array) – Non-uniform values to resample

• n_points_desired (int) – Number of desired uniformly sampled points

Returns:

Resampled uniform values

Return type:

jnp.ndarray [n_points_desired]