need to mask bouyancy forces as well

src/MiniKernels.jl

@@ -155,23 +155,23 @@ end
 ## Because mysum(::generator) does not work inside CUDA kernels...
 @inline mysum(A, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ranges...)
 
-@inline function mysum(f::F, A, ranges_i) where {F<:Function}
+@inline function mysum(f::F, A::AbstractArray, ranges_i) where {F<:Function}
     s = 0.0
     for i in ranges_i
         s += f(A[i])
     end
     return s
 end
 
-@inline function mysum(f::F, A, ranges_i, ranges_j) where {F<:Function}
+@inline function mysum(f::F, A::AbstractArray, ranges_i, ranges_j) where {F<:Function}
     s = 0.0
     for i in ranges_i, j in ranges_j
         s += f(A[i, j])
     end
     return s
 end
 
-@inline function mysum(f::F, A, ranges_i, ranges_j, ranges_k) where {F<:Function}
+@inline function mysum(f::F, A::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function}
     s = 0.0
     for i in ranges_i, j in ranges_j, k in ranges_k
         s += f(A[i, j, k])

src/variational_stokes/MiniKernels.jl

@@ -18,3 +18,44 @@ Base.@propagate_inbounds @inline _d_xi(A::T, ϕ::T, _dx, I::Vararg{Integer,N}) w
     (-front(A, ϕ, I...) + next(A, ϕ, I...)) * _dx
 Base.@propagate_inbounds @inline _d_yi(A::T, ϕ::T, _dy, I::Vararg{Integer,N}) where {N,T} =
     (-right(A, ϕ, I...) + next(A, ϕ, I...)) * _dy
+
+# averages
+Base.@propagate_inbounds @inline _av(A::T, ϕ::T, i, j) where {T<:T2} =
+    0.25 * mysum(A, ϕ, (i+1):(i+2), (j+1):(j+2))
+Base.@propagate_inbounds @inline _av_a(A::T, ϕ::T, i, j) where {T<:T2} =
+    0.25 * mysum(A, ϕ, (i):(i+1), (j):(j+1))
+Base.@propagate_inbounds @inline _av_xa(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} =
+    (center(A, ϕ, I...) + right(A, ϕ, I...)) * 0.5
+Base.@propagate_inbounds @inline _av_ya(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} =
+    (center(A, ϕ, I...) + front(A, ϕ, I...)) * 0.5
+Base.@propagate_inbounds @inline _av_xi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} =
+    (front(A, ϕ, I...), next(A, ϕ, I...)) * 0.5
+Base.@propagate_inbounds @inline _av_yi(A::T, ϕ::T, I::Vararg{Integer,2}) where {T<:T2} =
+    (right(A, ϕ, I...), next(A, ϕ, I...)) * 0.5
+
+## Because mysum(::generator) does not work inside CUDA kernels...
+@inline mysum(A, ϕ, ranges::Vararg{T,N}) where {T,N} = mysum(identity, A, ϕ, ranges...)
+
+@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i) where {F<:Function}
+    s = 0.0
+    for i in ranges_i
+        s += f(A[i]) * ϕ[i]
+    end
+    return s
+end
+
+@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j) where {F<:Function}
+    s = 0.0
+    for i in ranges_i, j in ranges_j
+        s += f(A[i, j]) * ϕ[i, j]
+    end
+    return s
+end
+
+@inline function mysum(f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i, ranges_j, ranges_k) where {F<:Function}
+    s = 0.0
+    for i in ranges_i, j in ranges_j, k in ranges_k
+        s += f(A[i, j, k]) * ϕ[i, j, k]
+    end
+    return s
+end

src/variational_stokes/VelocityKernels.jl

@@ -70,6 +70,8 @@ end
     d_xa(A, ϕ) = _d_xa(A, ϕ, _dx, i, j)
     d_yi(A, ϕ) = _d_yi(A, ϕ, _dy, i, j)
     d_ya(A, ϕ) = _d_ya(A, ϕ, _dy, i, j)
+    av_xa(A, ϕ) = _av_xa(A, ϕ, i, j)
+    av_ya(A, ϕ) = _av_ya(A, ϕ, i, j)
     av_xa(A) = _av_xa(A, i, j)
     av_ya(A) = _av_ya(A, i, j)
     harm_xa(A) = _av_xa(A, i, j)
@@ -80,7 +82,7 @@ end
             Rx[i, j] =
                 R_Vx = (
                     -d_xa(P, ϕ.center) + d_xa(τxx, ϕ.center) + d_yi(τxy, ϕ.vertex) -
-                    av_xa(ρgx)
+                    av_xa(ρgx, ϕ.center)
                 )
             Vx[i + 1, j + 1] += R_Vx * ηdτ / av_xa(ητ)
         else
@@ -94,7 +96,7 @@ end
             Ry[i, j] =
                 R_Vy =
                     -d_ya(P, ϕ.center) + d_ya(τyy, ϕ.center) + d_xi(τxy, ϕ.vertex) -
-                    av_ya(ρgy)
+                    av_ya(ρgy, ϕ.center)
             Vy[i + 1, j + 1] += R_Vy * ηdτ / av_ya(ητ)
         else
             Ry[i, j] = zero(T)

src/variational_stokes/mask.jl

@@ -45,12 +45,19 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`.
 - `phase_ratios`: The ratios of different phases present.
 - `air_phase`: The phase representing air.
 """
-function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, air_phase)
+function update_rock_ratio!(ϕ::JustRelax.RockRatio, phase_ratios, ratio_vel::NTuple{N}, air_phase) where N
     nvi = size_v(ϕ)
     @parallel (@idx nvi) update_rock_ratio_cv!(
         ϕ, phase_ratios.center, phase_ratios.vertex, air_phase
     )
-    @parallel (@idx nvi) update_rock_ratio_vel!(ϕ)
+    # @parallel (@idx nvi) update_rock_ratio_vel!(ϕ)
+
+    @parallel (@idx size(ϕ.Vx)) update_rock_ratio_vel!(ϕ.Vx, ratio_vel[1], air_phase)
+    @parallel (@idx size(ϕ.Vy)) update_rock_ratio_vel!(ϕ.Vy, ratio_vel[2], air_phase)
+    if N === 3
+        @parallel (@idx size(ϕ.Vz)) update_rock_ratio_vel!(ϕ.Vz, ratio_vel[3], air_phase)
+    end
+
     return nothing
 end
 
@@ -72,98 +79,11 @@ end
     return nothing
 end
 
-@parallel_indices (I...) function update_rock_ratio_vel!(
-    ϕ::JustRelax.RockRatio{T,N}
-) where {T,N}
-    # 2D
-    @inline av_x(A::AbstractArray{T,2}) where {T} = _av_xa(A, I...)
-    @inline av_y(A::AbstractArray{T,2}) where {T} = _av_ya(A, I...)
-    # 3D
-    @inline av_x(A::AbstractArray{T,3}) where {T} = _av_yz(A, I...)
-    @inline av_y(A::AbstractArray{T,3}) where {T} = _av_xz(A, I...)
-    @inline av_z(A::AbstractArray{T,3}) where {T} = _av_xy(A, I...)
-
-    all(I .≤ size(ϕ.Vx)) && (ϕ.Vx[I...] = av_y(ϕ.vertex))
-    all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex))
-    if N === 3 # control flow here, so that the branch can be removed by the compiler in the 2D case 
-        all(I .≤ size(ϕ.Vy)) && (ϕ.Vy[I...] = av_x(ϕ.vertex))
-    end
+@parallel_indices (I...) function update_rock_ratio_vel!(ϕ, ratio, air_phase)
+    ϕ[I...] = Float64(Float16(compute_rock_ratio(ratio, air_phase, I...)))
     return nothing
 end
 
-"""
-    isvalid_c(ϕ::JustRelax.RockRatio, inds...)
-
-Check if  `ϕ.center[inds...]` is a not a nullspace.
-
-# Arguments
-- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
-- `inds`: Cartesian indices to check.
-"""
-function isvalid_c(ϕ::JustRelax.RockRatio, i, j)
-    vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0)
-    vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0)
-    v = vx * vy
-    return v * (ϕ.center[i, j] > 0)
-end
-
-"""
-    isvalid_v(ϕ::JustRelax.RockRatio, inds...)
-
-Check if  `ϕ.vertex[inds...]` is a not a nullspace.
-
-# Arguments
-- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
-- `inds`: Cartesian indices to check.
-"""
-function isvalid_v(ϕ::JustRelax.RockRatio, i, j)
-    nx, ny = size(ϕ.Vx)
-    j_bot = max(j - 1, 1)
-    j0 = min(j, ny)
-    vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0)
-
-    nx, ny = size(ϕ.Vy)
-    i_left = max(i - 1, 1)
-    i0 = min(i, nx)
-    vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0)
-    v = vx * vy
-    return v * (ϕ.vertex[i, j] > 0)
-end
-
-"""
-    isvalid_vx(ϕ::JustRelax.RockRatio, inds...)
-
-Check if  `ϕ.Vx[inds...]` is a not a nullspace.
-
-# Arguments
-- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
-- `inds`: Cartesian indices to check.
-"""
-function isvalid_vx(ϕ::JustRelax.RockRatio, i, j)
-    c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0)
-    v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0)
-    cv = c * v
-    # return cv * (ϕ.Vx[i, j] > 0)
-    return ϕ.Vx[i, j] > 0
-end
-
-"""
-    isvalid_vy(ϕ::JustRelax.RockRatio, inds...)
-
-Check if  `ϕ.Vy[inds...]` is a not a nullspace.
-
-# Arguments
-- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
-- `inds`: Cartesian indices to check.
-"""
-function isvalid_vy(ϕ::JustRelax.RockRatio, i, j)
-    c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0)
-    v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0)
-    cv = c * v
-    # return cv * (ϕ.Vy[i, j] > 0)
-    return ϕ.Vy[i, j] > 0
-end
-
 
 
 # """
@@ -175,8 +95,11 @@ end
 # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
 # - `inds`: Cartesian indices to check.
 # """
-# function isvalid_c(ϕ::JustRelax.RockRatio, i, j)
-#     return (ϕ.center[i, j] > 0)
+# Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j)
+#     vx = (ϕ.Vx[i, j] > 0) * (ϕ.Vx[i + 1, j] > 0)
+#     vy = (ϕ.Vy[i, j] > 0) * (ϕ.Vy[i, j + 1] > 0)
+#     v = vx * vy
+#     return v * (ϕ.center[i, j] > 0)
 # end
 
 # """
@@ -188,8 +111,18 @@ end
 # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
 # - `inds`: Cartesian indices to check.
 # """
-# function isvalid_v(ϕ::JustRelax.RockRatio, i, j)
-#     return (ϕ.vertex[i, j] > 0)
+# Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j)
+#     nx, ny = size(ϕ.Vx)
+#     j_bot = max(j - 1, 1)
+#     j0 = min(j, ny)
+#     vx = (ϕ.Vx[i, j0] > 0) * (ϕ.Vx[i, j_bot] > 0)
+
+#     nx, ny = size(ϕ.Vy)
+#     i_left = max(i - 1, 1)
+#     i0 = min(i, nx)
+#     vy = (ϕ.Vy[i0, j] > 0) * (ϕ.Vy[i_left, j] > 0)
+#     v = vx * vy
+#     return v * (ϕ.vertex[i, j] > 0)
 # end
 
 # """
@@ -201,11 +134,12 @@ end
 # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
 # - `inds`: Cartesian indices to check.
 # """
-# function isvalid_vx(ϕ::JustRelax.RockRatio, i, j)
-#     c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0)
-#     v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0)
-#     cv = c || v
-#     return cv || (ϕ.Vx[i, j] > 0)
+# Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j)
+#     # c = (ϕ.center[i, j] > 0) * (ϕ.center[i - 1, j] > 0)
+#     # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i, j + 1] > 0)
+#     # cv = c * v
+#     # return cv * (ϕ.Vx[i, j] > 0)
+#     return (ϕ.Vx[i, j] > 0)
 # end
 
 # """
@@ -217,9 +151,72 @@ end
 # - `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
 # - `inds`: Cartesian indices to check.
 # """
-# function isvalid_vy(ϕ::JustRelax.RockRatio, i, j)
-#     c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0)
-#     v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0)
-#     cv = c || v
-#     return cv || (ϕ.Vy[i, j] > 0)
+# Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j)
+#     # c = (ϕ.center[i, j] > 0) * (ϕ.center[i, j - 1] > 0)
+#     # v = (ϕ.vertex[i, j] > 0) * (ϕ.vertex[i + 1, j] > 0)
+#     # cv = c * v
+#     # return cv * (ϕ.Vy[i, j] > 0)
+#     return (ϕ.Vy[i, j] > 0)
 # end
+
+#######
+
+"""
+    isvalid_c(ϕ::JustRelax.RockRatio, inds...)
+
+Check if  `ϕ.center[inds...]` is a not a nullspace.
+
+# Arguments
+- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
+- `inds`: Cartesian indices to check.
+"""
+Base.@propagate_inbounds @inline function isvalid_c(ϕ::JustRelax.RockRatio, i, j)
+    (ϕ.center[i, j] > 0)
+end
+
+"""
+    isvalid_v(ϕ::JustRelax.RockRatio, inds...)
+
+Check if  `ϕ.vertex[inds...]` is a not a nullspace.
+
+# Arguments
+- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
+- `inds`: Cartesian indices to check.
+"""
+Base.@propagate_inbounds @inline function isvalid_v(ϕ::JustRelax.RockRatio, i, j)
+    (ϕ.vertex[i, j] > 0)
+end
+
+"""
+    isvalid_vx(ϕ::JustRelax.RockRatio, inds...)
+
+Check if  `ϕ.Vx[inds...]` is a not a nullspace.
+
+# Arguments
+- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
+- `inds`: Cartesian indices to check.
+"""
+Base.@propagate_inbounds @inline function isvalid_vx(ϕ::JustRelax.RockRatio, i, j)
+    c = (ϕ.center[i, j] > 0) || (ϕ.center[i - 1, j] > 0)
+    v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i, j + 1] > 0)
+    cv = c || v
+    return cv || (ϕ.Vx[i, j] > 0)
+end
+
+"""
+    isvalid_vy(ϕ::JustRelax.RockRatio, inds...)
+
+Check if  `ϕ.Vy[inds...]` is a not a nullspace.
+
+# Arguments
+- `ϕ::JustRelax.RockRatio`: The `RockRatio` object to check against.
+- `inds`: Cartesian indices to check.
+"""
+Base.@propagate_inbounds @inline function isvalid_vy(ϕ::JustRelax.RockRatio, i, j)
+    c = (ϕ.center[i, j] > 0) || (ϕ.center[i, j - 1] > 0)
+    v = (ϕ.vertex[i, j] > 0) || (ϕ.vertex[i + 1, j] > 0)
+    cv = c || v
+    return cv || (ϕ.Vy[i, j] > 0)
+end
+
+Base.@propagate_inbounds @inline isvalid(ϕ, I::Vararg{Integer, N}) where N = ϕ[I...] > 0