Skip to content

Commit

Permalink
update miniapps
Browse files Browse the repository at this point in the history
  • Loading branch information
@albert-de-montserrat
albert-de-montserrat committed Nov 29, 2024
1 parent dc88598 commit f084316
Show file tree
Hide file tree
Showing 2 changed files with 388 additions and 62 deletions.
166 changes: 104 additions & 62 deletions miniapps/benchmarks/stokes2D/free_surface_stabilization/Crameri2D.jl
View file
Original file line number Diff line number Diff line change
@@ -1,39 +1,39 @@
using CUDA
using JustRelax, JustRelax.JustRelax2D
const backend_JR = CUDABackend
# const backend_JR = CPUBackend
const isCUDA = false
# const isCUDA = true

using JustPIC, JustPIC._2D
const backend = CUDABackend
# const backend = JustPIC.CPUBackend

using ParallelStencil, ParallelStencil.FiniteDifferences2D
@init_parallel_stencil(CUDA, Float64, 2)
@static if isCUDA
using CUDA
end

# Load script dependencies
using LinearAlgebra, GeoParams, GLMakie
using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO

# Velocity helper grids for the particle advection
function copyinn_x!(A, B)
@parallel function f_x(A, B)
@all(A) = @inn_x(B)
return nothing
end
@parallel f_x(A, B)
const backend = @static if isCUDA
CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
else
JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
end

import ParallelStencil.INDICES
const idx_j = INDICES[2]
macro all_j(A)
esc(:($A[$idx_j]))
using ParallelStencil

@static if isCUDA
@init_parallel_stencil(CUDA, Float64, 2)
else
@init_parallel_stencil(Threads, Float64, 2)
end

# Initial pressure profile - not accurate
@parallel function init_P!(P, ρg, z)
@all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0)
return nothing
using JustPIC, JustPIC._2D
# Threads is the default backend,
# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script,
# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script.
const backend_JP = @static if isCUDA
CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
else
JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
end

# Load script dependencies
using GeoParams, GLMakie

function init_phases!(phases, particles, A)
ni = size(phases)

Expand All @@ -57,8 +57,8 @@ function init_phases!(phases, particles, A)
@index phases[ip, i, j] = 2.0
end

if depth < (-cos(x * 2π/700e3) * 7e3 + 100e3)
# if depth < (-cos(x * 2π/2800e3) * 7e3 + 100e3)
# if depth < (-cos(x * 2π/700e3) * 7e3 + 100e3)
if depth < (-cos(x * 2π/2800e3) * 7e3 + 100e3)
@index phases[ip, i, j] = 1.0
end

Expand All @@ -72,7 +72,7 @@ end
## END OF HELPER FUNCTION ------------------------------------------------------------

# (Path)/folder where output data and figures are stored
n = 101
n = 64
nx = n
ny = n
igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid
Expand Down Expand Up @@ -101,8 +101,8 @@ function main(igg, nx, ny)
# Name = "Air",
SetMaterialParams(;
Phase = 1,
Density = ConstantDensity(; ρ=3.3e3),
CompositeRheology = CompositeRheology((LinearViscous(; η=1e23),)),
Density = ConstantDensity(; ρ=0e0),
CompositeRheology = CompositeRheology((LinearViscous(; η=1e22),)),
Gravity = ConstantGravity(; g=10),
),
# Name = "Crust",
Expand All @@ -125,7 +125,7 @@ function main(igg, nx, ny)
# Initialize particles -------------------------------
nxcell, max_xcell, min_xcell = 60, 80, 40
particles = init_particles(
backend, nxcell, max_xcell, min_xcell, xvi, di, ni
backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni
)
# velocity grids
grid_vx, grid_vy = velocity_grids(xci, xvi, di)
Expand All @@ -135,32 +135,33 @@ function main(igg, nx, ny)

# Elliptical temperature anomaly
A = 5e3 # Amplitude of the anomaly
phase_ratios = PhaseRatios(backend, length(rheology), ni)
phase_ratios = PhaseRatios(backend_JP, length(rheology), ni)
init_phases!(pPhases, particles, A)
update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases)
# ----------------------------------------------------


# rock ratios for variational stokes
# RockRatios
air_phase = 1
ϕ = RockRatio(backend, ni)
update_rock_ratio!(ϕ, phase_ratios, air_phase)
air_phase = 1
ϕ = RockRatio(backend, ni)
update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase)
# ----------------------------------------------------

# STOKES ---------------------------------------------
# Allocate arrays needed for every Stokes problem
stokes = StokesArrays(backend_JR, ni)
stokes = StokesArrays(backend, ni)
pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3e0, r=0.7, CFL = 0.98 / √2.1)
# ----------------------------------------------------

# TEMPERATURE PROFILE --------------------------------
thermal = ThermalArrays(backend_JR, ni)
thermal = ThermalArrays(backend, ni)
# ----------------------------------------------------

# Buoyancy forces & rheology
ρg = @zeros(ni...), @zeros(ni...)
args = (; T = thermal.Tc, P = stokes.P, dt = Inf)
compute_ρg!(ρg[2], phase_ratios, rheology, args)
@parallel init_P!(stokes.P, ρg[2], xci[2])
compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf))
stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2))
compute_viscosity!(stokes, phase_ratios, args, rheology, air_phase, (1e18, 1e24))

# Boundary conditions
flow_bcs = VelocityBoundaryConditions(;
Expand All @@ -178,21 +179,10 @@ function main(igg, nx, ny)
# Time loop
t, it = 0.0, 0
dt = 10e3 * (3600 * 24 * 365.25)
dt_max = 50e3 * (3600 * 24 * 365.25)

_di = inv.(di)
(; ϵ, r, θ_dτ, ηdτ) = pt_stokes
(; η, η_vep) = stokes.viscosity
ni = size(stokes.P)
iterMax = 15e3
nout = 1e3
viscosity_cutoff = (-Inf, Inf)
free_surface = false
ητ = @zeros(ni...)
while it < 20

## variational solver
# Stokes solver ----------------
# Variational Stokes solver ----------------
solve_VariationalStokes!(
stokes,
pt_stokes,
Expand All @@ -203,14 +193,14 @@ function main(igg, nx, ny)
ϕ,
rheology,
args,
dt,
Inf,
igg;
kwargs = (
iterMax = 50e3,
iterMin = 1e3,
viscosity_relaxation = 1e-2,
viscosity_relaxation = 1e-2,
nout = 2e3,
viscosity_cutoff = (-Inf, Inf)
viscosity_cutoff = (1e18, 1e24)
)
)
dt = compute_dt(stokes, di, dt_max)
Expand All @@ -225,7 +215,7 @@ function main(igg, nx, ny)
inject_particles_phase!(particles, pPhases, (), (), xvi)
# update phase ratios
update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases)
update_rock_ratio!(ϕ, phase_ratios, air_phase)
update_rock_ratio!(ϕ, phase_ratios, (phase_ratios.Vx, phase_ratios.Vy), air_phase)

@show it += 1
t += dt
Expand Down Expand Up @@ -256,14 +246,15 @@ function main(igg, nx, ny)
xci,
data_v,
data_c,
velocity_v
velocity_v;
t = t
)
# end

# if it == 1 || rem(it, 1) == 0
# px, py = particles.coords

# velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
# nt = 5
# fig = Figure(size = (900, 900), title = "t = $t")
# ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs")
Expand All @@ -279,6 +270,27 @@ function main(igg, nx, ny)
# fig
# save(joinpath(figdir, "$(it).png"), fig)

fig = Figure(size = (900, 900), title = "t = $t")
ax = Axis(fig[1,1], aspect = 1, title = " t=$(round.(t/(1e3 * 3600 * 24 *365.25); digits=3)) Kyrs")

# Make particles plottable
nt = 5
p = particles.coords
ppx, ppy = p
pxv = ppx.data[:]./1e3
pyv = ppy.data[:]./1e3
clr = pPhases.data[:]
idxv = particles.index.data[:];
heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.V.Vy))
scatter!(ax, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 2, colormap=:grayC)
arrows!(
ax,
xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))...,
lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)),
color = :red,
)
save(joinpath(figdir, "$(it).png"), fig)
fig
# end
end
return nothing
Expand All @@ -287,4 +299,34 @@ end
## END OF MAIN SCRIPT ----------------------------------------------------------------
main(igg, nx, ny)

# heatmap(Array(ρg[2]), colormap = :vikO)


# @parallel_indices (I...) function compute_P!(
# P,
# P0,
# RP,
# ∇V,
# η,
# rheology::NTuple{N,MaterialParams},
# phase_ratio,
# ϕ::JustRelax.RockRatio,
# dt,
# r,
# θ_dτ,
# ::Nothing,
# ) where {N}
# # if isvalid_c(ϕ, I...)
# # K = fn_ratio(get_bulk_modulus, rheology, @cell(phase_ratio[I...]))
# # RP[I...], P[I...] = _compute_P!(
# # P[I...], P0[I...], ∇V[I...], η[I...], K, dt, r, θ_dτ
# # )
# # else
# # RP[I...] = P[I...] = zero(eltype(P))
# # end
# return nothing
# end

# @parallel (@idx ni) compute_P_kernel!(
# stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, stokes.viscosityη,
# rheology, phase_ratio.center, ϕ, dt, pt_stokes.r, pt_stokes.θ_dτ, nothing
# )
Loading

0 comments on commit f084316

Please sign in to comment.