From 8f4600a24b9da8b4262967825323f9b185aa179b Mon Sep 17 00:00:00 2001
From: Albert de Montserrat
 <58044444+albert-de-montserrat@users.noreply.github.com>
Date: Mon, 17 Jun 2024 21:59:21 +0200
Subject: [PATCH] Add `Displacement` struct inside `StokesArrays` (#181)

* add displacement struct

* ext: displacements

* rework and test displacement

* format

* remove line

* fixup
---
 src/common.jl                    |  3 ++
 src/ext/AMDGPU/2D.jl             | 12 +++++++
 src/ext/AMDGPU/3D.jl             | 12 +++++++
 src/ext/CUDA/2D.jl               | 10 ++++++
 src/ext/CUDA/3D.jl               | 19 ++++++++++
 src/types/constructors/stokes.jl | 22 +++++++++++-
 src/types/displacement.jl        | 59 ++++++++++++++++++++++++++++++++
 src/types/stokes.jl              | 23 ++++++++++++-
 src/types/traits.jl              |  1 +
 src/types/type_conversions.jl    |  4 ++-
 test/test_types.jl               | 35 ++++++++++++++++---
 11 files changed, 193 insertions(+), 7 deletions(-)
 create mode 100644 src/types/displacement.jl

diff --git a/src/common.jl b/src/common.jl
index daf3244d..45117ee1 100644
--- a/src/common.jl
+++ b/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!
diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
index 46956a11..29159d42 100644
--- a/src/ext/AMDGPU/2D.jl
+++ b/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
diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
index df2c3bde..4f2422f3 100644
--- a/src/ext/AMDGPU/3D.jl
+++ b/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
diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
index f87d8086..526a0b58 100644
--- a/src/ext/CUDA/2D.jl
+++ b/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
diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
index 8bdf2423..bc9579b4 100644
--- a/src/ext/CUDA/3D.jl
+++ b/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))
diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl
index db12d424..da75e475 100644
--- a/src/types/constructors/stokes.jl
+++ b/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
diff --git a/src/types/displacement.jl b/src/types/displacement.jl
new file mode 100644
index 00000000..1988c02f
--- /dev/null
+++ b/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
diff --git a/src/types/stokes.jl b/src/types/stokes.jl
index e7b7cd48..2f088815 100644
--- a/src/types/stokes.jl
+++ b/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}
diff --git a/src/types/traits.jl b/src/types/traits.jl
index 0ae5bc93..7a9d6fe3 100644
--- a/src/types/traits.jl
+++ b/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)
diff --git a/src/types/type_conversions.jl b/src/types/type_conversions.jl
index 015091a5..bb97cd89 100644
--- a/src/types/type_conversions.jl
+++ b/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
 
diff --git a/test/test_types.jl b/test/test_types.jl
index cc625be4..12861e65 100644
--- a/test/test_types.jl
+++ b/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
\ No newline at end of file