Stop hard coding in gravity. (#47)

Stop hard coding in gravity.

Author: ali-ramadhan

src/models.jl

@@ -24,6 +24,7 @@ mutable struct CompressibleModel{A, FT, G, M, D, T, TS, B, C, TC, DF, MP, F, P,
                   forcing :: F
                parameters :: P
        reference_pressure :: FT
+                  gravity :: FT
             slow_forcings :: SF
          right_hand_sides :: RHS
         intermediate_vars :: IV
@@ -34,6 +35,9 @@ end
 ##### Constructor for compressible models
 #####
 
+const pₛ_Earth = 100000
+const  g_Earth = 9.80665
+
 function CompressibleModel(;
                      grid,
              architecture = CPU(),
@@ -50,7 +54,8 @@ function CompressibleModel(;
              microphysics = nothing,
                   forcing = ModelForcing(),
                parameters = nothing,
-       reference_pressure = 100000,
+       reference_pressure = pₛ_Earth,
+                  gravity = g_Earth,
             slow_forcings = ForcingFields(architecture, grid, tracernames(tracers)),
          right_hand_sides = RightHandSideFields(architecture, grid, tracernames(tracers)),
         intermediate_vars = RightHandSideFields(architecture, grid, tracernames(tracers)),
@@ -61,14 +66,15 @@ function CompressibleModel(;
     validate_microphysics(microphysics, tracers)
 
     reference_pressure = float_type(reference_pressure)
-    tracers = TracerFields(architecture, grid, tracers)
+               gravity = float_type(gravity)
 
+    tracers = TracerFields(architecture, grid, tracers)
     forcing = ModelForcing(tracernames(tracers), forcing)
     closure = with_tracers(tracernames(tracers), closure)
 
     return CompressibleModel(architecture, grid, clock, momenta, density, prognostic_temperature,
                              tracers, buoyancy, coriolis, closure, diffusivities, microphysics,
-                             forcing, parameters, reference_pressure, slow_forcings,
+                             forcing, parameters, reference_pressure, gravity, slow_forcings,
                              right_hand_sides, intermediate_vars, acoustic_time_stepper)
 end
 
@@ -93,9 +99,7 @@ function ForcingFields(arch, grid, tracernames)
     ρv = FaceFieldY(arch, grid)
     ρw = FaceFieldZ(arch, grid)
     momenta = (ρu=ρu, ρv=ρv, ρw=ρw)
-
     tracers = TracerFields(arch, grid, tracernames)
-
     return merge(momenta, tracers)
 end
 
@@ -105,8 +109,6 @@ function RightHandSideFields(arch, grid, tracernames)
     ρw = FaceFieldZ(arch, grid)
     ρ =  CellField(arch, grid)
     momenta_and_density = (ρu=ρu, ρv=ρv, ρw=ρw, ρ=ρ)
-
     tracers = TracerFields(arch, grid, tracernames)
-
     return merge(momenta_and_density, tracers)
 end

src/right_hand_sides.jl

@@ -1,9 +1,3 @@
-####
-#### Dirty hacks!
-####
-
-const grav = 9.80665
-
 ####
 #### Element-wise forcing and right-hand-side calculations
 ####
@@ -33,28 +27,28 @@ end
 @inline FC(i, j, k, grid, closure, ρᵈ, C, tracer_index, K̃) =
     @inbounds ρᵈ[i, j, k] * ∇_κ_∇c(i, j, k, grid, closure, ρᵈ, C, Val(tracer_index), K̃)
 
-@inline function RU(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C, FU)
+@inline function RU(i, j, k, grid, pt, b, mp, pₛ, ρᵈ, Ũ, C, FU)
     @inbounds begin
         return (- div_ρuũ(i, j, k, grid, ρᵈ, Ũ)
                 - ρᵈ_over_ρᵐ(i, j, k, grid, mp, ρᵈ, C) * ∂p∂x(i, j, k, grid, pt, b, ρᵈ, C)
                 + FU[i, j, k])
     end
 end
 
-@inline function RV(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C, FV)
+@inline function RV(i, j, k, grid, pt, b, mp, pₛ, ρᵈ, Ũ, C, FV)
     @inbounds begin
         return (- div_ρvũ(i, j, k, grid, ρᵈ, Ũ)
                 - ρᵈ_over_ρᵐ(i, j, k, grid, mp, ρᵈ, C) * ∂p∂y(i, j, k, grid, pt, b, ρᵈ, C)
                 + FV[i, j, k])
     end
 end
 
-@inline function RW(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C, FW)
+@inline function RW(i, j, k, grid, pt, b, mp, pₛ, gravity, ρᵈ, Ũ, C, FW)
     @inbounds begin
         return (- div_ρwũ(i, j, k, grid, ρᵈ, Ũ)
                 - ρᵈ_over_ρᵐ(i, j, k, grid, mp, ρᵈ, C) * (
                       ∂p∂z(i, j, k, grid, pt, b, ρᵈ, C)
-                    + buoyancy_perturbation(i, j, k, grid, grav, mp, ρᵈ, C))
+                    + buoyancy_perturbation(i, j, k, grid, gravity, mp, ρᵈ, C))
                 + FW[i, j, k])
     end
 end

src/time_stepping.jl

@@ -52,6 +52,7 @@ function time_step!(model::CompressibleModel; Δt, Nt=1)
     IV = model.intermediate_vars
 
     pₛ = model.reference_pressure
+    g  = model.gravity
 
     # On third RK3 step, we update Φ⁺ instead of model.intermediate_vars
     Φ⁺ = merge(Ũ, C̃, (ρ=ρᵈ,))
@@ -79,12 +80,12 @@ function time_step!(model::CompressibleModel; Δt, Nt=1)
 
             @debug "  Computing right hand sides..."
             if rk3_iter == 1
-                compute_rhs_args = (R, grid, prog_temp, buoyancy, pₛ, microphysics, ρᵈ, Ũ, C̃, F)
+                compute_rhs_args = (R, grid, prog_temp, buoyancy, microphysics, pₛ, g, ρᵈ, Ũ, C̃, F)
                 fill_halo_regions!(ρᵈ.data, hpbcs, arch, grid)
                 fill_halo_regions!(datatuple(merge(Ũ, C̃)), hpbcs, arch, grid)
                 fill_halo_regions!(Ũ.ρw.data, hpbcs_np, arch, grid)
             else
-                compute_rhs_args = (R, grid, prog_temp, buoyancy, pₛ, microphysics, IV.ρ, IV_Ũ, IV_C̃, F)
+                compute_rhs_args = (R, grid, prog_temp, buoyancy, microphysics, pₛ, g, IV.ρ, IV_Ũ, IV_C̃, F)
                 fill_halo_regions!(IV.ρ.data, hpbcs, arch, grid)
                 fill_halo_regions!(datatuple(merge(IV_Ũ, IV_C̃)), hpbcs, arch, grid)
                 fill_halo_regions!(IV_Ũ.ρw.data, hpbcs_np, arch, grid)

src/time_stepping_kernels.jl

@@ -23,14 +23,13 @@ end
 """
 Fast forcings include advection, pressure gradient, and buoyancy terms.
 """
-function compute_right_hand_sides!(R, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C̃, F)
+function compute_right_hand_sides!(R, grid, pt, b, mp, pₛ, g, ρᵈ, Ũ, C̃, F)
     @inbounds begin
         for k in 1:grid.Nz, j in 1:grid.Ny, i in 1:grid.Nx
-            R.ρu[i, j, k] = RU(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C̃, F.ρu)
-            R.ρv[i, j, k] = RV(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C̃, F.ρv)
-            R.ρw[i, j, k] = RW(i, j, k, grid, pt, b, pₛ, mp, ρᵈ, Ũ, C̃, F.ρw)
-
-            R.ρ[i, j, k] = Rρ(i, j, k, grid, Ũ)
+            R.ρu[i, j, k] = RU(i, j, k, grid, pt, b, mp, pₛ, ρᵈ, Ũ, C̃, F.ρu)
+            R.ρv[i, j, k] = RV(i, j, k, grid, pt, b, mp, pₛ, ρᵈ, Ũ, C̃, F.ρv)
+            R.ρw[i, j, k] = RW(i, j, k, grid, pt, b, mp, pₛ, g, ρᵈ, Ũ, C̃, F.ρw)
+             R.ρ[i, j, k] = Rρ(i, j, k, grid, Ũ)
         end
 
         for C_name in propertynames(C̃)