Numerical core for the advection-diffusion solver (climlab.dynamics.adv_diff_numerics)

The 1D advection-diffusion problem

The equation to be solved is

\[\begin{split}\frac{\partial}{\partial t} \psi(x,t) &= -\frac{1}{w(x)} \frac{\partial}{\partial x} \left[ w(x) ~ \mathcal{F}(x,t) \right] + \dot{\psi}\\ \mathcal{F} &= U(x) \psi(x) -K(x) ~ \frac{\partial \psi}{\partial x} + F(x)\end{split}\]

for the following quantities:

• state variable \(\psi(x,t)\)

• diffusivity \(K(x)\) in units of \(x^2 ~ t^{-1}\)

• advecting velocity \(U(x)\) in units of \(x ~ t^{-1}\)

• a prescribed flux \(F(x)\) (including boundary conditions) in units of \(\psi ~ x ~ t^{-1}\)

• a scalar source/sink \(\dot{\psi}(x)\) in units of \(\psi ~ t^{-1}\)

• weighting function \(w(x)\) for the divergence operator on curvilinear grids.

The boundary condition is a flux condition at the end points:

\[\begin{align} \label{eq:fluxcondition} \mathcal{F}(x_0) &= F(x_0) & \mathcal{F}(x_J) &= F(x_J) \end{align}\]

which requires that the advecting velocity \(u(x) = 0\) at the end points \(x_0, x_J\)

The solver is implemented on a 1D staggered grid, with J+1 flux points and J scalar points located somewhere between the flux points.

The solver does not assume the gridpoints are evenly spaced in \(x\).

Routines are provided to compute the following:

• Advective, diffusive, and total fluxes (the terms of \(\mathcal{F}\))

• Tridiagonal matrix operator for the flux convergence

• The actual flux convergence, or instantaneous scalar tendency given a current value of \(\psi(x)\)

• Future value of \(\psi(x)\) for an implicit timestep

Some details of the solver formulas are laid out below for reference.

Spatial discretization

We use a non-uniform staggered spatial grid with scalar \(\psi\) evaluated at \(J\) points, and flux \(\mathcal{F}\) evaluated at \(J+1\) flux points. The indexing will run from \(j=0\) to \(j=J\) for the flux points, and \(i=0\) to \(i=J-1\) for the scalar points. This notation is consistent with zero-indexed Python arrays.

We define the following arrays:

• \(\mathcal{X}_b[j]\) is a length J+1 array defining the location of the flux points.

• \(\mathcal{X}[i]\) is a length J array defining the location of the scalar points, where point \(\mathcal{X}[j]\) is somewhere between \(\mathcal{X}_b[j]\) and \(\mathcal{X}_b[j+1]\) for all \(j<J\).

• \(\psi[i], \dot{\psi}[i]\) are length J arrays defined on \(\mathcal{X}\).

• \(U[j], K[j], F[j]\) are all arrays of length J+1 defined on \(\mathcal{X}_b\).

• The grid weights are similarly in arrays \(W_b[j], W[i]\) respectively on \(\mathcal{X}_b`\) and \(\mathcal{X}\).

Centered difference formulas for the flux

We use centered differences in \(x\) to discretize the spatial derivatives. The diffusive component of the flux is thus

\[\begin{align*} \mathcal{F}_{diff}[j] &= - K[j] \frac{ \left( \psi[i] - \psi[i-1] \right) }{\left( \mathcal{X}[i] - \mathcal{X}[i-1] \right)} & j&=i=1,2,...,J-1 \end{align*}\]

The diffusive flux is assumed to be zero at the boundaries.

The advective term requires an additional approximation since the scalar \(\psi\) is not defined at the flux points. We use a linear interpolation to the flux points:

\[\begin{align*} \psi_b[j] &\equiv \psi[i-1] \left( \frac{\mathcal{X}[i] - \mathcal{X}_b[j]}{\mathcal{X}[i] - \mathcal{X}[i-1]} \right) + \psi[i] \left( \frac{ \mathcal{X}_b[j] - \mathcal{X}[i-1] }{\mathcal{X}[i] - \mathcal{X}[i-1]} \right) & j&=i=1,2,...,J-1 \end{align*}\]

Note that for an evenly spaced grid, this reduces to the simple average \(\frac{1}{2} \left( \psi[i-1] + \psi[i] \right)\).

With this interpolation, the advective flux is approximated by

\[\begin{align*} \mathcal{F}_{adv}[j] &= \frac{U[j] }{\mathcal{X}[i] - \mathcal{X}[i-1]} \left( \psi[i-1] (\mathcal{X}[i] - \mathcal{X}_b[j]) + \psi[i] (\mathcal{X}_b[j] - \mathcal{X}[i-1]) \right) & j&=i=1,2,...,J-1 \end{align*}\]

The total flux away from the boundaries (after some recombining terms) is thus:

\[\mathcal{F}[j] = F[j] + \psi[i-1] \left( \frac{K[j] + U[j] (\mathcal{X}[i] - \mathcal{X}_b[j]) }{ \mathcal{X}[i] - \mathcal{X}[i-1] } \right) - \psi[i] \left( \frac{K[j] - U[j] (\mathcal{X}_b[j] - \mathcal{X}[i-1]) }{\mathcal{X}[i] - \mathcal{X}[i-1] } \right)\]

which is valid for j=i=1,2,…,J-1.

Centered difference formulas for the flux convergence

Centered difference approximation of the flux convergence gives

\[\begin{align*} \frac{\partial }{\partial t} \psi[i] &= -\frac{ W_b[j+1] \mathcal{F}[j+1] - W_b[j] \mathcal{F}[j] }{W[i] ( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] )} + \dot{\psi}[i] & i&=j=0,1,...,J-1 \end{align*}\]

The flux convergences are best expressed together in matrix form:

\[\begin{equation} \frac{\partial \boldsymbol{\psi}}{\partial t} = \boldsymbol{T} ~ \boldsymbol{\psi} + \boldsymbol{S} \end{equation}\]

where \(\boldsymbol{\psi}\) is the \(J\times1\) column vector, \(\boldsymbol{S}\) is a \(J\times1\) column vector representing the prescribed flux convergence and source terms, whose elements are

\[\begin{align} S[i] &= \frac{-W_b[j+1] F[j+1] + W_b[j] F[j]}{W[i] ( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] )} + \dot{\psi}[i] & i&=j=0,1,...,J-1 \end{align}\]

and \(\boldsymbol{T}\) is a \(J\times J\) tridiagonal matrix:

\[\begin{split}\begin{equation} \boldsymbol{T} ~ \boldsymbol{\psi} = \left[\begin{array}{ccccccc} T_{m0} & T_{u1} & 0 & ... & 0 & 0 & 0 \\T_{l0} & T_{m1} & T_{u2} & ... & 0 & 0 & 0 \\ 0 & T_{l1} & T_{m2} & ... & 0 & 0 & 0 \\... & ... & ... & ... & ... & ... & ... \\0 & 0 & 0 & ... & T_{m(J-3)} & T_{u(J-2)} & 0 \\0 & 0 & 0 & ... & T_{l(J-3)} & T_{m(J-2)} & T_{u(J-1)} \\0 & 0 & 0 & ... & 0 & T_{l(J-2)} & T_{m(J-1)}\end{array}\right] \left[\begin{array}{c} \psi_0 \\ \psi_1 \\ \psi_2 \\... \\ \psi_{J-3} \\ \psi_{J-2} \\ \psi_{J-1} \end{array}\right] \end{equation}\end{split}\]

with vectors \(T_l, T_m, T_u\) representing respectively the lower, main, and upper diagonals of \(\boldsymbol{T}\). We will treat all three vectors as length J; the 0th element of \(T_u\) is ignored while the (J-1)th element of \(T_l\) is ignored (this is consistent with the expected inputs for the Python module scipy.linalg.solve_banded).

The instantanous tendency is then easily computed by matrix multiplication.

The elements of the main diagonal of \(\boldsymbol{\psi}\) can be computed from

\[\begin{split}\begin{align} \label{eq:maindiag} \begin{split} T_m[i] &= -\left( \frac{ W_b[j+1] \big( K[j+1] + U[j+1] (\mathcal{X}[i+1] - \mathcal{X}_b[j+1]) \big) }{ W[i] ( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] )(\mathcal{X}[i+1] - \mathcal{X}[i]) } \right) \\ & \qquad - \left( \frac{W_b[j] \big( K[j] - U[j] (\mathcal{X}_b[j] - \mathcal{X}[i-1]) \big) }{W[i] ( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] )(\mathcal{X}[i] - \mathcal{X}[i-1]) } \right) \\ i &=j=0,2,...,J-1 \end{split} \end{align}\end{split}\]

which is valid at the boundaries so long as we set \(W_b[0] = W_b[J] = 0\).

The lower diagonal (including the right boundary condition) is computed from

\[\begin{split}\begin{align} \label{eq:lowerdiag} \begin{split} T_l[i-1] &= \left( \frac{W_b[j]}{W[i] } \right) \left( \frac{ K[j] + U[j] (\mathcal{X}[i] - \mathcal{X}_b[j]) }{( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] ) (\mathcal{X}[i] - \mathcal{X}[i-1] )} \right) \\ i &=j =1,2,...,J-2, J-1 \end{split} \end{align}\end{split}\]

Finally the upper diagonal (including the left boundary condition) is computed from

\[\begin{split}\begin{align} \label{eq:upperdiag} \begin{split} T_u[i+1] &= \left( \frac{W_b[j+1]}{W[i]} \right) \left( \frac{K[j+1] - U[j+1] (\mathcal{X}_b[j+1] - \mathcal{X}[i]) }{( \mathcal{X}_b[j+1] - \mathcal{X}_b[j] )(\mathcal{X}[i+1] - \mathcal{X}[i] ) } \right) \\ i &= j=0,...,J-2 \end{split} \end{align}\end{split}\]

Implicit time discretization

The forward-time finite difference approximation to LHS of the flux-convergence equation is simply

\[\begin{equation} \frac{\partial \psi[i]}{\partial t} \approx \frac{\psi^{n+1}[i]- \psi^{n}[i]}{\Delta t} \end{equation}\]

where the superscript \(n\) indicates the time index.

We use the implicit-time method, in which the RHS is evaluated at the future time \(n+1\). Applying this to the matrix equation above and moving all the terms at time \(n+1\) over to the LHS yields

\[\begin{equation} \label{eq:implicit_tridiagonal} \left( \boldsymbol{I} - \boldsymbol{T} \Delta t \right) \boldsymbol{\psi}^{n+1} = \boldsymbol{\psi}^{n} + \boldsymbol{S} \Delta t \end{equation}\]

where \(\boldsymbol{I}\) is the \(J\times J\) identity matrix.

Solving for the future value \(\boldsymbol{\psi}^{n+1}\) is then accomplished by solving the \(J \times J\) tridiagonal linear system using standard routines.

Analytical benchmark

Here is an analytical case to be used for testing purposes to validate the numerical code. This is implemented in the CLIMLAB test suite.

• \(K=K_0\) is constant

• \(w(x) = 1\) everywhere (Cartesian coordinates)

• \(F = 0\) everywhere

• \(\psi(x,0) = \psi_0 \sin^2\left(\frac{\pi x}{L}\right)\)

• \(u(x) = U_0 \sin\left(\frac{\pi x}{L}\right)\)

for a domain with endpoints at \(x=0\) and \(x=L\).

The analytical solution is

\[\begin{split}\begin{align} \mathcal{F} &= \psi_0 \sin\left(\frac{\pi x}{L}\right) \left[U_0 \sin^2\left(\frac{\pi x}{L}\right) - 2K \frac{\pi}{L} \cos\left(\frac{\pi x}{L}\right) \right] \\ \frac{\partial \psi}{\partial t} &= -\psi_0 \frac{\pi}{L} \left\{ 3 U_0 \sin^2\left(\frac{\pi x}{L}\right) \cos\left(\frac{\pi x}{L}\right) -2K\frac{\pi}{L} \left[\cos^2\left(\frac{\pi x}{L}\right) -\sin^2\left(\frac{\pi x}{L}\right) \right] \right\} \end{align}\end{split}\]

which satisfies the boundary condition \(\mathcal{F} = 0\) at \(x=0\) and \(x=L\).

Module function reference

All the functions in climlab.dynamics.adv_diff_numerics are vectorized to handle multidimensional input. The key assumption is that advection-diffusion operates along the final dimension.

Inputs should be reshaped appropriately (e.g. with numpy.moveaxis()) before calling these functions.

climlab.dynamics.adv_diff_numerics.advdiff_tridiag(X, Xb, K, U, W=None, Wb=None, use_banded_solver=False)

Compute the tridiagonal matrix operator for the advective-diffusive flux convergence.

Input arrays of length J+1:

Xb, Wb, K, U

Input arrays of length J:

X, W

The 0th and Jth (i.e. first and last) elements of Wb are ignored; assuming boundary condition is a prescribed flux.

The return value depends on input flag use_banded_solver

If use_banded_solver==True, return a 3xJ array containing the elements of the tridiagonal. This version is restricted to 1D input arrays, but is suitable for use with the efficient banded solver.

If use_banded_solver=False (which it must be for multidimensional input), return an array (…,J,J) with the full tridiagonal matrix.

climlab.dynamics.adv_diff_numerics.advective_flux(X, Xb, U, field)

Return the advective flux on cell boundaries (length J+1)

climlab.dynamics.adv_diff_numerics.compute_source(X, Xb, prescribed_flux=None, prescribed_source=None, W=None, Wb=None)

Return the source array S consisting of the convergence of the prescribed flux plus the prescribed scalar source.

climlab.dynamics.adv_diff_numerics.compute_tendency(field, tridiag, source, use_banded_solver=False)

Return the instantaneous scalar tendency.

This is the sum of the convergence of advective+diffusive flux plus any prescribed convergence or scalar sources.

The convergence is computed by matrix multiplication:

\[\frac{\partial \psi}{\partial t} = T \times \psi + S\]

where \(T\) is the tridiagonal flux convergence matrix.

climlab.dynamics.adv_diff_numerics.diffusive_flux(X, Xb, K, field)

Return the diffusive flux on cell boundaries (length J+1)

climlab.dynamics.adv_diff_numerics.implicit_step_forward(initial_field, tridiag, source, timestep, use_banded_solver=False)

Return the field at future time using an implicit timestep.

The matrix problem is

\[(I - T \Delta t) \psi^{n+1} = \psi^n + S \Delta t\]

where \(T\) is the tridiagonal matrix for the flux convergence, \(psi\) is the state variable, the superscript \(n\) refers to the time index, and \(S \Delta t\) is the accumulated source over the timestep \(\Delta t\).

Input arguments:

• initial_field: the current state variable \(\psi^n\), dimensions (…,J)

• tridiag: the tridiagonal matrix \(T\), dimensions (…,J,J) or (…,3,J) depending on the value of use_banded_solver

• source: prescribed sources/sinks of \(\psi\), dimensions (…,J)

• timestep: the discrete timestep in time units

• use_banded_solver: switch to use the optional efficient banded solver (see below)

Returns the updated value of the state variable \(\psi^{n+1}\), dimensions (…,J)

The expected shape of tridiag depends on the switch use_banded_solver, which should be consistent with that used in the call to advdiff_tridiag(). If True, we use the efficient banded matrix solver scipy.linalg.solve_banded(). However this will probably only work for a 1D state variable.

The default is to use the general linear system solver numpy.linalg.solve().

climlab.dynamics.adv_diff_numerics.make_the_actual_tridiagonal_matrix(tridiag_banded)

Convert a (3xJ) banded array into full (JxJ) tridiagonal matrix form.

climlab.dynamics.adv_diff_numerics.total_flux(X, Xb, K, U, field, prescribed_flux=None)

Return the total (advective + diffusive + prescribed) flux on cell boundaries (length J+1)