format

PTsolvers · Oct 19, 2024 · b2a71af · b2a71af
1 parent d8a4e99
commit b2a71af
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 68 deletions.
diff --git a/src/MiniKernels.jl b/src/MiniKernels.jl
@@ -7,18 +7,18 @@ const T2 = AbstractArray{T,2} where {T}
 @inline _d_xi(A::T, i, j, _dx) where {T<:T2} = (-A[i, j + 1] + A[i + 1, j + 1]) * _dx
 @inline _d_yi(A::T, i, j, _dy) where {T<:T2} = (-A[i + 1, j] + A[i + 1, j + 1]) * _dy
 # averages
-@inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i+1):(i+2), (j+1):(j+2))
-@inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i+1), (j):(j+1))
+@inline _av(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i + 1):(i + 2), (j + 1):(j + 2))
+@inline _av_a(A::T, i, j) where {T<:T2} = 0.25 * mysum(A, (i):(i + 1), (j):(j + 1))
 @inline _av_xa(A::T, i, j) where {T<:T2} = (A[i, j] + A[i + 1, j]) * 0.5
 @inline _av_ya(A::T, i, j) where {T<:T2} = (A[i, j] + A[i, j + 1]) * 0.5
 @inline _av_xi(A::T, i, j) where {T<:T2} = (A[i, j + 1], A[i + 1, j + 1]) * 0.5
 @inline _av_yi(A::T, i, j) where {T<:T2} = (A[i + 1, j], A[i + 1, j + 1]) * 0.5
 # harmonic averages
 @inline function _harm(A::T, i, j) where {T<:T2}
-    return eltype(A)(4) * mysum(inv, A, (i+1):(i+2), (j+1):(j+2))
+    return eltype(A)(4) * mysum(inv, A, (i + 1):(i + 2), (j + 1):(j + 2))
 end
 @inline function _harm_a(A::T, i, j) where {T<:T2}
-    return eltype(A)(4) * mysum(inv, A, (i):(i+1), (j):(j+1))
+    return eltype(A)(4) * mysum(inv, A, (i):(i + 1), (j):(j + 1))
 end
 @inline function _harm_xa(A::T, i, j) where {T<:T2}
     return eltype(A)(2) * (inv(A[i + 1, j]) + inv(A[i, j]))
@@ -48,16 +48,16 @@ end
     return (-A[i + 1, j + 1, k] + A[i + 1, j + 1, k + 1]) * _dz
 end
 # averages
-@inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i+1), j:(j+1), k:(k+1))
+@inline _av(A::T, i, j, k) where {T<:T3} = 0.125 * mysum(A, i:(i + 1), j:(j + 1), k:(k + 1))
 @inline _av_x(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i + 1, j, k])
 @inline _av_y(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i, j + 1, k])
 @inline _av_z(A::T, i, j, k) where {T<:T3} = 0.5 * (A[i, j, k] + A[i, j, k + 1])
-@inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i+1), j:(j+1), k:k)
-@inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i+1), j:j, k:(k+1))
-@inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j+1), k:(k+1))
-@inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i-1):i, (j-1):j, k:k)
-@inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i-1):i, j:j, (k-1):k)
-@inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j-1):j, (k-1):k)
+@inline _av_xy(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:(j + 1), k:k)
+@inline _av_xz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:(i + 1), j:j, k:(k + 1))
+@inline _av_yz(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, j:(j + 1), k:(k + 1))
+@inline _av_xyi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, (j - 1):j, k:k)
+@inline _av_xzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, (i - 1):i, j:j, (k - 1):k)
+@inline _av_yzi(A::T, i, j, k) where {T<:T3} = 0.25 * mysum(A, i:i, (j - 1):j, (k - 1):k)
 # harmonic averages
 @inline function _harm_x(A::T, i, j, k) where {T<:T3}
     return eltype(A)(2) * inv(inv(A[i, j, k]) + inv(A[i + 1, j, k]))
@@ -69,22 +69,22 @@ end
     return eltype(A)(2) * inv(inv(A[i, j, k]) + inv(A[i, j, k + 1]))
 end
 @inline function _harm_xy(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, i:(i+1), j:(j+1), k:k))
+    return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:(j + 1), k:k))
 end
 @inline function _harm_xz(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, i:(i+1), j:j, k:(k+1)))
+    return eltype(A)(4) * inv(mysum(A, i:(i + 1), j:j, k:(k + 1)))
 end
 @inline function _harm_yz(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, i:i, j:(j+1), k:(k+1)))
+    return eltype(A)(4) * inv(mysum(A, i:i, j:(j + 1), k:(k + 1)))
 end
 @inline function _harm_xyi(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, (i-1):i, (j-1):j, k:k))
+    return eltype(A)(4) * inv(mysum(A, (i - 1):i, (j - 1):j, k:k))
 end
 @inline function _harm_xzi(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, (i-1):i, j:j, (k-1):k))
+    return eltype(A)(4) * inv(mysum(A, (i - 1):i, j:j, (k - 1):k))
 end
 @inline function _harm_yzi(A::T, i, j, k) where {T<:T3}
-    return eltype(A)(4) * inv(mysum(A, i:i, (j-1):j, (k-1):k))
+    return eltype(A)(4) * inv(mysum(A, i:i, (j - 1):j, (k - 1):k))
 end
 
 # others

diff --git a/src/Utils.jl b/src/Utils.jl
@@ -83,9 +83,8 @@ end
 
 @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy
 @inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz
-@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = A.xx,
-A.yy,
-A.xy_c
+@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} =
+    A.xx, A.yy, A.xy_c
 @inline function _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}}) where {T}
     return A.xx, A.yy, A.zz, A.yz_c, A.xz_c, A.xy_c
 end
@@ -102,9 +101,8 @@ macro velocity(A)
 end
 
 @inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy
-@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx,
-V.Vy,
-V.Vz
+@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} =
+    V.Vx, V.Vy, V.Vz
 
 """
     @displacement(U)
@@ -117,11 +115,10 @@ 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,3}}) where {T} = U.Ux,
-U.Uy,
-U.Uz
+@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
 
 """
     @qT(V)
@@ -135,9 +132,8 @@ macro qT(A)
 end
 
 @inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx, A.qTy
-@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = A.qTx,
-A.qTy,
-A.qTz
+@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} =
+    A.qTx, A.qTy, A.qTz
 
 """
     @qT2(V)
@@ -150,8 +146,8 @@ macro qT2(A)
     end
 end
 
-@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx2,
-A.qTy2
+@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} =
+    A.qTx2, A.qTy2
 @inline function unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T}
     return A.qTx2, A.qTy2, A.qTz2
 end
@@ -364,8 +360,8 @@ end
 
 @inline function _maxloc_window_clamped(A, I, J, width_x, width_y)
     nx, ny = size(A)
-    I_range = (I-width_x):(I+width_x)
-    J_range = (J-width_y):(J+width_y)
+    I_range = (I - width_x):(I + width_x)
+    J_range = (J - width_y):(J + width_y)
     x = -Inf
 
     for j in J_range
@@ -383,9 +379,9 @@ end
 
 @inline function _maxloc_window_clamped(A, I, J, K, width_x, width_y, width_z)
     nx, ny, nz = size(A)
-    I_range = (I-width_x):(I+width_x)
-    J_range = (J-width_y):(J+width_y)
-    K_range = (K-width_z):(K+width_z)
+    I_range = (I - width_x):(I + width_x)
+    J_range = (J - width_y):(J + width_y)
+    K_range = (K - width_z):(K + width_z)
     x = -Inf
 
     for k in K_range
@@ -432,21 +428,17 @@ function compute_dt(::CPUBackendTrait, S::JustRelax.StokesArrays, args...)
     return _compute_dt(S, args...)
 end
 
-@inline _compute_dt(S::JustRelax.StokesArrays, di) = _compute_dt(
-    @velocity(S), di, Inf, maximum
-)
+@inline _compute_dt(S::JustRelax.StokesArrays, di) =
+    _compute_dt(@velocity(S), di, Inf, maximum)
 
-@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = _compute_dt(
-    @velocity(S), di, dt_diff, maximum
-)
+@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) =
+    _compute_dt(@velocity(S), di, dt_diff, maximum)
 
-@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) = _compute_dt(
-    @velocity(S), di, dt_diff, maximum_mpi
-)
+@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, ::IGG) =
+    _compute_dt(@velocity(S), di, dt_diff, maximum_mpi)
 
-@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) = _compute_dt(
-    @velocity(S), di, Inf, maximum_mpi
-)
+@inline _compute_dt(S::JustRelax.StokesArrays, di, ::IGG) =
+    _compute_dt(@velocity(S), di, Inf, maximum_mpi)
 
 @inline function _compute_dt(V::NTuple, di, dt_diff, max_fun::F) where {F<:Function}
     n = inv(length(V) + 0.1)

diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl
@@ -91,20 +91,17 @@ Compute the buoyancy forces based on the given rheology, arguments, and phase ra
 end
 
 # without phase ratios
-@inline update_ρg!(ρg::AbstractArray, rheology, args) = update_ρg!(
-    isconstant(rheology), ρg, rheology, args
-)
+@inline update_ρg!(ρg::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
-)
+@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::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
-@inline update_ρg!(::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args) = compute_ρg!(
-    ρg, phase_ratios, rheology, args
-)
+@inline update_ρg!(
+    ::NonConstantDensityTrait, ρg, phase_ratios::JustPIC.PhaseRatios, rheology, args
+) = compute_ρg!(ρg, phase_ratios, rheology, args)
diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl
@@ -377,7 +377,7 @@ function _solve!(
             center2vertex!(stokes.τ.xy, stokes.τ.xy_c)
             update_halo!(stokes.τ.xy)
 
-            @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!(
+            @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!(
                 Vx_on_Vy, stokes.V.Vx
             )
 
@@ -601,7 +601,7 @@ function _solve!(
             )
             update_halo!(stokes.τ.xy)
 
-            @parallel (1:(size(stokes.V.Vy, 1)-2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!(
+            @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx∂ρ∂x_on_Vy!(
                 Vx_on_Vy, stokes.V.Vx, ρg[2], _di[1]
             )
 

diff --git a/src/thermal_diffusion/DiffusionExplicit.jl b/src/thermal_diffusion/DiffusionExplicit.jl
@@ -381,7 +381,7 @@ function JustRelax.solve!(
 
     # solve heat diffusion
     @parallel assign!(thermal.Told, thermal.T)
-    @parallel (1:(nx-1), 1:(ny-1)) compute_flux!(
+    @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!(
         thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy
     )
     @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _dx, _dy)
@@ -413,7 +413,7 @@ function JustRelax.solve!(
 
     # solve heat diffusion
     @parallel assign!(thermal.Told, thermal.T)
-    @parallel (1:(nx-1), 1:(ny-1)) compute_flux!(
+    @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!(
         thermal.qTx, thermal.qTy, thermal.T, rheology, phase_ratios.center, args, _di...
     )
     @parallel advect_T!(thermal.dT_dt, thermal.qTx, thermal.qTy, _di...)
@@ -442,7 +442,7 @@ function JustRelax.solve!(
     nx, ny = size(thermal.T)
     # solve heat diffusion
     @parallel assign!(thermal.Told, thermal.T)
-    @parallel (1:(nx-1), 1:(ny-1)) compute_flux!(
+    @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!(
         thermal.qTx, thermal.qTy, thermal.T, rheology, args, _dx, _dy
     )
     @parallel advect_T!(
@@ -481,7 +481,7 @@ function JustRelax.solve!(
     nx, ny = size(thermal.T)
     # solve heat diffusion
     @parallel assign!(thermal.Told, thermal.T)
-    @parallel (1:(nx-1), 1:(ny-1)) compute_flux!(
+    @parallel (1:(nx - 1), 1:(ny - 1)) compute_flux!(
         thermal.qTx, thermal.qTy, thermal.T, phases, rheology, args, _dx, _dy
     )
     @parallel advect_T!(