diff --git a/Project.toml b/Project.toml index b5d9c352a44..bfdc53b09ce 100644 --- a/Project.toml +++ b/Project.toml @@ -21,6 +21,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" @@ -61,6 +62,7 @@ DiffEqBase = "6" DiffEqCallbacks = "2" DocStringExtensions = "0.8, 0.9" FastGaussQuadrature = "0.4, 0.5, 1" +ForwardDiff = "0.10" Insolation = "=0.9.1" Interpolations = "0.14, 0.15" IntervalSets = "0.5, 0.6, 0.7" diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index 5cb222bb112..b2d99905e08 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -80,6 +80,15 @@ eisenstat_walker_forcing_alpha: jvp_step_adjustment: help: "Amount by which the step size of the forward difference approximation of the Jacobian-vector product in the Krylov method should be scaled (only used if `use_krylov_method` is `true`)" value: 1.0 +use_exact_jacobian: + help: "Whether to solve the linear system using the exact Jacobian, which is computed with forward-mode automatic differentiation (default is `false`)" + value: false +n_steps_update_exact_jacobian: + help: "Number of timesteps between updates to the exact Jacobian (default is 0, which corresponds to updating the exact Jacobian on each Newton iteration)" + value: 0 +debug_jacobian_approximation: + help: "Whether to check the accuracy of the Jacobian approximation against the exact Jacobian every `n_steps_update_exact_jacobian` timesteps (default is `false`)" + value: false # Radiation rad: help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]" diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 182c9efb941..1992770473e 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -306,7 +306,7 @@ weakdeps = ["SparseArrays"] ChainRulesCoreSparseArraysExt = "SparseArrays" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.20.1" diff --git a/examples/Manifest.toml b/examples/Manifest.toml index 796c1f8920f..5e7cc7c01e5 100644 --- a/examples/Manifest.toml +++ b/examples/Manifest.toml @@ -299,7 +299,7 @@ weakdeps = ["CairoMakie"] CairoMakieExt = "CairoMakie" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.20.1" diff --git a/perf/Manifest.toml b/perf/Manifest.toml index 311d98e7a07..ffb4cc2f949 100644 --- a/perf/Manifest.toml +++ b/perf/Manifest.toml @@ -304,7 +304,7 @@ weakdeps = ["CairoMakie"] CairoMakieExt = "CairoMakie" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.20.1" diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index e5409138b95..ae5692ec286 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -46,6 +46,9 @@ include(joinpath("prognostic_equations", "zero_velocity.jl")) include(joinpath("prognostic_equations", "implicit", "implicit_tendency.jl")) include(joinpath("prognostic_equations", "implicit", "implicit_solver.jl")) +include(joinpath("prognostic_equations", "implicit", "approx_jacobian.jl")) +include(joinpath("prognostic_equations", "implicit", "exact_jacobian.jl")) +include(joinpath("prognostic_equations", "implicit", "debug_jacobian.jl")) include(joinpath("prognostic_equations", "remaining_tendency.jl")) include(joinpath("prognostic_equations", "forcing", "large_scale_advection.jl")) # TODO: should this be in tendencies/? diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index 482edfa15d5..716167f5845 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -320,6 +320,23 @@ function gc_func(integrator) return nothing end +function update_exact_jacobian!(integrator) + @assert integrator.alg isa CTS.IMEXAlgorithm + γs = unique(filter(!iszero, diag(integrator.alg.tableau.a_imp))) + length(γs) == 1 || + error("The exact Jacobian must be updated on every Newton iteration, \ + rather than on every timestep (or every N steps), because the \ + specified IMEX algorithm has implicit stages with distinct \ + coefficients (i.e., it is not an SDIRK algorithm).") + update_exact_jacobian!( + integrator.cache.newtons_method_cache.j, + integrator.u, + integrator.p, + integrator.dt * γs[1], + integrator.t, + ) +end + """ maybe_graceful_exit(integrator) diff --git a/src/callbacks/get_callbacks.jl b/src/callbacks/get_callbacks.jl index 8b7988205bd..b382292c93e 100644 --- a/src/callbacks/get_callbacks.jl +++ b/src/callbacks/get_callbacks.jl @@ -60,6 +60,19 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start) ) end + if ( + parsed_args["use_exact_jacobian"] || + parsed_args["debug_jacobian_approximation"] + ) && parsed_args["n_steps_update_exact_jacobian"] != 0 + callbacks = ( + callbacks..., + call_every_n_steps( + update_exact_jacobian!, + parsed_args["n_steps_update_exact_jacobian"], + ), + ) + end + if atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode # TODO: better if-else criteria? dt_rad = if parsed_args["config"] == "column" diff --git a/src/prognostic_equations/edmfx_boundary_condition.jl b/src/prognostic_equations/edmfx_boundary_condition.jl index a54c7d2b208..bb5dcbe2e60 100644 --- a/src/prognostic_equations/edmfx_boundary_condition.jl +++ b/src/prognostic_equations/edmfx_boundary_condition.jl @@ -4,8 +4,8 @@ function sgs_scalar_first_interior_bc( ᶜz_int::FT, - ᶜρ_int::FT, - ᶜscalar_int::FT, + ᶜρ_int, + ᶜscalar_int, sfc_buoyancy_flux, sfc_ρ_flux_scalar, ustar, diff --git a/src/prognostic_equations/edmfx_closures.jl b/src/prognostic_equations/edmfx_closures.jl index 7552d9907f4..1fc01a304c3 100644 --- a/src/prognostic_equations/edmfx_closures.jl +++ b/src/prognostic_equations/edmfx_closures.jl @@ -189,17 +189,17 @@ Smagorinsky length scale. """ function mixing_length( params, - ustar::FT, + ustar, ᶜz::FT, - z_sfc::FT, - ᶜdz::FT, - sfc_tke::FT, - ᶜlinear_buoygrad::FT, - ᶜtke::FT, - obukhov_length::FT, - ᶜstrain_rate_norm::FT, - ᶜPr::FT, - ᶜtke_exch::FT, + z_sfc, + ᶜdz, + sfc_tke, + ᶜlinear_buoygrad, + ᶜtke, + obukhov_length, + ᶜstrain_rate_norm, + ᶜPr, + ᶜtke_exch, ) where {FT} turbconv_params = CAP.turbconv_params(params) diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index d7863307d8e..db613d2fa19 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -85,17 +85,17 @@ end function entrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - buoy_flux_surface::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, + z_sfc, + ᶜp, + ᶜρ, + buoy_flux_surface, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, ::GeneralizedEntrainment, ) where {FT} turbconv_params = CAP.turbconv_params(params) @@ -312,19 +312,19 @@ end function detrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - buoy_flux_surface::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜentr::FT, - ᶜvert_div::FT, + z_sfc, + ᶜp, + ᶜρ, + buoy_flux_surface, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜentr, + ᶜvert_div, ::ConstantAreaDetrainment, ) where {FT} detr = ᶜentr - ᶜvert_div diff --git a/src/prognostic_equations/implicit/approx_jacobian.jl b/src/prognostic_equations/implicit/approx_jacobian.jl new file mode 100644 index 00000000000..f00158be10f --- /dev/null +++ b/src/prognostic_equations/implicit/approx_jacobian.jl @@ -0,0 +1,471 @@ +abstract type DerivativeFlag end +struct UseDerivative <: DerivativeFlag end +struct IgnoreDerivative <: DerivativeFlag end +use_derivative(::UseDerivative) = true +use_derivative(::IgnoreDerivative) = false + +""" + ApproxJacobian(Y, p; approximate_solve_iters, diffusion_flag, topography_flag) + +An approximation of the exact `ImplicitEquationJacobian`, which is updated using +analytically derived tendency derivatives and inverted using a specialized +nested linear solver. + +# Keyword Arguments + +- `approximate_solve_iters::Int`: number of iterations to take for the + approximate linear solve required when `diffusion_flag = UseDerivative()` +- `diffusion_flag::DerivativeFlag`: whether the derivative of the + diffusion tendency with respect to the quantities being diffused should be + computed or approximated as 0; must be either `UseDerivative()` or + `Ignoreerivative()` instead of a `Bool` to ensure type-stability +- `topography_flag::DerivativeFlag`: whether the derivative of vertical + contravariant velocity with respect to horizontal covariant velocity should be + computed or approximated as 0; must be either `UseDerivative()` or + `IgnoreDerivative()` instead of a `Bool` to ensure type-stability +""" +struct ApproxJacobian{M, S, F1, F2, C} <: ImplicitEquationJacobian + matrix::M + solver::S + diffusion_flag::F1 + topography_flag::F2 + default_cache::C +end + +function ApproxJacobian( + Y, + p; + approximate_solve_iters = 1, + diffusion_flag = p.atmos.diff_mode == Implicit() ? UseDerivative() : + IgnoreDerivative(), + topography_flag = has_topography(axes(Y.c)) ? UseDerivative() : + IgnoreDerivative(), + kwargs..., +) + FT = Spaces.undertype(axes(Y.c)) + CTh = CTh_vector_type(axes(Y.c)) + + TridiagonalRow = TridiagonalMatrixRow{FT} + BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} + TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} + BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} + BidiagonalRow_C3xACTh = + BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} + TridiagonalRow_C3xACT3 = + TridiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} + + is_in_Y(name) = MatrixFields.has_field(Y, name) + + ρq_tot_if_available = is_in_Y(@name(c.ρq_tot)) ? (@name(c.ρq_tot),) : () + ρatke_if_available = + is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : () + + tracer_names = ( + @name(c.ρq_tot), + @name(c.ρq_liq), + @name(c.ρq_ice), + @name(c.ρq_rai), + @name(c.ρq_sno), + ) + available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names) + other_names = (@name(c.sgsʲs), @name(f.sgsʲs), @name(sfc)) + other_available_names = MatrixFields.unrolled_filter(is_in_Y, other_names) + + # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0, + # which means that multiplying inv(-1) by a Float32 will yield a Float64. + identity_blocks = MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (@name(c.ρ), other_available_names...), + ) + + active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...) + advection_blocks = ( + ( + use_derivative(topography_flag) ? + MatrixFields.unrolled_map( + name -> + (name, @name(c.uₕ)) => + similar(Y.c, TridiagonalRow_ACTh), + active_scalar_names, + ) : () + )..., + MatrixFields.unrolled_map( + name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + active_scalar_names, + )..., + MatrixFields.unrolled_map( + name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), + active_scalar_names, + )..., + (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh), + (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), + ) + + diffused_scalar_names = + (@name(c.ρe_tot), available_tracer_names..., ρatke_if_available...) + diffusion_blocks = if use_derivative(diffusion_flag) + ( + MatrixFields.unrolled_map( + name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow), + diffused_scalar_names, + )..., + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + diffused_scalar_names, + )..., + ( + is_in_Y(@name(c.ρq_tot)) ? + ( + (@name(c.ρe_tot), @name(c.ρq_tot)) => + similar(Y.c, TridiagonalRow), + ) : () + )..., + (@name(c.uₕ), @name(c.uₕ)) => + !isnothing(p.atmos.turbconv_model) || + diffuse_momentum(p.atmos.vert_diff) ? + similar(Y.c, TridiagonalRow) : FT(-1) * I, + ) + else + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (diffused_scalar_names..., @name(c.uₕ)), + ) + end + + matrix = MatrixFields.FieldMatrix( + identity_blocks..., + advection_blocks..., + diffusion_blocks..., + ) + + names₁_group₁ = (@name(c.ρ), other_available_names...) + names₁_group₂ = (available_tracer_names..., ρatke_if_available...) + names₁_group₃ = (@name(c.ρe_tot),) + names₁ = (names₁_group₁..., names₁_group₂..., names₁_group₃...) + + alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) + alg = if use_derivative(diffusion_flag) + alg₁_subalg₂ = + is_in_Y(@name(c.ρq_tot)) ? + (; + alg₂ = MatrixFields.BlockLowerTriangularSolve(names₁_group₂...) + ) : (;) + alg₁ = MatrixFields.BlockLowerTriangularSolve( + names₁_group₁...; + alg₁_subalg₂..., + ) + MatrixFields.ApproximateBlockArrowheadIterativeSolve( + names₁...; + alg₁, + alg₂, + P_alg₁ = MatrixFields.MainDiagonalPreconditioner(), + n_iters = approximate_solve_iters, + ) + else + MatrixFields.BlockArrowheadSolve(names₁...; alg₂) + end + + return ApproxJacobian( + matrix, + MatrixFields.FieldMatrixSolver(alg, matrix, Y), + diffusion_flag, + topography_flag, + default_jacobian_cache(Y; kwargs...), + ) +end + +NVTX.@annotate function approximate_jacobian!(A::ApproxJacobian, Y, p, dtγ, t) + # Remove unnecessary values from p to avoid allocations in bycolumn. + reduced_p = (; + p.precomputed.ᶜspecific, + p.precomputed.ᶜK, + p.precomputed.ᶜts, + p.precomputed.ᶜp, + ( + p.atmos.precip_model isa Microphysics1Moment ? + (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;) + )..., + p.precomputed.ᶜh_tot, + ( + use_derivative(A.diffusion_flag) ? + (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;) + )..., + ( + use_derivative(A.diffusion_flag) && + p.atmos.turbconv_model isa AbstractEDMF ? + (; p.precomputed.ᶜtke⁰, p.precomputed.ᶜmixing_length) : (;) + )..., + ( + use_derivative(A.diffusion_flag) && + p.atmos.turbconv_model isa PrognosticEDMFX ? + (; p.precomputed.ᶜρa⁰) : (;) + )..., + p.core.ᶜΦ, + p.core.ᶠgradᵥ_ᶜΦ, + p.core.ᶜρ_ref, + p.core.ᶜp_ref, + p.scratch.ᶜtemp_scalar, + p.scratch.∂ᶜK_∂ᶜuₕ, + p.scratch.∂ᶜK_∂ᶠu₃, + p.scratch.ᶠp_grad_matrix, + p.scratch.ᶜadvection_matrix, + p.scratch.ᶜdiffusion_h_matrix, + p.scratch.ᶜdiffusion_u_matrix, + p.dt, + p.params, + p.atmos, + ( + p.atmos.rayleigh_sponge isa RayleighSponge ? + (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;) + )..., + ) + Fields.bycolumn(axes(Y.c)) do colidx + approximate_column_jacobian!(A, Y, reduced_p, dtγ, colidx) + end +end + +function approximate_column_jacobian!(A, Y, p, dtγ, colidx) + (; matrix, diffusion_flag, topography_flag) = A + (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p + (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p + (; ᶜdiffusion_h_matrix, ᶜdiffusion_u_matrix, params) = p + + FT = Spaces.undertype(axes(Y.c)) + CTh = CTh_vector_type(axes(Y.c)) + one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' + + cv_d = FT(CAP.cv_d(params)) + Δcv_v = FT(CAP.cv_v(params)) - cv_d + T_tri = FT(CAP.T_triple(params)) + e_int_v0 = T_tri * Δcv_v - FT(CAP.e_int_v0(params)) + thermo_params = CAP.thermodynamics_params(params) + + ᶜρ = Y.c.ρ + ᶜuₕ = Y.c.uₕ + ᶠu₃ = Y.f.u₃ + ᶜJ = Fields.local_geometry_field(Y.c).J + ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ + ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ + + ᶜkappa_m = p.ᶜtemp_scalar + @. ᶜkappa_m[colidx] = + TD.gas_constant_air(thermo_params, ᶜts[colidx]) / + TD.cv_m(thermo_params, ᶜts[colidx]) + + if use_derivative(topography_flag) + @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow( + adjoint(CTh(ᶜuₕ[colidx])) + + adjoint(ᶜinterp(ᶠu₃[colidx])) * g³ʰ(ᶜgⁱʲ[colidx]), + ) + else + @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ[colidx]))) + end + @. ∂ᶜK_∂ᶠu₃[colidx] = + ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[colidx]))) + + DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix() + + @. ᶠp_grad_matrix[colidx] = + DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() + + @. ᶜadvection_matrix[colidx] = + -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) + + if use_derivative(topography_flag) + ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] + @. ∂ᶜρ_err_∂ᶜuₕ[colidx] = + dtγ * ᶜadvection_matrix[colidx] ⋅ + ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx])) + end + ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] + @. ∂ᶜρ_err_∂ᶠu₃[colidx] = + dtγ * ᶜadvection_matrix[colidx] ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ[colidx])) + + tracer_info = ( + (@name(c.ρe_tot), @name(ᶜh_tot)), + (@name(c.ρq_tot), @name(ᶜspecific.q_tot)), + ) + MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name) + MatrixFields.has_field(Y, ρχ_name) || return + ᶜχ = MatrixFields.get_field(p, χ_name) + if use_derivative(topography_flag) + ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] + end + ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] + use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] = + dtγ * ᶜadvection_matrix[colidx] ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅ + ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ + DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx])) + @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = + dtγ * ᶜadvection_matrix[colidx] ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * g³³(ᶠgⁱʲ[colidx])) + end + + ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] + ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] + @. ∂ᶠu₃_err_∂ᶜρ[colidx] = + dtγ * ( + ᶠp_grad_matrix[colidx] ⋅ DiagonalMatrixRow( + ᶜkappa_m[colidx] * (T_tri * cv_d - ᶜK[colidx] - ᶜΦ[colidx]), + ) + + DiagonalMatrixRow( + ( + ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - + ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] + ) / abs2(ᶠinterp(ᶜρ[colidx])), + ) ⋅ ᶠinterp_matrix() + ) + @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] = + dtγ * ᶠp_grad_matrix[colidx] ⋅ DiagonalMatrixRow(ᶜkappa_m[colidx]) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] + @. ∂ᶠu₃_err_∂ᶜρq_tot[colidx] = + dtγ * ᶠp_grad_matrix[colidx] ⋅ + DiagonalMatrixRow(ᶜkappa_m[colidx] * e_int_v0) + end + + ∂ᶠ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_∂ᶜuₕ[colidx] = + dtγ * ᶠp_grad_matrix[colidx] ⋅ + DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ ∂ᶜK_∂ᶜuₕ[colidx] + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = + dtγ * ( + ᶠp_grad_matrix[colidx] ⋅ + DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ + ∂ᶜK_∂ᶠu₃[colidx] + + DiagonalMatrixRow(-p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,)) + ) - (I_u₃,) + else + @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = + dtγ * ᶠp_grad_matrix[colidx] ⋅ + DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ + ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,) + end + + if use_derivative(diffusion_flag) + (; ᶜK_h, ᶜK_u) = p + @. ᶜdiffusion_h_matrix[colidx] = + ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(ᶜK_h[colidx])) ⋅ + ᶠgradᵥ_matrix() + if ( + MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) || + !isnothing(p.atmos.turbconv_model) || + diffuse_momentum(p.atmos.vert_diff) + ) + @. ᶜdiffusion_u_matrix[colidx] = + ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(ᶜK_u[colidx])) ⋅ + ᶠgradᵥ_matrix() + end + + ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] + ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρ[colidx] = + dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ DiagonalMatrixRow( + ( + -(1 + ᶜkappa_m[colidx]) * ᶜspecific.e_tot[colidx] - + ᶜkappa_m[colidx] * e_int_v0 * ᶜspecific.q_tot[colidx] + ) / ᶜρ[colidx], + ) + @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] = + dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ + DiagonalMatrixRow((1 + ᶜkappa_m[colidx]) / ᶜρ[colidx]) - (I,) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρq_tot[colidx] = + dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ + DiagonalMatrixRow(ᶜkappa_m[colidx] * e_int_v0 / ᶜρ[colidx]) + end + + tracer_info = ( + (@name(c.ρq_tot), @name(q_tot)), + (@name(c.ρq_liq), @name(q_liq)), + (@name(c.ρq_ice), @name(q_ice)), + (@name(c.ρq_rai), @name(q_rai)), + (@name(c.ρq_sno), @name(q_sno)), + ) + MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name) + MatrixFields.has_field(Y, ρq_name) || return + ᶜq = MatrixFields.get_field(ᶜspecific, q_name) + ∂ᶜρq_err_∂ᶜρ = matrix[ρq_name, @name(c.ρ)] + ∂ᶜρq_err_∂ᶜρq = matrix[ρq_name, ρq_name] + @. ∂ᶜρq_err_∂ᶜρ[colidx] = + dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ + DiagonalMatrixRow(-(ᶜq[colidx]) / ᶜρ[colidx]) + @. ∂ᶜρq_err_∂ᶜρq[colidx] = + dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ + DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,) + end + + if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) + turbconv_params = CAP.turbconv_params(params) + c_d = CAP.tke_diss_coeff(turbconv_params) + (; ᶜtke⁰, ᶜmixing_length, dt) = p + ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ + ᶜρatke⁰ = Y.c.sgs⁰.ρatke + + dissipation_rate(tke⁰, mixing_length) = + tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : 1 / dt + ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = + tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : 0 + + ᶜdissipation_matrix_diagonal = p.ᶜtemp_scalar + @. ᶜdissipation_matrix_diagonal[colidx] = + ᶜρatke⁰[colidx] * + ∂dissipation_rate_∂tke⁰(ᶜtke⁰[colidx], ᶜmixing_length[colidx]) + + ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] + ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = + matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] + @. ∂ᶜρatke⁰_err_∂ᶜρ[colidx] = + dtγ * ( + ᶜdiffusion_u_matrix[colidx] - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal[colidx]) + ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰[colidx]) / ᶜρa⁰[colidx]) + @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰[colidx] = + dtγ * ( + ( + ᶜdiffusion_u_matrix[colidx] - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal[colidx]) + ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰[colidx]) - + DiagonalMatrixRow( + dissipation_rate(ᶜtke⁰[colidx], ᶜmixing_length[colidx]), + ) + ) - (I,) + end + + if ( + !isnothing(p.atmos.turbconv_model) || + diffuse_momentum(p.atmos.vert_diff) + ) + ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)] + @. ∂ᶜuₕ_err_∂ᶜuₕ[colidx] = + dtγ * DiagonalMatrixRow(1 / ᶜρ[colidx]) ⋅ + ᶜdiffusion_u_matrix[colidx] - (I,) + end + + ᶠlg = Fields.local_geometry_field(Y.f) + precip_info = + ((@name(c.ρq_rai), @name(ᶜwᵣ)), (@name(c.ρq_sno), @name(ᶜwₛ))) + MatrixFields.unrolled_foreach(precip_info) do (ρqₚ_name, wₚ_name) + MatrixFields.has_field(Y, ρqₚ_name) || return + ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name] + ᶜwₚ = MatrixFields.get_field(p, wₚ_name) + @. ∂ᶜρqₚ_err_∂ᶜρqₚ[colidx] += + dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + CT3(unit_basis_vector_data(CT3, ᶠlg[colidx])) * + ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]), + ) ⋅ ᶠright_bias_matrix() ⋅ + DiagonalMatrixRow(-(ᶜwₚ[colidx]) / ᶜρ[colidx]) + end + end +end + +NVTX.@annotate invert_jacobian!(A::ApproxJacobian, x, b) = + MatrixFields.field_matrix_solve!(A.solver, x, A.matrix, b) diff --git a/src/prognostic_equations/implicit/debug_jacobian.jl b/src/prognostic_equations/implicit/debug_jacobian.jl new file mode 100644 index 00000000000..927df781ad4 --- /dev/null +++ b/src/prognostic_equations/implicit/debug_jacobian.jl @@ -0,0 +1,246 @@ +struct DebugJacobian{E <: ExactJacobian, A <: ApproxJacobian} <: + ImplicitEquationJacobian + exact_jacobian::E + approx_jacobian::A + use_exact_jacobian::Bool + exact_jacobian_just_updated::Base.RefValue{Bool} +end + +DebugJacobian(exact_jacobian, approx_jacobian; use_exact_jacobian) = + DebugJacobian( + exact_jacobian, + approx_jacobian, + use_exact_jacobian, + Ref(false), + ) + +always_update_exact_jacobian(A::DebugJacobian) = + always_update_exact_jacobian(A.exact_jacobian) + +function factorize_exact_jacobian!(A::DebugJacobian, Y, p, dtγ, t) + factorize_exact_jacobian!(A.exact_jacobian, Y, p, dtγ, t) + A.exact_jacobian_just_updated[] = true +end + +approximate_jacobian!(A::DebugJacobian, Y, p, dtγ, t) = + A.use_exact_jacobian || + approximate_jacobian!(A.approx_jacobian, Y, p, dtγ, t) + +invert_jacobian!(A::DebugJacobian, x, b) = + if A.exact_jacobian_just_updated[] + A.exact_jacobian_just_updated[] = false + else + if A.use_exact_jacobian + invert_jacobian!(A.exact_jacobian, x, b) + else + invert_jacobian!(A.approx_jacobian, x, b) + end + end + +function debug_jacobian_cache(approx_jacobian, Y, p) + FT = eltype(Y) + DA = ClimaComms.array_type(Y) + + @assert all(name -> name in (:c, :f, :sfc), propertynames(Y)) + n_c_levels = Spaces.nlevels(axes(Y.c)) + n_f_levels = Spaces.nlevels(axes(Y.f)) + n_c_fields = DataLayouts.ncomponents(Fields.field_values(Y.c)) + n_f_fields = DataLayouts.ncomponents(Fields.field_values(Y.f)) + n_sfc_fields = + hasproperty(Y, :sfc) ? + DataLayouts.ncomponents(Fields.field_values(Y.sfc)) : 0 + n_εs = n_c_levels * n_c_fields + n_f_levels * n_f_fields + n_sfc_fields + dual_T = ForwardDiff.Dual{Nothing, FT, n_εs} + + Y_dual = similar(Y, dual_T) # swap every FT in parent arrays with a dual_T + Y_dual_copy = similar(Y_dual) + Y_err_dual = similar(Y_dual) + + precomputed_dual = + Fields._values(similar(Fields.FieldVector(; p.precomputed...), dual_T)) + scratch_dual = + Fields._values(similar(Fields.FieldVector(; p.scratch...), dual_T)) + + column_matrices = map(_ -> DA{FT}(undef, n_εs, n_εs), colidx_iterator(Y.c)) + + exact_jacobian = Dict() + for (row_name, col_name) in keys(approx_jacobian) + entry_rows = Spaces.nlevels(axes(MatrixFields.get_field(Y, row_name))) + entry_cols = Spaces.nlevels(axes(MatrixFields.get_field(Y, col_name))) + exact_jacobian[row_name, col_name] = ntuple( + _ -> DA{FT}(undef, entry_rows, entry_cols), + length(column_matrices), + ) + end + + return (; + Y_dual, + precomputed_dual, + scratch_dual, + Y_dual_copy, + Y_err_dual, + column_matrices, + exact_jacobian, + ) +end + +function print_jacobian_debug_info(A, Y, p, t, dt) + (; + Y_dual, + precomputed_dual, + scratch_dual, + Y_dual_copy, + Y_err_dual, + column_matrices, + exact_jacobian, + ) = A.debug_cache + + FT = eltype(Y) + n_c_levels = Spaces.nlevels(axes(Y.c)) + n_f_levels = Spaces.nlevels(axes(Y.f)) + n_c_fields = DataLayouts.ncomponents(Fields.field_values(Y.c)) + n_f_fields = DataLayouts.ncomponents(Fields.field_values(Y.f)) + n_sfc_fields = + hasproperty(Y, :sfc) ? + DataLayouts.ncomponents(Fields.field_values(Y.sfc)) : 0 + n_εs = n_c_levels * n_c_fields + n_f_levels * n_f_fields + n_sfc_fields + + Y_dual .= Y + for level in 1:n_c_levels, field_index in 1:n_c_fields + ε_index = (field_index - 1) * n_c_levels + level + ε = ForwardDiff.Dual(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + c_level_values = DataLayouts.level(Fields.field_values(Y_dual.c), level) + @assert c_level_values isa DataLayouts.IJFH + view(parent(c_level_values), :, :, field_index, :) .+= ε + end + for level in 1:n_f_levels, field_index in 1:n_f_fields + ε_index = + n_c_fields * n_c_levels + (field_index - 1) * n_f_levels + level + ε = ForwardDiff.Dual(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + f_level_values = DataLayouts.level(Fields.field_values(Y_dual.f), level) + @assert f_level_values isa DataLayouts.IJFH + view(parent(f_level_values), :, :, field_index, :) .+= ε + end + for field_index in 1:n_sfc_fields + ε_index = + n_c_fields * n_c_levels + n_f_fields * n_f_levels + field_index + ε = ForwardDiff.Dual(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + sfc_values = Fields.field_values(Y_dual.sfc) + @assert sfc_values isa DataLayouts.IJFH + view(parent(sfc_values), :, :, field_index, :) .+= ε + end + + p_dual_args = ntuple(Val(fieldcount(typeof(p)))) do cache_field_index + cache_field_name = fieldname(typeof(p), cache_field_index) + if cache_field_name == :precomputed + precomputed_dual + elseif cache_field_name == :scratch + scratch_dual + else + getfield(p, cache_field_index) + end + end + p_dual = AtmosCache(p_dual_args...) + + # compute the exact Jacobian with dtγ = dt + Y_dual_copy .= Y_dual # need to copy because Y_dual gets changed at boundary + set_precomputed_quantities!(Y_dual, p_dual, t) # changes Y_dual at boundary + implicit_tendency!(Y_err_dual, Y_dual, p_dual, t) + @. Y_err_dual = dt * Y_err_dual - Y_dual_copy + + # compute the approximate Jacobian with dtγ = dt + set_precomputed_quantities!(Y, p, t) + Wfact!(A, Y, p, dt, t) + approx_jacobian = A.matrix + + # print the exact Jacobian at column 1 + column_index = 1 + colidx = Fields.ColumnIndex{2}((1, 1), 1) + exact_column_jacobian = column_matrices[column_index] + column_c_data = parent(Y_err_dual.c[colidx]) + column_f_data = parent(Y_err_dual.f[colidx]) + for c_field_index in 1:n_c_fields + for level in 1:n_c_levels + ε_index = (c_field_index - 1) * n_c_levels + level + exact_column_jacobian[ε_index, :] .= + column_c_data[level, c_field_index].partials.values + end + end + for f_field_index in 1:n_f_fields + for level in 1:n_f_levels + ε_index = + n_c_fields * n_c_levels + + (f_field_index - 1) * n_f_levels + + level + exact_column_jacobian[ε_index, :] .= + column_f_data[level, f_field_index].partials.values + end + end + @info "Full Jacobian at time $t" + display(round.(exact_column_jacobian; sigdigits = 1)) + + # print individual blocks of the two Jacobians at column 1 + column_index = 1 + colidx = Fields.ColumnIndex{2}((1, 1), 1) + scalar_row(row) = + map(row) do x + x isa Union{Geometry.AxisTensor, Geometry.AdjointAxisTensor} ? + x[1] : x + end + function abs_eps_error(approx, exact) + float_eps_error = norm(exact - approx) / eps(one(norm(exact))) + return float_eps_error > typemax(Int) ? Inf : ceil(Int, float_eps_error) + end + function rel_eps_error(approx, exact) + float_eps_error = norm(exact - approx) / eps(norm(exact)) + return float_eps_error > typemax(Int) ? Inf : ceil(Int, float_eps_error) + end + + for (row_name, col_name) in keys(approx_jacobian) + @info "Info for ($row_name, $col_name) at time $t" + exact_jacobian_row = + parent(MatrixFields.get_field(Y_err_dual[colidx], row_name)) + exact_entry = exact_jacobian[row_name, col_name][column_index] + + col_field = MatrixFields.get_field(Y[colidx], col_name) + is_c_col = axes(col_field).staggering isa Spaces.CellCenter + if length(parent(col_field).indices[4]) != 1 + @info "TODO: Handle blocks encompassing multiple fields correctly" + end + col_field_index = parent(col_field).indices[4][1] + if is_c_col + ε_start_index = (col_field_index - 1) * n_c_levels + 1 + ε_end_index = col_field_index * n_c_levels + else + ε_start_index = + n_c_fields * n_c_levels + (col_field_index - 1) * n_f_levels + 1 + ε_end_index = n_c_fields * n_c_levels + col_field_index * n_f_levels + end + ε_range = ε_start_index:ε_end_index + + n_row_levels = size(exact_entry, 1) + for level in 1:n_row_levels + exact_entry[level, :] .= + exact_jacobian_row[level].partials.values[ε_range] + end + display(exact_entry) + + _approx_entry = approx_jacobian[row_name, col_name] + if _approx_entry isa LinearAlgebra.UniformScaling + approx_entry = Matrix(_approx_entry, n_row_levels, n_row_levels) + else + approx_entry = MatrixFields.column_field2array( + scalar_row.(_approx_entry[colidx]), + ) + end + display(approx_entry) + + max_abs_eps_error = + maximum(abs_eps_error.(Array(approx_entry), exact_entry)) + @info "Max absolute error = $max_abs_eps_error eps" + + max_rel_eps_error = + maximum(rel_eps_error.(Array(approx_entry), exact_entry)) + @info "Max relative error = $max_rel_eps_error eps" + end +end diff --git a/src/prognostic_equations/implicit/exact_jacobian.jl b/src/prognostic_equations/implicit/exact_jacobian.jl new file mode 100644 index 00000000000..59d5a13426b --- /dev/null +++ b/src/prognostic_equations/implicit/exact_jacobian.jl @@ -0,0 +1,185 @@ +# Custom implementation of bycolumn for iterating over columns on both CPUs and +# GPUs (bycolum does not support GPUs). +function colidx_iterator(field) + qs = 1:Quadratures.degrees_of_freedom(Spaces.quadrature_style(axes(field))) + hs = Spaces.eachslabindex(Spaces.horizontal_space(axes(field))) + return Iterators.map((i, j, h) -> Fields.ColumnIndex((i, j), h), qs, qs, hs) +end + +struct ExactJacobian{Y, P, S, M, F, C} <: ImplicitEquationJacobian + Y_dual::Y + Y_dual_copy::Y + Y_err_dual::Y + precomputed_dual::P + scratch_dual::S + column_matrices::M + column_lu_factorizations::F + preserve_unfactorized_matrices::Bool + default_cache::C +end + +function ExactJacobian(Y, p; preserve_unfactorized_matrices = false, kwargs...) + FT = eltype(Y) + DA = ClimaComms.array_type(Y) + + @assert all(name -> name in (:c, :f, :sfc), propertynames(Y)) + n_c_levels = Spaces.nlevels(axes(Y.c)) + n_f_levels = Spaces.nlevels(axes(Y.f)) + n_c_fields = DataLayouts.ncomponents(Fields.field_values(Y.c)) + n_f_fields = DataLayouts.ncomponents(Fields.field_values(Y.f)) + n_sfc_fields = + hasproperty(Y, :sfc) ? + DataLayouts.ncomponents(Fields.field_values(Y.sfc)) : 0 + n_εs = n_c_levels * n_c_fields + n_f_levels * n_f_fields + n_sfc_fields + dual_T = ForwardDiff.Dual{ExactJacobian, FT, n_εs} + + Y_dual = similar(Y, dual_T) # swap every FT in parent arrays with a dual_T + Y_dual_copy = similar(Y_dual) + Y_err_dual = similar(Y_dual) + + precomputed_dual = + Fields._values(similar(Fields.FieldVector(; p.precomputed...), dual_T)) + scratch_dual = + Fields._values(similar(Fields.FieldVector(; p.scratch...), dual_T)) + + column_matrices = map(_ -> DA{FT}(I, n_εs, n_εs), colidx_iterator(Y.c)) + column_lu_factorizations = + map(preserve_unfactorized_matrices ? lu : lu!, column_matrices) + + return ExactJacobian( + Y_dual, + Y_dual_copy, + Y_err_dual, + precomputed_dual, + scratch_dual, + column_matrices, + column_lu_factorizations, + preserve_unfactorized_matrices, + default_jacobian_cache(Y; kwargs...), + ) +end + +NVTX.@annotate function factorize_exact_jacobian!( + A::ExactJacobian, + Y, + p, + dtγ, + t, +) + (; Y_dual, Y_dual_copy, Y_err_dual, precomputed_dual, scratch_dual) = A + (; column_matrices, column_lu_factorizations) = A + (; preserve_unfactorized_matrices) = A + + FT = eltype(Y) + n_c_levels = Spaces.nlevels(axes(Y.c)) + n_f_levels = Spaces.nlevels(axes(Y.f)) + n_c_fields = DataLayouts.ncomponents(Fields.field_values(Y.c)) + n_f_fields = DataLayouts.ncomponents(Fields.field_values(Y.f)) + n_sfc_fields = + hasproperty(Y, :sfc) ? + DataLayouts.ncomponents(Fields.field_values(Y.sfc)) : 0 + n_εs = n_c_levels * n_c_fields + n_f_levels * n_f_fields + n_sfc_fields + dual_T = ForwardDiff.Dual{ExactJacobian, FT, n_εs} + + Y_dual .= Y + for level in 1:n_c_levels, field_index in 1:n_c_fields + ε_index = (field_index - 1) * n_c_levels + level + ε = dual_T(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + c_values = Fields.field_values(Y_dual.c) + @assert c_values isa DataLayouts.VIJFH + view(parent(c_values), level, :, :, field_index, :) .+= ε + end + for level in 1:n_f_levels, field_index in 1:n_f_fields + ε_index = + n_c_fields * n_c_levels + (field_index - 1) * n_f_levels + level + ε = dual_T(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + f_values = Fields.field_values(Y_dual.f) + @assert f_values isa DataLayouts.VIJFH + view(parent(f_values), level, :, :, field_index, :) .+= ε + end + for field_index in 1:n_sfc_fields + ε_index = + n_c_fields * n_c_levels + n_f_fields * n_f_levels + field_index + ε = dual_T(0, ntuple(i -> i == ε_index ? 1 : 0, n_εs)...) + sfc_values = Fields.field_values(Y_dual.sfc) + @assert sfc_values isa DataLayouts.IJFH + view(parent(sfc_values), :, :, field_index, :) .+= ε + end + + p_dual_args = ntuple(Val(fieldcount(typeof(p)))) do cache_field_index + cache_field_name = fieldname(typeof(p), cache_field_index) + if cache_field_name == :precomputed + precomputed_dual + elseif cache_field_name == :scratch + scratch_dual + else + getfield(p, cache_field_index) + end + end + p_dual = AtmosCache(p_dual_args...) + + Y_dual_copy .= Y_dual # need a copy because Y_dual gets changed at boundary + set_precomputed_quantities!(Y_dual, p_dual, t) # changes Y_dual at boundary + implicit_tendency!(Y_err_dual, Y_dual, p_dual, t) + @. Y_err_dual = dtγ * Y_err_dual - Y_dual_copy + + for (colidx_index, colidx) in enumerate(colidx_iterator(Y.c)) + column_matrix = column_matrices[colidx_index] + for level in 1:n_c_levels, field_index in 1:n_c_fields + row_index = (field_index - 1) * n_c_levels + level + row_dual_value = parent(Y_err_dual.c[colidx])[level, field_index] + column_matrix[row_index, :] .= ForwardDiff.partials(row_dual_value) + end + for level in 1:n_f_levels, field_index in 1:n_f_fields + row_index = + n_c_fields * n_c_levels + (field_index - 1) * n_f_levels + level + row_dual_value = parent(Y_err_dual.f[colidx])[level, field_index] + column_matrix[row_index, :] .= ForwardDiff.partials(row_dual_value) + end + for field_index in 1:n_sfc_fields + row_index = + n_c_fields * n_c_levels + n_f_fields * n_f_levels + field_index + row_dual_value = parent(Y_err_dual.sfc[colidx])[field_index] + column_matrix[row_index, :] .= ForwardDiff.partials(row_dual_value) + end + + column_matrix_copy = column_lu_factorizations[colidx_index].factors + if preserve_unfactorized_matrices + column_matrix_copy .= column_matrix + else + @assert column_matrix_copy === column_matrix + end + column_lu_factorizations[colidx_index] = lu!(column_matrix_copy) + # FIXME: lu! allocates an array of length n_εs for pivoting + end +end + +NVTX.@annotate function invert_jacobian!(A::ExactJacobian, x, b) + (; column_lu_factorizations) = A + for (colidx_index, colidx) in enumerate(colidx_iterator(x.c)) + column_lu_factorization = column_lu_factorizations[colidx_index] + column_x_args = (; c = x.c[colidx], f = x.f[colidx]) + column_x_sfc_args = hasproperty(x, :sfc) ? (; sfc = x.sfc[colidx]) : (;) + column_x = Fields.FieldVector(; column_x_args..., column_x_sfc_args...) + column_b_args = (; c = b.c[colidx], f = b.f[colidx]) + column_b_sfc_args = hasproperty(b, :sfc) ? (; sfc = b.sfc[colidx]) : (;) + column_b = Fields.FieldVector(; column_b_args..., column_b_sfc_args...) + ldiv!(column_x, column_lu_factorization, column_b) + end +end + +# Set the derivative of `sqrt(x)` to `iszero(x) ? zero(x) : inv(2 * sqrt(x))` in +# order to properly handle derivatives of `x * sqrt(x)`. Without this change, +# the derivative of `x * sqrt(x)` is `NaN` when `x` is zero. This method +# specializes on the tag `ExactJacobian` in order to avoid breaking +# precompilation by overwriting the generic method for `Dual` in `ForwardDiff`, +# and to avoid type piracy by using the default tag `Nothing`. +# TODO: Should this be moved to a package extension? +@inline function Base.sqrt(d::ForwardDiff.Dual{ExactJacobian}) + tag = Val{ExactJacobian}() + x = ForwardDiff.value(d) + partials = ForwardDiff.partials(d) + val = sqrt(x) + deriv = iszero(x) ? zero(x) : inv(2 * val) + return ForwardDiff.dual_definition_retval(tag, val, deriv, partials) +end diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 65bba0bd0cf..14117433120 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -1,18 +1,17 @@ -import LinearAlgebra: I, Adjoint, ldiv! +import ForwardDiff +import LinearAlgebra: I, Adjoint, lu, lu!, ldiv! import ClimaCore.MatrixFields: @name -using ClimaCore.MatrixFields - -abstract type DerivativeFlag end -struct UseDerivative <: DerivativeFlag end -struct IgnoreDerivative <: DerivativeFlag end -use_derivative(::UseDerivative) = true -use_derivative(::IgnoreDerivative) = false """ - ImplicitEquationJacobian(Y, atmos; approximate_solve_iters, diffusion_flag, topography_flag, transform_flag) + ImplicitEquationJacobian A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the -implicit step with the state ``Y``. +implicit step with the state ``Y``. Concrete implementations of this abstract +type should define 4 methods: + - `always_update_exact_jacobian(A::ImplicitEquationJacobian)` + - `factorize_exact_jacobian!(A::ImplicitEquationJacobian, Y, p, dtγ, t)` + - `approximate_jacobian!(A::ImplicitEquationJacobian, Y, p, dtγ, t)` + - `invert_jacobian!(x, A::ImplicitEquationJacobian, b)` # Background @@ -66,528 +65,55 @@ steps is repeated, i.e., ``ΔY[1]`` is computed and used to update ``Y`` to when the error ``E(Y)`` is sufficiently close to 0 (according to the convergence condition passed to Newton's method), or when the maximum number of iterations is reached. - -# Arguments - -- `Y::FieldVector`: the state of the simulation -- `atmos::AtmosModel`: the model configuration -- `approximate_solve_iters::Int`: number of iterations to take for the - approximate linear solve required when `diffusion_flag = UseDerivative()` -- `diffusion_flag::DerivativeFlag`: whether the derivative of the - diffusion tendency with respect to the quantities being diffused should be - computed or approximated as 0; must be either `UseDerivative()` or - `Ignoreerivative()` instead of a `Bool` to ensure type-stability -- `topography_flag::DerivativeFlag`: whether the derivative of vertical - contravariant velocity with respect to horizontal covariant velocity should be - computed or approximated as 0; must be either `UseDerivative()` or - `IgnoreDerivative()` instead of a `Bool` to ensure type-stability -- `transform_flag::Bool`: whether the error should be transformed from ``E(Y)`` - to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from - OrdinaryDiffEq.jl """ -struct ImplicitEquationJacobian{ - M <: MatrixFields.FieldMatrix, - S <: MatrixFields.FieldMatrixSolver, - F1 <: DerivativeFlag, - F2 <: DerivativeFlag, - T <: Fields.FieldVector, - R <: Base.RefValue, -} - # stores the matrix E'(Y) = Δt * T_imp'(Y) - I - matrix::M - - # solves the linear equation E'(Y) * ΔY = E(Y) for ΔY - solver::S - - # flags that determine how E'(Y) is approximated - diffusion_flag::F1 - topography_flag::F2 - - # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs - temp_b::T - temp_x::T - - # required by OrdinaryDiffEq.jl to run non-Rosenbrock timestepping schemes - transform_flag::Bool - dtγ_ref::R -end - -function ImplicitEquationJacobian( - Y, - atmos; - approximate_solve_iters = 1, - diffusion_flag = atmos.diff_mode == Implicit() ? UseDerivative() : - IgnoreDerivative(), - topography_flag = has_topography(axes(Y.c)) ? UseDerivative() : - IgnoreDerivative(), - transform_flag = false, -) - FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) - - TridiagonalRow = TridiagonalMatrixRow{FT} - BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} - TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} - BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - BidiagonalRow_C3xACTh = - BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} - TridiagonalRow_C3xACT3 = - TridiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} +abstract type ImplicitEquationJacobian end - is_in_Y(name) = MatrixFields.has_field(Y, name) - - ρq_tot_if_available = is_in_Y(@name(c.ρq_tot)) ? (@name(c.ρq_tot),) : () - ρatke_if_available = - is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : () - - tracer_names = ( - @name(c.ρq_tot), - @name(c.ρq_liq), - @name(c.ρq_ice), - @name(c.ρq_rai), - @name(c.ρq_sno), - ) - available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names) - other_names = (@name(c.sgsʲs), @name(f.sgsʲs), @name(sfc)) - other_available_names = MatrixFields.unrolled_filter(is_in_Y, other_names) +# We only use A, but ClimaTimeSteppers.jl currently requires us to pass +# jac_prototype and then calls similar(jac_prototype) to obtain A. Since +# this was only done for compatibility with OrdinaryDiffEq.jl, which we are no +# longer using, we should fix this in a new version of ClimaTimeSteppers.jl. +Base.similar(A::ImplicitEquationJacobian) = A - # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0, - # which means that multiplying inv(-1) by a Float32 will yield a Float64. - identity_blocks = MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - (@name(c.ρ), other_available_names...), - ) +default_jacobian_cache(Y; always_update_exact_jacobian = true) = + (; always_update_exact_jacobian, temp_b = similar(Y), temp_x = similar(Y)) - active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...) - advection_blocks = ( - ( - use_derivative(topography_flag) ? - MatrixFields.unrolled_map( - name -> - (name, @name(c.uₕ)) => - similar(Y.c, TridiagonalRow_ACTh), - active_scalar_names, - ) : () - )..., - MatrixFields.unrolled_map( - name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), - active_scalar_names, - )..., - MatrixFields.unrolled_map( - name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), - active_scalar_names, - )..., - (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh), - (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), - ) +always_update_exact_jacobian(A) = A.default_cache.always_update_exact_jacobian - diffused_scalar_names = - (@name(c.ρe_tot), available_tracer_names..., ρatke_if_available...) - diffusion_blocks = if use_derivative(diffusion_flag) - ( - MatrixFields.unrolled_map( - name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow), - diffused_scalar_names, - )..., - MatrixFields.unrolled_map( - name -> (name, name) => similar(Y.c, TridiagonalRow), - diffused_scalar_names, - )..., - ( - is_in_Y(@name(c.ρq_tot)) ? - ( - (@name(c.ρe_tot), @name(c.ρq_tot)) => - similar(Y.c, TridiagonalRow), - ) : () - )..., - (@name(c.uₕ), @name(c.uₕ)) => - !isnothing(atmos.turbconv_model) || - diffuse_momentum(atmos.vert_diff) ? - similar(Y.c, TridiagonalRow) : FT(-1) * I, - ) - else - MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - (diffused_scalar_names..., @name(c.uₕ)), - ) - end +factorize_exact_jacobian!(A, Y, p, dtγ, t) = nothing - matrix = MatrixFields.FieldMatrix( - identity_blocks..., - advection_blocks..., - diffusion_blocks..., - ) +approximate_jacobian!(A, Y, p, dtγ, t) = nothing - names₁_group₁ = (@name(c.ρ), other_available_names...) - names₁_group₂ = (available_tracer_names..., ρatke_if_available...) - names₁_group₃ = (@name(c.ρe_tot),) - names₁ = (names₁_group₁..., names₁_group₂..., names₁_group₃...) +# This is called from a callback when always_update_exact_jacobian is false. +function update_exact_jacobian!(A, Y, p, dtγ, t) + FT = eltype(Y) + factorize_exact_jacobian!(A, Y, p, FT(dtγ), t) +end - alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) - alg = if use_derivative(diffusion_flag) - alg₁_subalg₂ = - is_in_Y(@name(c.ρq_tot)) ? - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve(names₁_group₂...) - ) : (;) - alg₁ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₁...; - alg₁_subalg₂..., - ) - MatrixFields.ApproximateBlockArrowheadIterativeSolve( - names₁...; - alg₁, - alg₂, - P_alg₁ = MatrixFields.MainDiagonalPreconditioner(), - n_iters = approximate_solve_iters, - ) - else - MatrixFields.BlockArrowheadSolve(names₁...; alg₂) +# This is passed to ClimaTimeSteppers.jl and called on each Newton iteration. +function Wfact!(A::ImplicitEquationJacobian, Y, p, dtγ, t) + FT = eltype(Y) + if always_update_exact_jacobian(A) + factorize_exact_jacobian!(A, Y, p, FT(dtγ), t) end - - return ImplicitEquationJacobian( - matrix, - MatrixFields.FieldMatrixSolver(alg, matrix, Y), - diffusion_flag, - topography_flag, - similar(Y), - similar(Y), - transform_flag, - Ref{FT}(), - ) + approximate_jacobian!(A, Y, p, FT(dtγ), t) end -# We only use A, but ClimaTimeSteppers.jl require us to -# pass jac_prototype and then call similar(jac_prototype) to -# obtain A. This is a workaround to avoid unnecessary allocations. -Base.similar(A::ImplicitEquationJacobian) = A - -# This method specifies how to solve the equation E'(Y) * ΔY = E(Y) for ΔY. -NVTX.@annotate function ldiv!( +# This is called by ClimaTimeSteppers.jl on each Newton iteration. +ldiv!( x::Fields.FieldVector, A::ImplicitEquationJacobian, b::Fields.FieldVector, -) - MatrixFields.field_matrix_solve!(A.solver, x, A.matrix, b) - if A.transform_flag - @. x *= A.dtγ_ref[] - end -end +) = invert_jacobian!(A, x, b) -# This method for ldiv! is called by Krylov.jl from inside ClimaTimeSteppers.jl. -# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a -# related issue that requires the same workaround. -NVTX.@annotate function ldiv!( +# This is called by Krylov.jl from inside ClimaTimeSteppers.jl. See +# https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a related +# issue that requires the same workaround. +function ldiv!( x::AbstractVector, A::ImplicitEquationJacobian, b::AbstractVector, ) - A.temp_b .= b - ldiv!(A.temp_x, A, A.temp_b) - x .= A.temp_x -end - -# This function is used by DiffEqBase.jl instead of ldiv!. -linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve! -_linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b) - -# This method specifies how to compute E'(Y), which is referred to as "Wfact" in -# DiffEqBase.jl. -NVTX.@annotate function Wfact!(A, Y, p, dtγ, t) - # Remove unnecessary values from p to avoid allocations in bycolumn. - p′ = (; - p.precomputed.ᶜspecific, - p.precomputed.ᶜK, - p.precomputed.ᶜts, - p.precomputed.ᶜp, - ( - p.atmos.precip_model isa Microphysics1Moment ? - (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;) - )..., - p.precomputed.ᶜh_tot, - ( - use_derivative(A.diffusion_flag) ? - (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;) - )..., - ( - use_derivative(A.diffusion_flag) && - p.atmos.turbconv_model isa AbstractEDMF ? - (; p.precomputed.ᶜtke⁰, p.precomputed.ᶜmixing_length) : (;) - )..., - ( - use_derivative(A.diffusion_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; p.precomputed.ᶜρa⁰) : (;) - )..., - p.core.ᶜΦ, - p.core.ᶠgradᵥ_ᶜΦ, - p.core.ᶜρ_ref, - p.core.ᶜp_ref, - p.scratch.ᶜtemp_scalar, - p.scratch.∂ᶜK_∂ᶜuₕ, - p.scratch.∂ᶜK_∂ᶠu₃, - p.scratch.ᶠp_grad_matrix, - p.scratch.ᶜadvection_matrix, - p.scratch.ᶜdiffusion_h_matrix, - p.scratch.ᶜdiffusion_u_matrix, - p.dt, - p.params, - p.atmos, - ( - p.atmos.rayleigh_sponge isa RayleighSponge ? - (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;) - )..., - ) - - # Convert dtγ from a Float64 to an FT. - FT = Spaces.undertype(axes(Y.c)) - dtγ′ = FT(dtγ) - - A.dtγ_ref[] = dtγ′ - Fields.bycolumn(axes(Y.c)) do colidx - update_implicit_equation_jacobian!(A, Y, p′, dtγ′, colidx) - end -end - -function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) - (; matrix, diffusion_flag, topography_flag) = A - (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p - (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p - (; ᶜdiffusion_h_matrix, ᶜdiffusion_u_matrix, params) = p - - FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) - one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' - - cv_d = FT(CAP.cv_d(params)) - Δcv_v = FT(CAP.cv_v(params)) - cv_d - T_tri = FT(CAP.T_triple(params)) - e_int_v0 = T_tri * Δcv_v - FT(CAP.e_int_v0(params)) - thermo_params = CAP.thermodynamics_params(params) - - ᶜρ = Y.c.ρ - ᶜuₕ = Y.c.uₕ - ᶠu₃ = Y.f.u₃ - ᶜJ = Fields.local_geometry_field(Y.c).J - ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ - ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ - - ᶜkappa_m = p.ᶜtemp_scalar - @. ᶜkappa_m[colidx] = - TD.gas_constant_air(thermo_params, ᶜts[colidx]) / - TD.cv_m(thermo_params, ᶜts[colidx]) - - if use_derivative(topography_flag) - @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow( - adjoint(CTh(ᶜuₕ[colidx])) + - adjoint(ᶜinterp(ᶠu₃[colidx])) * g³ʰ(ᶜgⁱʲ[colidx]), - ) - else - @. ∂ᶜK_∂ᶜuₕ[colidx] = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ[colidx]))) - end - @. ∂ᶜK_∂ᶠu₃[colidx] = - ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃[colidx]))) + - DiagonalMatrixRow(adjoint(CT3(ᶜuₕ[colidx]))) ⋅ ᶜinterp_matrix() - - @. ᶠp_grad_matrix[colidx] = - DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ[colidx])) ⋅ ᶠgradᵥ_matrix() - - @. ᶜadvection_matrix[colidx] = - -(ᶜadvdivᵥ_matrix()) ⋅ - DiagonalMatrixRow(ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx])) - - if use_derivative(topography_flag) - ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] - @. ∂ᶜρ_err_∂ᶜuₕ[colidx] = - dtγ * ᶜadvection_matrix[colidx] ⋅ - ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ - DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx])) - end - ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] - @. ∂ᶜρ_err_∂ᶠu₃[colidx] = - dtγ * ᶜadvection_matrix[colidx] ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ[colidx])) - - tracer_info = ( - (@name(c.ρe_tot), @name(ᶜh_tot)), - (@name(c.ρq_tot), @name(ᶜspecific.q_tot)), - ) - MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name) - MatrixFields.has_field(Y, ρχ_name) || return - ᶜχ = MatrixFields.get_field(p, χ_name) - if use_derivative(topography_flag) - ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] - end - ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] - use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] = - dtγ * ᶜadvection_matrix[colidx] ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅ - ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅ - DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx])) - @. ∂ᶜρχ_err_∂ᶠu₃[colidx] = - dtγ * ᶜadvection_matrix[colidx] ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * g³³(ᶠgⁱʲ[colidx])) - end - - ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] - ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] - @. ∂ᶠu₃_err_∂ᶜρ[colidx] = - dtγ * ( - ᶠp_grad_matrix[colidx] ⋅ DiagonalMatrixRow( - ᶜkappa_m[colidx] * (T_tri * cv_d - ᶜK[colidx] - ᶜΦ[colidx]), - ) + - DiagonalMatrixRow( - ( - ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) - - ᶠinterp(ᶜρ_ref[colidx]) * ᶠgradᵥ_ᶜΦ[colidx] - ) / abs2(ᶠinterp(ᶜρ[colidx])), - ) ⋅ ᶠinterp_matrix() - ) - @. ∂ᶠu₃_err_∂ᶜρe_tot[colidx] = - dtγ * ᶠp_grad_matrix[colidx] ⋅ DiagonalMatrixRow(ᶜkappa_m[colidx]) - if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] - @. ∂ᶠu₃_err_∂ᶜρq_tot[colidx] = - dtγ * ᶠp_grad_matrix[colidx] ⋅ - DiagonalMatrixRow(ᶜkappa_m[colidx] * e_int_v0) - end - - ∂ᶠ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_∂ᶜuₕ[colidx] = - dtγ * ᶠp_grad_matrix[colidx] ⋅ - DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ ∂ᶜK_∂ᶜuₕ[colidx] - if p.atmos.rayleigh_sponge isa RayleighSponge - @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = - dtγ * ( - ᶠp_grad_matrix[colidx] ⋅ - DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ - ∂ᶜK_∂ᶠu₃[colidx] + - DiagonalMatrixRow(-p.ᶠβ_rayleigh_w[colidx] * (one_C3xACT3,)) - ) - (I_u₃,) - else - @. ∂ᶠu₃_err_∂ᶠu₃[colidx] = - dtγ * ᶠp_grad_matrix[colidx] ⋅ - DiagonalMatrixRow(-(ᶜkappa_m[colidx]) * ᶜρ[colidx]) ⋅ - ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,) - end - - if use_derivative(diffusion_flag) - (; ᶜK_h, ᶜK_u) = p - @. ᶜdiffusion_h_matrix[colidx] = - ᶜadvdivᵥ_matrix() ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(ᶜK_h[colidx])) ⋅ - ᶠgradᵥ_matrix() - if ( - MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) || - !isnothing(p.atmos.turbconv_model) || - diffuse_momentum(p.atmos.vert_diff) - ) - @. ᶜdiffusion_u_matrix[colidx] = - ᶜadvdivᵥ_matrix() ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(ᶜK_u[colidx])) ⋅ - ᶠgradᵥ_matrix() - end - - ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] - ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] - @. ∂ᶜρe_tot_err_∂ᶜρ[colidx] = - dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m[colidx]) * ᶜspecific.e_tot[colidx] - - ᶜkappa_m[colidx] * e_int_v0 * ᶜspecific.q_tot[colidx] - ) / ᶜρ[colidx], - ) - @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] = - dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ - DiagonalMatrixRow((1 + ᶜkappa_m[colidx]) / ᶜρ[colidx]) - (I,) - if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] - @. ∂ᶜρe_tot_err_∂ᶜρq_tot[colidx] = - dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ - DiagonalMatrixRow(ᶜkappa_m[colidx] * e_int_v0 / ᶜρ[colidx]) - end - - tracer_info = ( - (@name(c.ρq_tot), @name(q_tot)), - (@name(c.ρq_liq), @name(q_liq)), - (@name(c.ρq_ice), @name(q_ice)), - (@name(c.ρq_rai), @name(q_rai)), - (@name(c.ρq_sno), @name(q_sno)), - ) - MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name) - MatrixFields.has_field(Y, ρq_name) || return - ᶜq = MatrixFields.get_field(ᶜspecific, q_name) - ∂ᶜρq_err_∂ᶜρ = matrix[ρq_name, @name(c.ρ)] - ∂ᶜρq_err_∂ᶜρq = matrix[ρq_name, ρq_name] - @. ∂ᶜρq_err_∂ᶜρ[colidx] = - dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ - DiagonalMatrixRow(-(ᶜq[colidx]) / ᶜρ[colidx]) - @. ∂ᶜρq_err_∂ᶜρq[colidx] = - dtγ * ᶜdiffusion_h_matrix[colidx] ⋅ - DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,) - end - - if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) - turbconv_params = CAP.turbconv_params(params) - c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜtke⁰, ᶜmixing_length, dt) = p - ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ - ᶜρatke⁰ = Y.c.sgs⁰.ρatke - - dissipation_rate(tke⁰, mixing_length) = - tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : 1 / dt - ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = - tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : 0 - - ᶜdissipation_matrix_diagonal = p.ᶜtemp_scalar - @. ᶜdissipation_matrix_diagonal[colidx] = - ᶜρatke⁰[colidx] * - ∂dissipation_rate_∂tke⁰(ᶜtke⁰[colidx], ᶜmixing_length[colidx]) - - ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] - ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = - matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] - @. ∂ᶜρatke⁰_err_∂ᶜρ[colidx] = - dtγ * ( - ᶜdiffusion_u_matrix[colidx] - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal[colidx]) - ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰[colidx]) / ᶜρa⁰[colidx]) - @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰[colidx] = - dtγ * ( - ( - ᶜdiffusion_u_matrix[colidx] - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal[colidx]) - ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰[colidx]) - - DiagonalMatrixRow( - dissipation_rate(ᶜtke⁰[colidx], ᶜmixing_length[colidx]), - ) - ) - (I,) - end - - if ( - !isnothing(p.atmos.turbconv_model) || - diffuse_momentum(p.atmos.vert_diff) - ) - ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)] - @. ∂ᶜuₕ_err_∂ᶜuₕ[colidx] = - dtγ * DiagonalMatrixRow(1 / ᶜρ[colidx]) ⋅ - ᶜdiffusion_u_matrix[colidx] - (I,) - end - - ᶠlg = Fields.local_geometry_field(Y.f) - precip_info = - ((@name(c.ρq_rai), @name(ᶜwᵣ)), (@name(c.ρq_sno), @name(ᶜwₛ))) - MatrixFields.unrolled_foreach(precip_info) do (ρqₚ_name, wₚ_name) - MatrixFields.has_field(Y, ρqₚ_name) || return - ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name] - ᶜwₚ = MatrixFields.get_field(p, wₚ_name) - @. ∂ᶜρqₚ_err_∂ᶜρqₚ[colidx] += - dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - CT3(unit_basis_vector_data(CT3, ᶠlg[colidx])) * - ᶠwinterp(ᶜJ[colidx], ᶜρ[colidx]), - ) ⋅ ᶠright_bias_matrix() ⋅ - DiagonalMatrixRow(-(ᶜwₚ[colidx]) / ᶜρ[colidx]) - end - end + A.default_cache.temp_b .= b + ldiv!(A.default_cache.temp_x, A, A.default_cache.temp_b) + x .= A.default_cache.temp_x end diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index d77719cdedf..2599293c9fd 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -10,6 +10,7 @@ import LinearAlgebra import ClimaCore.Fields import ClimaTimeSteppers as CTS import DiffEqCallbacks as DECB +import DiffEqBase: KeywordArgSilent function get_atmos(config::AtmosConfig, params) (; turbconv_params) = params @@ -341,55 +342,27 @@ function get_surface_setup(parsed_args) return getproperty(SurfaceConditions, Symbol(parsed_args["surface_setup"]))() end -is_explicit_CTS_algo_type(alg_or_tableau) = - alg_or_tableau <: CTS.ERKAlgorithmName - -is_imex_CTS_algo_type(alg_or_tableau) = - alg_or_tableau <: CTS.IMEXARKAlgorithmName - -is_implicit_type(alg_or_tableau) = is_imex_CTS_algo_type(alg_or_tableau) - -is_imex_CTS_algo(::CTS.IMEXAlgorithm) = true -is_imex_CTS_algo(::SciMLBase.AbstractODEAlgorithm) = false - -is_implicit(ode_algo) = is_imex_CTS_algo(ode_algo) - -is_rosenbrock(::SciMLBase.AbstractODEAlgorithm) = false -use_transform(ode_algo) = - !(is_imex_CTS_algo(ode_algo) || is_rosenbrock(ode_algo)) - -additional_integrator_kwargs(::SciMLBase.AbstractODEAlgorithm) = (; - adaptive = false, - progress = isinteractive(), - progress_steps = isinteractive() ? 1 : 1000, -) -import DiffEqBase -additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; - kwargshandle = DiffEqBase.KeywordArgSilent, # allow custom kwargs - adjustfinal = true, - # TODO: enable progress bars in ClimaTimeSteppers -) - -is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false -is_cts_algo(::CTS.DistributedODEAlgorithm) = true - -function jac_kwargs(ode_algo, Y, atmos, parsed_args) - if is_implicit(ode_algo) - A = ImplicitEquationJacobian( - Y, - atmos; - approximate_solve_iters = parsed_args["approximate_linear_solve_iters"], - transform_flag = use_transform(ode_algo), - ) - if use_transform(ode_algo) - return (; jac_prototype = A, Wfact_t = Wfact!) +jac_kwargs(ode_algo, Y, p, parsed_args) = + if ode_algo isa CTS.IMEXAlgorithm + use_exact_jacobian = parsed_args["use_exact_jacobian"] + always_update_exact_jacobian = + parsed_args["n_steps_update_exact_jacobian"] == 0 + approximate_solve_iters = parsed_args["approximate_linear_solve_iters"] + A = if parsed_args["debug_jacobian_approximation"] + DebugJacobian( + ExactJacobian(Y, p; always_update_exact_jacobian), + ApproxJacobian(Y, p; approximate_solve_iters); + use_exact_jacobian, + ) else - return (; jac_prototype = A, Wfact = Wfact!) + use_exact_jacobian ? + ExactJacobian(Y, p; always_update_exact_jacobian) : + ApproxJacobian(Y, p; approximate_solve_iters) end + (; jac_prototype = A, Wfact = Wfact!) else - return NamedTuple() + (;) end -end #= ode_configuration(Y, parsed_args) @@ -398,14 +371,13 @@ Returns the ode algorithm =# function ode_configuration(::Type{FT}, parsed_args) where {FT} ode_name = parsed_args["ode_algo"] - alg_or_tableau = getproperty(CTS, Symbol(ode_name)) - @info "Using ODE config: `$alg_or_tableau`" - - if is_explicit_CTS_algo_type(alg_or_tableau) - return CTS.ExplicitAlgorithm(alg_or_tableau()) - elseif !is_implicit_type(alg_or_tableau) - return alg_or_tableau() - elseif is_imex_CTS_algo_type(alg_or_tableau) + ode_algo_name = getproperty(CTS, Symbol(ode_name)) + @info "Using ODE config: `$ode_algo_name`" + + if ode_algo_name <: CTS.ERKAlgorithmName + return CTS.ExplicitAlgorithm(ode_algo_name()) + else + @assert ode_algo_name <: CTS.IMEXARKAlgorithmName newtons_method = CTS.NewtonsMethod(; max_iters = parsed_args["max_newton_iters_ode"], krylov_method = if parsed_args["use_krylov_method"] @@ -434,9 +406,7 @@ function ode_configuration(::Type{FT}, parsed_args) where {FT} nothing end, ) - return CTS.IMEXAlgorithm(alg_or_tableau(), newtons_method) - else - return alg_or_tableau(; linsolve = linsolve!) + return CTS.IMEXAlgorithm(ode_algo_name(), newtons_method) end end @@ -617,30 +587,26 @@ function get_diagnostics(parsed_args, atmos_model, cspace) end function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) - (; atmos, dt) = p + (; dt) = p dt_save_to_sol = time_to_seconds(parsed_args["dt_save_to_sol"]) s = @timed_str begin func = if parsed_args["split_ode"] implicit_func = SciMLBase.ODEFunction( implicit_tendency!; - jac_kwargs(ode_algo, Y, atmos, parsed_args)..., + jac_kwargs(ode_algo, Y, p, parsed_args)..., tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), ) - if is_cts_algo(ode_algo) - CTS.ClimaODEFunction(; - T_lim! = limited_tendency!, - T_exp! = remaining_tendency!, - T_imp! = implicit_func, - # Can we just pass implicit_tendency! and jac_prototype etc.? - lim! = limiters_func!, - dss!, - post_explicit! = set_precomputed_quantities!, - post_implicit! = set_precomputed_quantities!, - ) - else - SciMLBase.SplitFunction(implicit_func, remaining_tendency!) - end + CTS.ClimaODEFunction(; + T_lim! = limited_tendency!, + T_exp! = remaining_tendency!, + T_imp! = implicit_func, + # Can we just pass implicit_tendency! and jac_prototype etc.? + lim! = limiters_func!, + dss!, + post_explicit! = set_precomputed_quantities!, + post_implicit! = set_precomputed_quantities!, + ) else remaining_tendency! # should be total_tendency! end @@ -656,7 +622,13 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) end # ensure that tspan[2] is always saved @info "dt_save_to_sol: $dt_save_to_sol, length(saveat): $(length(saveat))" args = (problem, ode_algo) - kwargs = (; saveat, callback, dt, additional_integrator_kwargs(ode_algo)...) + kwargs = (; + saveat, + callback, + dt, + kwargshandle = KeywordArgSilent, # allow custom kwargs + adjustfinal = true, + ) return (args, kwargs) end diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl index b15bd0444c2..b8a15eb37c8 100644 --- a/src/utils/abbreviations.jl +++ b/src/utils/abbreviations.jl @@ -1,4 +1,5 @@ -using ClimaCore: Geometry, Operators, MatrixFields +using ClimaCore: Geometry, Operators +using ClimaCore.MatrixFields # Alternatively, we could use Vec₁₂₃, Vec³, etc., if that is more readable. const C1 = Geometry.Covariant1Vector diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 69626e83ca2..109f346b0f2 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -62,9 +62,9 @@ function as (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2`. """ sgs_weight_function(a, a_half) = - if a < 0 + if a <= 0 # autodiff generates NaNs when a is 0 zero(a) - elseif a > 1 + elseif a > min(1, 42 * a_half) # autodiff generates NaNs when a is large one(a) else (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2 diff --git a/test/Project.toml b/test/Project.toml index 05c4f7dbeb4..121eb55c8c5 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,7 +25,6 @@ DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Glob = "c27321d9-0574-5035-807b-f59d2c89b15c" Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" @@ -74,7 +73,6 @@ ClimaCoreTempestRemap = "0.3" ClimaCoreVTK = "0.7" Dierckx = "0.5" DiffEqNoiseProcess = "5" -ForwardDiff = "0.10" Glob = "1" JSON = "0.21" NCRegressionTests = "0.2"