format

src/IO/VTK.jl

@@ -70,7 +70,7 @@ function save_vtk(
             for (name_i, array_i) in zip(data_names_v, data_arrays_v)
                 vtk[name_i] = Array(array_i)
             end
-            vtk["Velocity"] = velocity_field
+            return vtk["Velocity"] = velocity_field
         end
     end
 

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

src/Utils.jl

@@ -31,9 +31,9 @@ function detect_args_size(A::NTuple{N,AbstractArray{T,Dims}}) where {N,T,Dims}
         Base.@_inline_meta
         s = ntuple(Val(N)) do j
             Base.@_inline_meta
-            size(A[j], i)
+            return size(A[j], i)
         end
-        maximum(s)
+        return maximum(s)
     end
 end
 
@@ -52,7 +52,7 @@ end
 Base.@propagate_inbounds @generated function unrolled_copy!(
     dst::NTuple{N,T}, src::NTuple{N,T}, I::Vararg{Int,NI}
 ) where {N,NI,T}
-    quote
+    return quote
         Base.@_inline_meta
         Base.@nexprs $N n -> begin
             if all(tuple(I...) .≤ size(dst[n]))
@@ -68,7 +68,7 @@ end
 Add `I` to the scalars in `args`
 """
 macro add(I, args...)
-    quote
+    return quote
         Base.@_inline_meta
         v = (; $(esc.(args)...))
         values(v) .+ $(esc(I))
@@ -83,8 +83,9 @@ 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
@@ -101,8 +102,9 @@ 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)
@@ -115,10 +117,11 @@ 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)
@@ -132,8 +135,9 @@ 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)
@@ -146,8 +150,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
@@ -250,7 +254,7 @@ end
     A::JustRelax.SymmetricTensor{<:AbstractArray{T,N}}
 ) where {T,N}
     syms = (:xx, :yy, :zz)
-    quote
+    return quote
         Base.@_inline_meta
         Base.@nexprs $N i -> f_i = getfield(A, $syms[i])
         Base.@ncall $N tuple f
@@ -360,8 +364,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
@@ -379,9 +383,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
@@ -428,17 +432,21 @@ 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)

src/rheology/BuoyancyForces.jl

@@ -91,17 +91,20 @@ 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
+)

src/rheology/StressUpdate.jl

@@ -24,7 +24,7 @@ function _compute_τ_nonlinear!(
     # visco-elastic strain rates
     εij_ve = ntuple(Val(N1)) do i
         Base.@_inline_meta
-        fma(0.5 * τij_o[i], _Gdt, εij[i])
+        return fma(0.5 * τij_o[i], _Gdt, εij[i])
     end
     # get plastic parameters (if any...)
     (; is_pl, C, sinϕ, cosϕ, η_reg, volume) = plastic_parameters
@@ -60,7 +60,7 @@ end
 
 # fill plastic strain rate tensor
 @generated function update_plastic_strain_rate!(ε_pl::NTuple{N,T}, λdQdτ, idx) where {N,T}
-    quote
+    return quote
         Base.@_inline_meta
         Base.@nexprs $N i -> ε_pl[i][idx...] = !isinf(λdQdτ[i]) * λdQdτ[i]
     end
@@ -76,7 +76,8 @@ function compute_stress_increment_and_trial(
 ) where {N,T}
     dτij = ntuple(Val(N)) do i
         Base.@_inline_meta
-        dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i]))
+        return dτ_r *
+               fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i]))
     end
     return dτij, second_invariant((τij .+ dτij)...)
 end
@@ -94,13 +95,13 @@ function compute_dτ_pl(
     λdQdτ = ntuple(Val(N)) do i
         Base.@_inline_meta
         # derivatives of the plastic potential
-        (τij[i] + dτij[i]) * λ_τII
+        return (τij[i] + dτij[i]) * λ_τII
     end
 
     dτ_pl = ntuple(Val(N)) do i
         Base.@_inline_meta
         # corrected stress
-        fma(-dτ_r * 2.0, ηij * λdQdτ[i], dτij[i])
+        return fma(-dτ_r * 2.0, ηij * λdQdτ[i], dτij[i])
     end
     return dτ_pl, λ, λdQdτ
 end
@@ -109,7 +110,7 @@ end
 @generated function correct_stress!(
     τ, τij::NTuple{N1}, idx::Vararg{Integer,N2}
 ) where {N1,N2}
-    quote
+    return quote
         Base.@_inline_meta
         Base.@nexprs $N1 i -> τ[i][idx...] = τij[i]
     end
@@ -179,7 +180,7 @@ end
 @generated function _plastic_params_phase(
     rheology::NTuple{N,AbstractMaterialParamsStruct}, EII, ratio
 ) where {N}
-    quote
+    return quote
         Base.@_inline_meta
         empty_args = false, 0.0, 0.0, 0.0, 0.0, 0.0
         Base.@nexprs $N i ->

src/rheology/Viscosity.jl

@@ -310,7 +310,7 @@ end
 @generated function compute_phase_viscosity_εII(
     rheology::NTuple{N,AbstractMaterialParamsStruct}, ratio, εII, args
 ) where {N}
-    quote
+    return quote
         Base.@_inline_meta
         η = 0.0
         Base.@nexprs $N i -> (

src/stokes/Stokes2D.jl

@@ -374,7 +374,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
             )
 
@@ -593,7 +593,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]
             )
 

src/stokes/StressKernels.jl

@@ -9,7 +9,8 @@ function compute_stress_increment(
 ) where {N}
     dτij = ntuple(Val(N)) do i
         Base.@_inline_meta
-        dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i]))
+        return dτ_r *
+               fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i]))
     end
     return dτij
 end

src/stokes/StressRotation.jl

@@ -145,7 +145,7 @@ Base.@propagate_inbounds function advect_stress(τxx, τyy, τxy, Vx, Vy, i, j,
     τ_adv = ntuple(Val(3)) do k
         Base.@_inline_meta
         dx_right, dx_left, dy_up, dy_down = upwind_derivatives(τ[k], i, j)
-        advection_term(Vx, Vy, dx_right, dx_left, dy_up, dy_down, _dx, _dy)
+        return advection_term(Vx, Vy, dx_right, dx_left, dy_up, dy_down, _dx, _dy)
     end
     return τ_adv
 end
@@ -160,7 +160,7 @@ Base.@propagate_inbounds function advect_stress(
         dx_right, dx_left, dy_back, dy_front, dz_up, dz_down = upwind_derivatives(
             τ[l], i, j, k
         )
-        advection_term(
+        return advection_term(
             Vx,
             Vy,
             Vz,