Skip to content

Commit

Permalink
thermal stress switch via multiple dispatch
Browse files Browse the repository at this point in the history
  • Loading branch information
albert-de-montserrat committed Apr 5, 2024
1 parent 08225b9 commit a318545
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 14 deletions.
31 changes: 28 additions & 3 deletions src/stokes/PressureKernels.jl
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,35 @@ Compute the pressure field `P` and the residual `RP` for the compressible case.
after the implementation of Kiss et al. (2023). The temperature difference `ΔTc` on the cell center is used to compute this
as well as α as the thermal expansivity.
"""
@parallel_indices (I...) function compute_P!(
P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, args
function compute_P!(
P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, dt, r, θ_dτ, kwargs::NamedTuple
) where N
compute_P!(P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ; kwargs...)
end

function compute_P!(
P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio, dt, r, θ_dτ; ΔTc = nothing, kwargs...
) where N
ni = size(P)
@parallel (@idx ni) compute_P_kernel!(
P, P0, RP, ∇V, η, rheology, phase_ratio, dt, r, θ_dτ, ΔTc
)
return nothing
end

@parallel_indices (I...) function compute_P_kernel!(
P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, ::Nothing
) where {N,C<:JustRelax.CellArray}
K = fn_ratio(get_bulk_modulus, rheology, phase_ratio[I...])
RP[I...], P[I...] = _compute_P!(
P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ
)
return nothing
end

@parallel_indices (I...) function compute_P_kernel!(
P, P0, RP, ∇V, η, rheology::NTuple{N,MaterialParams}, phase_ratio::C, dt, r, θ_dτ, ΔTc
) where {N,C<:JustRelax.CellArray}
ΔTc = args.ΔTc
K = fn_ratio(get_bulk_modulus, rheology, phase_ratio[I...])
α = fn_ratio(get_thermal_expansion, rheology, phase_ratio[I...])
RP[I...], P[I...] = _compute_P!(
Expand Down
6 changes: 3 additions & 3 deletions src/stokes/Stokes2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -561,7 +561,7 @@ function JustRelax.solve!(

@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...)

@parallel (@idx ni) compute_P!(
compute_P!(
stokes.P,
stokes.P0,
stokes.R.RP,
Expand Down Expand Up @@ -620,8 +620,8 @@ function JustRelax.solve!(
dt,
θ_dτ,
)
free_surface_bcs!(stokes)
# @views stokes.τ.yy[:, end] .= stokes.P[:, end]
free_surface_bcs!(stokes, flow_bcs)

@parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c)
update_halo!(stokes.τ.xy)

Expand Down
15 changes: 7 additions & 8 deletions src/stokes/Stokes3D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -417,19 +417,18 @@ function JustRelax.solve!(

# solver loop
wtime0 = 0.0
ητ = deepcopy(η)

while iter < 2 || (err > ϵ && iter iterMax)
wtime0 += @elapsed begin
# ~preconditioner
ητ = deepcopy(η)
compute_maxloc!(ητ, η)
update_halo!(ητ)
# @hide_communication b_width begin # communication/computation overlap
# @parallel compute_maxloc!(ητ, η)
# update_halo!(ητ)
# end
@hide_communication b_width begin # communication/computation overlap
@parallel compute_maxloc!(ητ, η)
update_halo!(ητ)
end

@parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...)
@parallel (@idx ni) compute_P!(
compute_P!(
stokes.P,
stokes.P0,
stokes.R.RP,
Expand Down

0 comments on commit a318545

Please sign in to comment.