Skip to content

Inline functions that return broadcasted objects #3798

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

Merged
merged 1 commit into from
May 10, 2025
Merged

Inline functions that return broadcasted objects #3798

Show file tree
Hide file tree
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
3 changes: 1 addition & 2 deletions src/cache/diagnostic_edmf_precomputed_quantities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -979,8 +979,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
ᶠu⁰ = p.scratch.ᶠtemp_C123
@. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_center(ᶠu⁰)
@. ᶜstrain_rate = bc_strain_rate
ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)

ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
Expand Down
6 changes: 2 additions & 4 deletions src/cache/precomputed_quantities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -270,8 +270,7 @@ the `turbconv_model` is EDMFX, the `Y.f.sgsʲs` are also modified so that each
"""
function set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
sfc_u₃ = Fields.level(Y.f.u₃.components.data.:1, half)
bc_sfc_u₃ = surface_velocity(Y.f.u₃, ᶠuₕ³)
@. sfc_u₃ = bc_sfc_u₃
sfc_u₃ .= surface_velocity(Y.f.u₃, ᶠuₕ³)
if turbconv_model isa PrognosticEDMFX
for j in 1:n_mass_flux_subdomains(turbconv_model)
sfc_u₃ʲ = Fields.level(Y.f.sgsʲs.:($j).u₃.components.data.:1, half)
Expand Down Expand Up @@ -316,8 +315,7 @@ end
function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶠuₕ³)
@. ᶜu = C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
@. ᶠu³ = ᶠuₕ³ + CT3(ᶠu₃)
bc_kinetic = compute_kinetic(ᶜuₕ, ᶠu₃)
@. ᶜK = bc_kinetic
ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
return nothing
end

Expand Down
3 changes: 1 addition & 2 deletions src/cache/prognostic_edmf_precomputed_quantities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -497,8 +497,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
ᶠu⁰ = p.scratch.ᶠtemp_C123
@. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰)
ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_center(ᶠu⁰)
@. ᶜstrain_rate = bc_strain_rate
ᶜstrain_rate .= compute_strain_rate_center(ᶠu⁰)
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)

ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar
Expand Down
3 changes: 1 addition & 2 deletions src/initial_conditions/initial_conditions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,8 +407,7 @@ function overwrite_initial_conditions!(
Y.c.uₕ .= C12.(Geometry.UVVector.(vel))
Y.f.u₃ .= ᶠinterp.(C3.(Geometry.WVector.(vel)))
e_kin = similar(ᶜT)
bc_kinetic = compute_kinetic(Y.c.uₕ, Y.f.u₃)
@. e_kin = bc_kinetic
e_kin .= compute_kinetic(Y.c.uₕ, Y.f.u₃)
e_pot = Fields.coordinate_field(Y.c).z .* thermo_params.grav
Y.c.ρe_tot .= TD.total_energy.(thermo_params, ᶜts, e_kin, e_pot) .* Y.c.ρ
if hasproperty(Y.c, :ρq_tot)
Expand Down
6 changes: 2 additions & 4 deletions src/prognostic_equations/edmfx_sgs_flux.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -356,8 +356,7 @@ function edmfx_sgs_diffusive_flux_tendency!(

# momentum
ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_face(ᶜu⁰)
@. ᶠstrain_rate = bc_strain_rate
ᶠstrain_rate .= compute_strain_rate_face(ᶜu⁰)
@. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
end
return nothing
Expand Down Expand Up @@ -419,8 +418,7 @@ function edmfx_sgs_diffusive_flux_tendency!(

# momentum
ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_face(ᶜu)
@. ᶠstrain_rate = bc_strain_rate
ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
@. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ)
end

Expand Down
5 changes: 2 additions & 3 deletions src/prognostic_equations/gm_sgs_closures.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,7 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
ᶠu = p.scratch.ᶠtemp_C123
@. ᶠu = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³)
ᶜstrain_rate = p.scratch.ᶜtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_center(ᶠu)
@. ᶜstrain_rate = bc_strain_rate
ᶜstrain_rate .= compute_strain_rate_center(ᶠu)
@. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate)

ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_2
Expand All @@ -64,6 +63,6 @@ NVTX.@annotate function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
sqrt(max(ᶜlinear_buoygrad, 0)), #N_eff
ᶜdz,
ᶜprandtl_nvec,
norm_sqr(ᶜstrain_rate),
ᶜstrain_rate_norm,
)
end
3 changes: 1 addition & 2 deletions src/prognostic_equations/vertical_diffusion_boundary_layer.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,7 @@ function vertical_diffusion_boundary_layer_tendency!(

if !disable_momentum_vertical_diffusion(p.atmos.vert_diff)
ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
bc_strain_rate = compute_strain_rate_face(ᶜu)
@. ᶠstrain_rate = bc_strain_rate
ᶠstrain_rate .= compute_strain_rate_face(ᶜu)
@. Yₜ.c.uₕ -= C12(
ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_u) * ᶠstrain_rate) / Y.c.ρ,
)
Expand Down
9 changes: 3 additions & 6 deletions src/utils/utilities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,7 @@ sort_files_by_time(files) =
permute!(files, sortperm(time_from_filename.(files)))

"""
bc_kinetic = compute_kinetic(uₕ::Field, uᵥ::Field)
@. κ = bc_kinetic
κ .= compute_kinetic(uₕ::Field, uᵥ::Field)

Compute the specific kinetic energy at cell centers, resulting in `κ` from
individual velocity components:
Expand Down Expand Up @@ -80,8 +79,7 @@ state.
compute_kinetic(Y::Fields.FieldVector) = compute_kinetic(Y.c.uₕ, Y.f.u₃)

"""
bc_ϵ = compute_strain_rate_center(u::Field)
@. ϵ = bc_ϵ
ϵ .= compute_strain_rate_center(u::Field)

Compute the strain_rate at cell centers from velocity at cell faces.
"""
Expand All @@ -97,8 +95,7 @@ function compute_strain_rate_center(u::Fields.Field)
end

"""
bc_ϵ = compute_strain_rate_face(u::Field)
@. ϵ = bc_ϵ
ϵ .= compute_strain_rate_face(u::Field)

Compute the strain_rate at cell faces from velocity at cell centers.
"""
Expand Down
11 changes: 3 additions & 8 deletions test/utilities.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,7 @@ end
CT123 = Geometry.Contravariant123Vector
# Exercise function
κ = zeros(cent_space)
bc_kinetic = CA.compute_kinetic(uₕ, uᵥ)
@. κ = bc_kinetic
κ .= CA.compute_kinetic(uₕ, uᵥ)
ᶜκ_exact = @. 1 // 2 *
cos(z)^2 *
((sin(x)^2) * (cos(y)^2) + (cos(x)^2) * (sin(y)^2))
Expand Down Expand Up @@ -87,12 +86,8 @@ end
ᶠu = @. UVW(Geometry.UVector(ᶠu)) +
UVW(Geometry.VVector(ᶠv)) +
UVW(Geometry.WVector(ᶠw))
bc_strain_rate =
CA.compute_strain_rate_center(Geometry.Covariant123Vector.(ᶠu))
@. ᶜϵ = bc_strain_rate
bc_strain_rate =
CA.compute_strain_rate_face(Geometry.Covariant123Vector.(ᶜu))
@. ᶠϵ = bc_strain_rate
ᶜϵ .= CA.compute_strain_rate_center(Geometry.Covariant123Vector.(ᶠu))
ᶠϵ .= CA.compute_strain_rate_face(Geometry.Covariant123Vector.(ᶜu))

# Center valued strain rate
@test ᶜϵ.components.data.:1 == ᶜϵ.components.data.:1 .* FT(0)
Expand Down
Loading