JaxNumPy functions for GridArrays/GridVariables about jax-cfd HOT 2 OPEN

Hi @gemmaellen -- thanks for your interest!

This was an intentional design choice -- GridVariables have boundary conditions, which we don't know how to propagate automatically (unless using periodic boundaries, which aren't really boundary conditions at all). So we only support math on GridArray objects.

Also, am I correct in thinking that the way to use a JaxNumPy function on a GridArray is to call it via NumPy?

This is correct, I agree it's strange. It's for the simple reason that NumPy supports overriding it's functions on new types but JAX doesn't.

from jax-cfd.

Oh, I see! Yes, addition and multiplication would also work for matching homogeneous Dirichet/Neumann BCs, but that's a special case, and it wouldn't extend to other functions like sines and cosines and so on. I assume this is also the reason why the "shift" method on a GridVariable returns a GridArray. Thanks for the explanation!

from jax-cfd.