format

PTsolvers · Jul 5, 2024 · d8a30a6 · d8a30a6
1 parent 9599ffa
commit d8a30a6
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 42 deletions.
diff --git a/src/JustRelax.jl b/src/JustRelax.jl
@@ -32,11 +32,11 @@ include("types/phases.jl")
 # export PhaseRatio
 
 include("boundaryconditions/types.jl")
-export     AbstractBoundaryConditions,
-           TemperatureBoundaryConditions,
-           AbstractFlowBoundaryConditions,
-           DisplacementBoundaryConditions,
-           VelocityBoundaryConditions
+export AbstractBoundaryConditions,
+    TemperatureBoundaryConditions,
+    AbstractFlowBoundaryConditions,
+    DisplacementBoundaryConditions,
+    VelocityBoundaryConditions
 
 include("types/traits.jl")
 export BackendTrait, CPUBackendTrait, NonCPUBackendTrait

diff --git a/src/Utils.jl b/src/Utils.jl
@@ -115,7 +115,8 @@ macro displacement(A)
     end
 end
 
-@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} = U.Ux, U.Uy
+@inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,2}}) where {T} =
+    U.Ux, U.Uy
 @inline unpack_displacement(U::JustRelax.Displacement{<:AbstractArray{T,3}}) where {T} =
     U.Ux, U.Uy, U.Uz
 

diff --git a/src/common.jl b/src/common.jl
@@ -32,15 +32,15 @@ include("types/displacement.jl")
 export velocity2displacement!, displacement2velocity!
 
 include("boundaryconditions/BoundaryConditions.jl")
-export  AbstractBoundaryConditions,
-        TemperatureBoundaryConditions,
-        AbstractFlowBoundaryConditions,
-        DisplacementBoundaryConditions,
-        VelocityBoundaryConditions,
-        flow_bcs!,
-        thermal_bcs!,
-        pureshear_bc!,
-        apply_free_slip!
+export AbstractBoundaryConditions,
+    TemperatureBoundaryConditions,
+    AbstractFlowBoundaryConditions,
+    DisplacementBoundaryConditions,
+    VelocityBoundaryConditions,
+    flow_bcs!,
+    thermal_bcs!,
+    pureshear_bc!,
+    apply_free_slip!
 
 include("MiniKernels.jl")
 

diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl
@@ -86,19 +86,31 @@ function JR2D.PTThermalCoeffs(
 end
 
 # Boundary conditions
-function JR2D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function JR2D.flow_bcs!(
+    ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function flow_bcs!(
+    ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function JR2D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function JR2D.flow_bcs!(
+    ::AMDGPUBackendTrait,
+    stokes::JustRelax.StokesArrays,
+    bcs::DisplacementBoundaryConditions,
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 
-function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function flow_bcs!(
+    ::AMDGPUBackendTrait,
+    stokes::JustRelax.StokesArrays,
+    bcs::DisplacementBoundaryConditions,
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 

diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl
@@ -22,11 +22,11 @@ import JustRelax:
     Geometry,
     @cell
 import JustRelax:
-AbstractBoundaryConditions,
-TemperatureBoundaryConditions,
-AbstractFlowBoundaryConditions,
-DisplacementBoundaryConditions,
-VelocityBoundaryConditions
+    AbstractBoundaryConditions,
+    TemperatureBoundaryConditions,
+    AbstractFlowBoundaryConditions,
+    DisplacementBoundaryConditions,
+    VelocityBoundaryConditions
 
 @init_parallel_stencil(AMDGPU, Float64, 3)
 
@@ -86,19 +86,31 @@ function JR3D.PTThermalCoeffs(
 end
 
 # Boundary conditions
-function JR3D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function JR3D.flow_bcs!(
+    ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function flow_bcs!(
+    ::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function JR3D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function JR3D.flow_bcs!(
+    ::AMDGPUBackendTrait,
+    stokes::JustRelax.StokesArrays,
+    bcs::DisplacementBoundaryConditions,
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 
-function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function flow_bcs!(
+    ::AMDGPUBackendTrait,
+    stokes::JustRelax.StokesArrays,
+    bcs::DisplacementBoundaryConditions,
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 

diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl
@@ -85,19 +85,27 @@ function JR2D.PTThermalCoeffs(
 end
 
 # Boundary conditions
-function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function JR2D.flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function JR2D.flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 
-function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 

diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl
@@ -27,7 +27,6 @@ import JustRelax:
     DisplacementBoundaryConditions,
     VelocityBoundaryConditions
 
-
 @init_parallel_stencil(CUDA, Float64, 3)
 
 include("../../common.jl")
@@ -97,20 +96,28 @@ function JR3D.update_pt_thermal_arrays!(
 end
 
 # Boundary conditions
-function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function JR3D.flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
-function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions)
+function flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::VelocityBoundaryConditions
+)
     return _flow_bcs!(bcs, @velocity(stokes))
 end
 
 # Boundary conditions
-function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function JR3D.flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 
-function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions)
+function flow_bcs!(
+    ::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs::DisplacementBoundaryConditions
+)
     return _flow_bcs!(bcs, @displacement(stokes))
 end
 

diff --git a/src/types/displacement.jl b/src/types/displacement.jl
@@ -1,7 +1,6 @@
 # Velocity to displacement interpolation
 velocity2displacement!(stokes, dt) = velocity2displacement!(backend(stokes), stokes, dt)
 
-
 function velocity2displacement!(::CPUBackendTrait, stokes::JustRelax.StokesArrays, dt)
     return _velocity2displacement!(stokes, dt)
 end
@@ -30,7 +29,7 @@ end
 
 # Displacement to velocity interpolation
 
-displacement2velocity!(stokes, dt) =  displacement2velocity!(backend(stokes), stokes, dt)
+displacement2velocity!(stokes, dt) = displacement2velocity!(backend(stokes), stokes, dt)
 
 function displacement2velocity!(::CPUBackendTrait, stokes::JustRelax.StokesArrays, dt)
     return _displacement2velocity!(stokes, dt)
@@ -39,11 +38,15 @@ end
 function _displacement2velocity!(stokes::JustRelax.StokesArrays, dt)
     ni = size(stokes.P)
     (; V, U) = stokes
-    @parallel (@idx ni .+ 2) _displacement2velocity_kernel!(U.Ux, U.Uy, U.Uz, V.Vx, V.Vy, V.Vz, 1 / dt)
+    @parallel (@idx ni .+ 2) _displacement2velocity_kernel!(
+        U.Ux, U.Uy, U.Uz, V.Vx, V.Vy, V.Vz, 1 / dt
+    )
     return nothing
 end
 
-@parallel_indices (I...) function _displacement2velocity_kernel!(Ux, Uy, Uz, Vx, Vy, Vz, _dt)
+@parallel_indices (I...) function _displacement2velocity_kernel!(
+    Ux, Uy, Uz, Vx, Vy, Vz, _dt
+)
     if all(I .≤ size(Ux))
         Vx[I...] = Ux[I...] * _dt
     end
@@ -56,7 +59,7 @@ end
     return nothing
 end
 
-function displacement2velocity!(stokes,dt, flow_bcs::AbstractFlowBoundaryConditions)
+function displacement2velocity!(stokes, dt, flow_bcs::AbstractFlowBoundaryConditions)
     if typeof(flow_bcs) <: DisplacementBoundaryConditions
         displacement2velocity!(backend(stokes), stokes, dt)
         return nothing