Update for ClimaCore geometry refactor

src/cache/temporary_quantities.jl

@@ -83,7 +83,7 @@ function temporary_quantities(Y, atmos)
             face_space,
         ), # ᶠω¹²ʲs
         ᶜbidiagonal_adjoint_matrix_c3 = Fields.Field(
-            BidiagonalMatrixRow{Adjoint{FT, C3{FT}}},
+            BidiagonalMatrixRow{typeof(C3(FT(0))')},
             center_space,
         ),
         ᶜtemp_bdmr = similar(Y.c, BidiagonalMatrixRow{FT}),
@@ -99,8 +99,8 @@ function temporary_quantities(Y, atmos)
         # TODO: Remove this hack
         sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ
         # Implicit solver cache:
-        ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{Adjoint{FT, CT12{FT}}}),
-        ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}),
+        ∂ᶜK_∂ᶜuₕ = similar(Y.c, DiagonalMatrixRow{typeof(CT12(FT(0), FT(0))')}),
+        ∂ᶜK_∂ᶠu₃ = similar(Y.c, BidiagonalMatrixRow{typeof(CT3(FT(0))')}),
         ᶠp_grad_matrix = similar(Y.f, BidiagonalMatrixRow{C3{FT}}),
         ᶠbidiagonal_matrix_ct3 = similar(Y.f, BidiagonalMatrixRow{CT3{FT}}),
         ᶠband_matrix_wvec = similar(
@@ -119,7 +119,7 @@ function temporary_quantities(Y, atmos)
                 ClimaCore.Geometry.WVector{FT},
             },
         ),
-        ᶜtracer_advection_matrix = similar(Y.c, BidiagonalMatrixRow{Adjoint{FT, C3{FT}}}),
+        ᶜtracer_advection_matrix = similar(Y.c, BidiagonalMatrixRow{typeof(C3(FT(0))')}),
         ᶠbidiagonal_matrix_ct3_2 = similar(Y.f, BidiagonalMatrixRow{CT3{FT}}),
         ᶠbidiagonal_matrix_ct3xct12 = similar(
             Y.f,
@@ -137,7 +137,7 @@ function temporary_quantities(Y, atmos)
         ),
         ᶜadvection_matrix = similar(
             Y.c,
-            BidiagonalMatrixRow{Adjoint{FT, C3{FT}}},
+            BidiagonalMatrixRow{typeof(C3(FT(0))')},
         ),
         ᶜtridiagonal_matrix = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶜdiffusion_h_matrix = similar(Y.c, TridiagonalMatrixRow{FT}),

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -66,8 +66,8 @@ function jacobian_cache(alg::ManualSparseJacobian, Y, atmos)
     DiagonalRow = DiagonalMatrixRow{FT}
     TridiagonalRow = TridiagonalMatrixRow{FT}
     BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}}
-    TridiagonalRow_ACT12 = TridiagonalMatrixRow{Adjoint{FT, CT12{FT}}}
-    BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
+    TridiagonalRow_ACT12 = TridiagonalMatrixRow{typeof(CT12(FT(0), FT(0))')}
+    BidiagonalRow_ACT3 = BidiagonalMatrixRow{typeof(CT3(FT(0))')}
     BidiagonalRow_C3xACT12 =
         BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT12{FT})')}
     TridiagonalRow_C3xACT3 =
@@ -381,6 +381,70 @@ end
 
 # TODO: There are a few for loops in this function. This is because
 # using unrolled_foreach allocates (breaks the flame tests)
+# Extracted from `update_jacobian!` as a function barrier. The parent function
+# is ~1000 lines, which exhausts Julia's per-function type-inference budget:
+# past the exhaustion point, matrix-product element types widen to `Any`, and
+# `multiply_matrix_at_index` then throws on `rzero(Any)`. A separate function
+# gets a fresh inference budget; `@noinline` guarantees the barrier is not
+# inlined back into the parent.
+@noinline function update_dycore_jacobian_blocks!(Y, p, cache, dtγ)
+    (; topography_flag) = cache.derivative_flags
+    (; matrix) = cache
+    (; ᶜh_tot) = p.precomputed
+    (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p.scratch
+    ᶜρ = Y.c.ρ
+    ᶜuₕ = Y.c.uₕ
+    ᶠu₃ = Y.f.u₃
+    ᶜJ = Fields.local_geometry_field(Y.c).J
+    ᶠJ = Fields.local_geometry_field(Y.f).J
+    ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ
+    ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ
+
+    if use_derivative(topography_flag)
+        @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(
+            adjoint(CT12(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ),
+        )
+    else
+        @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CT12(ᶜuₕ)))
+    end
+    @. ∂ᶜK_∂ᶠu₃ =
+        ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) +
+        DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix()
+
+    @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix()
+
+    @. ᶜadvection_matrix =
+        -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ)
+    @. p.scratch.ᶠbidiagonal_matrix_ct3xct12 =
+        ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ))
+    if use_derivative(topography_flag)
+        ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)]
+        @. ∂ᶜρ_err_∂ᶜuₕ =
+            dtγ * ᶜadvection_matrix ⋅ p.scratch.ᶠbidiagonal_matrix_ct3xct12
+    end
+    ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)]
+    @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ))
+
+    tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot))
+
+    MatrixFields.unrolled_foreach(tracer_info) do ρχ_name
+        MatrixFields.has_field(Y, ρχ_name) || return
+        ᶜχ = ρχ_name === @name(c.ρe_tot) ? ᶜh_tot : (@. lazy(specific(Y.c.ρq_tot, Y.c.ρ)))
+
+        if use_derivative(topography_flag)
+            ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)]
+            @. ∂ᶜρχ_err_∂ᶜuₕ =
+                dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅
+                p.scratch.ᶠbidiagonal_matrix_ct3xct12
+        end
+
+        ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)]
+        @. ∂ᶜρχ_err_∂ᶠu₃ =
+            dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ))
+    end
+    return nothing
+end
+
 function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
     (;
         topography_flag,
@@ -451,48 +515,12 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
     ᶜ∂p∂ρq_tot = p.scratch.ᶜtemp_scalar_2
     @. ᶜ∂p∂ρq_tot = ᶜkappa_m * (-e_int_v0 - R_d * T_0 - Δcv_v * (ᶜT - T_0)) + ΔR_v * ᶜT
 
-    if use_derivative(topography_flag)
-        @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(
-            adjoint(CT12(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ),
-        )
-    else
-        @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CT12(ᶜuₕ)))
-    end
-    @. ∂ᶜK_∂ᶠu₃ =
-        ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) +
-        DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix()
-
-    @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix()
-
-    @. ᶜadvection_matrix =
-        -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ)
-    @. p.scratch.ᶠbidiagonal_matrix_ct3xct12 =
-        ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ))
-    if use_derivative(topography_flag)
-        ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)]
-        @. ∂ᶜρ_err_∂ᶜuₕ =
-            dtγ * ᶜadvection_matrix ⋅ p.scratch.ᶠbidiagonal_matrix_ct3xct12
-    end
-    ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)]
-    @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ))
-
-    tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot))
-
-    MatrixFields.unrolled_foreach(tracer_info) do ρχ_name
-        MatrixFields.has_field(Y, ρχ_name) || return
-        ᶜχ = ρχ_name === @name(c.ρe_tot) ? ᶜh_tot : (@. lazy(specific(Y.c.ρq_tot, Y.c.ρ)))
-
-        if use_derivative(topography_flag)
-            ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)]
-            @. ∂ᶜρχ_err_∂ᶜuₕ =
-                dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅
-                p.scratch.ᶠbidiagonal_matrix_ct3xct12
-        end
-
-        ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)]
-        @. ∂ᶜρχ_err_∂ᶠu₃ =
-            dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ))
-    end
+    # Dycore matrix blocks (∂ᶜK/∂*, pressure-gradient, ρ/ρχ continuity advection)
+    # are computed in a separate function as an inference-budget barrier — see
+    # the `update_dycore_jacobian_blocks!` comment. The scratch fields and
+    # `matrix` blocks it fills are mutated in place, so the bindings below still
+    # see the results.
+    update_dycore_jacobian_blocks!(Y, p, cache, dtγ)
 
     ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)]
     ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)]

src/prognostic_equations/surface_flux.jl

@@ -21,7 +21,7 @@ Arguments:
 - `ᶜρ`: Cell-center air density field.
 - `ᶜuₕ`: Cell-center horizontal velocity field (used for type/structure, not value in flux calc).
 - `ρ_flux_uₕ_surface`: The vertical flux of horizontal momentum through the bottom
-  boundary. This is a `ClimaCore.Geometry.AxisTensor` of type
+  boundary. This is a `ClimaCore.Geometry.Tensor` of type
   `C3{FT} ⊗ C12{FT}` (e.g., representing surface stress `τ` as
   `e_3 ⊗ τ` if defined as flux into the domain, or simply
   the stress vector `τ` if the `SetValue` operator handles the normal).

src/prognostic_equations/zero_velocity.jl

@@ -87,7 +87,7 @@ to either a scaled negative identity matrix (if `row_name == col_name`, i.e.,
 a diagonal block) or zero (if `row_name != col_name`, i.e., an off-diagonal block).
 
 Specifically, for diagonal blocks, it sets the entry to represent `-I` (negative
-identity), where `I` is an `AxisTensor` identity of the appropriate type and structure.
+identity), where `I` is a `Tensor` identity of the appropriate type and structure.
 This is used within `zero_velocity_jacobian!` to modify Jacobian contributions
 related to velocity variables during an advection test.
 
@@ -99,10 +99,7 @@ Arguments:
 """
 function set_identity_matrix_entry!(matrix_entry, row_name, col_name)
     identity_matrix_entry_value = if row_name == col_name
-        # TODO: Add a method for one(::Axis2Tensor) to simplify this.
-        T = eltype(eltype(matrix_entry))
-        tensor_data = UniformScaling(one(eltype(T)))
-        -DiagonalMatrixRow(Geometry.AxisTensor(axes(T), tensor_data))
+        -DiagonalMatrixRow(one(eltype(eltype(matrix_entry))))
     else
         zero(eltype(matrix_entry))
     end

src/solver/type_getters.jl

@@ -742,6 +742,10 @@ function get_grid(parsed_args, params, context)
             bubble = false,
             periodic_x = true,
             periodic_y = true,
+            z_elem = parsed_args["z_elem"],
+            z_max = parsed_args["z_max"],
+            z_stretch = parsed_args["z_stretch"],
+            dz_bottom = parsed_args["dz_bottom"],
             kwargs...,
         )
     elseif config == "box"

src/utils/utilities.jl

@@ -230,7 +230,7 @@ end
     g³³_field(space)
 
 Extracts the value of `g³³`, the 3rd component of the metric terms that convert
-Covariant AxisTensors to Contravariant AxisTensors, from the given space.
+Covariant Tensors to Contravariant Tensors, from the given space.
 """
 function g³³_field(space)
     g_field = Fields.local_geometry_field(space).gⁱʲ.components.data
@@ -243,10 +243,13 @@ end
 
 Extracts the `g³³` sub-tensor from the `gⁱʲ` tensor.
 """
-g³³(gⁱʲ) = Geometry.AxisTensor(
-    (Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()),
-    Geometry.components(gⁱʲ)[end],
-)
+function g³³(gⁱʲ)
+    FT = eltype(parent(gⁱʲ))
+    Geometry.Tensor(
+        SMatrix{1, 1, FT, 1}(parent(gⁱʲ)[end]),
+        (Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()),
+    )
+end
 
 
 """
@@ -270,6 +273,9 @@ function g³ʰ(gⁱʲ)
     elseif full_CT_axis == Geometry.Contravariant23Axis()
         @inbounds val = gⁱʲ_components[N, 1]
         SMatrix{1, 2, FT, 2}(zero(FT), val)
+    elseif full_CT_axis == Geometry.Contravariant3Axis()
+        # Single-column geometry: no horizontal axes, no off-diagonal terms
+        SMatrix{1, 2, FT, 2}(zero(FT), zero(FT))
     else
         error("$full_CT_axis is missing either vertical or horizontal sub-axes")
     end

src/utils/variable_manipulations.jl

@@ -486,7 +486,7 @@ A wrapper for Julia's `mapreduce` function that automatically determines
 the initial value (`init`) for the reduction.
 
 This is useful for iterators whose elements are custom structs or
-`ClimaCore.Geometry.AxisTensor`s, where the zero element cannot be inferred
+`ClimaCore.Geometry.AbstractTensor`s, where the zero element cannot be inferred
 as a simple scalar. It uses `ClimaCore.RecursiveApply` tools (`rzero`,
 `rpromote_type`) to create a type-stable, correctly-structured zero element
 based on the output of the function `f` applied to the first elements of the
@@ -514,7 +514,7 @@ failures.
 
 It uses `ClimaCore.RecursiveApply` operators (`⊞` for addition, `⊠` for
 multiplication), which allows it to handle dot products of tuples containing
-complex, nested types such as `ClimaCore.Geometry.AxisTensor`s.
+complex, nested types such as `ClimaCore.Geometry.AbstractTensor`s.
 
 Arguments:
 - `a`: The first `Tuple`.