Add nonequilibrium blocks to Jacobian

Author: dennisYatunin

.buildkite/Manifest-v1.11.toml

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

.buildkite/ci_driver.jl

@@ -41,6 +41,7 @@ import Base.Filesystem: rm
 import Statistics: mean
 import LinearAlgebra: norm_sqr
 include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl"))
+include(joinpath(pkgdir(CA), "post_processing", "jacobian_plots.jl"))
 
 ref_job_id = config.parsed_args["reference_job_id"]
 reference_job_id = isnothing(ref_job_id) ? simulation.job_id : ref_job_id
@@ -187,6 +188,7 @@ if ClimaComms.iamroot(config.comms_ctx)
             "reproducibility_utils.jl",
         ),
     )
+
     @info "Plotting"
     paths = latest_comparable_dirs() # __build__ path (not job path)
     if isempty(paths)
@@ -218,6 +220,17 @@ if ClimaComms.iamroot(config.comms_ctx)
         end
         make_plots(Val(Symbol(reference_job_id)), paths)
     end
+    if (
+        isnothing(config.parsed_args["plot_jacobian"]) ?
+        config.parsed_args["debug_approximate_jacobian"] :
+        config.parsed_args["plot_jacobian"]
+    )
+        make_jacobian_plots(
+            simulation.output_dir,
+            integrator.u,
+            float(integrator.dt),
+        )
+    end
     @info "Plotting done"
 
     if islink(simulation.output_dir)

.buildkite/pipeline.yml

@@ -31,6 +31,9 @@ steps:
 
       - echo "--- Instantiate .buildkite"
       - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true); using CUDA; CUDA.precompile_runtime(); Pkg.status()'"
+      - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"dy/disable_shmem\"))'"
+      - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"Thermodynamics\", rev=\"dy/more_autodiff_fixes\"))'"
+      - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"CloudMicrophysics\", rev=\"main\"))'"
 
     agents:
       slurm_cpus_per_task: 8
@@ -119,7 +122,7 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_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:
@@ -153,6 +156,8 @@ steps:
           --config_file $CONFIG_PATH/box_density_current_test.yml
           --job_id box_density_current_test
         artifact_paths: "box_density_current_test/output_active/*"
+        agents:
+          slurm_mem: 20GB
 
       - label: ":computer: Box rcemipii with diagnostic edmf"
         command: >
@@ -168,7 +173,7 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_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:
@@ -859,7 +864,7 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_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
 
@@ -885,7 +890,7 @@ steps:
           julia --color=yes --project=.buildkite .buildkite/ci_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

@@ -16,6 +16,7 @@ ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
 CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
@@ -48,6 +49,7 @@ ClimaTimeSteppers = "0.8.2"
 ClimaUtilities = "0.1.22"
 CloudMicrophysics = "0.22.9"
 Dates = "1"
+ForwardDiff = "0.10.36"
 Insolation = "0.9.2"
 Interpolations = "0.15.1"
 LazyArtifacts = "1"

config/default_configs/default_config.yml

@@ -62,9 +62,6 @@ newton_rtol:
 surface_thermo_state_type:
   help: "Surface thermo state type [`GCMSurfaceThermoState` (default), `PrescribedThermoState`]"
   value: "GCMSurfaceThermoState"
-split_ode:
-  help: "Use split of ODE problem. Examples: [`true` (implicit, default), `false` (explicit)]"
-  value: true
 use_krylov_method:
   help: "Whether to use a Krylov method to solve the linear system in Newton's method (only for ClimaTimeSteppers.jl)"
   value: false
@@ -86,6 +83,24 @@ 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 once per `dt_update_exact_jacobian` (default is `false`)"
+  value: false
+debug_approximate_jacobian:
+  help: "Whether to compare the approximate Jacobian against the exact Jacobian once per `dt_update_exact_jacobian` (default is `false`)"
+  value: true # TODO: Change this to false
+only_debug_first_column_jacobian:
+  help: "Whether to debug only the Jacobian in the first column when `debug_approximate_jacobian` is `true` and `use_exact_jacobian` is `false` (default is equal to `n_cols > 10000`)"
+  value: ~
+plot_jacobian:
+  help: "Whether to save and plot the Jacobian, and also the difference between the approximate Jacobian and exact Jacobian when `debug_approximate_jacobian` is `true` (default is equal to `debug_approximate_jacobian`)"
+  value: ~
+dt_save_jacobian:
+  help: "Time between snapshots of the Jacobian when `plot_jacobian` is `true` (default is equal to `t_span / 5`)"
+  value: ~
+dt_update_exact_jacobian:
+  help: "Time between updates to the exact Jacobian (when either `use_exact_jacobian` or `debug_approximate_jacobian` is `true`), with any value smaller than `dt` corresponding to an update before each linear solve (default is equal to `use_exact_jacobian ? 0 : dt_save_jacobian`)"
+  value: ~
 # Radiation
 rad:
   help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]"

docs/src/api.md

@@ -46,6 +46,10 @@ ClimaAtmos.InitialConditions.Soares
 ### Implicit Solver
 
 ```@docs
+ClimaAtmos.JacobianAlgorithm
+ClimaAtmos.ApproxJacobian
+ClimaAtmos.ExactJacobian
+ClimaAtmos.DebugJacobian
 ClimaAtmos.ImplicitEquationJacobian
 ```