CliMA · dennisYatunin · Jun 14, 2025
diff --git a/.buildkite/Manifest-v1.11.toml b/.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", "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", "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.30.3"
@@ -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 = "c6ab151ea66f3756566abc039c76ae767e490446"
+git-tree-sha1 = "371c3801d16438113f22e3db3d6da69cee5b3a0b"
+repo-rev = "tr/refactor-fm-internal-index"
+repo-url = "https://github.com/CliMA/ClimaCore.jl.git"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.34"
+version = "0.14.33"
 weakdeps = ["CUDA", "Krylov"]
 
     [deps.ClimaCore.extensions]
@@ -503,11 +505,6 @@ git-tree-sha1 = "37ea44092930b1811e666c3bc38065d7d87fcc74"
 uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
 version = "0.13.1"
 
-[[deps.Combinatorics]]
-git-tree-sha1 = "8010b6bb3388abe68d95743dcbea77650bb2eddf"
-uuid = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
-version = "1.0.3"
-
 [[deps.CommonDataModel]]
 deps = ["CFTime", "DataStructures", "Dates", "Preferences", "Printf", "Statistics"]
 git-tree-sha1 = "358bf5a7d5c1387b995a43577673290c5d344758"
@@ -1027,12 +1024,6 @@ git-tree-sha1 = "53bb909d1151e57e2484c3d1b53e19552b887fb2"
 uuid = "42e2da0e-8278-4e71-bc24-59509adca0fe"
 version = "1.0.2"
 
-[[deps.HCubature]]
-deps = ["Combinatorics", "DataStructures", "LinearAlgebra", "QuadGK", "StaticArrays"]
-git-tree-sha1 = "19ef9f0cb324eed957b7fe7257ac84e8ed8a48ec"
-uuid = "19dc6840-f33b-545b-b366-655c7e3ffd49"
-version = "1.7.0"
-
 [[deps.HDF5]]
 deps = ["Compat", "HDF5_jll", "Libdl", "MPIPreferences", "Mmap", "Preferences", "Printf", "Random", "Requires", "UUIDs"]
 git-tree-sha1 = "e856eef26cf5bf2b0f95f8f4fc37553c72c8641c"
@@ -2306,6 +2297,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"

diff --git a/.buildkite/ci_driver.jl b/.buildkite/ci_driver.jl
@@ -45,6 +45,63 @@ 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 true # TODO: Add a debug_manual_jacobian flag.
+    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
+    γs = unique(filter(!iszero, CA.LinearAlgebra.diag(tableau_coefficients)))
+    dtγ = float(dt) * γs[end]
+
+    auto_jac_alg = integrator.cache.newtons_method_cache.j.alg
+    manual_jac_alg = auto_jac_alg.sparse_alg
+    auto_jac = CA.Jacobian(auto_jac_alg, Y_end, atmos)
+    manual_jac = CA.Jacobian(manual_jac_alg, Y_end, atmos)
+    CA.update_jacobian!(auto_jac, Y_end, p, dtγ, t_end)
+    CA.update_jacobian!(manual_jac, Y_end, p, dtγ, t_end)
+    auto_matrix = auto_jac.cache.matrix.matrix
+    manual_matrix = manual_jac.cache.matrix.matrix
+    auto_scalar_matrix = CA.MatrixFields.scalar_field_matrix(auto_matrix)
+    manual_scalar_matrix = CA.MatrixFields.scalar_field_matrix(manual_matrix)
+    difference_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]
+        difference_block = difference_scalar_matrix[scalar_block_name]
+
+        auto_block isa CA.Fields.Field || continue
+
+        println("$scalar_block_name:")
+        println("\t$(eltype(auto_block))")
+
+        (_, _, auto_lower_band, auto_upper_band) =
+            CA.MatrixFields.band_matrix_info(auto_block)
+        (_, _, manual_lower_band, manual_upper_band) =
+            CA.MatrixFields.band_matrix_info(manual_block)
+        for band in auto_lower_band:auto_upper_band
+            auto_band_index = band - auto_lower_band + 1
+            auto_band_average =
+                mean(abs, auto_block.entries.:($auto_band_index))
+            is_manual = band in manual_lower_band:manual_upper_band
+            println("\tAverage of$(is_manual ? "" : " padding") band $band:")
+            println("\t\t$auto_band_average (auto)")
+            is_manual || continue
+            manual_band_index = band - manual_lower_band + 1
+            manual_band_average =
+                mean(abs, manual_block.entries.:($manual_band_index))
+            difference_band_average =
+                mean(abs, difference_block.entries.:($auto_band_index))
+            println("\t\t$manual_band_average (manual)")
+            println("\t\t$difference_band_average (difference)")
+        end
+    end
+end
+
 if sol_res.ret_code == :simulation_crashed
     error(
         "The ClimaAtmos simulation has crashed. See the stack trace for details.",

diff --git a/.buildkite/pipeline.yml b/.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=\"tr/refactor-fm-internal-index\"))'"
 
     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:

diff --git a/Project.toml b/Project.toml
@@ -29,6 +29,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 @@ ArgParse = "1"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
 ClimaComms = "0.6.8"
-ClimaCore = "0.14.34"
+ClimaCore = "0.14.33"
 ClimaDiagnostics = "0.2.12"
 ClimaParams = "0.10.32"
 ClimaTimeSteppers = "0.8.2"
@@ -62,6 +63,7 @@ NullBroadcasts = "0.1"
 RRTMGP = "0.21.2"
 Random = "1"
 SciMLBase = "2.12"
+SparseMatrixColorings = "0.4.20"
 StaticArrays = "1.7"
 Statistics = "1"
 SurfaceFluxes = "0.11, 0.12"

diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
@@ -61,6 +61,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"))

diff --git a/src/prognostic_equations/implicit/auto_dense_jacobian.jl b/src/prognostic_equations/implicit/auto_dense_jacobian.jl
@@ -89,9 +89,10 @@ 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)
+    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_args = ntuple(Val(fieldcount(typeof(p)))) do cache_field_index
         cache_field_name = fieldname(typeof(p), cache_field_index)
@@ -105,29 +106,30 @@ function update_column_matrices!(alg::AutoDenseJacobian, cache, Y, p, dtγ, t)
     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)
+    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 +143,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 +164,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 +195,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 +222,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]