Add JustPIC compatibility (#237)

* update src * add JP dep to ext * update miniapps and `Project.toml` * update tests * fix test * format * oops * fix miniapps; address requested changes; ext still buggy due to JP type export * fix ext * format * comment `backend(x::JustPIC.PhaseRatios)` * fix default backend * adding version compat for CUDA with julia 1.9 * fix CI CUDA version * revert runtest changes; rm CI for 1.9; add `pre` to buildkite * typo * add `1.11` to pipeline rather than `pre` * requested changes Co-authored-by: Albert de Montserrat <[email protected]> * format --------- Co-authored-by: Albert de Montserrat <[email protected]>
PTsolvers · Sep 24, 2024 · 062f4bf · 062f4bf
1 parent 0e80412
commit 062f4bf
Show file tree

Hide file tree

Showing 80 changed files with 533 additions and 1,034 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -3,8 +3,8 @@ steps:
     matrix:
       setup:
         version:
-          - "1.9"
           - "1.10"
+          - "1.11"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"
@@ -30,6 +30,7 @@ steps:
       setup:
         version:
           - "1.10"
+          - "1.11"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -14,7 +14,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.9'
           - '1.10'
           - 'pre'
           # - 'nightly'

diff --git a/Project.toml b/Project.toml
@@ -10,6 +10,7 @@ GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
@@ -31,12 +32,13 @@ JustRelaxCUDAExt = "CUDA"
 [compat]
 AMDGPU = "0.8, 0.9, 1"
 Adapt = "4"
-CUDA = "~5.3.5"
+CUDA = "5"
 CellArrays = "0.2"
 GeoParams = "0.5, 0.6"
 HDF5 = "0.17.1"
 ImplicitGlobalGrid = "0.15.0"
 JLD2 = "0.4, 0.5"
+JustPIC = "0.4.2"
 MPI = "0.20"
 MuladdMacro = "0.2"
 ParallelStencil = "0.12, 0.13"

diff --git a/docs/src/man/Blankenbach.md b/docs/src/man/Blankenbach.md
@@ -94,11 +94,11 @@ function init_phases!(phases, particles)
     ni = size(phases)
 
     @parallel_indices (i, j) function init_phases!(phases, index)
-        @inbounds for ip in JustRelax.cellaxes(phases)
+        @inbounds for ip in cellaxes(phases)
             # quick escape if the ip-th element of the [i,j]-th cell is empty
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+            @index(index[ip, i, j]) == 0 && continue
             # all particles have phase number = 1.0
-            JustRelax.@cell phases[ip, i, j] = 1.0
+            @index phases[ip, i, j] = 1.0
         end
         return nothing
     end
@@ -120,8 +120,8 @@ map!(x -> isnan(x) ? NaN : 1.0, pPhase.data, particles.index.data)
 
 and finally we need the phase ratios at the cell centers:
 ```julia
-phase_ratios = PhaseRatio(backend_JR, ni, length(rheology))
-phase_ratios_center(phase_ratios, particles, grid, pPhases)
+phase_ratios = PhaseRatios(backend, length(rheology), ni)
+phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 ```
 
 ### Stokes and heat diffusion arrays
@@ -312,7 +312,7 @@ move_particles!(particles, xvi, particle_args)
 # check if we need to inject particles
 inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi)
 # update phase ratios
-phase_ratios_center(phase_ratios, particles, grid, pPhases)
+phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 ```
 5.  Interpolate `T` back to the grid
 ```julia

diff --git a/docs/src/man/ShearBands.md b/docs/src/man/ShearBands.md
@@ -92,12 +92,12 @@ function init_phases!(phase_ratios, xci, radius)
     @parallel_indices (i, j) function init_phases!(phases, xc, yc, o_x, o_y, radius)
         x, y = xc[i], yc[j]
         if ((x-o_x)^2 + (y-o_y)^2) > radius^2
-            JustRelax.@cell phases[1, i, j] = 1.0
-            JustRelax.@cell phases[2, i, j] = 0.0
+            @index phases[1, i, j] = 1.0
+            @index phases[2, i, j] = 0.0
 
         else
-            JustRelax.@cell phases[1, i, j] = 0.0
-            JustRelax.@cell phases[2, i, j] = 1.0
+            @index phases[1, i, j] = 0.0
+            @index phases[2, i, j] = 1.0
         end
         return nothing
     end
@@ -109,7 +109,7 @@ end
 
 and finally we need the phase ratios at the cell centers:
 ```julia
-phase_ratios = PhaseRatio(backend_JR, ni, length(rheology))
+phase_ratios = PhaseRatios(backend, length(rheology), ni)
 init_phases!(phase_ratios, xci, radius)
 ```
 

diff --git a/docs/src/man/subduction2D/subduction2D.md b/docs/src/man/subduction2D/subduction2D.md
@@ -78,9 +78,9 @@ particle_args       = (pT, pPhases)
 Now we assign the material phases from the arrays we computed with help of [GeophysicalModelGenerator.jl](https://github.com/JuliaGeodynamics/GeophysicalModelGenerator.jl)
 ```julia
 phases_device    = PTArray(backend)(phases_GMG)
-phase_ratios     = PhaseRatio(backend, ni, length(rheology))
+phase_ratios     = PhaseRatios(backend, length(rheology), ni);
 init_phases!(pPhases, phases_device, particles, xvi)
-phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 ```
 
 ## Temperature profile
@@ -227,7 +227,7 @@ move_particles!(particles, xvi, particle_args)
 # check if we need to inject particles
 inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi)
 # update phase ratios
-phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 ```
 
 6. **Optional:** Save data as VTK to visualize it later with [ParaView](https://www.paraview.org/)

diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl
@@ -83,9 +83,9 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
     )
     # temperature
     pPhases,            = init_cell_arrays(particles, Val(1))
-    phase_ratios        = PhaseRatio(backend_JR, ni, length(rheology))
+    phase_ratios = PhaseRatios(backend, length(rheology), ni)
     init_phases!(pPhases, particles)
-    phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+    phase_ratios_center!(phase_ratios, particles, xci, pPhases)
     # ----------------------------------------------------
 
     # STOKES ---------------------------------------------

diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl
@@ -84,9 +84,9 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
     # temperature
     pT, pT0, pPhases    = init_cell_arrays(particles, Val(3))
     particle_args       = (pT, pT0, pPhases)
-    phase_ratios        = PhaseRatio(backend_JR, ni, length(rheology))
+    phase_ratios = PhaseRatios(backend, length(rheology), ni)
     init_phases!(pPhases, particles)
-    phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+    phase_ratios_center!(phase_ratios, particles, xci, pPhases)
     # ----------------------------------------------------
 
     # STOKES ---------------------------------------------
@@ -257,7 +257,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
         # check if we need to inject particles
         inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi)
         # update phase ratios
-        phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+        phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 
         # Nusselt number, Nu = H/ΔT/L ∫ ∂T/∂z dx ----
         Nu_it   =   (ly / (1000.0*lx)) *

diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl
@@ -76,9 +76,9 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
     # temperature
     pT, pT0, pPhases    = init_cell_arrays(particles, Val(3))
     particle_args       = (pT, pT0, pPhases)
-    phase_ratios        = PhaseRatio(backend_JR, ni, length(rheology))
+    phase_ratios = PhaseRatios(backend, length(rheology), ni)
     init_phases!(pPhases, particles)
-    phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+    phase_ratios_center!(phase_ratios, particles, xci, pPhases)
     # ----------------------------------------------------
 
     # STOKES ---------------------------------------------
@@ -247,7 +247,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal
         # check if we need to inject particles
         inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi)
         # update phase ratios
-        phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+        phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 
         # Nusselt number, Nu = ∫ ∂T/∂z dx ----
         Nu_it   =   sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1])

diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl
@@ -13,22 +13,22 @@ function init_rheologies()
             CompositeRheology = CompositeRheology((LinearViscous(; η=1.0e23),)),
             RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
             Gravity           = ConstantGravity(; g=10.0),
-        ),        
+        ),
     )
 end
 
 function init_phases!(phases, particles)
     ni = size(phases)
 
     @parallel_indices (i, j) function init_phases!(phases, index)
-        @inbounds for ip in JustRelax.cellaxes(phases)
+        @inbounds for ip in cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
-            JustRelax.@cell phases[ip, i, j] = 1.0
+            @index(index[ip, i, j]) == 0 && continue
+            @index phases[ip, i, j] = 1.0
         end
         return nothing
     end
 
     @parallel (@idx ni) init_phases!(phases, particles.index)
     return nothing
-end
+end
diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl
@@ -11,18 +11,18 @@ function init_rheologies()
             CompositeRheology = CompositeRheology((LinearViscous(; η=1),)),
             RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
             Gravity           = ConstantGravity(; g = 1e4),
-        ),        
+        ),
     )
 end
 
 function init_phases!(phases, particles)
     ni = size(phases)
 
     @parallel_indices (i, j) function init_phases!(phases, index)
-        @inbounds for ip in JustRelax.cellaxes(phases)
+        @inbounds for ip in cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
-            JustRelax.@cell phases[ip, i, j] = 1.0
+            @index(index[ip, i, j]) == 0 && continue
+            @index phases[ip, i, j] = 1.0
         end
         return nothing
     end

diff --git a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl
@@ -3,7 +3,7 @@ using ParallelStencil
 
 using Printf, LinearAlgebra, GeoParams, CellArrays
 using JustRelax, JustRelax.JustRelax2D
-import JustRelax.@cell
+
 const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
 
 using GLMakie
@@ -23,18 +23,18 @@ function init_phases!(phases, particles)
     ni = size(phases)
 
     @parallel_indices (i, j) function init_phases!(phases, px, py, index)
-        @inbounds for ip in JustRelax.cellaxes(phases)
+        @inbounds for ip in cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+            @index(index[ip, i, j]) == 0 && continue
 
-            x = JustRelax.@cell px[ip, i, j]
-            y = JustRelax.@cell py[ip, i, j]
+            x = @index px[ip, i, j]
+            y = @index py[ip, i, j]
 
             # plume - rectangular
             if y > 0.2 + 0.02 * cos(π * x / λ)
-                JustRelax.@cell phases[ip, i, j] = 2.0
+                @index phases[ip, i, j] = 2.0
             else
-                JustRelax.@cell phases[ip, i, j] = 1.0
+                @index phases[ip, i, j] = 1.0
             end
         end
         return nothing
@@ -86,9 +86,9 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
     # temperature
     pPhases,             = init_cell_arrays(particles, Val(1))
     particle_args        = (pPhases, )
-    phase_ratios         = PhaseRatio(backend_JR, ni, length(rheology))
+    phase_ratios         = PhaseRatios(backend, length(rheology), ni)
     init_phases!(pPhases, particles)
-    phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+    phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 
     # STOKES ---------------------------------------------
     # Allocate arrays needed for every Stokes problem
@@ -172,7 +172,7 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs")
         # inject && break
         inject_particles_phase!(particles, pPhases, (), (), xvi)
         # update phase ratios
-        phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+        phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 
         @show it += 1
         t        += dt

diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl
@@ -14,7 +14,7 @@ function init_phases!(phase_ratios)
     ni = size(phase_ratios.center)
 
     @parallel_indices (i, j) function init_phases!(phases)
-        JustRelax.@cell phases[1, i, j] = 1.0
+        @index phases[1, i, j] = 1.0
 
         return nothing
     end
@@ -56,7 +56,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs")
 
     # Initialize phase ratios -------------------------------
     radius       = 0.1
-    phase_ratios = PhaseRatio(backend_JR, ni, length(rheology))
+    phase_ratios = PhaseRatios(backend, length(rheology), ni)
     init_phases!(phase_ratios)
 
     # STOKES ---------------------------------------------

diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl
@@ -38,22 +38,22 @@ function init_phases!(phases, particles)
         r=100e3
         f(x, A, λ) = A * sin(π*x/λ)
 
-        @inbounds for ip in JustRelax.cellaxes(phases)
+        @inbounds for ip in cellaxes(phases)
             # quick escape
-            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+            @index(index[ip, i, j]) == 0 && continue
 
-            x = JustRelax.@cell px[ip, i, j]
-            depth = -(JustRelax.@cell py[ip, i, j])
-            JustRelax.@cell phases[ip, i, j] = 2.0
+            x = @index px[ip, i, j]
+            depth = -(@index py[ip, i, j])
+            @index phases[ip, i, j] = 2.0
 
             if 0e0 ≤ depth ≤ 100e3
-                JustRelax.@cell phases[ip, i, j] = 1.0
+                @index phases[ip, i, j] = 1.0
 
             else
-                JustRelax.@cell phases[ip, i, j] = 2.0
+                @index phases[ip, i, j] = 2.0
 
                 if ((x - 250e3)^2 + (depth - 250e3)^2 ≤ r^2)
-                    JustRelax.@cell phases[ip, i, j] = 3.0
+                    @index phases[ip, i, j] = 3.0
                 end
             end
 
@@ -119,8 +119,8 @@ function main(igg, nx, ny)
 
     # Elliptical temperature anomaly
     init_phases!(pPhases, particles)
-    phase_ratios  = PhaseRatio(backend_JR, ni, length(rheology))
-    phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+    phase_ratios  = PhaseRatios(backend, length(rheology), ni)
+    phase_ratios_center!(phase_ratios, particles, xci, pPhases)
     # ----------------------------------------------------
 
     # STOKES ---------------------------------------------
@@ -186,7 +186,7 @@ function main(igg, nx, ny)
         # check if we need to inject particles
         inject_particles_phase!(particles, pPhases, (), (), xvi)
         # update phase ratios
-        phase_ratios_center!(phase_ratios, particles, grid, pPhases)
+        phase_ratios_center!(phase_ratios, particles, xci, pPhases)
 
         @show it += 1
         t        += dt
-Original file line number
+Diff line change
@@ Expand Up / @@ -14,7 +14,6 @@ jobs: @@
           fail-fast: false
           matrix:
             version:
-              - '1.9'
               - '1.10'
               - 'pre'
               # - 'nightly'
@@ Expand Down @@