Change reference state to MSE.

Project.toml

@@ -12,9 +12,11 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LambertW = "984bce1d-4616-540c-a9ee-88d1112d94c9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OperatorFlux = "47cfafe2-3833-4da3-8183-ce14c2b92cbd"
+RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"

docs/Project.toml

@@ -10,9 +10,11 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"

driver/dycore.jl

@@ -198,19 +198,24 @@ function compute_ref_state!(
     ts_g,
 ) where {PS}
     FT = eltype(grid)
+    kf_surf = TC.kf_surface(grid)
     qtg = TD.total_specific_humidity(param_set, ts_g)
-    θ_liq_ice_g = TD.liquid_ice_pottemp(param_set, ts_g)
     Pg = TD.air_pressure(param_set, ts_g)
+    Φ = TC.geopotential(param_set, grid.zf[kf_surf].z)
+    mse_g = TD.moist_static_energy(param_set, ts_g, Φ)
 
     # We are integrating the log pressure so need to take the log of the
     # surface pressure
     logp = log(Pg)
 
     # Form a right hand side for integrating the hydrostatic equation to
     # determine the reference pressure
-    function rhs(logp, u, z)
+    function minus_inv_scale_height(logp, u, z)
         p_ = exp(logp)
-        ts = TD.PhaseEquil_pθq(param_set, p_, θ_liq_ice_g, qtg)
+        grav = FT(CPP.grav(param_set))
+        Φ = TC.geopotential(param_set, z)
+        h = TC.enthalpy(mse_g, Φ)
+        ts = TC.PhaseEquil_phq(param_set, p_, h, qtg)
         R_m = TD.gas_constant_air(param_set, ts)
         T = TD.air_temperature(param_set, ts)
         return -FT(CPP.grav(param_set)) / (T * R_m)
@@ -219,7 +224,7 @@ function compute_ref_state!(
     # Perform the integration
     z_span = (grid.zmin, grid.zmax)
     @info "z_span = $z_span"
-    prob = ODE.ODEProblem(rhs, logp, z_span)
+    prob = ODE.ODEProblem(minus_inv_scale_height, logp, z_span)
     sol = ODE.solve(prob, ODE.Tsit5(), reltol = 1e-12, abstol = 1e-12)
 
     parent(p_f) .= sol.(vec(grid.zf))
@@ -230,12 +235,16 @@ function compute_ref_state!(
 
     # Compute reference state thermodynamic profiles
     @inbounds for k in TC.real_center_indices(grid)
-        ts = TD.PhaseEquil_pθq(param_set, p_c[k], θ_liq_ice_g, qtg)
+        Φ = TC.geopotential(param_set, grid.zc[k].z)
+        h = TC.enthalpy(mse_g, Φ)
+        ts = TC.PhaseEquil_phq(param_set, p_c[k], h, qtg)
         ρ_c[k] = TD.air_density(param_set, ts)
     end
 
     @inbounds for k in TC.real_face_indices(grid)
-        ts = TD.PhaseEquil_pθq(param_set, p_f[k], θ_liq_ice_g, qtg)
+        Φ = TC.geopotential(param_set, grid.zf[k].z)
+        h = TC.enthalpy(mse_g, Φ)
+        ts = TC.PhaseEquil_phq(param_set, p_f[k], h, qtg)
         ρ_f[k] = TD.air_density(param_set, ts)
     end
     return nothing