xinvert

1. Introduction

Researches on meteorology and oceanography usually encounter inversion problems that need to be solved numerically. One of the classical inversion problem is to solve Poisson equation for a streamfunction $\psi$ given the vertical component of vorticity $\zeta$ and proper boundary conditions.

$$\nabla^2\psi=\zeta$$

Nowadays xarray becomes a popular data structure commonly used in Big Data Geoscience. Since the whole 4D data, as well as the coordinate information, are all combined into xarray, solving the inversion problem become quite straightforward and the only input would be just one xarray.DataArray of vorticity. Inversion on the spherical earth, like some meteorological problems, could utilize the spherical harmonics like windspharm, which would be more efficient using FFT than SOR used here. However, in the case of ocean, SOR method is definitely a better choice in the presence of irregular land/sea mask.

More importantly, this could be generalized into a numerical solver for elliptical equation using SOR method, with spatially-varying coefficients. Various popular inversion problems in geofluid dynamics will be illustrated as examples.

One problem with SOR is that the speed of iteration using explicit loops in Python will be e-x-t-r-e-m-e-l-y ... s-l-o-w! A very suitable solution here is to use numba. In addition, through using xarray's dask backend for parallel computing, the speed of the inversion could be further improved using multiple CPU cores. Now, we could also make use of CUDA GPU for further acceleration, with a Red-Black SOR kernel achieving up to ~17× speedup over the CPU for large grids (measured on an RTX 3090; see the benchmark docs). See this notebook for more details of CPU/GPU parallel computing, and this docs for benchmark of the inversion speed.

Experimental GPU backend: CUDA support is available for testing and evaluation, but is not yet part of xinvert's stable API. Kernel behavior, configuration options, and performance characteristics may still change. Use the CPU backend for production workflows that require the stable path.

Classical problems include Gill-Matsuno model, Stommel-Munk model, QG omega model, PV inversion model, Swayer-Eliassen balance model... A complete list of the classical inversion problems can be found at this notebook.

Why xinvert?

• Thinking and coding in equations: User APIs are very close to the equations: unknowns are on the LHS of =, whereas the known forcings are on its RHS;
• Genearlize all the steady-state problems: All the known steady-state problems in geophysical fluid dynamics can be easily adapted to fit the solvers;
• Very short parameter list: Passing a single xarray forcing is enough for the inversion. Coordinates information is already encapsulated.
• Flexible model parameters: Model parameters can be either a constant, or varying with a specific dimension (like Coriolis $f$), or fully varying with space and time, due to the use of xarray's broadcasting capability;
• Parallel inverting: The use of xarray, and thus dask allow multi-core CPU parallel inverting, as well as GPU inverting. One only needs to change the parameter architect='CPU'/'GPU'/'CPU-mc', and the internal implementation is almost transparent to the user;
• Pure Python code for C-code speed: The use of numba allow pure python code in this package but native speed;

2. How to install

Requirements xinvert supports Python 3.9 and newer and requires xarray, dask, numpy, and numba. GPU acceleration additionally requires an NVIDIA GPU with a compatible driver/runtime and the maintained numba-cuda package.

Install via pip

pip install xinvert

Install GPU support

pip install "xinvert[gpu]"                 # existing CUDA runtime
pip install xinvert "numba-cuda[cu12]"    # install CUDA 12 runtime libraries

Install via conda

conda install -c conda-forge xinvert

Install from github

git clone https://github.com/miniufo/xinvert.git
cd xinvert
python -m pip install .

3. Examples:

The classic Poisson equation for streamfunction can be inverted by:

from xinvert import invert_Poisson

iParams = {
    'BCs'      : ['fixed', 'periodic'], # fixed north/south and periodic east/west BCs
    'mxLoop'   : 1000,                  # maximum loop count
    'tolerance': 1e-6,                  # for dtype float32
    'architect': 'GPU',                 # use GPU
}

sf = invert_Poisson(vorticiy, dims=['lat','lon'], iParams=iParams)

This is a list of the problems that can be solved by xinvert:

Gallery	Gallery
invert Poisson equation for horizontal streamfunction	invert Poisson equation for overturning streamfunction
invert geostrophic equation for balanced mass	invert Eliassen model for overturning streamfunction
invert PV balance equation for steady reference state	invert Gill-Matsuno model for wind and mass fields
invert Stommel-Munk model for wind-driven ocean circulation	invert Fofonoff model for inviscid/adiabatic steady state
invert Bretherton model for steady flow over topography	invert Omega equation for QG vertical velocity
invert geostrophic-adjustment problem	... add more examples ...

4 Animate the convergence of iteration

One can see the whole convergence process of SOR iteration as:

from xinvert import animate_iteration

# output has 1 more dimension (iter) than input, which could be animated over.
# Here 40 frames and loop 1 per frame (final state is after 40 iterations) is used.
psi = animate_iteration(invert_Poisson, vor, iParams=iParams,
                        loop_per_frame=1, max_frames=40)

See the animation at the top.

Xinvert uses a normalized, diagonally preconditioned residual for stopping. Its history can optionally be returned for analysis:

iParams = {
    'tolerance': 1e-8,
    'diag_interval': 20,  # diagnose every 20 SOR iterations
}
psi, diagnostics = invert_Poisson(
    vor, dims=['lat', 'lon'], iParams=iParams)

diagnostics.error.plot(x='iteration', yscale='log')

Through plotting the error, one can observe the convergence rate at regular interation interval. The default diag_interval=0 returns only the solution. See diagnostic analysis documentation for more details.

5 Cite

If you use the package in research, teaching, or other activities, we would be grateful if you mention xinvert and cite our paper in JOSS:

@article{Qian2023,
    doi = {10.21105/joss.05510},
    url = {https://doi.org/10.21105/joss.05510},
    year = {2023},
    publisher = {Journal of Open Source Software},
    volume = {8},
    number = {89},
    pages = {5510},
    author = {Yu-Kun Qian},
    title = {xinvert: A Python package for inversion problems in geophysical fluid dynamics}, journal = {Journal of Open Source Software}
}