Add missing blocks to the dycore Jacobian

src/cache/cache.jl

@@ -144,6 +144,7 @@ function build_cache(Y, atmos, params, surface_setup, dt, start_date)
         ᶜp_ref,
         ᶜT = similar(Y.c, FT),
         ᶜf,
+        ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CT12{FT}}}),
         ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
         # Used by diagnostics such as hfres, evspblw
         surface_ct3_unit = CT3.(

src/prognostic_equations/implicit/implicit_solver.jl

@@ -125,12 +125,15 @@ function ImplicitEquationJacobian(
     approximate_solve_iters,
 )
     FT = Spaces.undertype(axes(Y.c))
-    one_C3xACT3 = C3(FT(1)) * CT3(FT(1))'
     TridiagonalRow = TridiagonalMatrixRow{FT}
     BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}}
+    TridiagonalRow_ACT12 = TridiagonalMatrixRow{Adjoint{FT, CT12{FT}}}
     BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
     QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
-    TridiagonalRow_C3xACT3 = TridiagonalMatrixRow{typeof(one_C3xACT3)}
+    BidiagonalRow_C3xACT12 =
+        BidiagonalMatrixRow{typeof(C3(FT(1)) * CT12(FT(1), FT(1))')}
+    TridiagonalRow_C3xACT3 =
+        TridiagonalMatrixRow{typeof(C3(FT(1)) * CT3(FT(1))')}
 
     is_in_Y(name) = MatrixFields.has_field(Y, name)
 
@@ -172,26 +175,40 @@ function ImplicitEquationJacobian(
             name -> (name, name) => FT(-1) * I,
             other_available_names,
         )...,
-        (@name(c.ρ), @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
+        MatrixFields.unrolled_map(
+            name ->
+                (name, @name(c.uₕ)) => similar(Y.c, TridiagonalRow_ACT12),
+            (@name(c.ρ), @name(c.ρe_tot), available_tracer_names...),
+        )...,
+        MatrixFields.unrolled_map(
+            name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
+            (@name(c.ρ), available_tracer_names...),
+        )...,
         (@name(c.ρe_tot), @name(f.u₃)) =>
             use_derivative(enthalpy_flag) ?
             similar(Y.c, QuaddiagonalRow_ACT3) :
             similar(Y.c, BidiagonalRow_ACT3),
-        MatrixFields.unrolled_map(
-            name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
-            available_tracer_names,
-        )...,
         (@name(f.u₃), @name(c.ρ)) => similar(Y.f, BidiagonalRow_C3),
         (@name(f.u₃), @name(c.ρe_tot)) => similar(Y.f, BidiagonalRow_C3),
+        (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACT12),
         (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
     )
 
+    scalar_names = (
+        @name(c.ρ),
+        @name(c.ρe_tot),
+        available_tracer_names...,
+        other_available_names...,
+    ) # Everything in Y except for ᶜuₕ and ᶠu₃.
+
+    alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ))
     alg =
         use_derivative(diffusion_flag) ?
         MatrixFields.ApproximateBlockArrowheadIterativeSolve(
-            @name(c);
+            scalar_names...;
+            alg₂,
             n_iters = approximate_solve_iters,
-        ) : MatrixFields.BlockArrowheadSolve(@name(c))
+        ) : MatrixFields.BlockArrowheadSolve(scalar_names...; alg₂)
 
     # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration
     # and approximates the A[c, c] block with a MainDiagonalPreconditioner.
@@ -266,6 +283,7 @@ function Wfact!(A, Y, p, dtγ, t)
                 p.atmos.turbconv_model isa PrognosticEDMFX ?
                 (; p.precomputed.ᶜρa⁰) : (;)
             )...,
+            p.core.∂ᶜK_∂ᶜuₕ,
             p.core.∂ᶜK_∂ᶠu₃,
             p.core.ᶜΦ,
             p.core.ᶠgradᵥ_ᶜΦ,
@@ -294,15 +312,14 @@ end
 
 function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     (; matrix, diffusion_flag, enthalpy_flag) = A
-    (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶠu₃) = p
+    (; ᶜspecific, ᶠu³, ᶜK, ᶜp, ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃) = p
     (; ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref, params) = p
     (; energy_upwinding, density_upwinding) = p.atmos.numerics
     (; tracer_upwinding, precip_upwinding) = p.atmos.numerics
 
     FT = Spaces.undertype(axes(Y.c))
-    one_ATC3 = CT3(FT(1))'
+    one_AC3 = C3(FT(1))'
     one_C3xACT3 = C3(FT(1)) * CT3(FT(1))'
-    one_CT3xACT3 = CT3(FT(1)) * CT3(FT(1))'
     R_d = FT(CAP.R_d(params))
     cv_d = FT(CAP.cv_d(params))
     T_tri = FT(CAP.T_triple(params))
@@ -311,7 +328,8 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     ᶜuₕ = Y.c.uₕ
     ᶠu₃ = Y.f.u₃
     ᶜJ = Fields.local_geometry_field(Y.c).J
-    ᶠg³³ = g³³_field(Y.f)
+    ᶜg³ʰ = g³ʰ_tensor_field(Y.c)
+    ᶠg³³ = g³³_tensor_field(Y.f)
 
     # We can express the kinetic energy as
     # ᶜK =
@@ -322,6 +340,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     # ∂(ᶜK)/∂(ᶠu₃) =
     #     ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) +
     #     DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix().
+    @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow(
+        adjoint(CT12(ᶜuₕ[colidx])) +
+        adjoint(ᶜinterp(ᶠu₃[colidx])) * ᶜg³ʰ[colidx],
+    )
     @. ∂ᶜK_∂ᶠu₃[colidx] =
         ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[colidx]))) +
         DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix()
@@ -378,6 +400,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     )
     MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, χ_name, upwinding)
         MatrixFields.has_field(Y, ρχ_name) || return
+        ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)]
         ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)]
 
         ᶜχ = if χ_name == @name(ᶜ1)
@@ -389,24 +412,33 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
 
         if upwinding == Val(:none)
             if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot)
+                @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
+                    dtγ * -(ᶜadvdivᵥ_matrix()) ⋅
+                    DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ (
+                        DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅
+                        ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
+                        DiagonalMatrixRow(ᶜg³ʰ[colidx]) -
+                        (R_d / cv_d) * DiagonalMatrixRow(ᶠu³[colidx]) ⋅
+                        ᶠinterp_matrix() ⋅ ∂ᶜK_∂ᶜuₕ[colidx]
+                    )
                 @. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
                     dtγ * -(ᶜadvdivᵥ_matrix()) ⋅
                     DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ (
-                        DiagonalMatrixRow(
-                            ᶠg³³[colidx] *
-                            (one_CT3xACT3,) *
-                            ᶠinterp(ᶜχ[colidx]),
-                        ) +
-                        DiagonalMatrixRow(ᶠu³[colidx]) ⋅ ᶠinterp_matrix() ⋅
-                        ∂ᶜK_∂ᶠu₃[colidx] * (-R_d / cv_d)
+                        DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * ᶠg³³[colidx]) -
+                        (R_d / cv_d) * DiagonalMatrixRow(ᶠu³[colidx]) ⋅
+                        ᶠinterp_matrix() ⋅ ∂ᶜK_∂ᶠu₃[colidx]
                     )
             else
+                @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
+                    dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                        ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) * ᶠinterp(ᶜχ[colidx]),
+                    ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
+                    DiagonalMatrixRow(ᶜg³ʰ[colidx])
                 @. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
                     dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
                         ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) *
-                        ᶠg³³[colidx] *
-                        (one_CT3xACT3,) *
-                        ᶠinterp(ᶜχ[colidx]),
+                        ᶠinterp(ᶜχ[colidx]) *
+                        ᶠg³³[colidx],
                     )
             end
         else
@@ -426,6 +458,21 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
             ) # Need to wrap ᶠupwind_matrix in this for the same reason.
 
             if use_derivative(enthalpy_flag) && ρχ_name == @name(c.ρe_tot)
+                @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
+                    dtγ * -(ᶜadvdivᵥ_matrix()) ⋅
+                    DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ (
+                        DiagonalMatrixRow(
+                            u³_sign(ᶠu³[colidx]) *
+                            ᶠset_upwind_bcs(
+                                ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]),
+                            ) *
+                            (one_AC3,),
+                        ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
+                        DiagonalMatrixRow(ᶜg³ʰ[colidx]) -
+                        (R_d / cv_d) *
+                        ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³[colidx])) ⋅
+                        ∂ᶜK_∂ᶜuₕ[colidx]
+                    )
                 @. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
                     dtγ * -(ᶜadvdivᵥ_matrix()) ⋅
                     DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) ⋅ (
@@ -434,22 +481,33 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
                             ᶠset_upwind_bcs(
                                 ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]),
                             ) *
-                            ᶠg³³[colidx] *
-                            (one_ATC3,),
-                        ) +
+                            (one_AC3,) *
+                            ᶠg³³[colidx],
+                        ) -
+                        (R_d / cv_d) *
                         ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³[colidx])) ⋅
-                        ∂ᶜK_∂ᶠu₃[colidx] * (-R_d / cv_d)
+                        ∂ᶜK_∂ᶠu₃[colidx]
                     )
             else
+                @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
+                    dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                        ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) *
+                        u³_sign(ᶠu³[colidx]) *
+                        ᶠset_upwind_bcs(
+                            ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]),
+                        ) *
+                        (one_AC3,),
+                    ) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
+                    DiagonalMatrixRow(ᶜg³ʰ[colidx])
                 @. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
                     dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
                         ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]) *
                         u³_sign(ᶠu³[colidx]) *
                         ᶠset_upwind_bcs(
                             ᶠupwind(CT3(u³_sign(ᶠu³[colidx])), ᶜχ[colidx]),
                         ) *
-                        ᶠg³³[colidx] *
-                        (one_ATC3,),
+                        (one_AC3,) *
+                        ᶠg³³[colidx],
                     )
             end
         end
@@ -500,9 +558,10 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     # with I_u₃ = DiagonalMatrixRow(one_C3xACT3). We cannot use I * one_C3xACT3
     # because UniformScaling only supports Numbers as scale factors.
     ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)]
+    ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)]
+    ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)]
     ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)]
     I_u₃ = DiagonalMatrixRow(one_C3xACT3)
-    ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)]
     @. ∂ᶠu₃_err_∂ᶜρ[colidx] =
         dtγ * (
             DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅
@@ -519,6 +578,9 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] =
         dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() *
         R_d / cv_d
+    @. ∂ᶠu₃_err_∂ᶜuₕ[colidx] =
+        dtγ * DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() ⋅
+        DiagonalMatrixRow(-R_d * ᶜρ[colidx] / cv_d) ⋅ ∂ᶜK_∂ᶜuₕ[colidx]
     if p.atmos.rayleigh_sponge isa RayleighSponge
         @. ∂ᶠu₃_err_∂ᶠu₃[colidx] =
             dtγ * (

src/utils/abbreviations.jl

@@ -74,6 +74,7 @@ const ᶜinterp_matrix = MatrixFields.operator_matrix(ᶜinterp)
 const ᶜdivᵥ_matrix = MatrixFields.operator_matrix(ᶜdivᵥ)
 const ᶜadvdivᵥ_matrix = MatrixFields.operator_matrix(ᶜadvdivᵥ)
 const ᶠinterp_matrix = MatrixFields.operator_matrix(ᶠinterp)
+const ᶠwinterp_matrix = MatrixFields.operator_matrix(ᶠwinterp)
 const ᶠgradᵥ_matrix = MatrixFields.operator_matrix(ᶠgradᵥ)
 const ᶠupwind1_matrix = MatrixFields.operator_matrix(ᶠupwind1)
 const ᶠupwind3_matrix = MatrixFields.operator_matrix(ᶠupwind3)

src/utils/utilities.jl

@@ -126,6 +126,44 @@ function g³³_field(field)
     return g_field.:($end_index) # For both 2D and 3D spaces, g³³ = g[end].
 end
 
+"""
+    g³³_tensor_field(field)
+
+Extracts the `g³³` tensor field from `Fields.local_geometry_field(field)`.
+"""
+g³³_tensor_field(field) = Base.Broadcast.broadcasted(
+    Geometry.AxisTensor,
+    ((Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()),),
+    g³³_field(field),
+)
+
+"""
+    g³ʰ_tensor_field(field)
+
+Extracts the `g³ʰ` tensor field from `Fields.local_geometry_field(field)`.
+"""
+function g³ʰ_tensor_field(field)
+    full_axis = axes(eltype(Fields.local_geometry_field(field).gⁱʲ))[1]
+    all_components = Fields.local_geometry_field(field).gⁱʲ.components
+    if full_axis == Geometry.Contravariant123Axis()
+        horizontal_axis = Geometry.Contravariant12Axis()
+        n_components = 2
+    elseif full_axis == Geometry.Contravariant13Axis()
+        horizontal_axis = Geometry.Contravariant1Axis()
+        n_components = 1
+    elseif full_axis == Geometry.Contravariant23Axis()
+        horizontal_axis = Geometry.Contravariant2Axis()
+        n_components = 1
+    else
+        error("Local geometry does not specify a g³ʰ tensor field")
+    end
+    g³ʰ_axes = ((Geometry.Contravariant3Axis(), horizontal_axis),)
+    g³ʰ_field = Base.Broadcast.broadcasted(all_components) do g³ʰ
+        g³ʰ[end:end, 1:n_components]
+    end
+    return Base.Broadcast.broadcasted(Geometry.AxisTensor, g³ʰ_axes, g³ʰ_field)
+end
+
 """
     unit_basis_vector_data(type, local_geometry)