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

API

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