Use more grid abstractions

CliMA · Jul 27, 2021 · 06fbe15 · 06fbe15
1 parent 969df66
commit 06fbe15
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 29 deletions.
diff --git a/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl
@@ -1376,28 +1376,32 @@ function solve_updraft_scalars(self::EDMF_PrognosticTKE, GMV::GridMeanVariables)
 
     @inbounds for i in xrange(self.n_updrafts)
 
-        # at the surface
-        if self.UpdVar.Area.new[i, gw] >= self.minimum_area
-            self.UpdVar.H.new[i, gw] = self.h_surface_bc[i]
-            self.UpdVar.QT.new[i, gw] = self.qt_surface_bc[i]
-        else
-            self.UpdVar.H.new[i, gw] = GMV.H.values[gw]
-            self.UpdVar.QT.new[i, gw] = GMV.QT.values[gw]
-        end
+        # starting from the bottom do entrainment at each level
+        @inbounds for k in real_center_indicies(grid)
+            if k==gw
+                # at the surface
+                if self.UpdVar.Area.new[i, k] >= self.minimum_area
+                    self.UpdVar.H.new[i, k] = self.h_surface_bc[i]
+                    self.UpdVar.QT.new[i, k] = self.qt_surface_bc[i]
+                else
+                    self.UpdVar.H.new[i, k] = GMV.H.values[k]
+                    self.UpdVar.QT.new[i, k] = GMV.QT.values[k]
+                end
+
+                # saturation adjustment
+                sa = eos(
+                    self.UpdThermo.t_to_prog_fp,
+                    self.UpdThermo.prog_to_t_fp,
+                    ref_state.p0_half[k],
+                    self.UpdVar.QT.new[i, k],
+                    self.UpdVar.H.new[i, k],
+                )
+                self.UpdVar.QL.new[i, k] = sa.ql
+                self.UpdVar.T.new[i, k] = sa.T
+                continue
+            end
 
-        # saturation adjustment
-        sa = eos(
-            self.UpdThermo.t_to_prog_fp,
-            self.UpdThermo.prog_to_t_fp,
-            ref_state.p0_half[gw],
-            self.UpdVar.QT.new[i, gw],
-            self.UpdVar.H.new[i, gw],
-        )
-        self.UpdVar.QL.new[i, gw] = sa.ql
-        self.UpdVar.T.new[i, gw] = sa.T
 
-        # starting from the bottom do entrainment at each level
-        @inbounds for k in xrange(gw + 1, grid.nzg - gw)
             H_entr = self.EnvVar.H.values[k]
             QT_entr = self.EnvVar.QT.values[k]
 
@@ -2307,9 +2311,7 @@ function compute_tke_advection(self::EDMF_PrognosticTKE)
     ae = pyones(grid.nzg) .- self.UpdVar.Area.bulkvalues # area of environment
     drho_ae_we_e_plus = 0.0
 
-    # TODO: why minus 1?
-    # @inbounds for k in real_center_indicies(grid) # TODO: try changing to this
-    @inbounds for k in xrange(gw, get_grid(self).nzg - gw - 1)
+    @inbounds for k in real_face_indicies(grid)
         drho_ae_we_e_minus = drho_ae_we_e_plus
         drho_ae_we_e_plus =
             (
@@ -2338,9 +2340,7 @@ function compute_tke_transport(self::EDMF_PrognosticTKE)
     KM = diffusivity_m(self).values
     rho0_half = reference_state(self).rho0_half
 
-    # TODO: why minus 1?
-    # @inbounds for k in real_center_indicies(grid) # TODO: try changing to this
-    @inbounds for k in xrange(gw, grid.nzg - gw - 1)
+    @inbounds for k in real_face_indicies(grid)
         drho_ae_K_m_de_low = drho_ae_K_m_de_plus
         drho_ae_K_m_de_plus =
             (

diff --git a/src/Variables.jl b/src/Variables.jl
@@ -76,7 +76,7 @@ function zero_tendencies(self::GridMeanVariables)
 end
 
 function update(self::GridMeanVariables, TS::TimeStepping)
-    @inbounds for k in xrange(self.Gr.gw, self.Gr.nzg - self.Gr.gw)
+    @inbounds for k in center_indicies(self.Gr)
         self.U.values[k] += self.U.tendencies[k] * TS.dt
         self.V.values[k] += self.V.tendencies[k] * TS.dt
         self.H.values[k] += self.H.tendencies[k] * TS.dt
@@ -173,7 +173,7 @@ function mean_cloud_diagnostics(self)
     self.cloud_base = self.Gr.z_half[self.Gr.nzg - self.Gr.gw - 1]
     self.cloud_top = 0.0
 
-    @inbounds for k in xrange(self.Gr.gw, self.Gr.nzg - self.Gr.gw)
+    @inbounds for k in real_center_indicies(self.Gr)
         self.lwp += self.Ref.rho0_half[k] * self.QL.values[k] * self.Gr.dz
 
         if self.QL.values[k] > 1e-8
@@ -186,7 +186,7 @@ end
 
 function satadjust(self::GridMeanVariables)
     sa = eos_struct()
-    @inbounds for k in xrange(self.Gr.nzg)
+    @inbounds for k in center_indicies(self.Gr)
         h = self.H.values[k]
         qt = self.QT.values[k]
         p0 = self.Ref.p0_half[k]