Compatibility generic gravity fields from GeoParams (#275)

* add generic gravity field * bump GP * update Project.toml * format * bump GP
PTsolvers · Nov 26, 2024 · 8f2b8d6 · 8f2b8d6
1 parent ff9abba
commit 8f2b8d6
Show file tree

Hide file tree

Showing 9 changed files with 101 additions and 23 deletions.
diff --git a/Project.toml b/Project.toml
@@ -34,7 +34,7 @@ AMDGPU = "0.8, 0.9, 1"
 Adapt = "4"
 CUDA = "5"
 CellArrays = "0.2"
-GeoParams = "0.6.4"
+GeoParams = "0.6.7"
 HDF5 = "0.17.1"
 ImplicitGlobalGrid = "0.15"
 JLD2 = "0.4, 0.5"

diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -215,10 +215,15 @@ function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTens
 end
 
 ## Buoyancy forces
-function JR2D.compute_ρg!(ρg::ROCArray, rheology, args)
+function JR2D.compute_ρg!(ρg::Union{ROCArray,NTuple{N,ROCArray}}, rheology, args) where {N}
     return compute_ρg!(ρg, rheology, args)
 end
-function JR2D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args)
+function JR2D.compute_ρg!(
+    ρg::Union{ROCArray,NTuple{N,ROCArray}},
+    phase_ratios::JustPIC.PhaseRatios,
+    rheology,
+    args,
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
 

diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -215,15 +215,22 @@ function JR3D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTens
 end
 
 ## Buoyancy forces
-function JR3D.compute_ρg!(ρg::ROCArray, rheology, args)
+function JR3D.compute_ρg!(ρg::Union{ROCArray,NTuple{N,ROCArray}}, rheology, args) where {N}
     return compute_ρg!(ρg, rheology, args)
 end
 
-function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios::JustPIC.PhaseRatios, rheology, args)
+function JR3D.compute_ρg!(
+    ρg::Union{ROCArray,NTuple{N,ROCArray}},
+    phase_ratios::JustPIC.PhaseRatios,
+    rheology,
+    args,
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
 
-function JR3D.compute_ρg!(ρg::ROCArray, phase_ratios, rheology, args)
+function JR3D.compute_ρg!(
+    ρg::Union{ROCArray,NTuple{N,ROCArray}}, phase_ratios, rheology, args
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
 

diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -202,11 +202,13 @@ function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor
 end
 
 ## Buoyancy forces
-function JR2D.compute_ρg!(ρg::CuArray, rheology, args)
+function JR2D.compute_ρg!(ρg::Union{CuArray,NTuple{N,CuArray}}, rheology, args) where {N}
     return compute_ρg!(ρg, rheology, args)
 end
 
-function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args)
+function JR2D.compute_ρg!(
+    ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios::JustPIC.PhaseRatios, rheology, args
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
 

diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -219,13 +219,17 @@ function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor
 end
 
 ## Buoyancy forces
-function JR3D.compute_ρg!(ρg::CuArray, rheology, args)
+function JR3D.compute_ρg!(ρg::Union{CuArray,NTuple{N,CuArray}}, rheology, args) where {N}
     return compute_ρg!(ρg, rheology, args)
 end
-function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustPIC.PhaseRatios, rheology, args)
+function JR3D.compute_ρg!(
+    ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios::JustPIC.PhaseRatios, rheology, args
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
-function JR3D.compute_ρg!(ρg::CuArray, phase_ratios, rheology, args)
+function JR3D.compute_ρg!(
+    ρg::Union{CuArray,NTuple{N,CuArray}}, phase_ratios, rheology, args
+) where {N}
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end
 

diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl
@@ -4,7 +4,10 @@
 Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`.
 """
 function compute_ρg!(ρg, rheology, args)
-    ni = size(ρg)
+    _size(x::AbstractArray) = size(x)
+    _size(x::NTuple) = size(x[1])
+
+    ni = _size(ρg)
     @parallel (@idx ni) compute_ρg_kernel!(ρg, rheology, args)
     return nothing
 end
@@ -15,14 +18,27 @@ end
     return nothing
 end
 
+@parallel_indices (I...) function compute_ρg_kernel!(
+    ρg::NTuple{N,AbstractArray}, rheology, args
+) where {N}
+    args_ijk = ntuple_idx(args, I...)
+    gᵢ = compute_gravity(first(rheology))
+    ρgᵢ = compute_buoyancies(rheology, args_ijk, gᵢ, Val(N))
+    fill_density!(ρg, ρgᵢ, I...)
+    return nothing
+end
+
 """
     compute_ρg!(ρg, phase_ratios, rheology, 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::JustPIC.PhaseRatios, rheology, args)
-    ni = size(ρg)
+    _size(x::AbstractArray) = size(x)
+    _size(x::NTuple) = size(x[1])
+
+    ni = _size(ρg)
     @parallel (@idx ni) compute_ρg_kernel!(ρg, phase_ratios.center, rheology, args)
     return nothing
 end
@@ -33,7 +49,50 @@ end
     return nothing
 end
 
+@parallel_indices (I...) function compute_ρg_kernel!(
+    ρg::NTuple{N,AbstractArray}, phase_ratios, rheology, args
+) where {N}
+    args_ijk = ntuple_idx(args, I...)
+    gᵢ = compute_gravity(first(rheology))
+    ρgᵢ = compute_buoyancies(rheology, @cell(phase_ratios[I...]), args_ijk, gᵢ, Val(N))
+    fill_density!(ρg, ρgᵢ, I...)
+    return nothing
+end
+
 ## Inner buoyancy force kernels
+@generated function fill_density!(ρg::NTuple{N}, ρgᵢ::NTuple{N}, I::Vararg{Int,N}) where {N}
+    quote
+        Base.@nexprs $N i -> ρg[i][I...] = ρgᵢ[i]
+        return nothing
+    end
+end
+
+@inline fill_density!(ρg::NTuple{N}, ρgᵢ::Number, I::Vararg{Int,N}) where {N} =
+    setindex!(last(ρg), ρgᵢ, I...)
+
+@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::NTuple{3}, ::Val{2}) =
+    compute_density(rheology, args) .* (gᵢ[1], gᵢ[3])
+@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::NTuple{3}, ::Val{3}) =
+    compute_density(rheology, args) .* gᵢ
+@inline compute_buoyancies(rheology::MaterialParams, args, gᵢ::Number, ::Any) =
+    compute_density(rheology, args) * gᵢ
+
+@inline compute_buoyancies(
+    rheology::MaterialParams, phase_ratios, args, gᵢ::NTuple{3}, ::Val{2}
+) = compute_density_ratio(phase_ratios, rheology, args) .* (gᵢ[1], gᵢ[3])
+@inline compute_buoyancies(
+    rheology::MaterialParams, phase_ratios, args, gᵢ::NTuple{3}, ::Val{3}
+) = compute_density_ratio(phase_ratios, rheology, args) .* gᵢ
+@inline compute_buoyancies(
+    rheology::MaterialParams, phase_ratios, args, gᵢ::Number, ::Any
+) = compute_density_ratio(phase_ratios, rheology, args) * gᵢ
+
+@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::NTuple{3}, ::Val{2}) =
+    fn_ratio(compute_density, rheology, phase_ratios, args) .* (gᵢ[1], gᵢ[3])
+@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::NTuple{3}, ::Val{3}) =
+    fn_ratio(compute_density, rheology, phase_ratios, args) .* gᵢ
+@inline compute_buoyancies(rheology, phase_ratios, args, gᵢ::Number, ::Any) =
+    fn_ratio(compute_density, rheology, phase_ratios, args) * gᵢ
 
 """
     compute_buoyancy(rheology::MaterialParams, args)
@@ -91,14 +150,15 @@ Compute the buoyancy forces based on the given rheology, arguments, and phase ra
 end
 
 # without phase ratios
-@inline update_ρg!(ρg::AbstractArray, rheology, args) =
+@inline update_ρg!(ρg::Union{NTuple,AbstractArray}, rheology, args) =
     update_ρg!(isconstant(rheology), ρg, rheology, args)
 @inline update_ρg!(::ConstantDensityTrait, ρg, rheology, args) = nothing
 @inline update_ρg!(::NonConstantDensityTrait, ρg, rheology, args) =
     compute_ρg!(ρg, rheology, args)
 # with phase ratios
-@inline update_ρg!(ρg::AbstractArray, phase_ratios::JustPIC.PhaseRatios, rheology, args) =
-    update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args)
+@inline update_ρg!(
+    ρg::Union{NTuple,AbstractArray}, phase_ratios::JustPIC.PhaseRatios, rheology, args
+) = update_ρg!(isconstant(rheology), ρg, phase_ratios, rheology, args)
 @inline update_ρg!(
     ::ConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args
 ) = nothing

diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl
@@ -527,7 +527,7 @@ function _solve!(
     Vx_on_Vy = @zeros(size(stokes.V.Vy))
 
     # compute buoyancy forces and viscosity
-    compute_ρg!(ρg[end], phase_ratios, rheology, args)
+    compute_ρg!(ρg, phase_ratios, rheology, args)
     compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
     displacement2velocity!(stokes, dt, flow_bcs)
 
@@ -557,7 +557,7 @@ function _solve!(
                 args,
             )
 
-            update_ρg!(ρg[2], phase_ratios, rheology, args)
+            update_ρg!(ρg, phase_ratios, rheology, args)
 
             @parallel (@idx ni .+ 1) compute_strain_rate!(
                 @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di...

diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl
@@ -208,7 +208,7 @@ function _solve!(
     θ = @zeros(ni...)
 
     # compute buoyancy forces and viscosity
-    compute_ρg!(ρg[end], phase_ratios, rheology, args)
+    compute_ρg!(ρg, phase_ratios, rheology, args)
     compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
 
     # convert displacement to velocity
@@ -241,7 +241,7 @@ function _solve!(
             )
 
             # Update buoyancy
-            update_ρg!(ρg[3], rheology, args)
+            update_ρg!(ρg, rheology, args)
 
             update_viscosity!(
                 stokes,
@@ -422,7 +422,7 @@ function _solve!(
     ητ = deepcopy(η)
 
     # compute buoyancy forces and viscosity
-    compute_ρg!(ρg[end], phase_ratios, rheology, args)
+    compute_ρg!(ρg, phase_ratios, rheology, args)
     compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
 
     # convert displacement to velocity
@@ -457,7 +457,7 @@ function _solve!(
             )
 
             # Update buoyancy
-            update_ρg!(ρg[end], phase_ratios, rheology, args)
+            update_ρg!(ρg, phase_ratios, rheology, args)
 
             # Update viscosity
             update_viscosity!(

diff --git a/test/Project.toml b/test/Project.toml
@@ -21,7 +21,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 AMDGPU = "0.8, 0.9, 1"
 CUDA = "5"
 CellArrays = "0.2"
-GeoParams = "0.5, 0.6"
+GeoParams = "0.6.7"
 JustPIC = "0.5.3"
 MPI = "0.20"
 ParallelStencil = "0.11, 0.12, 0.13"