Introduction

PyMPDATA is a high-performance Numba-accelerated Pythonic implementation of the MPDATA algorithm of Smolarkiewicz et al. used in geophysical fluid dynamics and beyond. MPDATA numerically solves generalised transport equations - partial differential equations used to model conservation/balance laws, scalar-transport problems, convection-diffusion phenomena. As of the current version, PyMPDATA supports homogeneous transport in 1D, 2D and 3D using structured meshes, optionally generalised by employment of a Jacobian of coordinate transformation or a fluid density profile. PyMPDATA includes implementation of a set of MPDATA variants including the non-oscillatory option, infinite-gauge, divergent-flow, double-pass donor cell (DPDC) and third-order-terms options. It also features support for integration of Fickian-terms in advection-diffusion problems using the pseudo-transport velocity approach. In 2D and 3D simulations, domain-decomposition is used for multi-threaded parallelism.

PyMPDATA is engineered purely in Python targeting both performance and usability, the latter encompassing research users', developers' and maintainers' perspectives. From researcher's perspective, PyMPDATA offers hassle-free installation on multitude of platforms including Linux, OSX and Windows, and eliminates compilation stage from the perspective of the user. From developers' and maintainers' perspective, PyMPDATA offers a suite of unit tests, multi-platform continuous integration setup, seamless integration with Python development aids including debuggers and profilers.

PyMPDATA design features a custom-built multi-dimensional Arakawa-C grid layer allowing to concisely represent multi-dimensional stencil operations on both scalar and vector fields. The grid layer is built on top of NumPy's ndarrays (using "C" ordering) using the Numba's @jit(nopython=True) functionality for high-performance array traversals. It enables one to code once for multiple dimensions, and automatically handles (and hides from the user) any halo-filling logic related with boundary conditions. Numba prange() functionality is used for implementing multi-threading (it offers analogous functionality to OpenMP parallel loop execution directives). The Numba's deviation from Python semantics rendering closure variables as compile-time constants is extensively exploited within PyMPDATA code base enabling the just-in-time compilation to benefit from information on domain extents, algorithm variant used and problem characteristics (e.g., coordinate transformation used, or lack thereof).

Tutorial (in Python, Julia, Rust and Matlab)

Options class

The Options class groups both algorithm variant options as well as some implementation-related flags that need to be set at the first place. All are set at the time of instantiation using the following keyword arguments of the constructor (all having default values indicated below):

• n_iters: int = 2: number of iterations (2 means upwind + one corrective iteration)
• infinite_gauge: bool = False: flag enabling the infinite-gauge option (does not maintain sign of the advected field, thus in practice implies switching flux corrected transport on)
• divergent_flow: bool = False: flag enabling divergent-flow terms when calculating antidiffusive velocity
• nonoscillatory: bool = False: flag enabling the non-oscillatory or monotone variant (a.k.a flux-corrected transport option, FCT)
• third_order_terms: bool = False: flag enabling third-order terms
• epsilon: float = 1e-15: value added to potentially zero-valued denominators
• non_zero_mu_coeff: bool = False: flag indicating if code for handling the Fickian term is to be optimised out
• DPDC: bool = False: flag enabling double-pass donor cell option (recursive pseudovelocities)
• dimensionally_split: bool = False: flag disabling cross-dimensional terms in antidiffusive velocity
• dtype: np.floating = np.float64: floating point precision

For a discussion of the above options, see e.g., Smolarkiewicz & Margolin 1998, Jaruga, Arabas et al. 2015 and Olesik, Arabas et al. 2020 (the last with examples using PyMPDATA).

In most use cases of PyMPDATA, the first thing to do is to instantiate the Options class with arguments suiting the problem at hand, e.g.:

using Pkg
Pkg.add("PyCall")
using PyCall
Options = pyimport("PyMPDATA").Options
options = Options(n_iters=2)

Options = py.importlib.import_module('PyMPDATA').Options;
options = Options(pyargs('n_iters', 2));

use pyo3::prelude::*;
use pyo3::types::{IntoPyDict, PyDict, PyTuple};

fn main() -> PyResult<()> {
  Python::with_gil(|py| {
    let options_args = [("n_iters", 2)].into_py_dict_bound(py);
    let options = py.import_bound("PyMPDATA")?.getattr("Options")?.call((), Some(&options_args))?;

from PyMPDATA import Options
options = Options(n_iters=2)

Arakawa-C grid layer and boundary conditions

In PyMPDATA, the solution domain is assumed to extend from the first cell's boundary to the last cell's boundary (thus the first scalar field value is at $[\Delta x/2, \Delta y/2]$. The ScalarField and VectorField classes implement the Arakawa-C staggered grid logic in which:

• scalar fields are discretised onto cell centres (one value per cell),
• vector field components are discretised onto cell walls.

The schematic of the employed grid/domain layout in two dimensions is given below (with the Python code snippet generating the figure as a part of CI workflow):

import numpy as np
from matplotlib import pyplot

dx, dy = .2, .3
grid = (10, 5)

pyplot.scatter(*np.mgrid[
        dx / 2 : grid[0] * dx : dx,
        dy / 2 : grid[1] * dy : dy
    ], color='red',
    label='scalar-field values at cell centres'
)
pyplot.quiver(*np.mgrid[
        0 : (grid[0]+1) * dx : dx,
        dy / 2 : grid[1] * dy : dy
    ], 1, 0, pivot='mid', color='green', width=.005,
    label='vector-field x-component values at cell walls'
)
pyplot.quiver(*np.mgrid[
        dx / 2 : grid[0] * dx : dx,
        0: (grid[1] + 1) * dy : dy
    ], 0, 1, pivot='mid', color='blue', width=.005,
    label='vector-field y-component values at cell walls'
)
pyplot.xticks(np.linspace(0, grid[0]*dx, grid[0]+1))
pyplot.yticks(np.linspace(0, grid[1]*dy, grid[1]+1))
pyplot.title(f'staggered grid layout (grid={grid}, dx={dx}, dy={dy})')
pyplot.xlabel('x')
pyplot.ylabel('y')
pyplot.legend(bbox_to_anchor=(.1, -.1), loc='upper left', ncol=1)
pyplot.grid()
pyplot.savefig('readme_grid.png')

The __init__ methods of ScalarField and VectorField have the following signatures:

• ScalarField(data: np.ndarray, halo: int, boundary_conditions)
• VectorField(data: Tuple[np.ndarray, ...], halo: int, boundary_conditions) The data parameters are expected to be Numpy arrays or tuples of Numpy arrays, respectively. The halo parameter is the extent of ghost-cell region that will surround the data and will be used to implement boundary conditions. Its value (in practice 1 or 2) is dependent on maximal stencil extent for the MPDATA variant used and can be easily obtained using the Options.n_halo property.

As an example, the code below shows how to instantiate a scalar and a vector field given a 2D constant-velocity problem, using a grid of 24x24 points, Courant numbers of -0.5 and -0.25 in "x" and "y" directions, respectively, with periodic boundary conditions and with an initial Gaussian signal in the scalar field (settings as in Fig. 5 in Arabas et al. 2014):

ScalarField = pyimport("PyMPDATA").ScalarField
VectorField = pyimport("PyMPDATA").VectorField
Periodic = pyimport("PyMPDATA.boundary_conditions").Periodic

nx, ny = 24, 24
Cx, Cy = -.5, -.25
idx = CartesianIndices((nx, ny))
halo = options.n_halo
advectee = ScalarField(
    data=exp.(
        -(getindex.(idx, 1) .- .5 .- nx/2).^2 / (2*(nx/10)^2)
        -(getindex.(idx, 2) .- .5 .- ny/2).^2 / (2*(ny/10)^2)
    ),
    halo=halo,
    boundary_conditions=(Periodic(), Periodic())
)
advector = VectorField(
    data=(fill(Cx, (nx+1, ny)), fill(Cy, (nx, ny+1))),
    halo=halo,
    boundary_conditions=(Periodic(), Periodic())
)

ScalarField = py.importlib.import_module('PyMPDATA').ScalarField;
VectorField = py.importlib.import_module('PyMPDATA').VectorField;
Periodic = py.importlib.import_module('PyMPDATA.boundary_conditions').Periodic;

nx = int32(24);
ny = int32(24);

Cx = -.5;
Cy = -.25;

[xi, yi] = meshgrid(double(0:1:nx-1), double(0:1:ny-1));

halo = options.n_halo;
advectee = ScalarField(pyargs(...
    'data', py.numpy.array(exp( ...
        -(xi+.5-double(nx)/2).^2 / (2*(double(nx)/10)^2) ...
        -(yi+.5-double(ny)/2).^2 / (2*(double(ny)/10)^2) ...
    )), ...
    'halo', halo, ...
    'boundary_conditions', py.tuple({Periodic(), Periodic()}) ...
));
advector = VectorField(pyargs(...
    'data', py.tuple({ ...
        Cx * py.numpy.ones(int32([nx+1 ny])), ...
        Cy * py.numpy.ones(int32([nx ny+1])) ...
     }), ...
    'halo', halo, ...
    'boundary_conditions', py.tuple({Periodic(), Periodic()}) ...
));

    let vector_field = py.import_bound("PyMPDATA")?.getattr("VectorField")?;
    let scalar_field = py.import_bound("PyMPDATA")?.getattr("ScalarField")?;
    let periodic = py.import_bound("PyMPDATA.boundary_conditions")?.getattr("Periodic")?;

    let nx_ny = [24, 24];
    let cx_cy = [-0.5, -0.25];
    let boundary_con = PyTuple::new_bound(py, [periodic.call0()?, periodic.call0()?]).into_any();
    let halo = options.getattr("n_halo")?;

    let indices = PyDict::new_bound(py);
    Python::run_bound(py, &format!(r#"
import numpy as np
nx, ny = {}, {}
xi, yi = np.indices((nx, ny), dtype=float)
data=np.exp(
  -(xi+.5-nx/2)**2 / (2*(ny/10)**2)
  -(yi+.5-nx/2)**2 / (2*(ny/10)**2)
)
    "#, nx_ny[0], nx_ny[1]), None, Some(&indices)).unwrap();

    let advectee_arg = vec![("data", indices.get_item("data")?), ("halo", Some(halo.clone())), ("boundary_conditions", Some(boundary_con))].into_py_dict_bound(py);
    let advectee = scalar_field.call((), Some(&advectee_arg))?;
    let full = PyDict::new_bound(py);
    Python::run_bound(py, &format!(r#"
import numpy as np
nx, ny = {}, {}
Cx, Cy = {}, {}
data = (np.full((nx + 1, ny), Cx), np.full((nx, ny + 1), Cy))
    "#, nx_ny[0], nx_ny[1], cx_cy[0], cx_cy[1]), None, Some(&full)).unwrap();
    let boundary_con = PyTuple::new_bound(py, [periodic.call0()?, periodic.call0()?]).into_any();
    let advector_arg = vec![("data", full.get_item("data")?), ("halo", Some(halo.clone())), ("boundary_conditions", Some(boundary_con))].into_py_dict_bound(py);
    let advector = vector_field.call((), Some(&advector_arg))?;

from PyMPDATA import ScalarField
from PyMPDATA import VectorField
from PyMPDATA.boundary_conditions import Periodic
import numpy as np

nx, ny = 24, 24
Cx, Cy = -.5, -.25
halo = options.n_halo

xi, yi = np.indices((nx, ny), dtype=float)
advectee = ScalarField(
  data=np.exp(
    -(xi+.5-nx/2)**2 / (2*(nx/10)**2)
    -(yi+.5-ny/2)**2 / (2*(ny/10)**2)
  ),
  halo=halo,
  boundary_conditions=(Periodic(), Periodic())
)
advector = VectorField(
  data=(np.full((nx + 1, ny), Cx), np.full((nx, ny + 1), Cy)),
  halo=halo,
  boundary_conditions=(Periodic(), Periodic())
)

Note that the shapes of arrays representing components of the velocity field are different than the shape of the scalar field array due to employment of the staggered grid.

Besides the exemplified Periodic class representing periodic boundary conditions, PyMPDATA supports Extrapolated, Constant and Polar boundary conditions.

Stepper

The logic of the MPDATA iterative solver is represented in PyMPDATA by the Stepper class.

When instantiating the Stepper, the user has a choice of either supplying just the number of dimensions or specialising the stepper for a given grid:

Stepper = pyimport("PyMPDATA").Stepper

stepper = Stepper(options=options, n_dims=2)

Stepper = py.importlib.import_module('PyMPDATA').Stepper;

stepper = Stepper(pyargs(...
  'options', options, ...
  'n_dims', int32(2) ...
));

let n_dims: i32 = 2;
let stepper_arg = PyDict::new_bound(py);
let _ = PyDictMethods::set_item(&stepper_arg, "options", &options);
let _ = PyDictMethods::set_item(&stepper_arg, "n_dims", &n_dims);

from PyMPDATA import Stepper

stepper = Stepper(options=options, n_dims=2)

stepper = Stepper(options=options, grid=(nx, ny))

stepper = Stepper(pyargs(...
  'options', options, ...
  'grid', py.tuple({nx, ny}) ...
));

 let _stepper_arg_alternative = vec![("options", &options), ("grid", &PyTuple::new_bound(py, nx_ny).into_any())].into_py_dict_bound(py);
 let stepper_ = py.import_bound("PyMPDATA")?.getattr("Stepper")?;
 let stepper = stepper_.call((), Some(&stepper_arg))?; //or use stepper args alternative

stepper = Stepper(options=options, grid=(nx, ny))

Solver = pyimport("PyMPDATA").Solver
solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
solver.advance(n_steps=75)
state = solver.advectee.get()

Solver = py.importlib.import_module('PyMPDATA').Solver;
solver = Solver(pyargs('stepper', stepper, 'advectee', advectee, 'advector', advector));
solver.advance(pyargs('n_steps', 75));
state = solver.advectee.get();

    let solver_ = py.import_bound("PyMPDATA")?.getattr("Solver")?;
    let solver = solver_.call((), Some(&vec![("stepper", stepper), ("advectee", advectee), ("advector", advector)].into_py_dict_bound(py)))?;
    let _state_0 = solver.getattr("advectee")?.getattr("get")?.call0()?.getattr("copy")?.call0()?;
    solver.getattr("advance")?.call((), Some(&vec![("n_steps", 75)].into_py_dict_bound(py)))?;
    let _state = solver.getattr("advectee")?.getattr("get")?.call0()?;
    Ok(())
  })
}

from PyMPDATA import Solver

solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
state_0 = solver.advectee.get().copy()
solver.advance(n_steps=75)
state = solver.advectee.get()

def plot(psi, zlim, norm=None):
    xi, yi = np.indices(psi.shape)
    fig, ax = pyplot.subplots(subplot_kw={"projection": "3d"})
    pyplot.gca().plot_wireframe(
        xi+.5, yi+.5,
        psi, color='red', linewidth=.5
    )
    ax.set_zlim(zlim)
    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
        axis.pane.fill = False
        axis.pane.set_edgecolor('black')
        axis.pane.set_alpha(1)
    ax.grid(False)
    ax.set_zticks([])
    ax.set_xlabel('x/dx')
    ax.set_ylabel('y/dy')
    ax.set_proj_type('ortho')
    cnt = ax.contourf(xi+.5, yi+.5, psi, zdir='z', offset=-1, norm=norm)
    cbar = pyplot.colorbar(cnt, pad=.1, aspect=10, fraction=.04)
    return cbar.norm

zlim = (-1, 1)
norm = plot(state_0, zlim)
pyplot.savefig('readme_gauss_0.png')
plot(state, zlim, norm)
pyplot.savefig('readme_gauss.png')

ENV["NUMBA_DISABLE_JIT"] = "1"

setenv('NUMBA_DISABLE_JIT', '1');

import os
os.environ["NUMBA_DISABLE_JIT"] = "1"