Add an interface for debugging the implicit solver Jacobian

Project.toml

@@ -21,6 +21,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
@@ -61,6 +62,7 @@ DiffEqBase = "6"
 DiffEqCallbacks = "2"
 DocStringExtensions = "0.8, 0.9"
 FastGaussQuadrature = "0.4, 0.5, 1"
+ForwardDiff = "0.10"
 Insolation = "=0.9.1"
 Interpolations = "0.14, 0.15"
 IntervalSets = "0.5, 0.6, 0.7"

config/default_configs/default_config.yml

@@ -80,6 +80,15 @@ eisenstat_walker_forcing_alpha:
 jvp_step_adjustment:
   help: "Amount by which the step size of the forward difference approximation of the Jacobian-vector product in the Krylov method should be scaled (only used if `use_krylov_method` is `true`)"
   value: 1.0
+use_exact_jacobian:
+  help: "Whether to solve the linear system using the exact Jacobian, which is computed with forward-mode automatic differentiation (default is `false`)"
+  value: false
+n_steps_update_exact_jacobian:
+  help: "Number of timesteps between updates to the exact Jacobian (default is 0, which corresponds to updating the exact Jacobian on each Newton iteration)"
+  value: 0
+debug_jacobian_approximation:
+  help: "Whether to check the accuracy of the Jacobian approximation against the exact Jacobian every `n_steps_update_exact_jacobian` timesteps (default is `false`)"
+  value: false
 # Radiation
 rad:
   help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]"

docs/Manifest.toml

@@ -306,7 +306,7 @@ weakdeps = ["SparseArrays"]
     ChainRulesCoreSparseArraysExt = "SparseArrays"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.20.1"

examples/Manifest.toml

@@ -299,7 +299,7 @@ weakdeps = ["CairoMakie"]
     CairoMakieExt = "CairoMakie"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.20.1"

perf/Manifest.toml

@@ -304,7 +304,7 @@ weakdeps = ["CairoMakie"]
     CairoMakieExt = "CairoMakie"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaTimeSteppers", "CloudMicrophysics", "Colors", "Dates", "Dierckx", "DiffEqBase", "DiffEqCallbacks", "DocStringExtensions", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "IntervalSets", "Krylov", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "Printf", "RRTMGP", "Random", "RootSolvers", "SciMLBase", "StaticArrays", "Statistics", "StatsBase", "SurfaceFluxes", "Thermodynamics", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.20.1"

src/ClimaAtmos.jl

@@ -46,6 +46,9 @@ include(joinpath("prognostic_equations", "zero_velocity.jl"))
 
 include(joinpath("prognostic_equations", "implicit", "implicit_tendency.jl"))
 include(joinpath("prognostic_equations", "implicit", "implicit_solver.jl"))
+include(joinpath("prognostic_equations", "implicit", "approx_jacobian.jl"))
+include(joinpath("prognostic_equations", "implicit", "exact_jacobian.jl"))
+include(joinpath("prognostic_equations", "implicit", "debug_jacobian.jl"))
 
 include(joinpath("prognostic_equations", "remaining_tendency.jl"))
 include(joinpath("prognostic_equations", "forcing", "large_scale_advection.jl")) # TODO: should this be in tendencies/?

src/callbacks/callbacks.jl

@@ -320,6 +320,23 @@ function gc_func(integrator)
     return nothing
 end
 
+function update_exact_jacobian!(integrator)
+    @assert integrator.alg isa CTS.IMEXAlgorithm
+    γs = unique(filter(!iszero, diag(integrator.alg.tableau.a_imp)))
+    length(γs) == 1 ||
+        error("The exact Jacobian must be updated on every Newton iteration, \
+               rather than on every timestep (or every N steps), because the \
+               specified IMEX algorithm has implicit stages with distinct \
+               coefficients (i.e., it is not an SDIRK algorithm).")
+    update_exact_jacobian!(
+        integrator.cache.newtons_method_cache.j,
+        integrator.u,
+        integrator.p,
+        integrator.dt * γs[1],
+        integrator.t,
+    )
+end
+
 """
     maybe_graceful_exit(integrator)
 

src/callbacks/get_callbacks.jl

@@ -60,6 +60,19 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start)
         )
     end
 
+    if (
+        parsed_args["use_exact_jacobian"] ||
+        parsed_args["debug_jacobian_approximation"]
+    ) && parsed_args["n_steps_update_exact_jacobian"] != 0
+        callbacks = (
+            callbacks...,
+            call_every_n_steps(
+                update_exact_jacobian!,
+                parsed_args["n_steps_update_exact_jacobian"],
+            ),
+        )
+    end
+
     if atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode
         # TODO: better if-else criteria?
         dt_rad = if parsed_args["config"] == "column"

src/prognostic_equations/edmfx_boundary_condition.jl

@@ -4,8 +4,8 @@
 
 function sgs_scalar_first_interior_bc(
     ᶜz_int::FT,
-    ᶜρ_int::FT,
-    ᶜscalar_int::FT,
+    ᶜρ_int,
+    ᶜscalar_int,
     sfc_buoyancy_flux,
     sfc_ρ_flux_scalar,
     ustar,

src/prognostic_equations/edmfx_closures.jl

@@ -189,17 +189,17 @@ Smagorinsky length scale.
 """
 function mixing_length(
     params,
-    ustar::FT,
+    ustar,
     ᶜz::FT,
-    z_sfc::FT,
-    ᶜdz::FT,
-    sfc_tke::FT,
-    ᶜlinear_buoygrad::FT,
-    ᶜtke::FT,
-    obukhov_length::FT,
-    ᶜstrain_rate_norm::FT,
-    ᶜPr::FT,
-    ᶜtke_exch::FT,
+    z_sfc,
+    ᶜdz,
+    sfc_tke,
+    ᶜlinear_buoygrad,
+    ᶜtke,
+    obukhov_length,
+    ᶜstrain_rate_norm,
+    ᶜPr,
+    ᶜtke_exch,
 ) where {FT}
 
     turbconv_params = CAP.turbconv_params(params)

src/prognostic_equations/edmfx_entr_detr.jl

@@ -85,17 +85,17 @@ end
 function entrainment(
     params,
     ᶜz::FT,
-    z_sfc::FT,
-    ᶜp::FT,
-    ᶜρ::FT,
-    buoy_flux_surface::FT,
-    ᶜaʲ::FT,
-    ᶜwʲ::FT,
-    ᶜRHʲ::FT,
-    ᶜbuoyʲ::FT,
-    ᶜw⁰::FT,
-    ᶜRH⁰::FT,
-    ᶜbuoy⁰::FT,
+    z_sfc,
+    ᶜp,
+    ᶜρ,
+    buoy_flux_surface,
+    ᶜaʲ,
+    ᶜwʲ,
+    ᶜRHʲ,
+    ᶜbuoyʲ,
+    ᶜw⁰,
+    ᶜRH⁰,
+    ᶜbuoy⁰,
     ::GeneralizedEntrainment,
 ) where {FT}
     turbconv_params = CAP.turbconv_params(params)
@@ -312,19 +312,19 @@ end
 function detrainment(
     params,
     ᶜz::FT,
-    z_sfc::FT,
-    ᶜp::FT,
-    ᶜρ::FT,
-    buoy_flux_surface::FT,
-    ᶜaʲ::FT,
-    ᶜwʲ::FT,
-    ᶜRHʲ::FT,
-    ᶜbuoyʲ::FT,
-    ᶜw⁰::FT,
-    ᶜRH⁰::FT,
-    ᶜbuoy⁰::FT,
-    ᶜentr::FT,
-    ᶜvert_div::FT,
+    z_sfc,
+    ᶜp,
+    ᶜρ,
+    buoy_flux_surface,
+    ᶜaʲ,
+    ᶜwʲ,
+    ᶜRHʲ,
+    ᶜbuoyʲ,
+    ᶜw⁰,
+    ᶜRH⁰,
+    ᶜbuoy⁰,
+    ᶜentr,
+    ᶜvert_div,
     ::ConstantAreaDetrainment,
 ) where {FT}
     detr = ᶜentr - ᶜvert_div