From 1dc4033e946ffebc7c7d1dad55a473e6d37313fe Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Tue, 27 Feb 2024 17:02:24 -0800 Subject: [PATCH] Add an interface for debugging the implicit solver Jacobian --- Project.toml | 4 + config/default_configs/default_config.yml | 9 + ...prognostic_edmfx_bomex_implicit_column.yml | 2 +- docs/Manifest.toml | 2 +- examples/Manifest.toml | 2 +- perf/Manifest.toml | 2 +- src/ClimaAtmos.jl | 4 + src/cache/precomputed_quantities.jl | 175 ++-- .../prognostic_edmf_precomputed_quantities.jl | 57 -- src/callbacks/callbacks.jl | 18 + src/callbacks/get_callbacks.jl | 13 + .../microphysics/microphysics_wrappers.jl | 10 +- .../edmfx_boundary_condition.jl | 10 +- src/prognostic_equations/edmfx_closures.jl | 20 +- src/prognostic_equations/edmfx_entr_detr.jl | 134 +-- .../implicit/approx_jacobian.jl | 684 ++++++++++++++ .../implicit/debug_jacobian.jl | 106 +++ .../implicit/dual_fixes.jl | 192 ++++ .../implicit/exact_jacobian.jl | 256 ++++++ .../implicit/implicit_solver.jl | 864 ++---------------- src/solver/type_getters.jl | 127 +-- src/utils/abbreviations.jl | 3 +- src/utils/utilities.jl | 13 + src/utils/variable_manipulations.jl | 4 +- test/Project.toml | 2 - 25 files changed, 1623 insertions(+), 1090 deletions(-) create mode 100644 src/prognostic_equations/implicit/approx_jacobian.jl create mode 100644 src/prognostic_equations/implicit/debug_jacobian.jl create mode 100644 src/prognostic_equations/implicit/dual_fixes.jl create mode 100644 src/prognostic_equations/implicit/exact_jacobian.jl diff --git a/Project.toml b/Project.toml index 93a32025418..1b65ee6020a 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" ArtifactWrappers = "a14bc488-3040-4b00-9dc1-f6467924858a" Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" AtmosphericProfilesLibrary = "86bc3604-9858-485a-bdbe-831ec50de11d" +BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884" ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f" @@ -20,6 +21,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" 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" @@ -48,6 +50,7 @@ ArgParse = "1" ArtifactWrappers = "0.2" Artifacts = "1" AtmosphericProfilesLibrary = "0.1.6" +BlockArrays = "0, 1" ClimaComms = "0.6.2" ClimaCore = "0.14.11" ClimaDiagnostics = "0.2" @@ -59,6 +62,7 @@ Dates = "1" DiffEqBase = "6" DocStringExtensions = "0.8, 0.9" FastGaussQuadrature = "0.4, 0.5, 1" +ForwardDiff = "0.10" Insolation = "0.9.2" 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 57c84b544be..fda7bd46a47 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -89,6 +89,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: true +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_approximate_jacobian: + help: "Whether to check the accuracy of the approximate Jacobian 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/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml b/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml index a7e29317ff6..faea36866cf 100644 --- a/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml +++ b/config/model_configs/prognostic_edmfx_bomex_implicit_column.yml @@ -7,7 +7,7 @@ turbconv: "prognostic_edmfx" implicit_diffusion: true implicit_sgs_advection: true approximate_linear_solve_iters: 2 -max_newton_iters_ode: 3 +max_newton_iters_ode: 10 edmfx_upwinding: first_order edmfx_entr_model: "Generalized" edmfx_detr_model: "Generalized" diff --git a/docs/Manifest.toml b/docs/Manifest.toml index fb46e9b9851..745e5d28898 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -297,7 +297,7 @@ weakdeps = ["SparseArrays"] ChainRulesCoreSparseArraysExt = "SparseArrays" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "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.27.4" diff --git a/examples/Manifest.toml b/examples/Manifest.toml index db4a9dd457d..52852a46934 100644 --- a/examples/Manifest.toml +++ b/examples/Manifest.toml @@ -321,7 +321,7 @@ version = "0.5.6" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "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.27.4" diff --git a/perf/Manifest.toml b/perf/Manifest.toml index 88c49f309ab..afa50e51c3c 100644 --- a/perf/Manifest.toml +++ b/perf/Manifest.toml @@ -332,7 +332,7 @@ version = "0.5.6" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LazyArtifacts", "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.27.4" diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index d0c6cbda5ba..1523c2aa9dd 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -48,6 +48,10 @@ include(joinpath("prognostic_equations", "zero_tendency.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", "implicit", "dual_fixes.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/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index d0b9b0b560c..6a2290d8ad3 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -4,6 +4,43 @@ import Thermodynamics as TD import ClimaCore: Spaces, Fields +function implicit_precomputed_quantities(Y, atmos) + FT = eltype(Y) + TST = thermo_state_type(atmos.moisture_model, FT) + n = n_mass_flux_subdomains(atmos.turbconv_model) + gs_quantities = (; + ᶜspecific = specific_gs.(Y.c), + ᶠuₕ³ = similar(Y.f, CT3{FT}), + ᶜu = similar(Y.c, C123{FT}), + ᶠu³ = similar(Y.f, CT3{FT}), + ᶜK = similar(Y.c, FT), + ᶜts = similar(Y.c, TST), + ᶜp = similar(Y.c, FT), + ᶜh_tot = similar(Y.c, FT), + ) + advective_sgs_quantities = + atmos.turbconv_model isa PrognosticEDMFX ? + (; + ᶜtke⁰ = similar(Y.c, FT), + ᶜρa⁰ = similar(Y.c, FT), + ᶜmse⁰ = similar(Y.c, FT), + ᶜq_tot⁰ = similar(Y.c, FT), + ᶠu₃⁰ = similar(Y.f, C3{FT}), + ᶜu⁰ = similar(Y.c, C123{FT}), + ᶠu³⁰ = similar(Y.f, CT3{FT}), + ᶜK⁰ = similar(Y.c, FT), + ᶜts⁰ = similar(Y.c, TST), + ᶜρ⁰ = similar(Y.c, FT), + ᶜuʲs = similar(Y.c, NTuple{n, C123{FT}}), + ᶠu³ʲs = similar(Y.f, NTuple{n, CT3{FT}}), + ᶜKʲs = similar(Y.c, NTuple{n, FT}), + ᶠKᵥʲs = similar(Y.f, NTuple{n, FT}), + ᶜtsʲs = similar(Y.c, NTuple{n, TST}), + ᶜρʲs = similar(Y.c, NTuple{n, FT}), + ) : (;) + return (; gs_quantities..., advective_sgs_quantities...) +end + """ precomputed_quantities(Y, atmos) @@ -42,16 +79,8 @@ function precomputed_quantities(Y, atmos) TST = thermo_state_type(atmos.moisture_model, FT) SCT = SurfaceConditions.surface_conditions_type(atmos, FT) cspace = axes(Y.c) - fspace = axes(Y.f) n = n_mass_flux_subdomains(atmos.turbconv_model) gs_quantities = (; - ᶜspecific = specific_gs.(Y.c), - ᶜu = similar(Y.c, C123{FT}), - ᶠu³ = similar(Y.f, CT3{FT}), - ᶜK = similar(Y.c, FT), - ᶜts = similar(Y.c, TST), - ᶜp = similar(Y.c, FT), - ᶜh_tot = similar(Y.c, FT), ᶜmixing_length = similar(Y.c, FT), ᶜlinear_buoygrad = similar(Y.c, FT), ᶜstrain_rate_norm = similar(Y.c, FT), @@ -76,41 +105,25 @@ function precomputed_quantities(Y, atmos) advective_sgs_quantities = atmos.turbconv_model isa PrognosticEDMFX ? (; - ᶜtke⁰ = similar(Y.c, FT), - ᶜρa⁰ = similar(Y.c, FT), - ᶠu₃⁰ = similar(Y.f, C3{FT}), - ᶜu⁰ = similar(Y.c, C123{FT}), - ᶠu³⁰ = similar(Y.f, CT3{FT}), - ᶜK⁰ = similar(Y.c, FT), - ᶜmse⁰ = similar(Y.c, FT), - ᶜq_tot⁰ = similar(Y.c, FT), - ᶜts⁰ = similar(Y.c, TST), - ᶜρ⁰ = similar(Y.c, FT), ᶜmixing_length_tuple = similar(Y.c, MixingLength{FT}), ᶜK_u = similar(Y.c, FT), ᶜK_h = similar(Y.c, FT), ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}), - ᶜuʲs = similar(Y.c, NTuple{n, C123{FT}}), - ᶠu³ʲs = similar(Y.f, NTuple{n, CT3{FT}}), - ᶜKʲs = similar(Y.c, NTuple{n, FT}), - ᶠKᵥʲs = similar(Y.f, NTuple{n, FT}), - ᶜtsʲs = similar(Y.c, NTuple{n, TST}), bdmr_l = similar(Y.c, BidiagonalMatrixRow{FT}), bdmr_r = similar(Y.c, BidiagonalMatrixRow{FT}), bdmr = similar(Y.c, BidiagonalMatrixRow{FT}), - ᶜρʲs = similar(Y.c, NTuple{n, FT}), ᶜentrʲs = similar(Y.c, NTuple{n, FT}), ᶜdetrʲs = similar(Y.c, NTuple{n, FT}), ᶠnh_pressure₃ʲs = similar(Y.f, NTuple{n, C3{FT}}), - ᶜgradᵥ_θ_virt⁰ = Fields.Field(C3{FT}, cspace), - ᶜgradᵥ_q_tot⁰ = Fields.Field(C3{FT}, cspace), - ᶜgradᵥ_θ_liq_ice⁰ = Fields.Field(C3{FT}, cspace), + ᶜgradᵥ_θ_virt⁰ = similar(Y.c, C3{FT}), + ᶜgradᵥ_q_tot⁰ = similar(Y.c, C3{FT}), + ᶜgradᵥ_θ_liq_ice⁰ = similar(Y.c, C3{FT}), precipitation_sgs_quantities..., ) : (;) sgs_quantities = (; - ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace), - ᶜgradᵥ_q_tot = Fields.Field(C3{FT}, cspace), - ᶜgradᵥ_θ_liq_ice = Fields.Field(C3{FT}, cspace), + ᶜgradᵥ_θ_virt = similar(Y.c, C3{FT}), + ᶜgradᵥ_q_tot = similar(Y.c, C3{FT}), + ᶜgradᵥ_θ_liq_ice = similar(Y.c, C3{FT}), ) diagnostic_sgs_quantities = @@ -155,6 +168,7 @@ function precomputed_quantities(Y, atmos) ᶜqₛ = similar(Y.c, FT), ) : (;) return (; + implicit_precomputed_quantities(Y, atmos)..., gs_quantities..., sgs_quantities..., advective_sgs_quantities..., @@ -326,14 +340,14 @@ end function eddy_diffusivity_coefficient( z::FT, z₀, - f_b::FT, - h::FT, + f_b, + h, uₐ, - C_E::FT, - Ri::FT, - Ri_a::FT, - Ri_c::FT, - κ::FT, + C_E, + Ri, + Ri_a, + Ri_c, + κ, ) where {FT} # Equations (17), (18) if z <= f_b * h @@ -412,12 +426,12 @@ end function compute_surface_layer_diffusivity( z::FT, - z₀::FT, - κ::FT, - C_E::FT, - Ri::FT, - Ri_a::FT, - Ri_c::FT, + z₀, + κ, + C_E, + Ri, + Ri_a, + Ri_c, norm_uₐ, ) where {FT} # Equations (19), (20) @@ -452,15 +466,19 @@ function instead of recomputing the value yourself. Otherwise, it will be difficult to ensure that the duplicated computations are consistent. """ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) + set_implicit_precomputed_quantities!(Y, p, t) + set_explicit_precomputed_quantities!(Y, p, t) +end +function set_implicit_precomputed_quantities!(Y, p, t) (; moisture_model, turbconv_model, vert_diff, precip_model, cloud_model) = p.atmos (; call_cloud_diagnostics_per_stage) = p.atmos thermo_params = CAP.thermodynamics_params(p.params) + turbconv_params = CAP.turbconv_params(p.params) n = n_mass_flux_subdomains(turbconv_model) thermo_args = (thermo_params, moisture_model) (; ᶜΦ) = p.core - (; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed - ᶠuₕ³ = p.scratch.ᶠtemp_CT3 + (; ᶜspecific, ᶠuₕ³, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed @. ᶜspecific = specific_gs(Y.c) set_ᶠuₕ³!(ᶠuₕ³, Y) @@ -488,6 +506,63 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) end @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) + + if turbconv_model isa PrognosticEDMFX + @assert !(moisture_model isa DryModel) + (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed + (; ᶜtke⁰, ᶜρa⁰, ᶜmse⁰, ᶜq_tot⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰) = + p.precomputed + + for j in 1:n + ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃ + ᶜmseʲ = Y.c.sgsʲs.:($j).mse + ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot + + ᶜuʲ = ᶜuʲs.:($j) + ᶠu³ʲ = ᶠu³ʲs.:($j) + ᶜKʲ = ᶜKʲs.:($j) + ᶠKᵥʲ = ᶠKᵥʲs.:($j) + ᶜtsʲ = ᶜtsʲs.:($j) + ᶜρʲ = ᶜρʲs.:($j) + + set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³) + @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 + @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) + @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) + end + + @. ᶜρa⁰ = ρa⁰(Y.c) + @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) + @. ᶜmse⁰ = divide_by_ρa( + Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρ * (ᶜh_tot - ᶜK), + Y.c.ρ, + turbconv_model, + ) + @. ᶜq_tot⁰ = divide_by_ρa( + Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρq_tot, + Y.c.ρ, + turbconv_model, + ) + set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) + set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) + # @. ᶜK⁰ += ᶜtke⁰ + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) + @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + end + + return nothing +end +function set_explicit_precomputed_quantities!(Y, p, t) + (; moisture_model, turbconv_model, vert_diff, precip_model, cloud_model) = + p.atmos + (; call_cloud_diagnostics_per_stage) = p.atmos + thermo_params = CAP.thermodynamics_params(p.params) + (; ᶠuₕ³, ᶜts, ᶜp) = p.precomputed if turbconv_model isa AbstractEDMF @. p.precomputed.ᶜgradᵥ_θ_virt = @@ -498,9 +573,6 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) end - (; ᶜh_tot) = p.precomputed - @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) - if !isnothing(p.sfc_setup) SurfaceConditions.update_surface_conditions!(Y, p, t) end @@ -517,12 +589,11 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) if turbconv_model isa PrognosticEDMFX set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t) - set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t) set_prognostic_edmf_precomputed_quantities_precipitation!( Y, p, - p.atmos.precip_model, + precip_model, ) end @@ -535,7 +606,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) Y, p, t, - p.atmos.precip_model, + precip_model, ) end @@ -544,7 +615,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) interior_uₕ = Fields.level(Y.c.uₕ, 1) ᶜΔz_surface = Fields.Δz_field(interior_uₕ) @. ᶜK_h = eddy_diffusivity_coefficient( - p.atmos.vert_diff.C_E, + vert_diff.C_E, norm(interior_uₕ), ᶜΔz_surface / 2, ᶜp, diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 8ac5caf32a1..3cc2ea55a43 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -5,50 +5,6 @@ import NVTX import Thermodynamics as TD import ClimaCore: Spaces, Fields -""" - set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t) - -Updates the edmf environment precomputed quantities stored in `p` for edmfx. -""" -NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( - Y, - p, - ᶠuₕ³, - t, -) - @assert !(p.atmos.moisture_model isa DryModel) - - thermo_params = CAP.thermodynamics_params(p.params) - (; turbconv_model) = p.atmos - (; ᶜΦ,) = p.core - (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed - (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = - p.precomputed - - @. ᶜρa⁰ = ρa⁰(Y.c) - @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) - @. ᶜmse⁰ = divide_by_ρa( - Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_tot⁰ = divide_by_ρa( - Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_tot, - Y.c.ρ, - turbconv_model, - ) - set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) - set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) - # @. ᶜK⁰ += ᶜtke⁰ - @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) - return nothing -end - """ set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t) @@ -62,8 +18,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc! t, ) (; moisture_model, turbconv_model) = p.atmos - #EDMFX BCs only support total energy as state variable - @assert !(moisture_model isa DryModel) FT = Spaces.undertype(axes(Y.c)) n = n_mass_flux_subdomains(turbconv_model) @@ -78,21 +32,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc! (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions for j in 1:n - ᶜuʲ = ᶜuʲs.:($j) - ᶠu³ʲ = ᶠu³ʲs.:($j) - ᶜKʲ = ᶜKʲs.:($j) - ᶠKᵥʲ = ᶠKᵥʲs.:($j) - ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃ ᶜtsʲ = ᶜtsʲs.:($j) - ᶜρʲ = ᶜρʲs.:($j) ᶜmseʲ = Y.c.sgsʲs.:($j).mse ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot - set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³) - @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 - @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) - @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) - # EDMFX boundary condition: # We need field_values everywhere because we are mixing diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index df4a2cd0c2e..8bd156e354b 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -413,6 +413,24 @@ function gc_func(integrator) return nothing end +function update_exact_jacobian!(integrator) + integrator.alg.name isa CTS.IMEXARKAlgorithmName || return + tableau_coefficients = LinearAlgebra.diag(integrator.alg.tableau.a_imp) + γs = unique(filter(!iszero, tableau_coefficients)) + 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 \ + tableau 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 0b49ae5c395..697f77f4384 100644 --- a/src/callbacks/get_callbacks.jl +++ b/src/callbacks/get_callbacks.jl @@ -256,6 +256,19 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start) (callbacks..., call_every_dt(cloud_fraction_model_callback!, dt_cf)) end + if ( + parsed_args["use_exact_jacobian"] || + parsed_args["debug_approximate_jacobian"] + ) && 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/parameterized_tendencies/microphysics/microphysics_wrappers.jl b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl index 5379909bc2a..5fb0eecbdb1 100644 --- a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl +++ b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl @@ -119,14 +119,8 @@ end Returns the total energy source term multiplier from precipitation formation for the 0-moment scheme """ -function e_tot_0M_precipitation_sources_helper( - thp, - ts, - Φ::FT, -) where {FT <: Real} - - λ::FT = TD.liquid_fraction(thp, ts) - +function e_tot_0M_precipitation_sources_helper(thp, ts, Φ) + λ = TD.liquid_fraction(t_prs, ts) return λ * Iₗ(thp, ts) + (1 - λ) * Iᵢ(thp, ts) + Φ end diff --git a/src/prognostic_equations/edmfx_boundary_condition.jl b/src/prognostic_equations/edmfx_boundary_condition.jl index 20810c9df9a..5d781e19002 100644 --- a/src/prognostic_equations/edmfx_boundary_condition.jl +++ b/src/prognostic_equations/edmfx_boundary_condition.jl @@ -4,9 +4,9 @@ function sgs_scalar_first_interior_bc( ᶜz_int::FT, - ᶜρ_int::FT, - ᶜaʲ_int::FT, - ᶜscalar_int::FT, + ᶜρ_int, + ᶜaʲ_int, + ᶜscalar_int, sfc_buoyancy_flux, sfc_ρ_flux_scalar, ustar, @@ -32,9 +32,9 @@ end function get_first_interior_variance( flux, - ustar::FT, + ustar, z::FT, - oblength::FT, + oblength, local_geometry, ) where {FT} c_star = -flux / ustar diff --git a/src/prognostic_equations/edmfx_closures.jl b/src/prognostic_equations/edmfx_closures.jl index 707b6ef28ff..59a43d049c7 100644 --- a/src/prognostic_equations/edmfx_closures.jl +++ b/src/prognostic_equations/edmfx_closures.jl @@ -165,17 +165,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 5202cbf6256..94c2a074fa2 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -19,17 +19,18 @@ function entrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜtke⁰::FT, + z_sfc, + ᶜp, + ᶜρ, + buoy_flux_surface, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜtke⁰, ::NoEntrainment, ) where {FT} return FT(0) @@ -80,17 +81,17 @@ end function entrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜtke⁰::FT, + z_sfc, + ᶜp, + ᶜρ, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜtke⁰, ::GeneralizedEntrainment, ) where {FT} turbconv_params = CAP.turbconv_params(params) @@ -171,21 +172,21 @@ end function detrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - ᶜρaʲ::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜentr::FT, - ᶜvert_div::FT, - ᶜmassflux_vert_div::FT, - ᶜtke⁰::FT, + z_sfc, + ᶜp, + ᶜρ, + ᶜρaʲ, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜentr, + ᶜvert_div, + ᶜmassflux_vert_div, + ᶜtke⁰, ::NoDetrainment, ) where {FT} return FT(0) @@ -239,21 +240,21 @@ end function detrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - ᶜρaʲ::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜentr::FT, - ᶜvert_div::FT, - ᶜmassflux_vert_div::FT, - ᶜtke⁰::FT, + z_sfc, + ᶜp, + ᶜρ, + ᶜρaʲ, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜentr, + ᶜvert_div, + ᶜmassflux_vert_div, + ᶜtke⁰, ::GeneralizedDetrainment, ) where {FT} turbconv_params = CAP.turbconv_params(params) @@ -289,20 +290,21 @@ end function detrainment( params, ᶜz::FT, - z_sfc::FT, - ᶜp::FT, - ᶜρ::FT, - ᶜρaʲ::FT, - ᶜaʲ::FT, - ᶜwʲ::FT, - ᶜRHʲ::FT, - ᶜbuoyʲ::FT, - ᶜw⁰::FT, - ᶜRH⁰::FT, - ᶜbuoy⁰::FT, - ᶜentr::FT, - ᶜvert_div::FT, - ᶜmassflux_vert_div::FT, + z_sfc, + ᶜp, + ᶜρ, + ᶜρaʲ, + ᶜaʲ, + ᶜwʲ, + ᶜRHʲ, + ᶜbuoyʲ, + ᶜw⁰, + ᶜRH⁰, + ᶜbuoy⁰, + ᶜentr, + ᶜvert_div, + ᶜmassflux_vert_div, + ᶜtke⁰, ::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..caf3ba2153d --- /dev/null +++ b/src/prognostic_equations/implicit/approx_jacobian.jl @@ -0,0 +1,684 @@ +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 + +- `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 +- `sgs_advection_flag::DerivativeFlag`: whether the derivative of the + subgrid-scale advection tendency with respect to the updraft quantities should + be computed or approximated as 0; must be either `UseDerivative()` or + `IgnoreDerivative()` instead of a `Bool` to ensure type-stability +- `approximate_solve_iters::Int`: number of iterations to take for the + approximate linear solve required when `diffusion_flag = UseDerivative()` +""" +@kwdef struct ApproxJacobian{F1, F2, F3} <: JacobianAlgorithm + diffusion_flag::F1 = nothing + topography_flag::F2 = nothing + sgs_advection_flag::F3 = nothing + approximate_solve_iters::Int = 1 +end + +function jacobian_cache(alg::ApproxJacobian, Y, p) + diffusion_flag = if isnothing(alg.diffusion_flag) + p.atmos.diff_mode == Implicit() ? UseDerivative() : IgnoreDerivative() + else + alg.diffusion_flag + end + topography_flag = if isnothing(alg.topography_flag) + has_topography(axes(Y.c)) ? UseDerivative() : IgnoreDerivative() + else + alg.topography_flag + end + sgs_advection_flag = if isnothing(alg.sgs_advection_flag) + p.atmos.sgs_adv_mode == Implicit() ? UseDerivative() : IgnoreDerivative() + else + alg.sgs_advection_flag + end + + FT = Spaces.undertype(axes(Y.c)) + CTh = CTh_vector_type(axes(Y.c)) + + DiagonalRow = DiagonalMatrixRow{FT} + 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})')} + DiagonalRow_C3xACT3 = + DiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{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),) : () + sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : () + + 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) + + # 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.ρ), sfc_if_available...), + ) + + 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(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 + + sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX + @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 + sgs_scalar_names = ( + @name(c.sgsʲs.:(1).q_tot), + @name(c.sgsʲs.:(1).mse), + @name(c.sgsʲs.:(1).ρa) + ) + if use_derivative(sgs_advection_flag) + ( + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + sgs_scalar_names, + )..., + (@name(c.sgsʲs.:(1).mse), @name(c.ρ)) => + similar(Y.c, DiagonalRow), + (@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.c, DiagonalRow), + (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.c, TridiagonalRow), + (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)) => + similar(Y.c, TridiagonalRow), + (@name(f.sgsʲs.:(1).u₃), @name(c.ρ)) => + similar(Y.f, BidiagonalRow_C3), + (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.f, BidiagonalRow_C3), + (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)) => + similar(Y.f, BidiagonalRow_C3), + (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.f, TridiagonalRow_C3xACT3), + ) + else + ( + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + sgs_scalar_names, + )..., + (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => + !isnothing(atmos.rayleigh_sponge) ? + similar(Y.f, DiagonalRow_C3xACT3) : FT(-1) * I, + ) + end + else + () + end + + matrix = MatrixFields.FieldMatrix( + identity_blocks..., + sgs_advection_blocks..., + advection_blocks..., + diffusion_blocks..., + ) + + sgs_names_if_available = if atmos.turbconv_model isa PrognosticEDMFX + ( + @name(c.sgsʲs.:(1).q_tot), + @name(c.sgsʲs.:(1).mse), + @name(c.sgsʲs.:(1).ρa), + @name(f.sgsʲs.:(1).u₃), + ) + else + () + end + names₁_group₁ = (@name(c.ρ), sfc_if_available...) + names₁_group₂ = (available_tracer_names..., ρatke_if_available...) + names₁_group₃ = (@name(c.ρe_tot),) + names₁ = ( + names₁_group₁..., + sgs_names_if_available..., + names₁_group₂..., + names₁_group₃..., + ) + + alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) + alg = + if use_derivative(diffusion_flag) || use_derivative(sgs_advection_flag) + alg₁_subalg₂ = + if atmos.turbconv_model isa PrognosticEDMFX && + use_derivative(sgs_advection_flag) + diff_subalg = + use_derivative(diffusion_flag) ? + (; + alg₂ = MatrixFields.BlockLowerTriangularSolve( + names₁_group₂..., + ) + ) : (;) + (; + alg₂ = MatrixFields.BlockLowerTriangularSolve( + @name(c.sgsʲs.:(1).q_tot); + alg₂ = MatrixFields.BlockLowerTriangularSolve( + @name(c.sgsʲs.:(1).mse); + alg₂ = MatrixFields.BlockLowerTriangularSolve( + @name(c.sgsʲs.:(1).ρa), + @name(f.sgsʲs.:(1).u₃); + diff_subalg..., + ), + ), + ) + ) + else + is_in_Y(@name(c.ρq_tot)) ? + (; + alg₂ = MatrixFields.BlockLowerTriangularSolve( + names₁_group₂..., + ) + ) : (;) + end + 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 + + solver = MatrixFields.FieldMatrixSolver(alg, matrix, Y) + + return (; + diffusion_flag, + topography_flag, + sgs_advection_flag, + matrix, + solver, + ) +end + +always_update_exact_jacobian(::ApproxJacobian) = false + +factorize_exact_jacobian!(::ApproxJacobian, _, _, _, _, _) = nothing + +function approximate_jacobian!(::ApproxJacobian, cache, Y, p, dtγ, t) + (; diffusion_flag, topography_flag, sgs_advection_flag, matrix) = cache + (; params) = p + (; edmfx_upwinding) = p.atmos.numerics + (; ᶜΦ, ᶠgradᵥ_ᶜΦ) = p.core + (; ᶜspecific, ᶜK, ᶜts, ᶜp) = p.precomputed + (; + ᶜtemp_C3, + ∂ᶜK_∂ᶜuₕ, + ∂ᶜK_∂ᶠu₃, + ᶠp_grad_matrix, + ᶜadvection_matrix, + ᶜdiffusion_h_matrix, + ᶜdiffusion_u_matrix, + ᶠbidiagonal_matrix_ct3, + ᶠbidiagonal_matrix_ct3_2, + ᶠtridiagonal_matrix_c3, + ) = p.scratch + + 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 = + TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts) + + if use_derivative(topography_flag) + @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow( + adjoint(CTh(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), + ) + else + @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ))) + end + @. ∂ᶜK_∂ᶠu₃ = + ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + + DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix() + + @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() + + @. ᶜadvection_matrix = + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) + + if use_derivative(topography_flag) + ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] + @. ∂ᶜρ_err_∂ᶜuₕ = + dtγ * ᶜadvection_matrix ⋅ ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ + DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) + end + ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] + @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) + + 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ₕ = + dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅ + ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) + @. ∂ᶜρχ_err_∂ᶠu₃ = + dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ)) + end + + ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] + ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] + @. ∂ᶠu₃_err_∂ᶜρ = + dtγ * ( + ᶠp_grad_matrix ⋅ + DiagonalMatrixRow(ᶜkappa_m * (T_tri * cv_d - ᶜK - ᶜΦ)) + + DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / abs2(ᶠinterp(ᶜρ))) ⋅ + ᶠinterp_matrix() + ) + @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) + 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 = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m * 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ₕ = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃_err_∂ᶠu₃ = + dtγ * ( + ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ + ∂ᶜK_∂ᶠu₃ + DiagonalMatrixRow(-p.ᶠβ_rayleigh_w * (one_C3xACT3,)) + ) - (I_u₃,) + else + @. ∂ᶠu₃_err_∂ᶠu₃ = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ + ∂ᶜK_∂ᶠu₃ - (I_u₃,) + end + + if use_derivative(diffusion_flag) + (; ᶜK_h, ᶜK_u) = p + @. ᶜdiffusion_h_matrix = + ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅ + ᶠ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 = + ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠ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_∂ᶜρ = + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( + ( + -(1 + ᶜkappa_m) * ᶜspecific.e_tot - + ᶜkappa_m * e_int_v0 * ᶜspecific.q_tot + ) / ᶜρ, + ) + @. ∂ᶜρe_tot_err_∂ᶜρe_tot = + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) - + (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 = + dtγ * ᶜdiffusion_h_matrix ⋅ + DiagonalMatrixRow(ᶜkappa_m * e_int_v0 / ᶜρ) + 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_∂ᶜρ = + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(ᶜq) / ᶜρ) + @. ∂ᶜρq_err_∂ᶜρq = + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) - (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 + + @inline dissipation_rate(tke⁰, mixing_length) = + tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : 1 / dt + @inline ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = + tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : + typeof(tke⁰)(0) + + ᶜdissipation_matrix_diagonal = p.ᶜtemp_scalar + @. ᶜdissipation_matrix_diagonal = + ᶜρatke⁰ * ∂dissipation_rate_∂tke⁰(ᶜtke⁰, ᶜmixing_length) + + ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] + ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = + matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] + @. ∂ᶜρatke⁰_err_∂ᶜρ = + dtγ * ( + ᶜdiffusion_u_matrix - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) + ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰) + @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = + dtγ * ( + ( + ᶜdiffusion_u_matrix - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) + ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - + DiagonalMatrixRow(dissipation_rate(ᶜtke⁰, ᶜmixing_length)) + ) - (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ₕ = + dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_u_matrix - (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) + ᶠtmp = p.ᶠtemp_CT3 + @. ᶠtmp = CT3(unit_basis_vector_data(CT3, ᶠlg)) * ᶠwinterp(ᶜJ, ᶜρ) + @. ∂ᶜρqₚ_err_∂ᶜρqₚ += + dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠtmp) ⋅ + ᶠright_bias_matrix() ⋅ DiagonalMatrixRow(-(ᶜwₚ) / ᶜρ) + end + end + + if p.atmos.turbconv_model isa PrognosticEDMFX + if use_derivative(sgs_advection_flag) + (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p + (; bdmr_l, bdmr_r, bdmr) = p + is_third_order = edmfx_upwinding == Val(:third_order) + ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 + ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(CT3{FT})), + bottom = Operators.SetValue(zero(CT3{FT})), + ) # Need to wrap ᶠupwind in this for well-defined boundaries. + UpwindMatrixRowType = + is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow + ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix + ᶠset_upwind_matrix_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), + bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), + ) # Need to wrap ᶠupwind_matrix in this for well-defined boundaries. + ᶜkappa_mʲ = p.ᶜtemp_scalar + @. ᶜkappa_mʲ = + TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / + TD.cv_m(thermo_params, ᶜtsʲs.:(1)) + ∂ᶜq_totʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ = + dtγ * ( + DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - + ᶜadvdivᵥ_matrix() ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) + ) - (I,) + + ∂ᶜmseʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶜmseʲ_err_∂ᶜq_totʲ = + dtγ * ( + -DiagonalMatrixRow( + adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * + ᶜgradᵥ_ᶠΦ * + Y.c.ρ * + ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * e_int_v0, + ) + ) + ∂ᶜmseʲ_err_∂ᶜρ = matrix[@name(c.sgsʲs.:(1).mse), @name(c.ρ)] + @. ∂ᶜmseʲ_err_∂ᶜρ = + dtγ * ( + -DiagonalMatrixRow( + adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * ᶜgradᵥ_ᶠΦ / ᶜρʲs.:(1), + ) + ) + ∂ᶜmseʲ_err_∂ᶜmseʲ = + matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).mse)] + @. ∂ᶜmseʲ_err_∂ᶜmseʲ = + dtγ * ( + DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - + ᶜadvdivᵥ_matrix() ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - + DiagonalMatrixRow( + adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * + ᶜgradᵥ_ᶠΦ * + Y.c.ρ * + ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + ) - (I,) + + ∂ᶜρaʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)] + @. ᶠbidiagonal_matrix_ct3 = + DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind( + ᶠu³ʲs.:(1), + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ), + ), + ) ⋅ ᶠwinterp_matrix(ᶜJ) ⋅ DiagonalMatrixRow( + ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp) * + e_int_v0, + ) + @. ᶠbidiagonal_matrix_ct3_2 = + DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow( + Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * + e_int_v0, + ) + + @. ∂ᶜρaʲ_err_∂ᶜq_totʲ = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) + ∂ᶜρaʲ_err_∂ᶜmseʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)] + @. ᶠbidiagonal_matrix_ct3 = + DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind( + ᶠu³ʲs.:(1), + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ), + ), + ) ⋅ ᶠwinterp_matrix(ᶜJ) ⋅ DiagonalMatrixRow( + ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + @. ᶠbidiagonal_matrix_ct3_2 = + DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow( + Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + @. ∂ᶜρaʲ_err_∂ᶜmseʲ = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) + ∂ᶜρaʲ_err_∂ᶜρaʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)] + @. ᶜadvection_matrix = + -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) + @. ∂ᶜρaʲ_err_∂ᶜρaʲ = + dtγ * ᶜadvection_matrix ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow(1 / ᶜρʲs.:(1)) - (I,) + + ∂ᶠu₃ʲ_err_∂ᶜρ = matrix[@name(f.sgsʲs.:(1).u₃), @name(c.ρ)] + @. ∂ᶠu₃ʲ_err_∂ᶜρ = + dtγ * DiagonalMatrixRow(ᶠgradᵥ_ᶜΦ / ᶠinterp(ᶜρʲs.:(1))) ⋅ + ᶠinterp_matrix() + ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = + dtγ * DiagonalMatrixRow( + ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / (ᶠinterp(ᶜρʲs.:(1)))^2, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp) * + e_int_v0, + ) + ∂ᶠu₃ʲ_err_∂ᶜmseʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)] + @. ∂ᶠu₃ʲ_err_∂ᶜmseʲ = + dtγ * DiagonalMatrixRow( + ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / (ᶠinterp(ᶜρʲs.:(1)))^2, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] + ᶜu₃ʲ = ᶜtemp_C3 + @. ᶜu₃ʲ = ᶜinterp(Y.f.sgsʲs.:(1).u₃) + + @. bdmr_l = convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix()) + @. bdmr_r = convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix()) + @. bdmr = ifelse(ᶜu₃ʲ.components.data.:1 > 0, bdmr_l, bdmr_r) + + @. ᶠtridiagonal_matrix_c3 = -(ᶠgradᵥ_matrix()) ⋅ bdmr + if p.atmos.rayleigh_sponge isa RayleighSponge + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * ( + ᶠtridiagonal_matrix_c3 ⋅ + DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - + DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,)) + ) - (I_u₃,) + else + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * ᶠtridiagonal_matrix_c3 ⋅ + DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,) + end + elseif p.atmos.rayleigh_sponge isa RayleighSponge + ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * -DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,)) - + (I_u₃,) + end + end +end + +invert_jacobian!(::ApproxJacobian, cache, x, b) = + MatrixFields.field_matrix_solve!(cache.solver, x, cache.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..6441498c246 --- /dev/null +++ b/src/prognostic_equations/implicit/debug_jacobian.jl @@ -0,0 +1,106 @@ +struct DebugJacobian{E <: ExactJacobian, A <: ApproxJacobian} <: + JacobianAlgorithm + exact_jacobian_mode::E + approx_jacobian_mode::A + use_exact_jacobian::Bool +end +function DebugJacobian( + approx_jacobian_mode; + use_exact_jacobian = true, + exact_jacobian_mode_kwargs..., +) + exact_jacobian_mode = ExactJacobian(; + preserve_unfactorized_jacobian = true, + exact_jacobian_mode_kwargs..., + ) + return DebugJacobian( + exact_jacobian_mode, + approx_jacobian_mode, + use_exact_jacobian, + ) +end + +jacobian_cache(alg::DebugJacobian, Y, p) = (; + exact_jacobian_cache = jacobian_cache(alg.exact_jacobian_mode, Y, p), + approx_jacobian_cache = jacobian_cache(alg.approx_jacobian_mode, Y, p), + similar_to_x = similar(Y), + x_error = similar(Y), + x_rel_error = similar(Y), + compute_error_ref = Ref{Bool}(), + t_ref = Ref{typeof(p.t_end)}(), +) + +always_update_exact_jacobian(alg::DebugJacobian) = + always_update_exact_jacobian(alg.exact_jacobian_mode) + +function factorize_exact_jacobian!(alg::DebugJacobian, cache, Y, p, dtγ, t) + factorize_exact_jacobian!( + alg.exact_jacobian_mode, + cache.exact_jacobian_cache, + Y, + p, + dtγ, + t, + ) + cache.compute_error_ref[] = true + cache.t_ref[] = t +end + +approximate_jacobian!(alg::DebugJacobian, cache, Y, p, dtγ, t) = + (!alg.use_exact_jacobian || cache.compute_error_ref[]) && + approximate_jacobian!( + alg.approx_jacobian_mode, + cache.approx_jacobian_cache, + Y, + p, + dtγ, + t, + ) + +function invert_jacobian!(alg::DebugJacobian, cache, x, b) + (; exact_jacobian_mode, approx_jacobian_mode, use_exact_jacobian) = alg + (; exact_jacobian_cache, approx_jacobian_cache) = cache + (; similar_to_x, x_error, x_rel_error, compute_error_ref, t_ref) = cache + if compute_error_ref[] + x_exact = use_exact_jacobian ? x : similar_to_x + x_approx = use_exact_jacobian ? similar_to_x : x + x_exact .= x_approx .= NaN # TODO: Remove this sanity check + invert_jacobian!(exact_jacobian_mode, exact_jacobian_cache, x_exact, b) + invert_jacobian!( + approx_jacobian_mode, + approx_jacobian_cache, + x_approx, + b, + ) + compute_error_ref[] = false + + @. x_error = abs(x_approx - x_exact) + @. x_rel_error = ifelse( + x_approx == x_exact == 0, + 0, + abs((x_approx - x_exact) / x_exact), + ) + + rms(f, x) = sqrt(mean(abs2 ∘ f, x)) + rms(x) = rms(identity, x) + + field_vector_rms(x) = rms(rms ∘ Fields.backing_array, Fields._values(x)) + field_vector_max(x) = + maximum(maximum ∘ Fields.backing_array, Fields._values(x)) + + rel_rms_error = field_vector_rms(x_error) / field_vector_rms(x_exact) + rel_max_error = field_vector_max(x_error) / field_vector_max(x_exact) + rms_rel_error = field_vector_rms(x_rel_error) + max_rel_error = field_vector_max(x_rel_error) + @info "Error of approximate implicit solve at t = $(t_ref[]):\n \ + relative RMS error = $(@sprintf("%.3e", rel_rms_error))\n \ + relative max error = $(@sprintf("%.3e", rel_max_error))\n \ + RMS relative error = $(@sprintf("%.3e", rms_rel_error))\n \ + max relative error = $(@sprintf("%.3e", max_rel_error))" + @assert rel_rms_error < 1 # TODO: Make this test more rigorous + elseif use_exact_jacobian + invert_jacobian!(exact_jacobian_mode, exact_jacobian_cache, x, b) + else + invert_jacobian!(approx_jacobian_mode, approx_jacobian_cache, x, b) + end +end diff --git a/src/prognostic_equations/implicit/dual_fixes.jl b/src/prognostic_equations/implicit/dual_fixes.jl new file mode 100644 index 00000000000..e8dfbb5b2b2 --- /dev/null +++ b/src/prognostic_equations/implicit/dual_fixes.jl @@ -0,0 +1,192 @@ +# TODO: Move these to Thermodynamics.jl +TD.air_density(param_set, T, p, q) = + TD.air_density(param_set, TD.promote_phase_partition(T, p, q)...) +TD.air_pressure(param_set, T, ρ, q) = + TD.air_pressure(param_set, TD.promote_phase_partition(T, ρ, q)...) +TD.internal_energy(param_set, T, q) = + TD.internal_energy(param_set, TD.promote_phase_partition(T, T, q)[2:end]...) +TD.internal_energy_sat(param_set, T, ρ, q_tot, type) = + TD.internal_energy_sat(param_set, promote(T, ρ, q_tot)..., type) +TD.PhasePartition_equil(param_set, T, ρ, q_tot, type) = + TD.PhasePartition_equil(param_set, promote(T, ρ, q_tot)..., type) +TD.PhasePartition_equil_given_p(param_set, T, p, q_tot, type) = + TD.PhasePartition_equil_given_p(param_set, promote(T, p, q_tot)..., type) +TD.specific_enthalpy_sat(param_set, T, ρ, q_tot, type) = + TD.specific_enthalpy_sat(param_set, promote(T, ρ, q_tot)..., type) +TD.saturation_vapor_pressure(param_set, T, LH_0, Δcp) = + TD.saturation_vapor_pressure(param_set, promote(T, LH_0, Δcp)...) + +# TODO: Move this to ClimaCore.jl +ClimaComms.device(fv::Fields.FieldVector) = + ClimaComms.device(first(Fields._values(fv))) + +# Ensure that we can broadcast between FieldVectors and other BlockVectors, but +# without going through the type-unstable BlockArrays broadcasting code. +# TODO: Replace the FieldVector broadcasting code in ClimaCore.jl with this +struct NewFieldVectorStyle <: Base.Broadcast.AbstractArrayStyle{1} end +Base.Broadcast.BroadcastStyle(::Type{<:Fields.FieldVector}) = + NewFieldVectorStyle() +Base.Broadcast.BroadcastStyle( + fs::NewFieldVectorStyle, + as::Base.Broadcast.DefaultArrayStyle{N}, +) where {N} = N in (0, 1) ? fs : as # prevents an ambiguity with Base.Broadcast +Base.Broadcast.BroadcastStyle( + fs::NewFieldVectorStyle, + as::Base.Broadcast.AbstractArrayStyle{N}, +) where {N} = N in (0, 1) ? fs : as +Base.@propagate_inbounds function BlockArrays.viewblock( + fv::Fields.FieldVector, + block::BlockArrays.Block{1}, +) + array = Fields.backing_array(Fields._values(fv)[block.n...]) + # vec(array) allocates; see https://github.com/JuliaLang/julia/issues/36313 + return Base.ReshapedArray(array, (length(array),), ()) +end + +first_field_vector(x::Fields.FieldVector, args...) = x +@inline first_field_vector(x::Base.Broadcast.Broadcasted, args...) = + first_field_vector(x.args..., args...) +@inline first_field_vector(_, args...) = first_field_vector(args...) +check_broadcast_axes(fv, x::Base.Broadcast.Broadcasted) = + check_broadcast_args_axes(fv, x.args...) +@inline check_broadcast_axes(fv, x) = + isempty(axes(x)) || # scalar + axes(x) == axes(fv) || # AbstractBlockVector + x isa AbstractVector && length(x) == length(fv) || # fallback + error("Mismatched broadcast axes: $(axes(x)) != $(axes(fv))") +check_broadcast_args_axes(_) = true +@inline check_broadcast_args_axes(fv, arg, args...) = + check_broadcast_axes(fv, arg) && check_broadcast_args_axes(fv, args...) + +block_view(x, _, _) = x +Base.@propagate_inbounds block_view( + x::BlockArrays.AbstractBlockVector, + ::Val{index}, + _, +) where {index} = BlockArrays.viewblock(x, BlockArrays.Block(index)) +Base.@propagate_inbounds block_view(x::AbstractVector, val, fv) = + block_view(BlockArrays.PseudoBlockVector(x, axes(fv)), val, fv) +Base.@propagate_inbounds block_view(x::Base.Broadcast.Broadcasted, val, fv) = + Base.Broadcast.broadcasted(x.f, block_view_args(val, fv, x.args...)...) +block_view_args(_, _) = () +Base.@propagate_inbounds block_view_args(val, fv, arg, args...) = + (block_view(arg, val, fv), block_view_args(val, fv, args...)...) + +# Bypass the method for materialize! defined in BlockArrays, which results in +# type instabilities and allocations. +Base.materialize!( + dest::Fields.FieldVector, + bc::Base.Broadcast.Broadcasted{NewFieldVectorStyle}, +) = materialize_field_vector!(dest, bc) +Base.materialize!( + dest::AbstractVector, + bc::Base.Broadcast.Broadcasted{NewFieldVectorStyle}, +) = materialize_field_vector!(dest, bc) +Base.materialize!( + dest::Fields.FieldVector, + bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}}, +) = materialize_field_vector!(dest, bc) +Base.materialize!( + dest::Fields.FieldVector, + bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{1}}, +) = materialize_field_vector!(dest, bc) +Base.materialize!( + dest::Fields.FieldVector, + bc::Base.Broadcast.Broadcasted{<:BlockArrays.AbstractBlockStyle{1}}, +) = materialize_field_vector!(dest, bc) # prevents an ambiguity with BlockArrays +function materialize_field_vector!(dest, bc) + fv = first_field_vector(dest, bc) + check_broadcast_args_axes(fv, dest, bc) + block_index_vals = ntuple(Val, BlockArrays.blocksize(fv, 1)) + MatrixFields.unrolled_foreach(block_index_vals) do val + dest_view = block_view(dest, val, fv) + bc_view = block_view(bc, val, fv) + @inbounds Base.materialize!(dest_view, bc_view) + end + return dest +end +Base.@propagate_inbounds function Base.getindex( + fv::Fields.FieldVector, + i::Integer, +) + @warn "scalar indexing into FieldVector (this can be very slow)" maxlog = 1 + getindex(fv, BlockArrays.findblockindex(axes(fv, 1), i)) +end +Base.@propagate_inbounds function Base.setindex!( + fv::Fields.FieldVector, + val, + i::Integer, +) + @warn "scalar indexing into FieldVector (this can be very slow)" maxlog = 1 + setindex!(fv, val, BlockArrays.findblockindex(axes(fv, 1), i)) +end + +# TODO: Move all of the following to ClimaCore.jl + +filtered_names(f::F, x) where {F} = filtered_names_at_name(f, x, @name()) +function filtered_names_at_name(f::F, x, name) where {F} + field = MatrixFields.get_field(x, name) + f(field) && return (name,) + internal_names = MatrixFields.top_level_names(field) + isempty(internal_names) && return () + tuples_of_names = MatrixFields.unrolled_map(internal_names) do internal_name + Base.@_inline_meta + child_name = MatrixFields.append_internal_name(name, internal_name) + filtered_names_at_name(f, x, child_name) + end + return MatrixFields.unrolled_flatten(tuples_of_names) +end +if hasfield(Method, :recursion_relation) + dont_limit = (args...) -> true + for m in methods(filtered_names_at_name) + m.recursion_relation = dont_limit + end +end + +scalar_field_names(fv) = + filtered_names(x -> x isa Fields.Field && eltype(x) == eltype(fv), fv) + +function scalar_level_iterator(field_vector) + Iterators.flatmap(scalar_field_names(field_vector)) do name + field = MatrixFields.get_field(field_vector, name) + if field isa Fields.SpectralElementField + (field,) + else + Iterators.map(1:Spaces.nlevels(axes(field))) do v + Fields.level(field, v - 1 + Operators.left_idx(axes(field))) + end + end + end +end + +function column_iterator(field_vectors...) + n_field_vectors = length(field_vectors) + @assert n_field_vectors >= 1 + example_field = field_vectors[1].c + + horz_space = Spaces.horizontal_space(axes(example_field)) + qs = 1:Quadratures.degrees_of_freedom(Spaces.quadrature_style(horz_space)) + hs = Spaces.eachslabindex(horz_space) + colidx_iterator = if Fields.field_values(example_field) isa DataLayouts.VIFH + Iterators.map(Iterators.product(qs, hs)) do (i, h) + Fields.ColumnIndex((i,), h) + end + else + @assert Fields.field_values(example_field) isa DataLayouts.VIJFH + Iterators.map(Iterators.product(qs, qs, hs)) do (i, j, h) + Fields.ColumnIndex((i, j), h) + end + end + return Iterators.map(colidx_iterator) do colidx + field_vector_columns = map(field_vectors) do fv + @assert all(in((:c, :f, :sfc)), propertynames(fv)) + has_sfc = hasproperty(fv, :sfc) + Fields.FieldVector(; + c = fv.c[colidx], + f = fv.f[colidx], + (has_sfc ? (; sfc = fv.sfc[colidx]) : (;))..., + ) + end + n_field_vectors == 1 ? field_vector_columns[1] : field_vector_columns + 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..e6add94528a --- /dev/null +++ b/src/prognostic_equations/implicit/exact_jacobian.jl @@ -0,0 +1,256 @@ +@kwdef struct ExactJacobian <: JacobianAlgorithm + always_update_exact_jacobian::Bool = true + preserve_unfactorized_jacobian::Bool = false # only true for DebugJacobian +end + +# function jacobian_cache(alg::ExactJacobian, Y, p) +# FT = eltype(Y) +# DA = ClimaComms.array_type(Y) +# Y_columns = column_iterator(Y) +# n_εs = length(first(Y_columns)) +# dual_type = ForwardDiff.Dual{ExactJacobian, FT, n_εs} + +# similar_with_dual_fields(named_tuple) = +# Fields._values(similar(Fields.FieldVector(; named_tuple...), dual_type)) + +# Y_dual = similar(Y, dual_type) +# Y_dual_copy = similar(Y_dual) +# Y_err_dual = similar(Y_dual) +# implicit_precomputed_dual = +# similar_with_dual_fields(implicit_precomputed_quantities(Y, p.atmos)) +# scratch_dual = similar_with_dual_fields(p.scratch) + +# column_vectors = map(_ -> DA{FT}(undef, n_εs), Y_columns) +# column_partials_vectors = +# map(_ -> DA{ForwardDiff.Partials{n_εs, FT}}(undef, n_εs), Y_columns) +# column_matrices = map(_ -> DA{FT}(I, n_εs, n_εs), Y_columns) +# column_factorized_matrices = +# alg.preserve_unfactorized_jacobian ? deepcopy(column_matrices) : +# column_matrices + +# return (; +# Y_dual, +# Y_dual_copy, +# Y_err_dual, +# implicit_precomputed_dual, +# scratch_dual, +# column_vectors, +# column_partials_vectors, +# column_matrices, +# column_factorized_matrices, +# ) +# end +function jacobian_cache(alg::ExactJacobian, Y, p) + FT = eltype(Y) + DA = ClimaComms.array_type(Y) + Y_columns = column_iterator(Y) + n_εs = length(first(Y_columns)) + dual_type = ForwardDiff.Dual{ExactJacobian, FT, 1} + + similar_with_dual_fields(named_tuple) = + Fields._values(similar(Fields.FieldVector(; named_tuple...), dual_type)) + + Y_dual = similar(Y, dual_type) + Y_dual_copy = similar(Y_dual) + Y_err_dual = similar(Y_dual) + implicit_precomputed_dual = + similar_with_dual_fields(implicit_precomputed_quantities(Y, p.atmos)) + scratch_dual = similar_with_dual_fields(p.scratch) + + column_vectors = map(_ -> DA{FT}(undef, n_εs), Y_columns) + column_partials_vectors = + map(_ -> DA{ForwardDiff.Partials{1, FT}}(undef, n_εs), Y_columns) + column_matrices = map(_ -> DA{FT}(I, n_εs, n_εs), Y_columns) + column_factorized_matrices = + alg.preserve_unfactorized_jacobian ? deepcopy(column_matrices) : + column_matrices + + return (; + Y_dual, + Y_dual_copy, + Y_err_dual, + implicit_precomputed_dual, + scratch_dual, + column_vectors, + column_partials_vectors, + column_matrices, + column_factorized_matrices, + ) +end + +always_update_exact_jacobian(alg::ExactJacobian) = + alg.always_update_exact_jacobian + +# function factorize_exact_jacobian!(alg::ExactJacobian, cache, Y, p, dtγ, t) +# (; preserve_unfactorized_jacobian) = alg +# (; Y_dual, Y_dual_copy, Y_err_dual) = cache +# (; implicit_precomputed_dual, scratch_dual) = cache + +# 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 +# (; p.precomputed..., implicit_precomputed_dual...) +# elseif cache_field_name == :scratch +# scratch_dual +# else +# getfield(p, cache_field_index) +# end +# end +# p_dual = AtmosCache(p_dual_args...) + +# FT = eltype(Y) +# Y_dual_scalar_levels = scalar_level_iterator(Y_dual) +# Y_err_dual_columns = column_iterator(Y_err_dual) +# n_εs = length(first(Y_err_dual_columns)) + +# Y_dual .= Y +# zeros_tuple = ntuple(_ -> 0, n_εs) +# for (ε_index, Y_dual_scalar_level) in enumerate(Y_dual_scalar_levels) +# ε_partials = Base.setindex(zeros_tuple, 1, ε_index) +# ε = ForwardDiff.Dual{ExactJacobian}(0, ε_partials...) +# parent(Y_dual_scalar_level) .+= ε +# end # add a unique infinitesimal ε to the scalar values on each level of Y +# Y_dual_copy .= Y_dual # need a copy because Y_dual gets changed at boundary +# set_implicit_precomputed_quantities!(Y_dual, p_dual, t) # changes Y_dual +# implicit_tendency!(Y_err_dual, Y_dual, p_dual, t) # the tendency Jacobian +# @. Y_err_dual = dtγ * Y_err_dual - Y_dual_copy # the residual Jacobian + +# for (col_index, Y_err_dual_column) in enumerate(Y_err_dual_columns) +# column_partials_vector = cache.column_partials_vectors[col_index] +# column_matrix = cache.column_matrices[col_index] +# column_factorized_matrix = cache.column_factorized_matrices[col_index] + +# column_partials_vector .= ForwardDiff.partials.(Y_err_dual_column) +# column_matrix .= reinterpret(reshape, FT, column_partials_vector)' +# if preserve_unfactorized_jacobian +# column_factorized_matrix .= column_matrix +# end +# ClimaComms.allowscalar(ClimaComms.device(Y)) do +# lu_factorize!(column_factorized_matrix) +# end +# end # TODO: Parallelize this loop +# end +function factorize_exact_jacobian!(alg::ExactJacobian, cache, Y, p, dtγ, t) + (; preserve_unfactorized_jacobian) = alg + (; Y_dual, Y_dual_copy, Y_err_dual) = cache + (; implicit_precomputed_dual, scratch_dual) = cache + + 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 + (; p.precomputed..., implicit_precomputed_dual...) + elseif cache_field_name == :scratch + scratch_dual + else + getfield(p, cache_field_index) + end + end + p_dual = AtmosCache(p_dual_args...) + + FT = eltype(Y) + Y_dual_scalar_levels = scalar_level_iterator(Y_dual) + Y_err_dual_columns = column_iterator(Y_err_dual) + n_εs = length(first(Y_err_dual_columns)) + ε = ForwardDiff.Dual{ExactJacobian}(0, 1) + + Y_dual .= Y + for (ε_index, Y_dual_scalar_level) in enumerate(Y_dual_scalar_levels) + parent(Y_dual_scalar_level) .+= ε # add ε to this level of Y_dual + + Y_dual_copy .= Y_dual # make a copy because Y_dual gets changed at boundary + set_implicit_precomputed_quantities!(Y_dual, p_dual, t) # changes Y_dual + implicit_tendency!(Y_err_dual, Y_dual, p_dual, t) # the tendency Jacobian + @. Y_err_dual = dtγ * Y_err_dual - Y_dual_copy # the residual Jacobian + + for (col_index, Y_err_dual_column) in enumerate(Y_err_dual_columns) + column_partials_vector = cache.column_partials_vectors[col_index] + column_matrix = cache.column_matrices[col_index] + column_partials_vector .= ForwardDiff.partials.(Y_err_dual_column) + column_matrix[:, ε_index] .= + reinterpret(reshape, FT, column_partials_vector) + end # move the partials from Y_err_dual_column into the cached matrices + + parent(Y_dual_scalar_level) .-= ε # remove ε from this level of Y_dual + end + for col_index in 1:length(Y_err_dual_columns) + column_matrix = cache.column_matrices[col_index] + column_factorized_matrix = cache.column_factorized_matrices[col_index] + if preserve_unfactorized_jacobian + column_factorized_matrix .= column_matrix + end + ClimaComms.allowscalar(ClimaComms.device(Y)) do + lu_factorize!(column_factorized_matrix) + end + end +end + +approximate_jacobian!(::ExactJacobian, _, _, _, _, _) = nothing + +invert_jacobian!(::ExactJacobian, cache, x, b) = + for (col_index, (x_column, b_column)) in enumerate(column_iterator(x, b)) + column_vector = cache.column_vectors[col_index] + column_factorized_matrix = cache.column_factorized_matrices[col_index] + + column_vector .= b_column + ClimaComms.allowscalar(ClimaComms.device(x)) do + lu_solve!(column_vector, column_factorized_matrix) + end + x_column .= column_vector + end # TODO: Parallelize this loop + +# 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` because not specializing on any tag +# overwrites the generic method for `Dual` in `ForwardDiff` and breaks +# precompilation, while specializing on the default tag `Nothing` causes the +# type piracy Aqua test to fail. +@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 + +function lu_factorize!(A) + n = size(A, 1) + @assert size(A) == (n, n) + @inbounds for k in 1:n + A_kk = A[k, k] + (A_kk != 0 && isfinite(A_kk)) || + error("Cannot get LU factors due to $A_kk on the matrix diagonal") + A_kk_inv = inv(A_kk) + for i in (k + 1):n + A[i, k] *= A_kk_inv + end + for j in (k + 1):n + for i in (k + 1):n + A[i, j] -= A[i, k] * A[k, j] + end + end + end +end +function lu_solve!(x, A) + n = size(A, 1) + @assert size(A) == (n, n) && size(x) == (n,) + @inbounds begin + for i in 2:n + Δx_i = zero(eltype(x)) + for j in 1:(i - 1) + Δx_i += A[i, j] * x[j] + end + x[i] -= Δx_i + end + x[n] /= A[n, n] + for i in (n - 1):-1:1 + Δx_i = zero(eltype(x)) + for j in (i + 1):n + Δx_i += A[i, j] * x[j] + end + x[i] -= Δx_i + x[i] /= A[i, i] + end + end +end diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index d8f721e9679..c818aeb0fd2 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -1,28 +1,25 @@ -import LinearAlgebra: I, Adjoint, ldiv! +import BlockArrays, ForwardDiff +import LinearAlgebra: I, Adjoint, LU, 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, sgs_advection_flag, transform_flag - ) + JacobianAlgorithm -A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the -implicit step with the state ``Y``. +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 4 methods: + - `jacobian_cache(alg::JacobianAlgorithm, Y, p)` + - `always_update_exact_jacobian(alg::JacobianAlgorithm)` + - `factorize_exact_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)` + - `approximate_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)` + - `invert_jacobian!(alg::JacobianAlgorithm, cache, x, b)` # Background When we use an implicit or split implicit-explicit (IMEX) timestepping scheme, -we end up with a nonlinear equation of the form ``E(Y) = 0``, where +we end up with a nonlinear equation of the form ``R(Y) = 0``, where ```math - E(Y) = Y_{imp}(Y) - Y = \\hat{Y} + Δt * T_{imp}(Y) - Y. + R(Y) = Y_{imp}(Y) - Y = \\hat{Y} + Δt * T_{imp}(Y) - Y. ``` In this expression, ``Y_{imp}(Y)`` denotes the state at some time ``t + Δt``. This can be expressed as the sum of ``\\hat{Y}``, the contribution from the @@ -38,827 +35,90 @@ divided into several sub-steps or "stages", where the duration of stage ``i`` is ``Δt * γ_i`` for some constant ``γ_i`` between 0 and 1. In order to solve this equation using Newton's method, we must specify the -derivative ``∂E/∂Y``. Since ``\\hat{Y}`` does not depend on ``Y`` (it is only a +derivative ``∂R/∂Y``. Since ``\\hat{Y}`` does not depend on ``Y`` (it is only a function of the state at or before time ``t``), this derivative is ```math - E'(Y) = Δt * T_{imp}'(Y) - I. + R'(Y) = Δt * T_{imp}'(Y) - I. ``` -In addition, we must specify how to divide ``E(Y)`` by this derivative, i.e., +In addition, we must specify how to divide ``R(Y)`` by this derivative, i.e., how to solve the linear equation ```math - E'(Y) * ΔY = E(Y). + R'(Y) * ΔY = R(Y). ``` Note: This equation comes from assuming that there is some ``ΔY`` such that -``E(Y - ΔY) = 0`` and making the first-order approximation +``R(Y - ΔY) = 0`` and making the first-order approximation ```math - E(Y - ΔY) \\approx E(Y) - E'(Y) * ΔY. + R(Y - ΔY) \\approx R(Y) - R'(Y) * ΔY. ``` After initializing ``Y`` to ``Y[0] = \\hat{Y}``, Newton's method executes the following steps: -- Compute the derivative ``E'(Y[0])``. -- Compute the implicit tendency ``T_{imp}(Y[0])`` and use it to get ``E(Y[0])``. -- Solve the linear equation ``E'(Y[0]) * ΔY[0] = E(Y[0])`` for ``ΔY[0]``. +- Compute the derivative ``R'(Y[0])``. +- Compute the implicit tendency ``T_{imp}(Y[0])`` and use it to get ``R(Y[0])``. +- Solve the linear equation ``R'(Y[0]) * ΔY[0] = R(Y[0])`` for ``ΔY[0]``. - Update ``Y`` to ``Y[1] = Y[0] - ΔY[0]``. If the number of Newton iterations is limited to 1, this new value of ``Y`` is taken to be the solution of the implicit equation. Otherwise, this sequence of steps is repeated, i.e., ``ΔY[1]`` is computed and used to update ``Y`` to ``Y[2] = Y[1] - ΔY[1]``, then ``ΔY[2]`` is computed and used to update ``Y`` to ``Y[3] = Y[2] - ΔY[2]``, and so on. The iterative process is terminated either -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 - `IgnoreDerivative()` 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 -- `sgs_advection_flag::DerivativeFlag`: whether the derivative of the - subgrid-scale advection tendency with respect to the updraft quantities 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 +when the residual ``R(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. """ -struct ImplicitEquationJacobian{ - M <: MatrixFields.FieldMatrix, - S <: MatrixFields.FieldMatrixSolver, - F1 <: DerivativeFlag, - F2 <: DerivativeFlag, - F3 <: 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 +abstract type JacobianAlgorithm end - # flags that determine how E'(Y) is approximated - diffusion_flag::F1 - topography_flag::F2 - sgs_advection_flag::F3 - - # 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 +struct ImplicitEquationJacobian{A <: JacobianAlgorithm, C} + alg::A + cache::C 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(), - sgs_advection_flag = atmos.sgs_adv_mode == Implicit() ? UseDerivative() : - IgnoreDerivative(), - transform_flag = false, -) - FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) - - DiagonalRow = DiagonalMatrixRow{FT} - 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})')} - DiagonalRow_C3xACT3 = - DiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{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),) : () - sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : () - - 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) - - # 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.ρ), sfc_if_available...), +function ImplicitEquationJacobian(alg, Y, p) + cache = (; + jacobian_cache(alg, Y, p)..., + x_krylov = similar(Y), + b_krylov = similar(Y), ) + return ImplicitEquationJacobian(alg, cache) +end - 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(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 - - sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX - @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 - sgs_scalar_names = ( - @name(c.sgsʲs.:(1).q_tot), - @name(c.sgsʲs.:(1).mse), - @name(c.sgsʲs.:(1).ρa) - ) - if use_derivative(sgs_advection_flag) - ( - MatrixFields.unrolled_map( - name -> (name, name) => similar(Y.c, TridiagonalRow), - sgs_scalar_names, - )..., - (@name(c.sgsʲs.:(1).mse), @name(c.ρ)) => - similar(Y.c, DiagonalRow), - (@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.c, DiagonalRow), - (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.c, TridiagonalRow), - (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)) => - similar(Y.c, TridiagonalRow), - (@name(f.sgsʲs.:(1).u₃), @name(c.ρ)) => - similar(Y.f, BidiagonalRow_C3), - (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.f, BidiagonalRow_C3), - (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)) => - similar(Y.f, BidiagonalRow_C3), - (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.f, TridiagonalRow_C3xACT3), - ) - else - ( - MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - sgs_scalar_names, - )..., - (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => - !isnothing(atmos.rayleigh_sponge) ? - similar(Y.f, DiagonalRow_C3xACT3) : FT(-1) * I, - ) - end - else - () - end +# 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 - matrix = MatrixFields.FieldMatrix( - identity_blocks..., - sgs_advection_blocks..., - advection_blocks..., - diffusion_blocks..., - ) +# This is called from a callback when always_update_exact_jacobian is false. +NVTX.@annotate function update_exact_jacobian!(A, Y, p, dtγ, t) + FT = eltype(Y) + factorize_exact_jacobian!(A.alg, A.cache, Y, p, FT(dtγ), t) +end - sgs_names_if_available = if atmos.turbconv_model isa PrognosticEDMFX - ( - @name(c.sgsʲs.:(1).q_tot), - @name(c.sgsʲs.:(1).mse), - @name(c.sgsʲs.:(1).ρa), - @name(f.sgsʲs.:(1).u₃), - ) - else - () +# This is passed to ClimaTimeSteppers.jl and called on each Newton iteration. +NVTX.@annotate function update_jacobian!(A, Y, p, dtγ, t) + FT = eltype(Y) + if always_update_exact_jacobian(A.alg) + factorize_exact_jacobian!(A.alg, A.cache, Y, p, FT(dtγ), t) end - names₁_group₁ = (@name(c.ρ), sfc_if_available...) - names₁_group₂ = (available_tracer_names..., ρatke_if_available...) - names₁_group₃ = (@name(c.ρe_tot),) - names₁ = ( - names₁_group₁..., - sgs_names_if_available..., - names₁_group₂..., - names₁_group₃..., - ) - - alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ)) - alg = - if use_derivative(diffusion_flag) || use_derivative(sgs_advection_flag) - alg₁_subalg₂ = - if atmos.turbconv_model isa PrognosticEDMFX && - use_derivative(sgs_advection_flag) - diff_subalg = - use_derivative(diffusion_flag) ? - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₂..., - ) - ) : (;) - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.sgsʲs.:(1).q_tot); - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.sgsʲs.:(1).mse); - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.sgsʲs.:(1).ρa), - @name(f.sgsʲs.:(1).u₃); - diff_subalg..., - ), - ), - ) - ) - else - is_in_Y(@name(c.ρq_tot)) ? - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₂..., - ) - ) : (;) - end - 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 ImplicitEquationJacobian( - matrix, - MatrixFields.FieldMatrixSolver(alg, matrix, Y), - diffusion_flag, - topography_flag, - sgs_advection_flag, - similar(Y), - similar(Y), - transform_flag, - Ref{FT}(), - ) + approximate_jacobian!(A.alg, A.cache, 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. +NVTX.@annotate 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.alg, A.cache, 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⁰) : (;) - )..., - ( - use_derivative(A.sgs_advection_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; - p.core.ᶜgradᵥ_ᶠΦ, - p.precomputed.ᶜρʲs, - p.precomputed.ᶠu³ʲs, - p.precomputed.ᶜtsʲs, - p.precomputed.bdmr_l, - p.precomputed.bdmr_r, - p.precomputed.bdmr, - ) : (;) - )..., - p.core.ᶜΦ, - p.core.ᶠgradᵥ_ᶜΦ, - p.scratch.ᶜtemp_scalar, - p.scratch.ᶜtemp_C3, - p.scratch.ᶠtemp_CT3, - 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.scratch.ᶠbidiagonal_matrix_ct3, - p.scratch.ᶠbidiagonal_matrix_ct3_2, - p.scratch.ᶠtridiagonal_matrix_c3, - 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γ′ - update_implicit_equation_jacobian!(A, Y, p′, dtγ′) -end - -function update_implicit_equation_jacobian!(A, Y, p, dtγ) - (; matrix, diffusion_flag, sgs_advection_flag, topography_flag) = A - (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ) = p - (; - ᶜtemp_C3, - ∂ᶜK_∂ᶜuₕ, - ∂ᶜK_∂ᶠu₃, - ᶠp_grad_matrix, - ᶜadvection_matrix, - ᶠbidiagonal_matrix_ct3, - ᶠbidiagonal_matrix_ct3_2, - ᶠtridiagonal_matrix_c3, - ) = p - (; ᶜdiffusion_h_matrix, ᶜdiffusion_u_matrix, params) = p - (; edmfx_upwinding) = p.atmos.numerics - - 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 = - TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts) - - if use_derivative(topography_flag) - @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow( - adjoint(CTh(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), - ) - else - @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ))) - end - @. ∂ᶜK_∂ᶠu₃ = - ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + - DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix() - - @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() - - @. ᶜadvection_matrix = - -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρ)) - - if use_derivative(topography_flag) - ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] - @. ∂ᶜρ_err_∂ᶜuₕ = - dtγ * ᶜadvection_matrix ⋅ ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ - DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) - end - ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] - @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - - 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ₕ = - dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅ - ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) - @. ∂ᶜρχ_err_∂ᶠu₃ = - dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ)) - end - - ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] - ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] - @. ∂ᶠu₃_err_∂ᶜρ = - dtγ * ( - ᶠp_grad_matrix ⋅ - DiagonalMatrixRow(ᶜkappa_m * (T_tri * cv_d - ᶜK - ᶜΦ)) + - DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / abs2(ᶠinterp(ᶜρ))) ⋅ - ᶠinterp_matrix() - ) - @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) - 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 = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m * 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ₕ = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ - if p.atmos.rayleigh_sponge isa RayleighSponge - @. ∂ᶠu₃_err_∂ᶠu₃ = - dtγ * ( - ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ - ∂ᶜK_∂ᶠu₃ + DiagonalMatrixRow(-p.ᶠβ_rayleigh_w * (one_C3xACT3,)) - ) - (I_u₃,) - else - @. ∂ᶠu₃_err_∂ᶠu₃ = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ - ∂ᶜK_∂ᶠu₃ - (I_u₃,) - end - - if use_derivative(diffusion_flag) - (; ᶜK_h, ᶜK_u) = p - @. ᶜdiffusion_h_matrix = - ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅ - ᶠ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 = - ᶜadvdivᵥ_matrix() ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠ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_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m) * ᶜspecific.e_tot - - ᶜkappa_m * e_int_v0 * ᶜspecific.q_tot - ) / ᶜρ, - ) - @. ∂ᶜρe_tot_err_∂ᶜρe_tot = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) - - (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 = - dtγ * ᶜdiffusion_h_matrix ⋅ - DiagonalMatrixRow(ᶜkappa_m * e_int_v0 / ᶜρ) - 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_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(ᶜq) / ᶜρ) - @. ∂ᶜρq_err_∂ᶜρq = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) - (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 - - @inline dissipation_rate(tke⁰, mixing_length) = - tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : 1 / dt - @inline ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = - tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : - typeof(tke⁰)(0) - - ᶜdissipation_matrix_diagonal = p.ᶜtemp_scalar - @. ᶜdissipation_matrix_diagonal = - ᶜρatke⁰ * ∂dissipation_rate_∂tke⁰(ᶜtke⁰, ᶜmixing_length) - - ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] - ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = - matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] - @. ∂ᶜρatke⁰_err_∂ᶜρ = - dtγ * ( - ᶜdiffusion_u_matrix - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰) - @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = - dtγ * ( - ( - ᶜdiffusion_u_matrix - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - - DiagonalMatrixRow(dissipation_rate(ᶜtke⁰, ᶜmixing_length)) - ) - (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ₕ = - dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_u_matrix - (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) - ᶠtmp = p.ᶠtemp_CT3 - @. ᶠtmp = CT3(unit_basis_vector_data(CT3, ᶠlg)) * ᶠwinterp(ᶜJ, ᶜρ) - @. ∂ᶜρqₚ_err_∂ᶜρqₚ += - dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠtmp) ⋅ - ᶠright_bias_matrix() ⋅ DiagonalMatrixRow(-(ᶜwₚ) / ᶜρ) - end - end - - if p.atmos.turbconv_model isa PrognosticEDMFX - if use_derivative(sgs_advection_flag) - (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p - (; bdmr_l, bdmr_r, bdmr) = p - is_third_order = edmfx_upwinding == Val(:third_order) - ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 - ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; - top = Operators.SetValue(zero(CT3{FT})), - bottom = Operators.SetValue(zero(CT3{FT})), - ) # Need to wrap ᶠupwind in this for well-defined boundaries. - UpwindMatrixRowType = - is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow - ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix - ᶠset_upwind_matrix_bcs = Operators.SetBoundaryOperator(; - top = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), - bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), - ) # Need to wrap ᶠupwind_matrix in this for well-defined boundaries. - ᶜkappa_mʲ = p.ᶜtemp_scalar - @. ᶜkappa_mʲ = - TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / - TD.cv_m(thermo_params, ᶜtsʲs.:(1)) - ∂ᶜq_totʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ = - dtγ * ( - DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - - ᶜadvdivᵥ_matrix() ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - ) - (I,) - - ∂ᶜmseʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶜmseʲ_err_∂ᶜq_totʲ = - dtγ * ( - -DiagonalMatrixRow( - adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * - ᶜgradᵥ_ᶠΦ * - Y.c.ρ * - ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * e_int_v0, - ) - ) - ∂ᶜmseʲ_err_∂ᶜρ = matrix[@name(c.sgsʲs.:(1).mse), @name(c.ρ)] - @. ∂ᶜmseʲ_err_∂ᶜρ = - dtγ * ( - -DiagonalMatrixRow( - adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * ᶜgradᵥ_ᶠΦ / ᶜρʲs.:(1), - ) - ) - ∂ᶜmseʲ_err_∂ᶜmseʲ = - matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).mse)] - @. ∂ᶜmseʲ_err_∂ᶜmseʲ = - dtγ * ( - DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - - ᶜadvdivᵥ_matrix() ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - - DiagonalMatrixRow( - adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * - ᶜgradᵥ_ᶠΦ * - Y.c.ρ * - ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - ) - (I,) - - ∂ᶜρaʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)] - @. ᶠbidiagonal_matrix_ct3 = - DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind( - ᶠu³ʲs.:(1), - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ), - ), - ) ⋅ ᶠwinterp_matrix(ᶜJ) ⋅ DiagonalMatrixRow( - ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp) * - e_int_v0, - ) - @. ᶠbidiagonal_matrix_ct3_2 = - DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow( - Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * - e_int_v0, - ) - - @. ∂ᶜρaʲ_err_∂ᶜq_totʲ = - dtγ * ᶜadvdivᵥ_matrix() ⋅ - (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) - ∂ᶜρaʲ_err_∂ᶜmseʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)] - @. ᶠbidiagonal_matrix_ct3 = - DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind( - ᶠu³ʲs.:(1), - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ), - ), - ) ⋅ ᶠwinterp_matrix(ᶜJ) ⋅ DiagonalMatrixRow( - ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - @. ᶠbidiagonal_matrix_ct3_2 = - DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow( - Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - @. ∂ᶜρaʲ_err_∂ᶜmseʲ = - dtγ * ᶜadvdivᵥ_matrix() ⋅ - (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) - ∂ᶜρaʲ_err_∂ᶜρaʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)] - @. ᶜadvection_matrix = - -(ᶜadvdivᵥ_matrix()) ⋅ - DiagonalMatrixRow(ᶠwinterp(ᶜJ, ᶜρʲs.:(1))) - @. ∂ᶜρaʲ_err_∂ᶜρaʲ = - dtγ * ᶜadvection_matrix ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow(1 / ᶜρʲs.:(1)) - (I,) - - ∂ᶠu₃ʲ_err_∂ᶜρ = matrix[@name(f.sgsʲs.:(1).u₃), @name(c.ρ)] - @. ∂ᶠu₃ʲ_err_∂ᶜρ = - dtγ * DiagonalMatrixRow(ᶠgradᵥ_ᶜΦ / ᶠinterp(ᶜρʲs.:(1))) ⋅ - ᶠinterp_matrix() - ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = - dtγ * DiagonalMatrixRow( - ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / (ᶠinterp(ᶜρʲs.:(1)))^2, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp) * - e_int_v0, - ) - ∂ᶠu₃ʲ_err_∂ᶜmseʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)] - @. ∂ᶠu₃ʲ_err_∂ᶜmseʲ = - dtγ * DiagonalMatrixRow( - ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / (ᶠinterp(ᶜρʲs.:(1)))^2, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] - ᶜu₃ʲ = ᶜtemp_C3 - @. ᶜu₃ʲ = ᶜinterp(Y.f.sgsʲs.:(1).u₃) - - @. bdmr_l = convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix()) - @. bdmr_r = convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix()) - @. bdmr = ifelse(ᶜu₃ʲ.components.data.:1 > 0, bdmr_l, bdmr_r) - - @. ᶠtridiagonal_matrix_c3 = -(ᶠgradᵥ_matrix()) ⋅ bdmr - if p.atmos.rayleigh_sponge isa RayleighSponge - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * ( - ᶠtridiagonal_matrix_c3 ⋅ - DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - - DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,)) - ) - (I_u₃,) - else - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * ᶠtridiagonal_matrix_c3 ⋅ - DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,) - end - elseif p.atmos.rayleigh_sponge isa RayleighSponge - ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * -DiagonalMatrixRow(p.ᶠβ_rayleigh_w * (one_C3xACT3,)) - - (I_u₃,) - end - end - + A.cache.b_krylov .= b + ldiv!(A.cache.x_krylov, A, A.cache.b_krylov) + x .= A.cache.x_krylov end diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 5940a788df3..5cc8a00d391 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -9,8 +9,8 @@ import ClimaAtmos as CA import LinearAlgebra import ClimaCore.Fields import ClimaTimeSteppers as CTS - import ClimaDiagnostics +import DiffEqBase: KeywordArgSilent function get_atmos(config::AtmosConfig, params) (; turbconv_params) = params @@ -375,56 +375,29 @@ 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(::CTS.RosenbrockAlgorithm) = 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 Union{CTS.IMEXAlgorithm, CTS.RosenbrockAlgorithm} + 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"] + jacobian_algorithm = if parsed_args["debug_approximate_jacobian"] + DebugJacobian( + ApproxJacobian(; approximate_solve_iters); + use_exact_jacobian, + always_update_exact_jacobian, + ) else - return (; jac_prototype = A, Wfact = Wfact!) + use_exact_jacobian ? + ExactJacobian(; always_update_exact_jacobian) : + ApproxJacobian(; approximate_solve_iters) end + @info "Jacobian algorithm: $(dump_string(jacobian_algorithm))" + A = ImplicitEquationJacobian(jacobian_algorithm, Y, p) + (; jac_prototype = A, Wfact = update_jacobian!) else - return NamedTuple() + (;) end -end #= ode_configuration(Y, parsed_args) @@ -433,17 +406,14 @@ 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 ode_name == "SSPKnoth" - return CTS.RosenbrockAlgorithm(CTS.tableau(CTS.SSPKnoth())) - end - - 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`" + return if ode_algo_name <: CTS.RosenbrockAlgorithmName + CTS.RosenbrockAlgorithm(CTS.tableau(ode_algo_name())) + elseif ode_algo_name <: CTS.ERKAlgorithmName + 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"] @@ -472,9 +442,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!) + CTS.IMEXAlgorithm(ode_algo_name(), newtons_method) end end @@ -528,29 +496,20 @@ function get_sim_info(config::AtmosConfig) 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)..., - tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), - ) - if is_cts_algo(ode_algo) - CTS.ClimaODEFunction(; - T_exp_T_lim! = 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 + jacobian = jac_kwargs(ode_algo, Y, p, parsed_args) + CTS.ClimaODEFunction(; + T_exp_T_lim! = remaining_tendency!, + T_imp! = SciMLBase.ODEFunction(implicit_tendency!; jacobian...), + lim! = limiters_func!, + dss!, + post_explicit! = set_precomputed_quantities!, + post_implicit! = set_precomputed_quantities!, + ) # TODO: Separate the Wfact! and jac_prototype kwargs from T_imp!. else remaining_tendency! # should be total_tendency! end @@ -566,7 +525,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 8b686bc2e9d..6b5dea700df 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/utilities.jl b/src/utils/utilities.jl index b5c04f5f268..28a9b8b0cf9 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -333,6 +333,19 @@ macro timed_str(ex) end end +""" + dump_string(x) + +Returns a string that contains the output of `dump(x)`. +""" +function dump_string(x) + buffer = IOBuffer() + dump(buffer, x) + result = String(take!(buffer)) + close(buffer) + return result +end + struct AllNothing end const all_nothing = AllNothing() Base.getproperty(::AllNothing, ::Symbol) = nothing diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 721bc78b0e4..9a2dcd80360 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 994d80c8283..ca2ad9c4c14 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -25,7 +25,6 @@ DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" 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" @@ -72,7 +71,6 @@ ClimaCoreSpectra = "0.1" ClimaCoreTempestRemap = "0.3" ClimaCoreVTK = "0.7" DiffEqNoiseProcess = "5" -ForwardDiff = "0.10" Glob = "1" JSON = "0.21" NCRegressionTests = "0.2"