MatrixFields
ClimaCore.MatrixFields
— ModuleMatrixFields
This module adds support for defining and manipulating Field
s that represent matrices. Specifically, it adds the BandMatrixRow
type, which can be used to store the entries of a band matrix. A Field
of BandMatrixRow
s on a FiniteDifferenceSpace
can be interpreted as a band matrix by vertically concatenating the BandMatrixRow
s. Similarly, a Field
of BandMatrixRow
s on an ExtrudedFiniteDifferenceSpace
can be interpreted as a collection of band matrices, one for each column of the Field
. Such Field
s are called ColumnwiseBandMatrixField
s, and this module adds the following functionality for them:
- Constructors, e.g.,
matrix_field = @. BidiagonalMatrixRow(field1, field2)
- Linear combinations, e.g.,
@. 3 * matrix_field1 + matrix_field2 / 3
- Matrix-vector multiplication, e.g.,
@. matrix_field ⋅ field
- Matrix-matrix multiplication, e.g.,
@. matrix_field1 ⋅ matrix_field2
- Compatibility with
LinearAlgebra.I
, e.g.,@. matrix_field = (4I,)
or@. matrix_field - (4I,)
- Integration with
RecursiveApply
, e.g., the entries ofmatrix_field
can beTuple
s orNamedTuple
s instead of single values, which allowsmatrix_field
to represent multiple band matrices at the same time - Integration with
Operators
, e.g., thematrix_field
that gets applied to the argument of anyFiniteDifferenceOperator
op
can be obtained using theFiniteDifferenceOperator
operator_matrix(op)
- Conversions to native array types, e.g.,
field2arrays(matrix_field)
can convert each column ofmatrix_field
into aBandedMatrix
fromBandedMatrices.jl
- Custom printing, e.g.,
matrix_field
gets displayed as aBandedMatrix
, specifically, as theBandedMatrix
that corresponds to its first column
This module also adds support for defining and manipulating sparse block matrices of Field
s. Specifically, it adds the FieldMatrix
type, which is a dictionary that maps pairs of FieldName
s to ColumnwiseBandMatrixField
s or multiples of LinearAlgebra.I
. This comes with the following functionality:
- Addition and subtraction, e.g.,
@. field_matrix1 + field_matrix2
- Matrix-vector multiplication, e.g.,
@. field_matrix * field_vector
- Matrix-matrix multiplication, e.g.,
@. field_matrix1 * field_matrix2
- Integration with
RecursiveApply
, e.g., the entries offield_matrix
can be specified either as matrixField
s ofTuple
s orNamedTuple
s, or as separate matrixField
s of single values - The ability to solve linear equations using
FieldMatrixSolver
, which is a generalization ofldiv!
that is designed to optimize solver performance
Matrix Field Element Type
ClimaCore.MatrixFields.BandMatrixRow
— TypeBandMatrixRow{ld}(entries...)
Stores the nonzero entries in a row of a band matrix, starting with the lowest diagonal, which has index ld
. Supported operations include accessing the entry on the diagonal with index d
by calling row[d]
, taking linear combinations with other band matrix rows (and with LinearAlgebra.I
), and checking for equality with other band matrix rows (and with LinearAlgebra.I
). There are several aliases for commonly used subtypes of BandMatrixRow
:
DiagonalMatrixRow(entry_1)
BidiagonalMatrixRow(entry_1, entry_2)
TridiagonalMatrixRow(entry_1, entry_2, entry_3)
QuaddiagonalMatrixRow(entry_1, entry_2, entry_3, entry_4)
PentadiagonalMatrixRow(entry_1, entry_2, entry_3, entry_4, entry_5)
Matrix Field Multiplication
ClimaCore.MatrixFields.MultiplyColumnwiseBandMatrixField
— TypeMultiplyColumnwiseBandMatrixField()
An operator that multiplies a ColumnwiseBandMatrixField
by another Field
, i.e., matrix-vector or matrix-matrix multiplication. The ⋅
symbol is an alias for MultiplyColumnwiseBandMatrixField()
.
What follows is a derivation of the algorithm used by this operator with single-column Field
s. For Field
s on multiple columns, the same computation is done for each column.
In this derivation, we will use $M_1$ and $M_2$ to denote two ColumnwiseBandMatrixField
s, and we will use $V$ to denote a regular (vector-like) Field
. For both $M_1$ and $M_2$, we will use the array-like index notation $M[row, col]$ to denote $M[row][col-row]$, i.e., the entry in the BandMatrixRow
$M[row]$ located on the diagonal with index $col - row$. We will also use outer_indices
$($space
$)$ to denote the tuple $($left_idx
$($space
$),$right_idx
$($space
$))$.
1. Matrix-Vector Multiplication
From the definition of matrix-vector multiplication,
\[(M_1 ⋅ V)[i] = \sum_k M_1[i, k] * V[k].\]
To establish bounds on the values of $k$, let us define the following values:
- $li_1, ri_1 ={}$
outer_indices
$($column_axes
$(M_1))$ - $ld_1, ud_1 ={}$
outer_diagonals
$($eltype
$(M_1))$
Since $M_1[i, k]$ is only well-defined if $k$ is a valid column index and $k - i$ is a valid diagonal index, we know that
\[li_1 \leq k \leq ri_1 \quad \text{and} \quad ld_1 \leq k - i \leq ud_1.\]
Combining these into a single inequality gives us
\[\text{max}(li_1, i + ld_1) \leq k \leq \text{min}(ri_1, i + ud_1).\]
So, we can rewrite the expression for $(M_1 ⋅ V)[i]$ as
\[(M_1 ⋅ V)[i] = \sum_{k\ =\ \text{max}(li_1, i + ld_1)}^{\text{min}(ri_1, i + ud_1)} M_1[i, k] * V[k].\]
If we replace the variable $k$ with $d = k - i$ and switch from array-like indexing to Field
indexing, we find that
\[(M_1 ⋅ V)[i] = \sum_{d\ =\ \text{max}(li_1 - i, ld_1)}^{\text{min}(ri_1 - i, ud_1)} M_1[i][d] * V[i + d].\]
1.1 Interior vs. Boundary Indices
Now, suppose that the row index $i$ is such that
\[li_1 - ld_1 \leq i \leq ri_1 - ud_1.\]
If this is the case, then the bounds on $d$ can be simplified to
\[\text{max}(li_1 - i, ld_1) = ld_1 \quad \text{and} \quad \text{min}(ri_1 - i, ud_1) = ud_1.\]
The expression for $(M_1 ⋅ V)[i]$ then becomes
\[(M_1 ⋅ V)[i] = \sum_{d = ld_1}^{ud_1} M_1[i][d] * V[i + d].\]
The values of $i$ in this range are considered to be in the "interior" of the operator, while those not in this range (for which we cannot make the above simplification) are considered to be on the "boundary".
2. Matrix-Matrix Multiplication
From the definition of matrix-matrix multiplication,
\[(M_1 ⋅ M_2)[i, j] = \sum_k M_1[i, k] * M_2[k, j].\]
To establish bounds on the values of $j$ and $k$, let us define the following values:
- $li_1, ri_1 ={}$
outer_indices
$($column_axes
$(M_1))$ - $ld_1, ud_1 ={}$
outer_diagonals
$($eltype
$(M_1))$ - $li_2, ri_2 ={}$
outer_indices
$($column_axes
$(M_2))$ - $ld_2, ud_2 ={}$
outer_diagonals
$($eltype
$(M_2))$
In addition, let $ld_{prod}$ and $ud_{prod}$ denote the outer diagonal indices of the product matrix $M_1 ⋅ M_2$. We will derive the values of $ld_{prod}$ and $ud_{prod}$ in the last section.
Since $M_1[i, k]$ is only well-defined if $k$ is a valid column index and $k - i$ is a valid diagonal index, we know that
\[li_1 \leq k \leq ri_1 \quad \text{and} \quad ld_1 \leq k - i \leq ud_1.\]
Since $M_2[k, j]$ is only well-defined if $j$ is a valid column index and $j - k$ is a valid diagonal index, we also know that
\[li_2 \leq j \leq ri_2 \quad \text{and} \quad ld_2 \leq j - k \leq ud_2.\]
Finally, $(M_1 ⋅ M_2)[i, j]$ is only well-defined if $j - i$ is a valid diagonal index, so
\[ld_{prod} \leq j - i \leq ud_{prod}.\]
These inequalities can be combined to obtain
\[\begin{gather*} \text{max}(li_2, i + ld_{prod}) \leq j \leq \text{min}(ri_2, i + ud_{prod}) \\ \text{and} \\ \text{max}(li_1, i + ld_1, j - ud_2) \leq k \leq \text{min}(ri_1, i + ud_1, j - ld_2). \end{gather*}\]
So, we can rewrite the expression for $(M_1 ⋅ M_2)[i, j]$ as
\[\begin{gather*} (M_1 ⋅ M_2)[i, j] = \sum_{ k\ =\ \text{max}(li_1, i + ld_1, j - ud_2) }^{\text{min}(ri_1, i + ud_1, j - ld_2)} M_1[i, k] * M_2[k, j], \text{ where} \\[0.5em] \text{max}(li_2, i + ld_{prod}) \leq j \leq \text{min}(ri_2, i + ud_{prod}). \end{gather*}\]
If we replace the variable $k$ with $d = k - i$, replace the variable $j$ with $d_{prod} = j - i$, and switch from array-like indexing to Field
indexing, we find that
\[\begin{gather*} (M_1 ⋅ M_2)[i][d_{prod}] = \sum_{ d\ =\ \text{max}(li_1 - i, ld_1, d_{prod} - ud_2) }^{\text{min}(ri_1 - i, ud_1, d_{prod} - ld_2)} M_1[i][d] * M_2[i + d][d_{prod} - d], \text{ where} \\[0.5em] \text{max}(li_2 - i, ld_{prod}) \leq d_{prod} \leq \text{min}(ri_2 - i, ud_{prod}). \end{gather*}\]
2.1 Interior vs. Boundary Indices
Now, suppose that the row index $i$ is such that
\[\text{max}(li_1 - ld_1, li_2 - ld_{prod}) \leq i \leq \text{min}(ri_1 - ud_1, ri_2 - ud_{prod}).\]
If this is the case, then the bounds on $d_{prod}$ can be simplified to
\[\text{max}(li_2 - i, ld_{prod}) = ld_{prod} \quad \text{and} \quad \text{min}(ri_2 - i, ud_{prod}) = ud_{prod}.\]
Similarly, the bounds on $d$ can be simplified using the fact that
\[\text{max}(li_1 - i, ld_1) = ld_1 \quad \text{and} \quad \text{min}(ri_1 - i, ud_1) = ud_1.\]
The expression for $(M_1 ⋅ M_2)[i][d_{prod}]$ then becomes
\[\begin{gather*} (M_1 ⋅ M_2)[i][d_{prod}] = \sum_{ d\ =\ \text{max}(ld_1, d_{prod} - ud_2) }^{\text{min}(ud_1, d_{prod} - ld_2)} M_1[i][d] * M_2[i + d][d_{prod} - d], \text{ where} \\[0.5em] ld_{prod} \leq d_{prod} \leq ud_{prod}. \end{gather*}\]
The values of $i$ in this range are considered to be in the "interior" of the operator, while those not in this range (for which we cannot make these simplifications) are considered to be on the "boundary".
2.2 $ld_{prod}$ and $ud_{prod}$
We only need to compute $(M_1 ⋅ M_2)[i][d_{prod}]$ for values of $d_{prod}$ that correspond to a nonempty sum in the interior, i.e, those for which
\[\text{max}(ld_1, d_{prod} - ud_2) \leq \text{min}(ud_1, d_{prod} - ld_2).\]
This can be broken down into the four inequalities
\[ld_1 \leq ud_1, \qquad ld_1 \leq d_{prod} - ld_2, \qquad d_{prod} - ud_2 \leq ud_1, \quad \text{and} \quad d_{prod} - ud_2 \leq d_{prod} - ld_2.\]
By definition, $ld_1 \leq ud_1$ and $ld_2 \leq ud_2$, so the first and last inequality are always true. Rearranging the remaining two inequalities tells us that
\[ld_1 + ld_2 \leq d_{prod} \leq ud_1 + ud_2.\]
In other words, the outer diagonal indices of $M_1 ⋅ M_2$ are
\[ld_{prod} = ld_1 + ld_2 \quad \text{and} \quad ud_{prod} = ud_1 + ud_2.\]
This means that we can express the bounds on the interior values of $i$ as
\[\text{max}(li_1, li_2 - ld_2) - ld_1 \leq i \leq \text{min}(ri_1, ri_2 - ud_2) - ud_1.\]
Operator Matrices
ClimaCore.MatrixFields.operator_matrix
— Functionoperator_matrix(op)
Constructs a new operator (or operator-like object) that generates the matrix applied by op
to its final argument. If op_matrix = operator_matrix(op)
, we can use the following identities:
- When
op
takes one argument,@. op(arg) == @. op_matrix() ⋅ arg
. - When
op
takes multiple arguments,@. op(args..., arg) == @. op_matrix(args...) ⋅ arg
.
When op
takes more than one argument, operator_matrix(op)
constructs a FiniteDifferenceOperator
that generates the operator matrix. When op
only takes one argument, it instead constructs an AbstractLazyOperator
, which is internally converted into a FiniteDifferenceOperator
when used in a broadcast expression. Implementing op_matrix
as a lazy operator allows us to add an argument to the expression op_matrix.()
, and we then use this argument to infer the space and element type of the operator matrix.
As an example, the InterpolateF2C()
operator on a space with $n$ cell centers applies an $n \times (n + 1)$ bidiagonal matrix:
\[\textrm{interp}(arg) = \begin{bmatrix} 0.5 & 0.5 & 0 & \cdots & 0 & 0 & 0 \\ 0 & 0.5 & 0.5 & \cdots & 0 & 0 & 0 \\ 0 & 0 & 0.5 & \cdots & 0 & 0 & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & 0.5 & 0.5 & 0 \\ 0 & 0 & 0 & \cdots & 0 & 0.5 & 0.5 \end{bmatrix} ⋅ arg\]
The GradientF2C()
operator applies a similar matrix, but with different entries:
\[\textrm{grad}(arg) = \begin{bmatrix} -\textbf{e}^3 & \textbf{e}^3 & 0 & \cdots & 0 & 0 & 0 \\ 0 & -\textbf{e}^3 & \textbf{e}^3 & \cdots & 0 & 0 & 0 \\ 0 & 0 & -\textbf{e}^3 & \cdots & 0 & 0 & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & -\textbf{e}^3 & \textbf{e}^3 & 0 \\ 0 & 0 & 0 & \cdots & 0 & -\textbf{e}^3 & \textbf{e}^3 \end{bmatrix} ⋅ arg\]
The unit vector $\textbf{e}^3$, which can also be thought of as the differential along the third coordinate axis ($\textrm{d}\xi^3$), is implemented as a Geometry.Covariant3Vector(1)
.
Not all operators have well-defined operator matrices. For example, the operator GradientC2F(; bottom = SetGradient(grad_b), top = SetGradient(grad_t))
applies an affine transformation:
\[\textrm{grad}(arg) = \begin{bmatrix} grad_b \\ 0 \\ 0 \\ \vdots \\ 0 \\ 0 \\ grad_t \end{bmatrix} + \begin{bmatrix} 0 & 0 & 0 & \cdots & 0 & 0 \\ -\textbf{e}^3 & \textbf{e}^3 & 0 & \cdots & 0 & 0 \\ 0 & -\textbf{e}^3 & \textbf{e}^3 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & \textbf{e}^3 & 0 \\ 0 & 0 & 0 & \cdots & -\textbf{e}^3 & \textbf{e}^3 \\ 0 & 0 & 0 & \cdots & 0 & 0 \end{bmatrix} ⋅ arg\]
However, this simplifies to a linear transformation when $grad_b$ and $grad_t$ are both 0:
\[\textrm{grad}(arg) = \begin{bmatrix} 0 & 0 & 0 & \cdots & 0 & 0 \\ -\textbf{e}^3 & \textbf{e}^3 & 0 & \cdots & 0 & 0 \\ 0 & -\textbf{e}^3 & \textbf{e}^3 & \cdots & 0 & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & \textbf{e}^3 & 0 \\ 0 & 0 & 0 & \cdots & -\textbf{e}^3 & \textbf{e}^3 \\ 0 & 0 & 0 & \cdots & 0 & 0 \end{bmatrix} ⋅ arg\]
In general, when op
has nonzero boundary conditions that make it apply an affine transformation, operator_matrix(op)
will print out a warning and zero out the boundary conditions before computing the operator matrix.
In addition to affine transformations, there are also some operators that apply nonlinear transformations to their arguments; that is, transformations which cannot be accurately approximated without using more terms of the form
\[\textrm{op}(\textbf{0}) + \textrm{op}'(\textbf{0}) ⋅ arg + \textrm{op}''(\textbf{0}) ⋅ arg ⋅ arg + \ldots.\]
When op
is such an operator, operator_matrix(op)
will throw an error. In the future, we may want to modify operator_matrix(op)
so that it will instead return $\textrm{op}'(\textbf{0})$, where $\textbf{0} ={}$zero.(arg)
.
Linear Solvers
ClimaCore.MatrixFields.FieldMatrixSolverAlgorithm
— TypeFieldMatrixSolverAlgorithm
Description of how to solve an equation of the form A * x = b
for x
, where A
is a FieldMatrix
and where x
and b
are both FieldVector
s. Different algorithms can be nested inside each other, enabling the construction of specialized linear solvers that fully utilize the sparsity pattern of A
.
Interface
Every subtype of FieldMatrixSolverAlgorithm
must implement methods for the following functions:
ClimaCore.MatrixFields.FieldMatrixSolver
— TypeFieldMatrixSolver(alg, A, b)
Combination of a FieldMatrixSolverAlgorithm
alg
and the cache it requires to solve the equation A * x = b
for x
. The values of A
and b
that get passed to this constructor should be similar
to the ones that get passed to field_matrix_solve!
in order to ensure that the cache gets allocated correctly.
ClimaCore.MatrixFields.FieldMatrixWithSolver
— TypeFieldMatrixWithSolver(A, b, [alg])
A wrapper that combines a FieldMatrix
A
with a FieldMatrixSolver
that can be used to solve the equation A * x = b
for x
, where x
and b
are both FieldVector
s. Similar to a LinearAlgebra.Factorization
, this wrapper can be passed to ldiv!
, whereas a regular FieldMatrix
cannot be passed to ldiv!
.
By default, the FieldMatrixSolverAlgorithm
alg
is set to a BlockDiagonalSolve
, so a custom alg
must be specified when A
is not a block diagonal matrix.
ClimaCore.MatrixFields.field_matrix_solve!
— Functionfield_matrix_solve!(solver, x, A, b)
Solves the equation A * x = b
for x
using the FieldMatrixSolver
solver
.
ClimaCore.MatrixFields.BlockDiagonalSolve
— TypeBlockDiagonalSolve()
A FieldMatrixSolverAlgorithm
for a block diagonal matrix:
\[A = \begin{bmatrix} A_{11} & \mathbf{0} & \mathbf{0} & \cdots & \mathbf{0} \\ \mathbf{0} & A_{22} & \mathbf{0} & \cdots & \mathbf{0} \\ \mathbf{0} & \mathbf{0} & A_{33} & \cdots & \mathbf{0} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ \mathbf{0} & \mathbf{0} & \mathbf{0} & \cdots & A_{NN} \end{bmatrix}\]
This algorithm solves the N
block equations Aₙₙ * xₙ = bₙ
in sequence (though we might want to parallelize it in the future).
If Aₙₙ
is a diagonal matrix, the equation Aₙₙ * xₙ = bₙ
is solved by making a single pass over the data, setting each xₙ[i]
to inv(Aₙₙ[i, i]) * bₙ[i]
.
Otherwise, the equation Aₙₙ * xₙ = bₙ
is solved using Gaussian elimination (without pivoting), which makes two passes over the data. This is currently only implemented for tri-diagonal and penta-diagonal matrices Aₙₙ
. In Gaussian elimination, Aₙₙ
is effectively factorized into the product Lₙ * Dₙ * Uₙ
, where Dₙ
is a diagonal matrix, and where Lₙ
and Uₙ
are unit lower and upper triangular matrices, respectively. The first pass multiplies both sides of the equation by inv(Lₙ * Dₙ)
, replacing Aₙₙ
with Uₙ
and bₙ
with Uₙxₙ
, which is referred to as putting Aₙₙ
into "reduced row echelon form". The second pass solves Uₙ * xₙ = Uₙxₙ
for xₙ
with a unit upper triangular matrix solver, which is referred to as "back substitution". These operations can become numerically unstable when Aₙₙ
has entries with large disparities in magnitude, but avoiding this would require swapping the rows of Aₙₙ
(i.e., replacing Dₙ
with a partial pivoting matrix).
ClimaCore.MatrixFields.BlockLowerTriangularSolve
— TypeBlockLowerTriangularSolve(names₁...; [alg₁], [alg₂])
A FieldMatrixSolverAlgorithm
for a 2×2 block lower triangular matrix:
\[A = \begin{bmatrix} A_{11} & \mathbf{0} \\ A_{21} & A_{22} \end{bmatrix}\]
The FieldName
s in names₁
correspond to the subscript ₁
, while all other FieldName
s correspond to the subscript ₂
. This algorithm has 2 steps:
- Solve
A₁₁ * x₁ = b₁
forx₁
using the algorithmalg₁
, which is set to aBlockDiagonalSolve
by default. - Solve
A₂₂ * x₂ = b₂ - A₂₁ * x₁
forx₂
using the algorithmalg₂
, which is also set to aBlockDiagonalSolve
by default.
ClimaCore.MatrixFields.BlockArrowheadSolve
— TypeBlockArrowheadSolve(names₁...; [alg₂])
A FieldMatrixSolverAlgorithm
for a 2×2 block arrowhead matrix:
\[A = \begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix}, \quad \text{where } A_{11} \text{ is a diagonal matrix}\]
The FieldName
s in names₁
correspond to the subscript ₁
, while all other FieldName
s correspond to the subscript ₂
. This algorithm has only 1 step:
- Solve
(A₂₂ - A₂₁ * inv(A₁₁) * A₁₂) * x₂ = b₂ - A₂₁ * inv(A₁₁) * b₁
forx₂
using the algorithmalg₂
, which is set to aBlockDiagonalSolve
by default, and setx₁
toinv(A₁₁) * (b₁ - A₁₂ * x₂)
.
Since A₁₁
is a diagonal matrix, inv(A₁₁)
is easy to compute, which means that the Schur complement of A₁₁
in A
, A₂₂ - A₂₁ * inv(A₁₁) * A₁₂
, as well as the vectors b₂ - A₂₁ * inv(A₁₁) * b₁
and inv(A₁₁) * (b₁ - A₁₂ * x₂)
, are also easy to compute.
This algorithm is equivalent to block Gaussian elimination with all operations inlined into a single step.
ClimaCore.MatrixFields.SchurComplementReductionSolve
— TypeSchurComplementReductionSolve(names₁...; [alg₁], alg₂)
A FieldMatrixSolverAlgorithm
for any 2×2 block matrix:
\[A = \begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix}\]
The FieldName
s in names₁
correspond to the subscript ₁
, while all other FieldName
s correspond to the subscript ₂
. This algorithm has 3 steps:
- Solve
A₁₁ * x₁′ = b₁
forx₁′
using the algorithmalg₁
, which is set to aBlockDiagonalSolve
by default. - Solve
(A₂₂ - A₂₁ * inv(A₁₁) * A₁₂) * x₂ = b₂ - A₂₁ * x₁′
forx₂
using the algorithmalg₂
. - Solve
A₁₁ * x₁ = b₁ - A₁₂ * x₂
forx₁
using the algorithmalg₁
.
Since A₁₁
is not necessarily a diagonal matrix, inv(A₁₁)
will generally be a dense matrix, which means that the Schur complement of A₁₁
in A
, A₂₂ - A₂₁ * inv(A₁₁) * A₁₂
, cannot be computed efficiently. So, alg₂
must be set to a LazyFieldMatrixSolverAlgorithm
, which can evaluate the matrix-vector product (A₂₂ - A₂₁ * inv(A₁₁) * A₁₂) * x₂
without actually computing the Schur complement matrix. This involves representing the Schur complement matrix by a LazySchurComplement
, which uses alg₁
to invert A₁₁
when computing the matrix-vector product.
This algorithm is equivalent to block Gaussian elimination, where steps 1 and 2 put A
into reduced row echelon form, and step 3 performs back substitution. For more information on this algorithm, see Section 5 of Numerical solution of saddle point problems.
ClimaCore.MatrixFields.LazyFieldMatrixSolverAlgorithm
— TypeLazyFieldMatrixSolverAlgorithm
A FieldMatrixSolverAlgorithm
that does not require A
to be a FieldMatrix
, i.e., a "matrix-free" algorithm. Internally, a FieldMatrixSolverAlgorithm
(for example, SchurComplementReductionSolve
) might run a LazyFieldMatrixSolverAlgorithm
on a "lazy" representation of a FieldMatrix
(like a LazySchurComplement
).
The only operations used by a LazyFieldMatrixSolverAlgorithm
that depend on A
are lazy_mul
and, when required, lazy_preconditioner
. These and other lazy operations are used to minimize the number of calls to Base.materialize!
, since each call comes with a small performance penalty.
ClimaCore.MatrixFields.StationaryIterativeSolve
— TypeStationaryIterativeSolve(; [kwargs...])
A LazyFieldMatrixSolverAlgorithm
that solves A * x = b
by setting x
to some initial value x[0]
(usually just the zero vector, $\mathbf{0}$) and then iteratively updating it to
\[x[n] = x[n - 1] + \textrm{inv}(P) * (b - A * x[n - 1]).\]
The matrix P
is called a "left preconditioner" for A
. In general, this algorithm converges more quickly when P
is a close approximation of A
, although more complicated approximations often come with a performance penalty.
Background
Let x'
denote the value of x
for which A * x = b
. Replacing b
with A * x'
in the formula for x[n]
tells us that
\[x[n] = x' + (I - \textrm{inv}(P) * A) * (x[n - 1] - x').\]
In other words, the error on iteration n
, x[n] - x'
, can be expressed in terms of the error on the previous iteration, x[n - 1] - x'
, as
\[x[n] - x' = (I - \textrm{inv}(P) * A) * (x[n - 1] - x').\]
By induction, this means that the error on iteration n
is
\[x[n] - x' = (I - \textrm{inv}(P) * A)^n * (x[0] - x').\]
If we pick some norm $||\cdot||$, we find that the norm of the error is bounded by
\[||x[n] - x'|| ≤ ||(I - \textrm{inv}(P) * A)^n|| * ||x[0] - x'||.\]
For any matrix $M$, the spectral radius of $M$ is defined as
\[\rho(M) = \max\{|λ| : λ \text{ is an eigenvalue of } M\}.\]
The spectral radius has the property that
\[||M^n|| \sim \rho(M)^n, \quad \text{i.e.,} \quad \lim_{n \to \infty} \frac{||M^n||}{\rho(M)^n} = 1.\]
So, as the value of n
increases, the norm of the error becomes bounded by
\[||x[n] - x'|| \leq \rho(I - \textrm{inv}(P) * A)^n * ||x[0] - x'||.\]
This indicates that x[n]
will converge to x'
(i.e., that the norm of the error will converge to 0) when ρ(I - inv(P) * A) < 1
, and that the convergence rate is roughly bounded by ρ(I - inv(P) * A)
for large values of n
. More precisely, it can be shown that x[n]
will converge to x'
if and only if ρ(I - inv(P) * A) < 1
. In practice, though, the convergence eventually stops due to the limits of floating point precision.
Also, if we assume that x[n] ≈ x'
, we can use the formula for x[n]
to approximate the error on the previous iteration as
\[x[n - 1] - x' ≈ x[n - 1] - x[n] = \textrm{inv}(P) * (A * x[n - 1] - b).\]
Debugging
This algorithm has support for 2 debugging message group names, which can be passed to the environment variable JULIA_DEBUG
:
error_norm
: prints||x[n] - x'||₂
on every iteration, approximating the errorx[n] - x'
as described abovespectral_radius
: printsρ(I - inv(P) * A)
, approximating this value with theeigsolve
function from KrylovKit.jl
Because the eigsolve
function is not compatible with CUDA, debugging the spectral radius is currently not possible on GPUs.
Keyword Arguments
There are 4 values that can be included in kwargs...
:
P_alg = nothing
: aPreconditionerAlgorithm
that specifies how to computeP
and solveP * x = b
forx
, ornothing
if preconditioning is not required (in which caseP
is effectively set toone(A)
)n_iters = 1
: the number of iterationscorrelated_solves = false
: whether to setx[0]
to a value ofx
that was generated during an earlier call tofield_matrix_solve!
, instead of setting it to $\mathbf{0}$ (it is always set to $\mathbf{0}$ on the first call tofield_matrix_solve!
)eigsolve_kwargs = (;)
: keyword arguments for theeigsolve
function that can be used to tune its accuracy and speed (only applicable when debugging the spectral radius)debug = nothing
: keyword argument for printing debug information to@debug
. By default,debug
is set to true if"error_norm"
or"spectral_radius"
is inENV["JULIA_DEBUG"]
, which must be set by users.
ClimaCore.MatrixFields.ApproximateBlockArrowheadIterativeSolve
— FunctionApproximateBlockArrowheadIterativeSolve(names₁...; [P_alg₁], [alg₁], [alg₂], [kwargs...])
Shorthand for constructing a SchurComplementReductionSolve
whose alg₂
is set to a StationaryIterativeSolve
with a BlockArrowheadSchurComplementPreconditioner
. The keyword argument alg₁
is passed to the constructor for SchurComplementReductionSolve
, the keyword arguments P_alg₁
and alg₂
are passed to the constructor for BlockArrowheadSchurComplementPreconditioner
, and all other keyword arguments are passed to the constructor for StationaryIterativeSolve
.
This algorithm is somewhat similar to a StationaryIterativeSolve
with a BlockArrowheadPreconditioner
, but it usually converges much more quickly, i.e., the spectral radius of its iteration matrix (I - inv(P) * A
) tends to be significantly smaller. Roughly speaking, this is because it runs the iterative solver on an equation with fewer variables (the Schur complement equation, (A₂₂ - A₂₁ * inv(A₁₁) * A₁₂) * x₂ = b₂′
), which means that, on each iteration, it accumulates less error due to coupling between variables. However, even though it converges more quickly, its iterations take longer because they involve using alg₁
to invert A₁₁
. So, when only a few iterations are needed, a StationaryIterativeSolve
with a BlockArrowheadPreconditioner
might be more performant.
This algorithm is an example of a "segregated" solve, in contrast to the alternative "coupled" solve. In the context of computational fluid dynamics, this algorithm can also be viewed as a "SIMPLE" (Semi-Implicit Method for Pressure-Linked Equations) scheme. For more information, see Sections 4, 5, and 10 of Numerical solution of saddle point problems.
Preconditioners
ClimaCore.MatrixFields.PreconditionerAlgorithm
— TypePreconditionerAlgorithm
Description of how to approximate a FieldMatrix
or something similar like a LazySchurComplement
with a preconditioner P
for which P * x = b
is easy to solve for x
. If P
is a diagonal matrix, then x
can be computed as @. inv(P) * b
; otherwise, the PreconditionerAlgorithm
must specify a FieldMatrixSolverAlgorithm
that can be used to solve P * x = b
for x
.
Interface
Every subtype of PreconditionerAlgorithm
must implement methods for the following functions:
ClimaCore.MatrixFields.MainDiagonalPreconditioner
— TypeMainDiagonalPreconditioner()
A PreconditionerAlgorithm
that sets P
to the main diagonal of A
. This is also called a "Jacobi" preconditioner.
ClimaCore.MatrixFields.BlockDiagonalPreconditioner
— TypeBlockDiagonalPreconditioner()
A PreconditionerAlgorithm
that sets P
to the block diagonal entries of A
. This is also called a "block Jacobi" preconditioner.
ClimaCore.MatrixFields.BlockArrowheadPreconditioner
— TypeBlockArrowheadPreconditioner(names₁...; [P_alg₁], [alg₂])
A PreconditionerAlgorithm
for a 2×2 block matrix:
\[A = \begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix}\]
The FieldName
s in names₁
correspond to the subscript ₁
, while all other FieldName
s correspond to the subscript ₂
. The preconditioner P
is set to the following matrix:
\[P = \begin{bmatrix} P_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix}, \quad \text{where } P_{11} \text{ is a diagonal matrix}\]
The internal preconditioner P₁₁
is generated by the PreconditionerAlgorithm
P_alg₁
, which is set to a MainDiagonalPreconditioner
by default. The Schur complement of P₁₁
in P
, A₂₂ - A₂₁ * inv(P₁₁) * A₁₂
, is inverted using the FieldMatrixSolverAlgorithm
alg₂
, which is set to a BlockDiagonalSolve
by default.
ClimaCore.MatrixFields.BlockArrowheadSchurComplementPreconditioner
— TypeBlockArrowheadSchurComplementPreconditioner(; [P_alg₁], [alg₂])
A PreconditionerAlgorithm
that is equivalent to a BlockArrowheadPreconditioner
, but only applied to the Schur complement of A₁₁
in A
, A₂₂ - A₂₁ * inv(A₁₁) * A₁₂
, which is represented by a LazySchurComplement
. Specifically, the preconditioner this generates is the Schur complement of P₁₁
in P
, A₂₂ - A₂₁ * inv(P₁₁) * A₁₂
, where P₁₁
is generated by P_alg₁
. Unlike the BlockArrowheadPreconditioner
constructor, this constructor does not require names₁
because the block structure of A
can be inferred from the LazySchurComplement
.
ClimaCore.MatrixFields.WeightedPreconditioner
— TypeWeightedPreconditioner(M, unweighted_P_alg)
A PreconditionerAlgorithm
that sets P
to M * P′
, where M
is a diagonal FieldMatrix
and P′
is the preconditioner generated by unweighted_P_alg
. When the entries of M
are larger than 1, this is called "relaxation" or "damping"; when the entries are smaller than 1, this is called "extrapolation".
ClimaCore.MatrixFields.CustomPreconditioner
— TypeCustomPreconditioner(M, alg)
A PreconditionerAlgorithm
that sets P
to the FieldMatrix
M
and inverts P
using the FieldMatrixSolverAlgorithm
alg
.
Internals
ClimaCore.MatrixFields.outer_diagonals
— Functionouter_diagonals(::Type{<:BandMatrixRow})
Gets the indices of the lower and upper diagonals, ld
and ud
, of the given subtype of BandMatrixRow
.
ClimaCore.MatrixFields.band_matrix_row_type
— Functionband_matrix_row_type(ld, ud, T)
A shorthand for getting the subtype of BandMatrixRow
that has entries of type T
on the diagonals with indices in the range ld:ud
.
ClimaCore.MatrixFields.matrix_shape
— Functionmatrix_shape(matrix_field, [matrix_space])
Returns the matrix shape for a matrix field defined on the matrix_space
. By default, matrix_space
is set to axes(matrix_field)
.
When the matrixspace is a finite difference space (extruded or otherwise): the shape is either Square()
, FaceToCenter()
, or CenterToFace()
, depending on whether the diagonal indices of `matrixfieldare
Ints or
PlusHalfs and whether
matrix_space` is on cell centers or cell faces.
When the matrix_space is a spectral element or point space: only a Square() shape is supported.
ClimaCore.MatrixFields.column_axes
— Functioncolumn_axes(matrix_field, [matrix_space])
Returns the space that corresponds to the columns of matrix_field
, i.e., the axes
of the Field
s by which matrix_field
can be multiplied. The matrix_space
, on the other hand, is the space that corresponds to the rows of matrix_field
. By default, matrix_space
is set to axes(matrix_field)
.
ClimaCore.MatrixFields.AbstractLazyOperator
— TypeAbstractLazyOperator
Supertype for "lazy operators", i.e., operators that can be called without any arguments by users, as long as they appear in broadcast expressions that contain at least one Field
. If lazy_op
is an AbstractLazyOperator
, the expression lazy_op.()
will internally be translated to non_lazy_op.(fields...)
, as long as it appears in a broadcast expression with at least one Field
. This translation is done by the function replace_lazy_operator(space, lazy_op)
, which must be implemented by every subtype of AbstractLazyOperator
.
ClimaCore.MatrixFields.replace_lazy_operator
— Functionreplace_lazy_operator(space, lazy_op)
Generates an instance of Base.AbstractBroadcasted
that corresponds to the expression lazy_op.()
, where the broadcast in which this expression appears is being evaluated on the given space
. Note that the staggering (CellCenter
or CellFace
) of this space
depends on the specifics of the broadcast and is not predetermined.
ClimaCore.MatrixFields.FieldName
— TypeFieldName(name_chain...)
A singleton type that represents a chain of getproperty
calls, which can be used to access a property or sub-property of an object x
using the function get_field(x, name)
. The entire object x
can also be accessed with the empty FieldName()
.
ClimaCore.MatrixFields.@name
— Macro@name(expr)
Shorthand for constructing a FieldName
. Some examples include
name = @name()
, in which caseget_field(x, name)
returnsx
name = @name(a)
, in which caseget_field(x, name)
returnsx.a
name = @name(a.b.c)
, in which caseget_field(x, name)
returnsx.a.b.c
name = @name(a.b.c.:(1).d)
, in which caseget_field(x, name)
returnsx.a.b.c.:(1).d
This macro is preferred over the FieldName
constructor because it checks whether expr
is a syntactically valid chain of getproperty
calls before calling the constructor.
ClimaCore.MatrixFields.FieldNameTree
— TypeFieldNameTree(x)
Tree of FieldName
s that can be used to access x
with get_field(x, name)
. Check whether a name
is valid by calling is_valid_name(name, tree)
, and extract the children of name
by calling child_names(name, tree)
.
ClimaCore.MatrixFields.FieldNameSet
— TypeFieldNameSet{T}(values, [name_tree])
An AbstractSet
that contains values of type T
, which serves as an analogue of a KeySet
for a FieldNameDict
. There are two subtypes of FieldNameSet
:
FieldVectorKeys
, for whichT
is set toFieldName
FieldMatrixKeys
, for whichT
is set toTuple{FieldName, FieldName}
; each tuple of typeT
represents a pair of row-column indices
Since FieldName
s are singleton types, the result of almost any FieldNameSet
operation can be inferred during compilation. So, with the exception of map
, foreach
, and set_string
, functions of FieldNameSet
s do not have any performance cost at runtime (as long as their arguments are inferrable).
Unlike other AbstractSet
s, FieldNameSet
has special behavior for overlapping values. For example, the FieldName
s @name(a.b)
and @name(a.b.c)
overlap, so any set operation needs to first decompose @name(a.b)
into its child values before combining it with @name(a.b.c)
. In order to support this (and also to support the ability to compute set complements), FieldNameSet
stores a FieldNameTree
name_tree
, which it uses to infer child values. If name_tree
is not specified, it gets set to nothing
by default, which causes some FieldNameSet
operations to become disabled. For binary operations like union
or setdiff
, only one set needs to specify a name_tree
; if two sets both specify a name_tree
, the name_tree
s must be identical.
ClimaCore.MatrixFields.FieldNameDict
— TypeFieldNameDict(keys, entries)
FieldNameDict{T}(key_entry_pairs...)
An AbstractDict
with keys of type T
that are stored as a FieldNameSet{T}
. There are two subtypes of FieldNameDict
:
FieldMatrix
, which maps a set ofFieldMatrixKeys
to eitherColumnwiseBandMatrixField
s or multiples ofLinearAlgebra.I
; this is the only user-facing subtype ofFieldNameDict
FieldVectorView
, which maps a set ofFieldVectorKeys
toField
s; this subtype is automatically generated when aFieldVector
is used in the same operation as aFieldMatrix
(e.g., when both appear in the same broadcast expression, or when both are passed to aFieldMatrixSolver
)
A FieldNameDict
can also be "lazy", which means that it can store AbstractBroadcasted
objects that become Field
s when they are materialized. Many internal operations generate lazy FieldNameDict
s to reduce the number of calls to materialize!
, since each call comes with a small performance penalty.
The entry at a specific key can be extracted by calling dict[key]
, and the entries that correspond to all the keys in a FieldNameSet
can be extracted by calling dict[set]
. If dict
is a FieldMatrix
, the corresponding identity matrix can be computed by calling one(dict)
.
When broadcasting over FieldNameDict
s, the following operations are supported:
- Addition and subtraction
- Multiplication, where the first argument must be a
FieldMatrix
- Inversion, where the argument must be a diagonal
FieldMatrix
, i.e., one in which every entry is either aColumnwiseBandMatrixField
ofDiagonalMatrixRow
s or a multiple ofLinearAlgebra.I
ClimaCore.MatrixFields.field_vector_view
— Functionfield_vector_view(x, [name_tree])
Constructs a FieldVectorView
that contains all of the Field
s in the FieldVector
x
. The default name_tree
is FieldNameTree(x)
, but this can be modified if needed.
ClimaCore.MatrixFields.concrete_field_vector
— Functionconcrete_field_vector(vector)
Converts the FieldVectorView
vector
back into a FieldVector
.
ClimaCore.MatrixFields.is_lazy
— Functionis_lazy(dict)
Checks whether the FieldNameDict
dict
contains any un-materialized AbstractBroadcasted
entries.
ClimaCore.MatrixFields.lazy_main_diagonal
— Functionlazy_main_diagonal(matrix)
Constructs a lazy FieldMatrix
that contains the main diagonal of matrix
.
ClimaCore.MatrixFields.lazy_mul
— Functionlazy_mul(A, args...)
Constructs a lazy FieldMatrix
that represents the product @. *(A, args...)
. This involves regular broadcasting when A
is a FieldMatrix
, but it has more complex behavior for other objects like the LazySchurComplement
.
ClimaCore.MatrixFields.LazySchurComplement
— TypeLazySchurComplement(A₁₁, A₁₂, A₂₁, A₂₂, [alg₁, cache₁, A₁₂_x₂, invA₁₁_A₁₂_x₂])
An analogue of a FieldMatrix
that represents the Schur complement of A₁₁
in A
, A₂₂ - A₂₁ * inv(A₁₁) * A₁₂
. Since inv(A₁₁)
will generally be a dense matrix, it would not be efficient to directly compute the Schur complement. So, this object only supports the "lazy" functions lazy_mul
, which allows it to be multiplied by the vector x₂
, and lazy_preconditioner
, which allows it to be approximated with a FieldMatrix
.
The values alg₁
, cache₁
, A₁₂_x₂
, and invA₁₁_A₁₂_x₂
need to be specified in order for lazy_mul
to be able to compute inv(A₁₁) * A₁₂ * x₂
. When a LazySchurComplement
is not passed to lazy_mul
, these values can be omitted.
ClimaCore.MatrixFields.field_matrix_solver_cache
— Functionfield_matrix_solver_cache(alg, A, b)
Allocates the cache required by the FieldMatrixSolverAlgorithm
alg
to solve the equation A * x = b
.
ClimaCore.MatrixFields.check_field_matrix_solver
— Functioncheck_field_matrix_solver(alg, cache, A, b)
Checks that the sparsity structure of A
is supported by the FieldMatrixSolverAlgorithm
alg
, and that A
is compatible with b
in the equation A * x = b
.
ClimaCore.MatrixFields.run_field_matrix_solver!
— Functionrun_field_matrix_solver!(alg, cache, x, A, b)
Sets x
to the value that solves the equation A * x = b
using the FieldMatrixSolverAlgorithm
alg
.
ClimaCore.MatrixFields.solver_algorithm
— Functionsolver_algorithm(P_alg)
A FieldMatrixSolverAlgorithm
that can be used to solve P * x = b
for x
, where P
is the preconditioner generated by the PreconditionerAlgorithm
P_alg
. If P_alg
is nothing
instead of a PreconditionerAlgorithm
, or if P
is a diagonal matrix (and no solver is required to invert it), this returns nothing
.
ClimaCore.MatrixFields.lazy_preconditioner
— Functionlazy_preconditioner(P_alg, A)
Constructs a lazy FieldMatrix
(or a concrete one when possible) that approximates A
according to the PreconditionerAlgorithm
P_alg
. If P_alg
is nothing
instead of a PreconditionerAlgorithm
, this returns one(A)
.
ClimaCore.MatrixFields.preconditioner_cache
— Functionpreconditioner_cache(P_alg, A, b)
Allocates the cache required to solve the equation P * x = b
, where P
is the preconditioner generated by the PreconditionerAlgorithm
P_alg
for A
.
ClimaCore.MatrixFields.check_preconditioner
— Functioncheck_preconditioner(P_alg, P_cache, A, b)
Checks that P
is compatible with b
in the equation P * x = b
, where P
is the preconditioner generated by the PreconditionerAlgorithm
P_alg
for A
. If P_alg
requires a FieldMatrixSolverAlgorithm
alg
to solve the equation, this also calls check_field_matrix_solver
on alg
.
ClimaCore.MatrixFields.lazy_or_concrete_preconditioner
— Functionlazy_or_concrete_preconditioner(P_alg, P_cache, A)
A wrapper for lazy_preconditioner
that turns the lazy FieldMatrix
P
into a concrete FieldMatrix
when the PreconditionerAlgorithm
P_alg
requires a FieldMatrixSolverAlgorithm
to invert it.
ClimaCore.MatrixFields.apply_preconditioner
— Functionapply_preconditioner(P_alg, P_cache, P, lazy_b)
Constructs a lazy FieldMatrix
(or a concrete one when possible) that represents the product @. inv(P) * b
. Here, lazy_b
is a (possibly lazy) FieldVectorView
that represents b
.
Utilities
ClimaCore.MatrixFields.column_field2array
— Functioncolumn_field2array(field)
Converts a field defined on a FiniteDifferenceSpace
into either a Vector
or a BandedMatrix
, depending on whether the elements of the field are single values or BandMatrixRow
s. This involves copying the data stored in the field. Because BandedMatrix
does not currently support operations with CuArray
s, all GPU data is copied to the CPU.
ClimaCore.MatrixFields.column_field2array_view
— Functioncolumn_field2array_view(field)
Similar to column_field2array(field)
, except that this version avoids copying the data stored in the field.
ClimaCore.MatrixFields.field2arrays
— Functionfield2arrays(field)
Converts a field defined on a FiniteDifferenceSpace
or on an ExtrudedFiniteDifferenceSpace
into a collection of arrays, each of which corresponds to a column of the field. This is done by calling column_field2array
on each of the field's columns.
ClimaCore.MatrixFields.field2arrays_view
— Functionfield2arrays_view(field)
Similar to field2arrays(field)
, except that this version calls column_field2array_view
instead of column_field2array
.