Remove coriolis forces from the cache

CliMA · Nov 25, 2024 · 39ab56b · 39ab56b
1 parent 002b4f3
commit 39ab56b
Show file tree

Hide file tree

Showing 5 changed files with 77 additions and 49 deletions.
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
@@ -21,6 +21,7 @@ include(joinpath("utils", "read_gcm_driven_scm_data.jl"))
 include(joinpath("utils", "AtmosArtifacts.jl"))
 import .AtmosArtifacts as AA
 
+include(joinpath("core", "core_quantities.jl"))
 include(
     joinpath("parameterized_tendencies", "radiation", "radiation_utilities.jl"),
 )

diff --git a/src/cache/cache.jl b/src/cache/cache.jl
@@ -115,8 +115,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
     ᶜΦ = grav .* ᶜcoord.z
     ᶠΦ = grav .* ᶠcoord.z
 
-    (; ᶜf³, ᶠf¹²) = compute_coriolis(ᶜcoord, ᶠcoord, params)
-
     quadrature_style =
         Spaces.quadrature_style(Spaces.horizontal_space(axes(Y.c)))
     do_dss = quadrature_style isa Quadratures.GLL
@@ -150,8 +148,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         ᶜΦ,
         ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
         ᶜgradᵥ_ᶠΦ = ᶜgradᵥ.(ᶠΦ),
-        ᶜf³,
-        ᶠf¹²,
         # Used by diagnostics such as hfres, evspblw
         surface_ct3_unit = CT3.(
             unit_basis_vector_data.(CT3, sfc_local_geometry)
@@ -218,32 +214,3 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
 
     return AtmosCache{map(typeof, args)...}(args...)
 end
-
-
-function compute_coriolis(ᶜcoord, ᶠcoord, params)
-    if eltype(ᶜcoord) <: Geometry.LatLongZPoint
-        Ω = CAP.Omega(params)
-        global_geom = Spaces.global_geometry(axes(ᶜcoord))
-        if global_geom isa Geometry.DeepSphericalGlobalGeometry
-            @info "using deep atmosphere"
-            coriolis_deep(coord::Geometry.LatLongZPoint) = Geometry.LocalVector(
-                Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
-                global_geom,
-                coord,
-            )
-            ᶜf³ = @. CT3(CT123(coriolis_deep(ᶜcoord)))
-            ᶠf¹² = @. CT12(CT123(coriolis_deep(ᶠcoord)))
-        else
-            coriolis_shallow(coord::Geometry.LatLongZPoint) =
-                Geometry.WVector(2 * Ω * sind(coord.lat))
-            ᶜf³ = @. CT3(coriolis_shallow(ᶜcoord))
-            ᶠf¹² = nothing
-        end
-    else
-        f = CAP.f_plane_coriolis_frequency(params)
-        coriolis_f_plane(coord) = Geometry.WVector(f)
-        ᶜf³ = @. CT3(coriolis_f_plane(ᶜcoord))
-        ᶠf¹² = nothing
-    end
-    return (; ᶜf³, ᶠf¹²)
-end
diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -327,13 +327,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     microphys_params = CAP.microphysics_precipitation_params(params)
     turbconv_params = CAP.turbconv_params(params)
 
-    ᶠΦ = p.scratch.ᶠtemp_scalar
-    @. ᶠΦ = CAP.grav(params) * ᶠz
+    g = CAP.grav(params)
     ᶜ∇Φ³ = p.scratch.ᶜtemp_CT3
-    @. ᶜ∇Φ³ = CT3(ᶜgradᵥ(ᶠΦ))
+    @. ᶜ∇Φ₃ = ᶜgradᵥ(Φ(g, ᶠz))
+    @. ᶜ∇Φ³ = CT3(ᶜ∇Φ₃)
     @. ᶜ∇Φ³ += CT3(gradₕ(ᶜΦ))
     ᶜ∇Φ₃ = p.scratch.ᶜtemp_C3
-    @. ᶜ∇Φ₃ = ᶜgradᵥ(ᶠΦ)
 
     z_sfc_halflevel =
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))

diff --git a/src/core/core_quantities.jl b/src/core/core_quantities.jl
@@ -0,0 +1,54 @@
+import ClimaCore.Geometry
+
+function f³(
+    coord::Geometry.LatLongZPoint,
+    global_geom::Geometry.DeepSphericalGlobalGeometry,
+    params,
+)
+    Ω = CAP.Omega(params)
+    CT3(
+        CT123(
+            Geometry.LocalVector(
+                Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
+                global_geom,
+                coord,
+            ),
+        ),
+    )
+end
+
+function f¹²(
+    coord::Geometry.LatLongZPoint,
+    global_geom::Geometry.DeepSphericalGlobalGeometry,
+    params,
+)
+    Ω = CAP.Omega(params)
+    CT12(
+        CT123(
+            Geometry.LocalVector(
+                Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
+                global_geom,
+                coord,
+            ),
+        ),
+    )
+end
+
+# Shallow sphere
+function f³(coord::Geometry.LatLongZPoint, global_geom, params)
+    Ω = CAP.Omega(params)
+    CT3(Geometry.WVector(2 * Ω * sind(coord.lat)))
+end
+
+# Shallow cartesian
+function f³(coord, global_geom, params)
+    f = CAP.f_plane_coriolis_frequency(params)
+    CT3(Geometry.WVector(f))
+end
+
+f¹²(coord::Geometry.LatLongZPoint, global_geom, params) =
+    error("Not supported for $coord, $global_geom.")
+f¹²(coord, global_geom, p::CartesianCoriolisParams) =
+    error("Not supported for $coord, $global_geom.")
+
+Φ(g, ᶠz) = g * z
diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -73,13 +73,14 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     point_type = eltype(Fields.coordinate_field(Y.c))
     (; dt) = p
     ᶜJ = Fields.local_geometry_field(Y.c).J
-    (; ᶜf³, ᶠf¹², ᶜΦ) = p.core
     (; ᶜu, ᶠu³, ᶜK) = p.precomputed
     (; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing
     (; ᶜuʲs, ᶜKʲs, ᶠKᵥʲs) = n > 0 ? p.precomputed : all_nothing
     (; ᶠu³⁰) = advect_tke ? p.precomputed : all_nothing
     (; energy_upwinding, tracer_upwinding) = p.atmos.numerics
     (; ᶜspecific) = p.precomputed
+    ᶜcoords = Fields.coordinate_field(Y.c)
+    global_geom = Spaces.global_geometry(axes(ᶜcoord))
 
     ᶜρa⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρa⁰ : Y.c.ρ) : nothing
     ᶜρ⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρ⁰ : Y.c.ρ) : nothing
@@ -129,26 +130,32 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         end
     end
 
-    if isnothing(ᶠf¹²)
-        # shallow atmosphere
+    if global_geom isa Geometry.DeepSphericalGlobalGeometry
+        # deep atmosphere
         @. Yₜ.c.uₕ -=
-            ᶜinterp(ᶠω¹² × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) / (Y.c.ρ * ᶜJ) +
-            (ᶜf³ + ᶜω³) × CT12(ᶜu)
-        @. Yₜ.f.u₃ -= ᶠω¹² × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
+            ᶜinterp(
+                (f¹²(ᶠcoords, global_geom, params) + ᶠω¹²) ×
+                (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³),
+            ) / (Y.c.ρ * ᶜJ) +
+            (f³(ᶜcoords, global_geom, params) + ᶜω³) × CT12(ᶜu)
+        @. Yₜ.f.u₃ -=
+            (f¹²(ᶠcoords, global_geom, params) + ᶠω¹²) × ᶠinterp(CT12(ᶜu)) +
+            ᶠgradᵥ(ᶜK)
         for j in 1:n
             @. Yₜ.f.sgsʲs.:($$j).u₃ -=
-                ᶠω¹²ʲs.:($$j) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
+                (f¹²(ᶠcoords, global_geom, params) + ᶠω¹²ʲs.:($$j)) ×
+                ᶠinterp(CT12(ᶜuʲs.:($$j))) +
                 ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
         end
     else
-        # deep atmosphere
+        # shallow atmosphere
         @. Yₜ.c.uₕ -=
-            ᶜinterp((ᶠf¹² + ᶠω¹²) × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) /
-            (Y.c.ρ * ᶜJ) + (ᶜf³ + ᶜω³) × CT12(ᶜu)
-        @. Yₜ.f.u₃ -= (ᶠf¹² + ᶠω¹²) × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
+            ᶜinterp(ᶠω¹² × (ᶠinterp(Y.c.ρ * ᶜJ) * ᶠu³)) / (Y.c.ρ * ᶜJ) +
+            (f³(ᶜcoords, global_geom, params) + ᶜω³) × CT12(ᶜu)
+        @. Yₜ.f.u₃ -= ᶠω¹² × ᶠinterp(CT12(ᶜu)) + ᶠgradᵥ(ᶜK)
         for j in 1:n
             @. Yₜ.f.sgsʲs.:($$j).u₃ -=
-                (ᶠf¹² + ᶠω¹²ʲs.:($$j)) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
+                ᶠω¹²ʲs.:($$j) × ᶠinterp(CT12(ᶜuʲs.:($$j))) +
                 ᶠgradᵥ(ᶜKʲs.:($$j) - ᶜinterp(ᶠKᵥʲs.:($$j)))
         end
     end