diff --git a/Lithosphere_extension.jl b/Lithosphere_extension.jl deleted file mode 100644 index 6fa13cf2..00000000 --- a/Lithosphere_extension.jl +++ /dev/null @@ -1,931 +0,0 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) - -using JustPIC -using 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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - -using Printf, Statistics, LinearAlgebra, GeoParams, CairoMakie, CellArrays -using StaticArrays -using ImplicitGlobalGrid -using MPI: MPI -using WriteVTK - -# ----------------------------------------------------------------------------------------- -## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- - -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) - esc(:($A[$idx_j])) -end - -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z, sticky_air) - @all(P) = (abs(@all(ρg) * (@all_j(z))) - (@all(ρg)*sticky_air)) * <((@all_j(z)), 0.0) - return nothing -end - -function init_phases!(phases, particles, xc, yc, a, b, r, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!( - phases, px, py, index, xc, yc, a, b, r, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) - # quick escape - JustRelax.@cell(index[ip, i, j]) == 0 && continue - - x = JustRelax.@cell px[ip, i, j] - y = -(JustRelax.@cell py[ip, i, j]) - sticky_air - if 0e0 ≤ y ≤ 20e3 - @cell phases[ip, i, j] = 1.0 # crust - end - - # # chamber - elliptical - if (((x - xc)^2 / ((a)^2)) + ((y + yc)^2 / ((b)^2)) ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 2.0 - end - - # thermal anomaly - circular - if ((x - xc_anomaly)^2 + (y + yc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, i, j] = 3.0 - end - - if y < 0.0 - @cell phases[ip, i, j] = 4.0 - end - - end - return nothing - end - - @parallel (JustRelax.@idx ni) init_phases!( - phases, particles.coords..., particles.index, xc, yc, a, b, r, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) -end - -function update_depth!(depth_corrected, phases, topo_interp, x, y) - nx, ny = length(x), length(y) - - for i in 1:nx, j in 1:ny - - # vertex coordinates - xv, yv = x[i], y[j] - # topography at vertex - y_topo = topo_interp(xv) - # depth - depth = yv - - depth_corrected[i, j] = abs(depth - y_topo) - if depth > y_topo - phases[i, j] = 4.0 - # phases[i,j] = Int64(3) - else - # phases[i,j] = Int64(1) - phases[i, j] = 1.0 - end - end - return nothing -end - -# Initial thermal profile -@parallel_indices (i, j) function init_T!(T, y, sticky_air) - depth = -y[j] - sticky_air - - # (depth - 15e3) because we have 15km of sticky air - if depth < 0e0 - T[i + 1, j] = 273e0 - - - elseif 0e0 ≤ (depth) < 25e3 - dTdZ = (923-273)/35e3 - offset = 273e0 - T[i + 1, j] = (depth) * dTdZ + offset - - elseif 110e3 > (depth) ≥ 35e3 - dTdZ = (1492-923)/75e3 - offset = 923 - T[i + 1, j] = (depth - 35e3) * dTdZ + offset - - elseif (depth) ≥ 110e3 - dTdZ = (1837 - 1492)/590e3 - offset = 1492e0 - T[i + 1, j] = (depth - 110e3) * dTdZ + offset - - end - - return nothing -end - - -function circular_anomaly!(T, anomaly, xc, yc, r, xvi, sticky_air) - @parallel_indices (i, j) function _circular_anomaly!(T, anomaly, xc, yc, r, x, y, sticky_air) - depth = -y[j] - sticky_air - @inbounds if (((x[i] - xc)^2 + (depth[j] + yc)^2) ≤ r^2) - T[i + 1, j] = anomaly - end - return nothing - end - - ni = length.(xvi) - @parallel (@idx ni) _circular_anomaly!(T, anomaly, xc, yc, r, xvi..., sticky_air) - return nothing -end - -function elliptical_anomaly!(T, anomaly, xc, yc, a, b, r, xvi, sticky_air) - - @parallel_indices (i, j) function _elliptical_anomaly!( - T, anomaly, xc, yc, a, b, r, x, y, sticky_air - ) - depth = -y[j] - sticky_air - @inbounds if (((x[i] - xc)^2 / a^2) + ((depth[j] + yc)^2 / b^2) ≤ r^2) - T[i + 1, j ] = anomaly - end - return nothing - end - - ni = length.(xvi) - @parallel (@idx ni) _elliptical_anomaly!(T, anomaly, xc, yc, a, b, r, xvi..., sticky_air) - return nothing -end - -function elliptical_anomaly_gradient!(T, offset, xc, yc, a, b, r, xvi, sticky_air) - - @parallel_indices (i, j) function _elliptical_anomaly_gradient!( - T, offset, xc, yc, a, b, r, x, y, sticky_air - ) - depth = -y[j] - sticky_air - @inbounds if (((x[i] - xc)^2 / ((a)^2)) + ((depth[j] + yc)^2 / ((b)^2)) ≤ r^2) - T[i+1, j] += offset - # T[i + 1, j] = rand(0.55:0.001:offset) - end - return nothing - end - nx, ny = size(T) - - @parallel (1:nx-2, 1:ny) _elliptical_anomaly_gradient!(T, offset, xc, yc, a, b, r, xvi..., sticky_air) -end - - -function conduit_gradient!(T, offset, xc_conduit, yc_conduit, r_conduit, xvi, sticky_air) - - @parallel_indices (i, j) function _conduit_gradient!( - T, offset, xc_conduit, yc_conduit, r_conduit, x, y, sticky_air - ) - depth = -y[j] - sticky_air - @inbounds if ((x[i] - xc_conduit)^2 ≤ r_conduit^2) && (0.325*(depth[j] + yc_conduit)^2 ≤ r_conduit^2) - T[i+1, j] += offset - end - return nothing - end - nx, ny = size(T) - - @parallel (1:nx-2, 1:ny) _conduit_gradient!(T, offset, xc_conduit, yc_conduit, r_conduit, xvi..., sticky_air) -end - -function conduit_gradient_TBuffer!(T, offset, xc_conduit, yc_conduit, r_conduit, xvi, sticky_air) - - @parallel_indices (i, j) function _conduit_gradient_TBuffer!( - T, offset, xc_conduit, yc_conduit, r_conduit, x, y, sticky_air - ) - depth = -y[j] - sticky_air - @inbounds if ((x[i] - xc_conduit)^2 ≤ r_conduit^2) && (0.325*(depth[j] + yc_conduit)^2 ≤ r_conduit^2) - T[i, j] += offset - end - return nothing - end - nx, ny = size(T) - - @parallel (1:nx, 1:ny) _conduit_gradient_TBuffer!(T, offset, xc_conduit, yc_conduit, r_conduit, xvi..., sticky_air) -end - - - -function circular_perturbation!(T, δT, xc_anomaly, yc_anomaly, r_anomaly, xvi, sticky_air) - - @parallel_indices (i, j) function _circular_perturbation!(T, δT, xc_anomaly, yc_anomaly, r_anomaly, x, y, sticky_air) - depth = -y[j] - sticky_air - @inbounds if ((x[i] - xc_anomaly)^2 + (depth[j] + yc_anomaly)^2 ≤ r_anomaly^2) - T[i+1, j] *= δT / 100 + 1 - end - return nothing - end - - nx, ny = size(T) - - @parallel (1:nx-2, 1:ny) _circular_perturbation!(T, δT, xc_anomaly, yc_anomaly, r_anomaly, xvi..., sticky_air) -end - - - -@parallel_indices (i, j) function compute_melt_fraction!(ϕ, rheology, args) - ϕ[i, j] = compute_meltfraction(rheology, ntuple_idx(args, i, j)) - return nothing -end - - -@parallel_indices (I...) function compute_melt_fraction!(ϕ, phase_ratios, rheology, args) -args_ijk = ntuple_idx(args, I...) -ϕ[I...] = compute_melt_frac(rheology, args_ijk, phase_ratios[I...]) -return nothing -end - -@inline function compute_melt_frac(rheology, args, phase_ratios) - return GeoParams.compute_meltfraction_ratio(phase_ratios, rheology, args) -end - - -function phase_change!(phases, particles) - ni = size(phases) - @parallel_indices (I...) function _phase_change!(phases, px, py, index) - - @inbounds for ip in JustRelax.cellaxes(phases) - #quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue - - x = JustRelax.@cell px[ip,I...] - y = (JustRelax.@cell py[ip,I...]) - phase_ij = @cell phases[ip, I...] - if y > 0.0 && (phase_ij == 2.0 || phase_ij == 3.0) - @cell phases[ip, I...] = 4.0 - end - end - return nothing - end - - @parallel (JustRelax.@idx ni) _phase_change!( phases, particles.coords..., particles.index) -end - -function open_conduit!(phases, particles, xc_conduit, yc_conduit, r_conduit, sticky_air) - ni = size(phases) - @parallel_indices (I...) function _open_conduit!(phases, px, py, index, xc_conduit, yc_conduit, r_conduit, sticky_air) - - @inbounds for ip in JustRelax.cellaxes(phases) - #quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue - - x = JustRelax.@cell px[ip,I...] - y = -(JustRelax.@cell py[ip,I...])- sticky_air - - if ((x - xc_conduit)^2 ≤ r_conduit^2) && (0.25*(y - yc_conduit)^2 ≤ r_conduit^2) - JustRelax.@cell phases[ip, I...] = 2.0 - end - - end - return nothing - end - - @parallel (JustRelax.@idx ni) _open_conduit!(phases, particles.coords..., particles.index, xc_conduit, yc_conduit, r_conduit, sticky_air) -end - - -function pureshear!(stokes, εbg, xvi) - stokes.V.Vx .= PTArray([ x*εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray([-y*εbg for _ in 1:nx+2, y in xvi[2]]) - return nothing -end - - -function new_thermal_anomaly(phases, particles, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) - ni = size(phases) - - @parallel_indices (I...) function new_anomlay_particles(phases, px, py, index, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) - @inbounds for ip in JustRelax.cellaxes(phases) - @cell(index[ip, I...]) == 0 && continue - - x = JustRelax.@cell px[ip, I...] - y = -(JustRelax.@cell py[ip, I...]) - sticky_air - - # thermal anomaly - circular - if ((x - xc_anomaly)^2 + (y + yc_anomaly)^2 ≤ r_anomaly^2) - JustRelax.@cell phases[ip, I...] = 3.0 - end - end - return nothing - end - @parallel (JustRelax.@idx ni) new_anomlay_particles(phases, particles.coords..., particles.index, xc_anomaly, yc_anomaly, r_anomaly, sticky_air) -end - -function main2D(igg; figdir=figdir, nx=nx, ny=ny) - - #-------rheology parameters-------------------------------------------------------------- - # plasticity setup - do_DP = true # do_DP=false: Von Mises, do_DP=true: Drucker-Prager (friction angle) - η_reg = 1.0e14 # regularisation "viscosity" for Drucker-Prager - Coh = 10.0 # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) - ϕ = 30.0 * do_DP # friction angle - G0 = 25e9Pa # elastic shear modulus - G_magma = 10e9Pa # elastic shear modulus perturbation - εbg = 2.5e-14 # background strain rate - - soft_C = LinearSoftening((ustrip(Coh)/2, ustrip(Coh)), (0e0, 1e-1)) # softening law - pl = DruckerPrager_regularised(; C=Coh, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # plasticity - - el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring - el_magma = SetConstantElasticity(; G=G_magma, ν=0.3) # elastic spring - creep_rock = LinearViscous(; η=1e21 * Pa * s) - creep_magma = LinearViscous(; η=1e16 * Pa * s) - creep_air = LinearViscous(; η=1e16 * Pa * s) - cutoff_visc = (1e14, 1e24) - β_rock = inv(get_Kb(el)) - β_magma = inv(get_Kb(el_magma)) - Kb = get_Kb(el) - - #-------JustRelax parameters------------------------------------------------------------- - # Domain setup for JustRelax - sticky_air = 5e3 # thickness oif the sticky air layer - ly = 20e3 + sticky_air # domain length in y-direction - lx = 40e3 # domain length in x-direction - li = lx, ly # domain length in x- and y-direction - ni = nx, ny # number of grid points in x- and y-direction - di = @. li / ni # grid step in x- and y-direction - origin = 0.0, -ly # origin coordinates of the domain - grid = Geometry(ni, li; origin = origin) - (; xci, xvi)= grid # nodes at the center and vertices of the cells - #--------------------------------------------------------------------------------------- - - # Set material parameters - MatParam = ( - #Name="UpperCrust" - SetMaterialParams(; - Phase = 1, - Density = PT_Density(ρ0=2700kg/m^3, β=β_rock/Pa), - HeatCapacity = ConstantHeatCapacity(Cp=1050J/kg/K), - Conductivity = ConstantConductivity(k=3.0Watt/K/m), - LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - ShearHeat = ConstantShearheating(1.0NoUnits), - CompositeRheology = CompositeRheology((creep_rock, el, pl, )), - Melting = MeltingParam_Caricchi(), - Elasticity = el, - ), - - #Name="Magma" - SetMaterialParams(; - Phase = 2, - Density = PT_Density(ρ0=2600kg/m^3, β=β_magma/Pa), - HeatCapacity = ConstantHeatCapacity(Cp=1050J/kg/K), - Conductivity = ConstantConductivity(k=1.5Watt/K/m), - LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - ShearHeat = ConstantShearheating(0.0NoUnits), - CompositeRheology = CompositeRheology((creep_magma, el_magma)), - Melting = MeltingParam_Caricchi(), - Elasticity = el_magma, - ), - - #Name="Thermal Anomaly" - SetMaterialParams(; - Phase = 3, - Density = PT_Density(ρ0=2600kg/m^3, β=β_magma/Pa), - HeatCapacity = ConstantHeatCapacity(Cp=1050J/kg/K), - Conductivity = ConstantConductivity(k=1.5Watt/K/m), - LatentHeat = ConstantLatentHeat(Q_L=350e3J/kg), - ShearHeat = ConstantShearheating(0.0NoUnits), - CompositeRheology = CompositeRheology((creep_magma,el_magma)), - Melting = MeltingParam_Caricchi(), - Elasticity = el_magma, - ), - - #Name="Sticky Air" - SetMaterialParams(; - Phase = 4, - Density = ConstantDensity(ρ=10kg/m^3,), - HeatCapacity = ConstantHeatCapacity(Cp=1000J/kg/K), - Conductivity = ConstantConductivity(k=15Watt/K/m), - LatentHeat = ConstantLatentHeat(Q_L=0.0J/kg), - ShearHeat = ConstantShearheating(0.0NoUnits), - CompositeRheology = CompositeRheology((creep_air,)), - ), - ) - - #---------------------------------------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 40, 1 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(3)) - particle_args = (pT, pPhases) - - xc = lx/2 - yc = -ly/3 # origin of thermal anomaly - x_anomaly = lx * 0.5 - y_anomaly = -ly * 0.35 # origin of the small thermal anomaly - - radius = 5e2 # radius of perturbation - a = 15 - b = 5 - r_anomaly = 7.5e2 # radius of perturbation - # r_anomaly = nondimensionalize(1.5km,CharDim) # radius of perturbation - xc_conduit, yc_conduit, r_conduit = (lx*0.5),4.15e3,2.3e3 - - init_phases!(pPhases, particles, xc, yc, a, b, radius, x_anomaly, y_anomaly, r_anomaly,sticky_air) - phase_ratios = PhaseRatio(ni, length(MatParam)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - - # Physical Parameters - geotherm = GeoUnit(0.03K / m) - geotherm = ustrip(Value(geotherm)) - ΔT = geotherm * (ly - sticky_air) # temperature difference between top and bottom of the domain - tempoffset = 0.0 - η = MatParam[2].CompositeRheology[1][1].η.val # viscosity for the Rayleigh number - Cp0 = MatParam[2].HeatCapacity[1].Cp.val # heat capacity - ρ0 = MatParam[2].Density[1].ρ0.val # reference Density - k0 = MatParam[2].Conductivity[1] # Conductivity - G = MatParam[1].Elasticity[1].G.val # Shear Modulus - κ = ustrip(1.5Watt / K / m / (ρ0 * Cp0)) # thermal diffusivity - g = MatParam[1].Gravity[1].g.val # Gravity - - α = MatParam[1].Density[1].α.val # thermal expansion coefficient for PT Density - Ra = ρ0 * g * α * ΔT * 10^3 / (η * κ) # Rayleigh number - dt = dt_diff = 0.5 * min(di...)^2 / κ / 2.01 # diffusive CFL timestep limiter - - # Initialisation - thermal = ThermalArrays(ni) # initialise thermal arrays and boundary conditions - thermal_bc = TemperatureBoundaryConditions(; - no_flux=(left=true, right=true, top=false, bot=false), - periodicity=(left=false, right=false, top=false, bot=false), - ) - - @parallel (@idx ni .+ 1) init_T!( - thermal.T, xvi[2], sticky_air) - - thermal_bcs!(thermal.T, thermal_bc) - - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!( - thermal.Tc, thermal.T - ) - - stokes = StokesArrays(ni, ViscoElastic) # initialise stokes arrays with the defined regime - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL=0.99 / √2.1) #ϵ=1e-4, CFL=1 / √2.1 CFL=0.27 / √2.1 - - args = (; T=thermal.Tc, P=stokes.P, dt=Inf) - - pt_thermal = PTThermalCoeffs( - MatParam, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=5e-2 / √2.1 - ) - # Boundary conditions of the flow - stokes.V.Vx .= PTArray([ (x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray([ -(ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) - - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - periodicity = (left = false, right = false, top = false, bot = false), - ) - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy) - - η = @ones(ni...) # initialise viscosity - η_vep = deepcopy(η) # initialise viscosity for the VEP - G = @fill(MatParam[1].Elasticity[1].G.val, ni...) # initialise shear modulus - ϕ = similar(η) # melt fraction center - - # Buoyancy force - ρg = @zeros(ni...), @zeros(ni...) # ρg[1] is the buoyancy force in the x direction, ρg[2] is the buoyancy force in the y direction - - # Arguments for functions - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt) - @copy thermal.Told thermal.T - - for _ in 1:2 - @parallel (JustRelax.@idx ni) compute_ρg!( - ρg[2], phase_ratios.center, MatParam, (T=thermal.Tc, P=stokes.P) - ) - @parallel init_P!(stokes.P, ρg[2], xci[2], sticky_air) - # @parallel init_P!(stokes.P, ρg[2], xci[2]) - end - - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, MatParam, cutoff_visc - ) - η_vep = copy(η) - - anomaly = 1000.0 .+ 273.0 # temperature anomaly - offset = 350.0.+273.0 - δT = 20.0 # thermal perturbation (in %) - elliptical_anomaly_gradient!( - thermal.T, offset, xc, yc, a, b, radius, xvi, sticky_air - ) - # conduit_gradient!(thermal.T, offset, xc_conduit, yc_conduit, r_conduit, xvi, sticky_air) - circular_perturbation!(thermal.T, δT, x_anomaly, y_anomaly, r_anomaly, xvi, sticky_air) - - # make sure they are the same - thermal.Told .= thermal.T - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!( - thermal.Tc, thermal.T - ) - @parallel (@idx ni) compute_melt_fraction!( - ϕ, phase_ratios.center, MatParam, (T=thermal.Tc, P=stokes.P) - ) - - # Time loop - t, it = 0.0, 0 - local iters - - T_buffer = @zeros(ni .+ 1) - Told_buffer = similar(T_buffer) - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - - grid2particle!(pT, xvi, T_buffer, particles) - @copy stokes.P0 stokes.P - - - # IO ------------------------------------------------ - # if it does not exist, make folder where figures are stored - # if save_vtk - # vtk_dir = figdir*"\\vtk" - # take(vtk_dir) - # end - take(figdir) - # ---------------------------------------------------- - # Plot initial T and η profiles - 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:]), - Yv, - ) - lines!( - ax2, - Array(stokes.P[:]), - Y, - ) - hideydecorations!(ax2) - save(joinpath(figdir, "initial_profile.png"), fig) - fig - end - - while it < 1 #nt - - particle2grid!(T_buffer, pT, xvi, particles) - @views T_buffer[:, end] .= 273.0 - @views thermal.T[2:end-1, :] .= T_buffer - temperature2center!(thermal) - - # if rem(it, 100) == 0 - # # if it > 75 && rem(it, 75) == 0 - # x_anomaly, y_anomaly = lx * 0.6, -ly * 0.27 # Randomly vary center of dike - # r_anomaly = nondimensionalize(0.75km,CharDim) - # δT = 10.0 # thermal perturbation (in %) - # new_thermal_anomaly(pPhases, particles, x_anomaly, y_anomaly, r_anomaly, sticky_air) - # circular_perturbation!(thermal.T, δT, x_anomaly, -y_anomaly, r_anomaly, xvi, sticky_air) - # for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - # copyinn_x!(dst, src) - # end - # @views T_buffer[:, end] .= nondimensionalize(0.0C, CharDim) - # @views thermal.T[2:end-1, :] .= T_buffer - # temperature2center!(thermal) - # grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles.coords) - - # @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - # end - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt) - - # #open the conduit - # if it == 114 - # conduit_gradient_TBuffer!(T_buffer, offset, xc_conduit, -yc_conduit, r_conduit, xvi) - # open_conduit!(pPhases, particles, xc_conduit, yc_conduit, r_conduit) - # grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) - # end - - # @views T_buffer[:, end] .= nondimensionalize(0.0C, CharDim) - # @views T_buffer[args.sticky_air.==4.0] .= nondimensionalize(0.0C, CharDim) - # # @views T_buffer[:, 1] .= maximum(thermal.T) - # @views thermal.T[2:end-1, :] .= T_buffer - # temperature2center!(thermal) - - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, MatParam, cutoff_visc - ) - @parallel (@idx ni) compute_ρg!( - ρg[2], phase_ratios.center, MatParam, (T=thermal.Tc, P=stokes.P) - ) - @copy stokes.P0 stokes.P - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt) - # Stokes solver ----------------- - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - η, - η_vep, - phase_ratios, - MatParam, - args, - dt * 1e-1, - igg; - iterMax = 250e3, - nout = 5e3, - viscosity_cutoff=cutoff_visc, - ) - @parallel (JustRelax.@idx ni) tensor_invariant!(stokes.ε.II, @strain(stokes)...) - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) - ) - dt = compute_dt(stokes, di, dt_diff, igg) - # ------------------------------ - @show dt - - # @parallel (@idx ni) compute_shear_heating!( - # thermal.shear_heating, - # @tensor_center(stokes.τ), - # @tensor_center(stokes.τ_o), - # @strain(stokes), - # phase_ratios.center, - # MatParam, # needs to be a tuple - # dt, - # ) - # Thermal solver --------------- - heatdiffusion_PT!( - thermal, - pt_thermal, - thermal_bc, - MatParam, - args, - dt, - di; - igg=igg, - phase=phase_ratios, - iterMax=150e3, - nout=1e3, - verbose=true, - ) - - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - - # Advection -------------------- - # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # advect particles in memory - shuffle_particles!(particles, xvi, particle_args) - # JustPIC.clean_particles!(particles, xvi, particle_args) - grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) - # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) - - #phase change for particles - phase_change!(pPhases, particles) - # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - @parallel (@idx ni) compute_melt_fraction!( - ϕ, phase_ratios.center, MatParam, (T=thermal.Tc, P=stokes.P) - ) - - @show it += 1 - t += dt - - # # # Plotting ------------------------------------------------------- - if it == 1 || rem(it, 1) == 0 - checkpointing(figdir, stokes, thermal.T, η, t) - - if igg.me == 0 - if save_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - data_v = (; - T = Array(thermal.T[2:end-1, :]), - τxy = Array(stokes.τ.xy), - εxy = Array(stokes.ε.xy), - Vx = Array(Vx_v), - Vy = Array(Vy_v), - ) - data_c = (; - P = Array(stokes.P), - τxx = Array(stokes.τ.xx), - τyy = Array(stokes.τ.yy), - εxx = Array(stokes.ε.xx), - εyy = Array(stokes.ε.yy), - η = Array(η), - ) - save_vtk( - joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), - xvi, - xci, - data_v, - data_c - ) - end - - # Make particles plottable - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:]./1e3 - pyv = ppy.data[:]./1e3 - clr = pPhases.data[:] - idxv = particles.index.data[:]; - - pp = [argmax(p) for p in phase_ratios.center]; - @views η_vep[pp.==4.0] .= NaN - @views stokes.τ.II[pp.==4.0] .= NaN - @views stokes.ε.II[pp.==4.0] .= NaN - - # Make Makie figure - fig = Figure(; size=(2000, 1800), createmissing=true) - ar = li[1] / li[2] - # ar = DataAspect() - - ax0 = Axis( - fig[1, 1:2]; - aspect=ar, - title="t = $(t./1e3) Kyrs", - titlesize=50, - height=0.0, - ) - ax0.ylabelvisible = false - ax0.xlabelvisible = false - ax0.xgridvisible = false - ax0.ygridvisible = false - ax0.xticksvisible = false - ax0.yticksvisible = false - ax0.yminorticksvisible = false - ax0.xminorticksvisible = false - ax0.xgridcolor = :white - ax0.ygridcolor = :white - ax0.ytickcolor = :white - ax0.xtickcolor = :white - ax0.yticklabelcolor = :white - ax0.xticklabelcolor = :white - ax0.yticklabelsize = 0 - ax0.xticklabelsize = 0 - ax0.xlabelcolor = :white - ax0.ylabelcolor = :white - - ax1 = Axis( - fig[2, 1][1, 1]; - aspect=ar, - title=L"T [\mathrm{C}]", - titlesize=40, - yticklabelsize=25, - xticklabelsize=25, - xlabelsize=25, - ) - ax2 = Axis( - fig[2, 2][1, 1]; - aspect=ar, - title=L"\log_{10}(\eta_{vep}) [\mathrm{Pas}]", - titlesize=40, - yticklabelsize=25, - xticklabelsize=25, - xlabelsize=25, - ) - ax3 = Axis( - fig[3, 1][1, 1]; - aspect=ar, - title=L"Phases", - xlabel="Width [km]", - titlesize=40, - yticklabelsize=25, - xticklabelsize=25, - xlabelsize=25, - ) - ax4 = Axis( - fig[3, 2][1, 1]; - aspect=ar, - title=L"\log_{10}(\dot{\varepsilon}_{\textrm{II}}) [\mathrm{s}^{-1}]", - xlabel="Width [km]", - titlesize=40, - yticklabelsize=25, - xticklabelsize=25, - xlabelsize=25, - ) - ax5 = Axis( - fig[4, 1][1, 1]; - aspect=ar, - title=L"\tau_{\textrm{II}} [MPa]", - xlabel="Width [km]", - titlesize=40, - yticklabelsize=25, - xticklabelsize=25, - xlabelsize=25, - ) - # Plot temperature - p1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:batlow) - # Plot effective viscosity - p2 = heatmap!(ax2, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_vep)) , colormap=:glasgow, colorrange= (log10(1e14), log10(1e21))) - # Plot particles phase - p3 = scatter!(ax3, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv])) - # Plot 2nd invariant of strain rate - p4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:roma) - # Plot 2nd invariant of stress - p5 = heatmap!(ax5, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.τ.II)) , colormap=:batlow) - hidexdecorations!(ax1) - hidexdecorations!(ax2) - hidexdecorations!(ax3) - Colorbar( - fig[2, 1][1, 2], p1; height=Relative(0.7), ticklabelsize=25, ticksize=15 - ) - Colorbar( - fig[2, 2][1, 2], p2; height=Relative(0.7), ticklabelsize=25, ticksize=15 - ) - Colorbar( - fig[3, 1][1, 2], p3; height=Relative(0.7), ticklabelsize=25, ticksize=15 - ) - Colorbar( - fig[3, 2][1, 2], p4; height=Relative(0.7), ticklabelsize=25, ticksize=15 - ) - Colorbar( - fig[4, 1][1, 2], p5; height=Relative(0.7), ticklabelsize=25, ticksize=15 - ) - rowgap!(fig.layout, 1) - colgap!(fig.layout, 1) - colgap!(fig.layout, 1) - colgap!(fig.layout, 1) - fig - figsave = joinpath(figdir, @sprintf("%06d.png", it)) - save(figsave, fig) - - 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(; size=(1200, 900)) - ax1 = Axis(fig[1, 1]; aspect=2 / 3, title="T") - ax2 = Axis(fig[1, 2]; aspect=2 / 3, title="Pressure") - - scatter!( - ax1, - Array(thermal.T[2:(end - 1), :][:]), - Yv, - ) - lines!( - ax2, - Array(stokes.P[:]), - Y, - ) - - hideydecorations!(ax2) - save(joinpath(figdir, "pressure_profile_$it.png"), fig) - fig - end - end - end - end - # finalize_global_grid() - -end - - - -figdir = "New_benchmark_extension" -save_vtk = false # set to true to generate VTK files for ParaView -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, 1; init_MPI= true)...) -else - igg -end - -# run main script -main2D(igg; figdir=figdir, nx=nx, ny=ny); - -function plot_particles(particles, pPhases) - p = particles.coords - # pp = [argmax(p) for p in phase_ratios.center] #if you want to plot it in a heatmap rather than scatter - ppx, ppy = p - # pxv = ustrip.(dimensionalize(ppx.data[:], km, CharDim)) - # pyv = ustrip.(dimensionalize(ppy.data[:], km, CharDim)) - pxv = ppx.data[:] - pyv = ppy.data[:] - 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) - Colorbar(f[1,2], h) - f -end diff --git a/PlumeFreeSurface_2D copy.jl b/PlumeFreeSurface_2D copy.jl deleted file mode 100644 index 6617fd6a..00000000 --- a/PlumeFreeSurface_2D copy.jl +++ /dev/null @@ -1,323 +0,0 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - -using JustPIC, JustPIC._2D -const backend = CPUBackend - -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) - -# Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays - -# Velocity helper grids for the particle advection -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) - esc(:($A[$idx_j])) -end - -@parallel function smooth!( - A2::AbstractArray{T,2}, A::AbstractArray{T,2}, fact::Real -) where {T} - @inn(A2) = @inn(A) + 1.0 / 4.1 / fact * (@d2_xi(A) + @d2_yi(A)) - return nothing -end - -# Initial pressure profile - not accurate -# @parallel function init_P!(P, ρg, z) -# @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) -# return nothing -# end - -@parallel_indices (i, j) function init_P!(P, ρg, z) - P[i, j] = sum(abs(ρg[i, jj] * z[jj]) * z[jj] < 0.0 for jj in j:size(P, 2)) - return nothing -end - -function init_phases!(phases, particles) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index) - r = 25e3 - - @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]) - - if 0e0 ≤ depth ≤ 50e3 - @cell phases[ip, i, j] = 1.0 - - else - @cell phases[ip, i, j] = 2.0 - - if ((x - 125e3)^2 + (depth - 100e3)^2 ≤ r^2) - @cell phases[ip, i, j] = 3.0 - end - end - - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) -end -## END OF HELPER FUNCTION ------------------------------------------------------------ - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 50e3 # thickness of sticky air layer - ly = 200e3 + thick_air # domain length in y - lx = 250e3 # domain length in x - ni = nx, ny # number of cells - li = lx, ly # domain length in x- and y- - di = @. li / ni # grid step in x- and -y - origin = 0.0, -ly # origin coordinates (15km f sticky air layer) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - # Name = "Air", - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=1e1), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e17),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Mantle", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Plume", - SetMaterialParams(; - Phase = 3, - Density = ConstantDensity(; ρ=3.2e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e18),)), - Gravity = ConstantGravity(; g=9.81), - ) - ) - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 30, 40, 15 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - init_phases!(pPhases, particles) - 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-6, CFL = 1 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - η = @ones(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = 1e3) - @parallel (@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, (1e16, Inf) - ) - η_vep = copy(η) - - # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - # no_slip = (left = false, right=false, top=false, bot=true), - periodicity = (left = false, right = false, top = false, bot = false), - ) - # for _ in 1:10 - # @parallel smooth!(η_vep, η, 50) - # η_vep, η = η, η_vep - # @views η[1,:] .= η[2,:] - # @views η[end,:] .= η[end-1,:] - # @views η[:, 1] .= η[:, 2] - # @views η[:, end] .= η[:, end-1] - # end - - # Plot initial T and η profiles - let - Y = [y for x in xci[1], y in xci[2]][:] - fig = Figure(size = (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 - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "FreeSurfacePlume" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = dt_max = 10e3 * (3600 * 24 *365.25) - while it < 1 # run only for 5 Myrs - - # Stokes solver ---------------- - args = (; T = thermal.Tc, P = stokes.P, dt=dt) - - @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) - ) - - # for _ in 1:50 - # @parallel smooth!(η_vep, η, 2) - # η_vep, η = η, η_vep - # @views η[1,:] .= η[2,:] - # @views η[end,:] .= η[end-1,:] - # @views η[:, 1] .= η[:, 2] - # @views η[:, end] .= η[:, end-1] - # end - - η .= mean(η) - # η .= 1e21 - - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - η, - η_vep, - phase_ratios, - rheology, - args, - dt, - igg; - iterMax = 70e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) - ) - dt = min(compute_dt(stokes, di), dt_max) #/2 - # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # advect particles in memory - shuffle_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (), (), xvi) - # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - - @show it += 1 - t += dt - - if it == 1 || rem(it, 1) == 0 - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 2 - fig = Figure(size = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)), colormap = :grayC) - arrows!( - ax, - xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :red, - ) - fig - save(joinpath(figdir, "$(it).png"), fig) - - let - Y = [y for x in xci[1], y in xci[2]][:] - fig = Figure(size = (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) - save(joinpath(figdir, "profile_$(it).png"), fig) - fig - end - - end - end - return nothing -end -## END OF MAIN SCRIPT ---------------------------------------------------------------- - -# (Path)/folder where output data and figures are stored -n = 100 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -main(igg, nx, ny) - -# @inline function f1(x_new::T, x_old::T, ν) where {T} -# x_cont = exp((1 - ν) * log(x_old) + ν * log(x_new)) -# return isnan(x_cont) ? 0.0 : x_cont -# end -# @inline f2(x_new, x_old, ν) = (1 - ν) * x_old + ν * x_new - -# n=100 -# x0 = vcat(zeros(n>>>1), fill(100,n>>>1)) -# x1 = LinRange(0, 100, n) -# x2 = LinRange(0, 100, n) -# lines(x0) -# lines!(f1.(x0, x1, 1e-1)) -# lines!(f2.(x0, x1, 1e-1)) -# xn = x1 -# # for i in 1:10 -# # # xn = f1.(x0, xn, 0.5) -# # xn = f2.(x0, xn, 0.5) -# # lines!(xn) -# # end diff --git a/RayleighTaylor2D.jl b/RayleighTaylor2D.jl deleted file mode 100644 index 109f33a8..00000000 --- a/RayleighTaylor2D.jl +++ /dev/null @@ -1,314 +0,0 @@ -using CUDA -CUDA.allowscalar(false) # for safety - -using JustRelax, JustRelax.DataIO, JustPIC -import JustRelax.@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) -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 -------------------------------- -@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 - -# Velocity helper grids for the particle advection -function velocity_grids(xci, xvi, di) - dx, dy = di - yVx = LinRange(xci[2][1] - dy, xci[2][end] + dy, length(xci[2])+2) - xVy = LinRange(xci[1][1] - dx, xci[1][end] + dx, length(xci[1])+2) - grid_vx = xvi[1], yVx - grid_vy = xVy, xvi[2] - - 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) - esc(:($A[$idx_j])) -end - -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - -function init_phases!(phases, particles, A) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) - - f(x, A, λ) = A * sin(π*x/λ) - - @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 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 - - end - return nothing - end - - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, A) -end -## END OF HELPER FUNCTION ------------------------------------------------------------ - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function RT_2D(igg, nx, ny) - - # Physical domain ------------------------------------ - thick_air = 100e3 # thickness of sticky air layer - ly = 500e3 + thick_air # domain length in y - lx = 500e3 # domain length in x - ni = nx, ny # number of cells - li = lx, ly # domain length in x- and y- - di = @. li / ni # grid step in x- and -y - origin = 0.0, -ly # origin coordinates (15km f sticky air layer) - xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = rheology = ( - # Name = "Air", - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(; ρ=1e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e19),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Crust", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(; ρ=3.3e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "Mantle", - SetMaterialParams(; - Phase = 3, - Density = ConstantDensity(; ρ=3.2e3), - CompositeRheology = CompositeRheology((LinearViscous(; η=1e20),)), - Gravity = ConstantGravity(; g=9.81), - ) - ) - dt = 1 # diffusive CFL timestep limiter - # ---------------------------------------------------- - - # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 120, 1 - particles = init_particles_cellarrays( - nxcell, max_xcell, min_xcell, xvi[1], xvi[2], di[1], di[2], nx, ny - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_particle_fields_cellarrays(particles, Val(3)) - particle_args = (pT, pPhases) - - # Elliptical temperature 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) - # ---------------------------------------------------- - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) - # ---------------------------------------------------- - - # Buoyancy forces & rheology - ρg = @zeros(ni...), @zeros(ni...) - η = @zeros(ni...) - 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) - ) - η_vep = copy(η) - - # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - # no_slip = (left = false, right=false, top=false, bot=true), - 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 - - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - figdir = "FreeSurface" - take(figdir) - - # Time loop - t, it = 0.0, 0 - dt = 1 - dt_max = 1e3 * (3600 * 24 *365.25) - while it < 50 # run only for 5 Myrs - - # 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) - ) - - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - η, - η_vep, - phase_ratios, - rheology, - args, - dt, - igg; - iterMax = 10e3, - 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 -------------------- - # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # advect particles in memory - shuffle_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) - # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - - @show it += 1 - t += dt - - # if it == 1 || rem(it, 5) == 0 - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - nt = 5 - - fig = Figure(resolution = (900, 900), title = "t = $t") - ax = Axis(fig[1,1], aspect = ar, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)), colormap = :grayC) - arrows!( - ax, - xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., - lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), - color = :red, - ) - fig - save(joinpath(figdir, "$(it).png"), fig) - - # end - - dt = min(compute_dt(stokes, di) / 10, dt_max) - - end - return fig -end -## END OF MAIN SCRIPT ---------------------------------------------------------------- - -# (Path)/folder where output data and figures are stored -n = 64 -nx = n -ny = n -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) -else - igg -end - -RT_2D(igg, nx, ny)