Merge pull request #2140 from CliMA/ck/define_zero

Define zero on some MatrixField types

Author: charleskawczynski

src/MatrixFields/field_matrix_solver.jl

@@ -168,6 +168,9 @@ end
 LazySchurComplement(A₁₁, A₁₂, A₂₁, A₂₂) =
     LazySchurComplement(A₁₁, A₁₂, A₂₁, A₂₂, nothing, nothing, nothing, nothing)
 
+Base.zero(lsc::LazySchurComplement) =
+    LazySchurComplement(map(fn -> zero(getfield(lsc, fn)), fieldnames(lsc))...)
+
 NVTX.@annotate function lazy_mul(A₂₂′::LazySchurComplement, x₂)
     (; A₁₁, A₁₂, A₂₁, A₂₂, alg₁, cache₁, A₁₂_x₂, invA₁₁_A₁₂_x₂) = A₂₂′
     zero_rows = setdiff(keys(A₁₂_x₂), matrix_row_keys(keys(A₁₂)))
@@ -229,6 +232,8 @@ partial pivoting matrix).
 """
 struct BlockDiagonalSolve <: FieldMatrixSolverAlgorithm end
 
+Base.zero(alg::BlockDiagonalSolve) = alg
+
 function field_matrix_solver_cache(::BlockDiagonalSolve, A, b)
     caches = map(matrix_row_keys(keys(A))) do name
         single_field_solver_cache(A[name, name], b[name])
@@ -315,6 +320,9 @@ BlockLowerTriangularSolve(
     alg₂ = BlockDiagonalSolve(),
 ) = BlockLowerTriangularSolve(names₁, alg₁, alg₂)
 
+Base.zero(alg::BlockLowerTriangularSolve) =
+    BlockLowerTriangularSolve(alg.names₁, zero(alg.alg₁), zero(alg.alg₂))
+
 function field_matrix_solver_cache(alg::BlockLowerTriangularSolve, A, b)
     A₁₁, _, _, A₂₂, b₁, b₂ = partition_blocks(alg.names₁, A, b)
     cache₁ = field_matrix_solver_cache(alg.alg₁, A₁₁, b₁)
@@ -448,6 +456,9 @@ end
 SchurComplementReductionSolve(names₁...; alg₁ = BlockDiagonalSolve(), alg₂) =
     SchurComplementReductionSolve(names₁, alg₁, alg₂)
 
+Base.zero(alg::SchurComplementReductionSolve) =
+    SchurComplementReductionSolve(alg.names₁, zero(alg.alg₁), zero(alg.alg₂))
+
 function field_matrix_solver_cache(alg::SchurComplementReductionSolve, A, b)
     A₁₁, A₁₂, A₂₁, A₂₂, b₁, b₂ = partition_blocks(alg.names₁, A, b)
     b₁′ = similar(b₁)

src/MatrixFields/field_matrix_with_solver.jl

@@ -41,6 +41,9 @@ Base.:(==)(A1::FieldMatrixWithSolver, A2::FieldMatrixWithSolver) =
 Base.similar(A::FieldMatrixWithSolver) =
     FieldMatrixWithSolver(similar(A.matrix), A.solver)
 
+Base.zero(A::FieldMatrixWithSolver) =
+    FieldMatrixWithSolver(zero(A.matrix), A.solver)
+
 ldiv!(x::Fields.FieldVector, A::FieldMatrixWithSolver, b::Fields.FieldVector) =
     field_matrix_solve!(A.solver, x, A.matrix, b)
 

src/MatrixFields/field_name_dict.jl

@@ -182,6 +182,13 @@ function Base.similar(dict::FieldNameDict)
     return FieldNameDict(keys(dict), entries)
 end
 
+function Base.zero(dict::FieldNameDict)
+    entries = unrolled_map(values(dict)) do entry
+        entry isa UniformScaling ? entry : zero(entry)
+    end
+    return FieldNameDict(keys(dict), entries)
+end
+
 # Note: This assumes that the matrix has the same row and column units, since I
 # cannot be multiplied by anything other than a scalar.
 function Base.one(matrix::FieldMatrix)

test/MatrixFields/field_matrix_solvers.jl

@@ -18,10 +18,12 @@ function test_field_matrix_solver(; test_name, alg, A, b, use_rel_error = false)
     @testset "$test_name" begin
         x = similar(b)
         A′ = FieldMatrixWithSolver(A, b, alg)
+        @test zero(A′) isa typeof(A′)
         solve_time =
             @benchmark ClimaComms.@cuda_sync comms_device ldiv!(x, A′, b)
 
         b_test = similar(b)
+        @test zero(b) isa typeof(b)
         mul_time =
             @benchmark ClimaComms.@cuda_sync comms_device mul!(b_test, A′, x)