Merge pull request #3827 from CliMA/dy/auto_dense_jacobian

Add AutoDenseJacobian algorithm for implicit solver

Author: dennisYatunin

.buildkite/Manifest-v1.11.toml

@@ -363,7 +363,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.3"

.buildkite/pipeline.yml

@@ -320,6 +320,14 @@ steps:
           --job_id baroclinic_wave
         artifact_paths: "baroclinic_wave/output_active/*"
 
+      - label: ":computer: baroclinic wave dense autodiff"
+        key: baroclinic_wave_dense_autodiff
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/baroclinic_wave_dense_autodiff.yml
+          --job_id baroclinic_wave_dense_autodiff
+        artifact_paths: "baroclinic_wave_dense_autodiff/output_active/*"
+
       - label: ":computer: no lim baroclinic wave equilmoist"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -385,6 +393,16 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":umbrella: aquaplanet nonequil allsky monin_obukhov varying insol gravity wave (gfdl_restart) high top 1-moment dense autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/aquaplanet_nonequil_allsky_gw_res_dense_autodiff.yml
+          --job_id aquaplanet_nonequil_allsky_gw_res_dense_autodiff
+        artifact_paths: "aquaplanet_nonequil_allsky_gw_res_dense_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
       - label: ":computer: aquaplanet equil allsky monin_obukhov varying insol gravity wave (raw_topo) high top zonally asymmetric"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -731,6 +749,16 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":genie: Diagnostic EDMFX aquaplanet dense autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/diagnostic_edmfx_aquaplanet_dense_autodiff.yml
+          --job_id diagnostic_edmfx_aquaplanet_dense_autodiff
+        artifact_paths: "diagnostic_edmfx_aquaplanet_dense_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
   - group: "Prognostic EDMFX"
     steps:
 
@@ -896,6 +924,16 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":genie: Prognostic EDMFX aquaplanet dense autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_aquaplanet_dense_autodiff.yml
+          --job_id prognostic_edmfx_aquaplanet_dense_autodiff
+        artifact_paths: "prognostic_edmfx_aquaplanet_dense_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
   - group: "GPU"
     steps:
 
@@ -924,6 +962,19 @@ steps:
           slurm_gpus: 1
           slurm_mem: 16GB
 
+      - label: "GPU: baroclinic wave dense autodiff"
+        key: "baroclinic_wave_gpu_dense_autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/baroclinic_wave_dense_autodiff.yml
+          --job_id baroclinic_wave_gpu_dense_autodiff
+        artifact_paths: "baroclinic_wave_gpu_dense_autodiff/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+          slurm_mem: 16GB
+
       - label: "GPU: compare BW with CPU"
         command: >
           tar xvf baroclinic_wave/output_active/hdf5_files.tar -C baroclinic_wave
@@ -1137,6 +1188,15 @@ steps:
         agents:
           slurm_mem: 32GB
 
+      - label: ":fire: Flame graph: prognostic edmf dense autodiff"
+        command: >
+          julia --color=yes --project=.buildkite perf/flame.jl
+          --config_file $PERF_CONFIG_PATH/bm_aquaplanet_progedmf_dense_autodiff.yml
+          --job_id flame_aquaplanet_progedmf_dense_autodiff
+        artifact_paths: "flame_aquaplanet_progedmf_dense_autodiff/*"
+        agents:
+          slurm_mem: 32GB
+
       - label: ":fire: Flame graph: diffusion"
         command: >
           julia --color=yes --project=.buildkite perf/flame.jl

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"
@@ -40,14 +41,15 @@ Adapt = "4"
 ArgParse = "1"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
-ClimaComms = "0.6.6"
-ClimaCore = "0.14.24"
+ClimaComms = "0.6.8"
+ClimaCore = "0.14.34"
 ClimaDiagnostics = "0.2.12"
 ClimaParams = "0.10.31"
 ClimaTimeSteppers = "0.8.2"
 ClimaUtilities = "0.1.22"
 CloudMicrophysics = "0.23.0"
 Dates = "1"
+ForwardDiff = ">= 0.10"
 Insolation = "0.9.2"
 Interpolations = "0.15.1, 0.16"
 LazyArtifacts = "1"
@@ -64,6 +66,6 @@ StaticArrays = "1.7"
 Statistics = "1"
 SurfaceFluxes = "0.11, 0.12"
 Thermodynamics = "0.12.8"
-UnrolledUtilities = "0.1.6"
+UnrolledUtilities = "0.1.9"
 YAML = "0.4"
 julia = "1.9"

config/default_configs/default_config.yml

@@ -86,6 +86,12 @@ 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_dense_jacobian:
+  help: "Whether to use a dense Jacobian matrix that is computed using forward-mode automatic differentiation and inverted using LU factorization [`true`, `false` (default)]"
+  value: false
+update_jacobian_every:
+  help: "Frequency at which the Jacobian matrix should be updated (once per timestep, once per timestepper stage, or once per linear solve) [`dt`, `stage`, `solve` (default)]"
+  value: "solve"
 # Radiation
 rad:
   help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]"

config/model_configs/aquaplanet_nonequil_allsky_gw_res_dense_autodiff.yml

@@ -0,0 +1,30 @@
+z_max: 60000.0
+z_elem: 31
+dz_bottom: 50.0
+rayleigh_sponge: true
+dt: "360secs"
+t_end: "6hours" # This takes 2 hours to run in CI when t_end is 1 day.
+dt_save_state_to_disk: "6hours"
+vert_diff: "DecayWithHeightDiffusion"
+implicit_diffusion: true
+use_dense_jacobian: true
+update_jacobian_every: dt
+moist: "nonequil"
+precip_model: "1M"
+rad: "allskywithclear"
+aerosol_radiation: true
+prescribe_clouds_in_radiation: true
+radiation_reset_rng_seed: true
+insolation: "timevarying"
+non_orographic_gravity_wave: true
+orographic_gravity_wave: "gfdl_restart"
+surface_setup: "DefaultMoninObukhov"
+prescribe_ozone: true
+prescribed_aerosols: ["SO4", "CB1", "OC1", "DST01", "SSLT01"]
+toml: [toml/sphere_aquaplanet_1M.toml]
+diagnostics:
+  - short_name: [edt, evu, mmrso4, mmrbcpo, mmrocpo, mmrdust, mmrss, loadss, o3, od550aer, odsc550aer]
+    reduction_time: average
+    period: 6hours
+  - short_name: [reffclw, reffcli]
+    period: 6hours

config/model_configs/baroclinic_wave_dense_autodiff.yml

@@ -0,0 +1,13 @@
+dt_save_state_to_disk: "2days"
+initial_condition: "DryBaroclinicWave"
+use_dense_jacobian: true
+update_jacobian_every: dt
+dt: "400secs"
+t_end: "10days"
+disable_surface_flux_tendency: true
+diagnostics:
+  - short_name: [pfull, ua, wa, va, rv, ta, ke]
+    period: 1days
+  - short_name: [pfull, ua, wa, va, rv, ta, ke]
+    period: 1days
+    writer: h5

config/model_configs/diagnostic_edmfx_aquaplanet_dense_autodiff.yml

@@ -0,0 +1,24 @@
+z_max: 60000.0
+z_elem: 31
+dz_bottom: 50.0
+rayleigh_sponge: true
+surface_setup: DefaultMoninObukhov
+rad: clearsky
+co2_model: fixed
+turbconv: diagnostic_edmfx
+implicit_diffusion: true
+use_dense_jacobian: true
+update_jacobian_every: dt
+prognostic_tke: true
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_nh_pressure: true
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+moist: equil
+cloud_model: "quadrature_sgs"
+precip_model: 0M
+dt: 120secs
+t_end: 1hours
+toml: [toml/diagnostic_edmfx_0M.toml]

config/model_configs/prognostic_edmfx_aquaplanet_dense_autodiff.yml

@@ -0,0 +1,41 @@
+# This job takes too long so we use a lower top and coarser vertical resolution
+# z_max: 60000.0
+# z_elem: 31
+# dz_bottom: 50.0
+# rayleigh_sponge: true
+surface_setup: DefaultMoninObukhov
+rad: clearsky
+co2_model: fixed
+turbconv: prognostic_edmfx
+prognostic_tke: true
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_nh_pressure: true
+edmfx_filter: true
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+implicit_diffusion: true
+implicit_sgs_advection: true
+implicit_sgs_entr_detr: true
+implicit_sgs_nh_pressure: true
+implicit_sgs_mass_flux: true
+use_dense_jacobian: true
+update_jacobian_every: dt
+max_newton_iters_ode: 3
+moist: equil
+cloud_model: "quadrature_sgs"
+precip_model: 0M
+dt: 20secs
+t_end: 40mins
+dt_save_state_to_disk: 600secs
+toml: [toml/prognostic_edmfx.toml]
+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
+    period: 40mins
+  - short_name: [arup, waup, taup, thetaaup, haup, husup, hurup, clwup, cliup, waen, tke, lmix]
+    reduction_time: average
+    period: 40mins
+use_itime: true

config/perf_configs/bm_aquaplanet_progedmf_dense_autodiff.yml

@@ -0,0 +1,33 @@
+h_elem: 12
+z_elem: 25
+dt: 10secs
+t_end: 61mins
+dt_save_state_to_disk: "Inf"
+log_progress: false
+rayleigh_sponge: true
+viscous_sponge: true
+implicit_diffusion: true
+implicit_sgs_advection: true
+implicit_sgs_entr_detr: true
+implicit_sgs_nh_pressure: true
+implicit_sgs_mass_flux: true
+use_dense_jacobian: true
+update_jacobian_every: dt
+moist: equil
+surface_setup: DefaultMoninObukhov
+rad: allskywithclear
+dt_rad: 1hours
+dt_cloud_fraction: 1hours
+turbconv: prognostic_edmfx
+max_newton_iters_ode: 3
+prognostic_tke: true
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_nh_pressure: true
+edmfx_filter: true
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+cloud_model: "quadrature_sgs"
+precip_model: 0M
+toml: [toml/prognostic_edmfx.toml]

docs/src/api.md

@@ -43,16 +43,24 @@ ClimaAtmos.InitialConditions.Bomex
 ClimaAtmos.InitialConditions.Soares
 ```
 
-### Jacobian
+### Helper
+
+```@docs
+ClimaAtmos.InitialConditions.ColumnInterpolatableField
+```
+
+## Jacobian
 
 ```@docs
 ClimaAtmos.Jacobian
 ClimaAtmos.JacobianAlgorithm
 ClimaAtmos.ManualSparseJacobian
+ClimaAtmos.AutoDenseJacobian
 ```
 
-### Helper
+### Internals
 
 ```@docs
-ClimaAtmos.InitialConditions.ColumnInterpolatableField
+ClimaAtmos.parallel_lu_factorize!
+ClimaAtmos.parallel_lu_solve!
 ```