Merge pull request #3859 from CliMA/dy/auto_sparse_jacobian

Add AutoSparseJacobian algorithm for implicit solver

Author: dennisYatunin

.buildkite/Manifest-v1.11.toml

@@ -363,16 +363,16 @@ version = "0.5.18"
     Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaInterpolations", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "ForwardDiff", "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", "ClimaInterpolations", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "SparseMatrixColorings", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.31.0"
 
 [[deps.ClimaComms]]
 deps = ["Adapt", "Logging", "LoggingExtras"]
-git-tree-sha1 = "75b9d1a3b4e3efa2cbbae2eb7b52f14c0b38ccf0"
+git-tree-sha1 = "f3961fa943c1bbbc376af910cb98db67084dc642"
 uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
-version = "0.6.8"
+version = "0.6.9"
 weakdeps = ["CUDA", "MPI"]
 
     [deps.ClimaComms.extensions]
@@ -2318,6 +2318,20 @@ deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"]
 uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 version = "1.11.0"
 
+[[deps.SparseMatrixColorings]]
+deps = ["ADTypes", "DocStringExtensions", "LinearAlgebra", "PrecompileTools", "Random", "SparseArrays"]
+git-tree-sha1 = "ab958b4fec46d1f1d057bb8e2a99bfdb90744646"
+uuid = "0a514795-09f3-496d-8182-132a7b665d35"
+version = "0.4.20"
+
+    [deps.SparseMatrixColorings.extensions]
+    SparseMatrixColoringsCliqueTreesExt = "CliqueTrees"
+    SparseMatrixColoringsColorsExt = "Colors"
+
+    [deps.SparseMatrixColorings.weakdeps]
+    CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
+    Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+
 [[deps.SpecialFunctions]]
 deps = ["IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
 git-tree-sha1 = "41852b8679f78c8d8961eeadc8f62cef861a52e3"

.buildkite/ci_driver.jl

@@ -39,12 +39,21 @@ using Test
 import Tar
 import Base.Filesystem: rm
 import Statistics: mean
-import LinearAlgebra: norm_sqr
+import LinearAlgebra: norm_sqr, diag, UniformScaling
 include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl"))
 
 ref_job_id = config.parsed_args["reference_job_id"]
 reference_job_id = isnothing(ref_job_id) ? simulation.job_id : ref_job_id
 
+if (
+    config.parsed_args["debug_jacobian"] &&
+    !config.parsed_args["use_dense_jacobian"]
+)
+    @info "Debugging Jacobian in first column of final state"
+    include(joinpath(@__DIR__, "..", "post_processing", "jacobian_summary.jl"))
+    print_jacobian_summary(integrator)
+end
+
 if sol_res.ret_code == :simulation_crashed
     error(
         "The ClimaAtmos simulation has crashed. See the stack trace for details.",

.buildkite/pipeline.yml

@@ -292,6 +292,13 @@ steps:
           --job_id baroclinic_wave_equil_conservation_ft64
         artifact_paths: "baroclinic_wave_equil_conservation_ft64/output_active/*"
 
+      - label: ":computer: baroclinic wave moist check conservation float64 sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/baroclinic_wave_equil_conservation_ft64_sparse_autodiff.yml
+          --job_id baroclinic_wave_equil_conservation_ft64_sparse_autodiff
+        artifact_paths: "baroclinic_wave_equil_conservation_ft64_sparse_autodiff/output_active/*"
+        soft_fail: true
+
       - label: ":computer: baroclinic wave moist check conservation with sources"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/baroclinic_wave_equil_conservation_source.yml
@@ -328,6 +335,14 @@ steps:
           --job_id baroclinic_wave_dense_autodiff
         artifact_paths: "baroclinic_wave_dense_autodiff/output_active/*"
 
+      - label: ":computer: baroclinic wave sparse autodiff"
+        key: baroclinic_wave_sparse_autodiff
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/baroclinic_wave_sparse_autodiff.yml
+          --job_id baroclinic_wave_sparse_autodiff
+        artifact_paths: "baroclinic_wave_sparse_autodiff/output_active/*"
+
       - label: ":computer: no lim baroclinic wave equilmoist"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -379,6 +394,26 @@ 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: ":umbrella: aquaplanet nonequil allsky monin_obukhov varying insol gravity wave (gfdl_restart) high top 1-moment sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/aquaplanet_nonequil_allsky_gw_res_sparse_autodiff.yml
+          --job_id aquaplanet_nonequil_allsky_gw_res_sparse_autodiff
+        artifact_paths: "aquaplanet_nonequil_allsky_gw_res_sparse_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
       - label: ":umbrella: aquaplanet nonequil allsky monin_obukhov varying insol gravity wave (gfdl_restart) high top 2-moment"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -393,16 +428,6 @@ 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
@@ -633,6 +658,15 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: ":man_in_business_suit_levitating: AMIP Target EDOnly nonequil sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/amip_target_edonly_nonequil_sparse_autodiff.yml
+          --job_id amip_target_edonly_nonequil_sparse_autodiff
+        artifact_paths: "amip_target_edonly_nonequil_sparse_autodiff/output_active/*"
+        agents:
+          slurm_mem: 40GB
+
   - group: "Diagnostic EDMFX"
     steps:
 
@@ -759,6 +793,16 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":genie: Diagnostic EDMFX aquaplanet sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/diagnostic_edmfx_aquaplanet_sparse_autodiff.yml
+          --job_id diagnostic_edmfx_aquaplanet_sparse_autodiff
+        artifact_paths: "diagnostic_edmfx_aquaplanet_sparse_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
   - group: "Prognostic EDMFX"
     steps:
 
@@ -943,6 +987,16 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":genie: Prognostic EDMFX aquaplanet sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_aquaplanet_sparse_autodiff.yml
+          --job_id prognostic_edmfx_aquaplanet_sparse_autodiff
+        artifact_paths: "prognostic_edmfx_aquaplanet_sparse_autodiff/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
   - group: "GPU"
     steps:
 
@@ -984,6 +1038,19 @@ steps:
           slurm_gpus: 1
           slurm_mem: 16GB
 
+      - label: "GPU: baroclinic wave sparse autodiff"
+        key: "baroclinic_wave_gpu_sparse_autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/baroclinic_wave_sparse_autodiff.yml
+          --job_id baroclinic_wave_gpu_sparse_autodiff
+        artifact_paths: "baroclinic_wave_gpu_sparse_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
@@ -1085,6 +1152,30 @@ steps:
           slurm_gpus: 1
           slurm_mem: 20GB
 
+      - label: "GPU: Prognostic EDMFX aquaplanet dense autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_aquaplanet_gpu_dense_autodiff.yml
+          --job_id prognostic_edmfx_aquaplanet_gpu_dense_autodiff
+        artifact_paths: "prognostic_edmfx_aquaplanet_gpu_dense_autodiff/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+          slurm_mem: 20GB
+
+      - label: "GPU: Prognostic EDMFX aquaplanet sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_aquaplanet_gpu_sparse_autodiff.yml
+          --job_id prognostic_edmfx_aquaplanet_gpu_sparse_autodiff
+        artifact_paths: "prognostic_edmfx_aquaplanet_gpu_sparse_autodiff/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+          slurm_mem: 20GB
+
   - group: "Benchmarks"
     steps:
 
@@ -1218,6 +1309,15 @@ steps:
         agents:
           slurm_mem: 32GB
 
+      - label: ":fire: Flame graph: prognostic edmf sparse autodiff"
+        command: >
+          julia --color=yes --project=.buildkite perf/flame.jl
+          --config_file $PERF_CONFIG_PATH/bm_aquaplanet_progedmf_sparse_autodiff.yml
+          --job_id flame_aquaplanet_progedmf_sparse_autodiff
+        artifact_paths: "flame_aquaplanet_progedmf_sparse_autodiff/*"
+        agents:
+          slurm_mem: 32GB
+
       - label: ":fire: Flame graph: diffusion"
         command: >
           julia --color=yes --project=.buildkite perf/flame.jl

Project.toml

@@ -30,6 +30,7 @@ NullBroadcasts = "0d71be07-595a-4f89-9529-4065a4ab43a6"
 RRTMGP = "a01a1ee8-cea4-48fc-987c-fc7878d79da1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
@@ -42,7 +43,7 @@ Adapt = "4"
 ArgParse = "1"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
-ClimaComms = "0.6.8"
+ClimaComms = "0.6.9"
 ClimaCore = "0.14.37"
 ClimaDiagnostics = "0.2.12"
 ClimaInterpolations = "0.1.0"
@@ -64,6 +65,7 @@ NullBroadcasts = "0.1"
 RRTMGP = "0.21.3"
 Random = "1"
 SciMLBase = "2.12"
+SparseMatrixColorings = "0.4.20"
 StaticArrays = "1.9"
 Statistics = "1"
 SurfaceFluxes = "0.11, 0.12"

config/default_configs/default_config.yml

@@ -89,9 +89,18 @@ jvp_step_adjustment:
 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
+use_auto_jacobian:
+  help: "Whether to populate the entries of the sparse Jacobian matrix using forward-mode automatic differentiation with sparse matrix coloring (only used when `use_dense_jacobian` is `false`) [`true`, `false` (default)]"
+  value: false
+auto_jacobian_padding_bands:
+  help: "Target number of bands to add in every block of the sparse Jacobian matrix, eliminating errors from Jacobian entries that lie outside of the default sparsity pattern; when unspecified, each block gets a predetermined number of padding bands based on the typical magnitudes of its entries (only used when `use_auto_jacobian` is `true`)"
+  value: ~
 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"
+debug_jacobian:
+  help: "Whether to print summary information about the Jacobian matrix, including comparisons of different algorithms evaluated on the first column of the final state against the exact Jacobian [`true`, `false` (default)]"
+  value: false
 # Radiation
 rad:
   help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]"

config/model_configs/amip_target_edonly_nonequil_sparse_autodiff.yml

@@ -0,0 +1,34 @@
+h_elem: 16
+z_max: 60000.0
+z_elem: 63
+dz_bottom: 30.0
+rayleigh_sponge: true
+viscous_sponge: true
+dt_save_state_to_disk: "30days"
+cloud_model: "grid_scale"
+moist: "nonequil"
+implicit_noneq_cloud_formation: true
+use_auto_jacobian: true
+update_jacobian_every: dt
+debug_jacobian: true
+precip_model: "1M"
+rad: "allskywithclear"
+dt_rad: "1hours"
+dt_cloud_fraction: "1hours"
+insolation: "timevarying"
+co2_model: maunaloa
+prescribe_ozone: true
+aerosol_radiation: true
+edmfx_sgs_diffusive_flux: true
+prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "DST05", "OC1", "OC2", "SO4", "SSLT01", "SSLT02", "SSLT03", "SSLT04", "SSLT05"]
+surface_setup: "DefaultMoninObukhov"
+turbconv: "edonly_edmfx"
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
+dt: "120secs"
+t_end: "10mins"
+toml: [toml/short_aquaplanet_nonequil.toml]
+diagnostics:
+  - short_name: [ta, ua, wa, va, rhoa, hur, hus, hussn, husra, clw, cli]
+    period: 1hours
+netcdf_output_at_levels: true

config/model_configs/aquaplanet_nonequil_allsky_gw_res_sparse_autodiff.yml

@@ -0,0 +1,31 @@
+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_auto_jacobian: true
+update_jacobian_every: dt
+debug_jacobian: true
+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_equil_conservation_ft64_sparse_autodiff.yml

@@ -0,0 +1,16 @@
+FLOAT_TYPE: "Float64"
+initial_condition: "MoistBaroclinicWave"
+dt: "400secs"
+t_end: "5days"
+moist: "equil"
+disable_surface_flux_tendency: true
+use_auto_jacobian: true
+update_jacobian_every: dt
+debug_jacobian: true
+dt_save_state_to_disk: "5days"
+check_conservation: true
+diagnostics:
+  - short_name: [massa, energya]
+    period: 1days
+    writer: dict
+use_itime: true

config/model_configs/baroclinic_wave_sparse_autodiff.yml

@@ -0,0 +1,14 @@
+dt_save_state_to_disk: "2days"
+initial_condition: "DryBaroclinicWave"
+use_auto_jacobian: true
+update_jacobian_every: dt
+debug_jacobian: true
+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