Merge pull request #1788 from CliMA/dy/field_matrix_wrapper

Add wrapper for FieldMatrix and FieldMatrixSolver

NEWS.md

@@ -4,6 +4,12 @@ ClimaCore.jl Release Notes
 main
 -------
 
+v0.14.8
+-------
+
+- ![][badge-✨feature/enhancement] Added `FieldMatrixWithSolver`, a wrapper that helps defining implicit Jacobians. PR [#1788](https://github.com/CliMA/ClimaCore.jl/pull/1788)
+
+
 v0.14.6
 -------
 

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaCore"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 authors = ["CliMA Contributors <[email protected]>"]
-version = "0.14.7"
+version = "0.14.8"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/matrix_fields.md

@@ -31,6 +31,7 @@ operator_matrix
 ```@docs
 FieldMatrixSolverAlgorithm
 FieldMatrixSolver
+FieldMatrixWithSolver
 field_matrix_solve!
 BlockDiagonalSolve
 BlockLowerTriangularSolve

src/MatrixFields/MatrixFields.jl

@@ -43,7 +43,8 @@ multiples of `LinearAlgebra.I`. This comes with the following functionality:
 """
 module MatrixFields
 
-import LinearAlgebra: I, UniformScaling, Adjoint, AdjointAbsVec, mul!, inv, norm
+import LinearAlgebra: I, UniformScaling, Adjoint, AdjointAbsVec
+import LinearAlgebra: inv, norm, ldiv!, mul!
 import StaticArrays: SMatrix, SVector
 import BandedMatrices: BandedMatrix, band, _BandedMatrix
 import RecursiveArrayTools: recursive_bottom_eltype
@@ -73,7 +74,7 @@ export DiagonalMatrixRow,
     QuaddiagonalMatrixRow,
     PentadiagonalMatrixRow
 export FieldVectorKeys, FieldMatrixKeys, FieldVectorView, FieldMatrix
-export ⋅, FieldMatrixSolver, field_matrix_solve!
+export FieldMatrixWithSolver, ⋅
 
 # Types that are teated as single values when using matrix fields.
 const SingleValue =
@@ -106,6 +107,7 @@ include("single_field_solver.jl")
 include("multiple_field_solver.jl")
 include("field_matrix_solver.jl")
 include("field_matrix_iterative_solver.jl")
+include("field_matrix_with_solver.jl")
 
 function Base.show(io::IO, field::ColumnwiseBandMatrixField)
     print(io, eltype(field), "-valued Field")

src/MatrixFields/field_matrix_with_solver.jl

@@ -0,0 +1,48 @@
+"""
+    FieldMatrixWithSolver(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`](@ref), so a custom `alg` must be specified when `A` is
+not a block diagonal matrix.
+"""
+struct FieldMatrixWithSolver{M, S} <: AbstractDict{FieldNamePair, Any}
+    matrix::M
+    solver::S
+end
+FieldMatrixWithSolver(
+    A::FieldMatrix,
+    b::Fields.FieldVector,
+    alg::FieldMatrixSolverAlgorithm = BlockDiagonalSolve(),
+) = FieldMatrixWithSolver(A, FieldMatrixSolver(alg, A, b))
+
+# TODO: Find a simple way to make b an optional argument and add a method for
+# Base.one(::FieldMatrixWithSolver).
+
+Base.keys(A::FieldMatrixWithSolver) = keys(A.matrix)
+
+Base.values(A::FieldMatrixWithSolver) = values(A.matrix)
+
+Base.pairs(A::FieldMatrixWithSolver) = pairs(A.matrix)
+
+Base.length(A::FieldMatrixWithSolver) = length(A.matrix)
+
+Base.iterate(A::FieldMatrixWithSolver, index = 1) = iterate(A.matrix, index)
+
+Base.getindex(A::FieldMatrixWithSolver, key) = getindex(A.matrix, key)
+
+Base.:(==)(A1::FieldMatrixWithSolver, A2::FieldMatrixWithSolver) =
+    A1.matrix == A2.matrix && A1.solver.alg == A2.solver.alg
+
+Base.similar(A::FieldMatrixWithSolver) =
+    FieldMatrixWithSolver(similar(A.matrix), A.solver)
+
+ldiv!(x::Fields.FieldVector, A::FieldMatrixWithSolver, b::Fields.FieldVector) =
+    field_matrix_solve!(A.solver, x, A.matrix, b)
+
+mul!(b::Fields.FieldVector, A::FieldMatrixWithSolver, x::Fields.FieldVector) =
+    @. b = A.matrix * x

test/MatrixFields/field_matrix_solvers.jl

@@ -4,7 +4,7 @@ using Revise; include(joinpath("test", "MatrixFields", "field_matrix_solvers.jl"
 =#
 import Logging
 import Logging: Debug
-import LinearAlgebra: I, norm
+import LinearAlgebra: I, norm, ldiv!, mul!
 import ClimaComms
 import ClimaCore.Utilities: half
 import ClimaCore.RecursiveApply: ⊠
@@ -14,26 +14,16 @@ import ClimaCore:
 
 include("matrix_field_test_utils.jl")
 
-# This broadcast must be wrapped in a function to be tested with @test_opt.
-field_matrix_mul!(b, A, x) = @. b = A * x
-
 function test_field_matrix_solver(; test_name, alg, A, b, use_rel_error = false)
     @testset "$test_name" begin
         x = similar(b)
-        b_test = similar(b)
-        solver = FieldMatrixSolver(alg, A, b)
-        args = (solver, x, A, b)
-
+        A′ = FieldMatrixWithSolver(A, b, alg)
         solve_time =
-            @benchmark ClimaComms.@cuda_sync comms_device field_matrix_solve!(
-                args...,
-            )
+            @benchmark ClimaComms.@cuda_sync comms_device ldiv!(x, A′, b)
+
+        b_test = similar(b)
         mul_time =
-            @benchmark ClimaComms.@cuda_sync comms_device field_matrix_mul!(
-                b_test,
-                A,
-                x,
-            )
+            @benchmark ClimaComms.@cuda_sync comms_device mul!(b_test, A′, x)
 
         solve_time_rounded = round(solve_time; sigdigits = 2)
         mul_time_rounded = round(mul_time; sigdigits = 2)
@@ -70,14 +60,13 @@ function test_field_matrix_solver(; test_name, alg, A, b, use_rel_error = false)
             AnyFrameModule(Base.CoreLogging),
         )
         using_cuda ||
-            @test_opt ignored_modules = ignored FieldMatrixSolver(alg, A, b)
-        using_cuda ||
-            @test_opt ignored_modules = ignored field_matrix_solve!(args...)
-        @test_opt ignored_modules = ignored field_matrix_mul!(b, A, x)
+            @test_opt ignored_modules = ignored FieldMatrixWithSolver(A, b, alg)
+        using_cuda || @test_opt ignored_modules = ignored ldiv!(x, A′, b)
+        @test_opt ignored_modules = ignored mul!(b_test, A′, x)
 
         # TODO: fix broken test when Nv is added to the type space
-        using_cuda || @test @allocated(field_matrix_solve!(args...)) ≤ 1536
-        using_cuda || @test @allocated(field_matrix_mul!(b, A, x)) == 0
+        using_cuda || @test @allocated(ldiv!(x, A′, b)) ≤ 1536
+        using_cuda || @test @allocated(mul!(b_test, A′, x)) == 0
     end
 end
 
@@ -335,8 +324,7 @@ end
             b = Fields.FieldVector(; c = ᶜvec, f = ᶠvec)
 
             x = similar(b)
-            solver = FieldMatrixSolver(alg, A, b)
-            args = (solver, x, A, b)
+            A′ = FieldMatrixWithSolver(A, b, alg)
 
             # Compare the debugging logs to RegEx strings. Note that debugging the
             # spectral radius is currently not possible on GPUs.
@@ -348,9 +336,7 @@ end
                 (:debug, r"||x[2] - x'||₂ ≈"),
             )
             logs = (spectral_radius_logs..., error_norm_logs...)
-            @test_logs logs... min_level = Logging.Debug field_matrix_solve!(
-                args...,
-            )
+            @test_logs logs... min_level = Logging.Debug ldiv!(x, A′, b)
         end
     end
 end
@@ -567,9 +553,8 @@ end
     field = Fields.ones(space)
     b = Fields.FieldVector(; _ = field)
     x = similar(b)
-    solver =
-        MatrixFields.FieldMatrixSolver(MatrixFields.BlockDiagonalSolve(), A, b)
+    A′ = FieldMatrixWithSolver(A, b)
 
     # Run matrix solve
-    MatrixFields.field_matrix_solve!(solver, x, A, b)
+    ldiv!(x, A′, b)
 end