Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
1436: Add FieldMatrix and linear solvers r=dennisYatunin a=dennisYatunin ## Purpose Third PR of #1230. Adds an interface for specifying block matrices of matrix fields and solving linear systems with these matrices. This will replace what is currently in `ClimaAtmos.jl/src/prognostic_equations/implicit/schur_complement_W.jl`, generalizing it for implicit diffusion and implicit EDMF. ## Content ### Main Changes - Add the `FieldName` struct, which is a singleton type that represents a chain of `getproperty` calls - Add the ``@name`` macro for constructing `FieldName`s, which checks whether its input expression is a syntactically valid chain of `getproperty` calls before calling the default constructor. - A `name` can be used to access a property or sub-property of an object `x` by calling `get_field(x, name)`. - An `internal_name` can be appended onto another `name` in order to access a property or sub-property of `get_field(x, name)`. - Add the `FieldNameDict` struct, which maps each key in a set of `FieldVectorKeys` or `FieldMatrixKeys` (see below) to a `Field` or some other object. - There are currently four subtypes of `FieldNameDict`: - `FieldMatrix` (the only user-facing subtype), which maps `NTuple{2, FieldName}`s to `ColumnwiseBandMatrixField`s or multiples of `LinearAlgebra.I` - `FieldVectorView`, which maps `FieldName`s to `Field`s; this is used to wrap a `FieldVector` so that it can be used in conjunction with a `FieldMatrix` - `FieldVectorViewBroadcasted` and `FieldMatrixBroadcasted`, each of which can store unevaluated `Base.AbstractBroadcasted` objects, in addition to what `FieldVectorView` and `FieldMatrix` can already store - Supports standard `AbstractDict` functions like `keys` and `pairs`. - An individual block of a `FieldNameDict` can be accessed by calling `dict[key]`, and a range of blocks can be accessed by calling `dict[set]`, where `set` is a `FieldNameSet`. - Given a `FieldMatrix` `A`, a similar matrix that only contains identity matrix blocks can be constructed with `one(A)`. - `FieldNameDict`s can be used in broadcast expressions, which support the following operations: - `+`, `-`, or `*`, where each input is either a `FieldNameDict` or a `FieldVector` - `inv`, where the input is a diagonal `FieldMatrix` - The new methods for `Base.Broadcast.broadcasted` construct chains of `Field` broadcast expressions from `FieldNameDict` broadcast expressions on the fly, somewhat similarly to how broadcasting works for ClimaCore operators. - Add the `FieldMatrixSolver` struct, which solves an equation of the form `A * x = b`, where `A` is a `FieldMatrix` and where `x` and `b` are `FieldVector`s. - Add the `field_matrix_solve!` function, which works just like `ldiv!(x, A, b)`, except that it also takes a `FieldMatrixSolver` as its first argument. - Add four `FieldMatrixSolverAlgorithm`s, which can be nested inside of each other to build up more specialized algorithms: - `BlockDiagonalSolve`, which runs a "single field solver" for each block of the block diagonal matrix `A`; the single field solver can handle the four types of blocks: - Multiples of `LinearAlgebra.I` - Diagonal `ColumnwiseBandMatrixField`s - Tri-diagonal `ColumnwiseBandMatrixField`s (implementation of the Thomas algorithm) - Penta-diagonal `ColumnwiseBandMatrixField`s (implementation of the PTRANS-I algorithm) - `BlockLowerTriangularSolve`, which uses forward substitution to solve the equation for a block lower triangular matrix `A` - `SchurComplementSolve`, which generalizes what is currently in ClimaAtmos's `schur_complement_W.jl` file to any block matrix `A` with a diagonal block in the top-left corner - `ApproximateFactorizationSolve`, which lets us use "operator splitting" to approximately solve the equation for a diagonally dominant block matrix `A` - Add documentation for how to specify a `FieldMatrix` and use it in a linear solver, along with internal documentation for the new `FieldName`-based infrastructure. - Add unit tests for correctness, type stability, and allocations, and run them on both CPUs and GPUs through CI. - Test each single field solver on both a cell-center and a cell-face field. - Test each `FieldMatrixSolverAlgorithm` on block diagonal, block lower triangular, and block dense matrices. - Test solvers with identical structures to what we will use in ClimaAtmos for the following examples: - Dry dycore with implicit acoustic waves - Dry dycore with implicit acoustic waves and diffusion - Dry dycore + prognostic EDMF with implicit acoustic waves and SGS fluxes - Moist dycore + prognostic EDMF + tracers with implicit acoustic waves and SGS fluxes ### Internal Chagnes - Add a collection of "unrolled functions", whose return values can be inferred during compilation if their input values are all singleton types. - These are all implemented as combinations of `unrolled_zip`, `unrolled_map`, and `unrolled_foldl`. - Several of these need to have their recursion limits disabled for the unit tests to be type stable. - Add the `FieldNameTree` struct, which stores every `FieldName` that can be used to access `x` with `get_field(x, name)`. - A `name` can be checked for validity by calling `has_subtree_at_name(tree, name)`. - The children of `name` (the `FieldName`s that can be used to access the properties of `get_field(x, name)`) can be obtained by calling `child_names(name, tree)`. - Add the `FieldNameSet` struct, which stores a set of `FieldVectorKeys` (each of which is a `FieldName`) or a set of `FieldMatrixKeys` (each of which is an `NTuple{2, FieldName}`). - Roughly equivalent to the built-in `KeySet` for `AbstractDict`s, but specialized for `FieldNameDict`s. - Supports standard `AbstractSet` functions like `union` and `setdiff`, as well as custom functions like `set_complement` and `matrix_product_keys`. - Handles overlaps between `FieldName`s (that is, situations where one property of `x` lies inside another property of `x`) by storing a `FieldNameTree` that contains all available `FieldName`s. - Disable the recursion limits for several functions used to manipulate `FieldName`s, `FieldNameTree`s, and `FieldNameSet`s, as this is necessary in order for the unit tests to be type stable. - Remove the methods for `RecursiveApply.rmul` that specialize on `Number`, which is also necessary in order for the unit tests to be type stable. - These methods are no longer required, now that #1454 has been merged in. - Add support for calling `inv` on `BandMatrixRow`s. - Qualify the use of `CUDA.`@allowscalar`,` per Charlie's suggestion. - Fix some type instabilities in `matrix_field_test_utils.jl`. - Remove an unused variable name in `operator_matrices.jl`. Co-authored-by: Dennis Yatunin <[email protected]>
- Loading branch information
