Clean up notation in compute_w kernel (#4378)

* Clean up notation in compute_w kernel

* fix

Author: glwagner

src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl

@@ -41,15 +41,20 @@ compute_w_from_continuity!(velocities, arch, grid; parameters = w_kernel_paramet
 @kernel function _compute_w_from_continuity!(U, grid)
     i, j = @index(Global, NTuple)
 
-    @inbounds U.w[i, j, 1] = 0
-    for k in 2:grid.Nz+1
-        δh_u = flux_div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) / Azᶜᶜᶜ(i, j, k-1, grid) 
-        ∂tσ  = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_σ(i, j, k-1, grid)
+    u, v, w = U
+    wᵏ = zero(eltype(w))
+    @inbounds w[i, j, 1] = wᵏ
 
-        immersed = immersed_cell(i, j, k-1, grid)
-        Δw       = δh_u + ifelse(immersed, zero(grid), ∂tσ) # We do not account for grid changes in immersed cells
+    Nz = size(grid, 3)
+    for k in 2:Nz+1
+        δ = flux_div_xyᶜᶜᶜ(i, j, k-1, grid, u, v) / Azᶜᶜᶜ(i, j, k-1, grid) 
 
-        @inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δw
+        # We do not account for grid changes in immersed cells
+        not_immersed = !immersed_cell(i, j, k-1, grid)
+        w̃ = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_σ(i, j, k-1, grid) * not_immersed
+
+        wᵏ -= (δ + w̃) 
+        @inbounds w[i, j, k] = wᵏ
     end
 end