updates and fixes

PTsolvers · Nov 1, 2024 · 6744030 · 6744030
1 parent ef7e162
commit 6744030
Show file tree

Hide file tree

Showing 8 changed files with 56 additions and 10 deletions.
diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl
@@ -51,7 +51,8 @@ Base.@propagate_inbounds @inline _d_yi(
     A::AbstractArray, _dy, I::Vararg{Integer,N}
 ) where {N} = (-right(A, I...) + next(A, I...)) * _dy
 
-Base.@propagate_inbounds @inline div(Ax, Ay, _dx, _dy, I::Vararg{Integer, 2}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...)
+Base.@propagate_inbounds @inline div(Ax, Ay, _dx, _dy, I::Vararg{Integer,2}) =
+    _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...)
 
 # averages
 Base.@propagate_inbounds @inline _av(A::T, i, j) where {T<:T2} =
@@ -91,7 +92,8 @@ const T3 = AbstractArray{T,3} where {T}
 @inline function _d_zi(A::T, i, j, k, _dz) where {T<:T3}
     return (-A[i + 1, j + 1, k] + next(A, i, j, k)) * _dz
 end
-Base.@propagate_inbounds @inline div(Ax, Ay, Az, _dx, _dy, _dz, I::Vararg{Integer, 3}) = _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + _d_zi(Az, _dz, I...)
+Base.@propagate_inbounds @inline div(Ax, Ay, Az, _dx, _dy, _dz, I::Vararg{Integer,3}) =
+    _d_xi(Ax, _dx, I...) + _d_yi(Ay, _dy, I...) + _d_zi(Az, _dz, I...)
 
 # averages
 Base.@propagate_inbounds @inline _av(A::T, i, j, k) where {T<:T3} =

diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -408,4 +408,14 @@ function JR2D.update_rock_ratio!(
     return nothing
 end
 
+function JR2D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N},
+    phase_ratios,
+    ratio_vel::NTuple{N},
+    air_phase,
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -425,4 +425,14 @@ function JR3D.update_rock_ratio!(
     return nothing
 end
 
+function JR3D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{ROCArray{T,nD,D},N},
+    phase_ratios,
+    ratio_vel::NTuple{N},
+    air_phase,
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -185,6 +185,12 @@ function JR2D.compute_viscosity!(
     return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff)
 end
 
+function JR2D.compute_viscosity!(
+    ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, air_phase, cutoff
+)
+    return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, air_phase, cutoff)
+end
+
 function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff)
     return compute_viscosity!(η, ν, εII, args, rheology, cutoff)
 end
@@ -194,9 +200,9 @@ function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cuto
 end
 
 function compute_viscosity!(
-    ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff
+    ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, air_phase, cutoff
 )
-    return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff)
+    return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, air_phase, cutoff)
 end
 
 function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff)
@@ -394,4 +400,11 @@ function JR2D.update_rock_ratio!(
     return nothing
 end
 
+function JR2D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -421,4 +421,11 @@ function JR3D.update_rock_ratio!(
     return nothing
 end
 
+function JR3D.update_rock_ratio!(
+    ϕ::JustRelax.RockRatio{CuArray{T,nD,D},N}, phase_ratios, ratio_vel::NTuple{N}, air_phase
+) where {T,nD,N,D}
+    update_rock_ratio!(ϕ, phase_ratios, ratio_vel, air_phase)
+    return nothing
+end
+
 end
diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl
@@ -377,7 +377,7 @@ end
 #     end
 # end
 
-function correct_phase_ratio(air_phase, ratio::SVector{N, T}) where {N, T}
+function correct_phase_ratio(air_phase, ratio::SVector{N,T}) where {N,T}
     if iszero(air_phase)
         return ratio
     elseif ratio[air_phase] ≈ 1

diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl
@@ -622,7 +622,8 @@ end
         τxyv_old_ij = av_clamped_yz_z(τxyv_old, Ic...)
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1] + 1, I[2], I[3]]) * 0.5
+        # phase = @inbounds phase_vertex[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -697,7 +698,8 @@ end
         τxyv_old_ij = av_clamped_xz_z(τxyv_old, Ic...)
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1], I[2] + 1, I[3]]) * 0.5
+        # phase = @inbounds phase_vertex[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -774,7 +776,8 @@ end
         τxyv_old_ij = τxyv_old[I...]
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds (phase_vertex[I...] + phase_vertex[I[1], I[2], I[3] + 1]) * 0.5
+        # phase = @inbounds phase_vertex[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -969,7 +972,7 @@ end
         # yield function @ center
         F = τII_ij - C - Pr[I...] * sinϕ
 
-        if is_pl && !iszero(τII_ij) && F >0
+        if is_pl && !iszero(τII_ij) && F > 0
             # stress correction @ center
             λ[I...] =
                 (1.0 - relλ) * λ[I...] +

diff --git a/src/variational_stokes/Stokes2D.jl b/src/variational_stokes/Stokes2D.jl
@@ -177,7 +177,8 @@ function _solve_VS!(
                 length(stokes.R.Rx),
                 norm_mpi(@views stokes.R.Ry[2:(end - 1), 2:(end - 1)]) /
                 length(stokes.R.Ry),
-                norm_mpi(@views stokes.R.RP[ϕ.center .> 0]) / length(@views stokes.R.RP[ϕ.center .> 0]),
+                norm_mpi(@views stokes.R.RP[ϕ.center .> 0]) /
+                length(@views stokes.R.RP[ϕ.center .> 0]),
             )
             push!(norm_Rx, errs[1])
             push!(norm_Ry, errs[2])