Skip to content

Commit 73e33c1

Browse files
imreddyTejaimreddyTeja
committed
lazify bcs
1 parent 19ac42c commit 73e33c1

5 files changed

Lines changed: 55 additions & 36 deletions

File tree

‎examples/column/advect.jl‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import ClimaCore:
1111
Spaces,
1212
Utilities
1313

14+
import LazyBroadcast: lazy
1415
using OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33
1516
using ClimaCorePlots
1617
using Plots
@@ -121,15 +122,14 @@ for (fn, mesh) in zip(("sin", "step"), (mesh_sin, mesh_step))
121122
UB = Operators.UpwindBiasedProductC2F()
122123
∂ = Operators.DivergenceF2C()
123124
left_center = Fields.level(θ, 1)
125+
θ_top = Fields.level(θ, Fields.nlevels(θ))
126+
θ_top_m1 = Fields.level(θ, Fields.nlevels(θ) - 1)
124127
right_center_left_biased_grad =
125-
Geometry.Covariant3Vector.(
126-
Fields.level(θ, Fields.nlevels(θ)) .-
127-
Fields.level(θ, Fields.nlevels(θ) - 1)
128-
)
128+
@. lazy(Geometry.Covariant3Vector(θ_top - θ_top_m1))
129129
right_gradient_extrapolate = Operators.SetGradient(right_center_left_biased_grad)
130-
left_gradient_extrapolate = Operators.SetGradient(
131-
Geometry.Covariant3Vector.(Fields.level(θ, 2) .- left_center),
132-
)
130+
θ_2 = Fields.level(θ, 2)
131+
_left_lazy_grad2 = @. lazy(Geometry.Covariant3Vector(θ_2 - left_center))
132+
left_gradient_extrapolate = Operators.SetGradient(_left_lazy_grad2)
133133
gradc2f_fcc = Operators.GradientC2F(
134134
left = left_gradient_extrapolate,
135135
right = right_gradient_extrapolate,
@@ -144,13 +144,14 @@ for (fn, mesh) in zip(("sin", "step"), (mesh_sin, mesh_step))
144144
# use the advection operator
145145
function tendency3!(dθ, θ, _, t)
146146
left_center = Fields.level(θ, 1)
147+
θ_top = Fields.level(θ, Fields.nlevels(θ))
148+
θ_top_m1 = Fields.level(θ, Fields.nlevels(θ) - 1)
147149
right_center_left_biased_grad =
148-
Geometry.Covariant3Vector.(
149-
Fields.level(θ, Fields.nlevels(θ)) .-
150-
Fields.level(θ, Fields.nlevels(θ) - 1)
151-
)
150+
@. lazy(Geometry.Covariant3Vector(θ_top - θ_top_m1))
152151
left_gradient =
153-
Operators.SetGradient(@. Geometry.Covariant3Vector(2 * (left_center - sin(-t))))
152+
Operators.SetGradient(
153+
@. lazy(Geometry.Covariant3Vector(2 * (left_center - sin(-t))))
154+
)
154155
right_gradient = Operators.SetGradient(right_center_left_biased_grad)
155156
gradc2f = Operators.GradientC2F(left = left_gradient, right = right_gradient)
156157
interpf2c = Operators.InterpolateF2C()
@@ -160,20 +161,21 @@ for (fn, mesh) in zip(("sin", "step"), (mesh_sin, mesh_step))
160161
# use the advection operator
161162
function tendency4!(dθ, θ, _, t)
162163
left_center = Fields.level(θ, 1)
164+
θ_top = Fields.level(θ, Fields.nlevels(θ))
165+
θ_top_m1 = Fields.level(θ, Fields.nlevels(θ) - 1)
163166
right_center_left_biased_grad =
164-
Geometry.Covariant3Vector.(
165-
Fields.level(θ, Fields.nlevels(θ)) .-
166-
Fields.level(θ, Fields.nlevels(θ) - 1)
167-
)
167+
@. lazy(Geometry.Covariant3Vector(θ_top - θ_top_m1))
168168
left_gradient =
169-
Operators.SetGradient(@. Geometry.Covariant3Vector(2 * (left_center - sin(-t))))
169+
Operators.SetGradient(
170+
@. lazy(Geometry.Covariant3Vector(2 * (left_center - sin(-t))))
171+
)
170172
right_gradient_extrapolate = Operators.SetGradient(right_center_left_biased_grad)
171173
gradc2f =
172174
Operators.GradientC2F(left = left_gradient, right = right_gradient_extrapolate)
173175
interpf2c = Operators.InterpolateF2C()
174-
left_gradient_extrapolate = Operators.SetGradient(
175-
Geometry.Covariant3Vector.(Fields.level(θ, 2) .- left_center),
176-
)
176+
θ_2 = Fields.level(θ, 2)
177+
_left_lazy_grad4 = @. lazy(Geometry.Covariant3Vector(θ_2 - left_center))
178+
left_gradient_extrapolate = Operators.SetGradient(_left_lazy_grad4)
177179
gradc2f_fcc = Operators.GradientC2F(
178180
left = left_gradient_extrapolate,
179181
right = right_gradient_extrapolate,

‎examples/column/advect_diffusion.jl‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import ClimaCore:
1010
Geometry,
1111
Spaces
1212

13+
import LazyBroadcast: lazy
1314
using OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33
1415

1516
import Logging
@@ -63,10 +64,9 @@ function ∑tendencies!(dT, T, z, t)
6364
bc_gt = Operators.SetGradient(
6465
Geometry.WVector(FT(∇gaussian(z₁, t; ν = ν, δ = δ, 𝓌 = 𝓌, μ = μ))),
6566
)
66-
top_center_left_biased_grad =
67-
Geometry.Covariant3Vector.(
68-
Fields.level(T, Fields.nlevels(T)) .- Fields.level(T, Fields.nlevels(T) - 1)
69-
)
67+
T_top = Fields.level(T, Fields.nlevels(T))
68+
T_top_m1 = Fields.level(T, Fields.nlevels(T) - 1)
69+
top_center_left_biased_grad = @. lazy(Geometry.Covariant3Vector(T_top - T_top_m1))
7070

7171
bc_gt_lb = Operators.SetGradient(top_center_left_biased_grad)
7272
gradc2f = Operators.GradientC2F(bottom = bc_gb, top = bc_gt)

‎examples/column/ekman.jl‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import ClimaCore:
1313

1414
using OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33
1515

16+
import LazyBroadcast: lazy
1617
import Logging
1718
import TerminalLoggers
1819
Logging.global_logger(TerminalLoggers.TerminalLogger())
@@ -89,11 +90,14 @@ function tendency!(dY, Y, _, t)
8990
bcs_bottom = Operators.SetValue(Geometry.WVector(Cd * u_wind * u_1)) # Eq. 4.16
9091
u_top = Fields.level(u, Fields.nlevels(u))
9192
u_bottom = Fields.level(u, 1)
92-
bcs_top = Operators.SetGradient(@. Geometry.Covariant3Vector(2 * (FT(ug) - u_top))) # Eq. 4.18
93+
bcs_top =
94+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(2 * (FT(ug) - u_top)))) # Eq. 4.18
9395
gradc2f = Operators.GradientC2F(top = bcs_top)
9496
divf2c = Operators.DivergenceF2C(bottom = bcs_bottom)
95-
bcs_bottom_advection = Operators.SetGradient(@. Geometry.Covariant3Vector(2 * u_bottom))
96-
bcs_top_advection = Operators.SetGradient(@. Geometry.Covariant3Vector(-2 * u_top))
97+
bcs_bottom_advection =
98+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(2 * u_bottom)))
99+
bcs_top_advection =
100+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(-2 * u_top)))
97101
interpf2c = Operators.InterpolateF2C()
98102
gradc2f_advect =
99103
Operators.GradientC2F(top = bcs_top_advection, bottom = bcs_bottom_advection)
@@ -105,11 +109,14 @@ function tendency!(dY, Y, _, t)
105109
bcs_bottom = Operators.SetValue(Geometry.WVector(Cd * u_wind * v_1)) # Eq. 4.17
106110
v_top = Fields.level(v, Fields.nlevels(v))
107111
v_bottom = Fields.level(v, 1)
108-
bcs_top = Operators.SetGradient(@. Geometry.Covariant3Vector(2 * (FT(vg) - v_top))) # Eq. 4.19
112+
bcs_top =
113+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(2 * (FT(vg) - v_top)))) # Eq. 4.19
109114
gradc2f = Operators.GradientC2F(top = bcs_top)
110115
divf2c = Operators.DivergenceF2C(bottom = bcs_bottom)
111-
bcs_bottom_advection = Operators.SetGradient(@. Geometry.Covariant3Vector(2 * v_bottom))
112-
bcs_top_advection = Operators.SetGradient(@. Geometry.Covariant3Vector(-2 * v_top))
116+
bcs_bottom_advection =
117+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(2 * v_bottom)))
118+
bcs_top_advection =
119+
Operators.SetGradient(@. lazy(Geometry.Covariant3Vector(-2 * v_top)))
113120
gradc2f_advect =
114121
Operators.GradientC2F(top = bcs_top_advection, bottom = bcs_bottom_advection)
115122
@. dv =

‎examples/hybrid/plane/density_current_2dinvariant_rhoe.jl‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import ClimaCore:
1919
Operators,
2020
Utilities
2121
using ClimaCore.Geometry
22+
import LazyBroadcast: lazy
2223

2324
using Logging: global_logger
2425
using TerminalLoggers: TerminalLogger
@@ -256,10 +257,14 @@ function rhs_invariant!(dY, Y, _, t)
256257
lg_top_center = Fields.level(lg_field_centers, Fields.nlevels(lg_field_centers))
257258
ᶜ∇ᵥw_bottom = Fields.level(ᶜ∇ᵥw, 1)
258259
ᶜ∇ᵥw_top = Fields.level(ᶜ∇ᵥw, Fields.nlevels(ᶜ∇ᵥw))
259-
bottom_divergence = @. Geometry.Jcontravariant3(ᶜ∇ᵥw_bottom, lg_bottom_center) *
260-
(2 * inv(lg_bottom_face.J))
260+
bottom_divergence = @. lazy(
261+
Geometry.Jcontravariant3(ᶜ∇ᵥw_bottom, lg_bottom_center) *
262+
(2 * inv(lg_bottom_face.J)),
263+
)
261264
top_divergence =
262-
@. Geometry.Jcontravariant3(ᶜ∇ᵥw_top, lg_top_center) * (-2 * inv(lg_top_face.J))
265+
@. lazy(
266+
Geometry.Jcontravariant3(ᶜ∇ᵥw_top, lg_top_center) * (-2 * inv(lg_top_face.J)),
267+
)
263268
set_bcs = Operators.SetBoundaryOperator(
264269
bottom = Operators.SetValue(bottom_divergence),
265270
top = Operators.SetValue(top_divergence),

‎examples/hybrid/plane/topo_agnesi_nh.jl‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import ClimaCore:
1616
Operators,
1717
Hypsography
1818
using ClimaCore.Geometry
19+
import LazyBroadcast: lazy
1920

2021
using DiffEqCallbacks
2122

@@ -311,10 +312,14 @@ function rhs_invariant!(dY, Y, _, t)
311312
lg_top_center = Fields.level(lg_field_centers, Fields.nlevels(lg_field_centers))
312313
ᶜ∇ᵥw_bottom = Fields.level(ᶜ∇ᵥw, 1)
313314
ᶜ∇ᵥw_top = Fields.level(ᶜ∇ᵥw, Fields.nlevels(ᶜ∇ᵥw))
314-
bottom_divergence = @. Geometry.Jcontravariant3(ᶜ∇ᵥw_bottom, lg_bottom_center) *
315-
(2 * inv(lg_bottom_face.J))
315+
bottom_divergence = @. lazy(
316+
Geometry.Jcontravariant3(ᶜ∇ᵥw_bottom, lg_bottom_center) *
317+
(2 * inv(lg_bottom_face.J)),
318+
)
316319
top_divergence =
317-
@. Geometry.Jcontravariant3(ᶜ∇ᵥw_top, lg_top_center) * (-2 * inv(lg_top_face.J))
320+
@. lazy(
321+
Geometry.Jcontravariant3(ᶜ∇ᵥw_top, lg_top_center) * (-2 * inv(lg_top_face.J)),
322+
)
318323
set_bcs = Operators.SetBoundaryOperator(
319324
bottom = Operators.SetValue(bottom_divergence),
320325
top = Operators.SetValue(top_divergence),

0 commit comments

Comments
 (0)