From 46f7c2c7f5049fecb09c092e0a20ff0adf97ff20 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 30 Oct 2024 11:38:28 +0100 Subject: [PATCH] new subction benchmarks --- .../stokes2D/subduction.jl/Subduction2D.jl | 317 ++++++++++++++++++ .../subduction.jl/Subduction2D_rheology.jl | 69 ++++ .../subduction.jl/Subduction2D_setup.jl | 69 ++++ .../stokes3D/subduction/Subduction3D.jl | 209 ++++++++++++ .../subduction/Subduction3D_rheology.jl | 71 ++++ .../stokes3D/subduction/Subduction3D_setup.jl | 70 ++++ 6 files changed, 805 insertions(+) create mode 100644 miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D.jl create mode 100644 miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_rheology.jl create mode 100644 miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_setup.jl create mode 100644 miniapps/benchmarks/stokes3D/subduction/Subduction3D.jl create mode 100644 miniapps/benchmarks/stokes3D/subduction/Subduction3D_rheology.jl create mode 100644 miniapps/benchmarks/stokes3D/subduction/Subduction3D_setup.jl diff --git a/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D.jl b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D.jl new file mode 100644 index 00000000..ee7b5721 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D.jl @@ -0,0 +1,317 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO + +const backend = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences2D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 2) +else + @init_parallel_stencil(Threads, Float64, 2) +end + +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, CellArrays + +# Load file with all the rheology configurations +include("Subduction2D_setup.jl") +include("Subduction2D_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[2] +macro all_k(A) + esc(:($A[$idx_k])) +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 + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false) + + # Physical domain ------------------------------------ + ni = nx, ny # number of cells + di = @. li / ni # grid steps + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies() + dt = 10e3 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell = 40 + max_xcell = 60 + min_xcell = 20 + particles = init_particles( + backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # material phase & temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + + # particle fields for the stress rotation + pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress + # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress + pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity + particle_args = (pT, pPhases, pτ..., pω...) + particle_args_reduced = (pT, pτ..., pω...) + + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = phase_ratios = PhaseRatios(backend_JP, length(rheology), ni); + init_phases!(pPhases, phases_device, particles, xvi) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # # TEMPERATURE PROFILE -------------------------------- + # Ttop = 20 + 273 + # Tbot = maximum(T_GMG) + thermal = ThermalArrays(backend, ni) + # @views thermal.T[2:end-1, :] .= PTArray(backend)(T_GMG) + # thermal_bc = TemperatureBoundaryConditions(; + # no_flux = (left = true, right = true, top = false, bot = false), + # ) + # thermal_bcs!(thermal, thermal_bc) + # @views thermal.T[:, end] .= Ttop + # @views thermal.T[:, 1] .= Tbot + # temperature2center!(thermal) + # # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + stokes.P .= PTArray(backend)(reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2)) + + # Rheology + args0 = (T=thermal.Tc, P=stokes.P, dt = Inf) + viscosity_cutoff = (1e18, 1e23) + compute_viscosity!(stokes, phase_ratios, args0, rheology, viscosity_cutoff) + + # # PT coefficients for thermal diffusion + # pt_thermal = PTThermalCoeffs( + # backend, rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-8, CFL=0.95 / √2 + # ) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true), + free_surface = false, + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + end + + # T_buffer = @zeros(ni.+1) + # Told_buffer = similar(T_buffer) + # dt₀ = similar(stokes.P) + # for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + # copyinn_x!(dst, src) + # end + # grid2particle!(pT, xvi, T_buffer, particles) + + τxx_v = @zeros(ni.+1...) + τyy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + + # fig_iters = Figure(size=(1200, 800)) + # ax_iters1 = Axis(fig_iters[1,1], aspect = 1, title = "error") + # ax_iters2 = Axis(fig_iters[1,2], aspect = 1, title = "num iters / ny") + + while it < 1000 # run only for 5 Myrs + + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = ( + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = viscosity_cutoff, + free_surface = false, + viscosity_relaxation = 1e-2 + ) + ); + end + + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + tensor_invariant!(stokes.ε) + dt = compute_dt(stokes, di) * 0.8 + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 10) == 0 + # checkpointing(figdir, stokes, thermal.T, η, t) + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + ) + data_c = (; + P = Array(stokes.P), + η = Array(η_vep), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + # Make particles plottable + p = particles.coords + ppx, ppy = p + pxv = ppx.data[:]./1e3 + pyv = ppy.data[:]./1e3 + clr = pPhases.data[:] + # clr = pT.data[:] + idxv = particles.index.data[:]; + + # Make Makie figure + ar = 3 + fig = Figure(size = (1200, 900), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "log10(εII) (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Phase") + ax3 = Axis(fig[1,3], aspect = ar, title = "τII") + ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)") + # Plot temperature + # h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) + h1 = heatmap!(ax1, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) + # Plot 2nd invariant of strain rate + # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array((stokes.τ.II)) , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + fig + save(joinpath(figdir, "$(it).png"), fig) + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction2D" +n = 128 +nx, ny = 250, 100 +li, origin, phases_GMG, T_GMG = GMG_subduction_2D(nx+1, ny+1) +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); \ No newline at end of file diff --git a/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_rheology.jl b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_rheology.jl new file mode 100644 index 00000000..27401d5b --- /dev/null +++ b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_rheology.jl @@ -0,0 +1,69 @@ +function init_rheologies() + + # Define rheolgy struct + rheology = ( + # Name = "Asthenoshpere", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Oceanic lithosphere", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=0e0), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e19),)), + Gravity = ConstantGravity(; g=9.81), + ), + ) +end + +function init_phases!(phases, phase_grid, particles, xvi) + ni = size(phases) + @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) +end + +@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} + + ni = size(phases) + + for ip in cellaxes(phases) + # quick escape + @index(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1 + ii = I[1] + offi + jj = I[2] + offj + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + ) + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj] + end + end + @index phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_setup.jl b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_setup.jl new file mode 100644 index 00000000..6158fe3e --- /dev/null +++ b/miniapps/benchmarks/stokes2D/subduction.jl/Subduction2D_setup.jl @@ -0,0 +1,69 @@ +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + model_depth = 700 + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + Tbot = 1474.0 + x = range(0, 3000, nx); + air_thickness = 50.0 + z = range(-model_depth, air_thickness, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(Tbot, nx, 1, nz); + Tlab = 1300 + + # phases + # 1: asthenosphere + # 2: lithosphere + # 3: air + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 3000), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[], Phases=[0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(1000, 3000), + zlim =(-100, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[100], Phases=[1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add slab + add_box!( + Phases, + Temp, + Grid2D; + xlim =(1000, 1200), + zlim =(-100, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=90, + phase = LithosphericPhases(Layers=[200], Phases=[1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + Phases .+= 1 + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 20.0 + Phases[surf] .= 3 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] + T = Temp[:,1,:] + + return li, origin, ph, T +end \ No newline at end of file diff --git a/miniapps/benchmarks/stokes3D/subduction/Subduction3D.jl b/miniapps/benchmarks/stokes3D/subduction/Subduction3D.jl new file mode 100644 index 00000000..bbc5ee07 --- /dev/null +++ b/miniapps/benchmarks/stokes3D/subduction/Subduction3D.jl @@ -0,0 +1,209 @@ +using CUDA +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO +# const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend_JR = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +using ParallelStencil +using ParallelStencil.FiniteDifferences3D +@init_parallel_stencil(CUDA, Float64, 3) +# @init_parallel_stencil(Threads, Float64, 3) + +using JustPIC, JustPIC._3D +const backend_JP = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +# const backend_JP = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie + +# Load file with all the rheology configurations +include("Subduction3D_rheology.jl") +include("Subduction3D_setup.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[3] +macro all_k(A) + esc(:($A[$idx_k])) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main3D(li, origin, phases_GMG, igg; nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) + + # Physical domain ------------------------------------ + ni = nx, ny, nz # number of cells + di = @. li / ni # grid steps + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies() + dt = 10e3 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 50, 75, 25 + particles = init_particles(backend_JP, nxcell, max_xcell, min_xcell, xvi, di, ni) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vx, grid_vy, grid_vz = velocity_grids(xci, xvi, di) + # temperature + particle_args = pPhases, = init_cell_arrays(particles, Val(1)) + + # Assign particles phases anomaly + phases_device = PTArray(backend_JR)(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatios(backend_JP, length(rheology), ni) + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=5e-3, CFL = 0.99 / √3.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + temperature2center!(thermal) + # ---------------------------------------------------- + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(3)) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # @parallel (@idx ni) init_P!(stokes.P, ρg[3], xci[3]) + stokes.P .= PTArray(backend_JR)(reverse(cumsum(reverse((ρg[end]).* di[end], dims=3), dims=3), dims=3)) + # Rheology + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + viscosity_cutoff = (1e18, 1e24) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + + # Boundary conditions + flow_bcs = VelocityBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = false, front = true , back = true ), + no_slip = (left = false, right = false, top = false, bot = true, front = false, back = false), + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = joinpath(figdir, "vtk") + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + local Vx_v, Vy_v, Vz_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + Vz_v = @zeros(ni.+1...) + end + # Time loop + t, it = 0.0, 0 + # while it < 10000 # run only for 5 Myrs + while (t/(1e6 * 3600 * 24 *365.25)) < 10 # run only for 5 Myrs + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + Inf, + igg; + kwargs =(; + iterMax = 100e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) + ); + end + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + tensor_invariant!(stokes.ε) + dt = compute_dt(stokes, di) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy, grid_vz), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (), (), xvi) + # update phase ratios + update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 20) == 0 + # checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + data_v = (; + T = zeros(ni.+1...), + phase_vertex = [argmax(p) for p in Array(phase_ratios.center)] + ) + data_c = (; + # P = Array(stokes.P), + # τII = Array(stokes.τ.II), + # εII = Array(stokes.ε.II), + η = Array(stokes.viscosity.η), + phase_center = [argmax(p) for p in Array(phase_ratios.center)] + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + Array(Vz_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +# nx,ny,nz = 50, 50, 50 +nx,ny,nz = 125, 3, 50 +# nx,ny,nz = 128, 32, 64 +li, origin, phases_GMG, = GMG_only(nx+1, ny+1, nz+1) +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, nz; init_MPI= true)...) +else + igg +end + +# (Path)/folder where output data and figures are stored +figdir = "Subduction3D" +# main3D(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, nz = nz, do_vtk = do_vtk); + diff --git a/miniapps/benchmarks/stokes3D/subduction/Subduction3D_rheology.jl b/miniapps/benchmarks/stokes3D/subduction/Subduction3D_rheology.jl new file mode 100644 index 00000000..7afd971f --- /dev/null +++ b/miniapps/benchmarks/stokes3D/subduction/Subduction3D_rheology.jl @@ -0,0 +1,71 @@ +function init_rheologies() + + # Define rheolgy struct + rheology = ( + # Name = "Asthenoshpere", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e21),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Oceanic lithosphere", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.3e3), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e23),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=0e0), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e19),)), + Gravity = ConstantGravity(; g=9.81), + ), + ) +end + +function init_phases!(phases, phase_grid, particles, xvi) + ni = size(phases) + @parallel (@idx ni) _init_phases!(phases, phase_grid, particles.coords, particles.index, xvi) +end + +@parallel_indices (I...) function _init_phases!(phases, phase_grid, pcoords::NTuple{N, T}, index, xvi) where {N,T} + + ni = size(phases) + + for ip in cellaxes(phases) + # quick escape + @index(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @index pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1, offk in 0:1 + ii, jj, kk = I[1] + offi, I[2] + offj, I[3] + offk + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + !(kk ≤ ni[3]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + xvi[3][kk], + ) + # @show xvᵢ ii jj kk + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj, kk] + end + end + @index phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end \ No newline at end of file diff --git a/miniapps/benchmarks/stokes3D/subduction/Subduction3D_setup.jl b/miniapps/benchmarks/stokes3D/subduction/Subduction3D_setup.jl new file mode 100644 index 00000000..561d5376 --- /dev/null +++ b/miniapps/benchmarks/stokes3D/subduction/Subduction3D_setup.jl @@ -0,0 +1,70 @@ +using GeophysicalModelGenerator + +function GMG_only(nx, ny, nz) + model_depth = 700 + # Our starting basis is the example above with ridge and overriding slab + Tbot = 1474.0 + x = range(0, 3000, nx); + y = range(0, 100, ny); + air_thickness = 50.0 + z = range(-model_depth, air_thickness, nz); + Grid2D = CartData(xyz_grid(x, y, z)) + Phases = zeros(Int64, nx, ny, nz); + Temp = fill(Tbot, nx, ny, nz); + Tlab = 1300 + + # phases + # 1: asthenosphere + # 2: lithosphere + # 3: air + add_box!( + Phases, + Temp, + Grid2D; + xlim = (0, 3000), + zlim = (-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[], Phases=[0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim = (1000, 3000), + zlim = (-100, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[100], Phases=[1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add slab + add_box!( + Phases, + Temp, + Grid2D; + xlim = (1000, 1200), + zlim = (-100, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=90, + phase = LithosphericPhases(Layers=[200], Phases=[1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + Phases .+= 1 + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 20.0 + Phases[surf] .= 3 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = ( + abs(last(x)-first(x)), + abs(last(y)-first(y)), + abs(last(z)-first(z)) + ) .* 1e3 + origin = (x[1], y[1], z[1]) .* 1e3 + + return li, origin, Phases, Temp +end