diff --git a/src/common.jl b/src/common.jl index ea23836d0..e3c9318a4 100644 --- a/src/common.jl +++ b/src/common.jl @@ -46,7 +46,7 @@ export AbstractBoundaryConditions, include("MiniKernels.jl") include("phases/phases.jl") -export fn_ratio, phase_ratios_center!, numphases, nphases +export fn_ratio, phase_ratios_center!, phase_ratios_vertex!, numphases, nphases include("rheology/BuoyancyForces.jl") export compute_ρg! diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 57902dd90..fab57998c 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -182,6 +182,16 @@ function JR2D.phase_ratios_center!( 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 diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 4e4559b63..073d3d09a 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -182,6 +182,16 @@ function JR3D.phase_ratios_center!( 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 diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 551bfc8dc..2e4df75ee 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -177,6 +177,16 @@ function JR2D.phase_ratios_center!( 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 diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index e13239017..9a99d1e25 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -189,6 +189,16 @@ function JR3D.phase_ratios_center!( 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 + # Rheology ## viscosity diff --git a/src/phases/phases.jl b/src/phases/phases.jl index 81dea637c..4334e332d 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -48,6 +48,8 @@ end 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 @@ -88,6 +90,92 @@ 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)) + nx, ny = size(phases) + + for offsetᵢ in -1:0, offsetⱼ in -1:0 + offsets = offsetᵢ, 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}