Hoist some getproperty calls inside broadcast expr

CliMA · Jan 13, 2022 · 80edbf5 · 80edbf5
1 parent 1cde3b9
commit 80edbf5
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 62 deletions.
diff --git a/examples/3dsphere/baroclinic_wave.jl b/examples/3dsphere/baroclinic_wave.jl
@@ -185,6 +185,7 @@ function rhs!(dY, Y, _, t)
     dw = dY.w
     duₕ = dY.uₕ
     dρe = dY.Yc.ρe
+    z = coords.z
 
     # 0) update w at the bottom
     # fw = -g^31 cuₕ/ g^33
@@ -275,7 +276,7 @@ function rhs!(dY, Y, _, t)
         Geometry.Contravariant12Vector(Geometry.Covariant123Vector(cuₕ))
 
     ce = @. cρe / cρ
-    cp = @. pressure(cρ, ce, norm(cuvw), coords.z)
+    cp = @. pressure(cρ, ce, norm(cuvw), z)
 
     @. duₕ -= hgrad(cp) / cρ
     vgradc2f = Operators.GradientC2F(
@@ -284,7 +285,7 @@ function rhs!(dY, Y, _, t)
     )
     @. dw -= vgradc2f(cp) / Ic2f(cρ)
 
-    cE = @. (norm(cuvw)^2) / 2 + Φ(coords.z)
+    cE = @. (norm(cuvw)^2) / 2 + Φ(z)
     @. duₕ -= hgrad(cE)
     @. dw -= vgradc2f(cE)
 

diff --git a/examples/3dsphere/solid_body_rotation_3d.jl b/examples/3dsphere/solid_body_rotation_3d.jl
@@ -97,6 +97,7 @@ function rhs!(dY, Y, _, t)
     dw = dY.w
     duₕ = dY.uₕ
     dρe = dY.Yc.ρe
+    z = c_coords.z
 
     # # 0) update w at the bottom
     # fw = -g^31 cuₕ/ g^33 ????????
@@ -170,7 +171,7 @@ function rhs!(dY, Y, _, t)
         Geometry.Contravariant12Vector(Geometry.Covariant123Vector(cuₕ))
 
     ce = @. cρe / cρ
-    cp = @. pressure(cρ, ce, norm(cuvw), c_coords.z)
+    cp = @. pressure(cρ, ce, norm(cuvw), z)
 
     @. duₕ -= hgrad(cp) / cρ
     vgradc2f = Operators.GradientC2F(
@@ -179,7 +180,7 @@ function rhs!(dY, Y, _, t)
     )
     @. dw -= vgradc2f(cp) / Ic2f(cρ)
 
-    cE = @. (norm(cuvw)^2) / 2 + Φ(c_coords.z)
+    cE = @. (norm(cuvw)^2) / 2 + Φ(z)
     @. duₕ -= hgrad(cE)
     @. dw -= vgradc2f(cE)
 

diff --git a/examples/hybrid/bubble3d_invariant_totalenergy.jl b/examples/hybrid/bubble3d_invariant_totalenergy.jl
@@ -126,6 +126,7 @@ function rhs_invariant!(dY, Y, _, t)
     dw = dY.w
     duₕ = dY.uₕ
     dρe = dY.Yc.ρe
+    z = coords.z
 
     # 0) update w at the bottom
     # fw = -g^31 cuₕ/ g^33
@@ -219,8 +220,8 @@ function rhs_invariant!(dY, Y, _, t)
         cω³ × Geometry.Contravariant12Vector(Geometry.Covariant123Vector(cuₕ))
 
     ce = @. cρe / cρ
-    #cp = @. pressure(cρ, ce, cuvw, coords.z) #TODO: the pressure function doesn't work (broadcast error)
-    cI = @. ce - Φ(coords.z) - (norm(cuvw)^2) / 2
+    #cp = @. pressure(cρ, ce, cuvw, z) #TODO: the pressure function doesn't work (broadcast error)
+    cI = @. ce - Φ(z) - (norm(cuvw)^2) / 2
     cT = @. cI / C_v + T_0
     cp = @. cρ * R_d * cT
 
@@ -231,7 +232,7 @@ function rhs_invariant!(dY, Y, _, t)
     )
     @. dw -= vgradc2f(cp) / Ic2f(cρ)
 
-    cE = @. (norm(cuvw)^2) / 2 + Φ(coords.z)
+    cE = @. (norm(cuvw)^2) / 2 + Φ(z)
     @. duₕ -= hgrad(cE)
     @. dw -= vgradc2f(cE)
 

diff --git a/examples/hybrid/bubble_2d.jl b/examples/hybrid/bubble_2d.jl
@@ -154,6 +154,9 @@ function rhs!(dY, Y, _, t)
     dYc = dY.Yc
     dρuₕ = dY.ρuₕ
     dρw = dY.ρw
+    ρ = Yc.ρ
+    ρθ = Yc.ρθ
+    dρθ = dYc.ρθ
 
     # spectral horizontal operators
     hdiv = Operators.Divergence()
@@ -201,34 +204,34 @@ function rhs!(dY, Y, _, t)
         top = Operators.Extrapolate(),
     )
 
-    uₕ = @. ρuₕ / Yc.ρ
-    w = @. ρw / If(Yc.ρ)
-    wc = @. Ic(ρw) / Yc.ρ
-    p = @. pressure(Yc.ρθ)
-    θ = @. Yc.ρθ / Yc.ρ
-    Yfρ = @. If(Yc.ρ)
+    uₕ = @. ρuₕ / ρ
+    w = @. ρw / If(ρ)
+    wc = @. Ic(ρw) / ρ
+    p = @. pressure(ρθ)
+    θ = @. ρθ / ρ
+    Yfρ = @. If(ρ)
 
     ### HYPERVISCOSITY
     # 1) compute hyperviscosity coefficients
-    @. dYc.ρθ = hdiv(hgrad(θ))
+    @. dρθ = hdiv(hgrad(θ))
     @. dρuₕ = hdiv(hgrad(uₕ))
     @. dρw = hdiv(hgrad(w))
     Spaces.weighted_dss!(dYc)
     Spaces.weighted_dss!(dρuₕ)
     Spaces.weighted_dss!(dρw)
 
     κ₄ = 100.0 # m^4/s
-    @. dYc.ρθ = -κ₄ * hdiv(Yc.ρ * hgrad(dYc.ρθ))
-    @. dρuₕ = -κ₄ * hdiv(Yc.ρ * hgrad(dρuₕ))
+    @. dYc.ρθ = -κ₄ * hdiv(ρ * hgrad(dρθ))
+    @. dρuₕ = -κ₄ * hdiv(ρ * hgrad(dρuₕ))
     @. dρw = -κ₄ * hdiv(Yfρ * hgrad(dρw))
 
     # density
-    @. dYc.ρ = -∂(ρw)
-    @. dYc.ρ -= hdiv(ρuₕ)
+    @. dρ = -∂(ρw)
+    @. dρ -= hdiv(ρuₕ)
 
     # potential temperature
-    @. dYc.ρθ += -(∂(ρw * If(Yc.ρθ / Yc.ρ)))
-    @. dYc.ρθ -= hdiv(uₕ * Yc.ρθ)
+    @. dρθ += -(∂(ρw * If(ρθ / ρ)))
+    @. dρθ -= hdiv(uₕ * ρθ)
 
     # horizontal momentum
     Ih = Ref(
@@ -241,41 +244,40 @@ function rhs!(dY, Y, _, t)
     @. dρuₕ -= hdiv(ρuₕ ⊗ uₕ + p * Ih)
 
     # vertical momentum
+    z = coords.z
     @. dρw +=
         B(
-            Geometry.transform(
-                Geometry.WAxis(),
-                -(∂f(p)) - If(Yc.ρ) * ∂f(Φ(coords.z)),
-            ) - vvdivc2f(Ic(ρw ⊗ w)),
+            Geometry.transform(Geometry.WAxis(), -(∂f(p)) - If(ρ) * ∂f(Φ(z))) -
+            vvdivc2f(Ic(ρw ⊗ w)),
         ) + fcf(wc, ρw)
-    uₕf = @. If(ρuₕ / Yc.ρ) # requires boundary conditions
+    uₕf = @. If(ρuₕ / ρ) # requires boundary conditions
     @. dρw -= hdiv(uₕf ⊗ ρw)
 
     ### UPWIND FLUX CORRECTION
     upwind_correction = true
     if upwind_correction
-        @. dYc.ρ += fcc(w, Yc.ρ)
-        @. dYc.ρθ += fcc(w, Yc.ρθ)
+        @. dρ += fcc(w, ρ)
+        @. dρθ += fcc(w, ρθ)
         @. dρuₕ += fcc(w, ρuₕ)
         @. dρw += fcf(wc, ρw)
     end
 
     ### DIFFUSION
     κ₂ = 0.0 # m^2/s
     #  1a) horizontal div of horizontal grad of horiz momentun
-    @. dρuₕ += hdiv(κ₂ * (Yc.ρ * hgrad(ρuₕ / Yc.ρ)))
+    @. dρuₕ += hdiv(κ₂ * (ρ * hgrad(ρuₕ / ρ)))
     #  1b) vertical div of vertical grad of horiz momentun
-    @. dρuₕ += uvdivf2c(κ₂ * (Yfρ * ∂f(ρuₕ / Yc.ρ)))
+    @. dρuₕ += uvdivf2c(κ₂ * (Yfρ * ∂f(ρuₕ / ρ)))
 
     #  1c) horizontal div of horizontal grad of vert momentum
     @. dρw += hdiv(κ₂ * (Yfρ * hgrad(ρw / Yfρ)))
     #  1d) vertical div of vertical grad of vert momentun
-    @. dρw += vvdivc2f(κ₂ * (Yc.ρ * ∂c(ρw / Yfρ)))
+    @. dρw += vvdivc2f(κ₂ * (ρ * ∂c(ρw / Yfρ)))
 
     #  2a) horizontal div of horizontal grad of potential temperature
-    @. dYc.ρθ += hdiv(κ₂ * (Yc.ρ * hgrad(Yc.ρθ / Yc.ρ)))
+    @. dYc.ρθ += hdiv(κ₂ * (ρ * hgrad(ρθ / ρ)))
     #  2b) vertical div of vertial grad of potential temperature
-    @. dYc.ρθ += ∂(κ₂ * (Yfρ * ∂f(Yc.ρθ / Yc.ρ)))
+    @. dYc.ρθ += ∂(κ₂ * (Yfρ * ∂f(ρθ / ρ)))
 
     Spaces.weighted_dss!(dYc)
     Spaces.weighted_dss!(dρuₕ)

diff --git a/examples/hybrid/bubble_3d.jl b/examples/hybrid/bubble_3d.jl
@@ -162,6 +162,11 @@ function rhs!(dY, Y, _, t)
     Yc = Y.Yc
     dYc = dY.Yc
     dρw = dY.ρw
+    ρ = Yc.ρ
+    ρuₕ = Yc.ρuₕ
+    ρθ = Yc.ρθ
+    dρθ = dYc.ρθ
+    dρuₕ = dYc.ρuₕ
 
     # spectral horizontal operators
     hdiv = Operators.Divergence()
@@ -211,79 +216,78 @@ function rhs!(dY, Y, _, t)
         top = Operators.Extrapolate(),
     )
 
-    uₕ = @. Yc.ρuₕ / Yc.ρ
-    w = @. ρw / If(Yc.ρ)
-    wc = @. Ic(ρw) / Yc.ρ
-    p = @. pressure(Yc.ρθ)
-    θ = @. Yc.ρθ / Yc.ρ
-    Yfρ = @. If(Yc.ρ)
+    uₕ = @. ρuₕ / ρ
+    w = @. ρw / If(ρ)
+    wc = @. Ic(ρw) / ρ
+    p = @. pressure(ρθ)
+    θ = @. ρθ / ρ
+    Yfρ = @. If(ρ)
 
     ### HYPERVISCOSITY
     # 1) compute hyperviscosity coefficients
-    @. dYc.ρθ = hdiv(hgrad(θ))
-    @. dYc.ρuₕ = hdiv(hgrad(uₕ))
+    @. dρθ = hdiv(hgrad(θ))
+    @. dρuₕ = hdiv(hgrad(uₕ))
     @. dρw = hdiv(hgrad(w))
     Spaces.weighted_dss!(dYc)
     Spaces.weighted_dss!(dρw)
 
     κ₄ = 100.0 # m^4/s
-    @. dYc.ρθ = -κ₄ * hdiv(Yc.ρ * hgrad(dYc.ρθ))
-    @. dYc.ρuₕ = -κ₄ * hdiv(Yc.ρ * hgrad(dYc.ρuₕ))
+    @. dρθ = -κ₄ * hdiv(ρ * hgrad(dρθ))
+    @. dρuₕ = -κ₄ * hdiv(ρ * hgrad(dρuₕ))
     @. dρw = -κ₄ * hdiv(Yfρ * hgrad(dρw))
 
     # density
-    @. dYc.ρ = -∂(ρw)
-    @. dYc.ρ -= hdiv(Yc.ρuₕ)
+    @. dρ = -∂(ρw)
+    @. dρ -= hdiv(ρuₕ)
 
     # potential temperature
-    @. dYc.ρθ += -(∂(ρw * If(Yc.ρθ / Yc.ρ)))
-    @. dYc.ρθ -= hdiv(uₕ * Yc.ρθ)
+    @. dρθ += -(∂(ρw * If(ρθ / ρ)))
+    @. dρθ -= hdiv(uₕ * ρθ)
 
     # Horizontal momentum
-    @. dYc.ρuₕ += -uvdivf2c(ρw ⊗ If(uₕ))
+    @. dρuₕ += -uvdivf2c(ρw ⊗ If(uₕ))
     Ih = Ref(
         Geometry.Axis2Tensor(
             (Geometry.UVAxis(), Geometry.UVAxis()),
             @SMatrix [1.0 0.0; 0.0 1.0]
         ),
     )
-    @. dYc.ρuₕ -= hdiv(Yc.ρuₕ ⊗ uₕ + p * Ih)
+    @. dρuₕ -= hdiv(ρuₕ ⊗ uₕ + p * Ih)
 
     # vertical momentum
+    z = coords.z
     @. dρw += B(
-        Geometry.transform(
-            Geometry.WAxis(),
-            -(∂f(p)) - If(Yc.ρ) * ∂f(Φ(coords.z)),
-        ) - vvdivc2f(Ic(ρw ⊗ w)),
+        Geometry.transform(Geometry.WAxis(), -(∂f(p)) - If(ρ) * ∂f(Φ(z))) -
+        vvdivc2f(Ic(ρw ⊗ w)),
     )
-    uₕf = @. If(Yc.ρuₕ / Yc.ρ) # requires boundary conditions
+    uₕf = @. If(ρuₕ / ρ) # requires boundary conditions
     @. dρw -= hdiv(uₕf ⊗ ρw)
 
     ### UPWIND FLUX CORRECTION
     upwind_correction = true
     if upwind_correction
-        @. dYc.ρ += fcc(w, Yc.ρ)
-        @. dYc.ρθ += fcc(w, Yc.ρθ)
-        @. dYc.ρuₕ += fcc(w, Yc.ρuₕ)
+        @. dρ += fcc(w, ρ)
+        @. dρθ += fcc(w, ρθ)
+        @. dρuₕ += fcc(w, ρuₕ)
         @. dρw += fcf(wc, ρw)
     end
 
     ### DIFFUSION
     κ₂ = 0.0 # m^2/s
     #  1a) horizontal div of horizontal grad of horiz momentun
-    @. dYc.ρuₕ += hdiv(κ₂ * (Yc.ρ * hgrad(Yc.ρuₕ / Yc.ρ)))
+    @. dρuₕ += hdiv(κ₂ * (ρ * hgrad(ρuₕ / ρ)))
     #  1b) vertical div of vertical grad of horiz momentun
-    @. dYc.ρuₕ += uvdivf2c(κ₂ * (Yfρ * ∂f(Yc.ρuₕ / Yc.ρ)))
+    @. dρuₕ += uvdivf2c(κ₂ * (Yfρ * ∂f(ρuₕ / ρ)))
 
     #  1c) horizontal div of horizontal grad of vert momentum
     @. dρw += hdiv(κ₂ * (Yfρ * hgrad(ρw / Yfρ)))
     #  1d) vertical div of vertical grad of vert momentun
     @. dρw += vvdivc2f(κ₂ * (Yc.ρ * ∂c(ρw / Yfρ)))
 
     #  2a) horizontal div of horizontal grad of potential temperature
-    @. dYc.ρθ += hdiv(κ₂ * (Yc.ρ * hgrad(Yc.ρθ / Yc.ρ)))
+    @. dρθ += hdiv(κ₂ * (Yc.ρ * hgrad(ρθ / ρ)))
     #  2b) vertical div of vertial grad of potential temperature
-    @. dYc.ρθ += ∂(κ₂ * (Yfρ * ∂f(Yc.ρθ / Yc.ρ)))
+    @. dρθ += ∂(κ₂ * (Yfρ * ∂f(ρθ / ρ)))
 
     Spaces.weighted_dss!(dYc)
     Spaces.weighted_dss!(dρw)

diff --git a/examples/hybrid/bubble_invariant.jl b/examples/hybrid/bubble_invariant.jl
@@ -120,6 +120,8 @@ function rhs_invariant!(dY, Y, _, t)
     dw = dY.w
     duₕ = dY.uₕ
     dρθ = dY.Yc.ρθ
+    ρθ = Yc.ρθ
+    z = coords.z
 
 
     # 0) update w at the bottom
@@ -209,13 +211,13 @@ function rhs_invariant!(dY, Y, _, t)
     )
     @. dw -= vgradc2f(cp) / Ic2f(cρ)
 
-    cE = @. (norm(cuw)^2) / 2 + Φ(coords.z)
+    cE = @. (norm(cuw)^2) / 2 + Φ(z)
     @. duₕ -= hgrad(cE)
     @. dw -= vgradc2f(cE)
 
     # 3) potential temperature
 
-    @. dρθ -= hdiv(cuw * Yc.ρθ)
+    @. dρθ -= hdiv(cuw * ρθ)
     @. dρθ -= vdivf2c(fw * Ic2f(cρθ))
     @. dρθ -= vdivf2c(Ic2f(cuₕ * cρθ))