A functional API for the standard spatial operator types that we typically see in numerical methods. We will target the following methods:

• Finite Difference - just a standard wrapper for FiniteDiffX
• Finite Volume - needs interpolation and simple finite difference operators
• Spectral Methods - just a standard wrapper for FiniteDiffX

——
Grid Operators

These are operators that are useful for finite volume methods to do grid to grid interpolation and gradients. It is also useful for finite difference/volume methods for padding and boundary conditions.

API

• Grid.interp
• Grid.diff
• Grid.pad
• (Grid.cumsum)
• grid.boundary

u_avg = grid.interp(u, axis=0, method="mean", operation=“linear”)

du_dx = grid.difference(u, step_size=dx, axis=0, method="")

u_pad = grid.pad(u, axis=0, method="constant", values=0)

u_dirichlet = grid.boundary(u, axis=0, method=“dirichlet”)

In this library (findiff), they have a nice feature to generate matrices and stencils from coefficients.

It also might be helpful to add staggered coefficients as is done in jaxdf.

Background

Examples

• 1D Example - Example | Example
• North Atlantic Simulation - #31

Resources

https://github.com/dionhaefner/shallow-water/blob/master/shallow_water_simple.py

https://github.com/dionhaefner/shallow-water/blob/master/shallow_water_nonlinear.py

https://github.com/karnerb/semi-implicit-SWE/blob/main/swe.py

https://github.com/hongkangcarl/PDE-solver-shallow-water-equations/blob/main/swe.py

https://github.com/milankl/swm

https://github.com/Zanna-ResearchTeam/swe_stochastic_param/blob/master/shallowwater.py

https://github.com/Zanna-ResearchTeam/swe_stochastic_param/blob/master/shallowwaterParameterized.py

https://github.com/mrocklin/ShallowWater/blob/master/shallowwater_simple.py

http://basilisk.fr/sandbox/popinet/gulf-stream.md

https://github.com/jamesp/shallowwater

https://github.com/jostbr/shallow-water

https://github.com/daniel-koehn/Differential-equations-earth-system/tree/master/10_Shallow_Water_Equation_2D

https://github.com/leguillf/MASSH/blob/main/mapping/models/model_sw1l/jswm.py

https://github.com/jamesp/gfdbyexample/blob/master/geostrophic_adjustment.py

https://github.com/maul1609/shallow-water-model-practical/blob/master/python/shallow_water_model.py

https://github.com/mehradans92/Shallow-water-dynamics/blob/master/SWE_2D.py

This is the canonical 12 Steps to Navier-Stokes tutorial. All of this will be implemented using the package using the:

• Low-Level API - FiniteDiffX and Diffrax (PRIORITY)
• Mid-level API - kernex (kmap | kscan)
• High-level API - PDiffrax

Steps

1D Problems

• 1D Linear Advection
• 1D Non-Linear Advection
• 1D Diffusion
• 1D Burgers

2D Problems

• 2D Linear Advection
• 2D Non-Linear Advection
• 2D Diffusion
• 2D Burgers

Hyperbolic Problems

• Laplace Equation
• Poisson Equation

Navier-Stokes

• 2D Navier-Stokes - Cavity Flow
• 2D Navier-Stokes - Channel Flow

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.

Examples

• https://github.com/miniufo/xinvert/blob/master/xinvert/finitediffs.py
• https://github.com/team-ocean/veros

There is a new update in the kernex package that allows for padding keyword arguments (see here). I’d like to add those options to the kernel functions I constructed for moving across grids.

Add a simple Shallow water model implementation

Add sea surface height conversions.

Needed

• ssh to potential vorticity
• ssh to psi
• ssh to u, v
• psi to u,v
• kinetic energy
• enstrophy
• strain
• geostropic currents

Add the Quasi-Geostrophic Model

• Standard Model
• Test Case

An overview of the learning problems available

Parameter Estimation

$$ \boldsymbol{\theta}^* = \underset{\boldsymbol{\theta}}{\text{argmin}}\hspace{2mm}\mathcal{L}(\boldsymbol{\theta}) $$

• ODE Parameters (e.g. Lorenz63, Lorenz96)
• PDE Parameters (e.g. Diffusivity Coefficient)
• UDE (e.g. QG, SW)
  • Parameterizations
  • Hybrid Models
  • Surrogate Model

State Estimation

$$ \mathbf{x}^* = \underset{\mathbf{x}}{\text{argmin}}\hspace{2mm}\mathcal{U}(\mathbf{x}) $$

• Inverse Problem
  • Gradients - Unrolling, Implicit Diff
• Noisy Observations
• Missing Observations

Bi-Level Optimizations

• Inverse Problem
  • Gradients - Unrolling, Implicit Diff/Adjoint
• 4DVarNet (Learning-to-Learn)

• Add setitem to field - Example
• Handle Interior Points

Lorenz Family Models

• Lorenz 63
• Stochastic Lorenz 64
• Lorenz-84
• Lorenz 96 - 1 Level
• Lorenz 96 - 2 Level
• Stochastic Lorenz 96 - 2 Level
• Lorenz 96 - 3 Level

Resources

Lorenz 84

• See repo for a nice formulation.

Stochastic Lorenz-63

• https://arxiv.org/abs/1706.05882
• https://rmets.onlinelibrary.wiley.com/doi/10.1002/qj.3198

• Naive
• Energy-Conserving Scheme, Sadourny, 1975 - SWM rhs
• Arakawa and Lamb, 1981 - SWM RHS
• Upwind
• WENO

Showcase different ways to optimise for the hyperparameters

Schemes

• Standard (Unrolling)
• Adjoint
• Implicit Differentiation

This will feature some fairly generic PDE terms that keep cropping up that we may want to solve.

Diffusion

Some generic Diffusion equation ν (∂²η/∂x² + ∂²η/∂y²) that we may want to solve

$$ \begin{aligned} &=\nu \left( \frac{\partial^2 \eta}{\partial x^2} + \frac{\partial^2 \eta}{\partial y^2} + \frac{\partial^2 \eta}{\partial z^2}\right) \end{aligned} $$

• Naive
  • 1D
  • 2D

Advection

Some generic advection terms that we may want to solve, i.e. u ∂η/∂x + v ∂η/∂y + w ∂η/∂z

$$ \begin{aligned} &=u \frac{\partial \eta}{\partial x} + v \frac{\partial \eta}{\partial y} + w \frac{\partial \eta}{\partial z} \end{aligned} $$

• Naive
  • 1D
  • 2D
• Upwind
  • 1D
  • 2D
• #21

Note: I've also seen instantiations of this with the name "determinant Jacobian".

• Chebyshev Derivative
• Pseudospectral Derivative | jaxdf

Better intro to shallow water model linking the burgers equation and Navier stokes (particular depth Discretization)

Formulation

• Examples Lecture | Lecture

Examples

• Double Gyres - Notebooks
• Tsunamis - Notebooks

This is my mega-task list where I just dump all of my ideas into a single issue. At a later date, I will break these up into much smaller issues.

Package Elements

This will go over some specifics of the jaxsw package to showcase how it works and some things one can do. It will be taught all within the context of PDEs.

5 Levels of Granularity of PDE's

Background

• Theory
• Formulation - Example | Example

Examples

• 1 Layer Idealistic Gyre
• 1 Layer Idealistic Jet

——

https://github.com/SusanEAllen/numeric-1/blob/master/numlabs/lab8/qg.py

This showcases some of the different boundary conditions ones could encounter. This is easily the most underrated part of PDEs in general and cause most of the problems.

• Periodic
• Dirichlet
• Neumann
• Robin
• Open
• Custom

Comments From Thiery

• Check potential vorticity term - $q=\hat{k}\cdot \nabla\times u$
• Check Hyperdiffusivity term $\Delta^3$

In this tutorial, we will walk-through the different abstract components of a PDE and how they all come together using a minimal example.

• Domain
• State
• Equation of Motion (RHS)
• Parameters
• Boundary Conditions
• Spatial Discretization
• Time Stepper

Some simple tutorials for using the package for a data assimilation problem.

• Backwards-Forwards Nudging
• 4DVariational
  • Weak-Constrained (Step)
  • Strong-Constrained (Trajectory)

Show Step-by-Step way of using TimeSteppers with the framework.

• From Scratch - Jax + Scan
• Simplistic Jax Solver - odeint
• Generic Solvers - diffrax
  • Step Solver | KF
  • Step + Scan Solver
  • Full
  • Custom Solver - Crank-Nicholson | LeapFrog | Adam-Bashforth
• Probabilistic Solvers - probnum

Add some of the core finite difference operations. The full slicing version can be used in demos then then we can use the proper packages like serket and kernex.

Key Packages

• serket - has different FD operators
• kernex - has the convolution version for the FD operator

Core Functionality

• Gradient
• Laplacian
• Determinant Jacobian
• Divergence
• Curl
• kernex functionality

This is a high level API to represent the Arakawa C-grid where we have the variable of interest, the u-v velocities, and the tracer. It should be easy for people to move between the 4 grids and calculate differences. It will handle all of the ghost points and boundaries.

This tutorial showcases how granular/explicit one can be when defining a PDE. It will start with writing everything from scratch with loops and little by little start automating things here and there. It's a great exercise to teach people different levels of abstraction that are appropriate for different use cases.

• Scratch - Loops w/ Numba
• Functional
  • Spatial Discretisation - Slicing w/ FiniteDiffX
  • TimeStepper - Scan w/ JAX
• Automated w/ Kernex
  • Spatial Discretisation w/ kmap
  • TimeStepping w/ kscan
• Explicit
  • Spatial Discretisation w/ finitediffX
  • Time Stepping w/ diffrax
• Implicit
  • Discretisation w/ jaxdf-like
  • TimeStepping w/ jaxdf-like

I - Prebuilt Models

For just diving right in and using it!

II - Operator API

A medium level of granularity.

III - Functional API

The finest level of granularity available.

Need the obligatory Lorenz family of methods.

• Lorenz 63
• #23
• Lorenz 96

We go over the different ways to do the spatial operators.

• From Scratch (semantic) - kernex
• Slicing - FiniteDiffX
• Convolutions - Coefficients & Stencils | jax.conv, jaxdf
• Probabilistic Numerical Finite Differences - probfindiff
• Spectral - jaxdf
• Convolutional Neural Network - equinox | jaxdf

Currently the domain is regular. But it would be nice to add an optional staggering option when constructing the grid with some helpful methods to convert between the grids.

Grid Operators

This will make the functional grid operators to be compatible with fields and domains.

• Interpolate

——
Finite Differences (Slicing)

This will make the FiniteDiffX package compatible with the Field API.

• Difference
• Laplacian
• Divergence
• Curl 2D

—-
Finite Volume

This will make the functional grid operators to be compatible with fields and domains.

• Difference
• Laplacian
• Divergence
• Curl 2D

——
Finite Differences (Convolutions)

This will make the Convolutional finite difference API compatible with the Field API.

• Difference
• Laplacian - example
• Divergence
• Curl 2D

——
Spectral Differences (Convolutions)

This will make the Convolutional finite difference API compatible with the Field API.

• Difference
• Laplacian
• Divergence
• Curl 2D

Add a 2D Non-Linear Stacked Shallow Water Model

Resources:

• Basilisk Example
• Popinet, et al, 2020
  • Example Code in Basilisk

Some generic Hyperbolic functions that we may need to solve, e.g. Laplacian, Poisson. It's 2023 so we'll stick with iterative solvers that scale, e.g., the steepest descent methods, the conjugate gradient methods and the discrete sine transform methods.

$$ \begin{aligned} \boldsymbol{\nabla}^2\boldsymbol{u}&= 0 \\ \boldsymbol{\nabla}^2\boldsymbol{u}&= \boldsymbol{b} \end{aligned} $$

Methods From Scratch

I have don't a few methods from scratch but they are not well tested...

• Steepest Descent
• Conjugate Gradient
• Discrete Sine Transformation

External Packages

There are many external packages that can be used with general linear solvers

• Jaxopt
• Lineax (New)

This showcases some of the ways we can do linear solvers using the package.

• Exact Linear Solver
• Iterative Linear Solvers w/ jaxopt (Steepest Descent, Conjugate Gradient)
• Discrete Sine Transform
• Staggered Discrete Sine Transform
• Spectral

An overview of the different domains we can encounter for PDEs. The easy ones are the uniform grid-like scenarios and the more applied situations feature the spherical lat/lon/height cases. A very important improvement is the Staggered grid which will be useful for multivariate PDEs with velocities and such.

• Uniform Grid
• Latitude, Longitude (Meters)
• Latitude, Longitude (Spherical)
• Arakawa C-Grid

For long simulations, it might be good to make use of MPIJAX.

They have a SWM example.

Here, we showcase some examples of simple ocean models we can use within the package. These can be demonstrated with free-runs that are conditioned on real data.

Quasi-Geostrophic Equations

• 2D Model - Gyre
  • 1 Layer
  • 2 Layer
• 2D Model - Jet
  • 1 Layer - #26
  • 2 Layer - #26

Shallow Water Equations

• #30
• #8
• #31

Main Tutorials

These tutorials showcase how the user should think about PDEs using our abstraction.
Then they showcase how to use the software to comply with this abstraction.

• Anatomy of a PDE - #40
• Coding a PDE in 3 Ways - #38
• Spatial Operators - #41
• Boundary Value Operators - #42
• Grid Operations
• TimeSteppers - #39
• 12 Steps to Navier-Stokes - #11

Extended Deep Dives

These tutorials go over some of the nuances of the package that may help the user define special cases.

• Geometries - #43
• Modern Inversion Schemes - #24
• Experimental Scripts w/ Hydra
• Spherical Coordinates - #33
• FAQs & Gotchas
  • NANs
  • Non-Dimensionalization

Simple Ocean Models

These tutorials showcase the simplified ocean models.

• Lorenz Family - #23
• Quasi-Geostrophic Equation - #36
• Shallow Water Equations - #37

Learning Schemes - #13

These tutorials introduce the user to how one can do parameter and/or state estimation using the autodiff framework.

• Gradients & Optimization - #6
• State Estimation
• Parameter Estimation
• State & Parameter Estimation

• We could have more advanced averaging - example
• We could have a dedicated API - example
• A bit of theory - example
• Dedicated Arakawa C-Grid Module - example

More Case Studies

• Flood Mapping
• Tsunami Modeling