Merge pull request #3732 from CliMA/dy/zero_velocity

Fix Jacobian for zero_velocity_tendency

Author: dennisYatunin

src/prognostic_equations/implicit/implicit_solver.jl

@@ -550,10 +550,10 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
     dtγ′ = FT(float(dtγ))
 
     A.dtγ_ref[] = dtγ′
-    update_implicit_equation_jacobian!(A, Y, p′, dtγ′)
+    update_implicit_equation_jacobian!(A, Y, p′, dtγ′, t)
 end
 
-function update_implicit_equation_jacobian!(A, Y, p, dtγ)
+function update_implicit_equation_jacobian!(A, Y, p, dtγ, t)
     dtγ = float(dtγ)
     (;
         matrix,
@@ -1148,4 +1148,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
                 ) - (I_u₃,)
         end
     end
+
+    # NOTE: All velocity tendency derivatives should be set BEFORE this call.
+    zero_velocity_jacobian!(matrix, Y, p, t)
 end

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -46,7 +46,7 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
     # NOTE: All ρa tendencies should be applied before calling this function
     pressure_work_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 
-    # NOTE: This will zero out all monmentum tendencies in the edmfx advection test
+    # NOTE: This will zero out all momentum tendencies in the edmfx advection test
     # please DO NOT add additional velocity tendencies after this function
     zero_velocity_tendency!(Yₜ, Y, p, t)
 

src/prognostic_equations/zero_velocity.jl

@@ -1,20 +1,49 @@
+import ClimaCore: MatrixFields
+import LinearAlgebra: UniformScaling
+
 ###
 ### edmfx advection test 
 ###
 
+# Turn off all momentum tendencies in the advection test.
 function zero_velocity_tendency!(Yₜ, Y, p, t)
-    # turn off all monmentum tendencies in the advection test
-    if p.atmos.advection_test
-        FT = eltype(Y)
-        n = n_mass_flux_subdomains(p.atmos.turbconv_model)
+    p.atmos.advection_test || return nothing
+    @. Yₜ.c.uₕ = C12(0, 0)
+    @. Yₜ.f.u₃ = C3(0)
+    if p.atmos.turbconv_model isa PrognosticEDMFX
+        for j in 1:n_mass_flux_subdomains(p.atmos.turbconv_model)
+            @. Yₜ.f.sgsʲs.:($$j).u₃ = C3(0)
+        end
+    end
+    return nothing
+end
 
-        @. Yₜ.c.uₕ = C12(FT(0), FT(0))
-        @. Yₜ.f.u₃ = Geometry.Covariant3Vector(FT(0))
+# Turn off all momentum tendency derivatives in the advection test.
+function zero_velocity_jacobian!(∂Yₜ_err_∂Y, Y, p, t)
+    p.atmos.advection_test || return nothing
+    for ((row_name, col_name), matrix_entry) in pairs(∂Yₜ_err_∂Y)
+        matrix_entry isa Fields.Field || continue
+        if row_name in (MatrixFields.@name(c.uₕ), MatrixFields.@name(f.u₃))
+            set_identity_matrix_entry!(matrix_entry, row_name, col_name)
+        end
         if p.atmos.turbconv_model isa PrognosticEDMFX
-            for j in 1:n
-                @. Yₜ.f.sgsʲs.:($$j).u₃ = Geometry.Covariant3Vector(FT(0))
+            for j in 1:n_mass_flux_subdomains(p.atmos.turbconv_model)
+                if row_name == MatrixFields.FieldName(:f, :sgsʲs, j, :u₃)
+                    set_identity_matrix_entry!(matrix_entry, row_name, col_name)
+                end
             end
         end
     end
-    return nothing
+end
+
+function set_identity_matrix_entry!(matrix_entry, row_name, col_name)
+    identity_matrix_entry_value = if row_name == col_name
+        # TODO: Add a method for one(::Axis2Tensor) to simplify this.
+        T = eltype(eltype(matrix_entry))
+        tensor_data = UniformScaling(one(eltype(T)))
+        -DiagonalMatrixRow(Geometry.AxisTensor(axes(T), tensor_data))
+    else
+        zero(eltype(matrix_entry))
+    end
+    matrix_entry .= (identity_matrix_entry_value,)
 end