|
| 1 | +import LinearAlgebra: ldiv! |
| 2 | +using ClimaCore: Spaces, Fields, Operators |
| 3 | +using ClimaCore.Utilities: half |
| 4 | +using ClimaCore.MatrixFields |
| 5 | +using ClimaCore.MatrixFields: @name |
| 6 | + |
| 7 | +""" |
| 8 | + ImplicitEquationJacobian(Y, transform, [flags]) |
| 9 | +
|
| 10 | +Represents the Jacobian of the implicit residual, `R(Y)`, which is defined as |
| 11 | +either `Yᵖʳᵉᵛ + δt * γ * Yₜ - Y` or `(Y - Yᵖʳᵉᵛ) / (δt * γ) - Yₜ`, with the |
| 12 | +second version used when `transform` is `true`. In this expression, `δt` is the |
| 13 | +timestep, `γ` is a scalar quantity determined by the timestepping scheme, `Y` is |
| 14 | +the current approximation of the timestepper's implicit stage value, `Yₜ` is the |
| 15 | +value obtained by calling `implicit_tendency!(Yₜ, Y, p, t)`, and `Yᵖʳᵉᵛ` is a |
| 16 | +linear combination of the timestepper's previous stage values. The residual's |
| 17 | +Jacobian, `∂R/∂Y`, can be expressed in terms of the tendency's Jacobian, |
| 18 | +`∂Yₜ/∂Y`, as either `δt * γ * ∂Yₜ/∂Y - I` or `I / (δt * γ) - ∂Yₜ/∂Y`. |
| 19 | +
|
| 20 | +The `ImplicitEquationJacobian` allows the `staggered_nonhydrostatic_model.jl` |
| 21 | +file to be compatible with `ClimaTimeSteppers` and `OrdinaryDiffEq`. This file |
| 22 | +defines both `implicit_tendency!` and `implicit_equation_jacobian!`, where the |
| 23 | +latter sets `∂R/∂Y` based on the values of `Y` and `δt * γ`. An optional set of |
| 24 | +flags can also be passed to the `ImplicitEquationJacobian` constructor, which |
| 25 | +are then accessible from within `implicit_equation_jacobian!`. For all |
| 26 | +timestepping schemes from `ClimaTimeSteppers` and the Rosenbrock schemes from |
| 27 | +`OrdinaryDiffEq`, the `transform` flag should be set to `false`. For all other |
| 28 | +timestepping schemes from `OrdinaryDiffEq`, it should be set to `true`. |
| 29 | +
|
| 30 | +Within both `ClimaTimeSteppers` and `OrdinaryDiffEq`, this data structure is |
| 31 | +used to solve the linear equation `∂R/∂Y * δY = R`, which allows Newton's method |
| 32 | +to iteratively find the root of `R(Y) = 0`. In `ClimaTimeSteppers`, this is done |
| 33 | +by calling `ldiv!(δY, ∂R/∂Y, R)`, where `δY` and `R` are represeted as |
| 34 | +`FieldVector`s and `∂R/∂Y` is represented as an `ImplicitEquationJacobian`. When |
| 35 | +using Krylov methods from `ClimaTimeSteppers`, `δY` and `R` can also be |
| 36 | +represented as other `AbstractVector`s (which are internally converted to |
| 37 | +`FieldVector`s). For `OrdinaryDiffEq`, we use the `linsolve!` function instead |
| 38 | +of `ldiv!` by passing it to the timestepping scheme's constructor. |
| 39 | +
|
| 40 | +TODO: Compatibility with `OrdinaryDiffEq` is out of date and should be updated. |
| 41 | +""" |
| 42 | +struct ImplicitEquationJacobian{TJ, RJ, F, V} |
| 43 | + ∂Yₜ∂Y::TJ # nonzero blocks of the implicit tendency's Jacobian |
| 44 | + ∂R∂Y::RJ # the full implicit residual's Jacobian, and its linear solver |
| 45 | + transform::Bool # whether this struct is used to compute Wfact or Wfact_t |
| 46 | + flags::F # flags for computing the implicit tendency's Jacobian |
| 47 | + R_field_vector::V # cache that is used to evaluate ldiv! for Krylov methods |
| 48 | + δY_field_vector::V # cache that is used to evaluate ldiv! for Krylov methods |
| 49 | +end |
| 50 | + |
| 51 | +function ImplicitEquationJacobian(Y, transform, flags = (;)) |
| 52 | + FT = eltype(Y) |
| 53 | + |
| 54 | + ᶜρ_name = @name(c.ρ) |
| 55 | + ᶜ𝔼_name = if :ρθ in propertynames(Y.c) |
| 56 | + @name(c.ρθ) |
| 57 | + elseif :ρe in propertynames(Y.c) |
| 58 | + @name(c.ρe) |
| 59 | + elseif :ρe_int in propertynames(Y.c) |
| 60 | + @name(c.ρe_int) |
| 61 | + end |
| 62 | + ᶠ𝕄_name = @name(f.w) |
| 63 | + |
| 64 | + BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} |
| 65 | + BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} |
| 66 | + QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}} |
| 67 | + TridiagonalRow_C3xACT3 = |
| 68 | + TridiagonalMatrixRow{typeof(C3(FT(0)) * CT3(FT(0))')} |
| 69 | + ∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3 = |
| 70 | + flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact && :ρe in propertynames(Y.c) ? |
| 71 | + QuaddiagonalRow_ACT3 : BidiagonalRow_ACT3 |
| 72 | + ∂Yₜ∂Y = FieldMatrix( |
| 73 | + (ᶜρ_name, ᶠ𝕄_name) => Fields.Field(BidiagonalRow_ACT3, axes(Y.c)), |
| 74 | + (ᶜ𝔼_name, ᶠ𝕄_name) => Fields.Field(∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3, axes(Y.c)), |
| 75 | + (ᶠ𝕄_name, ᶜρ_name) => Fields.Field(BidiagonalRow_C3, axes(Y.f)), |
| 76 | + (ᶠ𝕄_name, ᶜ𝔼_name) => Fields.Field(BidiagonalRow_C3, axes(Y.f)), |
| 77 | + (ᶠ𝕄_name, ᶠ𝕄_name) => |
| 78 | + Fields.Field(TridiagonalRow_C3xACT3, axes(Y.f)), |
| 79 | + ) |
| 80 | + |
| 81 | + # When ∂Yₜ∂Y is sparse, one(∂Yₜ∂Y) doesn't contain every diagonal block. |
| 82 | + # To ensure that ∂R∂Y is invertible, we need to call identity_field_matrix. |
| 83 | + I = MatrixFields.identity_field_matrix(Y) |
| 84 | + |
| 85 | + δtγ = FT(1) |
| 86 | + ∂R∂Y = transform ? I ./ δtγ .- ∂Yₜ∂Y : δtγ .* ∂Yₜ∂Y .- I |
| 87 | + alg = MatrixFields.BlockArrowheadSolve(ᶜρ_name, ᶜ𝔼_name) |
| 88 | + |
| 89 | + return ImplicitEquationJacobian( |
| 90 | + ∂Yₜ∂Y, |
| 91 | + FieldMatrixWithSolver(∂R∂Y, Y, alg), |
| 92 | + transform, |
| 93 | + flags, |
| 94 | + similar(Y), |
| 95 | + similar(Y), |
| 96 | + ) |
| 97 | +end |
| 98 | + |
| 99 | +# Required for compatibility with OrdinaryDiffEq.jl |
| 100 | +Base.similar(j::ImplicitEquationJacobian) = ImplicitEquationJacobian( |
| 101 | + similar(j.∂Yₜ∂Y), |
| 102 | + similar(j.∂R∂Y), |
| 103 | + j.transform, |
| 104 | + j.flags, |
| 105 | + j.R_field_vector, |
| 106 | + j.δY_field_vector, |
| 107 | +) |
| 108 | + |
| 109 | +# Required for compatibility with ClimaTimeSteppers.jl |
| 110 | +Base.zero(j::ImplicitEquationJacobian) = ImplicitEquationJacobian( |
| 111 | + zero(j.∂Yₜ∂Y), |
| 112 | + zero(j.∂R∂Y), |
| 113 | + j.transform, |
| 114 | + j.flags, |
| 115 | + j.R_field_vector, |
| 116 | + j.δY_field_vector, |
| 117 | +) |
| 118 | + |
| 119 | +# This method for ldiv! is called by Newton's method from ClimaTimeSteppers.jl. |
| 120 | +# It solves ∂R∂Y * δY = R for δY, where R is the implicit residual. |
| 121 | +ldiv!( |
| 122 | + δY::Fields.FieldVector, |
| 123 | + j::ImplicitEquationJacobian, |
| 124 | + R::Fields.FieldVector, |
| 125 | +) = ldiv!(δY, j.∂R∂Y, R) |
| 126 | + |
| 127 | +# This method for ldiv! is called by Krylov.jl from ClimaTimeSteppers.jl. |
| 128 | +# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a |
| 129 | +# a similar way of handling the AbstractVectors generated by Krylov.jl. |
| 130 | +function ldiv!( |
| 131 | + δY::AbstractVector, |
| 132 | + j::ImplicitEquationJacobian, |
| 133 | + R::AbstractVector, |
| 134 | +) |
| 135 | + j.R_field_vector .= R |
| 136 | + ldiv!(j.δY_field_vector, j, j.R_field_vector) |
| 137 | + δY .= j.δY_field_vector |
| 138 | +end |
| 139 | + |
| 140 | +# This function can be called by Newton's method from OrdinaryDiffEq.jl. |
| 141 | +linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve! |
| 142 | +_linsolve!(δY, j, R, update_matrix = false; kwargs...) = ldiv!(δY, j, R) |
0 commit comments