One immediate feature that emerged from the chat with @jejjohnson is the ability to work with fields in a way that allows hiding them from the user, or at least not explicitly working with them.
A common pattern for achieving this is given by the following code:
def my_awesome_func(u: jax.ArrayLike):
# Declare fields
N = u.shape
dx = [0.1,] * len(N)
u_field = FourierSeries(u, Domain(N, dx))
# Perform the desired operation using jaxdf
v_field = some_operator(u_field)
# Return a simple jax array
return v_field.on_grid
To simplify the syntax and achieve a cleaner implementation, this pattern can be encapsulated in a decorator, as shown below:
One immediate feature that emerged from the chat with @jejjohnson is the ability to work with fields in a way that allows hiding them from the user, or at least not explicitly working with them.
A common pattern for achieving this is given by the following code:
To simplify the syntax and achieve a cleaner implementation, this pattern can be encapsulated in a decorator, as shown below:
Here, the
use_discretization
decorator takes care of packing and unpacking the fields:Potential issues and things to work out
dx
in this exampleOnGrid
fields?