Add AutoSparseJacobian algorithm for implicit solver

Author: dennisYatunin

.buildkite/Manifest-v1.11.toml

@@ -363,7 +363,7 @@ 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"
@@ -381,9 +381,11 @@ weakdeps = ["CUDA", "MPI"]
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
-git-tree-sha1 = "14d3d5810ce1e3c990450a2ce7abc6a1e162855f"
+git-tree-sha1 = "a273452127dbb052f2963e3c1095730a996d49a6"
+repo-rev = "main"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.35"
+version = "0.14.36"
 weakdeps = ["CUDA", "Krylov"]
 
     [deps.ClimaCore.extensions]
@@ -2318,6 +2320,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

@@ -45,6 +45,133 @@ 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_approximate_jacobian"] &&
+    !config.parsed_args["use_dense_jacobian"]
+)
+    Y_end = integrator.u
+    t_end = integrator.t
+    (; p, dt) = integrator
+    timestepper_alg = integrator.alg
+    tableau_coefficients =
+        timestepper_alg isa CA.CTS.RosenbrockAlgorithm ?
+        timestepper_alg.tableau.Γ : timestepper_alg.tableau.a_imp
+    jacobian =
+        timestepper_alg isa CA.CTS.RosenbrockAlgorithm ? integrator.cache.W :
+        integrator.cache.newtons_method_cache.j
+
+    FT = eltype(Y_end)
+    γs = filter(!iszero, CA.LinearAlgebra.diag(tableau_coefficients))
+    dtγ = FT(float(dt) * γs[end])
+    scalar_names = CA.scalar_field_names(Y_end)
+    block_keys = Iterators.product(scalar_names, scalar_names)
+
+    @info "Debugging Jacobian in first column of final state"
+    rms(block) = sqrt(mean(abs2.(block)))
+    highlighters = (
+        Highlighter((d, i, j) -> d[i, j] < 1e-12; foreground = :dark_gray),
+        Highlighter((d, i, j) -> 1e-12 <= d[i, j] < 1e-6; foreground = :green),
+        Highlighter((d, i, j) -> 1e-6 <= d[i, j] < 1e-2; foreground = :blue),
+        Highlighter((d, i, j) -> 1e-2 <= d[i, j] < 1e-1; foreground = :cyan),
+        Highlighter((d, i, j) -> 1e-1 <= d[i, j] < 1; foreground = :light_cyan),
+        Highlighter((d, i, j) -> d[i, j] == 1; foreground = :yellow),
+        Highlighter((d, i, j) -> d[i, j] > 1; foreground = :red),
+    )
+    table_kwargs = (;
+        columns_width = 5,
+        crop = :none,
+        formatters = ft_printf("%1.0e"),
+        highlighters,
+        row_labels = collect(scalar_names),
+        show_header = false,
+        tf = tf_matrix,
+        vlines = [1],
+    )
+
+    block_rescalings = CA.first_column_rescaling_arrays(Y_end, p, t_end)
+
+    exact_jacobian_alg = CA.AutoDenseJacobian()
+    exact_blocks =
+        CA.first_column_block_arrays(exact_jacobian_alg, Y_end, p, dtγ, t_end)
+    exact_rms_values = map(block_keys) do block_key
+        rms(exact_blocks[block_key])
+    end
+    exact_rescaled_rms_values = map(block_keys) do block_key
+        rms(exact_blocks[block_key] .* block_rescalings[block_key])
+    end
+    @info "exact, RMS per block [inconsistent units]:"
+    pretty_table(
+        exact_rms_values;
+        table_kwargs...,
+        highlighters = highlighters[1],
+    )
+    @info "exact, rescaled RMS per block [s^-1]:"
+    pretty_table(
+        exact_rescaled_rms_values;
+        table_kwargs...,
+        highlighters = highlighters[1],
+    )
+    println("<$('='^70)>")
+
+    all_approx_jacobian_algs =
+        jacobian.alg isa CA.AutoSparseJacobian ?
+        (; manual = jacobian.alg.sparse_jacobian_alg, auto = jacobian.alg) :
+        (; manual = jacobian.alg)
+    all_approx_blocks = map(all_approx_jacobian_algs) do jacobian_alg
+        CA.first_column_block_arrays(jacobian_alg, Y_end, p, dtγ, t_end)
+    end
+    if jacobian.alg isa CA.AutoSparseJacobian
+        approx_diff_rescaled_rms_values = map(block_keys) do block_key
+            (; manual, auto) = all_approx_blocks
+            rescaling = block_rescalings[block_key]
+            haskey(manual, block_key) &&
+                !(manual[block_key] isa CA.UniformScaling) ?
+            rms((manual[block_key] - auto[block_key]) .* rescaling) : FT(0)
+        end
+        @info "manual approx - auto approx, rescaled RMS per block [s^-1]:"
+        pretty_table(approx_diff_rescaled_rms_values; table_kwargs...)
+    end
+    for (approx_name, approx_blocks) in pairs(all_approx_blocks)
+        approx_error_rescaled_rms_values = map(block_keys) do block_key
+            approx_error =
+                haskey(approx_blocks, block_key) ?
+                approx_blocks[block_key] - exact_blocks[block_key] :
+                exact_blocks[block_key]
+            rescaling = block_rescalings[block_key]
+            rms(approx_error .* rescaling)
+        end
+        @info "$approx_name approx - exact, rescaled RMS per block [s^-1]:"
+        pretty_table(approx_error_rescaled_rms_values; table_kwargs...)
+    end
+    println("<$('='^70)>")
+    if jacobian.alg isa CA.AutoSparseJacobian
+        approx_diff_relative_rms_values = map(block_keys) do block_key
+            (; manual, auto) = all_approx_blocks
+            rescaling = block_rescalings[block_key]
+            approx_diff_rms_value =
+                haskey(manual, block_key) &&
+                !(manual[block_key] isa CA.UniformScaling) ?
+                rms((manual[block_key] - auto[block_key])) : FT(0)
+            approx_diff_rms_value == 0 ? FT(0) :
+            approx_diff_rms_value / rms(exact_blocks[block_key])
+        end
+        @info "manual approx - auto approx, relative RMS per block [unitless]:"
+        pretty_table(approx_diff_relative_rms_values; table_kwargs...)
+    end
+    for (approx_name, approx_blocks) in pairs(all_approx_blocks)
+        approx_error_relative_rms_values = map(block_keys) do block_key
+            approx_error_rms_value =
+                haskey(approx_blocks, block_key) ?
+                rms(approx_blocks[block_key] - exact_blocks[block_key]) :
+                rms(exact_blocks[block_key])
+            approx_error_rms_value == 0 ? FT(0) :
+            approx_error_rms_value / rms(exact_blocks[block_key])
+        end
+        @info "$approx_name approx - exact, relative RMS per block [unitless]:"
+        pretty_table(approx_error_relative_rms_values; table_kwargs...)
+    end
+end
+
 if sol_res.ret_code == :simulation_crashed
     error(
         "The ClimaAtmos simulation has crashed. See the stack trace for details.",

.buildkite/gpu_pipeline/pipeline.yml

@@ -24,6 +24,7 @@ steps:
       - julia --project=.buildkite -e 'using Pkg; Pkg.precompile()'
       - julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()'
       - julia --project=.buildkite -e 'using Pkg; Pkg.status()'
+      - julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name="ClimaCore", rev="main"))'
 
     agents:
       slurm_gpus: 1

.buildkite/longruns_gpu/pipeline.yml

@@ -25,6 +25,7 @@ steps:
       - julia --project=.buildkite -e 'using Pkg; Pkg.precompile()'
       - julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()'
       - julia --project=.buildkite -e 'using Pkg; Pkg.status()'
+      - julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name="ClimaCore", rev="main"))'
 
     agents:
       slurm_gpus: 1

.buildkite/pipeline.yml

@@ -28,6 +28,7 @@ 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=\"main\"))'"
 
     agents:
       slurm_cpus_per_task: 8
@@ -631,7 +632,7 @@ steps:
           --job_id amip_target_edonly_nonequil
         artifact_paths: "amip_target_edonly_nonequil/output_active/*"
         agents:
-          slurm_mem: 20GB
+          slurm_mem: 64GB
 
   - group: "Diagnostic EDMFX"
     steps:

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"
@@ -64,6 +65,7 @@ NullBroadcasts = "0.1"
 RRTMGP = "0.21.3"
 Random = "1"
 SciMLBase = "2.12"
+SparseMatrixColorings = "0.4.20"
 StaticArrays = "1.7"
 Statistics = "1"
 SurfaceFluxes = "0.11, 0.12"

config/default_configs/default_config.yml

@@ -89,6 +89,15 @@ 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: true # TODO: Change this to false
+auto_jacobian_padding_bands:
+  help: "Minimum 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 (only used when `use_auto_jacobian` is `true`; default is `0`)"
+  value: 2 # TODO: Change this to 0
+debug_approximate_jacobian:
+  help: "Whether to compare approximations of the Jacobian matrix in the first column of the final state against the exact Jacobian [`true`, `false` (default)]"
+  value: true # TODO: Change this to 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"

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-248
+249
 
 # **README**
 #

src/ClimaAtmos.jl

@@ -62,6 +62,7 @@ include(
     joinpath("prognostic_equations", "implicit", "manual_sparse_jacobian.jl"),
 )
 include(joinpath("prognostic_equations", "implicit", "auto_dense_jacobian.jl"))
+include(joinpath("prognostic_equations", "implicit", "auto_sparse_jacobian.jl"))
 include(joinpath("prognostic_equations", "implicit", "autodiff_utils.jl"))
 
 include(joinpath("prognostic_equations", "water_advection.jl"))