diff --git a/miniapps/benchmarks/stokes2D/Rayleigh_Taylor/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/Rayleigh_Taylor/RayleighTaylor2D.jl index 0e3c5570..f5fefe3e 100644 --- a/miniapps/benchmarks/stokes2D/Rayleigh_Taylor/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/Rayleigh_Taylor/RayleighTaylor2D.jl @@ -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) @@ -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 @@ -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) @@ -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 @@ -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", @@ -166,7 +146,6 @@ function RT_2D(igg, nx, ny) Gravity = ConstantGravity(; g=9.81), ) ) - dt = 1 # diffusive CFL timestep limiter # ---------------------------------------------------- # Initialize particles ------------------------------- @@ -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) @@ -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 -------------------------------- @@ -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, @@ -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 -------------------- @@ -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