Add sedimentation of rhoa, q_tot and moisture tracers to edmf. Working RICO

trontrytel · trontrytel · commit 19ce5d408cdc · 2025-06-05T17:20:27.000-07:00
diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml
@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.11.4"
+julia_version = "1.11.3"
 manifest_format = "2.0"
 project_hash = "04945e2a7116cfacb2d57340f477dbe23e208227"
 
@@ -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", "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", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "PrettyTables", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.30.3"
@@ -1808,7 +1808,7 @@ version = "2.5.4+0"
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
-version = "0.8.1+4"
+version = "0.8.1+2"
 
 [[deps.OpenMPI_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"]
diff --git a/Project.toml b/Project.toml
@@ -25,6 +25,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 NullBroadcasts = "0d71be07-595a-4f89-9529-4065a4ab43a6"
+PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 RRTMGP = "a01a1ee8-cea4-48fc-987c-fc7878d79da1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -57,6 +58,7 @@ Logging = "1"
 NCDatasets = "0.14.2"
 NVTX = "0.3, 1"
 NullBroadcasts = "0.1"
+PrettyTables = "2.4.0"
 RRTMGP = "0.21.2"
 Random = "1"
 SciMLBase = "2.12"
diff --git a/config/model_configs/prognostic_edmfx_rico_column.yml b/config/model_configs/prognostic_edmfx_rico_column.yml
@@ -15,7 +15,7 @@ edmfx_nh_pressure: true
 edmfx_filter: true
 prognostic_tke: true
 moist: "nonequil"
-cloud_model: "quadrature_sgs"
+cloud_model: "grid_scale" #"quadrature_sgs"
 call_cloud_diagnostics_per_stage: true
 precip_model: "1M"
 config: "column"
@@ -25,8 +25,8 @@ y_elem: 2
 z_elem: 100
 z_stretch: false
 perturb_initstate: false
-dt: "10secs"
-t_end: "6hours"
+dt: "2secs"
+t_end: "6hours" #"1200secs" #"6hours"
 dt_save_state_to_disk: "10mins"
 toml: [toml/prognostic_edmfx_1M.toml]
 netcdf_interpolation_num_points: [8, 8, 100]
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -151,6 +151,10 @@ function precomputed_quantities(Y, atmos)
             ᶜSqᵢᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqᵣᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqₛᵖʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₗʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwᵢʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwᵣʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₛʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqₗᵖ⁰ = similar(Y.c, FT),
             ᶜSqᵢᵖ⁰ = similar(Y.c, FT),
             ᶜSqᵣᵖ⁰ = similar(Y.c, FT),
diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -608,13 +608,60 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
     (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρ⁰, ᶜts⁰) = p.precomputed
     (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed
 
+    (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs) = p.precomputed
+
     # TODO - can I re-use them between js and env?
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
 
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
+    FT = eltype(params)
 
     for j in 1:n
+
+        # compute terminal velocity for precipitation
+        @. ᶜwᵣʲs.:($$j) = CM1.terminal_velocity(
+            cmp.pr,
+            cmp.tv.rain,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_rai),
+        )
+        @. ᶜwₛʲs.:($$j) = CM1.terminal_velocity(
+            cmp.ps,
+            cmp.tv.snow,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_sno),
+        )
+        # compute sedimentation velocity for cloud condensate [m/s]
+        @. ᶜwₗʲs.:($$j) = CMNe.terminal_velocity(
+            cmc.liquid,
+            cmc.Ch2022.rain,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_liq),
+        )
+        @. ᶜwᵢʲs.:($$j) = CMNe.terminal_velocity(
+            cmc.ice,
+            cmc.Ch2022.small_ice,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_ice),
+        )
+#=
+        header = (["CFL", "ρa", "mse", "q_tot", "q_liq", "q_ice", "q_rai", "q_sno", "ᶜwₗʲ", "ᶜwᵢʲ", "ᶜwᵣʲ", "ᶜwₛʲ",])
+        PRT.pretty_table(hcat(
+            parent(Fields.Δz_field(axes(Y.c)))[:] ./ dt,
+            parent(Y.c.sgsʲs.:($j).ρa)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).mse)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).q_tot)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).q_liq)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).q_ice)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).q_rai)[:, :, 1, 1, 1],
+            parent(Y.c.sgsʲs.:($j).q_sno)[:, :, 1, 1, 1],
+            parent(ᶜwₗʲs.:($j))[:, :, 1, 1, 1],
+            parent(ᶜwᵢʲs.:($j))[:, :, 1, 1, 1],
+            parent(ᶜwᵣʲs.:($j))[:, :, 1, 1, 1],
+            parent(ᶜwₛʲs.:($j))[:, :, 1, 1, 1],
+           ) , show_row_number = true, header = header, crop = :none)
+=#
         # Precipitation sources and sinks from the updrafts
         compute_precipitation_sources!(
             ᶜSᵖ,
diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -6,6 +6,8 @@ using LinearAlgebra: ×, dot
 import ClimaCore.Fields as Fields
 import ClimaCore.Geometry as Geometry
 
+import PrettyTables as PRT
+
 NVTX.@annotate function horizontal_advection_tendency!(Yₜ, Y, p, t)
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     (; ᶜΦ) = p.core
@@ -230,58 +232,111 @@ function edmfx_sgs_vertical_advection_tendency!(
     end
 
     for j in 1:n
+        # Flux form vertical advection of area farction with the grid mean velocity
         @. ᶜa_scalar = draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
-        vtt = vertical_transport(
-            ᶜρʲs.:($j),
-            ᶠu³ʲs.:($j),
-            ᶜa_scalar,
-            dt,
-            edmfx_upwinding,
-        )
+        vtt = vertical_transport(ᶜρʲs.:($j), ᶠu³ʲs.:($j), ᶜa_scalar, dt, edmfx_upwinding)
         @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
 
-        va = vertical_advection(
-            ᶠu³ʲs.:($j),
-            Y.c.sgsʲs.:($j).mse,
-            edmfx_upwinding,
-        )
+        # Advective form advection of mse and q_tot with the grid mean velocity
+        va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).mse, edmfx_upwinding)
         @. Yₜ.c.sgsʲs.:($$j).mse += va
-
-        va = vertical_advection(
-            ᶠu³ʲs.:($j),
-            Y.c.sgsʲs.:($j).q_tot,
-            edmfx_upwinding,
-        )
+        va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_tot, edmfx_upwinding)
         @. Yₜ.c.sgsʲs.:($$j).q_tot += va
+
         if p.atmos.moisture_model isa NonEquilMoistModel &&
            p.atmos.precip_model isa Microphysics1Moment
-            # TODO - add precipitation terminal velocity
-            # TODO - add cloud sedimentation velocity
-            # TODO - add their contributions to mean energy and mass
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_liq,
-                edmfx_upwinding,
-            )
+           # TODO - add precipitation terminal velocity in implicit solver/tendency with if/else
+           # TODO - add cloud sedimentation velocity in implicit solver/tendency with if/else
+
+            # Advective form advection of moisture tracers with the grid mean velocity
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_liq, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_liq += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_ice,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_ice, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_ice += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_rai,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_rai, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_rai += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_sno,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_sno, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_sno += va
-        end
+
+            FT = eltype(params)
+            (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜtsʲs) = p.precomputed
+
+            ᶜinv_ρ̂_∂ρ̂∂z = (@. lazy(divide_by_ρa(FT(1), Y.c.sgsʲs.:($$j).ρa, FT(0), Y.c.ρ, turbconv_model) * ᶜgradᵥ(ᶠinterp(Y.c.sgsʲs.:($$j).ρa))))
+            ᶠwₗ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwₗʲs.:($$j))))))
+            ᶠwᵢ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j))))))
+            ᶠwᵣ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j))))))
+            ᶠwₛ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwₛʲs.:($$j))))))
+
+            ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))))
+
+            # Flux form vertical advection of rho * area with sedimantation contributions
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwₗ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_liq)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwᵢ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_ice)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwᵣ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_rai)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwₛ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_sno)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+
+            # q_tot and moisture tracers advective form advection with sedimentation velocity
+            va = vertical_advection(ᶠwₗ³ʲs, Y.c.sgsʲs.:($j).q_liq, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_liq += va
+            va = vertical_advection(ᶠwᵢ³ʲs, Y.c.sgsʲs.:($j).q_ice, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_ice += va
+            va = vertical_advection(ᶠwᵣ³ʲs, Y.c.sgsʲs.:($j).q_rai, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_rai += va
+            va = vertical_advection(ᶠwₛ³ʲs, Y.c.sgsʲs.:($j).q_sno, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_sno += va
+
+            # q_tot and moisture tracers terms proportional to 1/ρ̂ ∂zρ̂
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwₗʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_liq * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_liq -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwₗʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_liq
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_ice * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_ice -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_ice
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_rai * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_rai -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_rai
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwₛʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_sno * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_sno -= dot(ᶜinv_ρ̂_∂ρ̂∂z, CT3(Geometry.WVector(-1 * ᶜwₛʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_sno
+
+            # q_tot and moisture tracer terms proportional to velocity gradients
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwₗ³ʲs) * Y.c.sgsʲs.:($$j).q_liq * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_liq -= ᶜdivᵥ(ᶠwₗ³ʲs) * Y.c.sgsʲs.:($$j).q_liq
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwᵢ³ʲs) * Y.c.sgsʲs.:($$j).q_ice * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_ice -= ᶜdivᵥ(ᶠwᵢ³ʲs) * Y.c.sgsʲs.:($$j).q_ice
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwᵣ³ʲs) * Y.c.sgsʲs.:($$j).q_rai * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_rai -= ᶜdivᵥ(ᶠwᵣ³ʲs) * Y.c.sgsʲs.:($$j).q_rai
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwₛ³ʲs) * Y.c.sgsʲs.:($$j).q_sno * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_sno -= ᶜdivᵥ(ᶠwₛ³ʲs) * Y.c.sgsʲs.:($$j).q_sno
+       end
     end
 end
+
+            #(; ᶜΦ) = p.core
+            #thp = CAP.thermodynamics_params(params)
+
+            #@. ᶜa_scalar = Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ
+            #vtt = vertical_transport(FT(1), (@. lazy( CT3(ᶠinterp(Geometry.WVector(ᶜwᵣʲs.:($$j)))))), ᶜa_scalar, dt, edmfx_upwinding)
+            #@. Yₜ.c.sgsʲs.:($$j).q_rai -= vtt
+            #@. Yₜ.c.sgsʲs.:($$j).q_rai -= ᶜinv_ρ̂_∂ρ̂∂z * ᶜwᵣʲs.:($$j) * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜϕ)
+
+            #(@. lazy(ᶠu³ʲs.:($$j) - CT3(ᶠinterp(Geometry.WVector(ᶜwₗʲs.:($$j)))))),
+
+            #@. Yₜ.c.sgsʲs.:($$j).mse -= ifelse(
+            #    Y.c.sgsʲs.:($$j).ρa <= eps(FT),
+            #    FT(0),
+            #    (
+            #     1 / Y.c.sgsʲs.:($$j).ρa *
+            #        ᶜdivᵥ(ᶠinterp(Y.c.sgsʲs.:($$j).ρa * (
+            #            Geometry.WVector(ᶜwₗʲs.:($$j)) * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ)  +
+            #            Geometry.WVector(ᶜwᵢʲs.:($$j)) * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) +
+            #            Geometry.WVector(ᶜwᵣʲs.:($$j)) * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ)  +
+            #            Geometry.WVector(ᶜwₛʲs.:($$j)) * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ)
+            #        )))
+            #   )
+            #)
+
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -70,32 +70,32 @@ end
 # expressions are less convoluted?
 
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:none})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠu³ * ᶠinterp(ᶜχ))))
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:first_order})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind1(ᶠu³, ᶜχ))))
 end
 @static if pkgversion(ClimaCore) ≥ v"0.14.22"
     function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:vanleer_limiter})
-        ᶜJ = Fields.local_geometry_field(ᶜρ).J
-        ᶠJ = Fields.local_geometry_field(ᶠu³).J
+        ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+        ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
         return @. lazy(
             -(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠlin_vanleer(ᶠu³, ᶜχ, dt))),
         )
     end
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:third_order})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind3(ᶠu³, ᶜχ))))
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:boris_book})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(
         -(ᶜadvdivᵥ(
             ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * (
@@ -109,8 +109,8 @@ function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:boris_book})
     )
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:zalesak})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(
         -(ᶜadvdivᵥ(
             ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * (
@@ -136,8 +136,8 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     (; moisture_model, turbconv_model, rayleigh_sponge, precip_model) = p.atmos
     (; dt) = p
     n = n_mass_flux_subdomains(turbconv_model)
-    ᶜJ = Fields.local_geometry_field(Y.c).J
-    ᶠJ = Fields.local_geometry_field(Y.f).J
+    ᶜJ = Fields.local_geometry_field(axes(Y.c)).J
+    ᶠJ = Fields.local_geometry_field(axes(Y.f)).J
     (; ᶠgradᵥ_ᶜΦ) = p.core
     (; ᶜh_tot, ᶜspecific, ᶠu³, ᶜp) = p.precomputed
 
