jax_sbgeom.jax_utils.splines module

bspline(x: Array, t: Array, c: Array, k: int, derivative: int)

Evaluate a B-spline at points x.

Note that it is required that c.shape[0] == t.shape[0] - k - 1

Parameters:

• x (Array) – Points at which to evaluate the B-spline.

• t (Array) – Knot vector

• c (Array) – Coefficients of the B-spline basis functions.

• k (int) – Degree of the B-spline.

• derivative (int) – Order of the derivative to compute

Returns:

Values of the B-spline at points x

Return type:

jnp.ndarray [M]

bspline_vectorized(x: Array, t: Array, c: Array, k: int, derivative: int)

Evaluate a B-spline at points x.

Note that it is required that c.shape[0] == t.shape[0] - k - 1

Parameters:

• x (Array) – Points at which to evaluate the B-spline.

• t (Array) – Knot vector

• c (Array) – Coefficients of the B-spline basis functions.

• k (int) – Degree of the B-spline.

• derivative (int) – Order of the derivative to compute

Returns:

Values of the B-spline at points x

Return type:

jnp.ndarray [M]

class BSpline(t: Array, c: Array, k: int)

Bases: object

Convenience class for representing an arbitrarily batched B-spline.

Data

tjnp.ndarray […, N]

Knot vector

cjnp.ndarray […, N - k - 1]

Coefficients of the B-spline basis functions.

kint

Degree of the B-spline.

t: Array

c: Array

k: int

periodic_interpolating_spline(x, y, k: int)

Fit a periodic interpolating B-spline to batched data points (x,y) of degree k. The resulting spline satisfies S(x[i]) = y[i] for all i and is periodic (including derivatives up to < k).

Parameters:

• x (jnp.ndarray [..., n]) – Data points to be interpolated.

• y (jnp.ndarray [..., n]) – Data values to be interpolated.

• k (int) – Degree of the B-spline.

Returns:

Fitted periodic interpolating B-spline.

Return type:

BSpline

make_periodic_knots(n: int, k: int, period: float = 6.283185307179586)

Build a periodic knot vector for n control points, degree k, on the domain [0, period). This is different than _periodic_knots, which builds a knot vector for periodic interpolation of given data points. This function is for building a knot vector for periodic B-spline parametrization.

Parameters:

• n (int) – Number of control points.

• k (int) – Degree of the B-spline.

• period (float) – Period of the B-spline. Default is 2*pi.

Returns:

Knot vector for periodic B-spline interpolation.

Return type:

Array

make_periodic_coeffs(c: Array, k: int)

Extend n coefficients to n+k by wrapping the first k, enforcing periodicity.

Parameters:

• c (Array) – Coefficients of the B-spline basis functions.

• k (int) – Degree of the B-spline.

Returns:

Extended coefficients for periodic B-spline interpolation.

Return type:

Array

periodic_bspline(theta: Array, c: Array, k: int, period: float = 6.283185307179586, derivative: int = 0)

Function to evaluate a periodic B-spline at angles theta given coefficients c and degree k.

Note that the coefficients c are the free parameters, which do not map directly to a control point given by their index. Use :func`greville_abscissa_periodic_bspline` to compute the control points location.

Parameters:

• theta (Array) – Angles at which to evaluate the B-spline, in radians.

• c (Array) – Coefficients of the B-spline basis functions. Must have length n, which is the number of free parameters. Internally extended to n+k for periodicity.

• k (int) – Degree of the B-spline.

• period (float) – Period of the B-spline. Default is 2*pi.

• derivative (int) – Order of the derivative to compute. Default is 0 (the function itself).

Returns:

Values of the periodic B-spline (or its derivative) at the angles theta.

Return type:

jnp.ndarray [M]

greville_abscissa_periodic_bspline(n: int, k: int, period: float = 6.283185307179586)

Return type:

Array

periodic_bspline_2d(theta, phi, c: Array, k_theta: int, k_phi: int, period_theta: float = 6.283185307179586, period_phi: float = 6.283185307179586)

Evaluate a 2D tensor-product periodic B-spline at a single point (theta, phi).

Parameters:

• theta (scalar) – Poloidal angle, in radians.

• phi (scalar) – Toroidal angle, in radians.

• c (Array) – Coefficient matrix.

• k_theta (int) – Degree in the theta direction.

• k_phi (int) – Degree in the phi direction.

• period_theta (float) – Period in the theta direction. Default is 2*pi.

• period_phi (float) – Period in the phi direction. Default is 2*pi.

Returns:

Value of the spline at (theta, phi).

Return type:

scalar

periodic_bspline_2d_vec(theta, phi, c: Array, k_theta: int, k_phi: int, period_theta: float = 6.283185307179586, period_phi: float = 6.283185307179586)

Vectorized version of periodic_bspline_2d for evaluating at multiple (theta, phi) points.

Just uses jnp.vectorize under the hood; this function wraps over it to ensure the correct number of arguments are always passed.

Parameters:

• theta (jnp.ndarray [...]) – Poloidal angles, in radians.

• phi (jnp.ndarray [...]) – Toroidal angles, in radians.

• c (Array) – Coefficient matrix.

• k_theta (int) – Degree in the theta direction.

• k_phi (int) – Degree in the phi direction.

• period_theta (float) – Period in the theta direction. Default is 2*pi.

• period_phi (float) – Period in the phi direction. Default is 2*pi.

Returns:

Values of the spline at each (theta[i], phi[j]).

Return type:

jnp.ndarray […]