fix stokes solve for non variational; typos; format

_typos.toml

@@ -2,3 +2,4 @@
 Ths = "Ths"
 oly = "oly"
 iy = "iy"
+pn = "pn"

miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl

@@ -110,7 +110,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs")
 
     # Rheology
     compute_viscosity!(
-        stokes, phase_ratios, args, rheology, (-Inf, Inf)
+        stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)
     )
     # Boundary conditions
     flow_bcs     = VelocityBoundaryConditions(;

src/boundaryconditions/Dirichlet.jl

@@ -27,14 +27,14 @@ end
 
 # struct ConstantArray{T,N} <: AbstractArray{T,N}
 #     val::T
-    
+
 #     ConstantArray(val::T, ::Val{N}) where {T<:Number,N} = new{T,N}(val)
 # end
 # Adapt.@adapt_structure ConstantArray
 
 struct ConstantArray{T}
     val::T
-    
+
     ConstantArray(val::T) where {T<:Number} = new{T}(val)
 end
 Adapt.@adapt_structure ConstantArray

src/ext/AMDGPU/2D.jl

@@ -446,8 +446,15 @@ end
 
 # marker chain
 
-function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{JustPIC.AMDGPUBackend}, particles::Particles{JustPIC.AMDGPUBackend}, origin, di, air_phase)
-    update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase)
+function JR2D.update_phases_given_markerchain!(
+    phase,
+    chain::MarkerChain{JustPIC.AMDGPUBackend},
+    particles::Particles{JustPIC.AMDGPUBackend},
+    origin,
+    di,
+    air_phase,
+)
+    return update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase)
 end
 
 end

src/ext/CUDA/2D.jl

@@ -427,8 +427,15 @@ end
 
 # marker chain
 
-function JR2D.update_phases_given_markerchain!(phase, chain::MarkerChain{CUDABackend}, particles::Particles{CUDABackend}, origin, di, air_phase)
-    update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase)
+function JR2D.update_phases_given_markerchain!(
+    phase,
+    chain::MarkerChain{CUDABackend},
+    particles::Particles{CUDABackend},
+    origin,
+    di,
+    air_phase,
+)
+    return update_phases_given_markerchain!(phase, chain, particles, origin, di, air_phase)
 end
 
 end

src/mask/mask.jl

@@ -35,21 +35,16 @@ Base.similar(m::Mask) = Mask(size(m)...)
 
 @inline dims(::Mask{A}) where {A<:AbstractArray{T,N}} where {T,N} = N
 
-@inline apply_mask!(A::AbstractArray, B::Any, m::Mask) =
-    (A .= inv(m) .* A .+ m.mask .* B)
+@inline apply_mask!(A::AbstractArray, B::Any, m::Mask) = (A .= inv(m) .* A .+ m.mask .* B)
 @inline apply_mask!(::AbstractArray, ::Any, ::Nothing) = nothing
-@inline apply_mask!(
-    A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}
-) where {N} = (A[inds...] = inv(m, inds...) * A[inds...] + m[inds...] * B[inds...])
-@inline apply_mask!(
-    ::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}
-) where {N} = nothing
+@inline apply_mask!(A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}) where {N} =
+    (A[inds...] = inv(m, inds...) * A[inds...] + m[inds...] * B[inds...])
+@inline apply_mask!(::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}) where {N} =
+    nothing
 
 @inline apply_mask(A::AbstractArray, B::Any, m::Mask) = inv(m) .* A .+ m.mask .* B
-@inline apply_mask(
-    A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}
-) where {N} = @muladd inv(m, inds...) * A[inds...] + m[inds...] * B[inds...]
+@inline apply_mask(A::AbstractArray, B::Any, m::Mask, inds::Vararg{Int,N}) where {N} =
+    @muladd inv(m, inds...) * A[inds...] + m[inds...] * B[inds...]
 @inline apply_mask(A::AbstractArray, ::Any, ::Nothing) = A
-@inline apply_mask(
-    A::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}
-) where {N} = A[inds...]
+@inline apply_mask(A::AbstractArray, ::Any, ::Nothing, inds::Vararg{Int,N}) where {N} =
+    A[inds...]

src/phases/topography_correction.jl

@@ -1,25 +1,34 @@
 import JustPIC._2D: cell_index, interp1D_inner, interp1D_extremas, distance
 using StaticArrays
 
-
-function update_phases_given_markerchain!(phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase) where {backend}
-    (; coords, index) = particles;
+function update_phases_given_markerchain!(
+    phase, chain::MarkerChain{backend}, particles::Particles{backend}, origin, di, air_phase
+) where {backend}
+    (; coords, index) = particles
     dy = di[2]
-    @parallel (1:size(index,1)) _update_phases_given_markerchain!(phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase)
+    @parallel (1:size(index, 1)) _update_phases_given_markerchain!(
+        phase, coords, index, chain.coords, chain.cell_vertices, origin, dy, air_phase
+    )
 end
 
-@parallel_indices (icell) function _update_phases_given_markerchain!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase)
-    _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords,cell_vertices, origin, dy, air_phase, icell)
+@parallel_indices (icell) function _update_phases_given_markerchain!(
+    phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase
+)
+    _update_phases_given_markerchain_kernel!(
+        phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell
+    )
     return nothing
 end
 
-function _update_phases_given_markerchain_kernel!(phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell)
-    T                       = eltype(eltype(phase))
-    chain_yi                = @cell chain_coords[2][icell]
-    min_cell_j, max_cell_j  = find_minmax_cell_indices(chain_yi, origin[2], dy)
-    min_cell_j              = max(1, min_cell_j - 10)
-    max_cell_j              = min(size(index, 2), max_cell_j + 10)
-    cell_range              = min_cell_j:max_cell_j
+function _update_phases_given_markerchain_kernel!(
+    phase, coords, index, chain_coords, cell_vertices, origin, dy, air_phase, icell
+)
+    T = eltype(eltype(phase))
+    chain_yi = @cell chain_coords[2][icell]
+    min_cell_j, max_cell_j = find_minmax_cell_indices(chain_yi, origin[2], dy)
+    min_cell_j = max(1, min_cell_j - 10)
+    max_cell_j = min(size(index, 2), max_cell_j + 10)
+    cell_range = min_cell_j:max_cell_j
 
     # iterate over cells with marker chain on them
     for j in cell_range
@@ -58,7 +67,7 @@ function extrema_CA(x::AbstractArray)
     max_val = x[1]
     min_val = x[1]
     for i in 2:length(x)
-        xᵢ = x[i] 
+        xᵢ = x[i]
         isnan(xᵢ) && continue
         if xᵢ > max_val
             max_val = xᵢ
@@ -89,10 +98,12 @@ end
 end
 
 # find closest phase different than the given `skip_phase`
-function closest_phase(coords, pn, index, current_particle, phases, skip_phase, I::Vararg{Int, N}) where N
+function closest_phase(
+    coords, pn, index, current_particle, phases, skip_phase, I::Vararg{Int,N}
+) where {N}
     new_phase = @index phases[current_particle, I...]
-    dist_min  = Inf
-    px, py    = coords
+    dist_min = Inf
+    px, py = coords
 
     for ip in cellaxes(index)
         # early escape conditions
@@ -108,7 +119,7 @@ function closest_phase(coords, pn, index, current_particle, phases, skip_phase,
         # update the closest phase
         if d < dist_min
             new_phase = phaseᵢ
-            dist_min  = d
+            dist_min = d
         end
     end
 

src/stokes/Stokes2D.jl

@@ -528,7 +528,7 @@ function _solve!(
 
     # compute buoyancy forces and viscosity
     compute_ρg!(ρg, phase_ratios, rheology, args)
-    compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff)
     displacement2velocity!(stokes, dt, flow_bcs)
 
     while iter ≤ iterMax
@@ -547,7 +547,7 @@ function _solve!(
                 stokes.∇V,
                 ητ,
                 rheology,
-                phase_ratios.center,
+                phase_ratios,
                 dt,
                 r,
                 θ_dτ,

src/stokes/Stokes3D.jl

@@ -211,7 +211,7 @@ function _solve!(
 
     # compute buoyancy forces and viscosity
     compute_ρg!(ρg, phase_ratios, rheology, args)
-    compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff)
 
     # convert displacement to velocity
     displacement2velocity!(stokes, dt, flow_bcs)
@@ -425,7 +425,7 @@ function _solve!(
 
     # compute buoyancy forces and viscosity
     compute_ρg!(ρg, phase_ratios, rheology, args)
-    compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff)
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, viscosity_cutoff)
 
     # convert displacement to velocity
     displacement2velocity!(stokes, dt, flow_bcs)
@@ -444,7 +444,7 @@ function _solve!(
                 stokes.∇V,
                 ητ,
                 rheology,
-                phase_ratios.center,
+                phase_ratios,
                 dt,
                 pt_stokes.r,
                 pt_stokes.θ_dτ,

src/variational_stokes/MiniKernels.jl

@@ -71,7 +71,8 @@ Base.@propagate_inbounds @inline _av_zi(A::T, ϕ::T, I::Vararg{Integer,3}) where
     (top(A, ϕ, I...) + next(A, ϕ, I...)) * 0.5
 
 ## Because mymaskedsum(::generator) does not work inside CUDA kernels...
-@inline mymaskedsum(A::AbstractArray, ϕ::AbstractArray, ranges::Vararg{T,N}) where {T,N} = mymaskedsum(identity, A, ϕ, ranges...)
+@inline mymaskedsum(A::AbstractArray, ϕ::AbstractArray, ranges::Vararg{T,N}) where {T,N} =
+    mymaskedsum(identity, A, ϕ, ranges...)
 
 @inline function mymaskedsum(
     f::F, A::AbstractArray, ϕ::AbstractArray, ranges_i

src/variational_stokes/mask.jl

@@ -45,9 +45,7 @@ Update the rock ratio `ϕ` based on the provided `phase_ratios` and `air_phase`.
 - `phase_ratios`: The ratios of different phases present.
 - `air_phase`: The phase representing air.
 """
-function update_rock_ratio!(
-    ϕ::JustRelax.RockRatio{T, 2}, phase_ratios, air_phase
-) where {T}
+function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,2}, phase_ratios, air_phase) where {T}
     nvi = size_v(ϕ)
     @parallel (@idx nvi) update_rock_ratio_cv!(
         ϕ, phase_ratios.center, phase_ratios.vertex, air_phase
@@ -63,27 +61,27 @@ function update_rock_ratio!(
     return nothing
 end
 
-function update_rock_ratio!(
-    ϕ::JustRelax.RockRatio{T, 3}, phase_ratios, air_phase
-) where {T}
+function update_rock_ratio!(ϕ::JustRelax.RockRatio{T,3}, phase_ratios, air_phase) where {T}
     nvi = size_v(ϕ)
     @parallel (@idx nvi) update_rock_ratio_cv!(
         ϕ, phase_ratios.center, phase_ratios.vertex, air_phase
     )
 
     dst = ϕ.Vx, ϕ.Vy, ϕ.Vz, ϕ.xy, ϕ.yz, ϕ.xz
-    src = phase_ratios.Vx, phase_ratios.Vy, phase_ratios.Vz, phase_ratios.xy, phase_ratios.yz, phase_ratios.xz
-
+    src = phase_ratios.Vx,
+    phase_ratios.Vy, phase_ratios.Vz, phase_ratios.xy, phase_ratios.yz,
+    phase_ratios.xz
+
     for (dstᵢ, srcᵢ) in zip(dst, src)
         @parallel (@idx size(dstᵢ)) _update_rock_ratio!(dstᵢ, srcᵢ, air_phase)
     end
 
     return nothing
 end
 
-@inline compute_rock_ratio(
-    phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}
-) where {N} = (x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5)
+@inline compute_rock_ratio(phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}) where {N} = (
+    x = 1 - @index phase_ratio[air_phase, I...]; x *= x > 1e-5
+)
 
 @inline compute_air_ratio(
     phase_ratio::CellArray, air_phase, I::Vararg{Integer,N}

test/test_Blankenbach.jl

@@ -142,7 +142,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10)
     η                = @ones(ni...)
     compute_ρg!(ρg[2], phase_ratios, rheology, args)
     compute_viscosity!(
-        stokes, phase_ratios, args, rheology, (-Inf, Inf)
+        stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)
     )
 
     # PT coefficients for thermal diffusion -------------
@@ -182,7 +182,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 10)
 
         # Update buoyancy and viscosity -
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
-        compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
+        compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf))
         compute_ρg!(ρg[2], phase_ratios, rheology, args)
         # ------------------------------
 

test/test_VanKeken.jl

@@ -127,7 +127,7 @@ function VanKeken2D(ny=32, nx=32)
     compute_ρg!(ρg[2], phase_ratios, rheology, args)
 
     # Rheology
-    compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf))
 
     # Boundary conditions
     flow_bcs             = VelocityBoundaryConditions(;

test/test_WENO5.jl

@@ -195,7 +195,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, thermal_perturbation =
 
     # Rheology
     viscosity_cutoff = (1e16, 1e24)
-    compute_viscosity!(stokes, args, rheology, viscosity_cutoff)
+    compute_viscosity!(stokes, args, rheology, 0, viscosity_cutoff)
 
     # PT coefficients for thermal diffusion
     pt_thermal       = PTThermalCoeffs(
@@ -225,7 +225,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, thermal_perturbation =
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
         compute_ρg!(ρg[end], rheology, (T=thermal.Tc, P=stokes.P))
         compute_viscosity!(
-            stokes, args, rheology, viscosity_cutoff
+            stokes, args, rheology, 0, viscosity_cutoff
         )
         # ------------------------------
         iters = solve!(

test/test_shearband2D.jl

@@ -130,7 +130,7 @@ function ShearBand2D()
 
     # Rheology
     compute_viscosity!(
-        stokes, phase_ratios, args, rheology, (-Inf, Inf)
+        stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)
     )
     # Boundary conditions
     flow_bcs     = VelocityBoundaryConditions(;

test/test_shearband2D_MPI.jl

@@ -126,7 +126,7 @@ function main(igg; nx=64, ny=64)
 
     # Rheology
     compute_viscosity!(
-        stokes, phase_ratios, args, rheology, (-Inf, Inf)
+        stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)
     )
 
     # Boundary conditions

test/test_shearband2D_softening.jl

@@ -135,7 +135,7 @@ function ShearBand2D()
 
     # Rheology
     compute_viscosity!(
-        stokes, phase_ratios, args, rheology, (-Inf, Inf)
+        stokes, phase_ratios, args, rheology, 0, (-Inf, Inf)
     )
 
     # Boundary conditions

test/test_shearband3D_MPI.jl

@@ -135,7 +135,7 @@ function main(igg; nx=64, ny=64, nz=64)
     args         = (; T = @zeros(ni...), P = stokes.P, dt = Inf)
     # Rheology
     cutoff_visc = -Inf, Inf
-    compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc)
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, cutoff_visc)
 
     # Boundary conditions
     flow_bcs     = VelocityBoundaryConditions(;

test/test_shearheating2D.jl

@@ -127,7 +127,7 @@ function Shearheating2D(igg; nx=32, ny=32)
 
     # Rheology
     args             = (; T = thermal.Tc, P = stokes.P, dt = Inf)
-    compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf))
 
     # PT coefficients for thermal diffusion
     pt_thermal       = PTThermalCoeffs(
@@ -226,7 +226,7 @@ function Shearheating2D(igg; nx=32, ny=32)
         @views T_buffer[:, end]      .= 273.0 + 400
         @views thermal.T[2:end-1, :] .= T_buffer
         temperature2center!(thermal)
-          
+
         @show it += 1
         t        += dt
 

test/test_shearheating3D.jl

@@ -125,7 +125,7 @@ function Shearheating3D(igg; nx=16, ny=16, nz=16)
 
     # Rheology
     args             = (; T = thermal.Tc, P = stokes.P, dt = Inf)
-    compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0,(-Inf, Inf))
 
     # PT coefficients for thermal diffusion
     pt_thermal       = PTThermalCoeffs(

test/test_sinking_block.jl

@@ -160,7 +160,7 @@ function Sinking_Block2D()
     # Viscosity
     args = (; T = @ones(ni...), P = stokes.P, dt=Inf)
     η_cutoff = -Inf, Inf
-    compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
+    compute_viscosity!(stokes, phase_ratios, args, rheology, 0, (-Inf, Inf))
     # ----------------------------------------------------
 
     # Boundary conditions