Remove coriolis forces from the cache

CliMA · Dec 30, 2024 · 2e8e66c · 2e8e66c
1 parent 8cdd3f8
commit 2e8e66c
Show file tree

Hide file tree

Showing 13 changed files with 177 additions and 116 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
@@ -33,7 +33,7 @@ struct AtmosCache{
     """ClimaAtmosParameters that have to be used"""
     params::CAP
 
-    """Variables that are used generally, such as ᶜΦ"""
+    """Variables that are used generally"""
     core::COR
 
     """Used by update_surface_conditions! in set_precomputed_quantities! and coupler"""
@@ -89,10 +89,8 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
     ᶜcoord = Fields.local_geometry_field(Y.c).coordinates
     ᶠcoord = Fields.local_geometry_field(Y.f).coordinates
     grav = FT(CAP.grav(params))
-    ᶜΦ = grav .* ᶜcoord.z
-    ᶠΦ = grav .* ᶠcoord.z
-
-    (; ᶜf³, ᶠf¹²) = compute_coriolis(ᶜcoord, ᶠcoord, params)
+    ᶜz = ᶜcoord.z
+    ᶠz = ᶠcoord.z
 
     ghost_buffer =
         !do_dss(axes(Y.c)) ? (;) :
@@ -121,11 +119,8 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
         Fields.level(Fields.local_geometry_field(Y.f), Fields.half)
 
     core = (
-        ᶜΦ,
-        ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(ᶜΦ),
-        ᶜgradᵥ_ᶠΦ = ᶜgradᵥ.(ᶠΦ),
-        ᶜf³,
-        ᶠf¹²,
+        ᶠgradᵥ_ᶜΦ = ᶠgradᵥ.(Φ.(grav, ᶜz)),
+        ᶜgradᵥ_ᶠΦ = ᶜgradᵥ.(Φ.(grav, ᶠz)),
         # Used by diagnostics such as hfres, evspblw
         surface_ct3_unit = CT3.(
             unit_basis_vector_data.(CT3, sfc_local_geometry)
@@ -185,32 +180,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
@@ -34,13 +34,14 @@ NVTX.@annotate function set_diagnostic_edmfx_draft_quantities_level!(
     mse_level,
     q_tot_level,
     p_level,
-    Φ_level,
+    z_level,
 )
     FT = eltype(thermo_params)
+    grav = TDP.grav(thermo_params)
     @. ts_level = TD.PhaseEquil_phq(
         thermo_params,
         p_level,
-        mse_level - Φ_level,
+        mse_level - Φ(grav, z_level),
         q_tot_level,
         8,
         FT(0.0003),
@@ -92,7 +93,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     (; turbconv_model) = p.atmos
     FT = eltype(Y)
     n = n_mass_flux_subdomains(turbconv_model)
-    (; ᶜΦ) = p.core
     (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) =
@@ -112,7 +112,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
     q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1))
 
     p_int_level = Fields.field_values(Fields.level(ᶜp, 1))
-    Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1))
 
     local_geometry_int_level =
         Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1))
@@ -187,7 +186,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!(
             mseʲ_int_level,
             q_totʲ_int_level,
             p_int_level,
-            Φ_int_level,
+            z_int_level,
         )
         @. ρaʲ_int_level = ρʲ_int_level * FT(turbconv_params.surface_area)
     end
@@ -296,7 +295,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     ᶜdz = Fields.Δz_field(axes(Y.c))
     (; params) = p
     (; dt) = p
-    (; ᶜΦ) = p.core
     (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed
     (; q_tot) = p.precomputed.ᶜspecific
     (;
@@ -324,13 +322,17 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
     microphys_1m_params = CAP.microphysics_1m_params(params)
     turbconv_params = CAP.turbconv_params(params)
 
-    ᶠΦ = p.scratch.ᶠtemp_scalar
-    @. ᶠΦ = CAP.grav(params) * ᶠz
+    g = CAP.grav(params)
+    ᶜcoords = Fields.coordinate_field(Y.c)
+    ᶜz = Fields.coordinate_field(Y.c).z
+    ᶠz = Fields.coordinate_field(Y.f).z
+    global_geom = Spaces.global_geometry(axes(ᶜcoord))
+
     ᶜ∇Φ³ = p.scratch.ᶜtemp_CT3
-    @. ᶜ∇Φ³ = CT3(ᶜgradᵥ(ᶠΦ))
-    @. ᶜ∇Φ³ += CT3(gradₕ(ᶜΦ))
+    @. ᶜ∇Φ³ = CT3(ᶜgradᵥ(Φ(g, ᶠz)))
+    @. ᶜ∇Φ³ += CT3(gradₕ(Φ(g, ᶜz)))
     ᶜ∇Φ₃ = p.scratch.ᶜtemp_C3
-    @. ᶜ∇Φ₃ = ᶜgradᵥ(ᶠΦ)
+    @. ᶜ∇Φ₃ = ᶜgradᵥ(Φ(g, ᶠz))
 
     z_sfc_halflevel =
         Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half))
@@ -344,7 +346,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i))
         q_tot_level = Fields.field_values(Fields.level(q_tot, i))
         p_level = Fields.field_values(Fields.level(ᶜp, i))
-        Φ_level = Fields.field_values(Fields.level(ᶜΦ, i))
+        z_level = Fields.field_values(Fields.level(ᶜz, i))
         local_geometry_level = Fields.field_values(
             Fields.level(Fields.local_geometry_field(Y.c), i),
         )
@@ -355,7 +357,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
         end_index = fieldcount(eltype(∂x∂ξ_level)) # This will be 4 in 2D and 9 in 3D.
         ∂x³∂ξ³_level = ∂x∂ξ_level.:($end_index)
 
-        Φ_prev_level = Fields.field_values(Fields.level(ᶜΦ, i - 1))
         ∇Φ³_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ³, i - 1))
         ∇Φ³_data_prev_level = ∇Φ³_prev_level.components.data.:1
         ∇Φ₃_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ₃, i - 1))
@@ -555,7 +556,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                     q_rai_prev_level,
                     q_sno_prev_level,
                     tsʲ_prev_level,
-                    Φ_prev_level,
                     dt,
                     microphys_1m_params,
                     thermo_params,
@@ -702,7 +702,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                             e_tot_0M_precipitation_sources_helper(
                                 thermo_params,
                                 tsʲ_prev_level,
-                                Φ_prev_level,
+                                Φ(g, z_prev_level),
                             )
                         )
                     )
@@ -807,7 +807,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
                 mseʲ_level,
                 q_totʲ_level,
                 p_level,
-                Φ_level,
+                Φ(g, z_level),
             )
         end
         ρaʲs_level = Fields.field_values(Fields.level(ᶜρaʲs, i))
@@ -1072,7 +1072,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita
         ᶜqᵣ,
         ᶜqₛ,
         ᶜts,
-        p.core.ᶜΦ,
         p.dt,
         microphys_1m_params,
         thermo_params,

diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -465,9 +465,10 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     thermo_params = CAP.thermodynamics_params(p.params)
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_args = (thermo_params, moisture_model)
-    (; ᶜΦ) = p.core
     (; ᶜspecific, ᶜu, ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed
     ᶠuₕ³ = p.scratch.ᶠtemp_CT3
+    ᶜz = Fields.coordinate_field(axes(Y.c)).z
+    grav = TDP.grav(thermo_params)
 
     @. ᶜspecific = specific_gs(Y.c)
     set_ᶠuₕ³!(ᶠuₕ³, Y)
@@ -493,7 +494,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
         # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ
         # TODO: We should think more about these increments before we use them.
     end
-    @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ)
+    @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, Φ(grav, ᶜz), Y.c.ρ)
     @. ᶜp = TD.air_pressure(thermo_params, ᶜts)
 
     if turbconv_model isa AbstractEDMF

diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -20,7 +20,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
 
     thermo_params = CAP.thermodynamics_params(p.params)
     (; turbconv_model) = p.atmos
-    (; ᶜΦ,) = p.core
+    ᶜz = Fields.coordinate_field(axes(Y.c)).z
+    grav = TDP.grav(thermo_params)
     (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
     (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) =
         p.precomputed
@@ -44,7 +45,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
     set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
     # @. ᶜK⁰ += ᶜtke⁰
-    @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
+    @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - Φ(grav, ᶜz), ᶜq_tot⁰)
     @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
     return nothing
 end
@@ -72,7 +73,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
     thermo_params = CAP.thermodynamics_params(params)
     turbconv_params = CAP.turbconv_params(params)
 
-    (; ᶜΦ,) = p.core
+    grav = TDP.grav(thermo_params)
+    ᶜz = Fields.coordinate_field(Y.c).z
     (; ᶜspecific, ᶜp, ᶜh_tot, ᶜK) = p.precomputed
     (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶠKᵥʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
     (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
@@ -90,7 +92,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
 
         set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
-        @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
+        @. ᶜtsʲ =
+            TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - Φ(grav, ᶜz), ᶜq_totʲ)
         @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
 
         # EDMFX boundary condition:
@@ -153,12 +156,11 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
         )
 
         # Then overwrite the prognostic variables at first inetrior point.
-        ᶜΦ_int_val = Fields.field_values(Fields.level(ᶜΦ, 1))
         ᶜtsʲ_int_val = Fields.field_values(Fields.level(ᶜtsʲ, 1))
         @. ᶜtsʲ_int_val = TD.PhaseEquil_phq(
             thermo_params,
             ᶜp_int_val,
-            ᶜmseʲ_int_val - ᶜΦ_int_val,
+            ᶜmseʲ_int_val - Φ(grav, ᶜz_int_val),
             ᶜq_totʲ_int_val,
         )
         sgsʲs_ρ_int_val = Fields.field_values(Fields.level(ᶜρʲs.:($j), 1))
@@ -424,7 +426,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
     @assert !(p.atmos.moisture_model isa DryModel)
 
     (; params, dt) = p
-    (; ᶜΦ,) = p.core
     thp = CAP.thermodynamics_params(params)
     cmp = CAP.microphysics_1m_params(params)
 
@@ -451,7 +452,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
             ᶜqᵣ,
             ᶜqₛ,
             ᶜtsʲs.:($j),
-            ᶜΦ,
             dt,
             cmp,
             thp,
@@ -470,7 +470,6 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
         ᶜqᵣ,
         ᶜqₛ,
         ᶜts⁰,
-        ᶜΦ,
         dt,
         cmp,
         thp,

diff --git a/src/core/core_quantities.jl b/src/core/core_quantities.jl
@@ -0,0 +1,65 @@
+import ClimaCore.Geometry
+
+f³(lg, global_geom, params) = f³(lg, lg.coordinates, global_geom, params)
+
+f¹²(lg, coord, global_geom, params) =
+    f¹²(lg, lg.coordinates, global_geom, params)
+
+function f³(
+    lg,
+    coord::Geometry.LatLongZPoint,
+    global_geom::Geometry.DeepSphericalGlobalGeometry,
+    params,
+)
+    Ω = CAP.Omega(params)
+    CT3(
+        CT123(
+            Geometry.LocalVector(
+                Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
+                global_geom,
+                coord,
+            ),
+            lg,
+        ),
+        lg,
+    )
+end
+
+function f¹²(
+    lg,
+    coord::Geometry.LatLongZPoint,
+    global_geom::Geometry.DeepSphericalGlobalGeometry,
+    params,
+)
+    Ω = CAP.Omega(params)
+    CT12(
+        CT123(
+            Geometry.LocalVector(
+                Geometry.Cartesian123Vector(zero(Ω), zero(Ω), 2 * Ω),
+                global_geom,
+                coord,
+            ),
+            lg,
+        ),
+        lg,
+    )
+end
+
+# Shallow sphere
+function f³(lg, coord::Geometry.LatLongZPoint, global_geom, params)
+    Ω = CAP.Omega(params)
+    CT3(Geometry.WVector(2 * Ω * sind(coord.lat), lg), lg)
+end
+
+# Shallow cartesian
+function f³(lg, coord, global_geom, params)
+    f = CAP.f_plane_coriolis_frequency(params)
+    CT3(Geometry.WVector(f, lg), lg)
+end
+
+f¹²(lg, coord::Geometry.LatLongZPoint, global_geom, params) =
+    error("Not supported for $coord, $global_geom.")
+f¹²(lg, coord, global_geom, p) =
+    error("Not supported for $coord, $global_geom.")
+
+Φ(g, z) = g * z
diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl
@@ -12,6 +12,7 @@ import ClimaCore.Utilities: half
 import Thermodynamics as TD
 
 import ..AtmosModel
+import ..Φ
 import ..AtmosCallback
 import ..EveryNSteps