diff --git a/Blober.jl b/Blober.jl new file mode 100644 index 00000000..5753639d --- /dev/null +++ b/Blober.jl @@ -0,0 +1,515 @@ +# const isCUDA = false +const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO +import JustRelax.@cell + +const backend_JR = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustRelax.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using ParallelStencil, ParallelStencil.FiniteDifferences3D + +@static if isCUDA + @init_parallel_stencil(CUDA, Float64, 3) +else + @init_parallel_stencil(Threads, Float64, 3) +end + +using JustPIC, JustPIC._3D +# 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 = @static if isCUDA + CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +else + JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +end + +using GeoParams, CairoMakie, CellArrays + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +function init_phases!(phases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, sticky_air,top, bottom) + ni = size(phases) + + @parallel_indices (I...) function init_phases!( + phases, px, py, pz, index, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, sticky_air, top, bottom + ) + @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + # quick escape + JustRelax.@cell(index[ip, I...]) == 0 && continue + + x = JustRelax.@cell px[ip, I...] + y = JustRelax.@cell py[ip, I...] + z = -(JustRelax.@cell pz[ip, I...]) - sticky_air + if top ≤ z ≤ bottom + @cell phases[ip, I...] = 1.0 # crust + end + + # thermal anomaly - circular + if ((x - xc_anomaly)^2 + (y - yc_anomaly)^2 + (z + zc_anomaly)^2 ≤ r_anomaly^2) + JustRelax.@cell phases[ip, I...] = 2.0 + end + + if z < top + JustRelax.@cell phases[ip, I...] = 3.0 + end + end + return nothing + end + + @parallel (@idx ni) init_phases!( + phases, + particles.coords..., + particles.index, + xc_anomaly, + yc_anomaly, + zc_anomaly, + r_anomaly, + sticky_air, + top, + bottom, + ) +end + +# Initial thermal profile +@parallel_indices (i, j, k) function init_T!(T, y, sticky_air, top, bottom, dTdz, offset) + I = i, j, k + depth = -y[k] - sticky_air + + if depth < top + T[I...] = offset + + else#if top ≤ (depth) < bottom + dTdZ = dTdz + offset = offset + T[I...] = (depth) * dTdZ + offset + end + + return nothing +end + + +function circular_perturbation!(T, δT, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, xvi, sticky_air) + + @parallel_indices (i, j, k) function _circular_perturbation!( + T, δT, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, x, y, z, sticky_air + ) + depth = -z[k] - sticky_air + @inbounds if ((x[i] - xc_anomaly)^2 + (y[j] - yc_anomaly)^2 + (depth + zc_anomaly)^2 ≤ r_anomaly^2) + # T[i, j, k] *= δT / 100 + 1 + T[i, j, k] += δT + end + return nothing + end + + ni = size(T) + + @parallel (@idx ni) _circular_perturbation!( + T, δT, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, xvi..., sticky_air + ) +end + +function init_rheology(CharDim; is_compressible = false) + # plasticity setup + G0 = 6.0e11Pa # elastic shear modulus + G_magma = 6.0e11Pa # elastic shear modulus perturbation + + if is_compressible == true + el = SetConstantElasticity(; G=G0, ν=0.25) # elastic spring + el_magma = SetConstantElasticity(; G=G_magma, ν=0.25)# elastic spring + else + el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring + el_magma = SetConstantElasticity(; G=G_magma, ν=0.5) # elastic spring + end + β_rock = inv(get_Kb(el)) + β_magma = inv(get_Kb(el_magma)) + creep_rock = LinearViscous(; η=1e20 * Pa * s) + creep_magma = LinearViscous(; η=1e20 * Pa * s) + g = 9.81m/s^2 + rheology = ( + #Name="UpperCrust" + SetMaterialParams(; + Phase = 1, + # Density = ConstantDensity(; ρ=2650kg / m^3), + Density = PT_Density(; ρ0=2650kg / m^3, α=3e-5 / K, T0=0.0C, β=β_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, )), + Melting = MeltingParam_Caricchi(), + Gravity = ConstantGravity(; g=g), + Elasticity = el, + CharDim = CharDim, + ), + + #Name="Magma" + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=2650kg / m^3, T0=0.0C, β=β_magma / Pa), + # Density = ConstantDensity(; ρ=2625kg / m^3), + 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, )), + Melting = MeltingParam_Caricchi(), + Gravity = ConstantGravity(; g=g), + Elasticity = el_magma, + CharDim = CharDim, + ), + + # #Name="Sticky Air" + # SetMaterialParams(; + # Phase = 3, + # Density = ConstantDensity(ρ=1kg/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,)), + # Gravity = ConstantGravity(; g=g), + # CharDim = CharDim, + # ), + ) + +end + +function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = false) + + # Characteristic lengths for non dimensionalization + L_km = 100 + CharDim = GEO_units(;length=L_km, viscosity=1e21, temperature = 1e3C) + + #-------JustRelax parameters------------------------------------------------------------- + # Domain setup for JustRelax + L = nondimensionalize(L_km * km,CharDim) + sticky_air = nondimensionalize(0km, CharDim) # thickness of the sticky air layer + lz = L + sticky_air # domain length in y-direction + lx = ly = L # domain length in x-direction + li = lx, ly, lz # domain length in x- and y-direction + ni = nx, ny, nz # number of grid points in x- and y-direction + di = @. li / ni # grid step in x- and y-direction + origin = nondimensionalize(0.0km,CharDim), nondimensionalize(0.0km,CharDim), -lz # origin coordinates of the domain + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + εbg = nondimensionalize(0.0 / s, CharDim) # background strain rate + #--------------------------------------------------------------------------------------- + + # Physical Parameters + rheology = init_rheology(CharDim; is_compressible=true) + cutoff_visc = nondimensionalize((1e16Pa*s, 1e24Pa*s),CharDim) + κ = (4 / (rheology[1].HeatCapacity[1].Cp * rheology[1].Density[1].ρ)) + dt = dt_diff = (0.5 * min(di...)^2 / κ / 2.01) # diffusive CFL timestep limiter + + # Initialize particles ------------------------------- + nxcell = 25 + max_xcell = 40 + min_xcell = 15 + particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi, di, ni) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Circular temperature anomaly-------------------------- + x_anomaly = lx * 0.5 + y_anomaly = ly * 0.5 + z_anomaly = -lz * 0.5 # origin of the small thermal anomaly + r_anomaly = nondimensionalize(10km, CharDim) # radius of perturbation + anomaly = nondimensionalize(50K, CharDim) # thermal perturbation (in K) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles, x_anomaly, y_anomaly, z_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), lz) + phase_ratios_center!(phase_ratios, particles, grid, pPhases) + + # Initialisation of thermal profile + thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, front = true, back = true, top = false, bot = false), + ) + @parallel (@idx ni .+ 1) init_T!( + thermal.T, + xvi[3], + sticky_air, + nondimensionalize(0e0km,CharDim), + lz, + nondimensionalize((723 - 273)K,CharDim) / lz, + nondimensionalize(273K,CharDim) + ) + circular_perturbation!( + thermal.T, anomaly, x_anomaly, y_anomaly, z_anomaly, r_anomaly, xvi, sticky_air + ) + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) # initialise stokes arrays with the defined regime + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL=0.9 / √3.1) + # ---------------------------------------------------- + + flow_bcs = FlowBoundaryConditions(; + free_slip = (left=true, right=true, front=true, back=true, top=true, bot=true), + no_slip = (left=false, right=false, front=false, back=false, top=false, bot=false), + free_surface = true, + ) + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) + + # Buoyancy force & viscosity + args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + ρbg = nondimensionalize(2650kg / m^3, CharDim) + ρg = @zeros(ni...), @zeros(ni...), @zeros(ni...) # ρg[1] is the buoyancy force in the x direction, ρg[2] is the buoyancy force in the y direction + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-2 / √3.1 + ) + Plitho = reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2) + + # Arguments for functions + @copy thermal.Told thermal.T + + # 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) + # ---------------------------------------------------- + + # Smooth out thermal field --------------------------- + for _ in 1:1 + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + nondimensionalize(50e3 * 3600 * 24 * 365.25 * s, CharDim), + di; + kwargs =(; + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e3, + verbose = true, + ) + ) + end + + # Plot initial T and η profiles + # let + # Zv = [z for _ in xvi[1], _ in xvi[2], z in xvi[3]][:] + # Z = [z for _ in xci[1], _ in xci[2], z in xci[3]][:] + # 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(ustrip.(dimensionalize(thermal.T[:], C, CharDim))), + # ustrip.(dimensionalize(Zv, km, CharDim)), + # ) + # scatter!( + # ax2, + # # Array(ustrip.(dimensionalize(stokes.P[:], MPa, CharDim))), + # Array(ρg[end][:]), + # ustrip.(dimensionalize(Z, km, CharDim)), + # ) + # hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + # fig + # end + + # Time loop + t, it = 0.0, 0 + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni .+ 1...) + Vy_v = @zeros(ni .+ 1...) + Vz_v = @zeros(ni .+ 1...) + end + + dt₀ = similar(stokes.P) + grid2particle!(pT, xvi, thermal.T, particles) + + @copy stokes.P0 stokes.P + @copy thermal.Told thermal.T + Tsurf, Tbot = extrema(thermal.T) + # Tbot = thermal.T[1, 1, 1] + + while it < 25 + + # # Update buoyancy and viscosity - + # args = (; T=thermal.Tc, P=stokes.P, dt=Inf) + # Update buoyancy and viscosity - + Plitho .= reverse(cumsum(reverse((ρg[3]).* di[2], dims=3), dims=3), dims=3) + Plitho .= stokes.P .+ Plitho .- minimum(stokes.P) + + args = (; T = thermal.Tc, P = Plitho, dt=Inf, ρbg = ρbg * rheology[1].Gravity[1].g) + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + + # Stokes solver ----------------- + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 50e3, + free_surface = false, + nout = 1e3, + viscosity_cutoff = cutoff_visc, + ) + ) + tensor_invariant!(stokes.ε) + dt = compute_dt(stokes, di, dt_diff, igg) + # -------------------------------- + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs =(; + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e3, + verbose = true, + ) + ) + # subgrid diffusion + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # 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, (pT, ), (thermal.T,), xvi) + # update phase ratios + phase_ratios_center!(phase_ratios, particles, grid, pPhases) + + particle2grid!(thermal.T, pT, xvi, particles) + @views thermal.T[:, :, end] .= Tsurf + @views thermal.T[:, :, 1] .= Tbot + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + thermal.ΔT .= thermal.T .- thermal.Told + vertex2center!(thermal.ΔTc, thermal.ΔT) + + @show it += 1 + t += dt + + # # # Plotting ------------------------------------------------------- + if it == 1 || rem(it, 10) == 0 + (; η) = stokes.viscosity + # checkpointing(figdir, stokes, thermal.T, η, t) + + if igg.me == 0 + if do_vtk + velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + + data_v = (; + T = Array(ustrip.(dimensionalize(thermal.T, C, CharDim))), + τxy= Array(ustrip.(dimensionalize(stokes.τ.xy, s^-1, CharDim))), + εxy= Array(ustrip.(dimensionalize(stokes.ε.xy, s^-1, CharDim))), + ) + data_c = (; + P = Array(ustrip.(dimensionalize(stokes.P,MPa,CharDim))), + τxx = Array(ustrip.(dimensionalize(stokes.τ.xx, MPa,CharDim))), + τyy = Array(ustrip.(dimensionalize(stokes.τ.yy,MPa,CharDim))), + τzz = Array(ustrip.(dimensionalize(stokes.τ.zz,MPa,CharDim))), + τII = Array(ustrip.(dimensionalize(stokes.τ.II, MPa, CharDim))), + εxx = Array(ustrip.(dimensionalize(stokes.ε.xx, s^-1,CharDim))), + εyy = Array(ustrip.(dimensionalize(stokes.ε.yy, s^-1,CharDim))), + εzz = Array(ustrip.(dimensionalize(stokes.ε.zz, s^-1,CharDim))), + εII = Array(ustrip.(dimensionalize(stokes.ε.II, s^-1,CharDim))), + η = Array(ustrip.(dimensionalize(η,Pa*s,CharDim))), + ) + velocity_v = ( + Array(ustrip.(dimensionalize(Vx_v,cm/yr,CharDim))), + Array(ustrip.(dimensionalize(Vy_v, cm/yr, CharDim))), + Array(ustrip.(dimensionalize(Vz_v, cm/yr, CharDim))), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + # let + # Zv = [z for _ in 1:nx+1, _ in 1:ny+1, z in ustrip.(dimensionalize(xvi[3],km,CharDim))][:] + # Z = [z for _ in 1:nx , _ in 1:ny , z in ustrip.(dimensionalize(xci[3],km,CharDim))][:] + # 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") + # a3 = Axis(fig[2, 1]; aspect = 2 / 3, title="τII") + + # scatter!(ax1, ustrip.(dimensionalize((Array(thermal.T)), C, CharDim))[:] , Zv, markersize = 4) + # scatter!(ax2, ustrip.(dimensionalize((Array(stokes.P)), MPa, CharDim))[:] , Z , markersize = 4) + # scatter!(a3, ustrip.(dimensionalize(Array(stokes.τ.II), MPa, CharDim))[:], Z , markersize = 4) + + # hideydecorations!(ax2) + # save(joinpath(figdir, "pressure_profile_$it.png"), fig) + # fig + # end + end + end + end + + finalize_global_grid() + return nothing +end + +figdir = "/scratch/snx3000/ademonts/Blober" +figdir = "Blober100" +do_vtk = true # set to true to generate VTK files for ParaView +n = 100 +nx = n +ny = n +nz = n +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, nz; init_MPI=true)...) +else + igg +end + +# run main script +main3D(igg; figdir=figdir, nx=nx, ny=ny, nz=nz, do_vtk = do_vtk); diff --git a/Project.toml b/Project.toml index d13de3db..8c3ba8dd 100644 --- a/Project.toml +++ b/Project.toml @@ -31,7 +31,7 @@ AMDGPU = "0.6, 0.7, 0.8" Adapt = "3" CUDA = "4.4.1, 5" CellArrays = "0.2" -GeoParams = "0.5.8" +GeoParams = "0.6" HDF5 = "0.17.1" ImplicitGlobalGrid = "0.15.0" MPI = "0.20" diff --git a/Subduction_GMG_2D.jl b/Subduction_GMG_2D.jl new file mode 100644 index 00000000..b40cc1fb --- /dev/null +++ b/Subduction_GMG_2D.jl @@ -0,0 +1,58 @@ +#= +# Creating 2D numerical model setups + +### Aim +The aim of this tutorial is to show you how to create 2D numerical model setups that can be used as initial setups for other codes. + +=# + + +#= +### 2D Subduction setup + +Lets start with creating a 2D model setup in cartesian coordinates, which uses the `CartData` data structure +=# +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + x = range(-1000,1000, nx); + z = range(-660,0, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(1350.0, nx, 1, nz); + lith = LithosphericPhases(Layers=[15 20 55], Phases=[3 4 5], Tlab=1250) + # mantle = LithosphericPhases(Phases=[1]) + + # add_box!(Phases, Temp, Grid2D; xlim=(-1000, 1000), zlim=(-600.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + # Phases .= 0 + + # Lets start with defining the horizontal part of the overriding plate. Note that we define this twice with different thickness to deal with the bending subduction area: + add_box!(Phases, Temp, Grid2D; xlim=(200,1000), zlim=(-150.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + add_box!(Phases, Temp, Grid2D; xlim=(0,200), zlim=(-50.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + + # The horizontal part of the oceanic plate is as before: + v_spread_cm_yr = 3 #spreading velocity + lith = LithosphericPhases(Layers=[15 55], Phases=[1 2], Tlab=1250) + add_box!(Phases, Temp, Grid2D; xlim=(-1000,0.0), zlim=(-150.0, 0.0), phase = lith, T=SpreadingRateTemp(SpreadingVel=v_spread_cm_yr)); + + # Yet, now we add a trench as well. The starting thermal age at the trench is that of the horizontal part of the oceanic plate: + AgeTrench_Myrs = 1000e3/(v_spread_cm_yr/1e2)/1e6 #plate age @ trench + + # We want to add a smooth transition from a halfspace cooling 1D thermal profile to a slab that is heated by the surrounding mantle below a decoupling depth `d_decoupling`. + T_slab = LinearWeightedTemperature( F1=HalfspaceCoolingTemp(Age=AgeTrench_Myrs), F2=McKenzie_subducting_slab(Tsurface=0,v_cm_yr=v_spread_cm_yr, Adiabat = 0.0)) + + # in this case, we have a more reasonable slab thickness: + trench = Trench(Start=(0.0,-100.0), End=(0.0,100.0), Thickness=100.0, θ_max=30.0, Length=600, Lb=200, + WeakzoneThickness=15, WeakzonePhase=6, d_decoupling=125); + add_slab!(Phases, Temp, Grid2D, trench, phase = lith, T=T_slab); + + # Lithosphere-asthenosphere boundary: + ind = findall(Temp .> 1250 .&& (Phases.==2 .|| Phases.==5)); + Phases[ind] .= 0; + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + +end + \ No newline at end of file diff --git a/_typos.toml b/_typos.toml index ac343ee7..6e0f2329 100644 --- a/_typos.toml +++ b/_typos.toml @@ -2,3 +2,4 @@ Ths = "Ths" oly = "oly" iy = "iy" +nd = "nd" diff --git a/docs/src/man/ShearBands.md b/docs/src/man/ShearBands.md index 33273296..3e3fca49 100644 --- a/docs/src/man/ShearBands.md +++ b/docs/src/man/ShearBands.md @@ -10,7 +10,6 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO const backend_JR = CPUBackend ``` - We will also use `ParallelStencil.jl` to write some device-agnostic helper functions: ```julia using ParallelStencil @@ -54,7 +53,7 @@ dt = η0/G0/4.0 # assumes Maxwell time of 4 el_bg = ConstantElasticity(; G=G0, Kb=4) el_inc = ConstantElasticity(; G=Gi, Kb=4) visc = LinearViscous(; η=η0) -pl = DruckerPrager_regularised(; # non-regularized plasticity +pl = DruckerPrager_regularised(; # regularized plasticity C = C, ϕ = ϕ, η_vp = η_reg, diff --git a/docs/src/man/equations.md b/docs/src/man/equations.md index 0154a877..303d75d3 100644 --- a/docs/src/man/equations.md +++ b/docs/src/man/equations.md @@ -1,7 +1,55 @@ +# Stokes equations + +Stokes equations for compressible flow are described by + +$$\nabla\cdot\boldsymbol{\tau} + \nabla p = \boldsymbol{f} $$ + +$$\nabla\cdot\boldsymbol{v} + \beta \frac{\partial p}{\partial t} + \alpha \frac{\partial T}{\partial t} = 0$$ + +where $\nabla = \left(\frac{\partial }{\partial x_i}, ..., \frac{\partial }{\partial x_n} \right)$ is the nabla operator, $\boldsymbol{\tau}$ is the deviatoric stress tensor, $p$ is the pressure, $\boldsymbol{f}$ is the body forces vector, $\boldsymbol{v}$ is the velocity field, and $\beta$ is in the inverse of the bulk modulus. In the simple case of a linear rheology, the constitutive equation for an isotropic flow is $\boldsymbol{\tau} = 2\eta\dot{\boldsymbol{\varepsilon}}$, where $\dot{\boldsymbol{\varepsilon}}$ is the deviatoric strain rate tensor. Heat diffusion + +$$\rho C_p \frac{\partial T}{\partial t} = - \nabla q + Q$$ + +$$q = - K \nabla T$$ + +where $\rho$ is the density, $C_p$ is the heat diffusion, $K$ is the heat conductivity, $Q$ is the sum of any amount of source terms, and $T$ is the temperature. + # Pseudo-transient iterative method +The pseudo-transient method consists in augmenting the right-hand-side of the target PDE with a pseudo-time derivative (where $\psi$ is the pseudo-time) of the primary variables. + +## Stokes + +The pseudo-transient formulation of the Stokes equations yields: + +$$\widetilde{\rho}\frac{\partial \boldsymbol{u}}{\partial \psi} + \nabla\cdot\boldsymbol{\tau} - \nabla p = \boldsymbol{f}$$ + +$$\frac{1}{\widetilde{K}}\frac{\partial p}{\partial \psi} + \nabla\cdot\boldsymbol{v} = \beta \frac{\partial p}{\partial t} + \alpha \frac{\partial T}{\partial t}$$ + +We also do a continuation on the constitutive law: + +$$\frac{1}{2\widetilde{G}} \frac{\partial\boldsymbol{\tau}}{\partial\psi}+ \frac{1}{2G}\frac{D\boldsymbol{\tau}}{Dt} + \frac{\boldsymbol{\tau}}{2\eta} = \dot{\boldsymbol{\varepsilon}}$$ + +where the wide tile denotes the effective damping coefficients and $\psi$ is the pseudo-time step. These are defined as in [Raess et al, 2022](https://doi.org/10.5194/gmd-2021-411): + +$$\widetilde{\rho} = Re\frac{\eta}{\widetilde{V}L}, \qquad + \widetilde{G} = \frac{\widetilde{\rho} \widetilde{V}^2}{r+2}, \qquad + \widetilde{K} = r \widetilde{G}$$ + +and + +$$\widetilde{V} = \sqrt{ \frac{\widetilde{K} +2\widetilde{G}}{\widetilde{\rho}}}, \qquad + r = \frac{\widetilde{K}}{\widetilde{G}}, \qquad + Re = \frac{\widetilde{\rho}\widetilde{V}L}{\eta}$$ + +where the P-wave $\widetilde{V}=V_p$ is the characteristic velocity scale for Stokes, and $Re$ is the Reynolds number. + ## Heat diffusion -## Stokes equations +The pseudo-transient heat-diffusion equation is: + +$$\widetilde{\rho}\frac{\partial T}{\partial \psi} + \rho C_p \frac{\partial T}{\partial t} = \nabla(K\nabla T) = -\nabla q$$ + +We use a second order pseudo-transient scheme were continuation is also done on the flux, so that: -## Constitutive equations +$$\widetilde{\theta}\frac{\partial q}{\partial \psi} + q = -K\nabla T$$ diff --git a/docs/src/man/subduction2d/rheology.md b/docs/src/man/subduction2d/rheology.md new file mode 100644 index 00000000..3a7b9282 --- /dev/null +++ b/docs/src/man/subduction2d/rheology.md @@ -0,0 +1,76 @@ +# Using GeoParams.jl to define the rheology of the material phases + +We will use the same physical parameters as the ones defined in [Hummel, et al 2024](https://doi.org/10.5194/se-15-567-2024). + +The thermal expansion coefficient $\alpha$ and heat capacity $C_p$ are the same for all the material phases + +```julia +α = 2.4e-5 # 1 / K +Cp = 750 # J / kg K +``` + +The density of all the phases is constant, except for the oceanic lithosphere, which uses the pressure and temperature dependent equation of state $\rho = \rho_0 \left(1 - \alpha (T-T_0) - \beta (P-P_0) \right)$, where $\rho_0 = \rho (T=1475 \text{C}^{\circ})=3200 \text{kg/m}^3$.which corresponds to the `PT_Density` object from `GeoParams.jl`: + +```julia +density_lithosphere = PT_Density(; ρ0=3.2e3, α = α, β = 0e0, T0 = 273+1474) +``` + +We will run the case where the rheology is given by a combination of dislocation and diffusion creep for wet olivine, + +```julia +disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) +diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) +``` + +and where plastic failure is given by the formulation from [Duretz et al, 2021](https://doi.org/10.1029/2021GC009675) +```julia +# non-regularized plasticity +ϕ = asind(0.1) +C = 1e6 # Pa +plastic_model = DruckerPrager_regularised(; C = C, ϕ = ϕ, η_vp=η_reg, Ψ=0.0) +``` + +Finally we define the rheology objects of `GeoParams.jl` +```julia +rheology = ( + SetMaterialParams(; + Name = "Asthenoshpere", + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ), + SetMaterialParams(; + Name = "Oceanic lithosphere", + Phase = 2, + Density = density_lithosphere, + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( + ( + disl_wet_olivine, + diff_wet_olivine, + plastic_model, + ) + ), + ), + SetMaterialParams(; + Name = "oceanic crust", + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e20),)), + ), + SetMaterialParams(; + Name = "StickyAir", + Phase = 4, + Density = ConstantDensity(; ρ=1), # water density + HeatCapacity = ConstantHeatCapacity(; Cp = 1e34), + Conductivity = ConstantConductivity(; k = 3), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e19),)), + ), +) +``` \ No newline at end of file diff --git a/docs/src/man/subduction2d/setup.md b/docs/src/man/subduction2d/setup.md new file mode 100644 index 00000000..35f65db2 --- /dev/null +++ b/docs/src/man/subduction2d/setup.md @@ -0,0 +1,101 @@ +# Model setup +As described in the original [paper](https://doi.org/10.5194/se-15-567-2024), the domain consists of a Cartesian box of $\Omega \in [0, 3000] \times [0, -660]$ km, with two 80km thick oceanic plates over the asthenospheric mantle. + +We will use GeophysicalModelGenerator.jl to generate the initial geometry, material phases, and thermal field of our models. We will start by defining the dimensions and resolution of our model, as well as initializing the `Grid2D` object and two arrays `Phases` and `Temp` that host the material phase (given by an integer) and the thermal field, respectively. + +```julia +nx, nz = 512, 218 # number of cells per dimension +Tbot = 1474.0 # [Celsius] +model_depth = 660 # [km] +air_thickness = 10 # [km] +Lx = 3000 # model length [km] +x = range(0, Lx, nx); +z = range(-model_depth, air_thickness, nz); +Grid2D = CartData(xyz_grid(x,0,z)) +Phases = zeros(Int64, nx, 1, nz); +Temp = fill(Tbot, nx, 1, nz); +``` + +In this model we have four material phases with their respective phase numbers: + +| Material | Phase number | +| :---------------- | :----------: | +| asthenosphere | 0 | +| oceanic lithosphere | 1 | +| oceanic crust | 3 | +| sticky air | 4 | + +We will start by initializing the model as asthenospheric mantle, with a thermal profile given by the half-space cooling model with an age of 80 Myrs. + +```julia +add_box!( + Phases, + Temp, + Grid2D; + xlim = (0, Lx), + zlim = (-model_depth, 0.0), + phase = LithosphericPhases(Layers=[], Phases=[0]), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=80,Adiabat=0.4) +) +``` +![](setup_1.png) + +Next we add a horizontal 80km thick oceanic lithosphere. Note that we leave a 100km buffer zone next to the vertical boundaries of the domain, to facilitate the sliding of the oceanic plates. +```julia +add_box!( + Phases, + Temp, + Grid2D; + xlim = (100, Lx-100), # 100 km buffer zones on both sides + zlim = (-model_depth, 0.0), + phase = LithosphericPhases(Layers=[80], Phases=[1 0]), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=80, Adiabat=0.4) +) +``` +![](setup_2.png) + +As in the original paper, we add a 8km thick crust on top of the subducting oceanic plate. +```julia +# Add right oceanic plate crust +add_box!( + Phases, + Temp, + Grid2D; + xlim = (Lx-1430, Lx-200), + zlim = (-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[8 72], Phases=[2 1 0]), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=80, Adiabat=0.4) +) +``` +![](setup_3.png) + +And finally we add the subducting slab, with the trench located at 1430km from the right-hand-side boundary. + +```julia +add_box!( + Phases, + Temp, + Grid2D; + xlim = (Lx-1430, Lx-1430-250), + zlim = (-80, 0.0), + Origin = (nothing, StrikeAngle=0, DipAngle=-30), + phase = LithosphericPhases(Layers=[8 72], Phases=[2 1 0]), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=80, Adiabat=0.4) +) +``` +![](setup_4.png) + +```julia +surf = Grid2D.z.val .> 0.0 +@views Temp[surf] .= 20.0 +@views Phases[surf] .= 3 +``` +![](setup_5.png) + +```julia +li = (abs(last(x)-first(x)), abs(last(z)-first(z))) .* 1e3 # in meters +origin = (x[1], z[1]) .* 1e3 # lower-left corner of the domain +Phases = Phases[:,1,:] .+ 1 # +1 because Julia is 1-indexed +Temp = Temp[:,1,:].+273 # in Kelvin +``` \ No newline at end of file diff --git a/docs/src/man/subduction2d/setup_1.png b/docs/src/man/subduction2d/setup_1.png new file mode 100644 index 00000000..13f5edb1 Binary files /dev/null and b/docs/src/man/subduction2d/setup_1.png differ diff --git a/docs/src/man/subduction2d/setup_2.png b/docs/src/man/subduction2d/setup_2.png new file mode 100644 index 00000000..6b30c071 Binary files /dev/null and b/docs/src/man/subduction2d/setup_2.png differ diff --git a/docs/src/man/subduction2d/setup_3.png b/docs/src/man/subduction2d/setup_3.png new file mode 100644 index 00000000..4db3da89 Binary files /dev/null and b/docs/src/man/subduction2d/setup_3.png differ diff --git a/docs/src/man/subduction2d/setup_4.png b/docs/src/man/subduction2d/setup_4.png new file mode 100644 index 00000000..63dc0664 Binary files /dev/null and b/docs/src/man/subduction2d/setup_4.png differ diff --git a/docs/src/man/subduction2d/setup_5.png b/docs/src/man/subduction2d/setup_5.png new file mode 100644 index 00000000..6910d24b Binary files /dev/null and b/docs/src/man/subduction2d/setup_5.png differ diff --git a/docs/src/man/subduction2d/subduction2D.md b/docs/src/man/subduction2d/subduction2D.md new file mode 100644 index 00000000..ed55ec69 --- /dev/null +++ b/docs/src/man/subduction2d/subduction2D.md @@ -0,0 +1,247 @@ +# 2D subduction + +Model setups taken from [Hummel, et al 2024](https://doi.org/10.5194/se-15-567-2024). + +# Model setup +We will use GeophysicalModelGenerator.jl to generate the initial geometry, material phases, and thermal field of our models. + + +# Initialize packages + +Load JustRelax necessary modules and define backend. +```julia +using CUDA +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = CUDABackend +``` + +For this benchmark we will use particles to track the advection of the material phases and their information. For this, we will use [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl) +```julia +using JustPIC, JustPIC._2D +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +``` + +We will also use `ParallelStencil.jl` to write some device-agnostic helper functions: +```julia +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +``` +and will use [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl/tree/main) to define and compute physical properties of the materials: +```julia +using GeoParams +``` + +# Script + +## Model domain +```julia +nx = ny = 64 # number of cells per dimension +igg = IGG( + init_global_grid(nx, ny, 1; init_MPI= true)... +) # initialize MPI grid +ly = 1.0 # domain length in y +lx = ly # 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, 0.0 # origin coordinates +grid = Geometry(ni, li; origin = origin) +(; xci, xvi) = grid # nodes at the center and vertices of the cells +dt = Inf +``` + +## Physical properties using GeoParams +```julia +τ_y = 1.6 # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) +ϕ = 30 # friction angle +C = τ_y # Cohesion +η0 = 1.0 # viscosity +G0 = 1.0 # elastic shear modulus +Gi = G0/(6.0-4.0) # elastic shear modulus perturbation +εbg = 1.0 # background strain-rate +η_reg = 8e-3 # regularisation "viscosity" +dt = η0/G0/4.0 # assumes Maxwell time of 4 +el_bg = ConstantElasticity(; G=G0, Kb=4) +el_inc = ConstantElasticity(; G=Gi, Kb=4) +visc = LinearViscous(; η=η0) +pl = DruckerPrager_regularised(; # non-regularized plasticity + C = C, + ϕ = ϕ, + η_vp = η_reg, + Ψ = 0 +) +``` +## Rheology +```julia + rheology = ( + # Low density phase + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + CompositeRheology = CompositeRheology((visc, el_bg, pl)), + Elasticity = el_bg, + + ), + # High density phase + SetMaterialParams(; + Density = ConstantDensity(; ρ = 0.0), + Gravity = ConstantGravity(; g = 0.0), + CompositeRheology = CompositeRheology((visc, el_inc, pl)), + Elasticity = el_inc, + ), + ) +``` + +# Phase anomaly + +Helper function to initialize material phases with `ParallelStencil.jl` +```julia +function init_phases!(phase_ratios, xci, radius) + ni = size(phase_ratios.center) + origin = 0.5, 0.5 + + @parallel_indices (i, j) function init_phases!(phases, xc, yc, o_x, o_y, radius) + x, y = xc[i], yc[j] + if ((x-o_x)^2 + (y-o_y)^2) > radius^2 + JustRelax.@cell phases[1, i, j] = 1.0 + JustRelax.@cell phases[2, i, j] = 0.0 + + else + JustRelax.@cell phases[1, i, j] = 0.0 + JustRelax.@cell phases[2, i, j] = 1.0 + end + return nothing + end + + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) +end + +``` + +and finally we need the phase ratios at the cell centers: +```julia +phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) +init_phases!(phase_ratios, xci, radius) +``` + +## Stokes arrays + +Stokes arrays object +```julia +stokes = StokesArrays(backend_JR, ni) +``` + +## Initialize viscosity fields + +We initialize the buoyancy forces and viscosity +```julia +ρg = @zeros(ni...), @zeros(ni...) +η = @ones(ni...) +args = (; T = thermal.Tc, P = stokes.P, dt = Inf) +compute_ρg!(ρg[2], phase_ratios, rheology, args) +compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) +``` +where `(-Inf, Inf)` is the viscosity cutoff. + +## Boundary conditions +```julia + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right = true, top = true, bot = true), + no_slip = (left = false, right = false, top = false, bot=false), + ) + 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]]) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(stokes.V.Vx, stokes.V.Vy) + +``` + +## Pseuo-transient coefficients +```julia +pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) +``` + +## Just before solving the problem... +In this benchmark we want to keep track of τII, the total time `ttot`, and the analytical elastic solution `sol` +```julia + solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) +``` +and store their time history in the vectors: +```julia + τII = Float64[] + sol = Float64[] + ttot = Float64[] +``` + +## Advancing one time step + +1. Solve stokes +```julia +solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 150e3, + nout = 200, + viscosity_cutoff = (-Inf, Inf), + verbose = true + ) +) +``` +2. calculate the second invariant and push to history vectors +```julia +tensor_invariant!(stokes.ε) +push!(τII, maximum(stokes.τ.xx)) + +@parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) +@parallel (@idx ni) multi_copy!( + @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) +) + +it += 1 +t += dt + +push!(sol, solution(εbg, t, G0, η0)) +push!(ttot, t) +``` +# Visualization +We will use `Makie.jl` to visualize the results +```julia +using GLMakie +``` + +## Fields +```julia + # visualisation of high density inclusion +th = 0:pi/50:3*pi; +xunit = @. radius * cos(th) + 0.5; +yunit = @. radius * sin(th) + 0.5; + +fig = Figure(size = (1600, 1600), title = "t = $t") +ax1 = Axis(fig[1,1], aspect = 1, title = L"\tau_{II}", titlesize=35) +ax2 = Axis(fig[2,1], aspect = 1, title = L"E_{II}", titlesize=35) +ax3 = Axis(fig[1,2], aspect = 1, title = L"\log_{10}(\varepsilon_{II})", titlesize=35) +ax4 = Axis(fig[2,2], aspect = 1) +heatmap!(ax1, xci..., Array(stokes.τ.II) , colormap=:batlow) +heatmap!(ax2, xci..., Array(log10.(stokes.EII_pl)) , colormap=:batlow) +heatmap!(ax3, xci..., Array(log10.(stokes.ε.II)) , colormap=:batlow) +lines!(ax2, xunit, yunit, color = :black, linewidth = 5) +lines!(ax4, ttot, τII, color = :black) +lines!(ax4, ttot, sol, color = :red) +hidexdecorations!(ax1) +hidexdecorations!(ax3) +save(joinpath(figdir, "$(it).png"), fig) +fig +``` + +### Final model +![](990.png) diff --git a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl index 4228ed60..01a6bd5d 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl @@ -198,7 +198,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv) - scatter!(ax2, Array(log10.(A[:])), Y) + scatter!(ax2, Array(log10.(η[:])), Y) ylims!(ax1, minimum(xvi[2]), 0) ylims!(ax2, minimum(xvi[2]), 0) hideydecorations!(ax2) @@ -390,4 +390,4 @@ else end # run main script -main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, do_vtk = do_vtk); +# main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/src/IO/VTK.jl b/src/IO/VTK.jl index 10dae088..33d61017 100644 --- a/src/IO/VTK.jl +++ b/src/IO/VTK.jl @@ -41,7 +41,9 @@ function append!(data_series, data::NamedTuple, time_step, seconds) return nothing end -function save_vtk(fname::String, xvi, xci, data_v::NamedTuple, data_c::NamedTuple) +function save_vtk( + fname::String, xvi, xci, data_v::NamedTuple, data_c::NamedTuple, velocity::NTuple{N,T} +) where {N,T} # unpack data names and arrays data_names_v = string.(keys(data_v)) @@ -49,6 +51,11 @@ function save_vtk(fname::String, xvi, xci, data_v::NamedTuple, data_c::NamedTupl data_names_c = string.(keys(data_c)) data_arrays_c = values(data_c) + velocity_field = rand(N, size(first(velocity))...) + for (i, v) in enumerate(velocity) + velocity_field[i, :, :, :] = v + end + vtk_multiblock(fname) do vtm # First block. # Variables stores in cell centers @@ -63,6 +70,7 @@ function save_vtk(fname::String, xvi, xci, data_v::NamedTuple, data_c::NamedTupl for (name_i, array_i) in zip(data_names_v, data_arrays_v) vtk[name_i] = Array(array_i) end + vtk["Velocity"] = velocity_field end end diff --git a/src/Interpolations.jl b/src/Interpolations.jl index f6b5a102..9eeb80f2 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -257,8 +257,10 @@ onto the pre-allocated `Vx_d`, `Vy_d`, `Vz_d` 3D arrays located at the grid vert """ function velocity2vertex!(Vx_v, Vy_v, Vz_v, Vx, Vy, Vz) # infer size of grid - nx, ny, nz = size(Vx) - nv_x, nv_y, nv_z = nx - 1, ny - 2, nz - 2 + nx1, ny1, nz1 = size(Vx) + nx2, ny2, nz2 = size(Vy) + nx3, ny3, nz3 = size(Vz) + nv_x, nv_y, nv_z = max(nx1, nx2, nx3), max(ny1, ny2, ny3), max(nz1, nz2, nz3) # interpolate to cell vertices @parallel (@idx nv_x nv_y nv_z) _velocity2vertex!(Vx_v, Vy_v, Vz_v, Vx, Vy, Vz) @@ -266,11 +268,28 @@ function velocity2vertex!(Vx_v, Vy_v, Vz_v, Vx, Vy, Vz) end @parallel_indices (i, j, k) function _velocity2vertex!(Vx_v, Vy_v, Vz_v, Vx, Vy, Vz) - @inbounds begin + i1, j1, k1 = (i, j, k) .+ 1 + # if all((i, j1, k1) .≤ size(Vx_v)) + # Vx_v[i, j1, k1] = + # 0.25 * (Vx[i, j, k] + Vx[i, j + 1, k] + Vx[i, j, k + 1] + Vx[i, j + 1, k + 1]) + # end + # if all((i1, j, k1) .≤ size(Vy_v)) + # Vy_v[i1, j, k1] = + # 0.25 * (Vy[i, j, k] + Vy[i + 1, j, k] + Vy[i, j, k + 1] + Vy[i + 1, j, k + 1]) + # end + # if all((i1, j1, k) .≤ size(Vz_v)) + # Vz_v[i1, j1, k] = + # 0.25 * (Vz[i, j, k] + Vz[i, j + 1, k] + Vz[i + 1, j, k] + Vz[i + 1, j + 1, k]) + # end + if all((i, j, k) .≤ size(Vx_v)) Vx_v[i, j, k] = 0.25 * (Vx[i, j, k] + Vx[i, j + 1, k] + Vx[i, j, k + 1] + Vx[i, j + 1, k + 1]) + end + if all((i, j, k) .≤ size(Vy_v)) Vy_v[i, j, k] = 0.25 * (Vy[i, j, k] + Vy[i + 1, j, k] + Vy[i, j, k + 1] + Vy[i + 1, j, k + 1]) + end + if all((i, j, k) .≤ size(Vz_v)) Vz_v[i, j, k] = 0.25 * (Vz[i, j, k] + Vz[i, j + 1, k] + Vz[i + 1, j, k] + Vz[i + 1, j + 1, k]) end diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl new file mode 100644 index 00000000..7edf21e8 --- /dev/null +++ b/src/MetaJustRelax.jl @@ -0,0 +1,175 @@ +struct PS_Setup{B,C} + device::Symbol + + function PS_Setup(device::Symbol, precision::DataType, nDim::Integer) + return new{precision,nDim}(device) + end +end + +function environment!(model::PS_Setup{T,N}) where {T,N} + + # call appropriate FD module + Base.eval(@__MODULE__, Meta.parse("using ParallelStencil.FiniteDifferences$(N)D")) + Base.eval(Main, Meta.parse("using ParallelStencil.FiniteDifferences$(N)D")) + + @eval const model_dims = $N + + # start ParallelStencil + if model.device == :CUDA + eval(:(using CUDA)) + eval(:(@init_parallel_stencil(CUDA, $T, $N))) + # Base.eval(Main, Meta.parse("using CUDA")) + if !isconst(Main, :PTArray) + eval(:(const PTArray = CUDA.CuArray{$T,$N,CUDA.Mem.DeviceBuffer})) + end + + # this is patchy, but it works for ParallelStencil 1.11 + @eval const backend = :CUDA + + elseif model.device == :AMDGPU + eval(:(using AMDGPU)) + eval(:(@init_parallel_stencil(AMDGPU, $T, $N))) + # Base.eval(Main, Meta.parse("using AMDGPU")) + if !isconst(Main, :PTArray) + eval(:(const PTArray = AMDGPU.ROCArray{$T,$N,AMDGPU.Runtime.Mem.HIPBuffer})) + end + + # this is patchy, but it works for ParallelStencil 1.11 + @eval const backend = :AMDGPU + + else + @eval begin + @init_parallel_stencil(Threads, $T, $N) + if !isconst(Main, :PTArray) + const PTArray = Array{$T,$N} + end + const backend = :Threads + end + end + + # CREATE ARRAY STRUCTS + make_velocity_struct!(N) # velocity + make_symmetrictensor_struct!(N) # (symmetric) tensors + make_vorticity_struct!(N) # (symmetric) tensors + ## Stokes + make_residual_struct!(N) # residuals + make_stokes_struct!() # Arrays for Stokes solver + make_PTstokes_struct!() + ## thermal diffusion + make_thermal_arrays!(N) # Arrays for Thermal Diffusion solver + make_PTthermal_struct!() # PT Thermal Diffusion coefficients + + # includes and exports + @eval begin + export USE_GPU, + PTArray, + Velocity, + SymmetricTensor, + Residual, + StokesArrays, + PTStokesCoeffs, + ThermalArrays, + PTThermalCoeffs, + compute_pt_thermal_arrays!, + AbstractStokesModel, + AbstractElasticModel, + Viscous, + ViscoElastic, + ViscoElastoPlastic, + solve! + + include(joinpath(@__DIR__, "Utils.jl")) + export @allocate, @add, @idx, @copy + export @velocity, + @strain, + @stress, + @tensor, + @shear, + @normal, + @stress_center, + @strain_center, + @tensor_center, + @qT, + @qT2, + @residuals, + compute_dt, + assign!, + tupleize, + compute_maxloc!, + continuation_log, + continuation_linear, + mean_mpi, + norm_mpi, + minimum_mpi, + maximum_mpi, + multi_copy!, + take + + include(joinpath(@__DIR__, "boundaryconditions/BoundaryConditions.jl")) + export pureshear_bc!, + FlowBoundaryConditions, + flow_bcs!, + TemperatureBoundaryConditions, + thermal_boundary_conditions!, + thermal_bcs!, + free_slip_x!, + free_slip_y!, + free_slip_z!, + apply_free_slip!, + free_surface_bcs! + + include(joinpath(@__DIR__, "phases/phases.jl")) + export PhaseRatio, fn_ratio, phase_ratios_center, phase_ratios_center! + + include(joinpath(@__DIR__, "rheology/BuoyancyForces.jl")) + export compute_ρg! + + include(joinpath(@__DIR__, "rheology/Viscosity.jl")) + export compute_viscosity! + + include(joinpath(@__DIR__, "stokes/Stokes2D.jl")) + export solve! + + include(joinpath(@__DIR__, "stokes/Stokes3D.jl")) + export solve! + + include(joinpath(@__DIR__, "thermal_diffusion/DiffusionExplicit.jl")) + export ThermalParameters + + include(joinpath(@__DIR__, "thermal_diffusion/DiffusionPT.jl")) + export heatdiffusion_PT! + + include(joinpath(@__DIR__, "particles/subgrid_diffusion.jl")) + export subgrid_characteristic_time! + + include(joinpath(@__DIR__, "thermal_diffusion/Shearheating.jl")) + export compute_shear_heating! + + include(joinpath(@__DIR__, "Interpolations.jl")) + export vertex2center!, center2vertex!, temperature2center!, interp_Vx_on_Vy! + + include(joinpath(@__DIR__, "advection/weno5.jl")) + export WENO5, WENO_advection! + end + + # conditional submodule load + module_names = if N === 1 + (Symbol("ThermalDiffusion$(N)D"),) + elseif N === 2 + (Symbol("Stokes$(N)D"), Symbol("ThermalDiffusion$(N)D")) + else + (Symbol("Stokes$(N)D"), Symbol("ThermalDiffusion$(N)D")) + end + + for m in module_names + Base.@eval begin + @reexport using .$m + end + end +end + +function ps_reset!() + Base.eval(Main, ParallelStencil.@reset_parallel_stencil) + Base.eval(@__MODULE__, ParallelStencil.@reset_parallel_stencil) + return nothing +end diff --git a/src/Utils.jl b/src/Utils.jl index 2be50507..b6785d8b 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -312,19 +312,6 @@ macro allocate(ni...) return esc(:(PTArray(undef, $(ni...)))) end -function indices(::NTuple{3,T}) where {T} - i = (blockIdx().x - 1) * blockDim().x + threadIdx().x - j = (blockIdx().y - 1) * blockDim().y + threadIdx().y - k = (blockIdx().z - 1) * blockDim().z + threadIdx().z - return i, j, k -end - -function indices(::NTuple{2,T}) where {T} - i = (blockIdx().x - 1) * blockDim().x + threadIdx().x - j = (blockIdx().y - 1) * blockDim().y + threadIdx().y - return i, j -end - """ maxloc!(B, A; window) @@ -428,6 +415,7 @@ Compute the time step `dt` for the velocity field `S.V` and the diffusive maximu dt_adv = mapreduce(x -> x[1] * inv(maximum(abs.(x[2]))), min, zip(di, V)) * n return min(dt_diff, dt_adv) end + """ compute_dt(S::JustRelax.StokesArrays, di, igg) @@ -532,6 +520,6 @@ for (f1, f2) in zip( ) @eval begin $f1(A::AbstractArray) = $f2(Array(A)) - $f1(A) = $f2(A) + $f1(A::Array) = $f2(A) end end diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index bdc9efeb..81968a6a 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -1,14 +1,15 @@ +# BOUNDARY CONDITIONS KERNELS +include("free_slip.jl") +include("free_surface.jl") +include("no_slip.jl") +include("pure_shear.jl") + @inline bc_index(x::NTuple{2,T}) where {T} = mapreduce(xi -> max(size(xi)...), max, x) -@inline bc_index(x::T) where {T<:AbstractArray{<:Any,2}} = max(size(x)...) +@inline bc_index(x::T) where {T<:AbstractArray} = max(size(x)...) @inline function bc_index(x::NTuple{3,T}) where {T} - nx, ny, nz = size(x[1]) - return max((nx, ny), (ny, nz), (nx, nz)) -end - -@inline function bc_index(x::T) where {T<:AbstractArray{<:Any,3}} - nx, ny, nz = size(x) - return max((nx, ny), (ny, nz), (nx, nz)) + n = mapreduce(xi -> max(size(xi)...), max, x) + return n, n end @inline do_bc(bc) = reduce(|, values(bc)) @@ -20,7 +21,7 @@ Apply the prescribed heat boundary conditions `bc` on the `T` """ thermal_bcs!(thermal, bcs) = thermal_bcs!(backend(thermal), thermal, bcs) function thermal_bcs!( - ::CPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs::FlowBoundaryConditions + ::CPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs::TemperatureBoundaryConditions ) return thermal_bcs!(thermal.T, bcs) end @@ -34,6 +35,8 @@ function thermal_bcs!(T::AbstractArray, bcs::TemperatureBoundaryConditions) return nothing end +# @inline thermal_bcs!(thermal::ThermalArrays, bcs::TemperatureBoundaryConditions) = thermal_bcs!(thermal.T, bcs) + """ flow_bcs!(stokes, bcs::FlowBoundaryConditions, di) @@ -56,9 +59,3 @@ function _flow_bcs!(bcs, V) return nothing end - -# BOUNDARY CONDITIONS KERNELS -include("free_slip.jl") -include("free_surface.jl") -include("no_slip.jl") -include("pure_shear.jl") diff --git a/src/common.jl b/src/common.jl index daf3244d..32216d31 100644 --- a/src/common.jl +++ b/src/common.jl @@ -35,7 +35,7 @@ export FlowBoundaryConditions, include("MiniKernels.jl") include("phases/phases.jl") -export fn_ratio, phase_ratios_center! +export phase_ratios_center!, fn_ratio include("rheology/BuoyancyForces.jl") export compute_ρg! @@ -59,12 +59,16 @@ export WENO5, WENO_advection! include("rheology/GeoParams.jl") include("rheology/StressUpdate.jl") -include("stokes/StressRotation.jl") +include("stokes/StressRotationParticles.jl") +export rotate_stress_particles! + include("stokes/StressKernels.jl") export tensor_invariant! include("stokes/PressureKernels.jl") include("stokes/VelocityKernels.jl") +include("stokes/VorticityKernels.jl") +export compute_vorticity! # thermal diffusion diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 74deef2d..29454e4d 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -168,6 +168,13 @@ function JR2D.center2vertex!( return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) end +function JR2D.velocity2vertex!( + Vx_v::CuArray, Vy_v::CuArray, Vx::CuArray, Vy::CuArray; ghost_nodes=true +) + velocity2vertex!(Vx_v, Vy_v, Vx, Vy; ghost_nodes=ghost_nodes) + return nothing +end + # Solvers function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index acf39f5a..fb0308ca 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -100,36 +100,37 @@ function JR3D.phase_ratios_center!( phases, ) return _phase_ratios_center!(phase_ratios, particles, grid, phases) + return _phase_ratios_center!(phase_ratios, particles, grid, phases) end # Rheology ## viscosity -function JR3D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) +function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end function JR3D.compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end -function JR2D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) +function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end -function compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) +function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end @@ -139,10 +140,10 @@ function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor end ## Buoyancy forces -function JR3D.compute_ρg!(ρg::RocArray, rheology, args) +function JR3D.compute_ρg!(ρg::CuArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::RocArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -155,20 +156,27 @@ function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArray return _temperature2center!(thermal) end -function JR3D.vertex2center!(center::T, vertex::T) where {T<:RocArray} +function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} return vertex2center!(center, vertex) end -function JR3D.center2vertex!(vertex::T, center::T) where {T<:RocArray} +function JR3D.center2vertex!(vertex::T, center::T) where {T<:CuArray} return center2vertex!(vertex, center) end function JR3D.center2vertex!( vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T -) where {T<:RocArray} +) where {T<:CuArray} return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) end +function JR3D.velocity2vertex!( + Vx_v::CuArray, Vy_v::CuArray, Vz_v::CuArray, Vx::CuArray, Vy::CuArray, Vz::CuArray +) + velocity2vertex!(Vx_v, Vy_v, Vz_v, Vx, Vy, Vz) + return nothing +end + # Solvers function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) @@ -183,10 +191,12 @@ function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) return _compute_dt(S, args...) end +# Subgrid diffusion + function JR3D.subgrid_characteristic_time!( subgrid_arrays, particles, - dt₀::RocArray, + dt₀::CuArray, phases::JustRelax.PhaseRatio, rheology, thermal::JustRelax.ThermalArrays, @@ -204,7 +214,7 @@ end function JR3D.subgrid_characteristic_time!( subgrid_arrays, particles, - dt₀::RocArray, + dt₀::CuArray, phases::AbstractArray{Int,N}, rheology, thermal::JustRelax.ThermalArrays, diff --git a/src/particles/subgrid_diffusion.jl b/src/particles/subgrid_diffusion.jl index bc457642..c4dbcdd2 100644 --- a/src/particles/subgrid_diffusion.jl +++ b/src/particles/subgrid_diffusion.jl @@ -1,9 +1,7 @@ -# import JustRelax.compute_ρCp - function subgrid_characteristic_time!( subgrid_arrays, particles, - dt₀, + dt₀::Array, phases::JustRelax.PhaseRatio, rheology, thermal::JustRelax.ThermalArrays, @@ -21,8 +19,8 @@ end function subgrid_characteristic_time!( subgrid_arrays, particles, - dt₀, - phases::AbstractArray{Int,N}, + dt₀::Array, + phases::Array{Int,N}, rheology, thermal::JustRelax.ThermalArrays, stokes::JustRelax.StokesArrays, @@ -51,81 +49,3 @@ end return nothing end - -# struct SubgridDiffusionArrays{T, CA} -# pT0::CA # CellArray -# pΔT::CA # CellArray -# ΔT::T # Array - -# function SubgridDiffusionArrays(particles, ni) -# pΔT, pT0 = init_cell_arrays(particles, Val(2)) -# ΔT = @zeros(ni.+1) -# CA, T = typeof(pΔT), typeof(ΔT) -# new{T, CA}(pT0, pΔT, ΔT) -# end -# end - -# @inline function init_cell_arrays(particles, ::Val{N}) where {N} -# return ntuple( -# _ -> @fill( -# 0.0, size(particles.coords[1])..., celldims = (cellsize(particles.index)) -# ), -# Val(N), -# ) -# end - -# function subgrid_diffusion!( -# pT, thermal, subgrid_arrays, pPhases, rheology, stokes, particles, T_buffer, xvi, di, dt -# ) -# (; pT0, pΔT) = subgrid_arrays -# # ni = size(pT) - -# @copy pT0.data pT.data -# grid2particle!(pT, xvi, T_buffer, particles) - -# @parallel (@idx ni) subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes.P, particles.index, di, dt) -# particle2grid!(subgrid_arrays.ΔT, pΔT, xvi, particles) - -# @parallel (@idx size(subgrid_arrays.ΔT)) update_ΔT_subgrid!(subgrid_arrays.ΔT, thermal.ΔT) -# grid2particle!(pΔT, xvi, subgrid_arrays.ΔT, particles) -# @. pT.data = pT0.data + pΔT.data -# return nothing -# end - -# @parallel_indices (i, j) function update_ΔT_subgrid!(ΔTsubgrid::_T, ΔT::_T) where _T<:AbstractMatrix -# ΔTsubgrid[i, j] = ΔT[i+1, j] - ΔTsubgrid[i, j] -# return nothing -# end - -# function subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes, particles, di, dt) -# ni = size(pT) -# @parallel (@idx ni) subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, stokes.P, particles.index, di, dt) -# end - -# @parallel_indices (I...) function subgrid_diffusion!(pT, pT0, pΔT, pPhases, rheology, P, index, di, dt) - -# P_cell = P[I...] - -# for ip in JustRelax.cellaxes(pT) -# # early escape if there is no particle in this memory locations -# doskip(index, ip, I...) && continue - -# pT0ᵢ = @cell pT0[ip, I...] -# pTᵢ = @cell pT[ip, I...] -# phase = Int(@cell(pPhases[ip, I...])) -# argsᵢ = (; T = pTᵢ, P = P_cell) -# # dimensionless numerical diffusion coefficient (0 ≤ d ≤ 1) -# d = 1 -# # Compute the characteristic timescale `dt₀` of the local cell -# ρCp = compute_ρCp(rheology, phase, argsᵢ) -# K = compute_conductivity(rheology, phase, argsᵢ) -# sum_dxi = mapreduce(x-> inv(x)^2, +, di) -# dt₀ = ρCp / (2 * K * sum_dxi) -# # subgrid diffusion of the i-th particle -# pΔTᵢ = (pTᵢ - pT0ᵢ) * (1-exp(-d * dt / dt₀)) -# @cell pT0[ip, I...] = pT0ᵢ + pΔTᵢ -# @cell pΔT[ip, I...] = pΔTᵢ -# end - -# return nothing -# end diff --git a/src/phases/phases.jl b/src/phases/phases.jl index a0d96ac6..81dea637 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -57,20 +57,20 @@ end function phase_ratios_center!( ::CPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) + return _phase_ratios_center!(phase_ratios, particles, grid, phases) end -function _phase_ratios_center( +function _phase_ratios_center!( phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases ) ni = size(phases) - @parallel (@idx ni) phase_ratios_center_kernel( + @parallel (@idx ni) phase_ratios_center_kernel!( phase_ratios.center, particles.coords, grid.xci, grid.di, phases ) return nothing end -@parallel_indices (I...) function phase_ratios_center_kernel( +@parallel_indices (I...) function phase_ratios_center_kernel!( ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases ) where {N,T1,T2,T3} diff --git a/src/rheology/GeoParams.jl b/src/rheology/GeoParams.jl index 55d54d47..eaf80996 100644 --- a/src/rheology/GeoParams.jl +++ b/src/rheology/GeoParams.jl @@ -25,3 +25,42 @@ end @inline get_α(p::MaterialParams) = get_α(p.Density[1]) @inline get_α(p::ConstantDensity) = 0.0 @inline get_α(p::Union{T_Density,PT_Density}) = GeoParams.get_α(p) + +# Check whether the material has constant density. If so, no need to calculate the density +# during PT iterations +@generated function is_constant_density( + rheology::NTuple{N,AbstractMaterialParamsStruct} +) where {N} + quote + Base.@_inline_meta + Base.@nexprs $N i -> !_is_constant_density(rheology[i].Density[1]) && return false + return true + end +end + +@inline _is_constant_density(::ConstantDensity) = true +@inline _is_constant_density(::AbstractDensity) = false + +# Check whether the material has a linear viscosity. If so, no need to calculate the viscosity +# during PT iterations +@generated function is_constant_viscosity( + rheology::NTuple{N,AbstractMaterialParamsStruct} +) where {N} + quote + Base.@_inline_meta + Base.@nexprs $N i -> + is_constant_viscosity(rheology[i].CompositeRheology[1].elements) && return false + return true + end +end + +@generated function is_constant_viscosity(creep_law::NTuple{N,AbstractCreepLaw}) where {N} + quote + Base.@_inline_meta + Base.@nexprs $N i -> _is_constant_viscosity(creep_law[i]) && return false + return true + end +end + +@inline _is_constant_viscosity(::Union{LinearViscous,ConstantElasticity}) = true +@inline _is_constant_viscosity(::AbstractCreepLaw) = false diff --git a/src/rheology/StressUpdate.jl b/src/rheology/StressUpdate.jl index 8de1aac6..c4fffd4b 100644 --- a/src/rheology/StressUpdate.jl +++ b/src/rheology/StressUpdate.jl @@ -1,7 +1,6 @@ # inner kernel to compute the plastic stress update within Pseudo-Transient stress continuation function _compute_τ_nonlinear!( τ::NTuple{N1,T}, - τII, τ_old::NTuple{N1,T}, ε::NTuple{N1,T}, ε_pl::NTuple{N1,T}, @@ -35,7 +34,8 @@ function _compute_τ_nonlinear!( # check if yielding; if so, compute plastic strain rate (λdQdτ), # plastic stress increment (dτ_pl), and update the plastic # multiplier (λ) - dτij, λdQdτ = if isyielding(is_pl, τII_trial, τy) + failure = isyielding(is_pl, τII_trial, τy) + dτij, λdQdτ = if failure # derivatives plastic stress correction dτ_pl, λ[idx...], λdQdτ = compute_dτ_pl( τij, dτij, τy, τII_trial, ηij, λ[idx...], η_reg, dτ_r, volume @@ -48,11 +48,12 @@ function _compute_τ_nonlinear!( end # fill plastic strain rate tensor - update_plastic_strain_rate!(ε_pl, λdQdτ, idx) + failure && update_plastic_strain_rate!(ε_pl, λdQdτ, idx) # update and correct stress - correct_stress!(τ, τij .+ dτij, idx...) + correct_stress!(τ, τij, dτij, idx...) - τII[idx...] = τII_ij = second_invariant(τij...) + # τII[idx...] = + τII_ij = second_invariant(τij...) η_vep[idx...] = τII_ij * 0.5 * inv(second_invariant(εij_ve...)) return nothing @@ -71,14 +72,33 @@ end @inline compute_dτ_r(θ_dτ, ηij, _Gdt) = inv(θ_dτ + fma(ηij, _Gdt, 1.0)) -function compute_stress_increment_and_trial( +# function compute_stress_increment_and_trial( +# τij::NTuple{N,T}, τij_o::NTuple{N,T}, ηij, εij::NTuple{N,T}, _Gdt, dτ_r +# ) where {N,T} +# dτij = ntuple(Val(N)) do i +# Base.@_inline_meta +# dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i])) +# end +# return dτij, second_invariant((τij .+ dτij)...) +# end + +# Fully unrolled version of the above function. Harder to read but faster +@generated function compute_stress_increment_and_trial( τij::NTuple{N,T}, τij_o::NTuple{N,T}, ηij, εij::NTuple{N,T}, _Gdt, dτ_r ) where {N,T} - dτij = ntuple(Val(N)) do i + quote Base.@_inline_meta - dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i])) + Base.@nexprs $N i -> begin + τij_n = τij[i] + dτ_i = + dτ_r * + fma(2.0 * ηij, εij[i], fma(-((τij_n - τij_o[i])) * ηij, _Gdt, -τij_n)) + pt_τ_i = τij_n + dτ_i + end + dτij = Base.@ncall $N tuple dτ + pt_τII = Base.@ncall $N second_invariant pt_τ + return dτij, pt_τII end - return dτij, second_invariant((τij .+ dτij)...) end function compute_dτ_pl( @@ -123,6 +143,25 @@ end return correct_stress!((τxx, τyy, τzz, τyz, τxz, τxy), τij, i, j, k) end +@generated function correct_stress!( + τ, τij::NTuple{N1,T}, dτij::NTuple{N1,T}, idx::Vararg{Integer,N2} +) where {N1,N2,T} + quote + Base.@_inline_meta + Base.@nexprs $N1 i -> τ[i][idx...] = τij[i] + dτij[i] + end +end + +@inline function correct_stress!(τxx, τyy, τxy, τij::NTuple, dτij::NTuple, i, j) + return correct_stress!((τxx, τyy, τxy), τij, dτij, i, j) +end + +@inline function correct_stress!( + τxx, τyy, τzz, τyz, τxz, τxy, τij::NTuple, dτij::NTuple, i, j, k +) + return correct_stress!((τxx, τyy, τzz, τyz, τxz, τxy), τij, dτij, i, j, k) +end + @inline isplastic(x::AbstractPlasticity) = true @inline isplastic(x) = false diff --git a/src/stokes/MetaStokes.jl b/src/stokes/MetaStokes.jl new file mode 100644 index 00000000..ce644e4d --- /dev/null +++ b/src/stokes/MetaStokes.jl @@ -0,0 +1,328 @@ +abstract type AbstractStokesModel end +abstract type AbstractViscosity end +abstract type Viscous <: AbstractStokesModel end +abstract type AbstractElasticModel <: AbstractStokesModel end +abstract type ViscoElastic <: AbstractElasticModel end +abstract type ViscoElastoPlastic <: AbstractElasticModel end + +function make_velocity_struct!(ndim::Integer; name::Symbol=:Velocity) + dims = (:Vx, :Vy, :Vz) + fields = [:($(dims[i])::T) for i in 1:ndim] + @eval begin + struct $(name){T} + $(fields...) + + function $(name)(ni::NTuple{2,T}) where {T} + return new{$PTArray}(@zeros(ni[1]...), @zeros(ni[2]...)) + end + + function $(name)(ni::NTuple{3,T}) where {T} + return new{$PTArray}(@zeros(ni[1]...), @zeros(ni[2]...), @zeros(ni[3]...)) + end + end + $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = $(name)(args...) + end +end + +abstract type AbstractVorticityTensor end + +function make_vorticity_struct!(ndim::Integer) + dims = (:xy, :xz, :yz) + fields_c = if ndim == 2 + [:($(Symbol(dims[1], :_c))::T)] + + elseif ndim == 3 + [:($(Symbol(dims[i], :_c))::T) for i in 1:ndim] + else + throw(ArgumentError("ndim must be 2 or 3")) + end + + fields_v = if ndim == 2 + [:($(Symbol(dims[1], :_v))::T)] + + elseif ndim == 3 + [:($(Symbol(dims[i], :_v))::T) for i in 1:ndim] + end + + @eval begin + struct VorticityTensor{T} <: AbstractVorticityTensor + $(fields_c...) # centers + $(fields_v...) # vertices + + function VorticityTensor(ni::NTuple{2,T}) where {T} + # centers + xy_c = @zeros(ni...) + # vertices + xy_v = @zeros(ni .+ 1...) + + return new{$PTArray}(xy_c, xy_v) + end + + function VorticityTensor(ni::NTuple{3,T}) where {T} + # centers + yz_c = @zeros(ni...) + xz_c = @zeros(ni...) + xy_c = @zeros(ni...) + # vertices + yz_v = @zeros(ni .+ 1...) + xz_v = @zeros(ni .+ 1...) + xy_v = @zeros(ni .+ 1...) + + return new{$PTArray}(yz_c, xz_c, xy_c, yz_v, xz_v, xy_v) + end + end + VorticityTensor(args::Vararg{T,N}) where {T<:AbstractArray,N} = + VorticityTensor(args...) + end +end + +function make_viscosity_struct!() + @eval begin + struct Viscosity{T} + η::T # with no plasticity + η_vep::T # with plasticity + ητ::T # PT viscosity + + function Viscosity(ni) + η = @zeros(ni...) + η_vep = @zeros(ni...) + ητ = @zeros(ni...) + return new{typeof(η)}(η, η_vep, ητ) + end + end + end +end + +function make_symmetrictensor_struct!(nDim::Integer; name::Symbol=:SymmetricTensor) + dims = (:x, :y, :z) + fields = [:($(Symbol((dims[i]), (dims[j])))::T) for i in 1:nDim, j in 1:nDim if j ≥ i] + + fields_c = if nDim == 2 + [:($(:xy_c)::T)] + elseif nDim == 3 + [:($(:yz_c)::T), :($(:xz_c)::T), :($(:xy_c)::T)] + end + + @eval begin + struct $(name){T} + $(fields...) + $(fields_c...) + II::T + + function $(name)(ni::NTuple{2,T}) where {T} + return new{$PTArray}( + @zeros(ni...), # xx + @zeros(ni[1] + 1, ni[2] + 1), # xy + @zeros(ni...), # yy + @zeros(ni...), # xy @ cell center + @zeros(ni...) # II (second invariant) + ) + end + + function $(name)(ni::NTuple{3,T}) where {T} + return new{$PTArray}( + @zeros(ni...), # xx + @zeros(ni[1] + 1, ni[2] + 1, ni[3]), # xy + @zeros(ni...), # yy + @zeros(ni[1] + 1, ni[2], ni[3] + 1), # xz + @zeros(ni[1], ni[2] + 1, ni[3] + 1), # yz + @zeros(ni...), # zz + @zeros(ni...), # yz @ cell center + @zeros(ni...), # xz @ cell center + @zeros(ni...), # xy @ cell center + @zeros(ni...), # II (second invariant) + ) + end + end + + $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = $(name)(args...) + end +end + +function make_residual_struct!(ndim; name::Symbol=:Residual) + dims = (:Rx, :Ry, :Rz) + fields = [:($(dims[i])::T) for i in 1:ndim] + @eval begin + struct $(name){T} + $(fields...) + RP::T + + function $(name)(ni::NTuple{3,T}) where {T} + Rx = @zeros(ni[1]...) + Ry = @zeros(ni[2]...) + RP = @zeros(ni[3]...) + return new{typeof(Rx)}(Rx, Ry, RP) + end + + function $(name)(ni::NTuple{4,T}) where {T} + Rx = @zeros(ni[1]...) + Ry = @zeros(ni[2]...) + Rz = @zeros(ni[3]...) + RP = @zeros(ni[4]...) + return new{typeof(Rx)}(Rx, Ry, Rz, RP) + end + end + $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = $(name)(args...) + end +end + +function make_stokes_struct!(; name::Symbol=:StokesArrays) + @eval begin + struct StokesArrays{M<:AbstractStokesModel,A,B,C,D,T,nDim} + P::T + P0::T + V::A + ∇V::T + τ::B + ω::D + ω_o::D + ε::B + ε_pl::B + EII_pl::T + τ_o::Union{B,Nothing} + R::C + + # 2D CONSTRUCTORS + + function StokesArrays(ni::NTuple{2,T}, model::Type{Viscous}) where {T} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(((ni[1] + 1, ni[2] + 2), (ni[1], ni[2] + 2))) + τ = SymmetricTensor(ni) + ε = SymmetricTensor(ni) + ω = VorticityTensor(ni) # vorticity + ω_o = VorticityTensor(ni) # vorticity + ε_pl = SymmetricTensor(ni) + EII_pl = @zeros(ni...) + R = Residual(((ni[1] - 1, ni[2]), (ni[1], ni[2] - 1)), ni) + + return new{model,typeof(V),typeof(τ),typeof(R),typeof(ω),typeof(P),2}( + P, P0, V, ∇V, τ, ω, ω_o, ε, ε_pl, EII_pl, nothing, R + ) + end + + function StokesArrays( + ni::NTuple{2,T}, model::Type{<:AbstractElasticModel} + ) where {T} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(((ni[1] + 1, ni[2] + 2), (ni[1] + 2, ni[2] + 1))) + τ = SymmetricTensor(ni) + ε = SymmetricTensor(ni) + ω = VorticityTensor(ni) # vorticity + ω_o = VorticityTensor(ni) # vorticity + ε_pl = SymmetricTensor(ni) + EII_pl = @zeros(ni...) + R = Residual(((ni[1] - 1, ni[2]), (ni[1], ni[2] - 1), ni)) + + return new{model,typeof(V),typeof(τ),typeof(R),typeof(ω),typeof(P),2}( + P, P0, V, ∇V, τ, ω, ω_o, ε, ε_pl, EII_pl, deepcopy(τ), R + ) + end + + # 3D CONSTRUCTORS + + function StokesArrays(ni::NTuple{3,T}, model::Type{Viscous}) where {T} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(( + (ni[1] + 1, ni[2], ni[3]), + (ni[1], ni[2] + 1, ni[3]), + (ni[1], ni[2], ni[3] + 1), + )) + τ = SymmetricTensor(ni) + ε = SymmetricTensor(ni) + ω = VorticityTensor(ni) # vorticity + ω_o = VorticityTensor(ni) # vorticity + ε_pl = SymmetricTensor(ni) + EII_pl = @zeros(ni...) + R = Residual(( + (ni[1] - 1, ni[2] - 2, ni[3] - 2), + (ni[1] - 2, ni[2] - 1, ni[3] - 2), + (ni[1] - 2, ni[2] - 2, ni[3] - 1), + ni, + )) + + return new{model,typeof(V),typeof(τ),typeof(R),typeof(ω),typeof(P),3}( + P, P0, V, ∇V, τ, ω, ω_o, ε, ε_pl, EII_pl, nothing, R + ) + end + + function StokesArrays( + ni::NTuple{3,T}, model::Type{<:AbstractElasticModel} + ) where {T} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(( + (ni[1] + 1, ni[2] + 2, ni[3] + 2), + (ni[1] + 2, ni[2] + 1, ni[3] + 2), + (ni[1] + 2, ni[2] + 2, ni[3] + 1), + )) + τ = SymmetricTensor(ni) + ε = SymmetricTensor(ni) + ω = VorticityTensor(ni) # vorticity + ω_o = VorticityTensor(ni) # vorticity + ε_pl = SymmetricTensor(ni) + EII_pl = @zeros(ni...) + R = Residual(( + (ni[1] - 1, ni[2], ni[3]), + (ni[1], ni[2] - 1, ni[3]), + (ni[1], ni[2], ni[3] - 1), + ni, + )) + + return new{model,typeof(V),typeof(τ),typeof(R),typeof(ω),typeof(P),3}( + P, P0, V, ∇V, τ, ω, ω_o, ε, ε_pl, EII_pl, deepcopy(τ), R + ) + end + + function $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} + return new{ + ViscoElastic, + typeof(args[4]), + typeof(args[5]), + typeof(args[end]), + typeof(args[6]), + typeof(args[1]), + N, + }( + args... + ) + end + end + end +end + +function make_PTstokes_struct!() + @eval begin + struct PTStokesCoeffs{T} + CFL::T + ϵ::T # PT tolerance + Re::T # Reynolds Number + r::T # + Vpdτ::T + θ_dτ::T + ηdτ::T + + function PTStokesCoeffs( + li::NTuple{N,T}, + di; + ϵ=1e-8, + Re=3π, + CFL=(N == 2 ? 0.9 / √2.1 : 0.9 / √3.1), + r=0.7, + ) where {N,T} + lτ = min(li...) + Vpdτ = min(di...) * CFL + θ_dτ = lτ * (r + 4 / 3) / (Re * Vpdτ) + ηdτ = Vpdτ * lτ / Re + + return new{typeof(r)}(CFL, ϵ, Re, r, Vpdτ, θ_dτ, ηdτ) + end + end + end +end diff --git a/src/stokes/PressureKernels.jl b/src/stokes/PressureKernels.jl index f1bc4379..03e8f295 100644 --- a/src/stokes/PressureKernels.jl +++ b/src/stokes/PressureKernels.jl @@ -1,3 +1,25 @@ + +""" + init_P!(P, ρg, z) + +Initialize the pressure field `P` based on the buoyancy forces `ρg` and the vertical coordinate `z`. + +# Arguments +- `P::Array`: Pressure field to be initialized. +- `ρg::Float64`: Buoyancy forces. +- `z::Array`: Vertical coordinate. +""" +function init_P!(P, ρg, z) + ni = size(P) + @parallel (@idx ni) init_P_kernel!(P, ρg, z) + return nothing +end + +@parallel_indices (I...) function init_P_kernel!(P, ρg, z) + P[I...] = abs(ρg[I...] * z[I[end]]) * (z[I[end]] < 0.0) + return nothing +end + # Continuity equation ## Incompressible diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 22c1ea75..bd27216b 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -345,7 +345,6 @@ function _solve!( @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), - stokes.τ.II, @tensor(stokes.τ_o), @strain(stokes), @tensor_center(stokes.ε_pl), @@ -506,9 +505,12 @@ function _solve!( end Vx_on_Vy = @zeros(size(stokes.V.Vy)) - # compute buoyancy forces and viscosity - compute_ρg!(ρg[end], phase_ratios, rheology, args) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + + (; ρbg) = args + ρgz_diff = ρg[2] .- ρbg + Plitho = reverse(cumsum(reverse((ρg[2]) .* di[2]; dims=2); dims=2); dims=2) while iter ≤ iterMax iterMin < iter && err < ϵ && break @@ -532,6 +534,7 @@ function _solve!( θ_dτ, args, ) + args.P .= stokes.P .+ Plitho .- minimum(stokes.P) # stokes.P[1, 1] = stokes.P[2, 1] # stokes.P[end, 1] = stokes.P[end - 1, 1] @@ -539,6 +542,7 @@ function _solve!( # stokes.P[end, end] = stokes.P[end - 1, end] update_ρg!(ρg[2], phase_ratios, rheology, args) + @. ρgz_diff = ρg[2] - ρbg @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... @@ -560,12 +564,12 @@ function _solve!( @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), - stokes.τ.II, @tensor_center(stokes.τ_o), @strain(stokes), @tensor_center(stokes.ε_pl), stokes.EII_pl, - stokes.P, + args.P, + # stokes.P, θ, η, η_vep, @@ -578,6 +582,8 @@ function _solve!( center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) + # stokes.τ.yy[:, end] .= Plitho[:, end] + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( Vx_on_Vy, stokes.V.Vx ) @@ -586,16 +592,19 @@ function _solve!( @parallel compute_V!( @velocity(stokes)..., Vx_on_Vy, - θ, + stokes.P, @stress(stokes)..., pt_stokes.ηdτ, - ρg..., + ρg[1], + ρgz_diff, ητ, _di..., dt * free_surface, ) # apply boundary conditions + # free_surface_bcs!(stokes, flow_bcs, args, η, rheology, phase_ratios, dt, di) free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) + # free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) flow_bcs!(stokes, flow_bcs) update_halo!(@velocity(stokes)...) end @@ -613,7 +622,8 @@ function _solve!( Vx_on_Vy, stokes.P, @stress(stokes)..., - ρg..., + ρg[1], + ρgz_diff, _di..., dt * free_surface, ) @@ -626,7 +636,7 @@ function _solve!( push!(norm_Rx, errs[1]) push!(norm_Ry, errs[2]) push!(norm_∇V, errs[3]) - err = maximum_mpi(errs) + err = maximum(errs) push!(err_evo1, err) push!(err_evo2, iter) @@ -648,13 +658,19 @@ function _solve!( end end - stokes.P .= θ + # stokes.P .= θ @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + # compute_vorticity!(stokes, di) + + # @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + # @parallel (@idx ni) multi_copy!( + # @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) + # ) return ( iter=iter, diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 3118d0e9..507516da 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -207,6 +207,10 @@ function _solve!( compute_ρg!(ρg[end], phase_ratios, rheology, args) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + (; ρbg) = args + ρgz_diff = ρg[3] .- ρbg + Plitho = reverse(cumsum(reverse((ρg[2]) .* di[2]; dims=3); dims=3); dims=3) + # solver loop wtime0 = 0.0 while iter < 2 || (err > ϵ && iter ≤ iterMax) @@ -226,12 +230,15 @@ function _solve!( pt_stokes.r, pt_stokes.θ_dτ, ) + args.P .= stokes.P .+ Plitho .- minimum(stokes.P) + @parallel (@idx ni) compute_strain_rate!( stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... ) # Update buoyancy - update_ρg!(ρg[3], rheology, args) + update_ρg!(ρg[end], rheology, args) + @. ρgz_diff = ρg[end] - ρbg update_viscosity!( stokes, @@ -244,7 +251,6 @@ function _solve!( @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), - stokes.τ.II, @tensor(stokes.τ_o), @strain(stokes), @tensor_center(stokes.ε_pl), @@ -276,7 +282,9 @@ function _solve!( stokes.R.Ry, stokes.R.Rz, stokes.P, - ρg..., + ρg[1], + ρg[2], + ρgz_diff, @stress(stokes)..., ητ, pt_stokes.ηdτ, @@ -284,7 +292,7 @@ function _solve!( ) # apply boundary conditions flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + update_halo!(@velocity(stokes)...) end end @@ -382,7 +390,11 @@ function _solve!( @copy stokes.P0 stokes.P λ = @zeros(ni...) θ = @zeros(ni...) + ητ = deepcopy(η) + cte_density = is_constant_density(rheology) + cte_viscosity = is_constant_viscosity(rheology) + do_halo_update = !all(isone, igg.dims) # solver loop wtime0 = 0.0 ητ = deepcopy(η) @@ -395,7 +407,11 @@ function _solve!( wtime0 += @elapsed begin # ~preconditioner compute_maxloc!(ητ, η) - update_halo!(ητ) + do_halo_update && update_halo!(ητ) + # @hide_communication b_width begin # communication/computation overlap + # @parallel compute_maxloc!(ητ, η) + # update_halo!(ητ) + # end @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) compute_P!( @@ -430,9 +446,9 @@ function _solve!( ) # update_stress!(stokes, θ, λ, phase_ratios, rheology, dt, pt_stokes.θ_dτ) + # if !cte_viscosity @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), - stokes.τ.II, @tensor_center(stokes.τ_o), @strain(stokes), @tensor_center(stokes.ε_pl), @@ -456,7 +472,19 @@ function _solve!( stokes.τ.xz_c, stokes.τ.xy_c, ) - update_halo!(stokes.τ.yz, stokes.τ.xz, stokes.τ.xy) + do_halo_update && update_halo!(stokes.τ.yz, stokes.τ.xz, stokes.τ.xy) + # else + # @parallel (@idx ni .+ 1) compute_τ!( + # @tensor(stokes.τ)..., + # @tensor(stokes.τ_o)..., + # @tensor(stokes.ε)..., + # η, + # phase_ratios.center, + # tupleize(rheology), # needs to be a tuple + # dt, + # pt_stokes.θ_dτ, + # ) + # end # @parallel (@idx ni .+ 1) compute_τ_vertex!( # @shear(stokes.τ)..., @shear(stokes.ε)..., η_vep, pt_stokes.θ_dτ @@ -476,7 +504,7 @@ function _solve!( # apply boundary conditions free_surface_bcs!(stokes, flow_bc, η, rheology, phase_ratios, dt, di) flow_bcs!(stokes, flow_bc) - update_halo!(@velocity(stokes)...) + do_halo_update && update_halo!(@velocity(stokes)...) end end diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index d06fdcba..1f97d93c 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -201,6 +201,123 @@ end return nothing end +@parallel_indices (i, j, k) function compute_τ!( + τxx, + τyy, + τzz, + τyz, + τxz, + τxy, + τxx_o, + τyy_o, + τzz_o, + τyz_o, + τxz_o, + τxy_o, + εxx, + εyy, + εzz, + εyz, + εxz, + εxy, + η, + phase_center, + rheology, + dt, + θ_dτ, +) + harm_xy(A) = _harm_xyi(A, i, j, k) + harm_xz(A) = _harm_xzi(A, i, j, k) + harm_yz(A) = _harm_yzi(A, i, j, k) + av_xy(A) = _av_xyi(A, i, j, k) + av_xz(A) = _av_xzi(A, i, j, k) + av_yz(A) = _av_yzi(A, i, j, k) + get(x) = x[i, j, k] + + @inbounds begin + if all((i, j, k) .≤ size(τxx)) + phase = phase_center[i, j, k] + _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + η_ij = get(η) + denominator = inv(θ_dτ + η_ij * _Gdt + 1.0) + # Compute τ_xx + τxx[i, j, k] += + ( + -(get(τxx) - get(τxx_o)) * η_ij * _Gdt - get(τxx) + + 2.0 * η_ij * get(εxx) + ) * denominator + # Compute τ_yy + τyy[i, j, k] += + ( + -(get(τyy) - get(τyy_o)) * η_ij * _Gdt - get(τyy) + + 2.0 * η_ij * get(εyy) + ) * denominator + # Compute τ_zz + τzz[i, j, k] += + ( + -(get(τzz) - get(τzz_o)) * η_ij * _Gdt - get(τzz) + + 2.0 * η_ij * get(εzz) + ) * denominator + end + # Compute τ_xy + if (1 < i < size(τxy, 1)) && (1 < j < size(τxy, 2)) && k ≤ size(τxy, 3) + G = + ( + fn_ratio(get_shear_modulus, rheology, phase_center[i, j, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i - 1, j, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i, j - 1, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i - 1, j - 1, k]) + ) * 0.25 + _Gdt = inv(G * dt) + η_ij = harm_xy(η) + denominator = inv(θ_dτ + η_ij * _Gdt + 1.0) + τxy[i, j, k] += + ( + -(get(τxy) - get(τxy_o)) * η_ij * _Gdt - get(τxy) + + 2.0 * η_ij * get(εxy) + ) * denominator + end + # Compute τ_xz + if (1 < i < size(τxz, 1)) && j ≤ size(τxz, 2) && (1 < k < size(τxz, 3)) + G = + ( + fn_ratio(get_shear_modulus, rheology, phase_center[i, j, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i - 1, j, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i, j, k - 1]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i - 1, j, k - 1]) + ) * 0.25 + _Gdt = inv(G * dt) + η_ij = harm_xz(η) + denominator = inv(θ_dτ + η_ij * _Gdt + 1.0) + τxz[i, j, k] += + ( + -(get(τxz) - get(τxz_o)) * η_ij * _Gdt - get(τxz) + + 2.0 * η_ij * get(εxz) + ) * denominator + end + # Compute τ_yz + if i ≤ size(τyz, 1) && (1 < j < size(τyz, 2)) && (1 < k < size(τyz, 3)) + G = + ( + fn_ratio(get_shear_modulus, rheology, phase_center[i, j, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i, j - 1, k]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i, j, k - 1]) + + fn_ratio(get_shear_modulus, rheology, phase_center[i, j - 1, k - 1]) + ) * 0.25 + _Gdt = inv(G * dt) + + η_ij = harm_yz(η) + denominator = inv(θ_dτ + η_ij * _Gdt + 1.0) + τyz[i, j, k] += + ( + -(get(τyz) - get(τyz_o)) * η_ij * _Gdt - get(τyz) + + 2.0 * η_ij * get(εyz) + ) * denominator + end + end + return nothing +end + @parallel_indices (i, j, k) function compute_τ_vertex!( τyz, τxz, τxy, εyz, εxz, εxy, ηvep, θ_dτ ) @@ -240,7 +357,6 @@ end @parallel_indices (I...) function compute_τ_nonlinear!( τ, # @ centers - τII, # @ centers τ_old, # @ centers ε, # @ vertices ε_pl, # @ centers @@ -269,7 +385,7 @@ end plastic_parameters = (; is_pl, C, sinϕ, cosϕ, η_reg, volume) _compute_τ_nonlinear!( - τ, τII, τ_old, ε, ε_pl, P, ηij, η_vep, λ, dτ_r, _Gdt, plastic_parameters, I... + τ, τ_old, ε, ε_pl, P, ηij, η_vep, λ, dτ_r, _Gdt, plastic_parameters, I... ) # augmented pressure with plastic volumetric strain over pressure @@ -281,7 +397,6 @@ end # multi phase visco-elasto-plastic flow, where phases are defined in the cell center @parallel_indices (I...) function compute_τ_nonlinear!( τ, # @ centers - τII, # @ centers τ_old, # @ centers ε, # @ vertices ε_pl, # @ centers @@ -311,7 +426,7 @@ end plastic_parameters = (; is_pl, C, sinϕ, cosϕ, η_reg, volume) _compute_τ_nonlinear!( - τ, τII, τ_old, ε, ε_pl, P, ηij, η_vep, λ, dτ_r, _Gdt, plastic_parameters, I... + τ, τ_old, ε, ε_pl, P, ηij, η_vep, λ, dτ_r, _Gdt, plastic_parameters, I... ) # augmented pressure with plastic volumetric strain over pressure @inbounds θ[I...] = P[I...] + (isinf(K) ? 0.0 : K * dt * λ[I...] * sinψ) diff --git a/src/stokes/StressRotation.jl b/src/stokes/StressRotation.jl index ea9b45da..f64d9250 100644 --- a/src/stokes/StressRotation.jl +++ b/src/stokes/StressRotation.jl @@ -2,15 +2,6 @@ using StaticArrays ## Stress Rotation on the particles -@parallel_indices (i, j) function compute_vorticity!(vorticity, Vx, Vy, _dx, _dy) - dx(A) = _d_xa(A, i, j, _dx) - dy(A) = _d_ya(A, i, j, _dy) - - vorticity[i, j] = 0.5 * (dx(Vy) - dy(Vx)) - - return nothing -end - @parallel_indices (i, j) function rotate_stress_particles_jaumann!(xx, yy, xy, ω, index, dt) cell = i, j @@ -96,7 +87,6 @@ end Jaumann derivative τij_o += v_k * ∂τij_o/∂x_k - ω_ij * ∂τkj_o + ∂τkj_o * ω_ij - """ Base.@propagate_inbounds function rotate_stress!( V, τ::NTuple{N,T}, idx, _di, dt diff --git a/src/stokes/StressRotationParticles.jl b/src/stokes/StressRotationParticles.jl new file mode 100644 index 00000000..5f695de6 --- /dev/null +++ b/src/stokes/StressRotationParticles.jl @@ -0,0 +1,122 @@ +using StaticArrays + +# ROTATION KERNELS + +function jaumann!(xx, yy, xy, ω, index, dt) + ni = size(xx) + @parallel (@idx ni) _jaumann!(xx, yy, xy, ω, index, dt) + return nothing +end + +@parallel_indices (i, j) function _jaumann!(xx, yy, xy, ω, index, dt) + cell = i, j + + for ip in JustRelax.cellaxes(index) + !@cell(index[ip, cell...]) && continue # no particle in this location + + ω_xy = @cell ω[ip, cell...] + τ_xx = @cell xx[ip, cell...] + τ_yy = @cell yy[ip, cell...] + τ_xy = @cell xy[ip, cell...] + + tmp = τ_xy * ω_xy * 2.0 + @cell xx[ip, cell...] = muladd(dt, tmp, τ_xx) + @cell yy[ip, cell...] = muladd(dt, tmp, τ_yy) + @cell xy[ip, cell...] = muladd(dt, (τ_xx - τ_yy) * ω_xy, τ_xy) + end + + return nothing +end + +function rotation_matrix!(xx, yy, xy, ω, index, dt) + ni = size(xx) + @parallel (@idx ni) _rotation_matrix!(xx, yy, xy, ω, index, dt) + return nothing +end + +@parallel_indices (i, j) function _rotation_matrix!(xx, yy, xy, ω, index, dt) + cell = i, j + + for ip in JustRelax.cellaxes(index) + !@cell(index[ip, cell...]) && continue # no particle in this location + + θ = dt * @cell ω[ip, cell...] + sinθ, cosθ = sincos(θ) + + τ_xx = @cell xx[ip, cell...] + τ_yy = @cell yy[ip, cell...] + τ_xy = @cell xy[ip, cell...] + + R = @SMatrix [ + cosθ -sinθ + sinθ cosθ + ] + + τ = @SMatrix [ + τ_xx τ_xy + τ_xy τ_yy + ] + + # this could be fully unrolled in 2D + τr = R * τ * R' + + @cell xx[ip, cell...] = τr[1, 1] + @cell yy[ip, cell...] = τr[2, 2] + @cell xy[ip, cell...] = τr[1, 2] + end + + return nothing +end + +# STRESS ROTATION ON THE PARTICLES + +function rotate_stress_particles!( + stokes::JustRelax.StokesArrays, + τxx_vertex::T, + τyy_vertex::T, + τxx_o_vertex::T, + τyy_o_vertex::T, + τxx_p::CA, + τyy_p::CA, + τxy_p::CA, + vorticity_p::CA, + particles, + grid::JustRelax.Geometry, + dt; + fn=rotation_matrix!, +) where {T<:AbstractArray,CA} + (; xvi, xci) = grid + nx, ny = size(τxx_p) + + # interpolate stresses to particles + for (src, src_o, dst_v, dst_v_o, dst_p) in zip( + @tensor_center(stokes.τ), + @tensor_center(stokes.τ_o), + (τxx_vertex, τyy_vertex, stokes.τ.xy), + (τxx_o_vertex, τyy_o_vertex, stokes.τ_o.xy), + (τxx_p, τyy_p, τxy_p), + ) + @parallel center2vertex!(dst_v, src) + @parallel center2vertex!(dst_v_o, src_o) + @parallel (1:(nx + 1)) free_slip_y!(dst_v) + @parallel (1:(ny + 1)) free_slip_x!(dst_v) + @parallel (1:(nx + 1)) free_slip_y!(dst_v_o) + @parallel (1:(ny + 1)) free_slip_x!(dst_v_o) + grid2particle_flip!(dst_p, xvi, dst_v, dst_v_o, particles) + end + + # interpolate vorticity to particles + @parallel center2vertex!(stokes.ω.xy_v, stokes.ω.xy_c) + @parallel center2vertex!(stokes.ω_o.xy_v, stokes.ω_o.xy_c) + @parallel (1:nx) free_slip_y!(stokes.ω.xy_c) + @parallel (1:ny) free_slip_x!(stokes.ω.xy_c) + @parallel (1:(nx + 1)) free_slip_y!(stokes.ω_o.xy_v) + @parallel (1:(ny + 1)) free_slip_x!(stokes.ω_o.xy_v) + grid2particle_flip!(vorticity_p, xvi, stokes.ω.xy_v, stokes.ω_o.xy_v, particles) + # rotate old stress + fn(τxx_p, τyy_p, τxy_p, vorticity_p, particles.index, dt) + # interpolate old stress to grid arrays + particle2grid_centroid!(stokes.τ_o.xx, τxx_p, xci, particles) + particle2grid_centroid!(stokes.τ_o.yy, τyy_p, xci, particles) + return particle2grid_centroid!(stokes.τ_o.xy_c, τxy_p, xci, particles) +end diff --git a/src/stokes/VorticityKernels.jl b/src/stokes/VorticityKernels.jl new file mode 100644 index 00000000..422d1a43 --- /dev/null +++ b/src/stokes/VorticityKernels.jl @@ -0,0 +1,19 @@ + +compute_vorticity!(stokes, di) = compute_vorticity!(backend(stokes), stokes, di) + +function compute_vorticity!(::CPUBackendTrait, stokes, di) + ω_xy = stokes.ω.xy + ni = size(ω_xy) + @parallel (@idx ni) compute_vorticity!(ω_xy, @velocity(stokes)..., inv.(di)...) + + return nothing +end + +@parallel_indices (i, j) function compute_vorticity!(ω_xy, Vx, Vy, _dx, _dy) + dx(A) = _d_xa(A, i, j, _dx) + dy(A) = _d_ya(A, i, j, _dy) + + ω_xy[i, j] = 0.5 * (dx(Vy) - dy(Vx)) + + return nothing +end diff --git a/src/thermal_diffusion/Rheology.jl b/src/thermal_diffusion/Rheology.jl index 3064c9f0..274f3f06 100644 --- a/src/thermal_diffusion/Rheology.jl +++ b/src/thermal_diffusion/Rheology.jl @@ -32,8 +32,6 @@ end return conductivity * inv(heatcapacity * ρ) end -# ρ*Cp - @inline function compute_ρCp(rheology, args) return compute_heatcapacity(rheology, args) * compute_density(rheology, args) end diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl index db12d424..47067bd5 100644 --- a/src/types/constructors/stokes.jl +++ b/src/types/constructors/stokes.jl @@ -4,7 +4,7 @@ function Velocity(nx::Integer, ny::Integer) nVx = (nx + 1, ny + 2) nVy = (nx + 2, ny + 1) - Vx, Vy = @zeros(nVx...), @zeros(nVy) + Vx, Vy = @zeros(nVx...), @zeros(nVy...) return JustRelax.Velocity(Vx, Vy, nothing) end @@ -13,10 +13,26 @@ function Velocity(nx::Integer, ny::Integer, nz::Integer) nVy = (nx + 2, ny + 1, nz + 2) nVz = (nx + 2, ny + 2, nz + 1) - Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) + Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy...), @zeros(nVz...) return JustRelax.Velocity(Vx, Vy, Vz) end +## Vorticity type + +function Vorticity(nx::Integer, ny::Integer) + xy = @zeros(nx, ny) + + return JustRelax.Vorticity(nothing, nothing, xy) +end + +function Vorticity(nx::Integer, ny::Integer, nz::Integer) + yz = @zeros(nx, ny, nz) + xz = @zeros(nx, ny, nz) + xy = @zeros(nx, ny, nz) + + return JustRelax.Vorticity(yz, xz, xy) +end + ## Viscosity type function Viscosity(ni::NTuple{N,Integer}) where {N} @@ -80,6 +96,7 @@ function StokesArrays(ni::NTuple{N,Integer}) where {N} P0 = @zeros(ni...) ∇V = @zeros(ni...) V = Velocity(ni...) + ω = Vorticity(ni...) τ = SymmetricTensor(ni...) τ_o = SymmetricTensor(ni...) ε = SymmetricTensor(ni...) @@ -88,5 +105,5 @@ function StokesArrays(ni::NTuple{N,Integer}) where {N} viscosity = Viscosity(ni) R = Residual(ni...) - return JustRelax.StokesArrays(P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R) + return JustRelax.StokesArrays(P, P0, V, ∇V, ω, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R) end diff --git a/src/types/stokes.jl b/src/types/stokes.jl index e7b7cd48..8cc1de59 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -12,10 +12,10 @@ Velocity(Vx::T, Vy::T) where {T} = Velocity(Vx, Vy, nothing) Velocity(ni::NTuple{N,Number}) where {N} = Velocity(ni...) function Velocity(::Number, ::Number) - throw(ArgumentError("Velocity dimensions must be given as integers")) + throw(ArgumentError("Dimensions must be given as integers")) end function Velocity(::Number, ::Number, ::Number) - throw(ArgumentError("Velocity dimensions must be given as integers")) + throw(ArgumentError("Dimensions must be given as integers")) end ## Viscosity type @@ -32,7 +32,24 @@ Viscosity(args...) = Viscosity(promote(args...)...) Viscosity(nx::T, ny::T) where {T<:Number} = Viscosity((nx, ny)) Viscosity(nx::T, ny::T, nz::T) where {T<:Number} = Viscosity((nx, ny, nz)) function Viscosity(::NTuple{N,Number}) where {N} - throw(ArgumentError("Viscosity dimensions must be given as integers")) + throw(ArgumentError("Dimensions must be given as integers")) +end + +## Vorticity type +struct Vorticity{T} + yz::Union{T,Nothing} + xz::Union{T,Nothing} + xy::T + + function Vorticity(yz::Union{T,Nothing}, xz::Union{T,Nothing}, xy::T) where {T} + return new{T}(yz, xz, xy) + end +end + +Vorticity(nx::T, ny::T) where {T<:Number} = Vorticity((nx, ny)) +Vorticity(nx::T, ny::T, nz::T) where {T<:Number} = Vorticity((nx, ny, nz)) +function Vorticity(::NTuple{N,Number}) where {N} + throw(ArgumentError("Dimensions must be given as integers")) end ## SymmetricTensor type @@ -101,11 +118,12 @@ end ## StokesArrays type -struct StokesArrays{A,B,C,D,T} +struct StokesArrays{A,B,C,D,E,T} P::T P0::T V::A ∇V::T + ω::E τ::B ε::B ε_pl::B diff --git a/subduction/2D/LAMEM2D.jl b/subduction/2D/LAMEM2D.jl new file mode 100644 index 00000000..13432ed4 --- /dev/null +++ b/subduction/2D/LAMEM2D.jl @@ -0,0 +1,369 @@ +# using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) +# @init_parallel_stencil(CUDA, 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 +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +model = PS_Setup(:cpu, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +# model = PS_Setup(:CUDA, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("LAMEM_rheology.jl") +include("LAMEM_setup2D.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, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + particle_args = (pPhases, pT) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3π, r=1e0, CFL = 1 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + Ttop = 20 + 273 + Tbot = 1565.0 + 273 + thermal = ThermalArrays(ni) + @views thermal.T[2:end-1, :] .= PTArray(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 + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + # for _ in 1:2 + # compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + # end + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # Plitho = reverse(cumsum(reverse(ρg[2] .+ (2700*9.81), dims=2), dims=2), dims=2) + + ρg_bg = 2700 * 9.81 + # args.P .= reverse(cumsum(reverse(ρg[2] .+ ρg_bg, dims=2), dims=2), dims=2) + Plitho = reverse(cumsum(reverse((ρg[2] .+ ρg_bg).* di[2], dims=2), dims=2), dims=2) + # args.P = stokes.P .+ Plitho .- minimum(stokes.P) + + # Rheology + η = @ones(ni...) + η_vep = similar(η) + args0 = (; T = thermal.Tc, P = Plitho, dt = Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args0, rheology, (1e18, 1e24) + ) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + 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)...) + + # 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 = "log10(η)") + scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + # scatter!(ax2, Array(log10.(η[:])), Y./1e3) + # scatter!(ax2, Array(stokes.P[:]), Y./1e3) + scatter!(ax2, Array(Plitho[:]), Y./1e3) + ylims!(ax1, minimum(xvi[2])./1e3, 0) + ylims!(ax2, minimum(xvi[2])./1e3, 0) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # 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) + + # Time loop + t, it = 0.0, 0 + + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # Update buoyancy and viscosity - + Plitho .= reverse(cumsum(reverse((ρg[2] .+ ρg_bg).* di[2], dims=2), dims=2), dims=2) + # Plitho .= -(ρg[2] .+ ρg_bg) .* xci[2]' + Plitho .= stokes.P .+ Plitho .- minimum(stokes.P) + # args.P .= 0 + + args = (; T = thermal.Tc, P = Plitho, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + dt, + igg; + iterMax = 100e3, + nout = 2e3, + viscosity_cutoff = (1e18, 1e24), + 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") + @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # 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, (pT, ), (T_buffer,), xvi) + # update phase ratios + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(T_buffer), + τ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(η_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 = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[1,3], aspect = ar, title = "log10(ε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) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv])) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # Plot effective viscosity + 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) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction_LAMEM_2D" +# nx, ny = 512, 256 +# nx, ny = 512, 128 +n = 64 +nx, ny = n*6, n +nx, ny = 512, 128 +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); diff --git a/subduction/2D/LAMEM2D_nondim.jl b/subduction/2D/LAMEM2D_nondim.jl new file mode 100644 index 00000000..d43357c9 --- /dev/null +++ b/subduction/2D/LAMEM2D_nondim.jl @@ -0,0 +1,368 @@ +# using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) +# @init_parallel_stencil(CUDA, 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 +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +model = PS_Setup(:cpu, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +# model = PS_Setup(:CUDA, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("LAMEM_rheology_nondim.jl") +include("LAMEM_setup2D.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) + + thickness = 660 * km + η0 = 1e20Pa*s + CharDim = GEO_units(; + length = thickness, viscosity = η0, temperature = 1e3C + ) + + # Physical domain ------------------------------------ + ni = nx, ny # number of cells + li_nd = nondimensionalize(li[1]m, CharDim), nondimensionalize(li[2]m, CharDim) + origin_nd = nondimensionalize(origin[1]m, CharDim), nondimensionalize(origin[2]m, CharDim) + di = @. li_nd / ni # grid steps + grid = Geometry(ni, li_nd; origin = origin_nd) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies(CharDim) + dt = nondimensionalize(10e3 * 3600 * 24 * 365 * s, CharDim) # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell = 40 + max_xcell = 60 + min_xcell = 20 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + particle_args = (pPhases, pT) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.1 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + Ttop = nondimensionalize(20e0C, CharDim) + Tbot = nondimensionalize(1565e0C, CharDim) + thermal = ThermalArrays(ni) + @views thermal.T[2:end-1, :] .= PTArray(nondimensionalize(T_GMG .* K, CharDim)) + 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 + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + # for _ in 1:2 + # compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + # end + + # Rheology + η = @ones(ni...) + η_vep = similar(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + viscosity_cutoff = nondimensionalize((1e16Pa*s, 1e24Pa*s), CharDim) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, viscosity_cutoff + ) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-2 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + 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)...) + + # 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 = "log10(η)") + scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + scatter!(ax2, Array(log10.(η[:])), Y./1e3) + # scatter!(ax2, Array(ρg[2][:]), Y./1e3) + ylims!(ax1, minimum(xvi[2])./1e3, 0) + ylims!(ax2, minimum(xvi[2])./1e3, 0) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # 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) + + # Time loop + t, it = 0.0, 0 + + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + # @views T_buffer[:, end] .= Ttop + # @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # interpolate fields from particle to grid vertices + # particle2grid!(thermal.T, pT, xvi, particles) + # temperature2center!(thermal) + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, viscosity_cutoff + ) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax = 1, + nout = 1, + viscosity_cutoff = viscosity_cutoff, + free_surface = false, + viscosity_relaxation = 1e-3 + ); + end + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.τ.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # 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, (pT, ), (T_buffer,), xvi) + # update phase ratios + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(thermal.T), + τ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(η), + ) + 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 = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[1,3], aspect = ar, title = "log10(ε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) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv])) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # Plot effective viscosity + 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) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction_LAMEM_2D" +# nx, ny = 512, 256 +# nx, ny = 512, 128 +n = 64 +nx, ny = n*6, n +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); diff --git a/subduction/2D/LAMEM2D_sticky.jl b/subduction/2D/LAMEM2D_sticky.jl new file mode 100644 index 00000000..3ecc01ef --- /dev/null +++ b/subduction/2D/LAMEM2D_sticky.jl @@ -0,0 +1,424 @@ +const isCUDA = false +# const isCUDA = true + +using TimerOutputs + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +import JustRelax.@cell + +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("LAMEM_rheology.jl") +include("LAMEM_setup2D_sticky.jl") +# include("LAMEM_setup2D.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 ---------------- + ρbg = 2.7e3 + rheology = init_rheologies(; ρbg = ρbg) + rheology_augmented = init_rheologies(; ρbg = 0e0) + dt = 50e3 * 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) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + τxx_p, τyy_p, τxy_p = init_cell_arrays(particles, Val(3)) + vorticity_p, = init_cell_arrays(particles, Val(1)) + particle_args = (pT, τxx_p, τyy_p, τxy_p, vorticity_p, pPhases) + + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios_center!(phase_ratios, particles, grid, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3π, r=1e0, CFL = 1 / √2.1) # Re=3π, r=0.7 + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + Ttop = 20 + 273 + Tbot = 1565.0 + 273 + 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)) + # for _ in 1:2 + # compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + # end + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + + ρg_bg = ρbg * 9.81 + Plitho = reverse(cumsum(reverse((ρg[2] .+ ρg_bg).* di[2], dims=2), dims=2), dims=2) + + # Rheology + args0 = (T=thermal.Tc, P=Plitho, dt = Inf) + compute_viscosity!(stokes, phase_ratios, args0, rheology, (1e18, 1e24)) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + backend, rheology_augmented, phase_ratios, args0, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + 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)...) + + # 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 = "log10(η)") + scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + # scatter!(ax2, Array(log10.(η[:])), Y./1e3) + # scatter!(ax2, Array(stokes.P[:]), Y./1e3) + scatter!(ax2, Array(Plitho[:]), Y./1e3) + ylims!(ax1, minimum(xvi[2])./1e3, 0) + ylims!(ax2, minimum(xvi[2])./1e3, 0) + hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # 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 + + # Smooth out thermal field --------------------------- + for _ in 1:10 + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology_augmented, + args, + 1e6 * 3600 * 24 * 365.25, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e2, + verbose = true, + ) + ) + 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) + + # Time loop + t, it = 0.0, 0 + + # Vertice arrays of normal components of the stress tensor + τxx_vertex, τyy_vertex = @zeros(ni.+1...), @zeros(ni.+1...) + τxx_o_vertex, τyy_o_vertex = @zeros(ni.+1...), @zeros(ni.+1...) + + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # Update buoyancy and viscosity - + Plitho .= reverse(cumsum(reverse((ρg[2] .+ ρg_bg).* di[2], dims=2), dims=2), dims=2) + Plitho .= stokes.P .+ Plitho .- minimum(stokes.P) + + args = (; T = thermal.Tc, P = Plitho, dt=Inf) + # args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + # compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + # compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = ( + iterMax = 50e3, + nout = 1e3, + viscosity_cutoff = (1e18, 1e24), + 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 + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology_augmented, + args, + dt, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e2, + verbose = true, + ) + ) + to = TimerOutput() + @timeit to "subgrid time" subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology_augmented, thermal, stokes, xci, di + ) + @timeit to "centroid2par" centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + @timeit to "subgrid diff" subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + @timeit to "advect" advection!(particles, RungeKutta2(2/3), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + @timeit to "move" move_particles!(particles, xvi, particle_args) + # JustRelax.Stokes2D.rotate_stress_particles!( + # stokes, + # τxx_vertex, + # τyy_vertex, + # τxx_o_vertex, + # τyy_o_vertex, + # τxx_p, + # τyy_p, + # τxy_p, + # vorticity_p, + # particles, + # grid, + # dt + # ) + + # check if we need to inject particles + @timeit to "inject" inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + # inject_particles_phase!(particles, pPhases, (pT, τxx_p, τyy_p, τxy_p, vorticity_p), (T_buffer, τxx_vertex, τyy_vertex, stokes.τ.xy, stokes.ω.xy_v), xvi) + # update phase ratios + @timeit to "phase ratios" phase_ratios_center!(phase_ratios, particles, grid, pPhases) + + @show to + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + # checkpointing(figdir, stokes, thermal.T, η, t) + (; η_vep, η) = stokes.viscosity + if do_vtk + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(T_buffer), + τ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(η_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 = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[1,3], aspect = ar, title = "log10(ε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) + # 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) + # Plot effective viscosity + 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) + 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 = "Subduction_LAMEM_2D" +# nx, ny = 512, 256 +# nx, ny = 512, 128 +n = 128 +nx, ny = n*6, n +nx, ny = 512, 128 +# nx, ny = 32*6, 32 +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); diff --git a/subduction/2D/LAMEM_rheology.jl b/subduction/2D/LAMEM_rheology.jl new file mode 100644 index 00000000..85571a5c --- /dev/null +++ b/subduction/2D/LAMEM_rheology.jl @@ -0,0 +1,182 @@ +using GeoParams.Dislocation +using GeoParams.Diffusion + +function init_rheologies(; ρbg = 0e0) + disl_dry_olivine = SetDislocationCreep(Dislocation.dry_olivine_Hirth_2003; V = 14.5e-6) + disl_oceanic_crust = SetDislocationCreep(Dislocation.plagioclase_An75_Ji_1993) + # disl_oceanic_litho = SetDislocationCreep(Dislocation.plagioclase_An75_Ji_1993) + disl_cont_crust = SetDislocationCreep(Dislocation.wet_quartzite_Kirby_1983) + + diff_dry_olivine = SetDiffusionCreep(Diffusion.dry_olivine_Hirth_2003; V = 14.5e-6) + + ϕ_dry_olivine = sind(20) + C_dry_olivine = 30e6 + + ϕ_oceanic_crust = sind(0) + C_oceanic_crust = 5e6 + + ϕ_oceanic_litho = sind(10) + C_oceanic_litho = 5e6 + + ϕ_cont_crust = sind(20) + C_cont_crust = 30e6 + + soft_C = LinearSoftening((C_oceanic_litho*0.05, C_oceanic_litho), (0.1, 0.5)) + + elasticity = ConstantElasticity(; G=5e10, ν=0.4) + # common physical properties + α = 3e-5 # 1 / K + Cp = 1000 # J / kg K + # C = 3e6 # Pa + η_reg = 1e20 + # ρbg = 2700 # kg / m^3 + + # Define rheolgy struct + rheology = ( + # Name = "dry olivine - Hirth_Kohlstedt_2003", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=3.3e3-ρbg, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3), + CompositeRheology = CompositeRheology( + ( + disl_dry_olivine, + diff_dry_olivine, + elasticity, + DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12), + Elasticity = elasticity, + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "oceanic crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3.3e3-ρbg, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + CompositeRheology = CompositeRheology( + ( + disl_oceanic_crust, + elasticity, + DruckerPrager_regularised(; C = C_oceanic_crust, ϕ = ϕ_oceanic_crust, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(2.333e-10), + ), + # Name = "oceanic lithosphere", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3.3e3-ρbg, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3), + CompositeRheology = CompositeRheology( + ( + disl_dry_olivine, + diff_dry_olivine, + elasticity, + DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12), + Elasticity = elasticity, + ), + # Name = "continental crust", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=2.7e3-ρbg, α = α, β = 0e0, T0 = 273), + RadioactiveHeat = ConstantRadioactiveHeat(5.3571e-10), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + CompositeRheology = CompositeRheology( + ( + disl_cont_crust, + elasticity, + DruckerPrager_regularised(; C = C_cont_crust, ϕ = ϕ_cont_crust, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + Elasticity = elasticity, + ), + # Name = "continental lithosphere", + SetMaterialParams(; + Phase = 5, + Density = PT_Density(; ρ0=3.3e3-ρbg, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3), + CompositeRheology = CompositeRheology( + ( + disl_dry_olivine, + diff_dry_olivine, + elasticity, + DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12), + Elasticity = elasticity, + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 6, + Density = ConstantDensity(; ρ=1-ρbg), # water density + HeatCapacity = ConstantHeatCapacity(; Cp=3e3), + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + Conductivity = ConstantConductivity(; k=3.0), + CompositeRheology = CompositeRheology( + ( + LinearViscous(; η=1e21), + # elasticity + ) + ), + # Elasticity = elasticity, + ), + ) +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 JustRelax.cellaxes(phases) + # quick escape + @cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @cell 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 + @cell phases[ip, I...] = Float64(particle_phase) + + if pᵢ[2] ≥ 0.0 + @cell phases[ip, I...] = 6.0 + end + + end + + return nothing +end diff --git a/subduction/2D/LAMEM_rheology_nondim.jl b/subduction/2D/LAMEM_rheology_nondim.jl new file mode 100644 index 00000000..d01b83b6 --- /dev/null +++ b/subduction/2D/LAMEM_rheology_nondim.jl @@ -0,0 +1,173 @@ +using GeoParams.Diffusion +using GeoParams.Dislocation + +function init_rheologies(CharDim) + disl_dry_olivine = SetDislocationCreep(Dislocation.dry_olivine_Hirth_2003; V = 14.5e-6m^3 / mol) + disl_oceanic_crust = SetDislocationCreep(Dislocation.plagioclase_An75_Ji_1993) + # disl_oceanic_litho = SetDislocationCreep(Dislocation.plagioclase_An75_Ji_1993) + disl_cont_crust = SetDislocationCreep(Dislocation.wet_quartzite_Kirby_1983) + + Transform_DislocationCreep(Dislocation.wet_quartzite_Kirby_1983) + + diff_dry_olivine = SetDiffusionCreep(Diffusion.dry_olivine_Hirth_2003; V = 14.5e-6m^3 / mol) + + ϕ_dry_olivine = sind(20) + C_dry_olivine = 30e6Pa + + ϕ_oceanic_crust = sind(0) + C_oceanic_crust = 5e6Pa + + ϕ_oceanic_litho = sind(0) + C_oceanic_litho = 5e6Pa + + ϕ_cont_crust = sind(20) + C_cont_crust = 30e6Pa + + soft_C = LinearSoftening((C_oceanic_litho.val*0.95, C_oceanic_litho.val), (0.1, 0.5)) + + # common physical properties + α = 3e-5 / K + Cp = 1000 * J / kg * K + η_reg = 1e18Pa * s + + + # Define rheolgy struct + rheology = ( + # Name = "dry olivine - Hirth_Kohlstedt_2003", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=3.3e3kg / m^3, α = α, β = 0e0 / Pa, T0 = 273K), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3Watt/m/K), + CompositeRheology = CompositeRheology( + ( + LinearViscous(;η=1e19Pa*s), + # disl_dry_olivine, + # diff_dry_olivine, + # ConstantElasticity(; G=5e10Pa, ν=0.5), + # DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12Watt/kg), + # Elasticity = ConstantElasticity(; G=5e10Pa, ν=0.5), + Gravity = ConstantGravity(; g=9.81m/s^2), + CharDim = CharDim + ), + # Name = "oceanic crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3.3e3kg / m^3, α = α, β = 0e0 / Pa, T0 = 273K), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3Watt/m/K), + CompositeRheology = CompositeRheology( + ( + LinearViscous(;η=1e20Pa*s), + # disl_oceanic_crust, + # ConstantElasticity(; G=5e10Pa, ν=0.5), + # DruckerPrager_regularised(; C = C_oceanic_crust, ϕ = ϕ_oceanic_crust, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(2.333e-10Watt/kg), + # Elasticity = ConstantElasticity(; G=5e10Pa, ν=0.5), + CharDim = CharDim + ), + # Name = "oceanic lithosphere", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3.3e3kg / m^3, α = α, β = 0e0 / Pa, T0 = 273K), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3Watt/m/K), + CompositeRheology = CompositeRheology( + ( + LinearViscous(;η=1e19Pa*s), + # disl_dry_olivine, + # diff_dry_olivine, + # ConstantElasticity(; G=5e10Pa, ν=0.5), + # DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12Watt/kg), + # Elasticity = ConstantElasticity(; G=5e10Pa, ν=0.5), + CharDim = CharDim + ), + # Name = "continental crust", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=2.7e3kg / m^3, α = α, β = 0e0 / Pa, T0 = 273K), + RadioactiveHeat = ConstantRadioactiveHeat(5.3571e-10Watt/kg), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3Watt/m/K), + CompositeRheology = CompositeRheology( + ( + LinearViscous(;η=1e21Pa*s), + # disl_cont_crust, + # ConstantElasticity(; G=5e10Pa, ν=0.5), + # DruckerPrager_regularised(; C = C_cont_crust, ϕ = ϕ_cont_crust, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + # Elasticity = ConstantElasticity(; G=5e10Pa, ν=0.5), + CharDim = CharDim + ), + # Name = "continental lithosphere", + SetMaterialParams(; + Phase = 5, + Density = PT_Density(; ρ0=3.3e3kg / m^3, α = α, β = 0e0 / Pa, T0 = 273K), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k = 3Watt/m/K), + CompositeRheology = CompositeRheology( + ( + LinearViscous(;η=1e19Pa*s), + # disl_dry_olivine, + # diff_dry_olivine, + # ConstantElasticity(; G=5e10Pa, ν=0.5), + # DruckerPrager_regularised(; C = C_dry_olivine, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(6.6667e-12Watt/kg), + # Elasticity = ConstantElasticity(; G=5e10Pa, ν=0.5), + CharDim = CharDim + ), + ) +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 JustRelax.cellaxes(phases) + # quick escape + @cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @cell 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 + @cell phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/subduction/2D/LAMEM_setup2D.jl b/subduction/2D/LAMEM_setup2D.jl new file mode 100644 index 00000000..3001a536 --- /dev/null +++ b/subduction/2D/LAMEM_setup2D.jl @@ -0,0 +1,123 @@ +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + x = range(-2000, 2000, nx); + z = range(-660, 0, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(1280.0, nx, 1, nz); + air_thickness = 20.0 + # lith = LithosphericPhases(Layers=[15 20 55], Phases=[3 4 5], Tlab=1250) + lith = LithosphericPhases(Layers=[20 80], Phases=[1 2 0]) + # mantle = LithosphericPhases(Phases=[1]) + + # add_box!(Phases, Temp, Grid2D; xlim=(-1000, 1000), zlim=(-600.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + # Phases .= 0 + + # Add left oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(-2000, 0), + zlim =(-660.0, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = lith, + T = SpreadingRateTemp( + Tsurface = 20, + Tmantle = 1280.0, + MORside = "left", + SpreadingVel= 0.5, + AgeRidge = 0.01; + maxAge = 80.0 + ) + ) + + # Add right oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(1500, 2000), + zlim =(-660.0, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = lith, + T = SpreadingRateTemp( + Tsurface = 20, + Tmantle = 1280.0, + MORside = "right", + SpreadingVel= 0.5, + AgeRidge = 0.01; + maxAge = 80.0 + ) + ) + + # Add overriding plate margin + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 400), + zlim =(-660.0, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[25 90], Phases=[3 4 0] ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 80.0 + ) + ) + + # Add overriding plate craton + add_box!( + Phases, + Temp, + Grid2D; + xlim =(400, 1500), + zlim =(-660.0, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[35 100], Phases=[3 4 0] ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 120.0 + ) + ) + # Add slab + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 300), + zlim =(-660.0, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=30, + phase = LithosphericPhases(Layers=[30 80], Phases=[1 2 0], Tlab=1250 ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 120.0 + ) + ) + + Adiabat = 0.4 + @. Temp = Temp - Grid2D.z.val .* Adiabat + + # Lithosphere-asthenosphere boundary: + # ind = findall(Temp .> 1250 .&& (Phases.==2 .|| Phases.==5)); + # Phases[ind] .= 0; + + # surf = Grid2D.z.val .> 0.0 + # Temp[surf] .= 0.0 + # # Phases[surf] .= 7 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] .+ 1 + + return li, origin, ph, Temp[:,1,:].+273 +end diff --git a/subduction/2D/LAMEM_setup2D_sticky.jl b/subduction/2D/LAMEM_setup2D_sticky.jl new file mode 100644 index 00000000..3621735e --- /dev/null +++ b/subduction/2D/LAMEM_setup2D_sticky.jl @@ -0,0 +1,119 @@ +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + model_depth = 660 + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + x = range(-2000, 2000, nx); + z = range(-model_depth, 20, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(1280.0, nx, 1, nz); + air_thickness = 20.0 + lith = LithosphericPhases(Layers=[20 80], Phases=[1 2 0]) + + # Add left oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(-2000, 0), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = lith, + T = SpreadingRateTemp( + Tsurface = 20, + Tmantle = 1280.0, + MORside = "left", + SpreadingVel= 0.5, + AgeRidge = 0.01; + maxAge = 80.0 + ) + ) + + # Add right oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(1500, 2000), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = lith, + T = SpreadingRateTemp( + Tsurface = 20, + Tmantle = 1280.0, + MORside = "right", + SpreadingVel= 0.5, + AgeRidge = 0.01; + maxAge = 80.0 + ) + ) + + # Add overriding plate margin + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 400), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[25 90], Phases=[3 4 0] ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 80.0 + ) + ) + + # Add overriding plate craton + add_box!( + Phases, + Temp, + Grid2D; + xlim =(400, 1500), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[35 100], Phases=[3 4 0] ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 120.0 + ) + ) + # Add slab + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 500), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=30, + phase = LithosphericPhases(Layers=[30 80], Phases=[1 2 0], Tlab=1250 ), + T = HalfspaceCoolingTemp( + Tsurface = 20, + Tmantle = 1280.0, + Age = 120.0 + ) + ) + + Adiabat = 0.4 + @. Temp = Temp - Grid2D.z.val .* Adiabat + + # Lithosphere-asthenosphere boundary: + # ind = findall(Temp .> 1250 .&& (Phases.==2 .|| Phases.==5)); + # Phases[ind] .= 0; + + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 20.0 + Phases[surf] .= 5 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] .+ 1 + + return li, origin, ph, Temp[:,1,:].+273 +end \ No newline at end of file diff --git a/subduction/Buiter/Buiter2D.jl b/subduction/Buiter/Buiter2D.jl new file mode 100644 index 00000000..7e6ce148 --- /dev/null +++ b/subduction/Buiter/Buiter2D.jl @@ -0,0 +1,439 @@ +const isCUDA = false +# const isCUDA = true + +@static if isCUDA + using CUDA +end + +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +import JustRelax.@cell + +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("Buiter_rheology.jl") +include("Buiter_setup2D_sticky.jl") +# include("LAMEM_setup2D.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 ---------------- + ρbg = 3.2e3 + 1 + rheology = init_rheology_linear(; ρbg = ρbg) + rheology_augmented = init_rheology_linear(; ρbg = 0e0) + + rheology = init_rheology_nonNewtonian(; ρbg = ρbg) + rheology_augmented = init_rheology_nonNewtonian(; ρbg = 0e0) + + rheology = init_rheology_nonNewtonian_plastic(; ρbg = ρbg) + rheology_augmented = init_rheology_nonNewtonian_plastic(; ρbg = 0e0) + + dt = 50e3 * 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) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + τxx_p, τyy_p, τxy_p = init_cell_arrays(particles, Val(3)) + vorticity_p, = init_cell_arrays(particles, Val(1)) + particle_args = (pT, τxx_p, τyy_p, τxy_p, vorticity_p, pPhases) + + # Assign particles phases anomaly + phases_device = PTArray(backend)(phases_GMG) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios_center!(phase_ratios, particles, grid, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3π, r=1e0, CFL = 1 / √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)) + # for _ in 1:2 + # compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + # JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + # end + compute_ρg!(ρg[2], phase_ratios, rheology_augmented, (T=thermal.Tc, P=stokes.P)) + Plitho = reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2) + + # ρg_bg = ρbg * 9.81 + # Plitho = reverse(cumsum(reverse((ρg[2] .+ ρg_bg).* di[2], dims=2), dims=2), dims=2) + + # Rheology + args0 = (T=thermal.Tc, P=Plitho, dt = Inf) + viscosity_cutoff = (1e17, 1e24) + compute_viscosity!(stokes, phase_ratios, args0, rheology, viscosity_cutoff) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + backend, rheology_augmented, phase_ratios, args0, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + 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)...) + + # 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 = "log10(η)") + # scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + # scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) + # # scatter!(ax2, Array(stokes.P[:]), Y./1e3) + # # scatter!(ax2, Array(Plitho[:]), Y./1e3) + # ylims!(ax1, minimum(xvi[2])./1e3, 0) + # ylims!(ax2, minimum(xvi[2])./1e3, 0) + # hideydecorations!(ax2) + # # save(joinpath(figdir, "initial_profile.png"), fig) + # fig + # end + + # 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 + + # Smooth out thermal field --------------------------- + for _ in 1:10 + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology_augmented, + args0, + 1e6 * 3600 * 24 * 365.25, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e2, + verbose = true, + ) + ) + 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) + + # Time loop + t, it = 0.0, 0 + + # Vertice arrays of normal components of the stress tensor + # τxx_vertex, τyy_vertex = @zeros(ni.+1...), @zeros(ni.+1...) + # τxx_o_vertex, τyy_o_vertex = @zeros(ni.+1...), @zeros(ni.+1...) + + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= Ttop + @views T_buffer[:, 1] .= Tbot + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # Update buoyancy and viscosity - + Plitho .= reverse(cumsum(reverse((ρg[2]).* di[2], dims=2), dims=2), dims=2) + Plitho .= stokes.P .+ Plitho .- minimum(stokes.P) + + args = (; T = thermal.Tc, P = Plitho, dt=Inf, ρbg = ρbg * 9.81) + # args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + # compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + # compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology_augmented, + args, + dt, + igg; + kwargs = ( + iterMax = 150e3, + nout = 1e3, + 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 + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology_augmented, + args, + dt, + di; + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 50e3, + nout = 1e2, + verbose = true, + ) + ) + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology_augmented, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), grid_vxi, dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # JustRelax.Stokes2D.rotate_stress_particles!( + # stokes, + # τxx_vertex, + # τyy_vertex, + # τxx_o_vertex, + # τyy_o_vertex, + # τxx_p, + # τyy_p, + # τxy_p, + # vorticity_p, + # particles, + # grid, + # dt + # ) + + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + # inject_particles_phase!(particles, pPhases, (pT, τxx_p, τyy_p, τxy_p, vorticity_p), (T_buffer, τxx_vertex, τyy_vertex, stokes.τ.xy, stokes.ω.xy_v), xvi) + # update phase ratios + phase_ratios_center!(phase_ratios, particles, grid, 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 = (; + T = Array(T_buffer), + τII = Array(stokes.τ.II), + εII = Array(stokes.ε.II), + Vx = Array(Vx_v), + Vy = Array(Vy_v), + ) + 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 = "T [K] (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 = "log10(ε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) + # 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) + # Plot effective viscosity + 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) + 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 = "Buiter_2D" +# nx, ny = 512, 256 +# nx, ny = 512, 128 +n = 128 +nx, ny = n*6, n +nx, ny = 512, 256 +# nx, ny = 128, 128 +# nx, ny = 32*6, 32 +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); + +const DeviceTrait = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU" + AMDGPUBackendTrait +elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA" + CUDABackendTrait +else + CPUBackendTrait +end + +const DeviceTrait = CPUBackendTrait + +@test bk(Array) === AMDGPUBackendTrait() + \ No newline at end of file diff --git a/subduction/Buiter/Buiter_rheology.jl b/subduction/Buiter/Buiter_rheology.jl new file mode 100644 index 00000000..2b8f269c --- /dev/null +++ b/subduction/Buiter/Buiter_rheology.jl @@ -0,0 +1,143 @@ +using GeoParams.Dislocation +using GeoParams.Diffusion + +function init_rheology_nonNewtonian(; ρbg = 0e0) + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + + lithosphere_rheology = CompositeRheology( (disl_wet_olivine, diff_wet_olivine)) + init_rheologies(lithosphere_rheology; ρbg = ρbg) +end + +function init_rheology_nonNewtonian_plastic(; ρbg = 0e0) + #dislocation laws + disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # diffusion laws + diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + # plasticity + ϕ_wet_olivine = asind(0.1) + C_wet_olivine = 1e6 + η_reg = 1e19 + + lithosphere_rheology = CompositeRheology( + ( + disl_wet_olivine, + diff_wet_olivine, + DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ) + init_rheologies(lithosphere_rheology; ρbg = ρbg) +end + +function init_rheology_linear(; ρbg = 0e0) + lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e20),)) + init_rheologies(lithosphere_rheology; ρbg = ρbg) +end + +function init_rheologies(lithosphere_rheology; ρbg = 0e0) + # #dislocation laws + # disl_wet_olivine = SetDislocationCreep(Dislocation.wet_olivine1_Hirth_2003) + # # diffusion laws + # diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) + + # ϕ_wet_olivine = asind(0.1) + # C_wet_olivine = 1e6 + + # ϕ_oceanic_crust_upper = asind(0.1) + # C_oceanic_crust_upper = 0.3e6 + + # # soft_C = LinearSoftening((C_oceanic_litho*0.05, C_oceanic_litho), (0.1, 0.5)) + + # elasticity = ConstantElasticity(; G=5e10, ν=0.5) + # common physical properties + α = 2.4e-5 # 1 / K + Cp = 750 # J / kg K + + # Define rheolgy struct + rheology = ( + # Name = "Asthenoshpere", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=3.2e3-ρbg), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e20),)), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "Oceanic lithosphere", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3.2e3-ρbg, α = α, β = 0e0, T0 = 273+1474), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + # CompositeRheology = CompositeRheology( + # ( + # disl_wet_olivine, + # diff_wet_olivine, + # DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + # ) + # ), + CompositeRheology = lithosphere_rheology, + ), + # Name = "oceanic crust", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3-ρbg), + HeatCapacity = ConstantHeatCapacity(; Cp = Cp), + Conductivity = ConstantConductivity(; k = 2.5), + CompositeRheology = CompositeRheology( (LinearViscous(; η=1e20),)), + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 4, + Density = ConstantDensity(; ρ=100-ρbg), # water density + HeatCapacity = ConstantHeatCapacity(; Cp=3e3), + Conductivity = ConstantConductivity(; k=1.0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e19),)), + ), + ) +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 JustRelax.cellaxes(phases) + # quick escape + @cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @cell 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 + @cell phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/subduction/Buiter/Buiter_setup2D_sticky.jl b/subduction/Buiter/Buiter_setup2D_sticky.jl new file mode 100644 index 00000000..b90564bf --- /dev/null +++ b/subduction/Buiter/Buiter_setup2D_sticky.jl @@ -0,0 +1,83 @@ +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + model_depth = 660 + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + Tbot = 1474.0 + x = range(0, 3000, nx); + air_thickness = 10.0 + z = range(-model_depth, air_thickness, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(Tbot, nx, 1, nz); + Tlab = 1300 + # lith = LithosphericPhases(Layers=[80], Phases=[1 0], Tlab=Tlab) + + # phases + # 0: asthenosphere + # 1: lithosphere + # 2: subduction lithosphere + # 3: oceanic crust + # 4: air + add_box!( + Phases, + Temp, + Grid2D; + xlim =(0, 3000), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[], Phases=[0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add left oceanic plate + add_box!( + Phases, + Temp, + Grid2D; + xlim =(100, 3000-100), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[80], Phases=[1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add right oceanic plate crust + add_box!( + Phases, + Temp, + Grid2D; + xlim =(3000-1430, 3000-200), + zlim =(-model_depth, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=0, + phase = LithosphericPhases(Layers=[8 72], Phases=[2 1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + # Add slab + add_box!( + Phases, + Temp, + Grid2D; + xlim = (3000-1430, 3000-1430-250), + zlim =(-80, 0.0), + Origin = nothing, StrikeAngle=0, DipAngle=-30, + phase = LithosphericPhases(Layers=[8 80], Phases=[2 1 0], Tlab=Tlab), + T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0) + ) + + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 20.0 + Phases[surf] .= 3 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] .+ 1 + T = Temp[:,1,:] + + return li, origin, ph, T +end \ No newline at end of file diff --git a/subduction/GMG_setup.jl b/subduction/GMG_setup.jl new file mode 100644 index 00000000..bb6a384a --- /dev/null +++ b/subduction/GMG_setup.jl @@ -0,0 +1,42 @@ +using GeophysicalModelGenerator + +function GMG_only(nx, ny, nz) + # nx,ny,nz = 99, 33, 66 + nx, ny, nz = (nx,ny,nz) .+ 1 + x = range(-3960, 500, nx); + y = range(0, 2640, ny); + z = range(-660,0, nz); + Grid = CartData(xyz_grid(x,y,z)); + + # Now we create an integer array that will hold the `Phases` information (which usually refers to the material or rock type in the simulation) + Phases = fill(3, nx, ny, nz); + + # In many (geodynamic) models, one also has to define the temperature, so lets define it as well + Temp = fill(1350.0, nx, ny, nz); + + #### Simple free subduction setup + + # Much of the options are explained in the 2D tutorial, which can directly be transferred to 3D. + # Therefore, we will start with a simple subduction setup, which consists of a horizontal part that has a mid-oceanic ridge on one explained + + # We use a lithospheric structure. Note that if the lowermost layer has the same phase as the mantle, you can define `Tlab` as the lithosphere-asthenosphere boundary which will automatically adjust the phase depending on temperature + lith = LithosphericPhases(Layers=[15 45 10], Phases=[1 2 3], Tlab=1250) + add_box!(Phases, Temp, Grid; xlim=(-3000,-1000), ylim=(0, 1000.0), zlim=(-60.0, 0.0), phase = lith, + Origin=(-0,0,0), + T=SpreadingRateTemp(SpreadingVel=3, MORside="left"), StrikeAngle=0); + + # And an an inclined part: + add_box!(Phases, Temp, Grid; xlim=(0,300).-1000, ylim=(0, 1000.0), zlim=(-60.0, 0.0), phase = lith, + # Origin=(-1000,0,0), + T=McKenzie_subducting_slab(Tsurface=0,v_cm_yr=3), DipAngle=15, StrikeAngle=0); + # Add them to the `CartData` dataset: + Grid = addfield(Grid,(;Phases, Temp)) + + # Which looks like + write_paraview(Grid,"Initial_Setup_Subduction"); + + li = abs(last(x)-first(x)), abs(last(y)-first(y)), abs(last(z)-first(z)) + origin = (x[1], y[1], z[1]) .* 1e3 + + return li, origin, Phases, Temp +end diff --git a/subduction/GMG_setup2D.jl b/subduction/GMG_setup2D.jl new file mode 100644 index 00000000..5bd10d44 --- /dev/null +++ b/subduction/GMG_setup2D.jl @@ -0,0 +1,87 @@ + +#= +# Creating 2D numerical model setups + +### Aim +The aim of this tutorial is to show you how to create 2D numerical model setups that can be used as initial setups for other codes. + +=# + + +#= +### 2D Subduction setup + +Lets start with creating a 2D model setup in cartesian coordinates, which uses the `CartData` data structure +=# +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + x = range(-1000, 1000, nx); + z = range(-660, 20, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(1350.0, nx, 1, nz); + air_thickness = 20.0 + lith = LithosphericPhases(Layers=[air_thickness+15 20 55], Phases=[3 4 5], Tlab=1250) + # mantle = LithosphericPhases(Phases=[1]) + + # add_box!(Phases, Temp, Grid2D; xlim=(-1000, 1000), zlim=(-600.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + # Phases .= 0 + + # Lets start with defining the horizontal part of the overriding plate. + # Note that we define this twice with different thickness to deal with the bending subduction area: + add_box!(Phases, Temp, Grid2D; xlim=(200,1000), zlim=(-150.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + add_box!(Phases, Temp, Grid2D; xlim=(0,200), zlim=(-50.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + + # The horizontal part of the oceanic plate is as before: + v_spread_cm_yr = 3 #spreading velocity + lith = LithosphericPhases(Layers=[15 55], Phases=[1 2], Tlab=1250) + add_box!(Phases, Temp, Grid2D; xlim=(-1000,0.0), zlim=(-150.0, 0.0), phase = lith, T=SpreadingRateTemp(SpreadingVel=v_spread_cm_yr)); + + # Yet, now we add a trench as well. The starting thermal age at the trench is that of the horizontal part of the oceanic plate: + AgeTrench_Myrs = 1000e3/(v_spread_cm_yr/1e2)/1e6 #plate age @ trench + + # We want to add a smooth transition from a halfspace cooling 1D thermal profile to a slab that is heated by the surrounding mantle below a decoupling depth `d_decoupling`. + T_slab = LinearWeightedTemperature( + F1=HalfspaceCoolingTemp(Age=AgeTrench_Myrs), + F2=McKenzie_subducting_slab(Tsurface=0,v_cm_yr=v_spread_cm_yr, Adiabat = 0.0) + ) + + # in this case, we have a more reasonable slab thickness: + trench = Trench(Start=(0.0, -100.0), End=(0.0, 100.0), Thickness=100.0, θ_max=30.0, Length=300, Lb=200, + WeakzoneThickness=15, WeakzonePhase=6, d_decoupling=125); + add_slab!(Phases, Temp, Grid2D, trench, phase = lith, T=T_slab); + + # Lithosphere-asthenosphere boundary: + ind = findall(Temp .> 1250 .&& (Phases.==2 .|| Phases.==5)); + Phases[ind] .= 0; + + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 0.0 + Phases[surf] .= 7 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] .+ 1; + ph2 = ph .== 2 + ph3 = ph .== 3 + ph4 = ph .== 4 + ph5 = ph .== 5 + ph6 = ph .== 6 + ph7 = ph .== 7 + ph8 = ph .== 8 + ph[ph2] .= 1 + ph[ph3] .= 1 + ph[ph4] .= 2 + ph[ph5] .= 3 + ph[ph6] .= 1 + ph[ph7] .= 4 + ph[ph8] .= 5 + + return li, origin, ph, Temp[:,1,:].+273 +end diff --git a/subduction/GMG_setup2D_noair.jl b/subduction/GMG_setup2D_noair.jl new file mode 100644 index 00000000..c80dcba6 --- /dev/null +++ b/subduction/GMG_setup2D_noair.jl @@ -0,0 +1,87 @@ + +#= +# Creating 2D numerical model setups + +### Aim +The aim of this tutorial is to show you how to create 2D numerical model setups that can be used as initial setups for other codes. + +=# + + +#= +### 2D Subduction setup + +Lets start with creating a 2D model setup in cartesian coordinates, which uses the `CartData` data structure +=# +using GeophysicalModelGenerator + +function GMG_subduction_2D(nx, ny) + # Our starting basis is the example above with ridge and overriding slab + nx, nz = nx, ny + x = range(-1000, 1000, nx); + z = range(-660, 0, nz); + Grid2D = CartData(xyz_grid(x,0,z)) + Phases = zeros(Int64, nx, 1, nz); + Temp = fill(1350.0, nx, 1, nz); + air_thickness = 20.0 + lith = LithosphericPhases(Layers=[15 20 55], Phases=[3 4 5], Tlab=1250) + # mantle = LithosphericPhases(Phases=[1]) + + # add_box!(Phases, Temp, Grid2D; xlim=(-1000, 1000), zlim=(-600.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + # Phases .= 0 + + # Lets start with defining the horizontal part of the overriding plate. + # Note that we define this twice with different thickness to deal with the bending subduction area: + add_box!(Phases, Temp, Grid2D; xlim=(200,1000), zlim=(-150.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + add_box!(Phases, Temp, Grid2D; xlim=(0,200), zlim=(-50.0, 0.0), phase = lith, T=HalfspaceCoolingTemp(Age=80)); + + # The horizontal part of the oceanic plate is as before: + v_spread_cm_yr = 3 #spreading velocity + lith = LithosphericPhases(Layers=[15 55], Phases=[1 2], Tlab=1250) + add_box!(Phases, Temp, Grid2D; xlim=(-1000,0.0), zlim=(-150.0, 0.0), phase = lith, T=SpreadingRateTemp(SpreadingVel=v_spread_cm_yr)); + + # Yet, now we add a trench as well. The starting thermal age at the trench is that of the horizontal part of the oceanic plate: + AgeTrench_Myrs = 1000e3/(v_spread_cm_yr/1e2)/1e6 #plate age @ trench + + # We want to add a smooth transition from a halfspace cooling 1D thermal profile to a slab that is heated by the surrounding mantle below a decoupling depth `d_decoupling`. + T_slab = LinearWeightedTemperature( + F1=HalfspaceCoolingTemp(Age=AgeTrench_Myrs), + F2=McKenzie_subducting_slab(Tsurface=0,v_cm_yr=v_spread_cm_yr, Adiabat = 0.0) + ) + + # in this case, we have a more reasonable slab thickness: + trench = Trench(Start=(0.0, -100.0), End=(0.0, 100.0), Thickness=100.0, θ_max=30.0, Length=300, Lb=200, + WeakzoneThickness=15, WeakzonePhase=6, d_decoupling=125); + add_slab!(Phases, Temp, Grid2D, trench, phase = lith, T=T_slab); + + # Lithosphere-asthenosphere boundary: + ind = findall(Temp .> 1250 .&& (Phases.==2 .|| Phases.==5)); + Phases[ind] .= 0; + + surf = Grid2D.z.val .> 0.0 + Temp[surf] .= 0.0 + # Phases[surf] .= 7 + + Grid2D = addfield(Grid2D,(;Phases, Temp)) + + li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3 + origin = (x[1], z[1]) .* 1e3 + + ph = Phases[:,1,:] .+ 1; + ph2 = ph .== 2 + ph3 = ph .== 3 + ph4 = ph .== 4 + ph5 = ph .== 5 + ph6 = ph .== 6 + ph7 = ph .== 7 + # ph8 = ph .== 8 + ph[ph2] .= 1 + ph[ph3] .= 1 + ph[ph4] .= 2 + ph[ph5] .= 3 + ph[ph6] .= 1 + ph[ph7] .= 4 + ph[ph8] .= 5 + + return li, origin, ph, Temp[:,1,:].+273 +end diff --git a/subduction/Subduction2D.jl b/subduction/Subduction2D.jl new file mode 100644 index 00000000..2a718ff7 --- /dev/null +++ b/subduction/Subduction2D.jl @@ -0,0 +1,351 @@ +using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(CUDA, 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 +const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +# model = PS_Setup(:Threads, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +model = PS_Setup(:CUDA, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Subduction_rheology2D.jl") +include("GMG_setup2D.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, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + particle_args = (pPhases, pT) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + @views thermal.T[2:end-1, :] .= PTArray(T_GMG) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + ) + thermal_bcs!(thermal, thermal_bc) + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + for _ in 1:2 + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + end + # Rheology + η = @ones(ni...) + η_vep = similar(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-2 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true), + free_surface = true, + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) + + # 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 = "log10(η)") + scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + scatter!(ax2, Array(log10.(η[:])), Y./1e3) + # scatter!(ax2, Array(ρg[2][:]), Y./1e3) + ylims!(ax1, minimum(xvi[2])./1e3, 0) + ylims!(ax2, minimum(xvi[2])./1e3, 0) + hideydecorations!(ax2) + save(joinpath(figdir, "initial_profile.png"), fig) + fig + end + + # 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) + + # Time loop + t, it = 0.0, 0 + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= 273.0 + @views T_buffer[:, 1] .= 1623.0 + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # interpolate fields from particle to grid vertices + # particle2grid!(thermal.T, pT, xvi, particles) + # temperature2center!(thermal) + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + dt, + igg; + iterMax = 150e3, + nout = 2e3, + viscosity_cutoff = (1e18, 1e24), + free_surface = true, + # viscosity_relaxation = 1e-5 + ); + end + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vxi..., dt, 2 / 3) + # advect particles in memory + move_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, particles.coords, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(thermal.T), + τ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(η), + ) + 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 = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[1,3], aspect = ar, title = "log10(ε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) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv])) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # Plot effective viscosity + 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) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction2D" +# nx, ny = 512, 256 +nx, ny = 512, 128 +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); + diff --git a/subduction/Subduction2D_noair.jl b/subduction/Subduction2D_noair.jl new file mode 100644 index 00000000..fe4ab939 --- /dev/null +++ b/subduction/Subduction2D_noair.jl @@ -0,0 +1,353 @@ +# using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) +# @init_parallel_stencil(CUDA, 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 +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +model = PS_Setup(:cpu, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +# model = PS_Setup(:CUDA, Float64, 2) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Subduction_rheology2D_noair.jl") +include("GMG_setup2D_noair.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, nxcell, max_xcell, min_xcell, xvi, di, ni + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vxi = velocity_grids(xci, xvi, di) + # temperature + pPhases, pT = init_cell_arrays(particles, Val(2)) + particle_args = (pPhases, pT) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.99 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + @views thermal.T[2:end-1, :] .= PTArray(T_GMG) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + ) + thermal_bcs!(thermal, thermal_bc) + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(2)) + for _ in 1:2 + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + JustRelax.Stokes2D.init_P!(stokes.P, ρg[2], xci[2]) + end + # Rheology + η = @ones(ni...) + η_vep = similar(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-2 / √3 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + 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)...) + + # # 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 = "log10(η)") + # scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) + # scatter!(ax2, Array(log10.(η[:])), Y./1e3) + # # scatter!(ax2, Array(ρg[2][:]), Y./1e3) + # ylims!(ax1, minimum(xvi[2])./1e3, 0) + # ylims!(ax2, minimum(xvi[2])./1e3, 0) + # hideydecorations!(ax2) + # # save(joinpath(figdir, "initial_profile.png"), fig) + # fig + # end + + # 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) + + # Time loop + t, it = 0.0, 0 + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= 273.0 + @views T_buffer[:, 1] .= 1623.0 + @views thermal.T[2:end-1, :] .= T_buffer + thermal_bcs!(thermal, thermal_bc) + temperature2center!(thermal) + + # interpolate fields from particle to grid vertices + # particle2grid!(thermal.T, pT, xvi, particles) + # temperature2center!(thermal) + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + dt, + igg; + iterMax = 50e3, + nout = 2e3, + viscosity_cutoff = (1e18, 1e24), + free_surface = false, + # viscosity_relaxation = 1e-5 + ); + end + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + ) + centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + subgrid_diffusion!( + pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vxi..., dt, 2 / 3) + # advect particles in memory + move_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, particles.coords, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + data_v = (; + T = Array(thermal.T), + τ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(η), + ) + 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 = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") + ax3 = Axis(fig[1,3], aspect = ar, title = "log10(ε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) + # Plot particles phase + h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv])) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # Plot effective viscosity + 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) + Colorbar(fig[2,2], h2) + Colorbar(fig[1,4], h3) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +figdir = "Subduction2D" +# nx, ny = 512, 256 +# nx, ny = 512, 128 +nx, ny = 256, 64 +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); + diff --git a/subduction/Subduction3D.jl b/subduction/Subduction3D.jl new file mode 100644 index 00000000..b3464cb8 --- /dev/null +++ b/subduction/Subduction3D.jl @@ -0,0 +1,306 @@ +using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(CUDA, Float64, 3) + +using JustPIC +using JustPIC._3D +# 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 +const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +# model = PS_Setup(:Threads, Float64, 3) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +model = PS_Setup(:CUDA, Float64, 3) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Subduction_rheology.jl") +include("GMG_setup.jl") +# include("../toy.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[3] +macro all_k(A) + esc(:($A[$idx_k])) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main3D(li, origin, phases_GMG, igg; nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) + + # li, origin, T_GMG, phases_GMG = generate_model() + + # Physical domain ------------------------------------ + # lz = 700e3 # domain length in z + # lx = ly = lz # domain length in x and y + ni = nx, ny, nz # number of cells + di = @. li / ni # grid steps + # origin = 0.0, 0.0, -lz # origin coordinates (15km of 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 = init_rheologies() + dt = 10e3 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 40, 60, 20 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vx, grid_vy, grid_vz = velocity_grids(xci, xvi, di) + # temperature + pPhases, = init_cell_arrays(particles, Val(1)) + particle_args = (pPhases, ) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=5e-3, CFL = 0.99 / √3.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + # thermal_bc = TemperatureBoundaryConditions() + # thermal.T .= T_GMG + # @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + phase_ratios.center[1,1,1] + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(3)) + for _ in 1:3 + compute_ρg!(ρg[3], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + JustRelax.Stokes3D.init_P!(stokes.P, ρg[3], xci[3]) + end + # Rheology + η = @ones(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + η_vep = deepcopy(η) + + # # PT coefficients for thermal diffusion + # pt_thermal = PTThermalCoeffs( + # rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + # ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = false, front = true , back = true ), + no_slip = (left = false, right = false, top = false, bot = true, front = false, back = false), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = figdir*"\\vtk" + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + local Vx_v, Vy_v, Vz_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + Vz_v = @zeros(ni.+1...) + end + # Time loop + t, it = 0.0, 0 + while it < 1000 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # # interpolate fields from particle to grid vertices + # particle2grid!(thermal.T, pT, xvi, particles) + # temperature2center!(thermal) + + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + compute_ρg!(ρg[3], phase_ratios, rheology, args) + + # Stokes solver ---------------- + t_stokes = @elapsed begin + out = solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = (1e18, 1e24) + ); + end + println("Stokes solver time ") + println(" Total time: $t_stokes s") + println(" Time/iteration: $(t_stokes / out.iter) s") + @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # # Thermal solver --------------- + # heatdiffusion_PT!( + # thermal, + # pt_thermal, + # thermal_bc, + # rheology, + # args, + # dt, + # di; + # igg = igg, + # phase = phase_ratios, + # iterMax = 10e3, + # nout = 1e2, + # verbose = true, + # ) + # subgrid_characteristic_time!( + # subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + # ) + # centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + # subgrid_diffusion!( + # pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + # ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject = check_injection(particles) + inject && inject_particles_phase!(particles, pPhases, tuple(), tuple(), xvi) + # update phase ratios + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + data_v = (; + T = Array(thermal.T), + τxy = Array(stokes.τ.xy), + εxy = Array(stokes.ε.xy), + Vx = Array(Vx_v), + Vy = Array(Vy_v), + Vz = Array(Vz_v), + ) + data_c = (; + P = Array(stokes.P), + τxx = Array(stokes.τ.xx), + τyy = Array(stokes.τ.yy), + εxx = Array(stokes.ε.xx), + εyy = Array(stokes.ε.yy), + η = Array(η), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + Array(Vz_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + slice_j = ny >>> 1 + # Make Makie figure + fig = Figure(size = (1400, 1800), title = "t = $t") + ax1 = Axis(fig[1,1], title = "P [GPA] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], title = "τII [MPa]") + ax3 = Axis(fig[1,3], title = "log10(εII)") + ax4 = Axis(fig[2,3], title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[3].*1e-3, Array(stokes.P[:, slice_j, :]./1e9) , colormap=:lajolla) + # Plot particles phase + h2 = heatmap!(ax2, xci[1].*1e-3, xci[3].*1e-3, Array(stokes.τ.II[:, slice_j, :].*1e-6) , colormap=:batlow) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(stokes.ε.II[:, slice_j, :])) , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(η[:, slice_j, :])) , colormap=:batlow) + hideydecorations!(ax3) + hideydecorations!(ax4) + 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) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end + +## END OF MAIN SCRIPT ---------------------------------------------------------------- +do_vtk = true # set to true to generate VTK files for ParaView +# nx = 126 +# ny = 33 +# nz = 63 +# nx = 165 +# ny = 222 +# nz = 54 +nx,ny,nz = 128, 35, 101 +li, origin, phases_GMG, = GMG_only(nx, ny, nz) +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, nz; init_MPI= true)...) +else + igg +end + +# (Path)/folder where output data and figures are stored +figdir = "Subduction3D_2" +main3D(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, nz = nz, do_vtk = do_vtk); diff --git a/subduction/Subduction3D_debug.jl b/subduction/Subduction3D_debug.jl new file mode 100644 index 00000000..c79c20ea --- /dev/null +++ b/subduction/Subduction3D_debug.jl @@ -0,0 +1,396 @@ +# using CUDA +using JustRelax, JustRelax.DataIO +import JustRelax.@cell +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 3) + +using JustPIC +using JustPIC._3D +# 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 +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# setup ParallelStencil.jl environment +model = PS_Setup(:Threads, Float64, 3) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +# model = PS_Setup(:CUDA, Float64, 3) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) +environment!(model) + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays + +# Load file with all the rheology configurations +include("Subduction_rheology.jl") +include("GMG_setup.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +import ParallelStencil.INDICES +const idx_k = INDICES[3] +macro all_k(A) + esc(:($A[$idx_k])) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_k(z)) * <(@all_k(z), 0.0) + return nothing +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main3D(li, origin, phases_GMG, igg; nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) + + # li, origin, T_GMG, phases_GMG = generate_model() + + # Physical domain ------------------------------------ + # lz = 700e3 # domain length in z + # lx = ly = lz # domain length in x and y + ni = nx, ny, nz # number of cells + di = @. li / ni # grid steps + # origin = 0.0, 0.0, -lz # origin coordinates (15km of 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 = init_rheologies() + dt = 10e3 * 3600 * 24 * 365 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 25, 35, 8 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + ) + subgrid_arrays = SubgridDiffusionCellArrays(particles) + # velocity grids + grid_vx, grid_vy, grid_vz = velocity_grids(xci, xvi, di) + # temperature + pPhases, = init_cell_arrays(particles, Val(1)) + particle_args = (pPhases, ) + + # Assign particles phases anomaly + phases_device = PTArray(phases_GMG) + init_phases!(pPhases, phases_device, particles, xvi) + phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-3, CFL = 0.95 / √3.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + # thermal_bc = TemperatureBoundaryConditions() + # thermal.T .= T_GMG + # @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(3)) + compute_ρg!(ρg[3], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[3], xci[3]) + # Rheology + η = @ones(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + η_vep = deepcopy(η) + + # # PT coefficients for thermal diffusion + # pt_thermal = PTThermalCoeffs( + # rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + # ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = false , front = true , back = true ), + no_slip = (left = false, right = false, top = false, bot = true, front = false, back = false), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + + # IO ------------------------------------------------- + # if it does not exist, make folder where figures are stored + if do_vtk + vtk_dir = figdir*"\\vtk" + take(vtk_dir) + end + take(figdir) + # ---------------------------------------------------- + + # # Plot initial T and η profiles + # fig = let + # Zv = [z for x in xvi[1], y in xvi[2], z in xvi[3]][:] + # Z = [z for x in xci[1], y in xci[2], z in xci[3]][:] + # 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(η)") + # lines!(ax1, Array(thermal.T[:]), Zv./1e3) + # lines!(ax2, Array(log10.(η[:])), Z./1e3) + # ylims!(ax1, minimum(xvi[3])./1e3, 0) + # ylims!(ax2, minimum(xvi[3])./1e3, 0) + # hideydecorations!(ax2) + # save(joinpath(figdir, "initial_profile.png"), fig) + # fig + # end + + # grid2particle!(pT, xvi, thermal.T, particles) + # dt₀ = similar(stokes.P) + + local Vx_v, Vy_v, Vz_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + Vz_v = @zeros(ni.+1...) + end + # Time loop + t, it = 0.0, 0 + while it < 150 # run only for 5 Myrs + # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + + # # interpolate fields from particle to grid vertices + # particle2grid!(thermal.T, pT, xvi, particles) + # temperature2center!(thermal) + + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_viscosity!( + η, 1.0, phase_ratios, stokes, args, rheology, (1e18, 1e24) + ) + compute_ρg!(ρg[3], phase_ratios, rheology, args) + + # Stokes solver ---------------- + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax = 1, + nout = 1, + viscosity_cutoff = (1e18, 1e24) + ); + @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di) + # ------------------------------ + + # # Thermal solver --------------- + # heatdiffusion_PT!( + # thermal, + # pt_thermal, + # thermal_bc, + # rheology, + # args, + # dt, + # di; + # igg = igg, + # phase = phase_ratios, + # iterMax = 10e3, + # nout = 1e2, + # verbose = true, + # ) + # subgrid_characteristic_time!( + # subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di + # ) + # centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) + # subgrid_diffusion!( + # pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi, di, dt + # ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # check if we need to inject particles + inject = check_injection(particles) + inject && inject_particles_phase!(particles, pPhases, tuple(), tuple(), xvi) + # update phase ratios + phase_ratios_center!(phase_ratios, particles, xci, di, pPhases) + + @show it += 1 + t += dt + + # Data I/O and plotting --------------------- + if it == 1 || rem(it, 5) == 0 + checkpointing(figdir, stokes, thermal.T, η, t) + + if do_vtk + JustRelax.velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + data_v = (; + T = Array(thermal.T), + τxy = Array(stokes.τ.xy), + εxy = Array(stokes.ε.xy), + Vx = Array(Vx_v), + Vy = Array(Vy_v), + Vz = Array(Vz_v), + ) + data_c = (; + P = Array(stokes.P), + τxx = Array(stokes.τ.xx), + τyy = Array(stokes.τ.yy), + εxx = Array(stokes.ε.xx), + εyy = Array(stokes.ε.yy), + η = Array(η), + ) + velocity_v = ( + Array(Vx_v), + Array(Vy_v), + Array(Vz_v), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c, + velocity_v + ) + end + + slice_j = ny >>> 1 + # Make Makie figure + fig = Figure(size = (1400, 1800), title = "t = $t") + ax1 = Axis(fig[1,1], title = "P [GPA] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], title = "τII [MPa]") + ax3 = Axis(fig[1,3], title = "log10(εII)") + ax4 = Axis(fig[2,3], title = "log10(η)") + # Plot temperature + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[3].*1e-3, Array(stokes.P[:, slice_j, :]./1e9) , colormap=:lajolla) + # Plot particles phase + h2 = heatmap!(ax2, xci[1].*1e-3, xci[3].*1e-3, Array(stokes.τ.II[:, slice_j, :].*1e-6) , colormap=:batlow) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(stokes.ε.II[:, slice_j, :])) , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(η[:, slice_j, :])) , colormap=:batlow) + hideydecorations!(ax3) + hideydecorations!(ax4) + 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) + save(joinpath(figdir, "$(it).png"), fig) + fig + end + # ------------------------------ + + end + + return nothing +end +## END OF MAIN SCRIPT ---------------------------------------------------------------- + +do_vtk = true # set to true to generate VTK files for ParaView +nx = 126 +ny = 33 +nz = 63 +li, origin, T_GMG, phases_GMG = generate_model() +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, nz; init_MPI= true)...) +else + igg +end + +# (Path)/folder where output data and figures are stored +figdir = "Subduction3D" +main3D(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, nz = nz, do_vtk = do_vtk); + +# @code_warntype main3D(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, nz = nz, do_vtk = do_vtk); + +ProfileCanvas.@profview begin + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax = 1e2, + nout = 1e2, + viscosity_cutoff = (1e18, 1e24) + ); +end + +# environment!(model) +# @code_warntype solve!( +# stokes, +# pt_stokes, +# di, +# flow_bcs, +# ρg, +# η, +# η_vep, +# phase_ratios, +# rheology, +# args, +# Inf, +# igg; +# iterMax = 2, +# nout = 1, +# viscosity_cutoff = (1e18, 1e24) +# ); + +# foo(stokes) = @copy stokes.P0 stokes.P +# @descend foo(stokes) + + +τij = rand(), rand(), rand(), rand(), rand(), rand() +τij_o = rand(), rand(), rand(), rand(), rand(), rand() +εij = rand(), rand(), rand(), rand(), rand(), rand() +ηij = rand() +_Gdt = rand() +dτ_r = rand() + +function foo( + τij::NTuple{N,T}, τij_o::NTuple{N,T}, ηij, εij::NTuple{N,T}, _Gdt, dτ_r +) where {N,T} + dτij = ntuple(Val(N)) do i + Base.@_inline_meta + dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i])) + end + return dτij, second_invariant((τij .+ dτij)...) +end + +@b foo($(τij, τij_o, ηij, εij, _Gdt, dτ_r)...) # 3.228ns + +@generated function foo2( + τij::NTuple{N,T}, τij_o::NTuple{N,T}, ηij, εij::NTuple{N,T}, _Gdt, dτ_r +) where {N,T} + quote + Base.@_inline_meta + Base.@nexprs $N i -> begin + τij_n = τij[i] + dτ_i = dτ_r * fma(2.0 * ηij, εij[i], fma(-((τij_n - τij_o[i])) * ηij, _Gdt, -τij_n)) + pt_τ_i = τij_n + dτ_i + end + dτij = Base.@ncall $N tuple dτ + pt_τII = Base.@ncall $N second_invariant pt_τ + return dτij, pt_τII + end +end +@be foo2($(τij, τij_o, ηij, εij, _Gdt, dτ_r)...) # 2.99ns + +foo2(τij, τij_o, ηij, εij, _Gdt, dτ_r) == foo2(τij, τij_o, ηij, εij, _Gdt, dτ_r) diff --git a/subduction/Subduction_rheology2D.jl b/subduction/Subduction_rheology2D.jl new file mode 100644 index 00000000..ca72d253 --- /dev/null +++ b/subduction/Subduction_rheology2D.jl @@ -0,0 +1,181 @@ +# from "Fingerprinting secondary mantle plumes", Cloetingh et al. 2022 + + + +function init_rheologies() + disl_dry_olivine = DislocationCreep(A=2.5e-17 , n=3.5, E=532e3, V=0e0 , r=0.0, R=8.3145) + disl_wet_olivine = DislocationCreep(A=9e-20 , n=3.5, E=480e3, V=11e-6, r=0.0, R=8.3145) + disl_wet_quartzite = DislocationCreep(A=1.97e-17, n=2.3, E=164e3, V=0e0 , r=0.0, R=8.3145) + disl_plagioclase = DislocationCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + disl_gabbro = DislocationCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + + diff_dry_olivine = DiffusionCreep(A=2.5e-17 , n=3.5, E=532e3, V=0e0 , r=0.0, R=8.3145) + diff_wet_olivine = DiffusionCreep(A=9e-20 , n=3.5, E=480e3, V=11e-6, r=0.0, R=8.3145) + diff_wet_quartzite = DiffusionCreep(A=1.97e-17, n=2.3, E=164e3, V=0e0 , r=0.0, R=8.3145) + diff_plagioclase = DiffusionCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + diff_gabbro = DiffusionCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + + ϕ_dry_olivine = asind(0.6) + ϕ_wet_olivine = asind(0.1) + ϕ_wet_quartzite = asind(0.3) + ϕ_plagioclase = asind(0.3) + + # common physical properties + α = 3e-5 # 1 / K + Cp = 1000 # J / kg K + C = 3e6 # Pa + η_reg = 1e18 + + # Define rheolgy struct + rheology = ( + # Name = "dry ol - lithospheric mantle", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=3.3e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_dry_olivine, + diff_dry_olivine, + DruckerPrager_regularised(; C = C, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(0.022), + # Elasticity = el_upper_crust, + Gravity = ConstantGravity(; g=9.81), + ), + # # Name = "gabbro - oceanic lithosphere", + # SetMaterialParams(; + # Phase = 2, + # Density = PT_Density(; ρ0=3e3, α = α, T0 = 273), + # CompositeRheology = CompositeRheology( + # ( + # disl_gabro, + # diff_gabro, + # DruckerPrager_regularised(; C = C, ϕ=ϕ_gabbro, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + # ) + # ), + # RadioactiveHeat = ConstantRadioactiveHeat(0.022), + # # Elasticity = el_upper_crust, + # Gravity = ConstantGravity(; g=9.81), + # ), + # # Name = "lower slab", + # SetMaterialParams(; + # Phase = 3, + # Density = ConstantDensity(; ρ=3.28e3), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e23), ) ), + # # Elasticity = el_upper_crust, + # ), + # Name = "wet qtz - upper continental crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=2.75e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_wet_quartzite, + diff_wet_quartzite, + DruckerPrager_regularised(; C = C, ϕ=ϕ_wet_quartzite, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(2), + # Elasticity = el_upper_crust, + ), + # Name = "plagioclase - lower continental crust", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_plagioclase, + diff_plagioclase, + DruckerPrager_regularised(; C = C, ϕ=ϕ_plagioclase, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(0.2), + # Elasticity = el_upper_crust, + ), + # # Name = "lithosphere", + # SetMaterialParams(; + # Phase = 6, + # Density = PT_Density(; ρ0=3.3e3, α = α, T0 = 273), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e23), ) ), + # # Elasticity = el_upper_crust, + # ), + # Name = "wet ol - weak zone", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=3.3e3, α = α, β = 0e0, T0 = 273), + RadioactiveHeat = ConstantRadioactiveHeat(0.022), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e21), ) ), + CompositeRheology = CompositeRheology( + ( + disl_wet_olivine, + diff_wet_olivine, + DruckerPrager_regularised(; C = 5e6, ϕ=ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + # Elasticity = el_upper_crust, + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 5, + Density = ConstantDensity(; ρ=1e3), # water density + HeatCapacity = ConstantHeatCapacity(; Cp=3e3), + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + Conductivity = ConstantConductivity(; k=1.0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + ), + ) +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 JustRelax.cellaxes(phases) + # quick escape + @cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @cell 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 + @cell phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/subduction/Subduction_rheology2D_noair.jl b/subduction/Subduction_rheology2D_noair.jl new file mode 100644 index 00000000..c9ef68c9 --- /dev/null +++ b/subduction/Subduction_rheology2D_noair.jl @@ -0,0 +1,172 @@ +# from "Fingerprinting secondary mantle plumes", Cloetingh et al. 2022 + + + +function init_rheologies() + disl_dry_olivine = DislocationCreep(A=2.5e-17 , n=3.5, E=532e3, V=0e0 , r=0.0, R=8.3145) + disl_wet_olivine = DislocationCreep(A=9e-20 , n=3.5, E=480e3, V=11e-6, r=0.0, R=8.3145) + disl_wet_quartzite = DislocationCreep(A=1.97e-17, n=2.3, E=164e3, V=0e0 , r=0.0, R=8.3145) + disl_plagioclase = DislocationCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + disl_gabbro = DislocationCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + + diff_dry_olivine = DiffusionCreep(A=2.5e-17 , n=3.5, E=532e3, V=0e0 , r=0.0, R=8.3145) + diff_wet_olivine = DiffusionCreep(A=9e-20 , n=3.5, E=480e3, V=11e-6, r=0.0, R=8.3145) + diff_wet_quartzite = DiffusionCreep(A=1.97e-17, n=2.3, E=164e3, V=0e0 , r=0.0, R=8.3145) + diff_plagioclase = DiffusionCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + diff_gabbro = DiffusionCreep(A=4.8e-22 , n=3.2, E=238e3, V=0e0 , r=0.0, R=8.3145) + + ϕ_dry_olivine = asind(0.6) + ϕ_wet_olivine = asind(0.1) + ϕ_wet_quartzite = asind(0.3) + ϕ_plagioclase = asind(0.3) + + # common physical properties + α = 3e-5 # 1 / K + Cp = 1000 # J / kg K + C = 3e6 # Pa + η_reg = 1e18 + + # Define rheolgy struct + rheology = ( + # Name = "dry ol - lithospheric mantle", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=3.3e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_dry_olivine, + diff_dry_olivine, + DruckerPrager_regularised(; C = C, ϕ=ϕ_dry_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(0.022), + # Elasticity = el_upper_crust, + Gravity = ConstantGravity(; g=9.81), + ), + # # Name = "gabbro - oceanic lithosphere", + # SetMaterialParams(; + # Phase = 2, + # Density = PT_Density(; ρ0=3e3, α = α, T0 = 273), + # CompositeRheology = CompositeRheology( + # ( + # disl_gabro, + # diff_gabro, + # DruckerPrager_regularised(; C = C, ϕ=ϕ_gabbro, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + # ) + # ), + # RadioactiveHeat = ConstantRadioactiveHeat(0.022), + # # Elasticity = el_upper_crust, + # Gravity = ConstantGravity(; g=9.81), + # ), + # # Name = "lower slab", + # SetMaterialParams(; + # Phase = 3, + # Density = ConstantDensity(; ρ=3.28e3), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e23), ) ), + # # Elasticity = el_upper_crust, + # ), + # Name = "wet qtz - upper continental crust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=2.75e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_wet_quartzite, + diff_wet_quartzite, + DruckerPrager_regularised(; C = C, ϕ=ϕ_wet_quartzite, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(2), + # Elasticity = el_upper_crust, + ), + # Name = "plagioclase - lower continental crust", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3e3, α = α, β = 0e0, T0 = 273), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e20), ) ), + CompositeRheology = CompositeRheology( + ( + disl_plagioclase, + diff_plagioclase, + DruckerPrager_regularised(; C = C, ϕ=ϕ_plagioclase, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + RadioactiveHeat = ConstantRadioactiveHeat(0.2), + # Elasticity = el_upper_crust, + ), + # # Name = "lithosphere", + # SetMaterialParams(; + # Phase = 6, + # Density = PT_Density(; ρ0=3.3e3, α = α, T0 = 273), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e23), ) ), + # # Elasticity = el_upper_crust, + # ), + # Name = "wet ol - weak zone", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=3.3e3, α = α, β = 0e0, T0 = 273), + RadioactiveHeat = ConstantRadioactiveHeat(0.022), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp), + Conductivity = ConstantConductivity(; k =3 ), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 1e21), ) ), + CompositeRheology = CompositeRheology( + ( + disl_wet_olivine, + diff_wet_olivine, + DruckerPrager_regularised(; C = 5e6, ϕ=ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + ) + ), + # Elasticity = el_upper_crust, + ), + ) +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 JustRelax.cellaxes(phases) + # quick escape + @cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + @cell 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 + @cell phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end diff --git a/subduction/Subduction_rheology3D.jl b/subduction/Subduction_rheology3D.jl new file mode 100644 index 00000000..76f2c59a --- /dev/null +++ b/subduction/Subduction_rheology3D.jl @@ -0,0 +1,76 @@ +# from "Fingerprinting secondary mantle plumes", Cloetingh et al. 2022 + +function init_rheologies() + # Define rheolgy struct + rheology = ( + # Name = "crust", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(; ρ=3.28e3), + CompositeRheology = CompositeRheology( (LinearViscous(η = 1e23), ) ), + # Elasticity = el_upper_crust, + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "slab", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(; ρ=3.28e3), + # CompositeRheology = CompositeRheology( (LinearViscous(η = 2e23), ) ), + CompositeRheology = CompositeRheology( (LinearViscous(η = 2e23), ) ), + # Elasticity = el_upper_crust, + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "mantle", + SetMaterialParams(; + Phase = 3, + Density = ConstantDensity(; ρ=3.2e3), + CompositeRheology = CompositeRheology( (LinearViscous(η = 1e21), ) ), + # Elasticity = el_upper_crust, + 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 JustRelax.cellaxes(phases) + # quick escape + JustRelax.@cell(index[ip, I...]) == 0 && continue + + pᵢ = ntuple(Val(N)) do i + JustRelax.@cell pcoords[i][ip, I...] + end + + d = Inf # distance to the nearest particle + particle_phase = -1 + for offi in 0:1, offj in 0:1, offk in 0:1 + ii, jj, kk = I[1] + offi, I[2] + offj, I[3] + offk + + !(ii ≤ ni[1]) && continue + !(jj ≤ ni[2]) && continue + !(kk ≤ ni[3]) && continue + + xvᵢ = ( + xvi[1][ii], + xvi[2][jj], + xvi[3][kk], + ) + # @show xvᵢ ii jj kk + d_ijk = √(sum((pᵢ[i] - xvᵢ[i])^2 for i in 1:N)) + if d_ijk < d + d = d_ijk + particle_phase = phase_grid[ii, jj, kk] + end + end + JustRelax.@cell phases[ip, I...] = Float64(particle_phase) + end + + return nothing +end \ No newline at end of file diff --git a/test/test_boundary_conditions2D.jl b/test/test_boundary_conditions2D.jl index 14e9d93d..2bcac240 100644 --- a/test/test_boundary_conditions2D.jl +++ b/test/test_boundary_conditions2D.jl @@ -68,3 +68,18 @@ const backend = CPUBackend @test @views stokes.V.Vx[: , end] == -stokes.V.Vx[: , end - 1] end end + +@testset "Temperature boundary conditions 2D" begin + ni = 5, 5 # number of elements + thermal = ThermalArrays(ni) + # free-slip + bcs = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = true, bot = true), + ) + thermal_bcs!(thermal, bcs) + + @test @views thermal.T[ :, 1] == thermal.T[ :, 2] + @test @views thermal.T[ :, end] == thermal.T[ :, end - 1] + @test @views thermal.T[ 1, :] == thermal.T[ 2, :] + @test @views thermal.T[end, :] == thermal.T[end - 1, :] +end diff --git a/test/test_boundary_conditions3D.jl b/test/test_boundary_conditions3D.jl new file mode 100644 index 00000000..9353b37e --- /dev/null +++ b/test/test_boundary_conditions3D.jl @@ -0,0 +1,153 @@ +using Test, JustRelax, ParallelStencil +@init_parallel_stencil(Threads, Float64, 3) + +model = PS_Setup(:cpu, Float64, 3) +environment!(model) + +@testset "Flow boundary conditions 3D" begin + n = 5 # number of elements + Vx, Vy, Vz = @rand(n + 1, n + 2, n + 2), @rand(n + 2, n + 1, n + 2), @rand(n + 2, n + 2, n + 1) + # free-slip + bcs = FlowBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true , front = true , back = true ), + no_slip = (left = false, right = false, top = false, bot = false, front = false, back = false), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + flow_bcs!(bcs, Vx, Vy, Vz) + + # Vx + @test @views Vx[:, :, 1] == Vx[:, :, 2] # bottom + @test @views Vx[:, :, end] == Vx[:, :, end-1] # top + @test @views Vx[:, 1, :] == Vx[:, 2, :] # left + @test @views Vx[:, end, :] == Vx[:, end-1, :] # right + # Vy + @test @views Vy[:, :, 1] == Vy[:, :, 2] # bottom + @test @views Vy[:, :, end] == Vy[:, :, end-1] # top + @test @views Vy[ 1, :, :] == Vy[ 2, :, :] # front + @test @views Vy[end, :, :] == Vy[end-1, :, :] # back + # Vz + @test @views Vz[:, 1, :] == Vz[:, 2, :] # left + @test @views Vz[:, end, :] == Vz[:, end-1, :] # right + @test @views Vz[ 1, :, :] == Vz[ 2, :, :] # front + @test @views Vz[end, :, :] == Vz[end-1, :, :] # back + + # no-slip + bcs = FlowBoundaryConditions(; + free_slip = (left = false, right = false, top = false, bot = false, front = false, back = false), + no_slip = (left = true , right = true , top = true , bot = true , front = true , back = true ), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + Vx, Vy, Vz = @rand(n + 1, n + 2, n + 2), @rand(n + 2, n + 1, n + 2), @rand(n + 2, n + 2, n + 1) + flow_bcs!(bcs, Vx, Vy, Vz) + + # Test the ones that are zero + @test sum(!iszero(Vx[ 1, i, j]) for i in axes(Vx, 2), j in axes(Vx, 3)) == 0 # left + @test sum(!iszero(Vx[end, i, j]) for i in axes(Vx, 2), j in axes(Vx, 3)) == 0 # right + @test sum(!iszero(Vy[i, 1, j]) for i in axes(Vy, 1), j in axes(Vy, 3)) == 0 # front + @test sum(!iszero(Vy[i, end, j]) for i in axes(Vy, 1), j in axes(Vy, 3)) == 0 # back + @test sum(!iszero(Vz[i, j, 1]) for i in axes(Vz, 1), j in axes(Vz, 2)) == 0 # bottom + @test sum(!iszero(Vz[i, j, end]) for i in axes(Vz, 1), j in axes(Vz, 2)) == 0 # top + + # Vx + @test @views Vx[:, :, 1] == -Vx[:, :, 2] # bottom + @test @views Vx[:, :, end] == -Vx[:, :, end-1] # top + @test @views Vx[:, 1, :] == -Vx[:, 2, :] # left + @test @views Vx[:, end, :] == -Vx[:, end-1, :] # right + # Vy + @test @views Vy[ :, :, 1] == -Vy[ :, :, 2] # bottom + @test @views Vy[ :, :, end] == -Vy[ :, :, end-1] # top + @test @views Vy[ 1, :, :] == -Vy[ 2, :, :] # front + @test @views Vy[end, :, :] == -Vy[end-1, :, :] # back + # Vz + @test @views Vz[ :, 1, :] == -Vz[ :, 2, :] # left + @test @views Vz[ :, end, :] == -Vz[ :, end-1, :] # right + @test @views Vz[ 1, :, :] == -Vz[ 2, :, :] # front + @test @views Vz[end, :, :] == -Vz[end-1, :, :] # back + + # test with StokesArrays + ni = 5, 5, 5 + stokes = StokesArrays(ni, ViscoElastic) + # free-slip + bcs = FlowBoundaryConditions(; + free_slip = (left = true , right = true , top = true , bot = true , front = true , back = true ), + no_slip = (left = false, right = false, top = false, bot = false, front = false, back = false), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + flow_bcs!(stokes, bcs) + + # Vx + @test @views stokes.V.Vx[:, :, 1] == stokes.V.Vx[:, :, 2] # bottom + @test @views stokes.V.Vx[:, :, end] == stokes.V.Vx[:, :, end-1] # top + @test @views stokes.V.Vx[:, 1, :] == stokes.V.Vx[:, 2, :] # left + @test @views stokes.V.Vx[:, end, :] == stokes.V.Vx[:, end-1, :] # right + # Vy + @test @views stokes.V.Vy[:, :, 1] ==stokes.V. Vy[:, :, 2] # bottom + @test @views stokes.V.Vy[:, :, end] ==stokes.V. Vy[:, :, end-1] # top + @test @views stokes.V.Vy[ 1, :, :] ==stokes.V. Vy[ 2, :, :] # front + @test @views stokes.V.Vy[end, :, :] ==stokes.V. Vy[end-1, :, :] # back + # Vz + @test @views stokes.V.Vz[:, 1, :] == stokes.V.Vz[:, 2, :] # left + @test @views stokes.V.Vz[:, end, :] == stokes.V.Vz[:, end-1, :] # right + @test @views stokes.V.Vz[ 1, :, :] == stokes.V.Vz[ 2, :, :] # front + @test @views stokes.V.Vz[end, :, :] == stokes.V.Vz[end-1, :, :] # back + + # no-slip + bcs = FlowBoundaryConditions(; + free_slip = (left = false, right = false, top = false, bot = false, front = false, back = false), + no_slip = (left = true , right = true , top = true , bot = true , front = true , back = true ), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + flow_bcs!(stokes, bcs) + + # Test the ones that are zero + @test sum(!iszero(stokes.V.Vx[ 1, i, j]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # left + @test sum(!iszero(stokes.V.Vx[end, i, j]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # right + @test sum(!iszero(stokes.V.Vy[i, 1, j]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # front + @test sum(!iszero(stokes.V.Vy[i, end, j]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # back + @test sum(!iszero(stokes.V.Vz[i, j, 1]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # bottom + @test sum(!iszero(stokes.V.Vz[i, j, end]) for i in axes(stokes.V.Vx, 2), j in axes(stokes.V.Vx, 3)) == 0 # top + + # Vx + @test @views stokes.V.Vx[:, :, 1] == -stokes.V.Vx[:, :, 2] # bottom + @test @views stokes.V.Vx[:, :, end] == -stokes.V.Vx[:, :, end-1] # top + @test @views stokes.V.Vx[:, 1, :] == -stokes.V.Vx[:, 2, :] # left + @test @views stokes.V.Vx[:, end, :] == -stokes.V.Vx[:, end-1, :] # right + # Vy + @test @views stokes.V.Vy[ :, :, 1] == -stokes.V.Vy[ :, :, 2] # bottom + @test @views stokes.V.Vy[ :, :, end] == -stokes.V.Vy[ :, :, end-1] # top + @test @views stokes.V.Vy[ 1, :, :] == -stokes.V.Vy[ 2, :, :] # front + @test @views stokes.V.Vy[end, :, :] == -stokes.V.Vy[end-1, :, :] # back + # Vz + @test @views stokes.V.Vz[ :, 1, :] == -stokes.V.Vz[ :, 2, :] # left + @test @views stokes.V.Vz[ :, end, :] == -stokes.V.Vz[ :, end-1, :] # right + @test @views stokes.V.Vz[ 1, :, :] == -stokes.V.Vz[ 2, :, :] # front + @test @views stokes.V.Vz[end, :, :] == -stokes.V.Vz[end-1, :, :] # back + +end + +@testset "Temperature boundary conditions 3D" begin + ni = 5, 5, 5 # number of elements + thermal = ThermalArrays(ni) + # free-slip + bcs = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = true, bot = true, front = true, back = true), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + thermal_bcs!(thermal, bcs) + + # Vx + @test @views thermal.T[:, :, 1] == thermal.T[:, :, 2] # bottom + @test @views thermal.T[:, :, end] == thermal.T[:, :, end-1] # top + @test @views thermal.T[:, 1, :] == thermal.T[:, 2, :] # left + @test @views thermal.T[:, end, :] == thermal.T[:, end-1, :] # right + # Vy + @test @views thermal.T[:, :, 1] == thermal.T[:, :, 2] # bottom + @test @views thermal.T[:, :, end] == thermal.T[:, :, end-1] # top + @test @views thermal.T[ 1, :, :] == thermal.T[ 2, :, :] # front + @test @views thermal.T[end, :, :] == thermal.T[end-1, :, :] # back + # Vz + @test @views thermal.T[:, 1, :] == thermal.T[:, 2, :] # left + @test @views thermal.T[:, end, :] == thermal.T[:, end-1, :] # right + @test @views thermal.T[ 1, :, :] == thermal.T[ 2, :, :] # front + @test @views thermal.T[end, :, :] == thermal.T[end-1, :, :] # back +end