jejjohnson / jaxsw

Simple differentiable approximate ocean models built with JAX.
https://jejjohnson.github.io/jaxsw/
MIT License
12 stars 1 forks source link

Arakawa C-Grid #18

Closed jejjohnson closed 1 year ago

jejjohnson commented 1 year ago

I want a working demo which showcase how one can define an Arakawa C-Grid for the SW and QG PDEs. I think this a minimum bang for buck improvement we can do for solving these PDEs with simple finite difference schemes.

Example API

I like the API seen in this codebase. However, we would simply define the state with the appropriate staggering. We can then use the transformations in the grid functions module to move the variables around, i.e., u --> v, v --> u, u,v --> H, etc.

Example PDEs

Shallow Water Equations

The most obvious choice is the SW equations as we have to solve a system of 3 states simultaneously. Lots of potential to blow up with instabilities so the C-Grid should be very helpful.

Quasi-Geostrophic Equations

Based on an implementation from Louis Thiry, we can define the QG model using the staggered grid for the velocities, potential vorticity and streamfunction/SSH.