Fix energy conservation with sparse autodiff

Author: dennisYatunin

.buildkite/pipeline.yml

@@ -286,6 +286,12 @@ steps:
           --job_id baroclinic_wave_conservation
         artifact_paths: "baroclinic_wave_conservation/output_active/*"
 
+      - label: ":computer: baroclinic wave check conservation float64 sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/baroclinic_wave_conservation_ft64_sparse_autodiff.yml
+          --job_id baroclinic_wave_conservation_ft64_sparse_autodiff
+        artifact_paths: "baroclinic_wave_conservation_ft64_sparse_autodiff/output_active/*"
+
       - label: ":computer: baroclinic wave moist check conservation"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/baroclinic_wave_equil_conservation.yml
@@ -303,7 +309,6 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/baroclinic_wave_equil_conservation_ft64_sparse_autodiff.yml
           --job_id baroclinic_wave_equil_conservation_ft64_sparse_autodiff
         artifact_paths: "baroclinic_wave_equil_conservation_ft64_sparse_autodiff/output_active/*"
-        soft_fail: true
 
       - label: ":computer: baroclinic wave moist check conservation with sources"
         command: >

config/model_configs/baroclinic_wave_conservation_ft64_sparse_autodiff.yml

@@ -0,0 +1,14 @@
+FLOAT_TYPE: "Float64"
+initial_condition: "DryBaroclinicWave"
+t_end: "5days"
+dt_save_state_to_disk: "5days"
+disable_surface_flux_tendency: true
+use_auto_jacobian: true
+update_jacobian_every: dt
+debug_jacobian: true
+check_conservation: true
+diagnostics:
+  - short_name: [massa, energya]
+    period: 1days
+    writer: dict
+use_itime: true

post_processing/jacobian_summary.jl

@@ -77,22 +77,24 @@ function print_jacobian_summary(integrator)
     dense_bandwidth_values = map(block_keys) do block_key
         bandwidth(dense_blocks[block_key])
     end
-    sparse_bandwidth_values = map(block_keys) do block_key
-        sparse_blocks = first(all_sparse_blocks)
-        haskey(sparse_blocks, block_key) ?
-        (
-            sparse_blocks[block_key] isa UniformScaling ? 1 :
-            bandwidth(sparse_blocks[block_key])
-        ) : 0
-    end
-    missing_bandwidth_values =
-        max.(dense_bandwidth_values .- sparse_bandwidth_values, 0)
-    @info "dense, number of nonzero bands per block:"
+    @info "dense, nonzero bands per block:"
     pretty_table(dense_bandwidth_values; bandwidth_table_kwargs...)
-    @info "sparse, number of nonzero bands per block:"
-    pretty_table(sparse_bandwidth_values; bandwidth_table_kwargs...)
-    @info "dense - sparse, number of missing nonzero bands per block:"
-    pretty_table(missing_bandwidth_values; bandwidth_table_kwargs...)
+    for (sparse_name, sparse_blocks) in pairs(all_sparse_blocks)
+        sparse_bandwidth_values = map(block_keys) do block_key
+            haskey(sparse_blocks, block_key) ?
+            (
+                sparse_blocks[block_key] isa UniformScaling ?
+                1 - iszero(sparse_blocks[block_key]) :
+                bandwidth(sparse_blocks[block_key])
+            ) : 0
+        end
+        missing_bandwidth_values =
+            max.(dense_bandwidth_values .- sparse_bandwidth_values, 0)
+        @info "$sparse_name sparse, nonzero bands per block:"
+        pretty_table(sparse_bandwidth_values; bandwidth_table_kwargs...)
+        @info "dense - $sparse_name sparse, missing nonzero bands per block:"
+        pretty_table(missing_bandwidth_values; bandwidth_table_kwargs...)
+    end
     println("<$('='^70)>\n")
 
     rms(block) = sqrt(mean(abs2.(block)))

src/prognostic_equations/implicit/auto_dense_jacobian.jl

@@ -48,7 +48,7 @@ AutoDenseJacobian(max_simultaneous_derivatives = 32) =
 # The number of derivatives computed simultaneously by AutoDenseJacobian.
 max_simultaneous_derivatives(::AutoDenseJacobian{S}) where {S} = S
 
-function jacobian_cache(alg::AutoDenseJacobian, Y, atmos)
+function jacobian_cache(alg::AutoDenseJacobian, Y, atmos; kwargs...)
     FT = eltype(Y)
     DA = ClimaComms.array_type(Y)
 

src/prognostic_equations/implicit/auto_sparse_jacobian.jl

@@ -23,7 +23,13 @@ end
 AutoSparseJacobian(sparse_jacobian_alg) =
     AutoSparseJacobian(sparse_jacobian_alg, nothing)
 
-function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
+function jacobian_cache(
+    alg::AutoSparseJacobian,
+    Y,
+    atmos;
+    verbose = true,
+    kwargs...,
+)
     (; sparse_jacobian_alg, padding_bands_per_block) = alg
 
     FT = eltype(Y)
@@ -38,10 +44,13 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
     scratch = implicit_temporary_quantities(Y, atmos)
 
     # Allocate ∂R/∂Y and its corresponding linear solver.
+    sparse_jacobian_cache =
+        jacobian_cache(sparse_jacobian_alg, Y, atmos; verbose, auto_pad = true)
+
     # TODO: Add FieldNameTree(Y) to the matrix in FieldMatrixWithSolver. The
     # tree is needed to evaluate scalar_tendency_matrix[autodiff_matrix_keys].
-    # (; matrix) = jacobian_cache(sparse_jacobian_alg, Y, atmos)
-    matrix_without_tree = jacobian_cache(sparse_jacobian_alg, Y, atmos).matrix
+    # (; matrix) = sparse_jacobian_cache
+    matrix_without_tree = sparse_jacobian_cache.matrix
     tree = MatrixFields.FieldNameTree(Y)
     matrix = MatrixFields.FieldMatrixWithSolver(
         MatrixFields.replace_name_tree(matrix_without_tree.matrix, tree),
@@ -137,6 +146,17 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
             (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno))
         max_padding_bands = if !isnothing(padding_bands_per_block)
             padding_bands_per_block
+        elseif (
+            block_row_name == @name(c.ρe_tot) &&
+            block_column_name in
+            (@name(c.ρ), @name(c.ρe_tot), uₕ_component_names...) &&
+            !(block_key in keys(autodiff_matrix))
+        )
+            # ‖∂ᶜρe_totₜ/∂ᶜχ‖ ≳ ‖∂ᶜρe_totₜ/∂ᶠu₃‖ / 10^12 for any center field ᶜχ.
+            # The ∂ᶜρe_totₜ/∂ᶠu₃ block is critical for conservation of energy,
+            # and even relative errors as small as one part in a trillion can
+            # noticeably degrade conservation with Float64 precision.
+            3
         elseif (
             (
                 block_row_name in uₕ_component_names &&
@@ -169,7 +189,7 @@ function jacobian_cache(alg::AutoSparseJacobian, Y, atmos; verbose = true)
             #   relatively smaller than ‖δᶜρaʲ‖.
             # - ‖∂ᶠu₃ʲₜ/∂ᶜu₁‖ and ‖∂ᶠu₃ʲₜ/∂ᶜu₁‖ ≳ ‖∂ᶠu₃ʲₜ/∂ᶠu₃ʲ‖, as long as
             #   ‖δᶜu₁‖ and ‖δᶜu₂‖ are relatively smaller than ‖δᶠu₃ʲ‖.
-            # Diagonal blocks are critical for conservation and stability, so
+            # Diagonal blocks are important for conservation and stability, so
             # these potential errors from off-diagonal blocks should be avoided.
             3
         elseif (

src/prognostic_equations/implicit/jacobian.jl

@@ -4,7 +4,7 @@
 A description of how to compute the matrix ``∂R/∂Y``, where ``R(Y)`` denotes the
 residual of an implicit step with the state ``Y``. Concrete implementations of
 this abstract type should define 3 methods:
- - `jacobian_cache(alg::JacobianAlgorithm, Y, atmos; [verbose])`
+ - `jacobian_cache(alg::JacobianAlgorithm, Y, atmos; kwargs...)`
  - `update_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)`
  - `invert_jacobian!(alg::JacobianAlgorithm, cache, ΔY, R)`
 To facilitate debugging, concrete implementations should also define
@@ -17,25 +17,22 @@ abstract type JacobianAlgorithm end
 abstract type SparseJacobian <: JacobianAlgorithm end
 
 """
-    Jacobian(alg, Y, atmos; [verbose])
+    Jacobian(alg, Y, atmos; kwargs...)
 
 Wrapper for a [`JacobianAlgorithm`](@ref) and its cache, which it uses to update
-and invert the Jacobian. The optional `verbose` flag specifies whether debugging
-information should be printed during initialization.
+and invert the Jacobian. Optional keyword arguments can also be passed to
+some algorithms during construction.
 """
 struct Jacobian{A <: JacobianAlgorithm, C}
     alg::A
     cache::C
 end
-function Jacobian(alg, Y, atmos; verbose = false)
+function Jacobian(alg, Y, atmos; kwargs...)
     krylov_cache = (; ΔY_krylov = similar(Y), R_krylov = similar(Y))
-    cache = (; jacobian_cache(alg, Y, atmos; verbose)..., krylov_cache...)
+    cache = (; jacobian_cache(alg, Y, atmos; kwargs...)..., krylov_cache...)
     return Jacobian(alg, cache)
 end
 
-# Ignore the verbose flag in jacobian_cache when it is not needed.
-jacobian_cache(alg, Y, atmos; verbose) = jacobian_cache(alg, Y, atmos)
-
 # ClimaTimeSteppers.jl calls zero(jac_prototype) to initialize the Jacobian, but
 # we don't need to allocate a second Jacobian for this (in particular, the exact
 # Jacobian can be very expensive to allocate).

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -58,7 +58,13 @@ struct ManualSparseJacobian{F1, F2, F3, F4, F5, F6} <: SparseJacobian
     approximate_solve_iters::Int
 end
 
-function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
+function jacobian_cache(
+    alg::ManualSparseJacobian,
+    Y,
+    atmos;
+    auto_pad = false,
+    kwargs...,
+)
     (;
         topography_flag,
         diffusion_flag,
@@ -74,6 +80,7 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
     BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}}
     TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}}
     BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
+    QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
     BidiagonalRow_C3xACTh =
         BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')}
     DiagonalRow_C3xACT3 =
@@ -128,7 +135,8 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
         (@name(c.ρ), sfc_if_available...),
     )
 
-    active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...)
+    active_scalar_names_except_ρe = (@name(c.ρ), ρq_tot_if_available...)
+    active_scalar_names = (@name(c.ρe_tot), active_scalar_names_except_ρe...)
     advection_blocks = (
         (
             use_derivative(topography_flag) ?
@@ -141,8 +149,12 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
         )...,
         MatrixFields.unrolled_map(
             name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
-            active_scalar_names,
+            active_scalar_names_except_ρe,
         )...,
+        (@name(c.ρe_tot), @name(f.u₃)) => similar(
+            Y.c,
+            auto_pad ? QuaddiagonalRow_ACT3 : BidiagonalRow_ACT3,
+        ),
         MatrixFields.unrolled_map(
             name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),
             active_scalar_names,