Add Displacement struct inside StokesArrays (#181)

* add displacement struct

* ext: displacements

* rework and test displacement

* format

* remove line

* fixup

src/common.jl

@@ -28,6 +28,9 @@ export @allocate,
     multi_copy!,
     take
 
+include("types/displacement.jl")
+export velocity2displacement!, displacement2velocity!
+
 include("boundaryconditions/BoundaryConditions.jl")
 export FlowBoundaryConditions,
     TemperatureBoundaryConditions, flow_bcs!, thermal_bcs!, pureshear_bc!, apply_free_slip!

src/ext/AMDGPU/2D.jl

@@ -34,6 +34,18 @@ function JR2D.StokesArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {
     return StokesArrays(ni)
 end
 
+function JR2D.velocity2displacement!(
+    stokes::JustRelax.StokesArrays, ::AMDGPUBackendTrait, dt
+)
+    return _velocity2displacement!(stokes, dt)
+end
+
+function JR2D.displacement2velocity!(
+    stokes::JustRelax.StokesArrays, ::AMDGPUBackendTrait, dt
+)
+    return _displacement2velocity!(stokes, dt)
+end
+
 function JR2D.ThermalArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Number}) where {N}
     return ThermalArrays(ni...)
 end

src/ext/AMDGPU/3D.jl

@@ -34,6 +34,18 @@ function JR3D.StokesArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {
     return StokesArrays(ni)
 end
 
+function JR3D.velocity2displacement!(
+    stokes::JustRelax.StokesArrays, ::AMDGPUBackendTrait, dt
+)
+    return _velocity2displacement!(stokes, dt)
+end
+
+function JR3D.displacement2velocity!(
+    stokes::JustRelax.StokesArrays, ::AMDGPUBackendTrait, dt
+)
+    return _displacement2velocity!(stokes, dt)
+end
+
 function JR3D.ThermalArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Number}) where {N}
     return ThermalArrays(ni...)
 end

src/ext/CUDA/2D.jl

@@ -33,6 +33,14 @@ function JR2D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N}
     return StokesArrays(ni)
 end
 
+function JR2D.velocity2displacement!(stokes::JustRelax.StokesArrays, ::CUDABackendTrait, dt)
+    return _velocity2displacement!(stokes, dt)
+end
+
+function JR2D.displacement2velocity!(stokes::JustRelax.StokesArrays, ::CUDABackendTrait, dt)
+    return _displacement2velocity!(stokes, dt)
+end
+
 function JR2D.ThermalArrays(::Type{CUDABackend}, ni::NTuple{N,Number}) where {N}
     return ThermalArrays(ni...)
 end
@@ -109,6 +117,7 @@ function JR2D.phase_ratios_center!(
 end
 
 # Rheology
+
 ## viscosity
 function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff)
     return _compute_viscosity!(stokes, ν, args, rheology, cutoff)
@@ -147,6 +156,7 @@ end
 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)
     return compute_ρg!(ρg, phase_ratios, rheology, args)
 end

src/ext/CUDA/3D.jl

@@ -33,6 +33,14 @@ function JR3D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N}
     return StokesArrays(ni)
 end
 
+function JR3D.velocity2displacement!(stokes::JustRelax.StokesArrays, ::CUDABackendTrait, dt)
+    return _velocity2displacement!(stokes, dt)
+end
+
+function JR3D.displacement2velocity!(stokes::JustRelax.StokesArrays, ::CUDABackendTrait, dt)
+    return _displacement2velocity!(stokes, dt)
+end
+
 function JR3D.ThermalArrays(::Type{CUDABackend}, ni::NTuple{N,Number}) where {N}
     return ThermalArrays(ni...)
 end
@@ -80,6 +88,17 @@ function JR3D.PTThermalCoeffs(
     return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL)
 end
 
+function JR3D.update_pt_thermal_arrays!(
+    pt_thermal::JustRelax.PTThermalCoeffs{T,<:CuArray},
+    phase_ratios::JustRelax.PhaseRatio,
+    rheology,
+    args,
+    _dt,
+) where {T}
+    update_pt_thermal_arrays!(pt_thermal, phase_ratios, rheology, args, _dt)
+    return nothing
+end
+
 # Boundary conditions
 function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs)
     return _flow_bcs!(bcs, @velocity(stokes))

src/types/constructors/stokes.jl

@@ -17,6 +17,25 @@ function Velocity(nx::Integer, ny::Integer, nz::Integer)
     return JustRelax.Velocity(Vx, Vy, Vz)
 end
 
+## Displacement type
+
+function Displacement(nx::Integer, ny::Integer)
+    nUx = (nx + 1, ny + 2)
+    nUy = (nx + 2, ny + 1)
+
+    Ux, Uy = @zeros(nUx...), @zeros(nUy)
+    return JustRelax.Displacement(Ux, Uy, nothing)
+end
+
+function Displacement(nx::Integer, ny::Integer, nz::Integer)
+    nUx = (nx + 1, ny + 2, nz + 2)
+    nUy = (nx + 2, ny + 1, nz + 2)
+    nUz = (nx + 2, ny + 2, nz + 1)
+
+    Ux, Uy, Uz = @zeros(nUx...), @zeros(nUy), @zeros(nUz)
+    return JustRelax.Displacement(Ux, Uy, Uz)
+end
+
 ## Viscosity type
 
 function Viscosity(ni::NTuple{N,Integer}) where {N}
@@ -80,6 +99,7 @@ function StokesArrays(ni::NTuple{N,Integer}) where {N}
     P0 = @zeros(ni...)
     ∇V = @zeros(ni...)
     V = Velocity(ni...)
+    U = Displacement(ni...)
     τ = SymmetricTensor(ni...)
     τ_o = SymmetricTensor(ni...)
     ε = SymmetricTensor(ni...)
@@ -88,5 +108,5 @@ function StokesArrays(ni::NTuple{N,Integer}) where {N}
     viscosity = Viscosity(ni)
     R = Residual(ni...)
 
-    return JustRelax.StokesArrays(P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R)
+    return JustRelax.StokesArrays(P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R, U)
 end

src/types/displacement.jl

@@ -0,0 +1,59 @@
+function velocity2displacement!(stokes::JustRelax.StokesArrays, dt)
+    velocity2displacement!(stokes, backend(stokes), dt)
+    return nothing
+end
+
+function velocity2displacement!(stokes::JustRelax.StokesArrays, ::CPUBackendTrait, dt)
+    return _velocity2displacement!(stokes, dt)
+end
+
+function _velocity2displacement!(stokes::JustRelax.StokesArrays, dt)
+    ni = size(stokes.P)
+    (; V, U) = stokes
+    @parallel (@idx ni .+ 2) _velocity2displacement!(
+        V.Vx, V.Vy, V.Vz, U.Ux, U.Uy, U.Uz, 1 / dt
+    )
+    return nothing
+end
+
+@parallel_indices (I...) function _velocity2displacement!(Vx, Vy, Vz, Ux, Uy, Uz, _dt)
+    if all(I .≤ size(Ux))
+        Ux[I...] = Vx[I...] * _dt
+    end
+    if all(I .≤ size(Uy))
+        Uy[I...] = Vy[I...] * _dt
+    end
+    if !isnothing(Vz) && all(I .≤ size(Uz))
+        Uz[I...] = Vz[I...] * _dt
+    end
+    return nothing
+end
+
+function displacement2velocity!(stokes::JustRelax.StokesArrays, dt)
+    displacement2velocity!(stokes, backend(stokes), dt)
+    return nothing
+end
+
+function displacement2velocity!(stokes::JustRelax.StokesArrays, ::CPUBackendTrait, dt)
+    return _displacement2velocity!(stokes, dt)
+end
+
+function _displacement2velocity!(stokes::JustRelax.StokesArrays, dt)
+    ni = size(stokes.P)
+    (; V, U) = stokes
+    @parallel (@idx ni .+ 2) _displacement2velocity!(U.Ux, U.Uy, U.Uz, V.Vx, V.Vy, V.Vz, dt)
+    return nothing
+end
+
+@parallel_indices (I...) function _displacement2velocity!(Ux, Uy, Uz, Vx, Vy, Vz, dt)
+    if all(I .≤ size(Ux))
+        Vx[I...] = Ux[I...] * dt
+    end
+    if all(I .≤ size(Uy))
+        Vy[I...] = Uy[I...] * dt
+    end
+    if !isnothing(Vz) && all(I .≤ size(Uz))
+        Vz[I...] = Uz[I...] * dt
+    end
+    return nothing
+end

src/types/stokes.jl

@@ -18,6 +18,26 @@ function Velocity(::Number, ::Number, ::Number)
     throw(ArgumentError("Velocity dimensions must be given as integers"))
 end
 
+## Displacement type
+
+struct Displacement{T}
+    Ux::T
+    Uy::T
+    Uz::Union{T,Nothing}
+
+    Displacement(Ux::T, Uy::T, Uz::Union{T,Nothing}) where {T} = new{T}(Ux, Uy, Uz)
+end
+
+Displacement(Ux::T, Uy::T) where {T} = Displacement(Ux, Uy, nothing)
+
+Displacement(ni::NTuple{N,Number}) where {N} = Displacement(ni...)
+function Displacement(::Number, ::Number)
+    throw(ArgumentError("Displacement dimensions must be given as integers"))
+end
+function Displacement(::Number, ::Number, ::Number)
+    throw(ArgumentError("Displacement dimensions must be given as integers"))
+end
+
 ## Viscosity type
 
 struct Viscosity{T}
@@ -101,7 +121,7 @@ end
 
 ## StokesArrays type
 
-struct StokesArrays{A,B,C,D,T}
+struct StokesArrays{A,B,C,D,E,T}
     P::T
     P0::T
     V::A
@@ -113,6 +133,7 @@ struct StokesArrays{A,B,C,D,T}
     viscosity::D
     τ_o::Union{B,Nothing}
     R::C
+    U::E
 end
 
 function StokesArrays(::Type{CPUBackend}, ni::Vararg{Integer,N}) where {N}

src/types/traits.jl

@@ -14,6 +14,7 @@ struct AMDGPUBackendTrait <: GPUBackendTrait end
 
 # Custom struct's
 @inline backend(::JustRelax.Velocity{T}) where {T} = backend(T)
+@inline backend(::JustRelax.Displacement{T}) where {T} = backend(T)
 @inline backend(::JustRelax.SymmetricTensor{T}) where {T} = backend(T)
 @inline backend(::JustRelax.Residual{T}) where {T} = backend(T)
 @inline backend(::JustRelax.Viscosity{T}) where {T} = backend(T)

src/types/type_conversions.jl

@@ -1,6 +1,8 @@
 import Base: Array, copy
 
-const JR_T = Union{StokesArrays,SymmetricTensor,ThermalArrays,Velocity,Residual,Viscosity}
+const JR_T = Union{
+    StokesArrays,SymmetricTensor,ThermalArrays,Velocity,Displacement,Residual,Viscosity
+}
 
 ## Conversion of structs to CPU
 

test/test_types.jl

@@ -34,6 +34,7 @@ const BackendArray = PTArray(backend)
     @test typeof(stokes.P0)     <: BackendArray
     @test typeof(stokes.∇V)     <: BackendArray
     @test stokes.V              isa JustRelax.Velocity
+    @test stokes.U              isa JustRelax.Displacement
     @test stokes.τ              isa JustRelax.SymmetricTensor
     @test stokes.τ_o            isa JustRelax.SymmetricTensor
     @test stokes.ε              isa JustRelax.SymmetricTensor
@@ -62,7 +63,6 @@ const BackendArray = PTArray(backend)
     @test typeof(visc.ητ)    <: BackendArray
     @test_throws MethodError JR2.Viscosity(10.0, 10.0)
 
-    v      = stokes.V
     tensor = stokes.τ
 
     @test size(tensor.xx)   == (nx, ny)
@@ -77,11 +77,23 @@ const BackendArray = PTArray(backend)
     @test typeof(tensor.xy_c) <: BackendArray
     @test typeof(tensor.II)   <: BackendArray
 
-
-
     @test_throws MethodError JR2.StokesArrays(backend, 10.0, 10.0)
 end
 
+@testset "2D Displacement" begin
+    ni = nx, ny = (2, 2)
+    stokes = JR2.StokesArrays(backend, ni)
+
+    stokes.V.Vx .= 1.0
+    stokes.V.Vy .= 1.0
+
+    JR2.velocity2displacement!(stokes, 10)
+    @test all(stokes.U.Ux.==0.1)
+
+    JR2.displacement2velocity!(stokes, 5)
+    @test all(stokes.V.Vx.==0.5)
+end
+
 @testset "3D allocators" begin
     ni = nx, ny, nz = (2, 2, 2)
 
@@ -96,6 +108,7 @@ end
     @test typeof(stokes.P0)     <: BackendArray
     @test typeof(stokes.∇V)     <: BackendArray
     @test stokes.V              isa JustRelax.Velocity
+    @test stokes.U              isa JustRelax.Displacement
     @test stokes.τ              isa JustRelax.SymmetricTensor
     @test stokes.τ_o            isa JustRelax.SymmetricTensor
     @test stokes.ε              isa JustRelax.SymmetricTensor
@@ -125,7 +138,6 @@ end
     @test typeof(visc.ητ)    <: BackendArray
     @test_throws MethodError JR3.Viscosity(1.0, 1.0, 1.0)
 
-    v      = stokes.V
     tensor = stokes.τ
 
     @test size(tensor.xx)   == ni
@@ -150,3 +162,18 @@ end
 
     @test_throws MethodError JR3.StokesArrays(backend, 10.0, 10.0)
 end
+
+@testset "3D Displacement" begin
+    ni = nx, ny, nz = (2, 2, 2)
+    stokes = JR3.StokesArrays(backend, ni)
+
+    stokes.V.Vx .= 1.0
+    stokes.V.Vy .= 1.0
+    stokes.V.Vz .= 1.0
+
+    JR3.velocity2displacement!(stokes, 10)
+    @test all(stokes.U.Ux.==0.1)
+
+    JR3.displacement2velocity!(stokes, 5)
+    @test all(stokes.V.Vx.==0.5)
+end