Merge branch 'main' into pa-mpi

PTsolvers · Nov 12, 2024 · 8e5a7a2 · 8e5a7a2
2 parents 6b2cb29 + c8106d1
commit 8e5a7a2
Show file tree

Hide file tree

Showing 8 changed files with 96 additions and 107 deletions.
diff --git a/Project.toml b/Project.toml
@@ -38,7 +38,7 @@ GeoParams = "0.6.4"
 HDF5 = "0.17.1"
 ImplicitGlobalGrid = "0.15"
 JLD2 = "0.4, 0.5"
-JustPIC = "0.5"
+JustPIC = "0.5.3"
 MPI = "0.20"
 MuladdMacro = "0.2"
 ParallelStencil = "0.13.6"

diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl
@@ -1,39 +1,42 @@
+using CUDA
 using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO
 
-const backend_JR = CPUBackend
-
+const backend_JR = CUDABackend
 using Printf, GeoParams, GLMakie, CellArrays
 
 using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 3)
+@init_parallel_stencil(CUDA, Float64, 3)
 
 using JustPIC, JustPIC._3D
-const backend = JustPIC.CPUBackend
+const backend = CUDABackend
 
 # HELPER FUNCTIONS ---------------------------------------------------------------
 solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η))
 
 # Initialize phases on the particles
-function init_phases!(phase_ratios, xci, xvi, radius)
-    ni      = size(phase_ratios.center)
+function init_phases!(phases, particles, radius)
+    ni      = size(phases)
     origin  = 0.5, 0.5, 0.5
 
-    @parallel_indices (i, j, k) function init_phases!(phases, xc, yc, zc, o_x, o_y, o_z)
-        x, y, z = xc[i], yc[j], zc[k]
-        if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius
-            @index phases[1, i, j, k] = 1.0
-            @index phases[2, i, j, k] = 0.0
+    @parallel_indices (I...) function init_phases!(phases, index, xc, yc, zc, o_x, o_y, o_z)
+
+        for ip in cellaxes(xc)
+            (@index index[ip, I...]) || continue
 
-        else
-            @index phases[1, i, j, k] = 0.0
-            @index phases[2, i, j, k] = 1.0
+            x = @index xc[ip, I...]
+            y = @index yc[ip, I...]
+            z = @index zc[ip, I...]
 
+            if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius^2
+                @index phases[ip, I...] = 1.0
+            else
+                @index phases[ip, I...] = 2.0
+            end
         end
         return nothing
     end
 
-    @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin...)
-    @parallel (@idx ni .+ 1) init_phases!(phase_ratios.vertex, xvi..., origin...)
+    @parallel (@idx ni) init_phases!(phases,  particles.index, particles.coords...,origin...)
     return nothing
 end
 
@@ -91,9 +94,15 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs")
     )
 
     # Initialize phase ratios -------------------------------
-    radius       = 0.1
+    nxcell, max_xcell, min_xcell = 125, 150, 75
+    particles                    = init_particles(backend, nxcell, max_xcell, min_xcell, xvi, di, ni)
+    radius                       = 0.1
+    phase_ratios                 = PhaseRatios(backend, length(rheology), ni)
+    pPhases,                     = init_cell_arrays(particles, Val(1))
+    # Assign particles phases anomaly
+    init_phases!(pPhases, particles, radius)
     phase_ratios = PhaseRatios(backend, length(rheology), ni)
-    init_phases!(phase_ratios, xci, xvi, radius)
+    update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases)
 
     # STOKES ---------------------------------------------
     # Allocate arrays needed for every Stokes problem
@@ -167,9 +176,9 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs")
 
         velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)
         data_v = (;
-            τII = Array(stokes.τ.II),
-            εII = Array(stokes.ε.II),
-            εII = Array(stokes.ε_pl.II),
+            τII    = Array(stokes.τ.II),
+            εII    = Array(stokes.ε.II),
+            εII_pl = Array(stokes.ε_pl.II),
         )
         data_c = (;
             P   = Array(stokes.P),
@@ -214,12 +223,12 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs")
     return nothing
 end
 
-n            = 128
+n            = 64
 nx = ny = nz = n
 figdir       = "ShearBand3D"
 igg          = if !(JustRelax.MPI.Initialized())
     IGG(init_global_grid(nx, ny, nz; init_MPI = true)...)
 else
     igg
 end
-main(igg; figdir = figdir, nx = nx, ny = ny, nz = nz);
+main(igg; figdir = figdir, nx = nx, ny = ny, nz = nz);
diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl
@@ -4,9 +4,11 @@ include("free_surface.jl")
 include("no_slip.jl")
 include("pure_shear.jl")
 
-@inline bc_index(x::NTuple{2,T}) where {T} = mapreduce(xi -> max(size(xi)...), max, x)
-@inline bc_index(x::T) where {T<:AbstractArray} = max(size(x)...)
+@inline bc_index(x::T) where {T<:AbstractArray{_T,2} where {_T}} = max(size(x)...)
+@inline bc_index(x::T) where {T<:AbstractArray{_T,3} where {_T}} =
+    max(size(x)...), max(size(x)...)
 
+@inline bc_index(x::NTuple{2,T}) where {T} = mapreduce(xi -> max(size(xi)...), max, x)
 @inline function bc_index(x::NTuple{3,T}) where {T}
     n = mapreduce(xi -> max(size(xi)...), max, x)
     return n, n

diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl
@@ -598,6 +598,9 @@ function _solve!(
                 rheology,
                 phase_ratios.center,
                 phase_ratios.vertex,
+                phase_ratios.xy,
+                phase_ratios.yz,
+                phase_ratios.xz,
             )
             update_halo!(stokes.τ.xy)
 

diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl
@@ -491,6 +491,9 @@ function _solve!(
                 rheology,
                 phase_ratios.center,
                 phase_ratios.vertex,
+                phase_ratios.xy,
+                phase_ratios.yz,
+                phase_ratios.xz,
             )
             update_halo!(stokes.τ.yz)
             update_halo!(stokes.τ.xz)

diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl
@@ -587,6 +587,9 @@ end
     rheology,
     phase_center,
     phase_vertex,
+    phase_xy,
+    phase_yz,
+    phase_xz,
 )
     τyzv, τxzv, τxyv = τshear_v
     τyzv_old, τxzv_old, τxyv_old = τshear_ov
@@ -622,7 +625,7 @@ end
         τxyv_old_ij = av_clamped_yz_z(τxyv_old, Ic...)
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds phase_yz[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -632,24 +635,12 @@ end
         dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0)
 
         # stress increments @ vertex
-        dτxxv =
-            (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) *
-            dτ_rv
-        dτyyv =
-            (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) *
-            dτ_rv
-        dτzzv =
-            (-(τzzv_ij - τzzv_old_ij) * ηv_ij * _Gvdt - τzzv_ij + 2.0 * ηv_ij * εzzv_ij) *
-            dτ_rv
-        dτyzv =
-            (-(τyzv_ij - τyzv_old_ij) * ηv_ij * _Gvdt - τyzv_ij + 2.0 * ηv_ij * εyzv_ij) *
-            dτ_rv
-        dτxzv =
-            (-(τxzv_ij - τxzv_old_ij) * ηv_ij * _Gvdt - τxzv_ij + 2.0 * ηv_ij * εxzv_ij) *
-            dτ_rv
-        dτxyv =
-            (-(τxyv_ij - τxyv_old_ij) * ηv_ij * _Gvdt - τxyv_ij + 2.0 * ηv_ij * εxyv_ij) *
-            dτ_rv
+        dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv)
+        dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv)
+        dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv)
+        dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv)
+        dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv)
+        dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv)
 
         dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv
         τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij
@@ -697,7 +688,7 @@ end
         τxyv_old_ij = av_clamped_xz_z(τxyv_old, Ic...)
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds phase_xz[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -707,24 +698,12 @@ end
         dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0)
 
         # stress increments @ vertex
-        dτxxv =
-            (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) *
-            dτ_rv
-        dτyyv =
-            (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) *
-            dτ_rv
-        dτzzv =
-            (-(τzzv_ij - τzzv_old_ij) * ηv_ij * _Gvdt - τzzv_ij + 2.0 * ηv_ij * εzzv_ij) *
-            dτ_rv
-        dτyzv =
-            (-(τyzv_ij - τyzv_old_ij) * ηv_ij * _Gvdt - τyzv_ij + 2.0 * ηv_ij * εyzv_ij) *
-            dτ_rv
-        dτxzv =
-            (-(τxzv_ij - τxzv_old_ij) * ηv_ij * _Gvdt - τxzv_ij + 2.0 * ηv_ij * εxzv_ij) *
-            dτ_rv
-        dτxyv =
-            (-(τxyv_ij - τxyv_old_ij) * ηv_ij * _Gvdt - τxyv_ij + 2.0 * ηv_ij * εxyv_ij) *
-            dτ_rv
+        dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv)
+        dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv)
+        dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv)
+        dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv)
+        dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv)
+        dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv)
 
         dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv
         τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij
@@ -774,7 +753,7 @@ end
         τxyv_old_ij = τxyv_old[I...]
 
         # vertex parameters
-        phase = @inbounds phase_vertex[I...]
+        phase = @inbounds phase_xy[I...]
         is_pl, Cv, sinϕv, cosϕv, sinψv, η_regv = plastic_params_phase(
             rheology, EIIv_ij, phase
         )
@@ -784,25 +763,12 @@ end
         dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0)
 
         # stress increments @ vertex
-        dτxxv =
-            (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) *
-            dτ_rv
-        dτyyv =
-            (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) *
-            dτ_rv
-        dτzzv =
-            (-(τzzv_ij - τzzv_old_ij) * ηv_ij * _Gvdt - τzzv_ij + 2.0 * ηv_ij * εzzv_ij) *
-            dτ_rv
-        dτyzv =
-            (-(τyzv_ij - τyzv_old_ij) * ηv_ij * _Gvdt - τyzv_ij + 2.0 * ηv_ij * εyzv_ij) *
-            dτ_rv
-        dτxzv =
-            (-(τxzv_ij - τxzv_old_ij) * ηv_ij * _Gvdt - τxzv_ij + 2.0 * ηv_ij * εxzv_ij) *
-            dτ_rv
-        dτxyv =
-            (-(τxyv_ij - τxyv_old_ij) * ηv_ij * _Gvdt - τxyv_ij + 2.0 * ηv_ij * εxyv_ij) *
-            dτ_rv
-
+        dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv)
+        dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv)
+        dτzzv = compute_stress_increment(τzzv_ij, τzzv_old_ij, ηv_ij, εzzv_ij, _Gvdt, dτ_rv)
+        dτyzv = compute_stress_increment(τyzv_ij, τyzv_old_ij, ηv_ij, εyzv_ij, _Gvdt, dτ_rv)
+        dτxzv = compute_stress_increment(τxzv_ij, τxzv_old_ij, ηv_ij, εxzv_ij, _Gvdt, dτ_rv)
+        dτxyv = compute_stress_increment(τxyv_ij, τxyv_old_ij, ηv_ij, εxyv_ij, _Gvdt, dτ_rv)
         dτijv = dτxxv, dτyyv, dτzzv, dτyzv, dτxzv, dτxyv
         τijv = τxxv_ij, τyyv_ij, τzzv_ij, τyzv_ij, τxzv_ij, τxyv_ij
         τIIv_ij = second_invariant(τijv .+ dτijv)
@@ -895,6 +861,9 @@ end
     rheology,
     phase_center,
     phase_vertex,
+    phase_xy,
+    phase_yz,
+    phase_xz,
 )
     τxyv = τshear_v[1]
     τxyv_old = τshear_ov[1]
@@ -921,15 +890,11 @@ end
     dτ_rv = inv(θ_dτ + ηv_ij * _Gvdt + 1.0)
 
     # stress increments @ vertex
-    dτxxv =
-        (-(τxxv_ij - τxxv_old_ij) * ηv_ij * _Gvdt - τxxv_ij + 2.0 * ηv_ij * εxxv_ij) * dτ_rv
-    dτyyv =
-        (-(τyyv_ij - τyyv_old_ij) * ηv_ij * _Gvdt - τyyv_ij + 2.0 * ηv_ij * εyyv_ij) * dτ_rv
-    dτxyv =
-        (
-            -(τxyv[I...] - τxyv_old[I...]) * ηv_ij * _Gvdt - τxyv[I...] +
-            2.0 * ηv_ij * ε[3][I...]
-        ) * dτ_rv
+    dτxxv = compute_stress_increment(τxxv_ij, τxxv_old_ij, ηv_ij, εxxv_ij, _Gvdt, dτ_rv)
+    dτyyv = compute_stress_increment(τyyv_ij, τyyv_old_ij, ηv_ij, εyyv_ij, _Gvdt, dτ_rv)
+    dτxyv = compute_stress_increment(
+        τxyv[I...], τxyv_old[I...], ηv_ij, ε[3][I...], _Gvdt, dτ_rv
+    )
     τIIv_ij = √(0.5 * ((τxxv_ij + dτxxv)^2 + (τyyv_ij + dτyyv)^2) + (τxyv[I...] + dτxyv)^2)
 
     # yield function @ center
@@ -964,7 +929,8 @@ end
         εij_ve = @. εij + 0.5 * τij_o * _Gdt
         εII_ve = GeoParams.second_invariant(εij_ve)
         # stress increments @ center
-        dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r
+        # dτij = @. (-(τij - τij_o) * ηij * _Gdt - τij .+ 2.0 * ηij * εij) * dτ_r
+        dτij = compute_stress_increment(τij, τij_o, ηij, εij, _Gdt, dτ_r)
         τII_ij = GeoParams.second_invariant(dτij .+ τij)
         # yield function @ center
         F = τII_ij - C - Pr[I...] * sinϕ