Add an interface for debugging the implicit solver Jacobian

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"
@@ -19,6 +20,7 @@ CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
@@ -41,6 +43,7 @@ ArgParse = "1"
 ArtifactWrappers = "0.2"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
+BlockArrays = "1"
 ClimaComms = "0.6.4"
 ClimaCore = "0.14.12"
 ClimaDiagnostics = "0.2.4"
@@ -51,6 +54,7 @@ CloudMicrophysics = "0.22.3"
 Dates = "1"
 DiffEqBase = "6.145"
 FastGaussQuadrature = "0.5, 1"
+ForwardDiff = "0.10.36"
 Insolation = "0.9.2"
 Interpolations = "0.15.1"
 LazyArtifacts = "1"

config/default_configs/default_config.yml

@@ -86,6 +86,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 # 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: 20 # 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`]"

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"

docs/Manifest.toml

@@ -295,7 +295,7 @@ weakdeps = ["SparseArrays"]
     ChainRulesCoreSparseArraysExt = "SparseArrays"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.5"

examples/Manifest.toml

@@ -319,7 +319,7 @@ version = "0.5.7"
     Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.5"

examples/hybrid/driver.jl

@@ -145,10 +145,17 @@ if ClimaComms.iamroot(config.comms_ctx)
     include(
         joinpath(pkgdir(CA), "regression_tests", "self_reference_or_path.jl"),
     )
+
+    # TODO: Improve design of Diagnostics to make Jacobian plotting less clunky.
+    (; dict_writer) = simulation.output_writers
+    Yₜ_end = similar(integrator.u)
+    CA.implicit_tendency!(Yₜ_end, integrator.u, integrator.p, integrator.t)
+
     @info "Plotting"
     path = self_reference_or_path() # __build__ path (not job path)
     if path == :self_reference
-        make_plots(Val(Symbol(reference_job_id)), simulation.output_dir)
+        paths = [simulation.output_dir]
+        make_plots(Val(Symbol(reference_job_id)), paths, dict_writer, Yₜ_end)
     else
         main_job_path = joinpath(path, reference_job_id)
         nc_dir = joinpath(main_job_path, "nc_files")
@@ -170,11 +177,11 @@ if ClimaComms.iamroot(config.comms_ctx)
         end
 
         paths = if isempty(readdir(nc_dir))
-            simulation.output_dir
+            [simulation.output_dir]
         else
             [simulation.output_dir, nc_dir]
         end
-        make_plots(Val(Symbol(reference_job_id)), paths)
+        make_plots(Val(Symbol(reference_job_id)), paths, dict_writer, Yₜ_end)
     end
     @info "Plotting done"
 

perf/Manifest.toml

@@ -330,7 +330,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", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
+deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "BlockArrays", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.27.5"