Improve free surface stabilization kernel (#189)

* improve free surface stabilization

* update Van Keken test

* update more miniapps

* format

miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl

@@ -6,6 +6,8 @@ using JustRelax, JustRelax.JustRelax2D
 import JustRelax.@cell
 const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
 
+using GLMakie
+
 using JustPIC, JustPIC._2D
 # Threads is the default backend,
 # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script,
@@ -75,9 +77,9 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
     )
 
     # Initialize particles -------------------------------
-    nxcell, max_p, min_p = 40, 80, 20
+    nxcell, max_p, min_p = 20, 30, 10
     particles            = init_particles(
-        backend, nxcell, max_p, min_p, xvi..., di..., nx, ny
+        backend, nxcell, max_p, min_p, xvi, di, ni
     )
     # velocity grids
     grid_vx, grid_vy     = velocity_grids(xci, xvi, di)
@@ -146,16 +148,16 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
             igg;
             kwargs = (
                 iterMax          = 10e3,
-                nout             = 50,
+                nout             = 1e3,
                 viscosity_cutoff = (-Inf, Inf)
             )
         )
-        dt = compute_dt(stokes, di) / 10
+        dt = compute_dt(stokes, di) * 0.8
         # ------------------------------
 
         # Compute U rms ---------------
         Urms_it = let
-            JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true)
+            velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true)
             @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude
             sum(Vx_v.^2) * prod(di) |> sqrt
         end
@@ -176,7 +178,7 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
         t        += dt
 
         # Plotting ---------------------
-        if it == 1 || rem(it, 1000) == 0 || t >= tmax
+        if it == 1 || rem(it, 25) == 0 || t >= tmax
             fig = Figure(size = (1000, 1000), title = "t = $t")
             ax1 = Axis(fig[1,1], aspect = 1/λ, title = "t=$t")
             heatmap!(ax1, xvi[1], xvi[2], Array(ρg[2]), colormap = :oleron)
@@ -193,9 +195,9 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
 end
 
 figdir = "VanKeken"
-n      = 128 + 2
-nx     = n - 2
-ny     = n - 2
+n      = 64
+nx     = n
+ny     = n
 igg  = if !(JustRelax.MPI.Initialized())
     IGG(init_global_grid(nx, ny, 1; init_MPI = true)...)
 else

miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl

@@ -106,7 +106,7 @@ function RT_2D(igg, nx, ny)
     # Initialize particles -------------------------------
     nxcell, max_xcell, min_xcell = 30, 40, 10
     particles = init_particles(
-        backend, nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny
+        backend, nxcell, max_xcell, min_xcell, xvi, di, ni
     )
     # velocity grids
     grid_vx, grid_vy = velocity_grids(xci, xvi, di)
@@ -237,7 +237,7 @@ end
 ## END OF MAIN SCRIPT ----------------------------------------------------------------
 
 # (Path)/folder where output data and figures are stored
-n        = 100
+n        = 50
 nx       = n
 ny       = n
 igg      = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid
@@ -246,4 +246,4 @@ else
     igg
 end
 
-# RT_2D(igg, nx, ny)
+RT_2D(igg, nx, ny)

src/Interpolations.jl

@@ -6,14 +6,24 @@ Interpolates the values of `Vx` onto the grid points of `Vy`.
 # Arguments
 - `Vx_on_Vy::AbstractArray`: `Vx` at `Vy` grid points.
 - `Vx::AbstractArray`: `Vx` at its staggered grid points.
-
-
 """
 @parallel_indices (i, j) function interp_Vx_on_Vy!(Vx_on_Vy, Vx)
     Vx_on_Vy[i + 1, j] = 0.25 * (Vx[i, j] + Vx[i + 1, j] + Vx[i, j + 1] + Vx[i + 1, j + 1])
     return nothing
 end
 
+@parallel_indices (i, j) function interp_Vx∂ρ∂x_on_Vy!(Vx_on_Vy, Vx, ρg, _dx)
+    nx, ny = size(ρg)
+    ii = clamp(i, 1, nx)
+    ii1 = clamp(i + 1, 1, nx)
+    jj = clamp(j, 1, ny)
+    Vx_on_Vy[i + 1, j] =
+        (0.25 * (Vx[i, j] + Vx[i + 1, j] + Vx[i, j + 1] + Vx[i + 1, j + 1])) *
+        (ρg[ii1, jj] - ρg[ii, jj]) *
+        _dx
+    return nothing
+end
+
 # From cell vertices to cell center
 
 temperature2center!(thermal) = temperature2center!(backend(thermal), thermal)

src/boundaryconditions/no_slip.jl

@@ -1,17 +1,21 @@
 @views function no_slip!(Ax, Ay, bc)
     if bc.left
         Ax[1, :] .= 0
+        Ay[2, :] .= Ay[3, :] / 3
         Ay[1, :] .= -Ay[2, :]
     end
     if bc.right
         Ax[end, :] .= 0
+        Ay[end - 1, :] .= Ay[end - 2, :] / 3
         Ay[end, :] .= -Ay[end - 1, :]
     end
     if bc.bot
+        Ax[:, 2] .= Ax[:, 3] / 3
         Ax[:, 1] .= -Ax[:, 2]
         Ay[:, 1] .= 0
     end
     if bc.top
+        Ax[:, end - 1] .= Ax[:, end - 2] / 3
         Ax[:, end] .= -Ax[:, end - 1]
         Ay[:, end] .= 0
     end
@@ -20,34 +24,46 @@ end
 @views function no_slip!(Ax, Ay, Az, bc)
     if bc.left
         Ax[1, :, :] .= 0
+        Ay[2, :, :] .= Ay[3, :, :] / 3
         Ay[1, :, :] .= -Ay[2, :, :]
+        Az[2, :, :] .= Az[3, :, :] / 3
         Az[1, :, :] .= -Az[2, :, :]
     end
     if bc.right
         Ax[end, :, :] .= 0
+        Ay[end - 1, :, :] .= Ay[end - 2, :, :] / 3
         Ay[end, :, :] .= -Ay[end - 1, :, :]
+        Az[end - 1, :, :] .= Az[end - 2, :, :] / 3
         Az[end, :, :] .= -Az[end - 1, :, :]
     end
 
     if bc.front
+        Ax[:, 2, :] .= -Ax[:, 3, :] / 3
         Ax[:, 1, :] .= -Ax[:, 2, :]
         Ay[:, 1, :] .= 0
+        Az[:, 2, :] .= -Az[:, 3, :] / 3
         Az[:, 1, :] .= -Az[:, 2, :]
     end
 
     if bc.back
+        Ax[:, end - 1, :] .= -Ax[:, end - 2, :] / 3
         Ax[:, end, :] .= -Ax[:, end - 1, :]
         Ay[:, end, :] .= 0
+        Az[:, end - 1, :] .= -Az[:, end - 2, :] / 3
         Az[:, end, :] .= -Az[:, end - 1, :]
     end
 
     if bc.bot
+        Ax[:, :, 1] .= -Ax[:, :, 3] / 3
         Ax[:, :, 1] .= -Ax[:, :, 2]
+        Ay[:, :, 1] .= -Ay[:, :, 3] / 3
         Ay[:, :, 1] .= -Ay[:, :, 2]
         Az[:, :, 1] .= 0
     end
     if bc.top
+        Ax[:, :, end - 1] .= -Ax[:, :, end - 2] / 3
         Ax[:, :, end] .= -Ax[:, :, end - 1]
+        Ay[:, :, end - 1] .= -Ay[:, :, end - 2] / 3
         Ay[:, :, end] .= -Ay[:, :, end - 1]
         Az[:, :, end] .= 0
     end

src/stokes/Stokes2D.jl

@@ -578,10 +578,14 @@ function _solve!(
                 args,
             )
             center2vertex!(stokes.τ.xy, stokes.τ.xy_c)
-            update_halo!(stokes.τ.xy)
+            # update_halo!(stokes.τ.xy)
 
-            @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!(
-                Vx_on_Vy, stokes.V.Vx
+            # @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!(
+            #     Vx_on_Vy, stokes.V.Vx
+            # )
+
+            @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!(
+                Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1]
             )
 
             @hide_communication b_width begin # communication/computation overlap
@@ -599,7 +603,7 @@ function _solve!(
                 # apply boundary conditions
                 free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di)
                 flow_bcs!(stokes, flow_bcs)
-                update_halo!(@velocity(stokes)...)
+                # update_halo!(@velocity(stokes)...)
             end
         end
 
@@ -621,8 +625,10 @@ function _solve!(
             )
             # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP)))
             errs = (
-                norm_mpi(stokes.R.Rx) / length(stokes.R.Rx),
-                norm_mpi(stokes.R.Ry) / length(stokes.R.Ry),
+                norm_mpi(@views stokes.R.Rx[2:(end - 1), 2:(end - 1)]) /
+                length(stokes.R.Rx),
+                norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) /
+                length(stokes.R.Ry),
                 norm_mpi(stokes.R.RP) / length(stokes.R.RP),
             )
             push!(norm_Rx, errs[1])

src/stokes/Stokes3D.jl

@@ -296,9 +296,21 @@ function _solve!(
             # push!(norm_Ry, maximum_mpi(abs.(stokes.R.Ry)))
             # push!(norm_Rz, maximum_mpi(abs.(stokes.R.Rz)))
             # push!(norm_∇V, maximum_mpi(abs.(stokes.R.RP)))
-            push!(norm_Rx, norm_mpi(stokes.R.Rx) / length(stokes.R.Rx))
-            push!(norm_Ry, norm_mpi(stokes.R.Ry) / length(stokes.R.Ry))
-            push!(norm_Rz, norm_mpi(stokes.R.Rz) / length(stokes.R.Rz))
+            push!(
+                norm_Rx,
+                norm_mpi(stokes.R.Rx[2:(end - 1), 2:(end - 1), 2:(end - 1)]) /
+                length(stokes.R.Rx),
+            )
+            push!(
+                norm_Ry,
+                norm_mpi(stokes.R.Ry[2:(end - 1), 2:(end - 1), 2:(end - 1)]) /
+                length(stokes.R.Ry),
+            )
+            push!(
+                norm_Rz,
+                norm_mpi(stokes.R.Rz[2:(end - 1), 2:(end - 1), 2:(end - 1)]) /
+                length(stokes.R.Rz),
+            )
             push!(norm_∇V, norm_mpi(stokes.R.RP) / length(stokes.R.RP))
 
             err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont])
@@ -493,7 +505,10 @@ function _solve!(
         if iter % nout == 0 && iter > 1
             cont += 1
             for (norm_Ri, Ri) in zip((norm_Rx, norm_Ry, norm_Rz), @residuals(stokes.R))
-                push!(norm_Ri, norm_mpi(Ri) / length(Ri))
+                push!(
+                    norm_Ri,
+                    norm_mpi(Ri[2:(end - 1), 2:(end - 1), 2:(end - 1)]) / length(Ri),
+                )
             end
             push!(norm_∇V, norm_mpi(stokes.R.RP) / length(stokes.R.RP))
             err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont])

src/stokes/VelocityKernels.jl

@@ -134,26 +134,51 @@ end
     end
 
     if all((i, j) .< size(Vy) .- 1)
+        # θ = 1.0
+        # # Interpolate Vx into Vy node
+        # Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1]
+        # # Vertical velocity
+        # Vyᵢⱼ = Vy[i + 1, j + 1]
+        # # Get necessary buoyancy forces
+        # i_W, i_E = max(i - 1, 1), min(i + 1, nx)
+        # j_N = min(j + 1, ny)
+        # ρg_stencil = (
+        #     ρgy[i_W, j], ρgy[i, j], ρgy[i_E, j], ρgy[i_W, j_N], ρgy[i, j_N], ρgy[i_E, j_N]
+        # )
+        # ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25
+        # ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25
+        # ρg_S = ρg_stencil[2]
+        # ρg_N = ρg_stencil[5]
+        # # Spatial derivatives
+        # ∂ρg∂x = (ρg_E - ρg_W) * _dx
+        # ∂ρg∂y = (ρg_N - ρg_S) * _dy
+        # # correction term
+        # ρg_correction = (Vxᵢⱼ * ∂ρg∂x + Vyᵢⱼ * ∂ρg∂y) * θ * dt
+
+        # Vy[i + 1, j + 1] +=
+        #     (-d_ya(P) + d_ya(τyy) + d_xi(τxy) - av_ya(ρgy) - ρg_correction) * ηdτ /
+        #     av_ya(ητ)
+
         θ = 1.0
-        # Interpolate Vx into Vy node
+        # Interpolated Vx into Vy node (includes density gradient)
         Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1]
         # Vertical velocity
         Vyᵢⱼ = Vy[i + 1, j + 1]
         # Get necessary buoyancy forces
-        i_W, i_E = max(i - 1, 1), min(i + 1, nx)
+        # i_W, i_E = max(i - 1, 1), min(i + 1, nx)
         j_N = min(j + 1, ny)
-        ρg_stencil = (
-            ρgy[i_W, j], ρgy[i, j], ρgy[i_E, j], ρgy[i_W, j_N], ρgy[i, j_N], ρgy[i_E, j_N]
-        )
-        ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25
-        ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25
-        ρg_S = ρg_stencil[2]
-        ρg_N = ρg_stencil[5]
+        # ρg_stencil = (
+        #     ρgy[i_W, j], ρgy[i, j], ρgy[i_E, j], ρgy[i_W, j_N], ρgy[i, j_N], ρgy[i_E, j_N]
+        # )
+        # ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25
+        # ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25
+        ρg_S = ρgy[i, j]
+        ρg_N = ρgy[i, j_N]
         # Spatial derivatives
-        ∂ρg∂x = (ρg_E - ρg_W) * _dx
+        # ∂ρg∂x = (ρg_E - ρg_W) * _dx
         ∂ρg∂y = (ρg_N - ρg_S) * _dy
         # correction term
-        ρg_correction = (Vxᵢⱼ * ∂ρg∂x + Vyᵢⱼ * ∂ρg∂y) * θ * dt
+        ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt
 
         Vy[i + 1, j + 1] +=
             (-d_ya(P) + d_ya(τyy) + d_xi(τxy) - av_ya(ρgy) + ρg_correction) * ηdτ /
@@ -265,33 +290,48 @@ end
             Rx[i, j] = d_xa(τxx) + d_yi(τxy) - d_xa(P) - av_xa(ρgx)
         end
         if all((i, j) .≤ size(Ry))
+            # θ = 1.0
+            # Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1]
+            # # Vertical velocity
+            # Vyᵢⱼ = Vy[i + 1, j + 1]
+            # # Get necessary buoyancy forces
+            # i_W, i_E = max(i - 1, 1), min(i + 1, nx)
+            # j_N = min(j + 1, ny)
+            # ρg_stencil = (
+            #     ρgy[i_W, j],
+            #     ρgy[i, j],
+            #     ρgy[i_E, j],
+            #     ρgy[i_W, j_N],
+            #     ρgy[i, j_N],
+            #     ρgy[i_E, j_N],
+            # )
+            # ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25
+            # ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25
+            # ρg_S = ρg_stencil[2]
+            # ρg_N = ρg_stencil[5]
+            # # Spatial derivatives
+            # ∂ρg∂x = (ρg_E - ρg_W) * _dx
+            # ∂ρg∂y = (ρg_N - ρg_S) * _dy
+            # # correction term
+            # ρg_correction = (Vxᵢⱼ * ∂ρg∂x + Vyᵢⱼ * ∂ρg∂y) * θ * dt
+            # Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy) + ρg_correction
+            # # Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy)
+
             θ = 1.0
-            #     0.25 * (Vx[i, j + 1] + Vx[i + 1, j + 1] + Vx[i, j + 2] + Vx[i + 1, j + 2])
+            # Interpolated Vx into Vy node (includes density gradient)
             Vxᵢⱼ = Vx_on_Vy[i + 1, j + 1]
             # Vertical velocity
             Vyᵢⱼ = Vy[i + 1, j + 1]
             # Get necessary buoyancy forces
-            i_W, i_E = max(i - 1, 1), min(i + 1, nx)
             j_N = min(j + 1, ny)
-            ρg_stencil = (
-                ρgy[i_W, j],
-                ρgy[i, j],
-                ρgy[i_E, j],
-                ρgy[i_W, j_N],
-                ρgy[i, j_N],
-                ρgy[i_E, j_N],
-            )
-            ρg_W = (ρg_stencil[1] + ρg_stencil[2] + ρg_stencil[4] + ρg_stencil[5]) * 0.25
-            ρg_E = (ρg_stencil[2] + ρg_stencil[3] + ρg_stencil[5] + ρg_stencil[6]) * 0.25
-            ρg_S = ρg_stencil[2]
-            ρg_N = ρg_stencil[5]
+            ρg_S = ρgy[i, j]
+            ρg_N = ρgy[i, j_N]
             # Spatial derivatives
-            ∂ρg∂x = (ρg_E - ρg_W) * _dx
             ∂ρg∂y = (ρg_N - ρg_S) * _dy
             # correction term
-            ρg_correction = (Vxᵢⱼ * ∂ρg∂x + Vyᵢⱼ * ∂ρg∂y) * θ * dt
+            ρg_correction = (Vxᵢⱼ + Vyᵢⱼ * ∂ρg∂y) * θ * dt
+
             Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy) + ρg_correction
-            # Ry[i, j] = d_ya(τyy) + d_xi(τxy) - d_ya(P) - av_ya(ρgy)
         end
     end