Merge #1399

1399: Add operator_matrix to MatrixFields, along with tests and docs r=dennisYatunin a=dennisYatunin

## Purpose

Second PR of #1230. Refactors what is currently in `src/operators/operator2stencil.jl`.

## Content

- Adds the function `operator_matrix(op)`, which will replace `Operator2Stencil(op)`. This function has a detailed docstring and throws descriptive errors for operators without well-defined operator matrices.
- Improves the usability of operator matrices. With this new interface, users will no longer need to specify intermediate fields to compute operator matrices. For example, with our pre-existing code, the matrix of the interpolation operator (a bidiagonal matrix whose entries are `-1/2` and `+1/2`) needs to be computed with ``@.` matrix_field = interp_matrix(ones_field)`, where `ones_field` is a field filled with the number `1` that is used to infer the space and entry type of the matrix. Now, this matrix can just be computed with ``@.` matrix_field = interp_matrix()`. In order to make this work, `interp_matrix` is now defined as a "lazy operator". When a broadcast expression containing lazy operators is evaluated, each lazy operator is replaced with an actual operator, and it is given one or more fields as input arguments. In this case, `interp_matrix` is given the local geometry field as an input argument, and this field is used to infer the space and entry type of the operator matrix.
- This usability improvement slightly changes the computation of derivative matrices. With our pre-existing code, ``@.` op(func(field))` is equivalent to ``@.` op_matrix(ones_field) ⋅ func(field)`, and the derivative of this expression with respect to `field` can be specified as ``@.` op_matrix(func'(field))`, where `func'` is the derivative of the point-wise function `func`. With this new interface, ``@.` op(func(field))` is equivalent to ``@.` op_matrix() ⋅ func(field)`, and the derivative can be specified as ``@.` op_matrix() ⋅ DiagonalMatrixRow(func'(field))`. Similarly, the derivative of ``@.` op2(func2(op1(func1(field))))` with respect to `field` is `op2_matrix(func2'(op1(func1(field)))) ⋅ op1_matrix(func1'(field))` with our pre-existing code and `op2_matrix() ⋅ DiagonalMatrixRow(func2'(op1(func1(field)))) ⋅ op1_matrix() ⋅ DiagonalMatrixRow(func1'(field))` with the new interface. Although the new interface leads to longer derivative expressions, those expressions are more similar to how the chain rule is usually written out, and they can be debugged/analyzed more incrementally.
- Adds support for computing operator matrices of multi-argument operators. For example, if `op` is the `Upwind3rdOrderBiasedProductC2F` operator, then ``@.` op(velocity_field, tracer_field)` is equivalent to ``@.` op_matrix(velocity_field) ⋅ tracer_field`. The implementation is similar to that of single-argument operators, except that it does not require the use of "lazy operators" (since there is already a field being passed to the operator matrix, the local geometry field can be obtained from that field during the evaluation of `Base.Broadcast.broadcasted`).
- Adds support for computing operator matrices of operators with `Extrapolate` boundary conditions. These boundary conditions cause the matrices to have larger bandwidths than other boundary conditions.
- Tests the `operator_matrix` function with every valid combination of finite difference operators and boundary conditions. The tests check for correctness, type stability, and lack of allocations. The tests are run on both CPUs and GPUs. In addition, the tests print out how the performance of ``@.` op_matrix() ⋅ field` compares to the performance of ``@.` op(field)`; the two expressions are similarly fast on GPUs (between `-70%` and `+40%` relative change in speed), though the operator matrix expressions tend to be slower on CPUs.
- Tests a few more complicated broadcast expressions involving products and linear combinations of operator matrices. These tests indicate that operator matrices are similarly performant to regular matrix fields.
- Modifies `test_field_broadcast` so that it also tests whether `get_result` generates the same result as `set_result!`.
- Modifies the `*` method for `BandMatrixRow` so that matrix fields can be scaled by vectors/covectors in addition to numbers. This simplifies a few of the complicated broadcast tests.
- Fixes a typo (`Geometery`) in the method of `stencil_left_boundary` for `GradientF2C`.
- Adds a missing method to `rpromote_type` that was preventing empty matrix rows from being constructed.
- Modifies the `Base.Broadcast.broadcasted` method for `FiniteDifferenceOperator` and `SpectralElementOperator` so that lazy operators can work correctly (the original versions of these methods would always overwrite the `LazyOperatorStyle` with `StencilStyle` or `SpectralStyle`, respectively).
- Unfortunately, this PR also adds 2 method invalidations. These are due to a new definition of `broadcasted` for lazy operators:
    ```
    Base.Broadcast.broadcasted(::LazyOperatorStyle, f::F, args...) where {F} =
        LazyOperatorBroadcasted(f, args)
    ```
    As explained in an accompanying code comment, removing this method requires modifying several other method definitions, and one of these modifications adds 11 invalidations. So, there doesn't seem to be a good way to avoid these 2 invalidations.

Co-authored-by: Dennis Yatunin <[email protected]>

.buildkite/pipeline.yml

@@ -533,6 +533,17 @@ steps:
           slurm_gpus: 1
           slurm_mem: 40GB
 
+      - label: "Unit: operator matrices (CPU)"
+        key: unit_operator_matrices_cpu
+        command: "julia --color=yes --check-bounds=yes --project=test test/MatrixFields/operator_matrices.jl"
+
+      - label: "Unit: operator matrices (GPU)"
+        key: unit_operator_matrices_gpu
+        command: "julia --color=yes --project=test test/MatrixFields/operator_matrices.jl"
+        agents:
+          slurm_gpus: 1
+          slurm_mem: 40GB
+
   - group: "Unit: Hypsography"
     steps:
 

docs/src/matrix_fields.md

@@ -20,6 +20,12 @@ BandMatrixRow
 MultiplyColumnwiseBandMatrixField
 ```
 
+## Operator Matrices
+
+```@docs
+operator_matrix
+```
+
 ## Internals
 
 ```@docs
@@ -31,6 +37,8 @@ mul_return_type
 rmul_return_type
 matrix_shape
 column_axes
+AbstractLazyOperator
+replace_lazy_operator
 ```
 
 ## Utilities

src/MatrixFields/MatrixFields.jl

@@ -48,10 +48,16 @@ export DiagonalMatrixRow,
     QuaddiagonalMatrixRow,
     PentadiagonalMatrixRow
 
+# Types that are teated as single values when using matrix fields.
+const SingleValue =
+    Union{Number, Geometry.AxisTensor, Geometry.AdjointAxisTensor}
+
 include("band_matrix_row.jl")
 include("rmul_with_projection.jl")
 include("matrix_shape.jl")
 include("matrix_multiplication.jl")
+include("lazy_operators.jl")
+include("operator_matrices.jl")
 include("field2arrays.jl")
 
 const ColumnwiseBandMatrixField{V, S} = Fields.Field{

src/MatrixFields/band_matrix_row.jl

@@ -128,9 +128,9 @@ Base.:-(row1::BandMatrixRow, row2::UniformScaling) =
 Base.:-(row1::UniformScaling, row2::BandMatrixRow) =
     map(rsub, promote(row1, row2)...)
 
-Base.:*(row::BandMatrixRow, value::Number) =
+Base.:*(row::BandMatrixRow, value::SingleValue) =
     map(entry -> rmul(entry, value), row)
-Base.:*(value::Number, row::BandMatrixRow) =
+Base.:*(value::SingleValue, row::BandMatrixRow) =
     map(entry -> rmul(value, entry), row)
 
 Base.:/(row::BandMatrixRow, value::Number) =

src/MatrixFields/lazy_operators.jl

@@ -0,0 +1,121 @@
+"""
+    AbstractLazyOperator
+
+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`.
+"""
+abstract type AbstractLazyOperator end
+
+struct LazyOperatorStyle <: Base.Broadcast.BroadcastStyle end
+
+Base.Broadcast.broadcasted(op::AbstractLazyOperator) =
+    Base.Broadcast.broadcasted(LazyOperatorStyle(), op)
+
+# Broadcasting over an AbstractLazyOperator and either a Ref, a Tuple, a Field,
+# an Operator, or another AbstractLazyOperator involves using LazyOperatorStyle.
+Base.Broadcast.BroadcastStyle(
+    ::LazyOperatorStyle,
+    ::Union{
+        Base.Broadcast.AbstractArrayStyle{0},
+        Base.Broadcast.Style{Tuple},
+        Fields.AbstractFieldStyle,
+        LazyOperatorStyle,
+    },
+) = LazyOperatorStyle()
+
+struct LazyOperatorBroadcasted{F, A} <:
+       Operators.OperatorBroadcasted{LazyOperatorStyle}
+    f::F
+    args::A
+end
+
+# TODO: This definition of Base.Broadcast.broadcasted results in 2 additional
+# method invalidations when using Julia 1.8.5. However, if we were to delete it,
+# we would also need to replace the following specializations on
+# LazyOperatorBroadcasted with specializations on Base.Broadcast.Broadcasted.
+# Specifically, we would need to modify Base.Broadcast.materialize so that it
+# specializes on Base.Broadcast.Broadcasted{LazyOperatorStyle}, and this would
+# result in 11 invalidations instead of 2.
+Base.Broadcast.broadcasted(::LazyOperatorStyle, f::F, args...) where {F} =
+    LazyOperatorBroadcasted(f, args)
+
+function Base.Broadcast.materialize(bc::LazyOperatorBroadcasted)
+    space = largest_space(bc)
+    isnothing(space) && error("Cannot materialize broadcast expression with \
+                               AbstractLazyOperator because it does not contain any Fields")
+    return Base.Broadcast.materialize(replace_lazy_operators(space, bc))
+end
+
+Base.Broadcast.materialize!(dest::Fields.Field, bc::LazyOperatorBroadcasted) =
+    Base.Broadcast.materialize!(dest, replace_lazy_operators(axes(dest), bc))
+
+replace_lazy_operators(_, arg) = arg
+replace_lazy_operators(space, bc::LazyOperatorBroadcasted) =
+    bc.f isa AbstractLazyOperator ? replace_lazy_operator(space, bc.f) :
+    Base.Broadcast.broadcasted(
+        bc.f,
+        replace_lazy_operators_args(space, bc.args...)...,
+    )
+
+replace_lazy_operators_args(_) = ()
+replace_lazy_operators_args(space, arg, args...) = (
+    replace_lazy_operators(space, arg),
+    replace_lazy_operators_args(space, args...)...,
+)
+
+"""
+    replace_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.
+"""
+replace_lazy_operator(_, ::AbstractLazyOperator) =
+    error("Every subtype of AbstractLazyOperator must implement a method for
+           replace_lazy_operator(space, lazy_op)")
+
+largest_space(_) = nothing
+largest_space(field::Fields.Field) = axes(field)
+largest_space(bc::Base.AbstractBroadcasted) = largest_space_args(bc.args...)
+
+largest_space_args() = nothing
+largest_space_args(arg, args...) =
+    larger_space(largest_space(arg), largest_space_args(args...))
+
+larger_space(::Nothing, ::Nothing) = nothing
+larger_space(space1, ::Nothing) = space1
+larger_space(::Nothing, space2) = space2
+larger_space(space1::S, ::S) where {S} = space1 # Neither space is larger.
+larger_space(
+    space1::Spaces.FiniteDifferenceSpace,
+    ::Spaces.FiniteDifferenceSpace,
+) = space1 # The staggering does not matter here, so neither space is larger.
+larger_space(
+    space1::Spaces.ExtrudedFiniteDifferenceSpace,
+    ::Spaces.ExtrudedFiniteDifferenceSpace,
+) = space1 # The staggering does not matter here, so neither space is larger.
+larger_space(
+    space1::Spaces.ExtrudedFiniteDifferenceSpace,
+    ::Spaces.FiniteDifferenceSpace,
+) = space1 # The types indicate that space2 is a subspace of space1.
+larger_space(
+    ::Spaces.FiniteDifferenceSpace,
+    space2::Spaces.ExtrudedFiniteDifferenceSpace,
+) = space2 # The types indicate that space1 is a subspace of space2.
+larger_space(
+    space1::Spaces.ExtrudedFiniteDifferenceSpace,
+    ::Spaces.AbstractSpectralElementSpace,
+) = space1 # The types indicate that space2 is a subspace of space1.
+larger_space(
+    ::Spaces.AbstractSpectralElementSpace,
+    space2::Spaces.ExtrudedFiniteDifferenceSpace,
+) = space2 # The types indicate that space1 is a subspace of space2.
+larger_space(::S1, ::S2) where {S1, S2} =
+    error("Mismatched spaces ($(S1.name.name) and $(S2.name.name))")