Try to reduce allocations

CliMA · Nov 7, 2021 · 8696123 · 8696123
1 parent 78658d1
commit 8696123
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 24 deletions.
diff --git a/src/Operators.jl b/src/Operators.jl
@@ -166,7 +166,8 @@ function upwind_advection_area(ρ0_half, a_up, w_up, grid, k)
     a_up_cut = ccut_upwind(a_up, grid, k)
     w_up_cut = daul_f2c_upwind(w_up, grid, k)
     m_cut = ρ_0_cut .* a_up_cut .* w_up_cut
-    ∇m = c∇_upwind(m_cut, grid, k; bottom = SetValue(0), top = SetGradient(0))
+    FT = eltype(grid)
+    ∇m = c∇_upwind(m_cut, grid, k; bottom = SetValue(FT(0)), top = SetGradient(FT(0)))
     return -∇m / ρ0_half[k]
 end
 
@@ -175,7 +176,8 @@ function upwind_advection_velocity(ρ0, a_up, w_up, grid, k; a_up_bcs)
     ρ_0_dual = fcut_upwind(ρ0, grid, k)
     w_up_dual = fcut_upwind(w_up, grid, k)
     adv_dual = a_dual .* ρ_0_dual .* w_up_dual .* w_up_dual
-    ∇ρaw = f∇_onesided(adv_dual, grid, k; bottom = FreeBoundary(), top = SetGradient(0))
+    FT = eltype(grid)
+    ∇ρaw = f∇_onesided(adv_dual, grid, k; bottom = FreeBoundary(), top = SetGradient(FT(0)))
     return ∇ρaw
 end
 
@@ -185,7 +187,8 @@ function upwind_advection_scalar(ρ0_half, a_up, w_up, var, grid, k)
     w_up_cut = daul_f2c_upwind(w_up, grid, k)
     var_cut = ccut_upwind(var, grid, k)
     m_cut = ρ_0_cut .* a_up_cut .* w_up_cut .* var_cut
-    ∇m = c∇_upwind(m_cut, grid, k; bottom = SetValue(0), top = SetGradient(0))
+    FT = eltype(grid)
+    ∇m = c∇_upwind(m_cut, grid, k; bottom = SetValue(FT(0)), top = SetGradient(FT(0)))
     return ∇m
 end
 

diff --git a/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl
@@ -737,14 +737,14 @@ function compute_covariance_detr(edmf::EDMF_PrognosticTKE, grid, state, covar_sy
     aux_up = center_aux_updrafts(state)
     aux_en = center_aux_environment(state)
     covar = getproperty(aux_en, covar_sym)
+    detr_loss = aux_covar.detr_loss
+    Ic = CCO.InterpolateF2C()
 
-    @inbounds for k in real_center_indices(grid)
-        aux_covar.detr_loss[k] = 0.0
-        @inbounds for i in 1:(up.n_updrafts)
-            w_up_c = interpf2c(aux_up_f[i].w, grid, k)
-            aux_covar.detr_loss[k] += aux_up[i].area[k] * abs(w_up_c) * aux_up[i].entr_sc[k]
-        end
-        aux_covar.detr_loss[k] *= ρ0_c[k] * covar[k]
+    parent(detr_loss) .= 0
+    @inbounds for i in 1:(up.n_updrafts)
+        aux_up_i = aux_up[i]
+        aux_up_f_i = aux_up_f[i]
+        @. detr_loss += ρ0_c * covar * aux_up_i.area * abs(Ic(aux_up_f_i.w)) * aux_up_i.entr_sc
     end
     return
 end
@@ -759,13 +759,12 @@ function compute_covariance_dissipation(edmf::EDMF_PrognosticTKE, grid, state, c
     aux_en_2m = center_aux_environment_2m(state)
     aux_covar = getproperty(aux_en_2m, covar_sym)
     covar = getproperty(aux_en, covar_sym)
+    dissipation = aux_covar.dissipation
+    area_en = aux_en.area
+    tke_en = aux_en.tke
+    mixing_length = aux_tc.mixing_length
 
-    @inbounds for k in real_center_indices(grid)
-        aux_covar.dissipation[k] = (
-            ρ0_c[k] * aux_en.area[k] * covar[k] * max(aux_en.tke[k], 0)^0.5 / max(aux_tc.mixing_length[k], 1.0e-3) *
-            c_d
-        )
-    end
+    @. dissipation = ρ0_c * area_en * covar * max(tke_en, 0)^0.5 / max(mixing_length, 1.0e-3) * c_d
     return
 end
 
@@ -798,6 +797,12 @@ function compute_en_tendencies!(edmf, grid::Grid, state, param_set, TS, covar_sy
     ae = 1 .- aux_bulk.area
     aeK = is_tke ? ae .* KM : ae .* KH
 
+    press = aux_covar.press
+    buoy = aux_covar.buoy
+    shear = aux_covar.shear
+    entr_gain = aux_covar.entr_gain
+    rain_src = aux_covar.rain_src
+
     cartvec = CC.Geometry.Cartesian3Vector(zero(FT))
     aeK_bcs = (; bottom = CCO.SetValue(aeK[kc_surf]), top = CCO.SetValue(aeK[kc_toa]))
     prog_bcs = (; bottom = CCO.SetGradient(cartvec), top = CCO.SetGradient(cartvec))
@@ -834,17 +839,13 @@ function compute_en_tendencies!(edmf, grid::Grid, state, param_set, TS, covar_sy
         end
         dissipation = ρ0_c[k] * aux_en.area[k] * c_d * sqrt(max(aux_en.tke[k], 0)) / max(mixing_length[k], 1)
         ρaew_en_ϕ_cut = ccut_downwind(ρaew_en_ϕ, grid, k)
-        ∇ρaew_en_ϕ = c∇_downwind(ρaew_en_ϕ_cut, grid, k; bottom = FreeBoundary(), top = SetGradient(0))
+        ∇ρaew_en_ϕ = c∇_downwind(ρaew_en_ϕ_cut, grid, k; bottom = FreeBoundary(), top = SetGradient(FT(0)))
         if is_surface_center(grid, k)
             prog_covar[k] = covar_surf
         else
-            prog_covar[k] = (
-                aux_covar.press[k] +
-                aux_covar.buoy[k] +
-                aux_covar.shear[k] +
-                aux_covar.entr_gain[k] +
-                aux_covar.rain_src[k] - D_env * covar[k] - dissipation * covar[k] - ∇ρaew_en_ϕ + ∇ρ_ae_K∇ϕ[k]
-            )
+            prog_covar[k] =
+                press[k] + buoy[k] + shear[k] + entr_gain[k] + rain_src[k] - D_env * covar[k] - dissipation * covar[k] -
+                ∇ρaew_en_ϕ + ∇ρ_ae_K∇ϕ[k]
         end
     end