Skip to content

Use LazyBroadcast in set_velocity_quantities #3608

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Use LazyBroadcast in set_velocity_quantities #3608

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions src/cache/precomputed_quantities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,11 +310,27 @@ function set_velocity_at_top!(Y, turbconv_model)
return nothing
end

"""
compute_ᶜu(ᶜuₕ, ᶠu₃)
Computes the cell-centered velocity, given
- `ᶜuₕ` the cell-centered horizontal velocity
- `ᶠu₃` the cell-face `Covariant3` velocity
"""
compute_ᶜu(ᶜuₕ, ᶠu₃) = @. lazy(C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃)))

"""
compute_ᶠu³(ᶠuₕ³, ᶠu₃)
Computes the cell-face contravariant velocity, given
- `ᶠuₕ³` the cell-face horizontal `Contravariant3` velocity
- `ᶠu₃` the cell-face `Covariant3` velocity
"""
compute_ᶠu³(ᶠuₕ³, ᶠu₃) = ᶠuₕ³ + CT3(ᶠu₃)

# This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on
# ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ.
function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³)
@. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
@. ᶠu³ = ᶠuₕ³ + CT3(ᶠu₃)
ᶜu .= compute_ᶜu(ᶜuₕ, ᶠu₃)
ᶠu³ .= compute_ᶠu³.(ᶠuₕ³, ᶠu₃)
ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
return nothing
end
Expand Down Expand Up @@ -451,7 +467,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t)
thermo_args = (thermo_params, moisture_model, precip_model)

@. ᶜspecific = specific_gs(Y.c)
@. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ)
ᶠuₕ³ .= compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ)

# TODO: We might want to move this to dss! (and rename dss! to something
# like enforce_constraints!).
Expand Down
Loading