Update examples to use new implicit solver interface

.buildkite/pipeline.yml

@@ -617,10 +617,6 @@ steps:
         key: unit_wfact
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/finitedifference/wfact.jl"
 
-      - label: "Unit: linsolve"
-        key: unit_linsolve
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/finitedifference/linsolve.jl"
-
       - label: "Unit: fd tensor"
         key: unit_fd_tensor
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/finitedifference/tensor.jl"

docs/src/matrix_fields.md

@@ -26,7 +26,16 @@ MultiplyColumnwiseBandMatrixField
 operator_matrix
 ```
 
-# Linear Solvers
+## Vectors and Matrices of Fields
+
+``@docs
+FieldNameDict
+identity_field_matrix
+field_vector_view
+concrete_field_vector
+``
+
+## Linear Solvers
 
 ```@docs
 FieldMatrixSolverAlgorithm
@@ -42,7 +51,7 @@ StationaryIterativeSolve
 ApproximateBlockArrowheadIterativeSolve
 ```
 
-# Preconditioners
+## Preconditioners
 
 ```@docs
 PreconditionerAlgorithm
@@ -67,9 +76,6 @@ FieldName
 @name
 FieldNameTree
 FieldNameSet
-FieldNameDict
-field_vector_view
-concrete_field_vector
 is_lazy
 lazy_main_diagonal
 lazy_mul

examples/hybrid/driver.jl

@@ -54,7 +54,6 @@ using JLD2
 
 const FT = get(ENV, "FLOAT_TYPE", "Float32") == "Float32" ? Float32 : Float64
 
-include("../implicit_solver_debugging_tools.jl")
 include("../ordinary_diff_eq_bug_fixes.jl")
 include("../common_spaces.jl")
 

examples/hybrid/implicit_equation_jacobian.jl

@@ -0,0 +1,113 @@
+import LinearAlgebra: ldiv!
+using ClimaCore: Spaces, Fields, Operators
+using ClimaCore.Utilities: half
+using ClimaCore.MatrixFields
+using ClimaCore.MatrixFields: @name
+
+struct ImplicitEquationJacobian{TJ, RJ, F, T1, T2}
+    # nonzero blocks of the implicit tendency's Jacobian
+    ∂Yₜ∂Y::TJ
+
+    # the full implicit residual's Jacobian, and its linear solver
+    ∂R∂Y::RJ
+
+    # whether this struct is used to compute Wfact_t or Wfact
+    transform::Bool
+
+    # flags for computing the Jacobian
+    flags::F
+
+    # cache that is used to evaluate ldiv!
+    temp1::T1
+    temp2::T2
+end
+
+function ImplicitEquationJacobian(Y, transform, flags)
+    FT = eltype(Y)
+
+    ᶜρ_name = @name(c.ρ)
+    ᶜ𝔼_name = if :ρθ in propertynames(Y.c)
+        @name(c.ρθ)
+    elseif :ρe in propertynames(Y.c)
+        @name(c.ρe)
+    elseif :ρe_int in propertynames(Y.c)
+        @name(c.ρe_int)
+    end
+    ᶠ𝕄_name = @name(f.w)
+
+    BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}}
+    BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
+    QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
+    TridiagonalRow_C3xACT3 =
+        TridiagonalMatrixRow{typeof(C3(FT(0)) * CT3(FT(0))')}
+    ∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3 =
+        flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact && :ρe in propertynames(Y.c) ?
+        QuaddiagonalRow_ACT3 : BidiagonalRow_ACT3
+    ∂Yₜ∂Y = FieldMatrix(
+        (ᶜρ_name, ᶠ𝕄_name) => Fields.Field(BidiagonalRow_ACT3, axes(Y.c)),
+        (ᶜ𝔼_name, ᶠ𝕄_name) => Fields.Field(∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3, axes(Y.c)),
+        (ᶠ𝕄_name, ᶜρ_name) => Fields.Field(BidiagonalRow_C3, axes(Y.f)),
+        (ᶠ𝕄_name, ᶜ𝔼_name) => Fields.Field(BidiagonalRow_C3, axes(Y.f)),
+        (ᶠ𝕄_name, ᶠ𝕄_name) =>
+            Fields.Field(TridiagonalRow_C3xACT3, axes(Y.f)),
+    )
+
+    dtγ = FT(1)
+    I = MatrixFields.identity_field_matrix(Y) # one(∂Yₜ∂Y) can't get every block
+    ∂R∂Y = transform ? I ./ dtγ .- ∂Yₜ∂Y : dtγ .* ∂Yₜ∂Y .- I
+    alg = MatrixFields.BlockArrowheadSolve(ᶜρ_name, ᶜ𝔼_name)
+
+    return ImplicitEquationJacobian(
+        ∂Yₜ∂Y,
+        FieldMatrixWithSolver(∂R∂Y, Y, alg),
+        transform,
+        flags,
+        similar(Y),
+        similar(Y),
+    )
+end
+
+# Required for compatibility with OrdinaryDiffEq.jl
+Base.similar(A::ImplicitEquationJacobian) = ImplicitEquationJacobian(
+    similar(A.∂Yₜ∂Y),
+    similar(A.∂R∂Y),
+    A.transform,
+    A.flags,
+    A.temp1,
+    A.temp2,
+)
+
+# Required for compatibility with ClimaTimeSteppers.jl
+Base.zero(A::ImplicitEquationJacobian) = ImplicitEquationJacobian(
+    zero(A.∂Yₜ∂Y),
+    zero(A.∂R∂Y),
+    A.transform,
+    A.flags,
+    A.temp1,
+    A.temp2,
+)
+
+# This method for ldiv! is called by Newton's method from ClimaTimeSteppers.jl.
+# It solves ∂R∂Y * x = b for x, where R is the implicit residual.
+ldiv!(
+    x::Fields.FieldVector,
+    A::ImplicitEquationJacobian,
+    b::Fields.FieldVector,
+) = ldiv!(x, A.∂R∂Y, b)
+
+# This method for ldiv! is called by Krylov.jl from ClimaTimeSteppers.jl.
+# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a
+# a similar way of handling the AbstractVectors generated by Krylov.jl.
+function ldiv!(
+    x::AbstractVector,
+    A::ImplicitEquationJacobian,
+    b::AbstractVector,
+)
+    A.temp_b .= b
+    ldiv!(A.temp_x, A, A.temp_b)
+    x .= A.temp_x
+end
+
+# This function is called by Newton's method from OrdinaryDiffEq.jl.
+linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve!
+_linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b)

examples/hybrid/ode_config.jl

@@ -14,7 +14,7 @@ use_transform(ode_algo) = !is_imex_CTS_algo(ode_algo)
 
 function jac_kwargs(ode_algo, Y, jacobi_flags)
     if is_imex_CTS_algo(ode_algo)
-        W = SchurComplementW(Y, use_transform(ode_algo), jacobi_flags)
+        W = ImplicitEquationJacobian(Y, use_transform(ode_algo), jacobi_flags)
         if use_transform(ode_algo)
             return (; jac_prototype = W, Wfact_t = Wfact!)
         else