From 65c80007e2d49ced2b4117d8dfee908981b711ee Mon Sep 17 00:00:00 2001 From: albert-de-montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Mon, 20 Nov 2023 13:35:20 +0100 Subject: [PATCH] Update convection mini-apps (#67) * 2D convection miniapp * rename advection miniapps * update constructors * 2D WENO5 convection mini app * new miniapps * format --- GlobalConvection2D_WENO5.jl | 322 ++++++++++++++++++ .../benchmarks/stokes2D/SinkingBlock2D.jl | 219 ++++++++++++ ...tion2D.jl => GlobalConvection2D_Upwind.jl} | 16 +- .../convection/GlobalConvection2D_WENO5.jl | 322 ++++++++++++++++++ ...tion3D.jl => GlobalConvection3D_Upwind.jl} | 2 +- src/stokes/Stokes2D.jl | 4 + src/stokes/StressKernels.jl | 9 +- src/thermal_diffusion/MetaDiffusion.jl | 59 +++- 8 files changed, 935 insertions(+), 18 deletions(-) create mode 100644 GlobalConvection2D_WENO5.jl create mode 100644 miniapps/benchmarks/stokes2D/SinkingBlock2D.jl rename miniapps/convection/{GlobalConvection2D.jl => GlobalConvection2D_Upwind.jl} (96%) create mode 100644 miniapps/convection/GlobalConvection2D_WENO5.jl rename miniapps/convection/{GlobalConvection3D.jl => GlobalConvection3D_Upwind.jl} (99%) diff --git a/GlobalConvection2D_WENO5.jl b/GlobalConvection2D_WENO5.jl new file mode 100644 index 00000000..d6012d1a --- /dev/null +++ b/GlobalConvection2D_WENO5.jl @@ -0,0 +1,322 @@ +using JustRelax + +# setup ParallelStencil.jl environment +model = PS_Setup(:CUDA, Float64, 2) +environment!(model) + +using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions + +# function to compute strain rate (compulsory) +@inline function custom_εII(a::CustomRheology, TauII; args...) + η = custom_viscosity(a; args...) + return TauII / η * 0.5 +end + +# function to compute deviatoric stress (compulsory) +@inline function custom_τII(a::CustomRheology, EpsII; args...) + η = custom_viscosity(a; args...) + return 2.0 * η * EpsII +end + +# helper function (optional) +@inline function custom_viscosity(a::CustomRheology; P=0.0, T=273.0, depth=0.0, kwargs...) + (; η0, Ea, Va, T0, R, cutoff) = a.args + η = η0 * exp((Ea + P * Va) / (R * T) - Ea / (R * T0)) + correction = (depth ≤ 660e3) + (2740e3 ≥ depth > 660e3) * 1e1 + (depth > 2700e3) * 1e-1 + η = clamp(η * correction, cutoff...) +end + +# HELPER FUNCTIONS --------------------------------------------------------------- + +import ParallelStencil.INDICES +const idx_j = INDICES[2] +macro all_j(A) + esc(:($A[$idx_j])) +end + +@parallel function init_P!(P, ρg, z) + @all(P) = @all(ρg)*abs(@all_j(z)) + return nothing +end + +# Half-space-cooling model +@parallel_indices (i, j) function init_T!(T, z, κ, Tm, Tp, Tmin, Tmax) + yr = 3600*24*365.25 + dTdz = (Tm-Tp)/2890e3 + zᵢ = abs(z[j]) + Tᵢ = Tp + dTdz*(zᵢ) + time = 100e6 * yr + Ths = Tmin + (Tm -Tmin) * erf((zᵢ)*0.5/(κ*time)^0.5) + T[i, j] = min(Tᵢ, Ths) + return +end + +function circular_perturbation!(T, δT, xc, yc, r, xvi) + + @parallel_indices (i, j) function _circular_perturbation!(T, δT, xc, yc, r, x, y) + @inbounds if (((x[i] - xc))^2 + ((y[j] - yc))^2) ≤ r^2 + T[i, j] *= δT / 100 + 1 + end + return nothing + end + + @parallel _circular_perturbation!(T, δT, xc, yc, r, xvi...) +end + +function random_perturbation!(T, δT, xbox, ybox, xvi) + + @parallel_indices (i, j) function _random_perturbation!(T, δT, xbox, ybox, x, y) + @inbounds if (xbox[1] ≤ x[i] ≤ xbox[2]) && (abs(ybox[1]) ≤ abs(y[j]) ≤ abs(ybox[2])) + δTi = δT * (rand() - 0.5) # random perturbation within ±δT [%] + T[i, j] *= δTi / 100 + 1 + end + return nothing + end + + @parallel (@idx size(T)) _random_perturbation!(T, δT, xbox, ybox, xvi...) +end + +# -------------------------------------------------------------------------------- +# BEGIN MAIN SCRIPT +# -------------------------------------------------------------------------------- +function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) + + # Physical domain ------------------------------------ + ly = 2890e3 + lx = ly * ar + origin = 0.0, -ly # origin coordinates + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / (nx_g(), ny_g()) # grid step in x- and -y + xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Weno model ----------------------------------------- + weno = WENO5(ni= ni .+ 1, method=Val{2}()) # ni.+1 for Temp + # ---------------------------------------------------- + + # create rheology struct + v_args = (; η0=5e20, Ea=200e3, Va=2.6e-6, T0=1.6e3, R=8.3145, cutoff=(1e16, 1e25)) + creep = CustomRheology(custom_εII, custom_τII, v_args) + + # Physical properties using GeoParams ---------------- + η_reg = 1e16 + G0 = 70e9 # shear modulus + cohesion = 30e6 + friction = asind(0.01) + pl = DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring + β = inv(get_Kb(el)) + + rheology = SetMaterialParams(; + Name = "Mantle", + Phase = 1, + Density = PT_Density(; ρ0=3.1e3, β=β, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.2e3), + Conductivity = ConstantConductivity(; k=3.0), + CompositeRheology = CompositeRheology((creep, el, )), + Elasticity = el, + Gravity = ConstantGravity(; g=9.81), + ) + rheology_plastic = SetMaterialParams(; + Name = "Mantle", + Phase = 1, + Density = PT_Density(; ρ0=3.5e3, β=β, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.2e3), + Conductivity = ConstantConductivity(; k=3.0), + CompositeRheology = CompositeRheology((creep, el, pl)), + Elasticity = el, + Gravity = ConstantGravity(; g=9.81), + ) + # heat diffusivity + κ = (rheology.Conductivity[1].k / (rheology.HeatCapacity[1].cp * rheology.Density[1].ρ0)).val + dt = dt_diff = 0.5 * min(di...)^2 / κ / 2.01 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + periodicity = (left = false, right = false, top = false, bot = false), + ) + # initialize thermal profile - Half space cooling + adiabat = 0.3 # adiabatic gradient + Tp = 1900 + Tm = Tp + adiabat * 2890 + Tmin, Tmax = 300.0, 3.5e3 + @parallel init_T!(thermal.T, xvi[2], κ, Tm, Tp, Tmin, Tmax) + thermal_bcs!(thermal.T, thermal_bc) + # Temperature anomaly + if thermal_perturbation == :random + δT = 5.0 # thermal perturbation (in %) + random_perturbation!(thermal.T, δT, (lx*1/8, lx*7/8), (-2000e3, -2600e3), xvi) + + elseif thermal_perturbation == :circular + δT = 10.0 # thermal perturbation (in %) + xc, yc = 0.5*lx, -0.75*ly # center of the thermal anomaly + r = 150e3 # radius of perturbation + circular_perturbation!(thermal.T, δT, xc, yc, r, xvi) + end + @views thermal.T[:, 1] .= Tmax + @views thermal.T[:, end] .= Tmin + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + for _ in 1:2 + @parallel (@idx ni) compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + end + + # Rheology + η = @ones(ni...) + depth = PTArray([y for x in xci[1], y in xci[2]]) + args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf) + η_cutoff = 1e18, 1e23 + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, η_cutoff + ) + η_vep = deepcopy(η) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), + periodicity = (left = false, right = false, top = false, bot = false), + ) + # ---------------------------------------------------- + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + !isdir(figdir) && mkpath(figdir) + # ---------------------------------------------------- + + # Plot initial T and η profiles + fig0 = let + Yv = [y for x in xvi[1], y in xvi[2]][:] + 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(η)") + lines!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + lines!(ax2, Array(log10.(η[:])), Y./1e3) + ylims!(ax1, -2890, 0) + ylims!(ax2, -2890, 0) + hideydecorations!(ax2) + save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # WENO arrays + T_WENO = @zeros(ni.+1) + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + local iters + while (t / (1e6 * 3600 * 24 * 365.25)) < 4.5e3 + # Stokes solver ---------------- + args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt=Inf) + @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, η_cutoff + ) + + iters = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + rheology, + args, + dt, + igg; + iterMax=50e3, + nout=1e3, + viscosity_cutoff = η_cutoff + ); + dt = compute_dt(stokes, di, dt_diff, igg) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + # Weno advection + T_WENO .= @views thermal.T[2:end-1, :] + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) + @views thermal.T[2:end-1, :] .= T_WENO + # ------------------------------ + + it += 1 + t += dt + + println("\n") + println("Time step number $it") + println(" time = $(t/(1e6 * 3600 * 24 *365.25)) Myrs, dt = $(dt/(1e6 * 3600 * 24 *365.25)) Myrs") + println("\n") + + # Plotting --------------------- + if it == 1 || rem(it, 10) == 0 + fig = Figure(resolution = (1000, 1000), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[3,1], aspect = ar, title = "τII [MPa]") + ax4 = Axis(fig[4,1], aspect = ar, title = "log10(η)") + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T) , colormap=:batlow) + h2 = heatmap!(ax2, xci[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy[2:end-1,:]) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II.*1e-6) , colormap=:batlow) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_vep)) , colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1, height=100) + Colorbar(fig[2,2], h2, height=100) + Colorbar(fig[3,2], h3, height=100) + Colorbar(fig[4,2], h4, height=100) + save( joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return (ni=ni, xci=xci, li=li, di=di), thermal +end + +figdir = "figs2D_test_weno" +ar = 8 # aspect ratio +n = 128 +nx = n*ar - 2 +ny = n - 2 +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) +else + igg +end + +thermal_convection2D(igg; figdir=figdir, ar=ar,nx=nx, ny=ny); diff --git a/miniapps/benchmarks/stokes2D/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/SinkingBlock2D.jl new file mode 100644 index 00000000..e0aa0333 --- /dev/null +++ b/miniapps/benchmarks/stokes2D/SinkingBlock2D.jl @@ -0,0 +1,219 @@ +using JustRelax + +# setup ParallelStencil.jl environment +model = PS_Setup(:threads, Float64, 2) +environment!(model) + +using Printf, LinearAlgebra, GeoParams, GLMakie, JustPIC, CellArrays + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +@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) + ni = nx, ny + ncells = nx * ny + np = max_xcell * ncells + px, py = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(2)) + inject = @fill(false, nx, ny, eltype=Bool) + index = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool) + + @parallel_indices (i, j) function fill_coords_index(px, py, index) + # lower-left corner of the cell + 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 + end + return nothing + end + + @parallel (1:nx, 1:ny) fill_coords_index(px, py, index) + + return Particles( + (px, py), index, inject, nxcell, max_xcell, min_xcell, np, (nx, ny) + ) +end + +# Thermal rectangular perturbation +function rectangular_perturbation!(T, xc, yc, r, xvi) + + @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + depth = abs(y[j]) + dTdZ = (2047 - 2017) / 50e3 + offset = 2017 + T[i + 1, j] = (depth - 585e3) * dTdZ + offset + end + return nothing + end + ni = length.(xvi) + @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) + + return nothing +end + +function init_phases!(phases, particles, xc, yc, r) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, xc, yc, r) + @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]) + # plume - rectangular + JustRelax.@cell phases[ip, i, j] = if ((x -xc)^2 ≤ r^2) && ((depth - yc)^2 ≤ r^2) + 2.0 + else + 1.0 + end + end + return nothing + end + + @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, xc, yc, r) +end + +import ParallelStencil.INDICES +const idx_j = INDICES[2] +macro all_j(A) + esc(:($A[$idx_j])) +end + +@parallel function init_P!(P, ρg, z) + @all(P) = @all(ρg)*abs(@all_j(z)) + return nothing +end + +# -------------------------------------------------------------------------------- +# BEGIN MAIN SCRIPT +# -------------------------------------------------------------------------------- +function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) + + # Physical domain ------------------------------------ + ly = 500e3 + lx = ly * ar + origin = 0.0, -ly # origin coordinates + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / (nx_g(), ny_g()) # grid step in x- and -y + xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + δρ = 100 + rheology = ( + SetMaterialParams(; + Name = "Mantle", + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology((LinearViscous(; η = 1e21), )), + Gravity = ConstantGravity(; g=9.81), + ), + SetMaterialParams(; + Name = "Block", + Phase = 2, + Density = ConstantDensity(; ρ=3.2e3 + δρ), + CompositeRheology = CompositeRheology((LinearViscous(; η = 1e23), )), + Gravity = ConstantGravity(; g=9.81), + ) + ) + # heat diffusivity + dt = 1 + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 20, 20, 1 + particles = init_particles_cellarrays( + nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny + ) + # temperature + pPhases, = init_particle_fields_cellarrays(particles, Val(1)) + particle_args = (pPhases, ) + # Rectangular density anomaly + xc_anomaly = 250e3 # origin of thermal anomaly + yc_anomaly = -(ly-400e3) # origin of thermal anomaly + r_anomaly = 50e3 # radius of perturbation + init_phases!(pPhases, particles, xc_anomaly, abs(yc_anomaly), r_anomaly) + phase_ratios = PhaseRatio(ni, length(rheology)) + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, CFL = 0.95 / √2.1) + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=@ones(ni...), P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + # ---------------------------------------------------- + + # Viscosity + η = @ones(ni...) + args = (; dt = Inf) + η_cutoff = -Inf, Inf + @parallel (@idx ni) compute_viscosity!( + η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, η_cutoff + ) + η_vep = deepcopy(η) + # ---------------------------------------------------- + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + periodicity = (left = false, right = false, top = false, bot = false), + ) + + # Stokes solver ---------------- + args = (; T = @ones(ni...), P = stokes.P, dt=Inf) + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax=150e3, + nout=1e3, + viscosity_cutoff = η_cutoff + ); + dt = compute_dt(stokes, di, igg) + # ------------------------------ + + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity = @. √(Vx_v^2 + Vy_v^2 ) + + # Plotting --------------------- + f, _, h = heatmap(velocity, colormap=:vikO) + Colorbar(f[1,2], h) + f + # ------------------------------ + + return nothing +end + +ar = 1 # aspect ratio +n = 128 +nx = n*ar - 2 +ny = n - 2 +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) +else + igg +end + +sinking_block2D(igg; ar=ar, nx=nx, ny=ny); diff --git a/miniapps/convection/GlobalConvection2D.jl b/miniapps/convection/GlobalConvection2D_Upwind.jl similarity index 96% rename from miniapps/convection/GlobalConvection2D.jl rename to miniapps/convection/GlobalConvection2D_Upwind.jl index 73b861de..95c85006 100644 --- a/miniapps/convection/GlobalConvection2D.jl +++ b/miniapps/convection/GlobalConvection2D_Upwind.jl @@ -1,7 +1,7 @@ using JustRelax # setup ParallelStencil.jl environment -model = PS_Setup(:gpu, Float64, 2) +model = PS_Setup(:threads, Float64, 2) environment!(model) using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions @@ -79,7 +79,7 @@ end # -------------------------------------------------------------------------------- # BEGIN MAIN SCRIPT # -------------------------------------------------------------------------------- -function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :random) +function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) # initialize MPI igg = IGG(init_global_grid(nx, ny, 0; init_MPI = JustRelax.MPI.Initialized() ? false : true)...) @@ -164,7 +164,7 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1.0 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) for _ in 1:2 @@ -176,7 +176,10 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p η = @ones(ni...) depth = PTArray([y for x in xci[1], y in xci[2]]) args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf) - @parallel (@idx ni) compute_viscosity!(η, 1, 1e-15, args, rheology) + viscosity_cutoff = 1e18, 1e23 + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, viscosity_cutoff + ) η_vep = deepcopy(η) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -220,12 +223,13 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p ρg, η, η_vep, - rheology_plastic, + rheology, args, dt, igg; - iterMax=250e3, + iterMax=50e3, nout=1e3, + viscosity_cutoff = viscosity_cutoff ); dt = compute_dt(stokes, di, dt_diff, igg) # ------------------------------ diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl new file mode 100644 index 00000000..869ab7ae --- /dev/null +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -0,0 +1,322 @@ +using JustRelax + +# setup ParallelStencil.jl environment +model = PS_Setup(:threads, Float64, 2) +environment!(model) + +using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions + +# function to compute strain rate (compulsory) +@inline function custom_εII(a::CustomRheology, TauII; args...) + η = custom_viscosity(a; args...) + return TauII / η * 0.5 +end + +# function to compute deviatoric stress (compulsory) +@inline function custom_τII(a::CustomRheology, EpsII; args...) + η = custom_viscosity(a; args...) + return 2.0 * η * EpsII +end + +# helper function (optional) +@inline function custom_viscosity(a::CustomRheology; P=0.0, T=273.0, depth=0.0, kwargs...) + (; η0, Ea, Va, T0, R, cutoff) = a.args + η = η0 * exp((Ea + P * Va) / (R * T) - Ea / (R * T0)) + correction = (depth ≤ 660e3) + (2740e3 ≥ depth > 660e3) * 1e1 + (depth > 2700e3) * 1e-1 + η = clamp(η * correction, cutoff...) +end + +# HELPER FUNCTIONS --------------------------------------------------------------- + +import ParallelStencil.INDICES +const idx_j = INDICES[2] +macro all_j(A) + esc(:($A[$idx_j])) +end + +@parallel function init_P!(P, ρg, z) + @all(P) = @all(ρg)*abs(@all_j(z)) + return nothing +end + +# Half-space-cooling model +@parallel_indices (i, j) function init_T!(T, z, κ, Tm, Tp, Tmin, Tmax) + yr = 3600*24*365.25 + dTdz = (Tm-Tp)/2890e3 + zᵢ = abs(z[j]) + Tᵢ = Tp + dTdz*(zᵢ) + time = 100e6 * yr + Ths = Tmin + (Tm -Tmin) * erf((zᵢ)*0.5/(κ*time)^0.5) + T[i, j] = min(Tᵢ, Ths) + return +end + +function circular_perturbation!(T, δT, xc, yc, r, xvi) + + @parallel_indices (i, j) function _circular_perturbation!(T, δT, xc, yc, r, x, y) + @inbounds if (((x[i] - xc))^2 + ((y[j] - yc))^2) ≤ r^2 + T[i, j] *= δT / 100 + 1 + end + return nothing + end + + @parallel _circular_perturbation!(T, δT, xc, yc, r, xvi...) +end + +function random_perturbation!(T, δT, xbox, ybox, xvi) + + @parallel_indices (i, j) function _random_perturbation!(T, δT, xbox, ybox, x, y) + @inbounds if (xbox[1] ≤ x[i] ≤ xbox[2]) && (abs(ybox[1]) ≤ abs(y[j]) ≤ abs(ybox[2])) + δTi = δT * (rand() - 0.5) # random perturbation within ±δT [%] + T[i, j] *= δTi / 100 + 1 + end + return nothing + end + + @parallel (@idx size(T)) _random_perturbation!(T, δT, xbox, ybox, xvi...) +end + +# -------------------------------------------------------------------------------- +# BEGIN MAIN SCRIPT +# -------------------------------------------------------------------------------- +function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) + + # Physical domain ------------------------------------ + ly = 2890e3 + lx = ly * ar + origin = 0.0, -ly # origin coordinates + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / (nx_g(), ny_g()) # grid step in x- and -y + xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Weno model ----------------------------------------- + weno = WENO5(ni= ni .+ 1, method=Val{2}()) # ni.+1 for Temp + # ---------------------------------------------------- + + # create rheology struct + v_args = (; η0=5e20, Ea=200e3, Va=2.6e-6, T0=1.6e3, R=8.3145, cutoff=(1e16, 1e25)) + creep = CustomRheology(custom_εII, custom_τII, v_args) + + # Physical properties using GeoParams ---------------- + η_reg = 1e16 + G0 = 70e9 # shear modulus + cohesion = 30e6 + friction = asind(0.01) + pl = DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring + β = inv(get_Kb(el)) + + rheology = SetMaterialParams(; + Name = "Mantle", + Phase = 1, + Density = PT_Density(; ρ0=3.1e3, β=β, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.2e3), + Conductivity = ConstantConductivity(; k=3.0), + CompositeRheology = CompositeRheology((creep, el, )), + Elasticity = el, + Gravity = ConstantGravity(; g=9.81), + ) + rheology_plastic = SetMaterialParams(; + Name = "Mantle", + Phase = 1, + Density = PT_Density(; ρ0=3.5e3, β=β, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.2e3), + Conductivity = ConstantConductivity(; k=3.0), + CompositeRheology = CompositeRheology((creep, el, pl)), + Elasticity = el, + Gravity = ConstantGravity(; g=9.81), + ) + # heat diffusivity + κ = (rheology.Conductivity[1].k / (rheology.HeatCapacity[1].cp * rheology.Density[1].ρ0)).val + dt = dt_diff = 0.5 * min(di...)^2 / κ / 2.01 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + periodicity = (left = false, right = false, top = false, bot = false), + ) + # initialize thermal profile - Half space cooling + adiabat = 0.3 # adiabatic gradient + Tp = 1900 + Tm = Tp + adiabat * 2890 + Tmin, Tmax = 300.0, 3.5e3 + @parallel init_T!(thermal.T, xvi[2], κ, Tm, Tp, Tmin, Tmax) + thermal_bcs!(thermal.T, thermal_bc) + # Temperature anomaly + if thermal_perturbation == :random + δT = 5.0 # thermal perturbation (in %) + random_perturbation!(thermal.T, δT, (lx*1/8, lx*7/8), (-2000e3, -2600e3), xvi) + + elseif thermal_perturbation == :circular + δT = 10.0 # thermal perturbation (in %) + xc, yc = 0.5*lx, -0.75*ly # center of the thermal anomaly + r = 150e3 # radius of perturbation + circular_perturbation!(thermal.T, δT, xc, yc, r, xvi) + end + @views thermal.T[:, 1] .= Tmax + @views thermal.T[:, end] .= Tmin + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + for _ in 1:2 + @parallel (@idx ni) compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + end + + # Rheology + η = @ones(ni...) + depth = PTArray([y for x in xci[1], y in xci[2]]) + args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf) + η_cutoff = 1e18, 1e23 + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, η_cutoff + ) + η_vep = deepcopy(η) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), + periodicity = (left = false, right = false, top = false, bot = false), + ) + # ---------------------------------------------------- + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + !isdir(figdir) && mkpath(figdir) + # ---------------------------------------------------- + + # Plot initial T and η profiles + fig0 = let + Yv = [y for x in xvi[1], y in xvi[2]][:] + 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(η)") + lines!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + lines!(ax2, Array(log10.(η[:])), Y./1e3) + ylims!(ax1, -2890, 0) + ylims!(ax2, -2890, 0) + hideydecorations!(ax2) + save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # WENO arrays + T_WENO = @zeros(ni.+1) + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + + # Time loop + t, it = 0.0, 0 + local iters + while (t / (1e6 * 3600 * 24 * 365.25)) < 4.5e3 + # Stokes solver ---------------- + args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt=Inf) + @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, η_cutoff + ) + + iters = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + rheology, + args, + dt, + igg; + iterMax=50e3, + nout=1e3, + viscosity_cutoff = η_cutoff + ); + dt = compute_dt(stokes, di, dt_diff, igg) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + # Weno advection + T_WENO .= @views thermal.T[2:end-1, :] + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) + @views thermal.T[2:end-1, :] .= T_WENO + # ------------------------------ + + it += 1 + t += dt + + println("\n") + println("Time step number $it") + println(" time = $(t/(1e6 * 3600 * 24 *365.25)) Myrs, dt = $(dt/(1e6 * 3600 * 24 *365.25)) Myrs") + println("\n") + + # Plotting --------------------- + if it == 1 || rem(it, 10) == 0 + fig = Figure(resolution = (1000, 1000), title = "t = $t") + ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[3,1], aspect = ar, title = "τII [MPa]") + ax4 = Axis(fig[4,1], aspect = ar, title = "log10(η)") + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T) , colormap=:batlow) + h2 = heatmap!(ax2, xci[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy[2:end-1,:]) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II.*1e-6) , colormap=:batlow) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_vep)) , colormap=:batlow) + hidexdecorations!(ax1) + hidexdecorations!(ax2) + hidexdecorations!(ax3) + Colorbar(fig[1,2], h1, height=100) + Colorbar(fig[2,2], h2, height=100) + Colorbar(fig[3,2], h3, height=100) + Colorbar(fig[4,2], h4, height=100) + save( joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return (ni=ni, xci=xci, li=li, di=di), thermal +end + +figdir = "figs2D_test_weno" +ar = 8 # aspect ratio +n = 128 +nx = n*ar - 2 +ny = n - 2 +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) +else + igg +end + +thermal_convection2D(igg; figdir=figdir, ar=ar,nx=nx, ny=ny); diff --git a/miniapps/convection/GlobalConvection3D.jl b/miniapps/convection/GlobalConvection3D_Upwind.jl similarity index 99% rename from miniapps/convection/GlobalConvection3D.jl rename to miniapps/convection/GlobalConvection3D_Upwind.jl index da2ac546..da3fac16 100644 --- a/miniapps/convection/GlobalConvection3D.jl +++ b/miniapps/convection/GlobalConvection3D_Upwind.jl @@ -1,7 +1,7 @@ using JustRelax, JustRelax.DataIO # setup ParallelStencil.jl environment -model = PS_Setup(:gpu, Float64, 3) +model = PS_Setup(:threads, Float64, 3) environment!(model) using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 8ebf031d..176b2a5f 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -355,6 +355,7 @@ function JustRelax.solve!( # solver loop wtime0 = 0.0 λ = @zeros(ni...) + θ = @zeros(ni...) while iter < 2 || (err > ϵ && iter ≤ iterMax) wtime0 += @elapsed begin @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) @@ -381,6 +382,7 @@ function JustRelax.solve!( @tensor(stokes.τ_o), @strain(stokes), stokes.P, + θ, η, η_vep, λ, @@ -436,6 +438,8 @@ function JustRelax.solve!( end end + stokes.P .= θ + return ( iter=iter, err_evo1=err_evo1, diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 93c42fc8..fa383c6b 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -263,6 +263,7 @@ end τ_old, # shear @ centers ε, # shear @ vertices P, + θ, η, η_vep, λ, @@ -277,12 +278,16 @@ end dτ_r = compute_dτ_r(θ_dτ, ηij, _Gdt) # get plastic paremeters (if any...) - is_pl, C, sinϕ, η_reg = plastic_params(rheology[1]) - plastic_parameters = (; is_pl, C, sinϕ, η_reg) + is_pl, C, sinϕ, cosϕ, sinψ, η_reg = plastic_params_phase(rheology, 1) + # plastic volumetric change K*dt*sinϕ*sinψ + K = get_bulkmodulus(rheology[1]) + volume = isinf(K) ? 0.0 : K * dt * sinϕ * sinψ + plastic_parameters = (; is_pl, C, sinϕ, cosϕ, η_reg, volume) _compute_τ_nonlinear!( τ, τII, τ_old, ε, P, ηij, η_vep, λ, dτ_r, _Gdt, plastic_parameters, I... ) + θ[I...] = P[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) return nothing end diff --git a/src/thermal_diffusion/MetaDiffusion.jl b/src/thermal_diffusion/MetaDiffusion.jl index 969ce8a4..faf527db 100644 --- a/src/thermal_diffusion/MetaDiffusion.jl +++ b/src/thermal_diffusion/MetaDiffusion.jl @@ -97,6 +97,7 @@ function make_PTthermal_struct!() end end + # with phase ratios function PTThermalCoeffs( rheology, phase_ratios, @@ -112,28 +113,53 @@ function make_PTthermal_struct!() max_lxyz = max(li...) θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) - @parallel (1:ni[1], 1:ni[2]) compute_pt_thermal_arrays!( + idx = ntuple(i -> 1:ni[i], Val(nDim)) + @parallel (idx) compute_pt_thermal_arrays!( θr_dτ, dτ_ρ, rheology, phase_ratios.center, args, max_lxyz, Vpdτ, inv(dt) ) return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) end - @parallel_indices (i, j) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray{T,2}, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt - ) where {T} + # without phase ratios + function PTThermalCoeffs( + # function PTThermalCoeffs( + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + Vpdτ = min(di...) * CFL + max_lxyz = max(li...) + θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) + + idx = ntuple(i -> 1:ni[i], Val(nDim)) + @parallel (idx) compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, inv(dt) + ) + + return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) + end + + @parallel_indices (I...) function compute_pt_thermal_arrays!( + θr_dτ::AbstractArray, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt + ) _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, i, j + θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, I... ) return nothing end - @parallel_indices (i, j, k) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray{T,3}, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt - ) where {T} + @parallel_indices (I...) function compute_pt_thermal_arrays!( + θr_dτ::AbstractArray, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt + ) _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, i, j, k + θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, I... ) return nothing @@ -153,5 +179,20 @@ function make_PTthermal_struct!() return nothing end + + function _compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} + ) where {N} + args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) + + ρCp = JustRelax.compute_ρCp(rheology, args_ij) + _K = inv(compute_conductivity(rheology, args_ij)) + + _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number + θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re + dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re + + return nothing + end end end