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"
@@ -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

@@ -45,6 +45,71 @@ 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
 
+# TODO: Add support for manual Jacobian debugging with the AutoDenseJacobian.
+integrator_has_auto_sparse_jacobian =
+    config.parsed_args["use_auto_jacobian"] &&
+    !config.parsed_args["use_dense_jacobian"]
+debug_manual_jacobian =
+    isnothing(config.parsed_args["debug_manual_jacobian"]) ?
+    config.parsed_args["use_auto_jacobian"] :
+    config.parsed_args["debug_manual_jacobian"]
+if integrator_has_auto_sparse_jacobian && debug_manual_jacobian
+    Y_end = integrator.u
+    t_end = integrator.t
+    dt = integrator.dt
+    timestepper_algorithm = integrator.alg
+    tableau_coefficients =
+        timestepper_algorithm isa CA.CTS.RosenbrockAlgorithm ?
+        timestepper_algorithm.tableau.Γ : timestepper_algorithm.tableau.a_imp
+
+    FT = eltype(Y_end)
+    γs = filter(!iszero, CA.LinearAlgebra.diag(tableau_coefficients))
+    dtγ = FT(float(dt) * γs[end])
+
+    auto_jacobian =
+        timestepper_algorithm isa CA.CTS.RosenbrockAlgorithm ?
+        integrator.cache.W : integrator.cache.newtons_method_cache.j
+    manual_jacobian = CA.Jacobian(auto_jacobian.alg.sparse_alg, Y_end, atmos)
+    CA.update_jacobian!(auto_jacobian, Y_end, p, dtγ, t_end)
+    CA.update_jacobian!(manual_jacobian, Y_end, p, dtγ, t_end)
+    auto_matrix = auto_jacobian.cache.matrix.matrix
+    manual_matrix = manual_jacobian.cache.matrix.matrix
+    auto_scalar_matrix = CA.MatrixFields.scalar_field_matrix(auto_matrix)
+    manual_scalar_matrix = CA.MatrixFields.scalar_field_matrix(manual_matrix)
+    diff_scalar_matrix = auto_scalar_matrix .- manual_scalar_matrix
+
+    @info "Debugging manual Jacobian"
+    for scalar_block_name in keys(auto_scalar_matrix)
+        auto_block = auto_scalar_matrix[scalar_block_name]
+        manual_block = manual_scalar_matrix[scalar_block_name]
+        diff_block = diff_scalar_matrix[scalar_block_name]
+
+        auto_block isa CA.Fields.Field || continue
+
+        println("$scalar_block_name:")
+        println("\t$(eltype(auto_block))")
+        (_, _, lower_band, upper_band) =
+            CA.MatrixFields.band_matrix_info(auto_block)
+        for band in lower_band:upper_band
+            band_index = band - lower_band + 1
+            auto_band_average = mean(abs, auto_block.entries.:($band_index))
+            manual_band_average = mean(abs, manual_block.entries.:($band_index))
+            diff_band_average = mean(abs, diff_block.entries.:($band_index))
+            averages =
+                (auto_band_average, manual_band_average, diff_band_average)
+            percent_error =
+                diff_band_average == auto_band_average == 0 ? FT(0) :
+                diff_band_average / auto_band_average * 100
+            rounded_averages = map(x -> round(x; sigdigits = 2), averages)
+            rounded_percent_error = round(percent_error; sigdigits = 2)
+            println(
+                "\tBand $band averages (auto/manual/difference) and error: \
+                   $(join(rounded_averages, '/')), $rounded_percent_error%",
+            )
+        end
+    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_manual_jacobian:
+  help: "Whether to compare the bands of the Jacobian matrix that is computed manually to those of the Jacobian matrix generated with automatic differentiation (default is equal to `use_auto_jacobian`) [`true`, `false`]"
+  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"

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"))

src/prognostic_equations/implicit/auto_dense_jacobian.jl

@@ -89,45 +89,36 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, dtγ, t)
     device = ClimaComms.device(Y.c)
     column_indices = column_index_iterator(Y)
     scalar_names = scalar_field_names(Y)
-    scalar_level_indices = scalar_level_index_pairs(Y)
-    batch_size = max_simultaneous_derivatives(alg)
-    batch_size_val = Val(batch_size)
-
-    p_dual_args = ntuple(Val(fieldcount(typeof(p)))) do cache_field_index
-        cache_field_name = fieldname(typeof(p), cache_field_index)
-        if cache_field_name == :precomputed
-            (; p.precomputed..., precomputed_dual...)
-        elseif cache_field_name == :scratch
-            scratch_dual
-        else
-            getfield(p, cache_field_index)
-        end
-    end
-    p_dual = AtmosCache(p_dual_args...)
-
-    batches = Iterators.partition(scalar_level_indices, batch_size)
-    for batch_scalar_level_indices in ClimaComms.threadable(device, batches)
+    jacobian_axis_index_to_field_vector_index_map =
+        enumerate(field_vector_index_iterator(Y))
+    n_εs = max_simultaneous_derivatives(alg)
+    n_εs_val = Val(n_εs)
+    p_dual = replace_precomputed_and_scratch(p, precomputed_dual, scratch_dual)
+
+    batches =
+        Iterators.partition(jacobian_axis_index_to_field_vector_index_map, n_εs)
+    for indices_for_Y_axis in ClimaComms.threadable(device, batches)
         Y_dual .= Y
 
         # Add a unique ε to Y for each scalar level index in this batch. With
         # Y_col and Yᴰ_col denoting the columns of Y and Y_dual at column_index,
-        # set Yᴰ_col to Y_col + I[:, batch_scalar_level_indices] * εs, where I
-        # is the identity matrix for Y_col (i.e., the value of ∂Y_col/∂Y_col),
-        # εs is a vector of batch_size dual number components, and
-        # batch_scalar_level_indices are the batch's indices into Y_col.
+        # set Yᴰ_col to Y_col + I[:, indices_for_Y_axis] * εs, where I is the
+        # identity matrix for Y_col (i.e., the value of ∂Y_col/∂Y_col), εs is a
+        # vector of n_εs dual number components, and indices_for_Y_axis are the
+        # batch's indices into Y_col.
         ClimaComms.@threaded device begin
             # On multithreaded devices, assign one thread to each combination of
             # spatial column index and scalar level index in this batch.
             for column_index in column_indices,
                 (ε_index, (_, (scalar_index, level_index))) in
-                enumerate(batch_scalar_level_indices)
+                enumerate(indices_for_Y_axis)
 
-                Y_partials = ntuple(i -> i == ε_index ? 1 : 0, batch_size_val)
-                Y_dual_increment = ForwardDiff.Dual{Jacobian}(0, Y_partials...)
+                Y_partials = ntuple(==(ε_index), n_εs_val)
+                Y_dual_εs_value = ForwardDiff.Dual{Jacobian}(0, Y_partials)
                 unrolled_applyat(scalar_index, scalar_names) do name
                     field = MatrixFields.get_field(Y_dual, name)
                     @inbounds point(field, level_index, column_index...)[] +=
-                        Y_dual_increment
+                        Y_dual_εs_value
                 end
             end
         end
@@ -141,19 +132,19 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, dtγ, t)
         # with col_matrix denoting the matrix at the corresponding matrix_index
         # in column_matrices, copy the coefficients of the εs in Yₜᴰ_col into
         # col_matrix, where the previous steps have set Yₜᴰ_col to
-        # Yₜ_col + (∂Yₜ_col/∂Y_col)[:, batch_scalar_level_indices] * εs.
-        # Specifically, set col_matrix[scalar_level_index1, scalar_level_index2]
-        # to ∂Yₜ_col[scalar_level_index1]/∂Y_col[scalar_level_index2], obtaining
+        # Yₜ_col + (∂Yₜ_col/∂Y_col)[:, indices_for_Y_axis] * εs. Specifically, set
+        # col_matrix[scalar_level_index1, scalar_level_index2] to
+        # ∂Yₜ_col[scalar_level_index1]/∂Y_col[scalar_level_index2], obtaining
         # this derivative from the coefficient of εs[ε_index] in
         # Yₜᴰ_col[scalar_level_index1], where ε_index is the index of
-        # scalar_level_index2 in batch_scalar_level_indices. After all batches
-        # have been processed, col_matrix is the full Jacobian ∂Yₜ_col/∂Y_col.
+        # scalar_level_index2 in indices_for_Y_axis. After all batches have been
+        # processed, col_matrix is the full Jacobian ∂Yₜ_col/∂Y_col.
         ClimaComms.@threaded device begin
             # On multithreaded devices, assign one thread to each combination of
             # spatial column index and scalar level index.
             for (matrix_index, column_index) in enumerate(column_indices),
                 (scalar_level_index1, (scalar_index1, level_index1)) in
-                scalar_level_indices
+                jacobian_axis_index_to_field_vector_index_map
 
                 Yₜ_dual_value =
                     unrolled_applyat(scalar_index1, scalar_names) do name
@@ -162,7 +153,7 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, dtγ, t)
                     end
                 Yₜ_partials = ForwardDiff.partials(Yₜ_dual_value)
                 for (ε_index, (scalar_level_index2, _)) in
-                    enumerate(batch_scalar_level_indices)
+                    enumerate(indices_for_Y_axis)
                     cartesian_index =
                         (scalar_level_index1, scalar_level_index2, matrix_index)
                     @inbounds column_matrices[cartesian_index...] =
@@ -193,14 +184,15 @@ function invert_jacobian!(::AutoDenseJacobian, cache, ΔY, R)
     device = ClimaComms.device(ΔY.c)
     column_indices = column_index_iterator(ΔY)
     scalar_names = scalar_field_names(ΔY)
-    scalar_level_indices = scalar_level_index_pairs(ΔY)
+    vector_index_to_field_vector_index_map =
+        enumerate(field_vector_index_iterator(ΔY))
 
     # Copy all scalar values from R into column_lu_vectors.
     ClimaComms.@threaded device begin
         # On multithreaded devices, assign one thread to each index into R.
         for (vector_index, column_index) in enumerate(column_indices),
             (scalar_level_index, (scalar_index, level_index)) in
-            scalar_level_indices
+            vector_index_to_field_vector_index_map
 
             value = unrolled_applyat(scalar_index, scalar_names) do name
                 field = MatrixFields.get_field(R, name)
@@ -219,7 +211,7 @@ function invert_jacobian!(::AutoDenseJacobian, cache, ΔY, R)
         # On multithreaded devices, assign one thread to each index into ΔY.
         for (vector_index, column_index) in enumerate(column_indices),
             (scalar_level_index, (scalar_index, level_index)) in
-            scalar_level_indices
+            vector_index_to_field_vector_index_map
 
             @inbounds value =
                 column_lu_vectors[scalar_level_index, vector_index]