astronomix._stencil_operations._stencil_operations

Convenience functions for operations that combine multiple elements of an array based on some stencil, e.g. b_i <- a_{i + 1} + a_{i - 1}. Allows for code “closer to the math”.

Module Contents

Functions

custom_roll

Periodic roll of input_array by shift along axis.

API

astronomix._stencil_operations._stencil_operations.custom_roll(input_array: jax.numpy.ndarray, shift: int, axis: int) → jax.numpy.ndarray

Periodic roll of input_array by shift along axis.

Equivalent to jnp.roll but expressed via two static slices and a concatenate, which keeps shift / axis compile-time constants so the stencil helpers built on top of it fuse cleanly.

Args:

input_array: The array to roll. shift: The (signed) number of positions to roll by. axis: The axis along which to roll.

Returns:

The rolled array.