Skip to content

Commit

Permalink
update miniapp
Browse files Browse the repository at this point in the history
  • Loading branch information
@albert-de-montserrat
albert-de-montserrat committed Dec 6, 2023
1 parent 6a0dacd commit 3cc5e4b
Showing 1 changed file with 28 additions and 69 deletions.
97 changes: 28 additions & 69 deletions miniapps/benchmarks/stokes2D/Rayleigh_Taylor/RayleighTaylor2D.jl
View file
Original file line number Diff line number Diff line change
@@ -1,41 +1,33 @@
using CUDA
CUDA.allowscalar(false) # for safety

using JustRelax, JustRelax.DataIO, JustPIC
import JustRelax.@cell
import @cell

## NOTE: need to run one of the lines below if one wishes to switch from one backend to another
# set_backend("Threads_Float64_2D")
# set_backend("CUDA_Float64_2D")

# setup ParallelStencil.jl environment
model = PS_Setup(:CUDA, Float64, 2)
model = PS_Setup(:Threads, Float64, 2)
environment!(model)

# Load script dependencies
using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays

# Load file with all the rheology configurations
include("Layered_rheology.jl")


## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT --------------------------------
function plot_particles(particles, pPhases)
p = particles.coords
p = particles.coords
ppx, ppy = p
pxv = ppx.data[:] ./ 1e3
pyv = ppy.data[:] ./ 1e3
clr = pPhases.data[:]
# clrT = pT.data[:]
idxv = particles.index.data[:]
f,ax,h=scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma)
pxv = ppx.data[:] ./ 1e3
pyv = ppy.data[:] ./ 1e3
clr = pPhases.data[:]
idxv = particles.index.data[:]
f, _, h = scatter(Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:roma)
Colorbar(f[1,2], h)
f
end

@inline init_particle_fields(particles) = @zeros(size(particles.coords[1])...)
@inline init_particle_fields(particles, nfields) = tuple([zeros(particles.coords[1]) for i in 1:nfields]...)
@inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
@inline init_particle_fields(particles) = @zeros(size(particles.coords[1])...)
@inline init_particle_fields(particles, nfields) = tuple([zeros(particles.coords[1]) for i in 1:nfields]...)
@inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N))
@inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N))

function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, nx, ny)
Expand All @@ -51,9 +43,9 @@ function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, dx, dy, n
x0, y0 = x[i], y[j]
# fill index array
for l in 1:nxcell
JustRelax.@cell px[l, i, j] = x0 + dx * rand(0.05:1e-5:0.95)
JustRelax.@cell py[l, i, j] = y0 + dy * rand(0.05:1e-5:0.95)
JustRelax.@cell index[l, i, j] = true
@cell px[l, i, j] = x0 + dx * rand(0.05:1e-5:0.95)
@cell py[l, i, j] = y0 + dy * rand(0.05:1e-5:0.95)
@cell index[l, i, j] = true
end
return nothing
end
Expand All @@ -76,16 +68,6 @@ function velocity_grids(xci, xvi, di)
return grid_vx, grid_vy
end

function copyinn_x!(A, B)

@parallel function f_x(A, B)
@all(A) = @inn_x(B)
return nothing
end

@parallel f_x(A, B)
end

import ParallelStencil.INDICES
const idx_j = INDICES[2]
macro all_j(A)
Expand All @@ -107,15 +89,13 @@ function init_phases!(phases, particles, A)

@inbounds for ip in JustRelax.cellaxes(phases)
# quick escape
JustRelax.@cell(index[ip, i, j]) == 0 && continue

x = JustRelax.@cell px[ip, i, j]
depth = -(JustRelax.@cell py[ip, i, j])
@cell(index[ip, i, j]) == 0 && continue
x = @cell px[ip, i, j]
depth = -(@cell py[ip, i, j])
@cell phases[ip, i, j] = 2.0

if 0e0 ≤ depth ≤ 100e3
@cell phases[ip, i, j] = 1.0

elseif depth > (-f(x, A, 500e3) + (200e3 + A))
@cell phases[ip, i, j] = 3.0
end
Expand Down Expand Up @@ -147,8 +127,8 @@ function RT_2D(igg, nx, ny)
# Name = "Air",
SetMaterialParams(;
Phase = 1,
Density = ConstantDensity(; ρ=1e1),
CompositeRheology = CompositeRheology((LinearViscous(; η=1e19),)),
Density = ConstantDensity(; ρ=1e0),
CompositeRheology = CompositeRheology((LinearViscous(; η=1e16),)),
Gravity = ConstantGravity(; g=9.81),
),
# Name = "Crust",
Expand All @@ -166,7 +146,6 @@ function RT_2D(igg, nx, ny)
Gravity = ConstantGravity(; g=9.81),
)
)
dt = 1 # diffusive CFL timestep limiter
# ----------------------------------------------------

# Initialize particles -------------------------------
Expand All @@ -181,7 +160,7 @@ function RT_2D(igg, nx, ny)
particle_args = (pT, pPhases)

# Elliptical temperature anomaly
A = 15e3 # Amplitude of the anomaly
A = 5e3 # Amplitude of the anomaly
init_phases!(pPhases, particles, A)
phase_ratios = PhaseRatio(ni, length(rheology))
@parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases)
Expand All @@ -190,7 +169,7 @@ function RT_2D(igg, nx, ny)
# STOKES ---------------------------------------------
# Allocate arrays needed for every Stokes problem
stokes = StokesArrays(ni, ViscoElastic)
pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1)
pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, CFL = 0.95 / √2.1)
# ----------------------------------------------------

# TEMPERATURE PROFILE --------------------------------
Expand All @@ -215,43 +194,28 @@ function RT_2D(igg, nx, ny)
periodicity = (left = false, right = false, top = false, bot = false),
)

# Plot initial T and η profiles
let
Y = [y for x in xci[1], y in xci[2]][:]
fig = Figure(resolution = (1200, 900))
ax1 = Axis(fig[1,1], aspect = 2/3, title = "T")
ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)")
scatter!(ax1, Array(ρg[2][:]./9.81), Y./1e3)
scatter!(ax2, Array(log10.(η[:])), Y./1e3)
# scatter!(ax2, Array(stokes.P[:]), Y./1e3)
ylims!(ax1, minimum(xvi[2])./1e3, 0)
ylims!(ax2, minimum(xvi[2])./1e3, 0)
hideydecorations!(ax2)
fig
end

# For plotting: velocity arrays at cell vertices
Vx_v = @zeros(ni.+1...)
Vy_v = @zeros(ni.+1...)

# Folder where to store figures
figdir = "FreeSurface"
take(figdir)

# Time loop
t, it = 0.0, 0
dt = 1
t, it = 0.0, 0
dt = 5e3 * (3600 * 24 *365.25)
dt_max = 25e3 * (3600 * 24 *365.25)
dt = dt_max
while it < 10 # run only for 5 Myrs
while t < (6 * (1e6 * 3600 * 24 *365.25))

# Stokes solver ----------------
args = (; T = thermal.Tc, P = stokes.P, dt=Inf)

@parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P))
@parallel init_P!(stokes.P, ρg[2], xci[2])
@parallel (@idx ni) compute_viscosity!(
η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf)
)

# Stokes solver ----------------
solve!(
stokes,
pt_stokes,
Expand All @@ -269,11 +233,6 @@ function RT_2D(igg, nx, ny)
nout=1e3,
viscosity_cutoff=(-Inf, Inf)
)
@parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ))
@parallel (@idx ni) multi_copy!(
@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)
)
@parallel (JustRelax.@idx ni) tensor_invariant!(stokes.ε.II, @strain(stokes)...)
# ------------------------------

# Advection --------------------
Expand Down Expand Up @@ -309,12 +268,12 @@ function RT_2D(igg, nx, ny)
end

dt = min(compute_dt(stokes, di) / 1, dt_max)
dt = compute_dt(stokes, di)

end
# return fig
end
## END OF MAIN SCRIPT ----------------------------------------------------------------
# RT_2D(igg, nx, ny)

# (Path)/folder where output data and figures are stored
n = 128
Expand Down

0 comments on commit 3cc5e4b

Please sign in to comment.