xinvert: A Python package for inversion problems in geophysical fluid dynamics

Once ζ is given as a known, one needs to get the unknown ψ with proper boundary conditions, which is essentially an inversion problem (Figure 1). Many geophysical fluid dynamical (GFD) problems can be described by such a balanced model, which is generally of second (or fourth) order in spatial derivatives and does not depend explicitly on time. Therefore, it is also known as a steady-state model. Early scientists tried to find analytical solutions by simplifying the parameters of these models (e.g., assuming constant coefficients). Nowadays, with the new developments in numerical algorithms and parallel-computing programming, one may need a modern solver, written in a popular programming language like Python, to invert all these models in a numerical fashion. More specifically, the following needs should be satisfied:

Once is given as a known, one needs to get the unknown with proper boundary conditions, which is essentially an inversion problem ( Figure 1). Many geophysical fluid dynamical (GFD) problems can be described by such a balanced model, which is generally of second (or fourth) order in spatial derivatives and does not depend explicitly on time. Therefore, it is also known as a steady-state model. Early scientists tried to find analytical solutions by simplifying the parameters of these models (e.g., assuming constant coefficients). Nowadays, with the new developments in numerical algorithms and parallel-computing programming, one may need a modern solver, written in a popular programming language like Python, to invert all these models in a numerical fashion. More specifically, the following needs should be satisfied: • A unified numerical solver: It can solve all of the classical balanced GFD models, even in a domain with irregular boundaries like the ocean. New models can also be adapted straightforwardly to fit the solver; • Thinking and coding in equations: Users focus naturally on the key inputs and outputs of the GFD models, just like thinking of the knowns and unknowns of the PDEs; • Flexible parameter specification: Coefficients of the models can be either constant, 1D vectors, or ND arrays. This allows an easy reproduction of early simplified results and also an extension to more general/realistic results; • Fast and efficient: The algorithm should be fast and efficient. Also, the code can be compiled first and then executed as fast as C or FORTRAN, instead of executed in the slow pure Python interpreter. In addition, it can leverage the multi-core and out-of-core computational capabilities of modern computers; xinvert is designed to satisfy all of the above needs, based on the Python ecosystem.

State of the field
There are also several PDE solvers written in Python, like windspharm (Dawson, 2016) and Dedalus (Burns et al., 2020). While they are efficient and accurate in double-periodic or global domains using the spectral method, they may not be suitable for arbitrary domains or boundaries like the ocean (Gibbs phenomenon will arise near ocean boundaries). Also, it is not easy to apply these solvers to more general elliptic equations with arbitrary coefficients (sometimes given in a numerical or discretized fashion). In these cases, the successive over relaxation (SOR) method should be a better choice.

Mathematics
This package, xinvert, is designed to invert or solve the following PDE in an abstract form: where is a second-order partial differential operator, the unknown to be inverted for, and a prescribed forcing function. There could be also some coefficients or parameters in the definition of , which should be specified before inverting .
For the 2D case, the general form of Eq. (1) is: where coefficients 1 − 6 are all known. When the condition 4 1 3 − 2 2 > 0 is met everywhere in the domain, the above equation is an elliptic-type equation. In this case, one can invert using the SOR iteration method. When 4 1 3 − 2 2 = 0 or 4 1 3 − 2 2 < 0, it is a parabolic or hyperbolic equation. In either case, SOR would fail to converge to the solution.
Sometimes the general form of Eq. (2) can be transformed into the standard form (i.e., standardization): In this case, 1 4 − 2 3 > 0 should be met to ensure its ellipticity. The elliptic condition has its own physical meaning in the problems of interest. That is, the system is in a steady (or balanced) state that is stable to any small perturbation.
Many problems in meteorology and oceanography can be cast into the forms of either Eq.
(2) or Eq. (3). However, some of them are formulated in a 3D form (like the QG-omega equation): So we implement four basic solvers to take into account the above equations (2) to (5). Most of the problems fit one of these four types of solver. However, it is not clear which form, the general form Eq.
(2) or the standard form Eq. (3), is preferred for the inversion if a problem can be cast into either one.

Summary
xinvert is an open-source and user-friendly Python package that enables GFD scientists or interested amateurs to solve all possible GFD problems in a numerical fashion. With the ecosystem of open-source Python packages, in particular xarray (Hoyer & Hamman, 2017), dask (Rocklin, 2015), and numba (Lam et al., 2015), it is able to satisfy the above requirements: • All the classical balanced GFD models can be inverted by this unified numerical solver; • User APIs (Table 1) are very close to the equations: unknowns are on the left-hand side of the equal sign =, whereas the known forcing functions are on its right-hand side (other known coefficients are also on the left-hand side but are passed in through mParams); • Passing a single xarray.DataArray is usually enough for the inversion. The fact that coordinate information is already encapsulated reduces the length of the parameter list.
In addition, parameters in mParams can be either a constant, or varying with a specific dimension (like the Coriolis parameter ), or fully varying with space and time, due to the use of xarray's (Hoyer & Hamman, 2017) broadcasting capability; • This package leverages numba (Lam et al., 2015) and dask (Rocklin, 2015) to support Just-In-Time (JIT) compilation, multi-core, and out-of-core computations, and therefore greatly increases the speed and efficiency of the inversion.
Here we summarize some inversion problems in meteorology and oceanography into Table 1.  The table can be extended further if one finds more problems that fit the abstract form of Eq.