Skip to content

Commit

Permalink
Merge #2274
Browse files Browse the repository at this point in the history
2274: separate precomputed quantities in 3 funcs r=trontrytel a=trontrytel

Following what we discussed `@szy21` , `@Sbozzolo`  this PR splits prognostic edmfx precomputed quantities into 3 steps
1) updating environment
2) updating updraft thermo state and BC
3) updating the closures.

I'm not sure if it actually makes anything better though. Or if you want to split the individual closures. 

Co-authored-by: Anna Jaruga <ajaruga@caltech.edu>
  • Loading branch information
bors[bot] and trontrytel authored Oct 25, 2023
2 parents 9fcf2d9 + 594c8fd commit 768aeb0
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 20 deletions.
6 changes: 4 additions & 2 deletions src/cache/precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ end
"""
set_velocity_at_top!(Y, turbconv_model)
Modifies `Y.f.u₃` so that `u₃` is 0 at the model top.
Modifies `Y.f.u₃` so that `u₃` is 0 at the model top.
"""
function set_velocity_at_top!(Y, turbconv_model)
top_u₃ = Fields.level(
Expand Down Expand Up @@ -340,7 +340,9 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
SurfaceConditions.update_surface_conditions!(Y, p, t)

if turbconv_model isa PrognosticEDMFX
set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
end

if turbconv_model isa DiagnosticEDMFX
Expand Down
73 changes: 55 additions & 18 deletions src/cache/prognostic_edmf_precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,28 +5,17 @@ import Thermodynamics as TD
import ClimaCore: Spaces, Fields

"""
set_prognostic_edmf_precomputed_quantities!(Y, p, t)
set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
Updates the precomputed quantities stored in `p` for edmfx.
Updates the edmf environment precomputed quantities stored in `p` for edmfx.
"""
function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
(; energy_form, moisture_model, turbconv_model) = p.atmos
#EDMFX BCs only support total energy as state variable
@assert energy_form isa TotalEnergy
@assert !(moisture_model isa DryModel)

FT = Spaces.undertype(axes(Y.c))
(; params) = p
(; dt) = p.simulation
thermo_params = CAP.thermodynamics_params(params)
n = n_mass_flux_subdomains(turbconv_model)
thermo_args = (thermo_params, energy_form, moisture_model)
function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t)
@assert !(p.atmos.moisture_model isa DryModel)

(; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot, ᶜρ_ref) = p
thermo_params = CAP.thermodynamics_params(p.params)
(; turbconv_model) = p.atmos
(; ᶜp, ᶜΦ, ᶜh_tot) = p
(; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p
(; ᶜmixing_length, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜK_u, ᶜK_h) = p
(; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p
(; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions

@. ᶜρa⁰ = ρa⁰(Y.c)
@. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
Expand All @@ -49,6 +38,30 @@ function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
@. ᶜK⁰ += ᶜtke⁰
@. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜh_tot⁰ - ᶜK⁰ - ᶜΦ, ᶜq_tot⁰)
@. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
return nothing
end

"""
set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
Updates the draft thermo state and boundary conditions
precomputed quantities stored in `p` for edmfx.
"""
function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)
(; energy_form, moisture_model, turbconv_model) = p.atmos
#EDMFX BCs only support total energy as state variable
@assert energy_form isa TotalEnergy
@assert !(moisture_model isa DryModel)

FT = Spaces.undertype(axes(Y.c))
n = n_mass_flux_subdomains(turbconv_model)

(; params) = p
thermo_params = CAP.thermodynamics_params(params)

(; ᶜspecific, ᶜp, ᶜΦ, ᶜh_tot) = p
(; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p
(; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions

for j in 1:n
ᶜuʲ = ᶜuʲs.:($j)
Expand Down Expand Up @@ -135,6 +148,30 @@ function set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
$(FT(turbconv_params.surface_area)) *
TD.air_density(thermo_params, ᶜtsʲ_int_val)
end
return nothing
end

"""
set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
Updates the precomputed quantities stored in `p` for edmfx closures.
"""
function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)

(; moisture_model, turbconv_model) = p.atmos
@assert !(moisture_model isa DryModel)

(; params) = p
(; dt) = p.simulation
thermo_params = CAP.thermodynamics_params(params)

n = n_mass_flux_subdomains(turbconv_model)

(; ᶜspecific, ᶜp, ᶜρ_ref) = p
(; ᶜtke⁰, ᶜρa⁰, ᶜu⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) = p
(; ᶜmixing_length, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜK_u, ᶜK_h) = p
(; ᶜuʲs, ᶜtsʲs, ᶜρʲs, ᶜentrʲs, ᶜdetrʲs) = p
(; ustar, obukhov_length, buoyancy_flux) = p.sfc_conditions

ᶜz = Fields.coordinate_field(Y.c).z
z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
Expand Down

0 comments on commit 768aeb0

Please sign in to comment.