Add an interface for debugging the implicit solver Jacobian

.buildkite/pipeline.yml

@@ -134,7 +134,7 @@ steps:
           julia --color=yes --project=examples examples/hybrid/driver.jl
           --config_file $CONFIG_PATH/single_column_nonorographic_gravity_wave.yml
           --job_id single_column_nonorographic_gravity_wave
-        artifact_paths: "single_column_nonorographic_gravity_wave/*"
+        artifact_paths: "single_column_nonorographic_gravity_wave/output_active/*"
 
   - group: "Column Examples"
     steps:
@@ -177,7 +177,7 @@ steps:
           julia --color=yes --project=examples examples/hybrid/driver.jl
           --config_file $CONFIG_PATH/les_isdac_box.yml
           --job_id les_isdac_box
-        artifact_paths: "les_isdac_box/*"
+        artifact_paths: "les_isdac_box/output_active/*"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:
@@ -745,7 +745,7 @@ steps:
           julia --color=yes --project=examples examples/hybrid/driver.jl
           --config_file $CONFIG_PATH/prognostic_edmfx_bomex_box.yml
           --job_id prognostic_edmfx_bomex_box
-        artifact_paths: "prognostic_edmfx_bomex_box/*"
+        artifact_paths: "prognostic_edmfx_bomex_box/output_active/*"
         agents:
           slurm_mem: 20GB
 
@@ -766,7 +766,7 @@ steps:
           julia --color=yes --project=examples examples/hybrid/driver.jl
           --config_file $CONFIG_PATH/sphere_nonorographic_gravity_wave.yml
           --job_id sphere_nonorographic_gravity_wave
-        artifact_paths: "sphere_nonorographic_gravity_wave/*"
+        artifact_paths: "sphere_nonorographic_gravity_wave/output_active/*"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:

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: 10 # 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/aquaplanet_diagedmf.yml

@@ -1,6 +1,6 @@
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"
-output_default_diagnostics: false
+# output_default_diagnostics: false
 h_elem: 30
 z_max: 55000.0
 z_elem: 63

config/model_configs/aquaplanet_progedmf.yml

@@ -1,6 +1,6 @@
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"
-output_default_diagnostics: false
+# output_default_diagnostics: false
 h_elem: 30
 z_max: 55000.0
 z_elem: 63

config/model_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yml

@@ -1,6 +1,6 @@
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"
-output_default_diagnostics: false
+# output_default_diagnostics: false
 dt: "150secs"
 t_end: "2days"
 h_elem: 16

config/model_configs/central_gpu_hs_rhoe_equil_55km_nz63_0M.yml

@@ -1,6 +1,6 @@
 dt_save_state_to_disk: "Inf"
 dt_save_to_sol: "Inf"
-output_default_diagnostics: false
+# output_default_diagnostics: false
 dt: "150secs"
 t_end: "2days"
 h_elem: 16

config/model_configs/plane_density_current_test.yml

@@ -9,7 +9,7 @@ z_stretch: false
 x_elem: 80
 config: "plane"
 z_max: 6400.0
-output_default_diagnostics: false
+# output_default_diagnostics: false
 diagnostics:
   - short_name: thetaa
     period: 900secs

config/model_configs/prognostic_edmfx_aquaplanet.yml

@@ -16,7 +16,7 @@ dt: 10secs
 t_end: 1hours
 dt_save_state_to_disk: 600secs
 toml: [toml/prognostic_edmfx.toml]
-output_default_diagnostics: false
+# output_default_diagnostics: false
 diagnostics:
   - short_name: [ts, ta, thetaa, ha, pfull, rhoa, ua, va, wa, hfes, hur, hus, cl, clw, cli, evspsbl, pr, rsd, rsu, rld, rlu]
     reduction_time: average

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"

config/model_configs/prognostic_edmfx_gcmdriven_column.yml

@@ -32,7 +32,7 @@ call_cloud_diagnostics_per_stage : true
 toml: [toml/prognostic_edmfx_gcmdriven.toml]
 netcdf_output_at_levels: true
 netcdf_interpolation_num_points: [2, 2, 60]
-output_default_diagnostics: false
+# output_default_diagnostics: false
 rad: allskywithclear
 insolation: "gcmdriven"
 diagnostics:

config/model_configs/prognostic_edmfx_simpleplume_column.yml

@@ -19,14 +19,16 @@ cloud_model: "quadrature_sgs"
 call_cloud_diagnostics_per_stage: true
 config: "column"
 z_max: 4e3
-z_elem: 80
+z_elem: 40
 z_stretch: false
 perturb_initstate: false
-dt: 20secs"
-t_end: "12hours"
+dt: 60secs"
+t_end: "4hours"
 dt_save_state_to_disk: "10mins"
 toml: [toml/prognostic_edmfx_simpleplume.toml]
 netcdf_interpolation_num_points: [2, 2, 80]
+# use_exact_jacobian: false
+n_steps_update_exact_jacobian: 0
 diagnostics:
   - short_name: [ts, ta, thetaa, ha, pfull, rhoa, ua, va, wa, hur, hus, cl, clw, cli, hussfc, evspsbl, pr]
     period: 10mins

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"