new subction benchmarks

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);

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