diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d14e236356..b4af9f507e 100644 --- a/.buildkite/pipeline.yml +++ b/.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: > diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index d49e00e50d..90fcc0bb6f 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -91,7 +91,7 @@ use_dense_jacobian: value: false use_auto_jacobian: help: "Whether to populate the entries of the sparse Jacobian matrix using forward-mode automatic differentiation with sparse matrix coloring (only used when `use_dense_jacobian` is `false`) [`true`, `false` (default)]" - value: false + value: true # TODO: Change this to false. auto_jacobian_padding_bands: help: "Target number of bands to add in every block of the sparse Jacobian matrix, eliminating errors from Jacobian entries that lie outside of the default sparsity pattern; when unspecified, each block gets a predetermined number of padding bands based on the typical magnitudes of its entries (only used when `use_auto_jacobian` is `true`)" value: ~ @@ -100,7 +100,7 @@ update_jacobian_every: value: "solve" debug_jacobian: help: "Whether to print summary information about the Jacobian matrix, including comparisons of different algorithms evaluated on the first column of the final state against the exact Jacobian [`true`, `false` (default)]" - value: false + value: true # TODO: Change this to false. # Radiation rad: help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]" diff --git a/config/model_configs/baroclinic_wave_conservation_ft64_sparse_autodiff.yml b/config/model_configs/baroclinic_wave_conservation_ft64_sparse_autodiff.yml new file mode 100644 index 0000000000..68740b4cc6 --- /dev/null +++ b/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 diff --git a/post_processing/jacobian_summary.jl b/post_processing/jacobian_summary.jl index a5d5b06677..ef62ccdcb1 100644 --- a/post_processing/jacobian_summary.jl +++ b/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))) diff --git a/src/prognostic_equations/implicit/auto_dense_jacobian.jl b/src/prognostic_equations/implicit/auto_dense_jacobian.jl index 9a119f5c2f..cbdd650030 100644 --- a/src/prognostic_equations/implicit/auto_dense_jacobian.jl +++ b/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) diff --git a/src/prognostic_equations/implicit/auto_sparse_jacobian.jl b/src/prognostic_equations/implicit/auto_sparse_jacobian.jl index ffc6e24584..655232ab89 100644 --- a/src/prognostic_equations/implicit/auto_sparse_jacobian.jl +++ b/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 ( diff --git a/src/prognostic_equations/implicit/jacobian.jl b/src/prognostic_equations/implicit/jacobian.jl index 151c96de33..4a659b5a39 100644 --- a/src/prognostic_equations/implicit/jacobian.jl +++ b/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). diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index e2297554d7..b20d657cdf 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/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,