From 062f4bf5b6dc71fecdfdbc910c66a6440bc68fd7 Mon Sep 17 00:00:00 2001 From: Pascal Aellig <93140290+aelligp@users.noreply.github.com> Date: Tue, 24 Sep 2024 16:02:06 +0200 Subject: [PATCH] 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 * format --------- Co-authored-by: Albert de Montserrat --- .buildkite/pipeline.yml | 3 +- .github/workflows/ci.yml | 1 - Project.toml | 4 +- docs/src/man/Blankenbach.md | 12 +- docs/src/man/ShearBands.md | 10 +- docs/src/man/subduction2D/subduction2D.md | 6 +- .../Blankenbach2D/Benchmark2D_WENO5.jl | 4 +- .../stokes2D/Blankenbach2D/Benchmark2D_sgd.jl | 6 +- .../Blankenbach2D/Benchmark2D_sgd_scaled.jl | 6 +- .../Blankenbach2D/Blankenbach_Rheology.jl | 10 +- .../Blankenbach_Rheology_scaled.jl | 8 +- .../stokes2D/VanKeken.jl/VanKeken.jl | 20 +- .../Elastic_BuildUp_phases_incompressible.jl | 4 +- .../PlumeFreeSurface_2D.jl | 22 +-- .../RayleighTaylor2D.jl | 20 +- .../stokes2D/shear_band/ShearBand2D.jl | 10 +- .../stokes2D/shear_band/ShearBand2D_MPI.jl | 10 +- .../shear_band/ShearBand2D_displacement.jl | 10 +- .../shear_band/ShearBand2D_softening.jl | 10 +- .../stokes2D/shear_heating/Shearheating2D.jl | 6 +- .../shear_heating/Shearheating_rheology.jl | 12 +- .../stokes2D/sinking_block/SinkingBlock2D.jl | 14 +- .../stokes3D/shear_band/ShearBand3D.jl | 10 +- .../stokes3D/shear_band/ShearBand3D_MPI.jl | 10 +- .../stokes3D/shear_heating/Shearheating3D.jl | 10 +- .../shear_heating/Shearheating_rheology.jl | 14 +- .../diffusion/diffusion2D.jl | 2 +- .../diffusion/diffusion2D_MPI.jl | 2 +- .../diffusion/diffusion2D_multiphase.jl | 18 +- .../diffusion/diffusion3D_multiphase.jl | 20 +- .../Thermal_Stress_Magma_Chamber_nondim.jl | 22 +-- .../Thermal_Stress_Magma_Chamber_nondim3D.jl | 24 +-- .../convection/GlobalConvection2D_WENO5.jl | 2 +- .../convection/GlobalConvection3D_Upwind.jl | 2 +- .../Particles2D/Layered_convection2D.jl | 8 +- .../Particles2D/Layered_rheology.jl | 22 +-- .../Layered_convection2D.jl | 8 +- .../Particles2D_nonDim/Layered_rheology.jl | 83 ++++---- .../Particles3D/Layered_convection3D.jl | 6 +- .../Particles3D/Layered_rheology.jl | 22 +-- miniapps/convection/RisingBlob3D/Blob3D.jl | 24 +-- miniapps/convection/WENO5/Layered_rheology.jl | 20 +- .../convection/WENO5/WENO_convection2D.jl | 8 +- miniapps/subduction/2D/Subd2D.jl | 8 +- miniapps/subduction/2D/Subd2D_rheology.jl | 22 +-- src/JustRelax.jl | 7 +- src/JustRelax_CPU.jl | 6 + src/common.jl | 5 +- src/ext/AMDGPU/2D.jl | 45 +---- src/ext/AMDGPU/3D.jl | 65 ++++--- src/ext/CUDA/2D.jl | 46 +---- src/ext/CUDA/3D.jl | 86 ++++---- src/particles/subgrid_diffusion.jl | 80 +------- src/phases/CellArrays.jl | 137 ------------- src/phases/phases.jl | 184 +----------------- src/rheology/BuoyancyForces.jl | 10 +- src/rheology/Viscosity.jl | 4 +- src/stokes/Stokes2D.jl | 2 +- src/stokes/Stokes3D.jl | 2 +- src/stokes/StressRotation.jl | 32 +-- src/thermal_diffusion/DiffusionExplicit.jl | 4 +- .../DiffusionPT_GeoParams.jl | 6 +- src/thermal_diffusion/ShearHeating.jl | 2 +- src/types/phases.jl | 8 - src/types/traits.jl | 2 +- test/Project.toml | 4 +- test/test_Blankenbach.jl | 8 +- test/test_CellArrays2D.jl | 36 ---- test/test_CellArrays3D.jl | 40 ---- test/test_VanKeken.jl | 20 +- test/test_WENO5.jl | 2 +- test/test_diffusion2D.jl | 2 +- test/test_diffusion2D_multiphase.jl | 18 +- test/test_diffusion3D_multiphase.jl | 20 +- test/test_shearband2D.jl | 28 ++- test/test_shearband2D_softening.jl | 28 ++- test/test_shearheating2D.jl | 9 +- test/test_shearheating3D.jl | 9 +- test/test_sinking_block.jl | 14 +- test/test_thermalstresses.jl | 21 +- 80 files changed, 533 insertions(+), 1034 deletions(-) delete mode 100644 src/phases/CellArrays.jl delete mode 100644 src/types/phases.jl delete mode 100644 test/test_CellArrays2D.jl delete mode 100644 test/test_CellArrays3D.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 068ab44a8..a83c19c79 100644 --- 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 index e62d480fe..136cb7e69 100644 --- 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 index 4de9d39c6..2df56c3ff 100644 --- 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 index fda37d4e4..fa50938a8 100644 --- 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 index 99f4165b4..fb9c381f3 100644 --- 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 index a09f361c1..03393f031 100644 --- 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 index e563bb6aa..c64e88c49 100755 --- 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 index 1e383b54b..25a6914ab 100755 --- 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 index 00ff92452..0b3b023b6 100755 --- 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 index 77da2cfdb..e77714ab3 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl @@ -13,7 +13,7 @@ function init_rheologies() CompositeRheology = CompositeRheology((LinearViscous(; η=1.0e23),)), RadioactiveHeat = ConstantRadioactiveHeat(0.0), Gravity = ConstantGravity(; g=10.0), - ), + ), ) end @@ -21,14 +21,14 @@ 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 \ No newline at end of file +end diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl index 15954fa3b..56f3aba6b 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl @@ -11,7 +11,7 @@ function init_rheologies() CompositeRheology = CompositeRheology((LinearViscous(; η=1),)), RadioactiveHeat = ConstantRadioactiveHeat(0.0), Gravity = ConstantGravity(; g = 1e4), - ), + ), ) end @@ -19,10 +19,10 @@ 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 index aec05ea30..4fe38e9f8 100644 --- 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 index f8c694ebc..c6fa0d089 100644 --- 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 index 65fe52543..d0b38dbb2 100644 --- 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 diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 1aceaa8c3..440f642c4 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -38,19 +38,19 @@ function init_phases!(phases, particles, A) 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 elseif depth > (-f(x, A, 500e3) + (200e3 - A)) - JustRelax.@cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 end @@ -116,9 +116,9 @@ function RT_2D(igg, nx, ny) # Elliptical temperature anomaly A = 5e3 # Amplitude of the anomaly - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, A) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -197,7 +197,7 @@ function RT_2D(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 diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 61728b781..8c23d7c78 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -16,12 +16,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 @@ -88,7 +88,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl index 6fb866448..4e23d4416 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl @@ -16,12 +16,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 @@ -83,7 +83,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_displacement.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_displacement.jl index 486d8bc8c..1c34f84fd 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_displacement.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_displacement.jl @@ -16,12 +16,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 @@ -88,7 +88,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl index 94a8b2640..9572b2812 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl @@ -15,12 +15,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 @@ -85,7 +85,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index 544b20714..22c440e1e 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -79,9 +79,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx / 2 # origin of thermal anomaly yc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -245,7 +245,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # 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) @show it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl index 40fb2c8a5..b0cf2fe08 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl @@ -57,17 +57,17 @@ function init_phases!(phases, particles, Lx, d, r) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx, d) - @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] = 1.0 # matrix + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) + @index phases[ip, i, j] = 1.0 # matrix # thermal anomaly - circular if ((x - Lx)^2 + (depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 end end return nothing diff --git a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl index 3b3d52a69..dc8458a15 100644 --- a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl +++ b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl @@ -33,14 +33,14 @@ function init_phases!(phases, particles, xc, yc, r) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, xc, yc, r) - @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]) + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) # plume - rectangular - JustRelax.@cell phases[ip, i, j] = if ((x -xc)^2 ≤ r^2) && ((depth - yc)^2 ≤ r^2) + @index phases[ip, i, j] = if ((x -xc)^2 ≤ r^2) && ((depth - yc)^2 ≤ r^2) 2.0 else 1.0 @@ -113,9 +113,9 @@ function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_per xc_anomaly = 250e3 # origin of thermal anomaly yc_anomaly = -(ly-400e3) # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, xc_anomaly, abs(yc_anomaly), r_anomaly) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl index 48256186f..f767c476a 100644 --- a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend using Printf, GeoParams, GLMakie, CellArrays, CSV, DataFrames @@ -18,12 +18,12 @@ function init_phases!(phase_ratios, xci, radius) @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 - JustRelax.@cell phases[1, i, j, k] = 1.0 - JustRelax.@cell phases[2, i, j, k] = 0.0 + @index phases[1, i, j, k] = 1.0 + @index phases[2, i, j, k] = 0.0 else - JustRelax.@cell phases[1, i, j, k] = 0.0 - JustRelax.@cell phases[2, i, j, k] = 1.0 + @index phases[1, i, j, k] = 0.0 + @index phases[2, i, j, k] = 1.0 end return nothing diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D_MPI.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D_MPI.jl index 45bae0a2b..ccd9cfffc 100644 --- a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D_MPI.jl +++ b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D_MPI.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend using GeoParams, GLMakie @@ -18,12 +18,12 @@ function init_phases!(phase_ratios, xci, radius) @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 - @cell phases[1, i, j, k] = 1.0 - @cell phases[2, i, j, k] = 0.0 + @index phases[1, i, j, k] = 1.0 + @index phases[2, i, j, k] = 0.0 else - @cell phases[1, i, j, k] = 0.0 - @cell phases[2, i, j, k] = 1.0 + @index phases[1, i, j, k] = 0.0 + @index phases[2, i, j, k] = 1.0 end return nothing end diff --git a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl index b3c7ef3cd..839fd5501 100644 --- a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl @@ -1,7 +1,7 @@ # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend using ParallelStencil @@ -73,9 +73,9 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal yc_anomaly = ly/2 # origin of thermal anomaly zc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - init_phases!(backend_JR, pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) + phase_ratios = PhaseRatio(backend, length(rheology), ni) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -227,7 +227,7 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal # check if we need to inject particles inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), 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/stokes3D/shear_heating/Shearheating_rheology.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl index 81ac67dd2..babea6918 100644 --- a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl +++ b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl @@ -57,18 +57,18 @@ function init_phases!(phases, particles, Lx, Ly, d, r) ni = size(phases) @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, r, Lx, Ly, d) - @inbounds for ip in JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - depth = -(JustRelax.@cell pz[ip, I...]) - @cell phases[ip, I...] = 1.0 # matrix + x = @index px[ip, I...] + y = @index py[ip, I...] + depth = -(@index pz[ip, I...]) + @index phases[ip, I...] = 1.0 # matrix # thermal anomaly - circular if ((x - Lx)^2 + (y - Ly)^2 + (depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 end end return nothing diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl index caeacec1b..ab26199f5 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.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 JustPIC, JustPIC._2D diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl index 85f244fdf..a1fc263ad 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl @@ -5,7 +5,7 @@ const backend_JR = CPUBackend using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -import JustRelax.@cell + using GeoParams using GLMakie diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl index 01d3bb4a2..1efa43b88 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl @@ -4,7 +4,7 @@ const backend_JR = CPUBackend using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -import JustRelax.@cell + using GeoParams distance(p1, p2) = mapreduce(x->(x[1]-x[2])^2, +, zip(p1, p2)) |> sqrt @@ -37,19 +37,19 @@ function init_phases!(phases, particles, xc, yc, r) center = xc, yc @parallel_indices (i, j) function init_phases!(phases, px, py, index, center, r) - @inbounds for ip in JustRelax.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 (((x - center[1] ))^2 + ((y - center[2]))^2) ≤ r^2 - 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 @@ -124,9 +124,9 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) backend, nxcell, max_xcell, min_xcell, xvi... ) 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, center_perturbation..., r) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- @parallel (@idx ni) compute_temperature_source_terms!(thermal.H, rheology, phase_ratios.center, args) diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl index 29901841d..6e1147c8e 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D -import JustRelax.@cell + const backend_JR = CPUBackend @@ -42,20 +42,20 @@ function init_phases!(phases, particles, xc, yc, zc, r) center = xc, yc, zc @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, center, r) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - z = JustRelax.@cell pz[ip, I...] + x = @index px[ip, I...] + y = @index py[ip, I...] + z = @index pz[ip, I...] # plume - rectangular if (((x - center[1]))^2 + ((y - center[2]))^2 + ((z - center[3]))^2) ≤ r^2 - JustRelax.@cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 else - JustRelax.@cell phases[ip, I...] = 1.0 + @index phases[ip, I...] = 1.0 end end return nothing @@ -143,9 +143,9 @@ function diffusion_3D(; backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni ) 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, center_perturbation..., r) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # PT coefficients for thermal diffusion diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl index 87ef9b0d7..81c166b35 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @@ -46,23 +46,23 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, r_anomaly, stic @parallel_indices (i, j) function init_phases!( phases, px, py, index, xc_anomaly, yc_anomaly, r_anomaly, sticky_air, top, bottom ) - @inbounds for ip in JustRelax.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]) - sticky_air + x = @index px[ip, i, j] + y = -(@index py[ip, i, j]) - sticky_air if top ≤ y ≤ bottom - @cell phases[ip, i, j] = 1.0 # crust + @index phases[ip, i, j] = 1.0 # crust end # thermal anomaly - circular if ((x - xc_anomaly)^2 + (y + yc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 end if y < top - @cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 end end return nothing @@ -238,8 +238,8 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) init_phases!(pPhases, particles, x_anomaly, y_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), nondimensionalize(20km,CharDim)) - 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) # Initialisation of thermal profile thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions @@ -432,7 +432,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) # 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) particle2grid!(T_buffer, pT, xvi, particles) @views T_buffer[:, end] .= Tsurf diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl index dafd77054..05dc630f2 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences3D @@ -38,24 +38,24 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_a @parallel_indices (I...) function init_phases!( phases, px, py, pz, index, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, sticky_air, top, bottom ) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - z = -(JustRelax.@cell pz[ip, I...]) - sticky_air + x = @index px[ip, I...] + y = @index py[ip, I...] + z = -(@index pz[ip, I...]) - sticky_air if top ≤ z ≤ bottom - @cell phases[ip, I...] = 1.0 # crust + @index phases[ip, I...] = 1.0 # crust end # thermal anomaly - circular if ((x - xc_anomaly)^2 + (y - yc_anomaly)^2 + (z + zc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 end if z < top - JustRelax.@cell phases[ip, I...] = 3.0 + @index phases[ip, I...] = 3.0 end end return nothing @@ -240,9 +240,9 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals z_anomaly = nondimensionalize(-5km,CharDim) # origin of the small thermal anomaly r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, x_anomaly, y_anomaly, z_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), nondimensionalize(20km,CharDim)) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # Initialisation of thermal profile thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions @@ -417,7 +417,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # check if we need to inject particles inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), xvi) # update phase ratios - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) particle2grid!(thermal.T, pT, xvi, particles) @views thermal.T[:, :, end] .= Tsurf diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index 6e217d923..1be51fcb0 100644 --- a/miniapps/convection/GlobalConvection2D_WENO5.jl +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend = CPUBackend diff --git a/miniapps/convection/GlobalConvection3D_Upwind.jl b/miniapps/convection/GlobalConvection3D_Upwind.jl index 41dd75446..8eb7bc76a 100644 --- a/miniapps/convection/GlobalConvection3D_Upwind.jl +++ b/miniapps/convection/GlobalConvection3D_Upwind.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index a99906281..6c6d1d0f5 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend @@ -126,9 +126,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx/2 # origin of thermal anomaly yc_anomaly = -610e3 # origin of thermal anomaly r_anomaly = 25e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly, thick_air) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -292,7 +292,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # 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) @show it += 1 t += dt diff --git a/miniapps/convection/Particles2D/Layered_rheology.jl b/miniapps/convection/Particles2D/Layered_rheology.jl index 502dc3cbe..fab0a894a 100644 --- a/miniapps/convection/Particles2D/Layered_rheology.jl +++ b/miniapps/convection/Particles2D/Layered_rheology.jl @@ -114,32 +114,32 @@ function init_phases!(phases, particles, Lx, d, r, thick_air) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) - @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]) - thick_air + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) - thick_air if 0e0 ≤ depth ≤ 21e3 - JustRelax.@cell phases[ip, i, j] = 1.0 + @index phases[ip, i, j] = 1.0 elseif 35e3 ≥ depth > 21e3 - JustRelax.@cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 elseif 90e3 ≥ depth > 35e3 - JustRelax.@cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth > 90e3 - JustRelax.@cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth < 0e0 - JustRelax.@cell phases[ip, i, j] = 5.0 + @index phases[ip, i, j] = 5.0 end # plume - rectangular - if ((x - Lx * 0.5)^2 ≤ r^2) && (((JustRelax.@cell py[ip, i, j]) - d - thick_air)^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 4.0 + if ((x - Lx * 0.5)^2 ≤ r^2) && (((@index py[ip, i, j]) - d - thick_air)^2 ≤ r^2) + @index phases[ip, i, j] = 4.0 end end return nothing diff --git a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl index 6e24bd164..0b532d538 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend @@ -133,9 +133,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx/2 # origin of thermal anomaly yc_anomaly = nondimensionalize(-610km, CharDim) # origin of thermal anomaly r_anomaly = nondimensionalize(25km, CharDim) # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly, thick_air, CharDim) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -298,7 +298,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # 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) @show it += 1 t += dt diff --git a/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl b/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl index c23d33856..25e93a25f 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl @@ -4,52 +4,52 @@ function init_rheologies(CharDim; is_plastic = true) # Dislocation and Diffusion creep disl_upper_crust = DislocationCreep( - A=5.07e-18Pa^(-23//10) / s, # units are Pa^(-n) / s - n=2.3NoUnits, - E=154e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + A=5.07e-18Pa^(-23//10) / s, # units are Pa^(-n) / s + n=2.3NoUnits, + E=154e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, R=8.3145J/mol/K ) disl_lower_crust = DislocationCreep( - A=2.08e-23Pa^(-32//10) / s, # units are Pa^(-n) / s - n=3.2NoUnits, - E=238e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + A=2.08e-23Pa^(-32//10) / s, # units are Pa^(-n) / s + n=3.2NoUnits, + E=238e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, R=8.3145J/mol/K, ) disl_lithospheric_mantle = DislocationCreep( - A=2.51e-17Pa^(-35//10) / s, # units are Pa^(-n) / s - n=3.5NoUnits, - E=530e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + A=2.51e-17Pa^(-35//10) / s, # units are Pa^(-n) / s + n=3.5NoUnits, + E=530e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, R=8.3145J/mol/K ) disl_sublithospheric_mantle = DislocationCreep( A=2.51e-17Pa^(-35//10) / s, # units are Pa^(-n) / s - n=3.5NoUnits, - E=530e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + n=3.5NoUnits, + E=530e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, R=8.3145J/mol/K ) diff_lithospheric_mantle = DiffusionCreep( A=2.51e-17Pa^(-1) / s, # units are Pa^(-n - r) / s * m^(-p) - n=1.0NoUnits, - E=530e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + n=1.0NoUnits, + E=530e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, p=0.0NoUnits, R=8.3145J/mol/K ) diff_sublithospheric_mantle = DiffusionCreep( A=2.51e-17Pa^(-1) / s, # units are Pa^(-n - r) / s * m^(-p) - n=1.0NoUnits, - E=530e3J/mol, - V=6e-6m^3/mol, - r=0.0NoUnits, + n=1.0NoUnits, + E=530e3J/mol, + V=6e-6m^3/mol, + r=0.0NoUnits, p=0.0NoUnits, R=8.3145J/mol/K ) @@ -79,7 +79,7 @@ function init_rheologies(CharDim; is_plastic = true) else DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity end - + K_crust = TP_Conductivity(; a = 0.64Watt / K / m , b = 807e00Watt / m , @@ -92,7 +92,7 @@ function init_rheologies(CharDim; is_plastic = true) c = 0.77K, d = 0.00004/ MPa, ) - + g = 9.81m/s^2 # Define rheolgy struct @@ -162,32 +162,32 @@ function init_phases!(phases, particles, Lx, d, r, thick_air, CharDim) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx, CharDim) - @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]) - nondimensionalize(thick_air * km, CharDim) + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) - nondimensionalize(thick_air * km, CharDim) if nondimensionalize(0e0km, CharDim) ≤ depth ≤ nondimensionalize(21km, CharDim) - JustRelax.@cell phases[ip, i, j] = 1.0 + @index phases[ip, i, j] = 1.0 elseif nondimensionalize(35km, CharDim) ≥ depth > nondimensionalize(21km, CharDim) - JustRelax.@cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 elseif nondimensionalize(90km, CharDim) ≥ depth > nondimensionalize(35km, CharDim) - JustRelax.@cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth > nondimensionalize(90km, CharDim) - JustRelax.@cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth < nondimensionalize(0e0km, CharDim) - JustRelax.@cell phases[ip, i, j] = 5.0 + @index phases[ip, i, j] = 5.0 - end + end # plume - rectangular - if ((x - Lx * 0.5)^2 ≤ r^2) && (((JustRelax.@cell py[ip, i, j]) - d - nondimensionalize(thick_air * km, CharDim))^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 4.0 + if ((x - Lx * 0.5)^2 ≤ r^2) && (((@index py[ip, i, j]) - d - nondimensionalize(thick_air * km, CharDim))^2 ≤ r^2) + @index phases[ip, i, j] = 4.0 end end return nothing @@ -195,4 +195,3 @@ function init_phases!(phases, particles, Lx, d, r, thick_air, CharDim) @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, CharDim) end - diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl index 7adf58db8..8bf346615 100644 --- a/miniapps/convection/Particles3D/Layered_convection3D.jl +++ b/miniapps/convection/Particles3D/Layered_convection3D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend @@ -115,7 +115,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) zc_anomaly = -610e3 # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation init_phases!(pPhases, particles, lx, ly; d=abs(zc_anomaly), r=r_anomaly) - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) @parallel (@idx ni) phase_ratios_center!(phase_ratios.center, particles.coords, xci, di, pPhases) # ---------------------------------------------------- @@ -265,7 +265,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) # 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) @show it += 1 t += dt diff --git a/miniapps/convection/Particles3D/Layered_rheology.jl b/miniapps/convection/Particles3D/Layered_rheology.jl index 10c9b840d..e041a7331 100644 --- a/miniapps/convection/Particles3D/Layered_rheology.jl +++ b/miniapps/convection/Particles3D/Layered_rheology.jl @@ -121,34 +121,34 @@ function init_phases!(phases, particles, Lx, Ly; d=650e3, r=50e3) @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, r, Lx, Ly) - @inbounds for ip in JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - depth = -(JustRelax.@cell pz[ip, I...]) + x = @index px[ip, I...] + y = @index py[ip, I...] + depth = -(@index pz[ip, I...]) if 0e0 ≤ depth ≤ 21e3 - JustRelax.@cell phases[ip, I...] = 1.0 + @index phases[ip, I...] = 1.0 elseif 35e3 ≥ depth > 21e3 - JustRelax.@cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 elseif 90e3 ≥ depth > 35e3 - JustRelax.@cell phases[ip, I...] = 3.0 + @index phases[ip, I...] = 3.0 elseif depth > 90e3 - JustRelax.@cell phases[ip, I...] = 3.0 + @index phases[ip, I...] = 3.0 elseif 0e0 > depth - JustRelax.@cell phases[ip, I...] = 5.0 + @index phases[ip, I...] = 5.0 end # plume - rectangular if ((x - Lx * 0.5)^2 ≤ r^2) && ((y - Ly * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, I...] = 4.0 + @index phases[ip, I...] = 4.0 end end return nothing diff --git a/miniapps/convection/RisingBlob3D/Blob3D.jl b/miniapps/convection/RisingBlob3D/Blob3D.jl index 5c186be60..93d148a9e 100644 --- a/miniapps/convection/RisingBlob3D/Blob3D.jl +++ b/miniapps/convection/RisingBlob3D/Blob3D.jl @@ -6,7 +6,7 @@ const isCUDA = true end using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = @static if isCUDA CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend @@ -54,24 +54,24 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_a @parallel_indices (I...) function init_phases!( phases, px, py, pz, index, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, sticky_air, top, bottom ) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - z = -(JustRelax.@cell pz[ip, I...]) - sticky_air + x = @index px[ip, I...] + y = @index py[ip, I...] + z = -(@index pz[ip, I...]) - sticky_air if top ≤ z ≤ bottom - @cell phases[ip, I...] = 1.0 # crust + @index phases[ip, I...] = 1.0 # crust end # thermal anomaly - circular if ((x - xc_anomaly)^2 + (y - yc_anomaly)^2 + (z + zc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 end if z < top - JustRelax.@cell phases[ip, I...] = 3.0 + @index phases[ip, I...] = 3.0 end end return nothing @@ -251,9 +251,9 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # z_anomaly = nondimensionalize(-5km,CharDim) # origin of the small thermal anomaly r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, x_anomaly, y_anomaly, z_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), nondimensionalize(20km,CharDim)) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # Initialisation of thermal profile thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions @@ -417,7 +417,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # check if we need to inject particles inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), xvi) # update phase ratios - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) particle2grid!(thermal.T, pT, xvi, particles) # @views thermal.T[:, :, end] .= Tsurf diff --git a/miniapps/convection/WENO5/Layered_rheology.jl b/miniapps/convection/WENO5/Layered_rheology.jl index 085bd0013..d919d283d 100644 --- a/miniapps/convection/WENO5/Layered_rheology.jl +++ b/miniapps/convection/WENO5/Layered_rheology.jl @@ -125,32 +125,32 @@ function init_phases!(phases, particles, Lx; d=650e3, r=50e3) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) - @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]) + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) if 0e0 ≤ depth ≤ 21e3 - @cell phases[ip, i, j] = 1.0 + @index phases[ip, i, j] = 1.0 elseif 35e3 ≥ depth > 21e3 - @cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 elseif 90e3 ≥ depth > 35e3 - @cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth > 90e3 - @cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 elseif depth < 0e0 - @cell phases[ip, i, j] = 5.0 + @index phases[ip, i, j] = 5.0 end # plume - rectangular if ((x - Lx * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 4.0 + @index phases[ip, i, j] = 4.0 end end return nothing diff --git a/miniapps/convection/WENO5/WENO_convection2D.jl b/miniapps/convection/WENO5/WENO_convection2D.jl index 9b405a3d3..c549fcfba 100644 --- a/miniapps/convection/WENO5/WENO_convection2D.jl +++ b/miniapps/convection/WENO5/WENO_convection2D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend_JR = CPUBackend @@ -117,8 +117,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) yc_anomaly = -610e3 # origin of thermal anomaly r_anomaly = 25e3 # radius of perturbation init_phases!(pPhases, particles, lx; d=abs(yc_anomaly), r=r_anomaly) - 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 --------------------------------------------- @@ -259,7 +259,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (pT, ), (T_WENO,), 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/subduction/2D/Subd2D.jl b/miniapps/subduction/2D/Subd2D.jl index 0d924324c..e64aca162 100644 --- a/miniapps/subduction/2D/Subd2D.jl +++ b/miniapps/subduction/2D/Subd2D.jl @@ -6,7 +6,7 @@ const isCUDA = true end using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -import JustRelax.@cell + const backend = @static if isCUDA CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend @@ -94,9 +94,9 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Assign particles phases anomaly phases_device = PTArray(backend)(phases_GMG) - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = 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) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -264,7 +264,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # 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) @show it += 1 t += dt diff --git a/miniapps/subduction/2D/Subd2D_rheology.jl b/miniapps/subduction/2D/Subd2D_rheology.jl index 36057edc8..d7a8c5f3f 100644 --- a/miniapps/subduction/2D/Subd2D_rheology.jl +++ b/miniapps/subduction/2D/Subd2D_rheology.jl @@ -18,15 +18,15 @@ function init_rheology_nonNewtonian_plastic() diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) # plasticity ϕ_wet_olivine = asind(0.1) - C_wet_olivine = 1e6 + C_wet_olivine = 1e6 η_reg = 1e19 - lithosphere_rheology = CompositeRheology( + lithosphere_rheology = CompositeRheology( ( disl_wet_olivine, diff_wet_olivine, DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity - ) + ) ) init_rheologies(lithosphere_rheology) end @@ -88,12 +88,12 @@ end ni = size(phases) - for ip in JustRelax.cellaxes(phases) + for ip in cellaxes(phases) # quick escape - @cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - pᵢ = ntuple(Val(N)) do i - @cell pcoords[i][ip, I...] + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] end d = Inf # distance to the nearest particle @@ -106,8 +106,8 @@ end !(jj ≤ ni[2]) && continue xvᵢ = ( - xvi[1][ii], - xvi[2][jj], + xvi[1][ii], + xvi[2][jj], ) d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) if d_ijk < d @@ -115,8 +115,8 @@ end particle_phase = phase_grid[ii, jj] end end - @cell phases[ip, I...] = Float64(particle_phase) + @index phases[ip, I...] = Float64(particle_phase) end return nothing -end \ No newline at end of file +end diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 0c0491565..51459fe7a 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -10,6 +10,7 @@ using HDF5 using CellArrays using StaticArrays using Statistics +@reexport using JustPIC function solve!() end @@ -29,9 +30,6 @@ include("types/stokes.jl") include("types/heat_diffusion.jl") # export ThermalArrays, PTThermalCoeffs -include("types/phases.jl") -# export PhaseRatio - include("boundaryconditions/types.jl") export TemperatureBoundaryConditions, DisplacementBoundaryConditions, VelocityBoundaryConditions @@ -42,9 +40,6 @@ export BackendTrait, CPUBackendTrait, NonCPUBackendTrait include("topology/Topology.jl") export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g -include("phases/CellArrays.jl") -export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! - include("JustRelax_CPU.jl") include("IO/DataIO.jl") diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 874fc2670..cb4134eee 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -1,6 +1,7 @@ module JustRelax2D using ..JustRelax +using JustPIC, JustPIC._2D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences2D @@ -18,6 +19,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._2D: numphases, nphases + @init_parallel_stencil(Threads, Float64, 2) include("common.jl") @@ -29,6 +32,7 @@ end module JustRelax3D using ..JustRelax +using JustPIC, JustPIC._3D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences3D @@ -46,6 +50,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._3D: numphases, nphases + @init_parallel_stencil(Threads, Float64, 3) include("common.jl") diff --git a/src/common.jl b/src/common.jl index e3c9318a4..0326a2324 100644 --- a/src/common.jl +++ b/src/common.jl @@ -4,9 +4,6 @@ export StokesArrays, PTStokesCoeffs include("types/constructors/heat_diffusion.jl") export ThermalArrays, PTThermalCoeffs -include("types/constructors/phases.jl") -export PhaseRatio - include("Utils.jl") export @allocate, @add, @@ -46,7 +43,7 @@ export AbstractBoundaryConditions, include("MiniKernels.jl") include("phases/phases.jl") -export fn_ratio, phase_ratios_center!, phase_ratios_vertex!, numphases, nphases +export fn_ratio include("rheology/BuoyancyForces.jl") export compute_ρg! diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index fab57998c..96eb34d38 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -2,6 +2,7 @@ module JustRelax2D using JustRelax: JustRelax using AMDGPU +using JustPIC, JustPIC._2D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences2D @@ -19,8 +20,8 @@ import JustRelax: backend, CPUBackend, AMDGPUBackend, - Geometry, - @cell + Geometry + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, @@ -28,6 +29,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._2D: nphases, numphases + @init_parallel_stencil(AMDGPU, Float64, 2) include("../../common.jl") @@ -46,10 +49,6 @@ function JR2D.ThermalArrays(::Type{AMDGPUBackend}, ni::Vararg{Number,N}) where { return ThermalArrays(ni...) end -function JR2D.PhaseRatio(::Type{AMDGPUBackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) -end - function JR2D.PTThermalCoeffs( ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -171,34 +170,8 @@ function thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bc return thermal_bcs!(thermal.T, bcs) end -# Phases -function JR2D.phase_ratios_center!( - ::AMDGPUBackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_center!(phase_ratios, particles, grid, phases) -end - -function JR2D.phase_ratios_vertex!( - ::AMDGPUBackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) -end - # Rheology -## viscosity -function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) -end - function JR2D.compute_viscosity!( ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) @@ -232,7 +205,7 @@ end function JR2D.compute_ρg!(ρg::ROCArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR2D.compute_ρg!(ρg::ROCArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR2D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -242,7 +215,7 @@ function JR2D.compute_melt_fraction!(ϕ::ROCArray, rheology, args) end function JR2D.compute_melt_fraction!( - ϕ::ROCArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ϕ::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args ) return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) end @@ -317,7 +290,7 @@ function JR2D.subgrid_characteristic_time!( subgrid_arrays, particles, dt₀::ROCArray, - phases::JustRelax.PhaseRatio, + phases::JustPIC.PhaseRatios, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -365,7 +338,7 @@ function JR2D.compute_shear_heating!(::AMDGPUBackendTrait, thermal, stokes, rheo end function JR2D.compute_shear_heating!( - ::AMDGPUBackendTrait, thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt + ::AMDGPUBackendTrait, thermal, stokes, phase_ratios::JustPIC.PhaseRatios, rheology, dt ) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 073d3d09a..1790bf59d 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -2,6 +2,7 @@ module JustRelax3D using JustRelax: JustRelax using AMDGPU +using JustPIC, JustPIC._3D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences3D @@ -19,8 +20,8 @@ import JustRelax: backend, CPUBackend, AMDGPUBackend, - Geometry, - @cell + Geometry + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, @@ -28,6 +29,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._3D: numphases, nphases + @init_parallel_stencil(AMDGPU, Float64, 3) include("../../common.jl") @@ -46,10 +49,6 @@ function JR3D.ThermalArrays(::Type{AMDGPUBackend}, ni::Vararg{Number,N}) where { return ThermalArrays(ni...) end -function JR3D.PhaseRatio(::Type{AMDGPUBackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) -end - function JR3D.PTThermalCoeffs( ::Type{AMDGPUBackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -171,27 +170,6 @@ function thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bc return thermal_bcs!(thermal.T, bcs) end -# Phases -function JR3D.phase_ratios_center!( - ::AMDGPUBackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_center!(phase_ratios, particles, grid, phases) -end - -function JR3D.phase_ratios_vertex!( - ::AMDGPUBackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) -end - # Rheology ## viscosity @@ -232,7 +210,12 @@ end function JR3D.compute_ρg!(ρg::ROCArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) + +function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end + +function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -242,11 +225,15 @@ function JR3D.compute_melt_fraction!(ϕ::ROCArray, rheology, args) end function JR3D.compute_melt_fraction!( - ϕ::ROCArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ϕ::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args ) return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) end +function JR3D.compute_melt_fraction!(ϕ::ROCArray, phase_ratios, rheology, args) + return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) +end + # Interpolations function JR3D.temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays) return _temperature2center!(thermal) @@ -314,7 +301,7 @@ function JR3D.subgrid_characteristic_time!( subgrid_arrays, particles, dt₀::ROCArray, - phases::JustRelax.PhaseRatio, + phases::JustPIC.PhaseRatios, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -362,7 +349,23 @@ function JR3D.compute_shear_heating!(::AMDGPUBackendTrait, thermal, stokes, rheo end function JR3D.compute_shear_heating!( - ::AMDGPUBackendTrait, thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt + ::AMDGPUBackendTrait, thermal, stokes, phase_ratios::JustPIC.PhaseRatios, rheology, dt +) + ni = size(thermal.shear_heating) + @parallel (@idx ni) compute_shear_heating_kernel!( + thermal.shear_heating, + @tensor_center(stokes.τ), + @tensor_center(stokes.τ_o), + @strain(stokes), + phase_ratios.center, + rheology, + dt, + ) + return nothing +end + +function JR3D.compute_shear_heating!( + ::AMDGPUBackendTrait, thermal, stokes, phase_ratios, rheology, dt ) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 2e4df75ee..05d8f3f3d 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -2,6 +2,7 @@ module JustRelax2D using JustRelax: JustRelax using CUDA +using JustPIC, JustPIC._2D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences2D @@ -12,14 +13,8 @@ using MPI import JustRelax.JustRelax2D as JR2D import JustRelax: - IGG, - BackendTrait, - CPUBackendTrait, - CUDABackendTrait, - backend, - CPUBackend, - Geometry, - @cell + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, @@ -27,6 +22,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._2D: numphases, nphases + @init_parallel_stencil(CUDA, Float64, 2) include("../../common.jl") @@ -45,10 +42,6 @@ function JR2D.ThermalArrays(::Type{CUDABackend}, ni::Vararg{Number,N}) where {N} return ThermalArrays(ni...) end -function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) -end - function JR2D.PTThermalCoeffs( ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -166,27 +159,6 @@ function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end -# Phases -function JR2D.phase_ratios_center!( - ::CUDABackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_center!(phase_ratios, particles, grid, phases) -end - -function JR2D.phase_ratios_vertex!( - ::CUDABackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) -end - # Rheology ## viscosity @@ -228,7 +200,7 @@ function JR2D.compute_ρg!(ρg::CuArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -238,7 +210,7 @@ function JR2D.compute_melt_fraction!(ϕ::CuArray, rheology, args) end function JR2D.compute_melt_fraction!( - ϕ::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ϕ::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args ) return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) end @@ -309,7 +281,7 @@ function JR2D.subgrid_characteristic_time!( subgrid_arrays, particles, dt₀::CuArray, - phases::JustRelax.PhaseRatio, + phases::JustPIC.PhaseRatios, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -357,7 +329,7 @@ function JR2D.compute_shear_heating!(::CUDABackendTrait, thermal, stokes, rheolo end function JR2D.compute_shear_heating!( - ::CUDABackendTrait, thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt + ::CUDABackendTrait, thermal, stokes, phase_ratios::JustPIC.PhaseRatios, rheology, dt ) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 9a99d1e25..c2a8962d1 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -2,6 +2,7 @@ module JustRelax3D using JustRelax: JustRelax using CUDA +using JustPIC, JustPIC._3D using StaticArrays using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences3D @@ -12,14 +13,7 @@ using MPI import JustRelax.JustRelax3D as JR3D import JustRelax: - IGG, - BackendTrait, - CPUBackendTrait, - CUDABackendTrait, - backend, - CPUBackend, - Geometry, - @cell + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, @@ -27,6 +21,8 @@ import JustRelax: DisplacementBoundaryConditions, VelocityBoundaryConditions +import JustPIC._3D: numphases, nphases + @init_parallel_stencil(CUDA, Float64, 3) include("../../common.jl") @@ -45,10 +41,6 @@ function JR3D.ThermalArrays(::Type{CUDABackend}, ni::Vararg{Number,N}) where {N} return ThermalArrays(ni...) end -function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) -end - function JR3D.PTThermalCoeffs( ::Type{CUDABackend}, K, ρCp, dt, di::NTuple, li::NTuple; ϵ=1e-8, CFL=0.9 / √3 ) @@ -86,7 +78,7 @@ end function JR3D.update_pt_thermal_arrays!( pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, rheology, args, _dt, @@ -95,6 +87,13 @@ function JR3D.update_pt_thermal_arrays!( return nothing end +function JR3D.update_pt_thermal_arrays!( + pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, phase_ratios, rheology, args, _dt +) where {T} + update_pt_thermal_arrays!(pt_thermal, phase_ratios, rheology, args, _dt) + return nothing +end + function JR3D.update_thermal_coeffs!( pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray}, rheology, phase_ratios, args, dt ) where {T} @@ -178,26 +177,18 @@ function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end -# Phases -function JR3D.phase_ratios_center!( - ::CUDABackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_center!(phase_ratios, particles, grid, phases) -end +# # Phases +# function JR3D.phase_ratios_center!( +# ::CUDABackendTrait, phase_ratios::JustPIC.PhaseRatios, particles, grid::Geometry, phases +# ) +# return _phase_ratios_center!(phase_ratios, particles, grid, phases) +# end -function JR3D.phase_ratios_vertex!( - ::CUDABackendTrait, - phase_ratios::JustRelax.PhaseRatio, - particles, - grid::Geometry, - phases, -) - return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) -end +# function JR3D.phase_ratios_vertex!( +# ::CUDABackendTrait, phase_ratios::JustPIC.PhaseRatios, particles, grid::Geometry, phases +# ) +# return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) +# end # Rheology @@ -239,7 +230,10 @@ end function JR3D.compute_ρg!(ρg::CuArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end +function JR3D.compute_ρg!(ρg::CuArray, phase_ratios, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -249,11 +243,15 @@ function JR3D.compute_melt_fraction!(ϕ::CuArray, rheology, args) end function JR3D.compute_melt_fraction!( - ϕ::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ϕ::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args ) return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) end +function JR3D.compute_melt_fraction!(ϕ::CuArray, phase_ratios, rheology, args) + return compute_melt_fraction!(ϕ, phase_ratios, rheology, args) +end + # Interpolations function JR3D.temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) return _temperature2center!(thermal) @@ -318,7 +316,7 @@ function JR3D.subgrid_characteristic_time!( subgrid_arrays, particles, dt₀::CuArray, - phases::JustRelax.PhaseRatio, + phases::JustPIC.PhaseRatios, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -366,7 +364,23 @@ function JR3D.compute_shear_heating!(::CUDABackendTrait, thermal, stokes, rheolo end function JR3D.compute_shear_heating!( - ::CUDABackendTrait, thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt + ::CUDABackendTrait, thermal, stokes, phase_ratios::JustPIC.PhaseRatios, rheology, dt +) + ni = size(thermal.shear_heating) + @parallel (@idx ni) compute_shear_heating_kernel!( + thermal.shear_heating, + @tensor_center(stokes.τ), + @tensor_center(stokes.τ_o), + @strain(stokes), + phase_ratios.center, + rheology, + dt, + ) + return nothing +end + +function JR3D.compute_shear_heating!( + ::CUDABackendTrait, thermal, stokes, phase_ratios, rheology, dt ) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( diff --git a/src/particles/subgrid_diffusion.jl b/src/particles/subgrid_diffusion.jl index bc457642a..4aef58d49 100644 --- a/src/particles/subgrid_diffusion.jl +++ b/src/particles/subgrid_diffusion.jl @@ -4,7 +4,7 @@ function subgrid_characteristic_time!( subgrid_arrays, particles, dt₀, - phases::JustRelax.PhaseRatio, + phases::JustPIC.PhaseRatios, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -51,81 +51,3 @@ end return nothing end - -# struct SubgridDiffusionArrays{T, CA} -# pT0::CA # CellArray -# pΔT::CA # CellArray -# ΔT::T # Array - -# function SubgridDiffusionArrays(particles, ni) -# pΔT, pT0 = init_cell_arrays(particles, Val(2)) -# ΔT = @zeros(ni.+1) -# CA, T = typeof(pΔT), typeof(ΔT) -# new{T, CA}(pT0, pΔT, ΔT) -# end -# end - -# @inline function init_cell_arrays(particles, ::Val{N}) where {N} -# return ntuple( -# _ -> @fill( -# 0.0, size(particles.coords[1])..., celldims = (cellsize(particles.index)) -# ), -# Val(N), -# ) -# end - -# function subgrid_diffusion!( -# pT, thermal, subgrid_arrays, pPhases, rheology, stokes, particles, T_buffer, xvi, di, dt -# ) -# (; pT0, pΔT) = subgrid_arrays -# # ni = size(pT) - -# @copy pT0.data pT.data -# grid2particle!(pT, xvi, T_buffer, particles) - -# @parallel (@idx ni) subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes.P, particles.index, di, dt) -# particle2grid!(subgrid_arrays.ΔT, pΔT, xvi, particles) - -# @parallel (@idx size(subgrid_arrays.ΔT)) update_ΔT_subgrid!(subgrid_arrays.ΔT, thermal.ΔT) -# grid2particle!(pΔT, xvi, subgrid_arrays.ΔT, particles) -# @. pT.data = pT0.data + pΔT.data -# return nothing -# end - -# @parallel_indices (i, j) function update_ΔT_subgrid!(ΔTsubgrid::_T, ΔT::_T) where _T<:AbstractMatrix -# ΔTsubgrid[i, j] = ΔT[i+1, j] - ΔTsubgrid[i, j] -# return nothing -# end - -# function subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes, particles, di, dt) -# ni = size(pT) -# @parallel (@idx ni) subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes.P, particles.index, di, dt) -# end - -# @parallel_indices (I...) function subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, P, index, di, dt) - -# P_cell = P[I...] - -# for ip in JustRelax.cellaxes(pT) -# # early escape if there is no particle in this memory locations -# doskip(index, ip, I...) && continue - -# pT0ᵢ = @cell pT0[ip, I...] -# pTᵢ = @cell pT[ip, I...] -# phase = Int(@cell(pPhases[ip, I...])) -# argsᵢ = (; T = pTᵢ, P = P_cell) -# # dimensionless numerical diffusion coefficient (0 ≤ d ≤ 1) -# d = 1 -# # Compute the characteristic timescale `dt₀` of the local cell -# ρCp = compute_ρCp(rheology, phase, argsᵢ) -# K = compute_conductivity(rheology, phase, argsᵢ) -# sum_dxi = mapreduce(x-> inv(x)^2, +, di) -# dt₀ = ρCp / (2 * K * sum_dxi) -# # subgrid diffusion of the i-th particle -# pΔTᵢ = (pTᵢ - pT0ᵢ) * (1-exp(-d * dt / dt₀)) -# @cell pT0[ip, I...] = pT0ᵢ + pΔTᵢ -# @cell pΔT[ip, I...] = pΔTᵢ -# end - -# return nothing -# end diff --git a/src/phases/CellArrays.jl b/src/phases/CellArrays.jl deleted file mode 100644 index 6df31d279..000000000 --- a/src/phases/CellArrays.jl +++ /dev/null @@ -1,137 +0,0 @@ - -@inline cellnum(A::CellArray) = prod(cellsize(A)) -@inline cellaxes(A) = map(Base.oneto, cellnum(A)) -@inline new_empty_cell(::CellArray{T,N}) where {T,N} = zeros(T) - -import Base.setindex! - -Base.@propagate_inbounds @inline function Base.getindex( - x::CellArray{SVector{Nv,T},N1,N2,T}, I::Vararg{Int,N} -) where {Nv,N,N1,N2,T} - idx_cell = cart2ind(x.dims, I...) - return SVector{Nv,T}(@inbounds x.data[idx_cell, i, 1] for i in 1:Nv) -end - -Base.@propagate_inbounds @inline function Base.getindex( - x::CPUCellArray{SVector{Nv,T},N1,N2,T}, I::Vararg{Int,N} -) where {Nv,N,N1,N2,T} - idx_cell = cart2ind(x.dims, I...) - return SVector{Nv,T}(@inbounds x.data[1, i, idx_cell] for i in 1:Nv) -end - -Base.@propagate_inbounds @inline function setindex!( - A::CellArray, x, cell::Int, I::Vararg{Int,N} -) where {N} - Base.@propagate_inbounds @inline f(A::Array, x, cell, idx) = A[1, cell, idx] = x - Base.@propagate_inbounds @inline f(A, x, cell, idx) = A[idx, cell, 1] = x - - idx = LinearIndices(n)[CartesianIndex(I...)] - - return f(A.data, x, cell, idx) -end - -""" - element(A, element_indices..., cell_indices...) - -Return a the element with `element_indices` of the Cell with `cell_indices` of the CellArray `A`. - -## Arguments -- `element_indices::Int|NTuple{N,Int}`: the `element_indices` that designate the field in accordance with `A`'s cell type. -- `cell_indices::Int|NTuple{N,Int}`: the `cell_indices` that designate the cell in accordance with `A`'s cell type. -""" -Base.@propagate_inbounds @inline function element( - A::CellArray{SVector,N,D,T_elem}, i::Int, icell::Vararg{Int,Nc} -) where {T_elem,N,Nc,D} - return viewelement(A, i, icell...) -end - -Base.@propagate_inbounds @inline function element( - A::CellArray, i::T, j::T, icell::Vararg{Int,Nc} -) where {Nc,T<:Int} - return viewelement(A, i, j, icell...) -end - -Base.@propagate_inbounds @inline function viewelement( - A::CellArray{SMatrix{Ni,Nj,T,N_array},N,D,T_elem}, i, j, icell::Vararg{Int,Nc} -) where {Nc,Ni,Nj,N_array,T,N,T_elem,D} - idx_element = cart2ind((Ni, Nj), i, j) - idx_cell = cart2ind(A.dims, icell...) - return _viewelement(A.data, idx_element, idx_cell) -end - -Base.@propagate_inbounds @inline function viewelement( - A::CellArray{SVector{Ni,T},N,D,T_elem}, i, icell::Vararg{Int,Nc} -) where {Nc,Ni,N,T,T_elem,D} - idx_cell = cart2ind(A.dims, icell...) - return _viewelement(A.data, i, idx_cell) -end - -Base.@propagate_inbounds @inline _viewelement(A::Array, idx, icell) = A[1, idx, icell] -Base.@propagate_inbounds @inline _viewelement(A, idx, icell) = A[icell, idx, 1] - -""" - setelement!(A, x, element_indices..., cell_indices...) - -Store the given value `x` at the given element with `element_indices` of the cell with the indices `cell_indices` - -## Arguments -- `x::Number`: value to be stored in the index `element_indices` of the cell with `cell_indices`. -- `element_indices::Int|NTuple{N,Int}`: the `element_indices` that designate the field in accordance with `A`'s cell type. -- `cell_indices::Int|NTuple{N,Int}`: the `cell_indices` that designate the cell in accordance with `A`'s cell type. -""" -Base.@propagate_inbounds @inline function setelement!( - A::CellArray{SMatrix{Ni,Nj,T,N_array},N,D,T_elem}, x::T, i, j, icell::Vararg{Int,Nc} -) where {Nc,Ni,Nj,N_array,T,N,T_elem,D} - idx_element = cart2ind((Ni, Nj), i, j) - idx_cell = cart2ind(A.dims, icell...) - return _setelement!(A.data, x, idx_element, idx_cell) -end - -Base.@propagate_inbounds @inline function setelement!( - A::CellArray{SVector{Ni,T},N,D,T_elem}, x::T, i, icell::Vararg{Int,Nc} -) where {Nc,Ni,T,N,T_elem,D} - idx_cell = cart2ind(A.dims, icell...) - return _setelement!(A.data, x, i, idx_cell) -end - -Base.@propagate_inbounds @inline function _setelement!(A::Array, x, idx::Int, icell::Int) - return (A[1, idx, icell] = x) -end - -Base.@propagate_inbounds @inline function _setelement!(A, x, idx::Int, icell::Int) - return (A[icell, idx, 1] = x) -end - -## Helper functions - -""" - cart2ind(A) - -Return the linear index of a `n`-dimensional array corresponding to the cartesian indices `I` -""" -@inline function cart2ind(n::NTuple{N1,Int}, I::Vararg{Int,N2}) where {N1,N2} - return LinearIndices(n)[CartesianIndex(I...)] -end -@inline cart2ind(ni::T, nj::T, i::T, j::T) where {T<:Int} = cart2ind((ni, nj), i, j) -@inline function cart2ind(ni::T, nj::T, nk::T, i::T, j::T, k::T) where {T<:Int} - return cart2ind((ni, nj, nk), i, j, k) -end - -## convenience macros - -macro cell(ex) - ex = if ex.head === (:(=)) - _set(ex) - else - _get(ex) - end - return :($(esc(ex))) -end - -@inline _get(ex) = Expr(:call, element, ex.args...) - -@inline function _set(ex) - return Expr( - :call, setelement!, ex.args[1].args[1], ex.args[2], ex.args[1].args[2:end]... - ) -end diff --git a/src/phases/phases.jl b/src/phases/phases.jl index b9912c200..70b5d8511 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -1,26 +1,7 @@ -""" - nphases(x::JustRelax.PhaseRatio) - -Return the number of phases in `x::JustRelax.PhaseRatio`. -""" -@inline nphases(x::JustRelax.PhaseRatio) = nphases(x.center) - -@inline function nphases( - ::CellArray{StaticArraysCore.SArray{Tuple{N},T,N1,N},N2,N3,T_Array} -) where {N,T,N1,N2,N3,T_Array} - return Val(N) -end - -@inline function numphases( - ::CellArray{StaticArraysCore.SArray{Tuple{N},T,N1,N},N2,N3,T_Array} -) where {N,T,N1,N2,N3,T_Array} - return N -end - """ fn_ratio(fn::F, rheology::NTuple{N, AbstractMaterialParamsStruct}, ratio) where {N, F} -Average the function `fn` over the material phases in `rheology` using the phase ratios `ratio`. +Average the function `fn` over the material phases in `rheology` using the phase ratios `ratio`. """ @generated function fn_ratio( fn::F, rheology::NTuple{N,AbstractMaterialParamsStruct}, ratio @@ -47,166 +28,3 @@ end return x end end - -## Kernels to compute phase ratios at the centers - -function phase_ratios_center!(phase_ratios, particles, grid, phases) - return phase_ratios_center!( - backend(phase_ratios), phase_ratios, particles, grid, phases - ) -end - -function phase_ratios_center!( - ::CPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases -) - return _phase_ratios_center!(phase_ratios, particles, grid, phases) -end - -function _phase_ratios_center!( - phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases -) - ni = size(phases) - @parallel (@idx ni) phase_ratios_center_kernel!( - phase_ratios.center, particles.coords, grid.xci, grid.di, phases - ) - return nothing -end - -@parallel_indices (I...) function phase_ratios_center_kernel!( - ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases -) where {N,T1,T2,T3} - - # index corresponding to the cell center - cell_center = ntuple(i -> xci[i][I[i]], Val(N)) - # phase ratios weights (∑w = 1.0) - w = phase_ratio_weights( - getindex.(pxi, I...), phases[I...], cell_center, di, nphases(ratio_centers) - ) - # update phase ratios array - for k in 1:numphases(ratio_centers) - @cell ratio_centers[k, I...] = w[k] - end - - return nothing -end - -## Kernels to compute phase ratios at the vertices - -function phase_ratios_vertex!(phase_ratios, particles, grid, phases) - return phase_ratios_vertex!( - backend(phase_ratios), phase_ratios, particles, grid, phases - ) -end - -function phase_ratios_vertex!( - ::CPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases -) - return _phase_ratios_vertex!(phase_ratios, particles, grid, phases) -end - -function _phase_ratios_vertex!( - phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases -) - ni = size(phases) .+ 1 - @parallel (@idx ni) phase_ratios_vertex_kernel!( - phase_ratios.vertex, particles.coords, grid.xvi, grid.di, phases - ) - return nothing -end - -@parallel_indices (I...) function phase_ratios_vertex_kernel!( - ratio_vertices, pxi::NTuple{3,T1}, xvi::NTuple{3,T2}, di::NTuple{3,T3}, phases -) where {T1,T2,T3} - - # # index corresponding to the cell center - cell_vertex = ntuple(i -> xvi[i][I[i]], Val(3)) - ni = size(phases) - - for offsetᵢ in -1:0, offsetⱼ in -1:0, offsetₖ in -1:0 - offsets = offsetᵢ, offsetⱼ, offsetₖ - cell_index = ntuple(Val(3)) do i - clamp(I[i] + offsets[i], 1, ni[i]) - end - # phase ratios weights (∑w = 1.0) - w = phase_ratio_weights( - getindex.(pxi, cell_index...), - phases[cell_index...], - cell_vertex, - di, - nphases(ratio_vertices), - ) - # update phase ratios array - for k in 1:numphases(ratio_vertices) - @cell ratio_vertices[k, I...] = w[k] - end - end - - return nothing -end - -@parallel_indices (I...) function phase_ratios_vertex_kernel!( - ratio_vertices, pxi::NTuple{2,T1}, xvi::NTuple{2,T2}, di::NTuple{2,T3}, phases -) where {T1,T2,T3} - - # index corresponding to the cell center - cell_vertex = ntuple(i -> xvi[i][I[i]], Val(2)) - ni = size(phases) - - for offsetᵢ in -1:0, offsetⱼ in -1:0 - offsets = offsetᵢ, offsetⱼ - cell_index = ntuple(Val(2)) do i - clamp(I[i] + offsets[i], 1, ni[i]) - end - # phase ratios weights (∑w = 1.0) - w = phase_ratio_weights( - getindex.(pxi, cell_index...), - phases[cell_index...], - cell_vertex, - di, - nphases(ratio_vertices), - ) - # update phase ratios array - for k in 1:numphases(ratio_vertices) - @cell ratio_vertices[k, I...] = w[k] - end - end - - return nothing -end - -## interpolation kernels - -function phase_ratio_weights( - pxi::NTuple{NP,C}, ph::SVector{N1,T}, cell_center, di, ::Val{NC} -) where {N1,NC,NP,T,C} - - # Initiaze phase ratio weights (note: can't use ntuple() here because of the @generated function) - w = ntuple(_ -> zero(T), Val(NC)) - sumw = zero(T) - - for i in eachindex(ph) - # bilinear weight (1-(xᵢ-xc)/dx)*(1-(yᵢ-yc)/dy) - p = getindex.(pxi, i) - isnan(first(p)) && continue - x = @inline bilinear_weight(cell_center, p, di) - sumw += x # reduce - ph_local = ph[i] - # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta - # w = w .+ x .* ntuple(j -> δ(Int(ph_local), j), Val(NC)) - w = w .+ x .* ntuple(j -> (ph_local == j), Val(NC)) - end - w = w .* inv(sumw) - return w -end - -@generated function bilinear_weight( - a::NTuple{N,T}, b::NTuple{N,T}, di::NTuple{N,T} -) where {N,T} - quote - Base.@_inline_meta - val = one($T) - Base.Cartesian.@nexprs $N i -> - @inbounds val *= muladd(-abs(a[i] - b[i]), inv(di[i]), one($T)) - return val - end -end diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 5fe2d6d1b..d3f493e96 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -18,10 +18,10 @@ end """ compute_ρg!(ρg, phase_ratios, rheology, args) -Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. +Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. The `phase_ratios` are used to compute the density of the composite rheology. """ -function compute_ρg!(ρg, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function compute_ρg!(ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args) ni = size(ρg) @parallel (@idx ni) compute_ρg_kernel!(ρg, phase_ratios.center, rheology, args) return nothing @@ -97,11 +97,11 @@ end @inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) = compute_ρg!(ρg, rheology, args) # with phase ratios -@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) = +@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) = update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args) @inline update_ρg!( - ::ConstantDensityTrait, ρg, phase_ratios::JustRelax.PhaseRatio, rheology, args + ::ConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args ) = nothing @inline update_ρg!( - ::NonConstantDensityTrait, ρg, phase_ratios::JustRelax.PhaseRatio, rheology, args + ::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args ) = compute_ρg!(ρg, phase_ratios, rheology, args) diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index c09960a39..c574b2737 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,4 +1,4 @@ -# # Traits +# # Traits # without phase ratios @inline function update_viscosity!( @@ -171,7 +171,7 @@ end function _compute_viscosity!( stokes::JustRelax.StokesArrays, ν, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, args, rheology, cutoff, diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index c455c3359..27b9077b9 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -462,7 +462,7 @@ function _solve!( di::NTuple{2,T}, flow_bcs::AbstractFlowBoundaryConditions, ρg, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, rheology, args, dt, diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 4df97ea78..53778974b 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -371,7 +371,7 @@ function _solve!( di::NTuple{3,T}, flow_bcs::AbstractFlowBoundaryConditions, ρg, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, rheology::NTuple{N,AbstractMaterialParamsStruct}, args, dt, diff --git a/src/stokes/StressRotation.jl b/src/stokes/StressRotation.jl index ea9b45daa..27baa0c7c 100644 --- a/src/stokes/StressRotation.jl +++ b/src/stokes/StressRotation.jl @@ -15,17 +15,17 @@ end cell = i, j for ip in JustRelax.cellaxes(index) - !@cell(index[ip, cell...]) && continue # no particle in this location + !@index(index[ip, cell...]) && continue # no particle in this location - ω_xy = @cell ω[ip, cell...] - τ_xx = @cell xx[ip, cell...] - τ_yy = @cell yy[ip, cell...] - τ_xy = @cell xy[ip, cell...] + ω_xy = @index ω[ip, cell...] + τ_xx = @index xx[ip, cell...] + τ_yy = @index yy[ip, cell...] + τ_xy = @index xy[ip, cell...] tmp = τ_xy * ω_xy * 2.0 - @cell xx[ip, cell...] = fma(dt, cte, τ_xx) - @cell yy[ip, cell...] = fma(dt, cte, τ_yy) - @cell xy[ip, cell...] = fma(dt, (τ_xx - τ_yy) * ω_xy, τ_xy) + @index xx[ip, cell...] = fma(dt, cte, τ_xx) + @index yy[ip, cell...] = fma(dt, cte, τ_yy) + @index xy[ip, cell...] = fma(dt, (τ_xx - τ_yy) * ω_xy, τ_xy) end return nothing @@ -37,14 +37,14 @@ end cell = i, j for ip in JustRelax.cellaxes(index) - !@cell(index[ip, cell...]) && continue # no particle in this location + !@index(index[ip, cell...]) && continue # no particle in this location - θ = dt * @cell ω[ip, cell...] + θ = dt * @index ω[ip, cell...] sinθ, cosθ = sincos(θ) - τ_xx = @cell xx[ip, cell...] - τ_yy = @cell yy[ip, cell...] - τ_xy = @cell xy[ip, cell...] + τ_xx = @index xx[ip, cell...] + τ_yy = @index yy[ip, cell...] + τ_xy = @index xy[ip, cell...] R = @SMatrix [ cosθ -sinθ @@ -59,9 +59,9 @@ end # this could be fully unrolled in 2D τr = R * τ * R' - @cell xx[ip, cell...] = τr[1, 1] - @cell yy[ip, cell...] = τr[2, 2] - @cell xy[ip, cell...] = τr[1, 2] + @index xx[ip, cell...] = τr[1, 1] + @index yy[ip, cell...] = τr[2, 2] + @index xy[ip, cell...] = τr[1, 2] end return nothing diff --git a/src/thermal_diffusion/DiffusionExplicit.jl b/src/thermal_diffusion/DiffusionExplicit.jl index 5241e2359..2fa055fb8 100644 --- a/src/thermal_diffusion/DiffusionExplicit.jl +++ b/src/thermal_diffusion/DiffusionExplicit.jl @@ -401,7 +401,7 @@ function JustRelax.solve!( thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology::NTuple{N,AbstractMaterialParamsStruct}, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, args::NamedTuple, di::NTuple{2,_T}, dt, @@ -874,7 +874,7 @@ function JustRelax.solve!( thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology::NTuple{N,AbstractMaterialParamsStruct}, - phase_ratios::JustRelax.PhaseRatio, + phase_ratios::JustPIC.PhaseRatios, args::NamedTuple, di::NTuple{3,_T}, dt; diff --git a/src/thermal_diffusion/DiffusionPT_GeoParams.jl b/src/thermal_diffusion/DiffusionPT_GeoParams.jl index 9e942f578..f9299cc63 100644 --- a/src/thermal_diffusion/DiffusionPT_GeoParams.jl +++ b/src/thermal_diffusion/DiffusionPT_GeoParams.jl @@ -1,12 +1,12 @@ ## Phases -@inline get_phase(x::JustRelax.PhaseRatio) = x.center +@inline get_phase(x::JustPIC.PhaseRatios) = x.center @inline get_phase(x) = x # update_pt_thermal_arrays!(::Vararg{Any,N}) where {N} = nothing function update_pt_thermal_arrays!( - pt_thermal, phase_ratios::JustRelax.PhaseRatio, rheology, args, _dt + pt_thermal, phase_ratios::JustPIC.PhaseRatios, rheology, args, _dt ) ni = size(phase_ratios.center) @@ -51,7 +51,7 @@ end @inline getindex_phase(::Nothing, I::Vararg{Int,N}) where {N} = nothing -# Diffusivity +# Diffusivity @inline function compute_diffusivity(rheology, args) return compute_conductivity(rheology, args) * diff --git a/src/thermal_diffusion/ShearHeating.jl b/src/thermal_diffusion/ShearHeating.jl index 8991bff2e..d4ad63201 100644 --- a/src/thermal_diffusion/ShearHeating.jl +++ b/src/thermal_diffusion/ShearHeating.jl @@ -26,7 +26,7 @@ end end function compute_shear_heating!( - ::CPUBackendTrait, thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt + ::CPUBackendTrait, thermal, stokes, phase_ratios::JustPIC.PhaseRatios, rheology, dt ) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( diff --git a/src/types/phases.jl b/src/types/phases.jl deleted file mode 100644 index 8e7bd9ddf..000000000 --- a/src/types/phases.jl +++ /dev/null @@ -1,8 +0,0 @@ -struct PhaseRatio{T} - vertex::T - center::T - - PhaseRatio(vertex::T, center::T) where {T<:AbstractArray} = new{T}(vertex, center) -end - -@inline PhaseRatio(::Type{CPUBackend}, ni, num_phases) = PhaseRatio(ni, num_phases) diff --git a/src/types/traits.jl b/src/types/traits.jl index 746717244..3fb32ceb2 100644 --- a/src/types/traits.jl +++ b/src/types/traits.jl @@ -27,7 +27,7 @@ for type in ( end @inline backend(x::JustRelax.StokesArrays) = backend(x.P) -@inline backend(x::JustRelax.PhaseRatio) = backend(x.center.data) +# @inline backend(x::JustPIC.PhaseRatios) = backend(x.center.data) # Error handling @inline backend(::T) where {T} = throw(ArgumentError("$(T) is not a supported backend")) diff --git a/test/Project.toml b/test/Project.toml index 3a35bb387..2a30fb77e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,5 +1,4 @@ [deps] -AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" @@ -20,8 +19,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" [compat] AMDGPU = "0.8, 0.9, 1" -AllocCheck = "0.1" -CUDA = "~5.3.5" +CUDA = "5" CellArrays = "0.2" GeoParams = "0.5, 0.6" JustPIC = "0.3.4, 0.3.5, 0.3.6, 0.4" diff --git a/test/test_Blankenbach.jl b/test/test_Blankenbach.jl index 0fe75b7bd..34dbaca79 100644 --- a/test/test_Blankenbach.jl +++ b/test/test_Blankenbach.jl @@ -36,7 +36,7 @@ else JustPIC.CPUBackend end -import JustRelax.@cell + # Load script dependencies using Printf, LinearAlgebra, CellArrays @@ -111,9 +111,9 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10) # 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 --------------------------------------------- @@ -255,7 +255,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10) # 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/test/test_CellArrays2D.jl b/test/test_CellArrays2D.jl deleted file mode 100644 index 20e344348..000000000 --- a/test/test_CellArrays2D.jl +++ /dev/null @@ -1,36 +0,0 @@ -@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" - using AMDGPU - AMDGPU.allowscalar(true) -elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" - using CUDA - CUDA.allowscalar(true) -end - -using Test, StaticArrays -using JustRelax, JustRelax.JustRelax2D -using ParallelStencil - -@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" - @init_parallel_stencil(AMDGPU, Float64, 2) -elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" - @init_parallel_stencil(CUDA, Float64, 2) -else - @init_parallel_stencil(Threads, Float64, 2) -end - -@testset "CellArrays 2D" begin - ni = 5, 5 - A = @fill(false, ni..., celldims=(2,), eltype=Bool) - - @test cellaxes(A) === Base.OneTo(2) - @test cellnum(A) == 2 - @test new_empty_cell(A) === SA[false, false] - - @test @cell(A[1, 1, 1]) === false - # @test (@allocated @cell A[1, 1, 1]) === 0 - - @cell A[1, 1, 1] = true - @test @cell(A[1, 1, 1]) === true - # @test (@allocated @cell A[1, 1, 1] = true) === 0 - @test A[1, 1] === SA[true, false] -end diff --git a/test/test_CellArrays3D.jl b/test/test_CellArrays3D.jl deleted file mode 100644 index 522794832..000000000 --- a/test/test_CellArrays3D.jl +++ /dev/null @@ -1,40 +0,0 @@ - -using Test, StaticArrays, AllocCheck -using Suppressor - -@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" - using AMDGPU - AMDGPU.allowscalar(true) -elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" - using CUDA - CUDA.allowscalar(true) -end - -using JustRelax, JustRelax.JustRelax3D -using ParallelStencil - -@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" - @init_parallel_stencil(AMDGPU, Float64, 3) -elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" - @init_parallel_stencil(CUDA, Float64, 3) -else - @init_parallel_stencil(Threads, Float64, 3) -end - -@testset "CellArrays 3D" begin - @suppress begin - ni = 5, 5, 5 - A = @fill(false, ni..., celldims=(2,), eltype=Bool) - - @test @cell(A[1, 1, 1, 1]) === false - # @test (@allocated @cell A[1, 1, 1, 1]) == 0 - ; - @cell A[1, 1, 1, 1] = true - @test @cell(A[1, 1, 1, 1]) === true - # @test (@allocated @cell A[1, 1, 1, 1] = true) == 0 - - @test A[1, 1, 1] == SA[true, false] - # allocs = check_allocs(getindex, (typeof(A), Int64, Int64, Int64)) - # @test isempty(allocs) - end -end diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index cc762699e..32dd5b1b7 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -9,7 +9,7 @@ end using Printf, LinearAlgebra, GeoParams, CellArrays using JustRelax, JustRelax.JustRelax2D -import JustRelax.@cell + using ParallelStencil const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" @@ -44,18 +44,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 @@ -112,9 +112,9 @@ function VanKeken2D(ny=32, nx=32) # 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 @@ -191,7 +191,7 @@ function VanKeken2D(ny=32, nx=32) # 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/test/test_WENO5.jl b/test/test_WENO5.jl index 96fa9cf8d..a0fd32b9d 100644 --- a/test/test_WENO5.jl +++ b/test/test_WENO5.jl @@ -23,7 +23,7 @@ else CPUBackend end -import JustRelax.@cell + using GeoParams, SpecialFunctions # function to compute strain rate (compulsory) diff --git a/test/test_diffusion2D.jl b/test/test_diffusion2D.jl index e93541883..bd0ef0b77 100644 --- a/test/test_diffusion2D.jl +++ b/test/test_diffusion2D.jl @@ -21,7 +21,7 @@ else CPUBackend end -import JustRelax.@cell + using GeoParams # HELPER FUNCTIONS --------------------------------------------------------------- diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl index cc39b725f..9ff440cc4 100644 --- a/test/test_diffusion2D_multiphase.jl +++ b/test/test_diffusion2D_multiphase.jl @@ -34,7 +34,7 @@ else JustPIC.CPUBackend end -import JustRelax.@cell + distance(p1, p2) = mapreduce(x->(x[1]-x[2])^2, +, zip(p1, p2)) |> sqrt @@ -66,19 +66,19 @@ function init_phases!(phases, particles, xc, yc, r) center = xc, yc @parallel_indices (i, j) function init_phases!(phases, px, py, index, center, r) - @inbounds for ip in JustRelax.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 (((x - center[1] ))^2 + ((y - center[2]))^2) ≤ r^2 - 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 @@ -156,8 +156,8 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) # temperature pPhases, = init_cell_arrays(particles, Val(1)) init_phases!(pPhases, particles, center_perturbation..., r) - 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) # ---------------------------------------------------- @parallel (@idx ni) compute_temperature_source_terms!(thermal.H, rheology, phase_ratios.center, args) diff --git a/test/test_diffusion3D_multiphase.jl b/test/test_diffusion3D_multiphase.jl index 13bfcd8d4..2e3dda984 100644 --- a/test/test_diffusion3D_multiphase.jl +++ b/test/test_diffusion3D_multiphase.jl @@ -38,7 +38,7 @@ else JustPIC.CPUBackend end -import JustRelax.@cell + @parallel_indices (i, j, k) function init_T!(T, z) if z[k] == maximum(z) @@ -68,20 +68,20 @@ function init_phases!(phases, particles, xc, yc, zc, r) center = xc, yc, zc @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, center, r) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in cellaxes(phases) # quick escape - @cell(index[ip, I...]) == 0 && continue + @index(index[ip, I...]) == 0 && continue - x = @cell px[ip, I...] - y = @cell py[ip, I...] - z = @cell pz[ip, I...] + x = @index px[ip, I...] + y = @index py[ip, I...] + z = @index pz[ip, I...] # plume - rectangular if (((x - center[1]))^2 + ((y - center[2]))^2 + ((z - center[3]))^2) ≤ r^2 - @cell phases[ip, I...] = 2.0 + @index phases[ip, I...] = 2.0 else - @cell phases[ip, I...] = 1.0 + @index phases[ip, I...] = 1.0 end end return nothing @@ -170,9 +170,9 @@ function diffusion_3D(; ) # 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, center_perturbation..., r) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # PT coefficients for thermal diffusion diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index c23d512b7..c02a48399 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -11,7 +11,7 @@ using GeoParams using JustRelax, JustRelax.JustRelax2D using ParallelStencil -const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" +const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" @init_parallel_stencil(AMDGPU, Float64, 2) AMDGPUBackend elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" @@ -22,6 +22,16 @@ else CPUBackend end +using JustPIC, JustPIC._2D + +const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + JustPIC.AMDGPUBackend +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" + CUDABackend +else + JustPIC.CPUBackend +end + # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -33,12 +43,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 @@ -105,12 +115,12 @@ function ShearBand2D() # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend, ni) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces @@ -126,8 +136,8 @@ function ShearBand2D() free_slip = (left = true, right = true, top = true, bot = true), no_slip = (left = false, right = false, top = false, bot=false), ) - stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray(backend)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) + stokes.V.Vx .= PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend_JR)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(@velocity(stokes)...) diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl index 429ea3dd9..1cdb0669e 100644 --- a/test/test_shearband2D_softening.jl +++ b/test/test_shearband2D_softening.jl @@ -12,7 +12,7 @@ using GeoParams, CellArrays using JustRelax, JustRelax.JustRelax2D using ParallelStencil -const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" +const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" @init_parallel_stencil(AMDGPU, Float64, 2) AMDGPUBackend elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" @@ -23,6 +23,16 @@ else CPUBackend end +using JustPIC, JustPIC._2D + +const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + JustPIC.AMDGPUBackend +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" + CUDABackend +else + JustPIC.CPUBackend +end + # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -35,12 +45,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 @@ -111,12 +121,12 @@ function ShearBand2D() # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend, ni) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces @@ -133,8 +143,8 @@ function ShearBand2D() free_slip = (left = true, right = true, top = true, bot = true), no_slip = (left = false, right = false, top = false, bot=false), ) - stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray(backend)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) + stokes.V.Vx .= PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend_JR)([-y*εbg for _ in 1:nx+2, y in xvi[2]]) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(@velocity(stokes)...) diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index d3d6c6f0d..4b7dab0c3 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -99,9 +99,9 @@ function Shearheating2D(igg; nx=32, ny=32) xc_anomaly = lx / 2 # origin of thermal anomaly yc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -227,14 +227,14 @@ function Shearheating2D(igg; nx=32, ny=32) # 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) @show it += 1 t += dt end - finalize_global_grid(; finalize_MPI=true) + # finalize_global_grid(; finalize_MPI=true) return iters, thermal end @@ -267,5 +267,6 @@ end # Ensure iters is defined before running the test @test iters != nothing && iters.err_evo1[end] < 1e-4 + @test any(x -> x < 0, thermal.shear_heating) == false end end diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index 0b56dd5ee..a60aff681 100644 --- a/test/test_shearheating3D.jl +++ b/test/test_shearheating3D.jl @@ -37,7 +37,7 @@ elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" else JustPIC.CPUBackend end -import JustRelax.@cell + # Load script dependencies using Printf, GeoParams @@ -95,9 +95,9 @@ function Shearheating3D(igg; nx=16, ny=16, nz=16) yc_anomaly = ly/2 # origin of thermal anomaly zc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -218,7 +218,7 @@ function Shearheating3D(igg; nx=16, ny=16, nz=16) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (pT,), (thermal.T,), 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 @@ -259,5 +259,6 @@ end # Ensure iters is defined before running the test @test iters != nothing && iters.err_evo1[end] < 1e-4 + @test any(x -> x < 0, thermal.shear_heating) == false end end diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index e381db90f..47faad54d 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -59,14 +59,14 @@ function init_phases!(phases, particles, xc, yc, r) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, xc, yc, r) - @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]) + x = @index px[ip, i, j] + depth = -(@index py[ip, i, j]) # plume - rectangular - JustRelax.@cell phases[ip, i, j] = if ((x -xc)^2 ≤ r^2) && ((depth - yc)^2 ≤ r^2) + @index phases[ip, i, j] = if ((x -xc)^2 ≤ r^2) && ((depth - yc)^2 ≤ r^2) 2.0 else 1.0 @@ -145,9 +145,9 @@ function Sinking_Block2D() xc_anomaly = 250e3 # origin of thermal anomaly yc_anomaly = -(ly-400e3) # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatios(backend, length(rheology), ni) init_phases!(pPhases, particles, xc_anomaly, abs(yc_anomaly), r_anomaly) - phase_ratios_center!(phase_ratios, particles, grid, pPhases) + phase_ratios_center!(phase_ratios, particles, xci, pPhases) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem diff --git a/test/test_thermalstresses.jl b/test/test_thermalstresses.jl index 4017b3e89..71cec6184 100644 --- a/test/test_thermalstresses.jl +++ b/test/test_thermalstresses.jl @@ -36,7 +36,6 @@ else JustPIC.CPUBackend end -import JustRelax.@cell # Load script dependencies using Printf, Statistics, LinearAlgebra, CellArrays, StaticArrays @@ -70,23 +69,23 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, r_anomaly, stic @parallel_indices (i, j) function init_phases!( phases, px, py, index, xc_anomaly, yc_anomaly, r_anomaly, sticky_air, top, bottom ) - @inbounds for ip in JustRelax.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]) - sticky_air + x = @index px[ip, i, j] + y = -(@index py[ip, i, j]) - sticky_air if top ≤ y ≤ bottom - @cell phases[ip, i, j] = 1.0 # crust + @index phases[ip, i, j] = 1.0 # crust end # thermal anomaly - circular if ((x - xc_anomaly)^2 + (y + yc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, i, j] = 2.0 + @index phases[ip, i, j] = 2.0 end if y < top - @cell phases[ip, i, j] = 3.0 + @index phases[ip, i, j] = 3.0 end end return nothing @@ -267,8 +266,8 @@ function main2D(; nx=32, ny=32) r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) init_phases!(pPhases, particles, x_anomaly, y_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), nondimensionalize(20km,CharDim)) - 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) # Initialisation of thermal profile thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions @@ -427,7 +426,7 @@ function main2D(; nx=32, ny=32) # 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) particle2grid!(T_buffer, pT, xvi, particles) @views T_buffer[:, end] .= Tsurf