From c02d41a233467e2f6d12e950bc9cf483cd087405 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat <58044444+albert-de-montserrat@users.noreply.github.com> Date: Fri, 20 Oct 2023 13:01:22 +0200 Subject: [PATCH] Coupling 3D solvers with JustPIC (#62) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Pin to ParallelStencik 0.9.0 * update density kernel syntax * generalize Delta function * use triple equality in Delta function * explicitly add fma's * update indexing syntax * move Delta to Utils.jl * format & add @inbounds * format changes * 3d multiphase sxript * more updates to 3d script * generalize compute_maxloc! * generalize phase_ratios_center * fix 3D viscosity kernel * some cosmetic changes * include StressKernels.jl; export tensor_invariant! * export only tensor_invariant! * generalize multi_copyto! * update PS syntax of pressure kernels * syntax updates in stress kernels * syntax updates in pressure kernels * do not consider P in isyielding * move @idx to the top * include StressKernels.jl in the appropiate file. * some fixes in 2D convection miniapp * updates in 3D scripts * minor fixes in Stokes2D * re-structuring and cleaning of Stokes3D.jl * formatting * minimal changes to some Utils funs * dimension agnostic stress update * fix NaNs on eff visc when strain is zero * some fixes * update 3d scripts * make plasticity great again * Plasticity is great again * fixed shear stress at the vertices * 3D script * call correct viscosity kernel in the 3D case * 3D script * update 3D stress kernels * update Stokes3D * format * update scripts * fix 3D vertex stress kernel * simplify _compute_τ_nonlinear! * format and remove unused arguments * update P kernels syntax * updates * improve multi_copy! * fix 3D center2vertex tensor interpolation * fix/improve few thins * update PS syntax * fix 3D stokes * format * add volumetric plasticity component * fix voluemtric plasticity * copy P into P0 * fix small bug * few tweaks * force fma * format and remove commented lines * format * add 1.10 to CI * fix conflicts * replace delta by allzero * update convection miniapp * update scripts * fix merge errors * rework @idx * fix merge errors * take function * fix stress update * update viscosity calls * update 3D scripts * remove dead code * fix 3D stuff * move 3D scripts to own folder * tweak compats * format * fix IGG dimensions * remove JustPIC from Project.toml * remove duplicated lines * format --- .github/workflows/ci.yml | 1 + Project.toml | 3 +- .../stokes2D/shear_band/ShearBand2D.jl | 2 +- .../stokes3D/shear_band/ShearBand3D.jl | 2 +- miniapps/convection/GlobalConvection3D.jl | 14 +- .../Particles2D/Layered_convection2D.jl | 25 +- .../Particles2D/Layered_rheology.jl | 30 +- .../Particles3D/Layered_convection3D.jl | 392 ++++++++++++++++++ .../Particles3D/Layered_rheology.jl | 158 +++++++ src/MetaJustRelax.jl | 6 +- src/Utils.jl | 83 ++-- src/phases/CallArrays.jl | 13 +- src/phases/phases.jl | 42 +- src/rheology/BuoyancyForces.jl | 43 +- src/rheology/StressUpdate.jl | 18 +- src/rheology/Viscosity.jl | 115 ++--- src/stokes/PressureKernels.jl | 14 +- src/stokes/Stokes2D.jl | 46 +- src/stokes/Stokes3D.jl | 62 +-- src/stokes/StressKernels.jl | 32 +- 20 files changed, 823 insertions(+), 278 deletions(-) create mode 100644 miniapps/convection/Particles3D/Layered_convection3D.jl create mode 100644 miniapps/convection/Particles3D/Layered_rheology.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f74e050f..0873578d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,6 +15,7 @@ jobs: matrix: version: - '1.9' + - '1.10' # - 'nightly' os: - ubuntu-latest diff --git a/Project.toml b/Project.toml index d7dbded9..de4addb0 100644 --- a/Project.toml +++ b/Project.toml @@ -20,7 +20,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" [compat] -julia = "1.9" AMDGPU = "0.5.4" CUDA = "4.4.1" CellArrays = "0.1.3" @@ -31,7 +30,9 @@ MPI = "0.20.14" ParallelStencil = "0.9.0" Reexport = "1.2.2" StaticArrays = "1.6.5" +Statistics = "1.9" WriteVTK = "1.18.0" +julia = "1.9" [extras] CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 5cf83aa1..a134b6ae 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -100,7 +100,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Rheology η = @ones(ni...) η_vep = similar(η) # effective visco-elasto-plastic viscosity - compute_viscosity!( + @parallel (@idx ni) compute_viscosity!( η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) ) diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl index dba58dac..a7032d04 100644 --- a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl @@ -213,7 +213,7 @@ n = 64 + 2 nx = ny = nz = n - 2 figdir = "ShearBand3D" igg = if !(JustRelax.MPI.Initialized()) - IGG(init_global_grid(nx, ny, 0; init_MPI = true)...) + IGG(init_global_grid(nx, ny, nz; init_MPI = true)...) else igg end diff --git a/miniapps/convection/GlobalConvection3D.jl b/miniapps/convection/GlobalConvection3D.jl index 4f178d17..da2ac546 100644 --- a/miniapps/convection/GlobalConvection3D.jl +++ b/miniapps/convection/GlobalConvection3D.jl @@ -182,7 +182,9 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th η = @ones(ni...) depth = PTArray([abs(z) for x in xci[1], y in xci[2], z in xci[3]]) args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf) - @parallel (@idx ni) compute_viscosity!(η, 1, 1e-15, args, rheology) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, @strain(stokes)..., args, rheology, (1e18, 1e24) + ) η_vep = deepcopy(η) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -225,7 +227,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th args = (; T=thermal.Tc, P=stokes.P, depth=depth, dt=Inf) # Stokes solver ---------------- - iters = solve!( + solve!( stokes, pt_stokes, di, @@ -299,7 +301,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th return (ni=ni, xci=xci, li=li, di=di), thermal end -function run() +# function run() figdir = "figs3D_test" ar = 3 # aspect ratio n = 32 @@ -307,7 +309,7 @@ function run() ny = nx nz = n - 2 - thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz); -end +# thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz); +# end -run() \ No newline at end of file +# run() \ No newline at end of file diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index c94acc8c..b12f5b0f 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -1,5 +1,5 @@ -# using CUDA -# CUDA.allowscalar(false) # for safety +using CUDA +CUDA.allowscalar(false) # for safety using JustRelax, JustRelax.DataIO, JustPIC import JustRelax.@cell @@ -9,12 +9,11 @@ import JustRelax.@cell # set_backend("CUDA_Float64_2D") # setup ParallelStencil.jl environment -# model = PS_Setup(:CUDA, Float64, 2) -model = PS_Setup(:Threads, Float64, 2) +model = PS_Setup(:CUDA, Float64, 2) environment!(model) # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions, CellArrays # Load file with all the rheology configurations include("Layered_rheology.jl") @@ -90,7 +89,7 @@ end @parallel_indices (i, j) function init_T!(T, z) depth = -z[j] - # (depth) because we have 15km of sticky air + # (depth - 15e3) because we have 15km of sticky air if depth < 0e0 T[i + 1, j] = 273e0 @@ -201,14 +200,14 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false) # Rheology η = @zeros(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!( + @parallel (@idx ni) compute_viscosity!( η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e16, 1e24) ) η_vep = copy(η) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2 + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 ) # Boundary conditions @@ -221,9 +220,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false) # if it does not exist, make folder where figures are stored if save_vtk vtk_dir = figdir*"\\vtk" - !isdir(vtk_dir) && mkpath(vtk_dir) + take(vtk_dir) end - !isdir(figdir) && mkpath(figdir) + take(figdir) # ---------------------------------------------------- # Plot initial T and η profiles @@ -261,7 +260,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false) # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_viscosity!( + @parallel (@idx ni) compute_viscosity!( η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) ) @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) @@ -407,7 +406,7 @@ end figdir = "Plume2D" save_vtk = false # set to true to generate VTK files for ParaView ar = 1 # aspect ratio -n = 128 +n = 64 nx = n*ar - 2 ny = n - 2 igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid @@ -417,4 +416,4 @@ else end # run main script -main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, save_vtk = save_vtk); \ No newline at end of file +# main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, save_vtk = save_vtk); \ No newline at end of file diff --git a/miniapps/convection/Particles2D/Layered_rheology.jl b/miniapps/convection/Particles2D/Layered_rheology.jl index c840b8da..2c935162 100644 --- a/miniapps/convection/Particles2D/Layered_rheology.jl +++ b/miniapps/convection/Particles2D/Layered_rheology.jl @@ -90,18 +90,18 @@ function init_rheologies(; is_plastic = true) Elasticity = el_lithospheric_mantle, ), # Name = "SubLithosphericMantle", - SetMaterialParams(; - Phase = 4, - Density = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), - HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), - Conductivity = K_mantle, - RadioactiveHeat = ConstantRadioactiveHeat(0.0), - CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), - Elasticity = el_sublithospheric_mantle, - ), + # SetMaterialParams(; + # Phase = 4, + # Density = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), + # HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + # Conductivity = K_mantle, + # RadioactiveHeat = ConstantRadioactiveHeat(0.0), + # CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), + # Elasticity = el_sublithospheric_mantle, + # ), # Name = "Plume", SetMaterialParams(; - Phase = 5, + Phase = 4, Density = PT_Density(; ρ0=3.3e3-50, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), Conductivity = K_mantle, @@ -111,7 +111,7 @@ function init_rheologies(; is_plastic = true) ), # Name = "StickyAir", SetMaterialParams(; - Phase = 6, + Phase = 5, Density = ConstantDensity(; ρ=1e3), HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), RadioactiveHeat = ConstantRadioactiveHeat(0.0), @@ -125,7 +125,7 @@ function init_phases!(phases, particles, Lx; d=650e3, r=50e3) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue @@ -143,14 +143,14 @@ function init_phases!(phases, particles, Lx; d=650e3, r=50e3) elseif depth > 90e3 @cell phases[ip, i, j] = 3.0 - elseif 0e0 > depth - @cell phases[ip, i, j] = 6.0 + elseif depth < 0e0 + @cell phases[ip, i, j] = 5.0 end # plume - rectangular if ((x - Lx * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 5.0 + JustRelax.@cell phases[ip, i, j] = 4.0 end end return nothing diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl new file mode 100644 index 00000000..baa4688e --- /dev/null +++ b/miniapps/convection/Particles3D/Layered_convection3D.jl @@ -0,0 +1,392 @@ +using JustRelax, JustRelax.DataIO, JustPIC +import JustRelax.@cell + +## NOTE: need to run one of the lines below if one wishes to switch from one backend to another +# set_backend("Threads_Float64_2D") +# set_backend("CUDA_Float64_2D") + +# setup ParallelStencil.jl environment +model = PS_Setup(:CUDA, Float64, 3) # or (:Threads, 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("Layered_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +@inline init_particle_fields(particles) = @zeros(size(particles.coords[1])...) +@inline init_particle_fields(particles, nfields) = tuple([zeros(particles.coords[1]) for i in 1:nfields]...) +@inline init_particle_fields(particles, ::Val{N}) where N = ntuple(_ -> @zeros(size(particles.coords[1])...) , Val(N)) +@inline init_particle_fields_cellarrays(particles, ::Val{N}) where N = ntuple(_ -> @fill(0.0, size(particles.coords[1])..., celldims=(cellsize(particles.index))), Val(N)) + +function init_particles_cellarrays(nxcell, max_xcell, min_xcell, x, y, z, dx, dy, dz, ni::NTuple{3, Int}) + ncells = prod(ni) + np = max_xcell * ncells + px, py, pz = ntuple(_ -> @fill(NaN, ni..., celldims=(max_xcell,)) , Val(3)) + inject = @fill(false, ni..., eltype=Bool) + index = @fill(false, ni..., celldims=(max_xcell,), eltype=Bool) + + @parallel_indices (i, j, k) function fill_coords_index(px, py, pz, index) + @inline r()= rand(0.05:1e-5:0.95) + I = i, j, k + # lower-left corner of the cell + x0, y0, z0 = x[i], y[j], z[k] + # fill index array + for l in 1:nxcell + JustRelax.@cell px[l, I...] = x0 + dx * r() + JustRelax.@cell py[l, I...] = y0 + dy * r() + JustRelax.@cell pz[l, I...] = z0 + dz * r() + JustRelax.@cell index[l, I...] = true + end + return nothing + end + + @parallel (@idx ni) fill_coords_index(px, py, pz, index) + + return Particles( + (px, py, pz), index, inject, nxcell, max_xcell, min_xcell, np, ni + ) +end + +# Velocity helper grids for the particle advection +function velocity_grids(xci, xvi, di) + xghost = ntuple(Val(3)) do i + LinRange(xci[i][1] - di[i], xci[i][end] + di[i], length(xci[i])+2) + end + grid_vx = xvi[1] , xghost[2], xghost[3] + grid_vy = xghost[1], xvi[2] , xghost[3] + grid_vz = xghost[1], xghost[2], xvi[3] + + return grid_vx, grid_vy, grid_vz +end + +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 + +# Initial thermal profile +@parallel_indices (I...) function init_T!(T, z) + depth = -z[I[3]] + + if depth < 0e0 + T[I...] = 273.0 + + elseif 0e0 ≤ (depth) < 35e3 + dTdZ = (923-273)/35e3 + offset = 273e0 + T[I...] = (depth) * dTdZ + offset + + elseif 110e3 > (depth) ≥ 35e3 + dTdZ = (1492-923)/75e3 + offset = 923 + T[I...] = (depth - 35e3) * dTdZ + offset + + elseif (depth) ≥ 110e3 + dTdZ = (1837 - 1492)/590e3 + offset = 1492e0 + T[I...] = (depth - 110e3) * dTdZ + offset + + end + + return nothing +end + +# Thermal rectangular perturbation +function rectangular_perturbation!(T, xc, yc, zc, r, xvi) + + @parallel_indices (i, j, k) function _rectangular_perturbation!(T, xc, yc, zc, r, x, y, z) + @inbounds if (abs(x[i]-xc) ≤ r) && (abs(y[j] - yc) ≤ r) && (abs(z[k] - zc) ≤ r) + depth = abs(z[k]) + dTdZ = (2047 - 2017) / 50e3 + offset = 2017 + T[i, j, k] = (depth - 585e3) * dTdZ + offset + end + return nothing + end + + @parallel _rectangular_perturbation!(T, xc, yc, zc, r, xvi...) +end +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) + + # Physical domain ------------------------------------ + lz = 700e3 # domain length in z + lx = ly = lz * ar # domain length in x and y + ni = nx, ny, nz # number of cells + li = lx, ly, lz # domain length + di = @. li / ni # grid steps + origin = 0.0, 0.0, -lz # origin coordinates (15km of sticky air layer) + xci, xvi = lazy_grid( + di, + li, + ni; + origin = origin + ) # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies(; is_plastic = true) + κ = (10 / (rheology[1].HeatCapacity[1].cp * rheology[1].Density[1].ρ0)) + dt = dt_diff = 0.5 * min(di...)^3 / κ / 3.01 # diffusive CFL timestep limiter + # ---------------------------------------------------- + + # Initialize particles ------------------------------- + nxcell, max_xcell, min_xcell = 20, 20, 1 + particles = init_particles_cellarrays( + nxcell, max_xcell, min_xcell, xvi..., di..., ni + ) + # velocity grids + grid_vx, grid_vy, grid_vz = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_particle_fields_cellarrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + xc_anomaly = lx/2 # origin of thermal anomaly + yc_anomaly = ly/2 # origin of thermal anomaly + zc_anomaly = -610e3 # origin of thermal anomaly + r_anomaly = 25e3 # radius of perturbation + init_phases!(pPhases, particles, lx, ly; d=abs(zc_anomaly), r=r_anomaly) + phase_ratios = PhaseRatio(ni, length(rheology)) + @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(ni, ViscoElastic) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.5 / √3.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), + periodicity = (left = false, right = false, top = false, bot = false, front = false, back = false), + ) + # initialize thermal profile - Half space cooling + @parallel init_T!(thermal.T, xvi[3]) + thermal_bcs!(thermal.T, thermal_bc) + + rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, xvi) + @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = ntuple(_ -> @zeros(ni...), Val(3)) + for _ in 1:1 + @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[3], xci[3]) + end + # Rheology + η = @ones(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, phase_ratios.center, @strain(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 = 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), + ) + + # 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(resolution = (1200, 900)) + ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") + ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") + lines!(ax1, Array(thermal.T[:]), 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.coords) + + 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 (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + @parallel (@idx ni) compute_viscosity!( + η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) + ) + @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + + # Stokes solver ---------------- + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios, + rheology, + args, + Inf, + igg; + iterMax = 100e3, + nout = 1e3, + viscosity_cutoff = (1e18, 1e24) + ); + @parallel (JustRelax.@idx ni) tensor_invariant!(stokes.ε.II, @strain(stokes)...) + dt = compute_dt(stokes, di, dt_diff) / 2 + # ------------------------------ + + # interpolate fields from particle to grid vertices + particle2grid!(thermal.T, pT, xvi, particles.coords) + temperature2center!(thermal) + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3) + # advect particles in memory + shuffle_particles!(particles, xvi, particle_args) + # interpolate fields from grid vertices to particles + grid2particle_flip!(pT, xvi, thermal.T, thermal.Told, particles.coords) + # check if we need to inject particles + inject = check_injection(particles) + inject && inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), 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, 1) == 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), + ) + data_c = (; + Tc = Array(thermal.Tc), + P = Array(stokes.P), + τxx = Array(stokes.τ.xx), + τyy = Array(stokes.τ.yy), + εxx = Array(stokes.ε.xx), + εyy = Array(stokes.ε.yy), + η = Array(log10.(η)), + ) + save_vtk( + joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), + xvi, + xci, + data_v, + data_c + ) + end + + xz_slice = ny >>> 1 + # Make Makie figure + fig = Figure(resolution = (1400, 1800), 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 = "τII [MPa]") + 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[:, xz_slice, :]) , colormap=:batlow) + # Plot particles phase + h2 = heatmap!(ax2, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II[:, xz_slice, :]./1e6) , colormap=:batlow) + # Plot 2nd invariant of strain rate + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II[:, xz_slice, :])) , colormap=:batlow) + # Plot effective viscosity + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_vep[:, xz_slice, :])) , 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 +ar = 1 # aspect ratio +n = 128 +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 + +# (Path)/folder where output data and figures are stored +figdir = "Plume3D_$n" +main3D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/miniapps/convection/Particles3D/Layered_rheology.jl b/miniapps/convection/Particles3D/Layered_rheology.jl new file mode 100644 index 00000000..96f49cf2 --- /dev/null +++ b/miniapps/convection/Particles3D/Layered_rheology.jl @@ -0,0 +1,158 @@ +# from "Fingerprinting secondary mantle plumes", Cloetingh et al. 2022 + +function init_rheologies(; is_plastic = true) + + # Dislocation and Diffusion creep + disl_upper_crust = DislocationCreep(A=5.07e-18, n=2.3, E=154e3, V=6e-6, r=0.0, R=8.3145) + disl_lower_crust = DislocationCreep(A=2.08e-23, n=3.2, E=238e3, V=6e-6, r=0.0, R=8.3145) + disl_lithospheric_mantle = DislocationCreep(A=2.51e-17, n=3.5, E=530e3, V=6e-6, r=0.0, R=8.3145) + disl_sublithospheric_mantle = DislocationCreep(A=2.51e-17, n=3.5, E=530e3, V=6e-6, r=0.0, R=8.3145) + diff_lithospheric_mantle = DislocationCreep(A=2.51e-17, n=1.0, E=530e3, V=6e-6, r=0.0, R=8.3145) + diff_sublithospheric_mantle = DislocationCreep(A=2.51e-17, n=1.0, E=530e3, V=6e-6, r=0.0, R=8.3145) + + # Elasticity + el_upper_crust = SetConstantElasticity(; G=25e9, ν=0.5) + el_lower_crust = SetConstantElasticity(; G=25e9, ν=0.5) + el_lithospheric_mantle = SetConstantElasticity(; G=67e9, ν=0.5) + el_sublithospheric_mantle = SetConstantElasticity(; G=67e9, ν=0.5) + β_upper_crust = inv(get_Kb(el_upper_crust)) + β_lower_crust = inv(get_Kb(el_lower_crust)) + β_lithospheric_mantle = inv(get_Kb(el_lithospheric_mantle)) + β_sublithospheric_mantle = inv(get_Kb(el_sublithospheric_mantle)) + + # Physical properties using GeoParams ---------------- + η_reg = 1e16 + cohesion = 3e6 + friction = asind(0.2) + pl_crust = if is_plastic + DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + friction = asind(0.3) + pl = if is_plastic + DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + pl_wz = if is_plastic + DruckerPrager_regularised(; C = 2e6, ϕ=2.0, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + else + DruckerPrager_regularised(; C = Inf, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity + end + + # crust + K_crust = TP_Conductivity(; + a = 0.64, + b = 807e0, + c = 0.77, + d = 0.00004*1e-6, + ) + + K_mantle = TP_Conductivity(; + a = 0.73, + b = 1293e0, + c = 0.77, + d = 0.00004*1e-6, + ) + + # Define rheolgy struct + rheology = ( + # Name = "UpperCrust", + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0=2.75e3, β=β_upper_crust, T0=0.0, α = 3.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=7.5e2), + Conductivity = K_crust, + CompositeRheology = CompositeRheology((disl_upper_crust, el_upper_crust, pl_crust)), + Elasticity = el_upper_crust, + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "LowerCrust", + SetMaterialParams(; + Phase = 2, + Density = PT_Density(; ρ0=3e3, β=β_lower_crust, T0=0.0, α = 3.5e-5), + HeatCapacity = ConstantHeatCapacity(; cp=7.5e2), + Conductivity = K_crust, + CompositeRheology = CompositeRheology((disl_lower_crust, el_lower_crust, pl_crust)), + Elasticity = el_lower_crust, + ), + # Name = "LithosphericMantle", + SetMaterialParams(; + Phase = 3, + Density = PT_Density(; ρ0=3.3e3, β=β_lithospheric_mantle, T0=0.0, α = 3e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + Conductivity = K_mantle, + CompositeRheology = CompositeRheology((disl_lithospheric_mantle, diff_lithospheric_mantle, el_lithospheric_mantle, pl)), + Elasticity = el_lithospheric_mantle, + ), + # # Name = "SubLithosphericMantle", + # SetMaterialParams(; + # Phase = 4, + # Density = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), + # HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + # Conductivity = K_mantle, + # RadioactiveHeat = ConstantRadioactiveHeat(0.0), + # CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), + # Elasticity = el_sublithospheric_mantle, + # ), + # Name = "Plume", + SetMaterialParams(; + Phase = 4, + Density = PT_Density(; ρ0=3.3e3-50, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + Conductivity = K_mantle, + CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)), + Elasticity = el_sublithospheric_mantle, + ), + # Name = "StickyAir", + SetMaterialParams(; + Phase = 5, + Density = ConstantDensity(; ρ=1e3), + HeatCapacity = ConstantHeatCapacity(; cp=1.25e3), + Conductivity = ConstantConductivity(; k=15.0), + CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)), + ), + ) +end + +function init_phases!(phases, particles, Lx, Ly; d=650e3, r=50e3) + ni = size(phases) + + @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, r, Lx, Ly) + + @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...] + depth = -(JustRelax.@cell pz[ip, I...]) + + if 0e0 ≤ depth ≤ 21e3 + @cell phases[ip, I...] = 1.0 + + elseif 35e3 ≥ depth > 21e3 + @cell phases[ip, I...] = 2.0 + + elseif 90e3 ≥ depth > 35e3 + @cell phases[ip, I...] = 3.0 + + elseif depth > 90e3 + @cell phases[ip, I...] = 3.0 + + elseif 0e0 > depth + @cell phases[ip, I...] = 5.0 + + end + + # plume - rectangular + if ((x - Lx * 0.5)^2 ≤ r^2) && ((y - Ly * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2) + JustRelax.@cell phases[ip, I...] = 4.0 + end + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, Ly) +end diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl index ca75af8f..b0d0472e 100644 --- a/src/MetaJustRelax.jl +++ b/src/MetaJustRelax.jl @@ -89,7 +89,8 @@ function environment!(model::PS_Setup{T,N}) where {T,N} norm_mpi, minimum_mpi, maximum_mpi, - multi_copy! + multi_copy!, + take include(joinpath(@__DIR__, "boundaryconditions/BoundaryConditions.jl")) export pureshear_bc!, @@ -112,6 +113,9 @@ function environment!(model::PS_Setup{T,N}) where {T,N} include(joinpath(@__DIR__, "rheology/Viscosity.jl")) export compute_viscosity! + include(joinpath(@__DIR__, "stokes/StressKernels.jl")) + export tensor_invariant! + include(joinpath(@__DIR__, "stokes/Stokes2D.jl")) export solve! diff --git a/src/Utils.jl b/src/Utils.jl index 1a49cbc1..b3073132 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -1,4 +1,17 @@ # MACROS +""" + @idx(args...) + +Make a linear range from `1` to `args[i]`, with `i ∈ [1, ..., n]` +""" +macro idx(args...) + return quote + _idx($(esc.(args)...)) + end +end + +@inline _idx(args::NTuple{N,Int}) where {N} = ntuple(i -> 1:args[i], Val(N)) +@inline _idx(args::Vararg{Int,N}) where {N} = ntuple(i -> 1:args[i], Val(N)) """ copy(B, A) @@ -13,9 +26,14 @@ end multi_copyto!(B::AbstractArray, A::AbstractArray) = copyto!(B, A) -function multi_copyto!(B::NTuple{N,AbstractArray}, A::NTuple{N,AbstractArray}) where {N} - for (Bi, Ai) in zip(B, A) - copyto!(Bi, Ai) +function detect_arsg_size(A::NTuple{N,AbstractArray{T,Dims}}) where {N,T,Dims} + ntuple(Val(Dims)) do i + Base.@_inline_meta + s = ntuple(Val(N)) do j + Base.@_inline_meta + size(A[j], i) + end + maximum(s) end end @@ -31,6 +49,19 @@ end return nothing end +Base.@propagate_inbounds @generated function unrolled_copy!( + dst::NTuple{N,T}, src::NTuple{N,T}, I::Vararg{Int,NI} +) where {N,NI,T} + quote + Base.@_inline_meta + Base.@nexprs $N n -> begin + if all(tuple(I...) .≤ size(dst[n])) + dst[n][I...] = src[n][I...] + end + end + end +end + """ @add(I, args...) @@ -257,20 +288,6 @@ macro allocate(ni...) return esc(:(PTArray(undef, $(ni...)))) end -""" - @idx(args...) - -Make a linear range from `1` to `args[i]`, with `i ∈ [1, ..., n]` -""" -macro idx(args...) - return quote - _idx(tuple($(esc.(args)...))...) - end -end - -@inline Base.@pure _idx(args::Vararg{Int,N}) where {N} = ntuple(i -> 1:args[i], Val(N)) -@inline Base.@pure _idx(args::NTuple{N,Int}) where {N} = ntuple(i -> 1:args[i], Val(N)) - function indices(::NTuple{3,T}) where {T} i = (blockIdx().x - 1) * blockDim().x + threadIdx().x j = (blockIdx().y - 1) * blockDim().y + threadIdx().y @@ -306,10 +323,11 @@ end I_range = (I - width_x):(I + width_x) J_range = (J - width_y):(J + width_y) x = -Inf - for i in I_range - ii = clamp(i, 1, nx) - for j in J_range - jj = clamp(j, 1, ny) + + for j in J_range + jj = clamp(j, 1, ny) # handle boundary cells + for i in I_range + ii = clamp(i, 1, nx) # handle boundary cells Aij = A[ii, jj] if Aij > x x = Aij @@ -325,12 +343,13 @@ end J_range = (J - width_y):(J + width_y) K_range = (K - width_z):(K + width_z) x = -Inf - for i in I_range - ii = clamp(i, 1, nx) + + for k in K_range + kk = clamp(k, 1, nz) # handle boundary cells for j in J_range - jj = clamp(j, 1, ny) - for k in K_range - kk = clamp(k, 1, nz) + jj = clamp(j, 1, ny) # handle boundary cells + for i in I_range + ii = clamp(i, 1, nx) # handle boundary cells Aijk = A[ii, jj, kk] if Aijk > x x = Aijk @@ -405,13 +424,23 @@ end @inline tupleize(v) = (v,) @inline tupleize(v::Tuple) = v +# Delta function +@inline allzero(x::Vararg{T,N}) where {T,N} = all(x -> x == 0, x) + +""" + take(fldr::String) + +Create folder `fldr` if it does not exist. +""" +take(fldr::String) = !isdir(fldr) && mkpath(fldr) + """ continuation_log(x_new, x_old, ν) Do a continuation step `exp((1-ν)*log(x_old) + ν*log(x_new))` with damping parameter `ν` """ # @inline continuation_log(x_new, x_old, ν) = exp((1 - ν) * log(x_old) + ν * log(x_new)) -@inline continuation_log(x_new, x_old, ν) = (1 - ν) * x_old + ν * x_new +@inline continuation_log(x_new, x_old, ν) = muladd((1 - ν), x_old, ν * x_new) # (1 - ν) * x_old + ν * x_new # Others diff --git a/src/phases/CallArrays.jl b/src/phases/CallArrays.jl index 8108b0f7..03ed22c0 100644 --- a/src/phases/CallArrays.jl +++ b/src/phases/CallArrays.jl @@ -9,14 +9,14 @@ Base.@propagate_inbounds @inline function Base.getindex( x::CellArray{SVector{Nv,T},N1,N2,T}, I::Vararg{Int,N} ) where {Nv,N,N1,N2,T} idx_cell = cart2ind(x.dims, I...) - return SVector{Nv,T}(x.data[idx_cell, i, 1] for i in 1:Nv) + return SVector{Nv,T}(@inbounds x.data[idx_cell, i, 1] for i in 1:Nv) end Base.@propagate_inbounds @inline function Base.getindex( x::CPUCellArray{SVector{Nv,T},N1,N2,T}, I::Vararg{Int,N} ) where {Nv,N,N1,N2,T} idx_cell = cart2ind(x.dims, I...) - return SVector{Nv,T}(x.data[1, i, idx_cell] for i in 1:Nv) + return SVector{Nv,T}(@inbounds x.data[1, i, idx_cell] for i in 1:Nv) end Base.@propagate_inbounds @inline function setindex!( @@ -45,6 +45,7 @@ Base.@propagate_inbounds @inline function element( ) where {T_elem,N,Nc,D} return viewelement(A, i, icell...) end + Base.@propagate_inbounds @inline function element( A::CellArray, i::T, j::T, icell::Vararg{Int,Nc} ) where {Nc,T<:Int} @@ -97,6 +98,7 @@ end Base.@propagate_inbounds @inline function _setelement!(A::Array, x, idx::Int, icell::Int) return (A[1, idx, icell] = x) end + Base.@propagate_inbounds @inline function _setelement!(A, x, idx::Int, icell::Int) return (A[icell, idx, 1] = x) end @@ -107,7 +109,6 @@ end cart2ind(A) Return the linear index of a `n`-dimensional array corresponding to the cartesian indices `I` - """ @inline function cart2ind(n::NTuple{N1,Int}, I::Vararg{Int,N2}) where {N1,N2} return LinearIndices(n)[CartesianIndex(I...)] @@ -120,8 +121,9 @@ end ## Fallbacks import Base: getindex, setindex! -@inline element(A::Union{Array,CuArray,ROCArray}, I::Vararg{Int,N}) where {N} = - getindex(A, I...) +@inline function element(A::Union{Array,CuArray,ROCArray}, I::Vararg{Int,N}) where {N} + return getindex(A, I...) +end @inline function setelement!( A::Union{Array,CuArray,ROCArray}, x::Number, I::Vararg{Int,N} ) where {N} @@ -140,6 +142,7 @@ macro cell(ex) end @inline _get(ex) = Expr(:call, element, ex.args...) + @inline function _set(ex) return Expr( :call, setelement!, ex.args[1].args[1], ex.args[2], ex.args[1].args[2:end]... diff --git a/src/phases/phases.jl b/src/phases/phases.jl index b1583597..9a7890a5 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -109,7 +109,6 @@ end quote Base.@_inline_meta x = 0.0 - # Base.@nexprs $N i -> x += iszero(ratio[i]) ? 0.0 : fn(rheology[i], args) * ratio[i] Base.@nexprs $N i -> x += begin r = ratio[i] isone(r) && return fn(rheology[i], args) * r @@ -120,31 +119,34 @@ end end # ParallelStencil launch kernel for 2D -@parallel_indices (i, j) function phase_ratios_center( - ratio_centers, px, py, xc, yc, di, phases -) - I = i, j +@parallel_indices (I...) function phase_ratios_center( + ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases +) where {N,T1,T2,T3} + # index corresponding to the cell center + cell_center = ntuple(i -> xci[i][I[i]], Val(N)) + # phase ratios weights (∑w = 1.0) w = phase_ratio_weights( - px[I...], py[I...], phases[I...], xc[i], yc[j], di, JustRelax.nphases(ratio_centers) + getindex.(pxi, I...), + phases[I...], + cell_center, + di, + JustRelax.nphases(ratio_centers), ) - + # update phase ratios array for k in 1:numphases(ratio_centers) JustRelax.@cell ratio_centers[k, I...] = w[k] end + return nothing end -@inline δ(i, j) = i == j - function phase_ratio_weights( - pxi::SVector{N1,T}, pyi::SVector{N1,T}, ph::SVector{N1,T}, xc, yc, di, ::Val{NC} -) where {N1,NC,T} + pxi::NTuple{NP,C}, ph::SVector{N1,T}, cell_center, di, ::Val{NC} +) where {N1,NC,NP,T,C} if @generated quote Base.@_inline_meta - # wrap cell center coordinates into a tuple - cell_center = xc, yc # Initiaze phase ratio weights (note: can't use ntuple() here because of the @generated function) Base.@nexprs $NC i -> w_i = zero($T) w = Base.@ncall $NC tuple w @@ -153,11 +155,12 @@ function phase_ratio_weights( sumw = zero($T) Base.@nexprs $N1 i -> begin # bilinear weight (1-(xᵢ-xc)/dx)*(1-(yᵢ-yc)/dy) - x = bilinear_weight(cell_center, (pxi[i], pyi[i]), di) + x = bilinear_weight(cell_center, getindex.(pxi, i), di) sumw += x # reduce ph_local = ph[i] # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta - Base.@nexprs $NC j -> tmp_j = w[j] + x * δ(ph_local, j) + # Base.@nexprs $NC j -> tmp_j = w[j] + x * δ(Int(ph_local), j) + Base.@nexprs $NC j -> tmp_j = w[j] + x * (ph_local == j) w = Base.@ncall $NC tuple tmp end @@ -168,8 +171,6 @@ function phase_ratio_weights( return w end else - # wrap cell center coordinates into a tuple - cell_center = xc, yc # Initiaze phase ratio weights (note: can't use ntuple() here because of the @generated function) w = ntuple(_ -> zero(T), Val(NC)) # initialie sum of weights @@ -177,11 +178,12 @@ function phase_ratio_weights( for i in eachindex(pxi) # bilinear weight (1-(xᵢ-xc)/dx)*(1-(yᵢ-yc)/dy) - x = @inline bilinear_weight(cell_center, (pxi[i], pyi[i]), di) + x = @inline bilinear_weight(cell_center, getindex.(pxi, i), di) sumw += x # reduce ph_local = ph[i] # this is doing sum(w * δij(i, phase)), where δij is the Kronecker delta - w = w .+ x .* ntuple(j -> δ(ph_local, j), Val(NC)) + # w = w .+ x .* ntuple(j -> δ(Int(ph_local), j), Val(NC)) + w = w .+ x .* ntuple(j -> (ph_local == j), Val(NC)) end w = w .* inv(sum(w)) return w @@ -195,7 +197,7 @@ end Base.@_inline_meta val = one($T) Base.Cartesian.@nexprs $N i -> - @inbounds val *= one($T) - abs(a[i] - b[i]) * inv(di[i]) + @inbounds val *= muladd(-abs(a[i] - b[i]), inv(di[i]), one($T)) return val end end diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 73dea296..2bf8784f 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -3,56 +3,25 @@ Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. """ -@parallel_indices (i, j) function compute_ρg!( - ρg::_T, rheology, args -) where {_T<:AbstractArray{M,2} where {M<:Real}} - _compute_ρg!(ρg, rheology, args, i, j) - return nothing -end -@parallel_indices (i, j, k) function compute_ρg!( - ρg::_T, rheology, args -) where {_T<:AbstractArray{M,3} where {M<:Real}} - _compute_ρg!(ρg, rheology, args, i, j, k) +@parallel_indices (I...) function compute_ρg!(ρg, rheology, args) # index arguments for the current cell cell center + args_ijk = ntuple_idx(args, I...) + ρg[I...] = JustRelax.compute_buoyancy(rheology, args_ijk) return nothing end -function _compute_ρg!(ρg, rheology, args, I::Vararg{Int,N}) where {N} - # index arguments for the current cell cell center - T = args.T[I...] - 273.0 - P = args.P[I...] - args_ijk = (; T, P) - return ρg[I...] = compute_buoyancy(rheology, args_ijk) -end - """ compute_ρg!(ρg, phase_ratios, rheology, args) Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. The `phase_ratios` are used to compute the density of the composite rheology. """ -@parallel_indices (i, j) function compute_ρg!( - ρg::_T, phase_ratios, rheology, args -) where {_T<:AbstractArray{M,2} where {M<:Real}} - _compute_ρg!(ρg, phase_ratios, rheology, args, i, j) +@parallel_indices (I...) function compute_ρg!(ρg, phase_ratios, rheology, args) # index arguments for the current cell cell center + args_ijk = ntuple_idx(args, I...) + ρg[I...] = JustRelax.compute_buoyancy(rheology, args_ijk, phase_ratios[I...]) return nothing end -@parallel_indices (i, j, k) function compute_ρg!( - ρg::_T, phase_ratios, rheology, args -) where {_T<:AbstractArray{M,3} where {M<:Real}} - _compute_ρg!(ρg, phase_ratios, rheology, args, i, j, k) - return nothing -end - -function _compute_ρg!(ρg, phase_ratios, rheology, args, I::Vararg{Int,N}) where {N} - # index arguments for the current cell cell center - T = args.T[I...] - 273.0 - P = args.P[I...] - args_ijk = (; T, P) - return ρg[I...] = compute_buoyancy(rheology, args_ijk, phase_ratios[I...]) -end - ## Inner buoyancy force kernels @inline function compute_buoyancy(rheology::MaterialParams, args) diff --git a/src/rheology/StressUpdate.jl b/src/rheology/StressUpdate.jl index 4b74dc9e..373c8452 100644 --- a/src/rheology/StressUpdate.jl +++ b/src/rheology/StressUpdate.jl @@ -26,8 +26,7 @@ function _compute_τ_nonlinear!( end # get plastic paremeters (if any...) (; is_pl, C, sinϕ, cosϕ, η_reg, volume) = plastic_parameters - τy = C + P[idx...] * sinϕ - # τy = C * cosϕ + P[idx...] * sinϕ + τy = C * cosϕ + P[idx...] * sinϕ dτij = if isyielding(is_pl, τII_trial, τy) # derivatives plastic stress correction @@ -35,6 +34,7 @@ function _compute_τ_nonlinear!( τij, dτij, τy, τII_trial, ηij, λ[idx...], η_reg, dτ_r, volume ) dτ_pl + else dτij end @@ -50,14 +50,16 @@ end # check if plasticity is active @inline isyielding(is_pl, τII_trial, τy) = is_pl && τII_trial > τy -@inline compute_dτ_r(θ_dτ, ηij, _Gdt) = inv(θ_dτ + ηij * _Gdt + 1.0) +@inline compute_dτ_r(θ_dτ, ηij, _Gdt) = inv(θ_dτ + muladd(ηij, _Gdt, 1.0)) 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 - @inbounds dτ_r * (-(τij[i] - τij_o[i]) * ηij * _Gdt - τij[i] + 2.0 * ηij * εij[i]) + @inbounds dτ_r * muladd( + 2.0 * ηij, εij[i], muladd(-((τij[i] - τij_o[i])) * ηij, _Gdt, -τij[i]) + ) end return dτij, second_invariant((τij .+ dτij)...) end @@ -68,17 +70,15 @@ function compute_dτ_pl( # yield function F = τII_trial - τy # Plastic multiplier - ν = 0.05 - λ = (1 - ν) * λ0 + ν * (F > 0.0) * F * inv(dτ_r * ηij + η_reg + volume) + ν = 0.5 + λ = ν * λ0 + (1 - ν) * (F > 0.0) * F * inv(ηij * dτ_r + η_reg + volume) λ_τII = λ * 0.5 * inv(τII_trial) dτ_pl = ntuple(Val(N)) do i Base.@_inline_meta # derivatives of the plastic potential λdQdτ = (τij[i] + dτij[i]) * λ_τII - # corrected stress: dτ_r * (-(τij[i] - τij_o[i]) * ηij * _Gdt - τij[i] + 2.0 * ηij * (εij[i] - λdQdτ)) - # NOTE: dτij[i] already contains the non-plastic stress increment - # dτij[i] - dτ_r * 2.0 * ηij * λdQdτ + # corrected stress muladd(-dτ_r * 2.0, ηij * λdQdτ, dτij[i]) end return dτ_pl, λ diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 55f390da..6fb0799c 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,80 +1,74 @@ ## 2D KERNELS -@parallel_indices (i, j) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity!( η, ν, εxx, εyy, εxyv, args, rheology, cutoff ) # convinience closure - @inline gather(A) = _gather(A, i, j) + @inline gather(A) = _gather(A, I...) @inbounds begin # we need strain rate not to be zero, otherwise we get NaNs - εII_0 = (εxx[i, j] == εyy[i, j] == 0) * 1e-15 + εII_0 = allzero(εxx[I...], εyy[I...]) * 1e-15 # argument fields at local index - args_ij = local_viscosity_args(args, i, j) + args_ij = local_viscosity_args(args, I...) # compute second invariant of strain rate tensor - εij = εII_0 + εxx[i, j], -εII_0 + εyy[i, j], gather(εxyv) + εij = εII_0 + εxx[I...], -εII_0 + εyy[I...], gather(εxyv) εII = second_invariant(εij...) # compute and update stress viscosity ηi = compute_viscosity_εII(rheology, εII, args_ij) - ηi = continuation_log(ηi, η[i, j], ν) - η[i, j] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing end -@parallel_indices (i, j) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) +@parallel_indices (I...) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) @inbounds begin # argument fields at local index - args_ij = local_viscosity_args(args, i, j) + args_ij = local_viscosity_args(args, I...) # compute second invariant of strain rate tensor - εII_ij = εII[i, j] + εII_ij = εII[I...] # compute and update stress viscosity ηi = compute_viscosity_εII(rheology, εII_ij, args_ij) - ηi = continuation_log(ηi, η[i, j], ν) - η[i, j] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing end -function compute_viscosity!(η, x...) - ni = size(η) - @parallel (JustRelax.@idx ni) _compute_viscosity!(η, x...) - return nothing -end - -@parallel_indices (i, j) function _compute_viscosity!( +@parallel_indices (I...) function compute_viscosity!( η, ν, ratios_center, εxx, εyy, εxyv, args, rheology, cutoff ) # convinience closure - @inline gather(A) = _gather(A, i, j) + @inline gather(A) = _gather(A, I...) @inbounds begin # we need strain rate not to be zero, otherwise we get NaNs - εII_0 = (εxx[i, j] == εyy[i, j] == 0) * 1e-15 + εII_0 = allzero(εxx[I...], εyy[I...]) * 1e-18 # argument fields at local index - args_ij = local_viscosity_args(args, i, j) + args_ij = local_viscosity_args(args, I...) # local phase ratio - ratio_ij = ratios_center[i, j] + ratio_ij = ratios_center[I...] # compute second invariant of strain rate tensor - εij = εII_0 + εxx[i, j], -εII_0 + εyy[i, j], gather(εxyv) + εij = εII_0 + εxx[I...], -εII_0 + εyy[I...], gather(εxyv) εII = second_invariant(εij...) # compute and update stress viscosity ηi = compute_phase_viscosity_εII(rheology, ratio_ij, εII, args_ij) - ηi = continuation_log(ηi, η[i, j], ν) - η[i, j] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing @@ -82,24 +76,24 @@ end ## 3D KERNELS -@parallel_indices (i, j, k) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity!( η, ν, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, cutoff ) # convinience closures - @inline gather_yz(A) = _gather_yz(A, i, j, k) - @inline gather_xz(A) = _gather_xz(A, i, j, k) - @inline gather_xy(A) = _gather_xy(A, i, j, k) + @inline gather_yz(A) = _gather_yz(A, I...) + @inline gather_xz(A) = _gather_xz(A, I...) + @inline gather_xy(A) = _gather_xy(A, I...) @inbounds begin # we need strain rate not to be zero, otherwise we get NaNs - εII_0 = (εxx[i, j, k] == εyy[i, j, k] == εzz[i, j, k] == 0) * 1e-18 + εII_0 = allzero(εxx[I...], εyy[I...], εzz[I...]) * 1e-18 # # argument fields at local index - args_ijk = local_viscosity_args(args, i, j, k) + args_ijk = local_viscosity_args(args, I...) # compute second invariant of strain rate tensor - εij_normal = εxx[i, j, k], εyy[i, j, k], εzz[i, j, k] + εij_normal = εxx[I...], εyy[I...], εzz[I...] εij_normal = εij_normal .+ (εII_0, -εII_0 * 0.5, -εII_0 * 0.5) εij_shear = gather_yz(εyzv), gather_xz(εxzv), gather_xy(εxyv) εij = (εij_normal..., εij_shear...) @@ -107,51 +101,51 @@ end # update stress and effective viscosity ηi = compute_viscosity_εII(rheology, εII, args_ijk) - ηi = continuation_log(ηi, η[i, j, k], ν) - η[i, j, k] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing end -@parallel_indices (i, j, k) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) +@parallel_indices (I...) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) @inbounds begin # # argument fields at local index - args_ijk = local_viscosity_args(args, i, j, k) + args_ijk = local_viscosity_args(args, I...) # compute second invariant of strain rate tensor - εII_ij = εII[i, j, k] + εII_ij = εII[I...] # update stress and effective viscosity ηi = compute_viscosity_εII(rheology, εII_ij, args_ijk) - ηi = continuation_log(ηi, η[i, j, k], ν) - η[i, j, k] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing end -@parallel_indices (i, j, k) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity!( η, ν, ratios_center, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, cutoff ) # convinience closures - @inline gather_yz(A) = _gather_yz(A, i, j, k) - @inline gather_xz(A) = _gather_xz(A, i, j, k) - @inline gather_xy(A) = _gather_xy(A, i, j, k) + @inline gather_yz(A) = _gather_yz(A, I...) + @inline gather_xz(A) = _gather_xz(A, I...) + @inline gather_xy(A) = _gather_xy(A, I...) @inbounds begin # we need strain rate not to be zero, otherwise we get NaNs - εII_0 = (εxx[i, j, k] == εyy[i, j, k] == εzz[i, j, k] == 0) * 1e-18 + εII_0 = allzero(εxx[I...], εyy[I...], εzz[I...]) * 1e-18 # # argument fields at local index - args_ijk = local_viscosity_args(args, i, j, k) + args_ijk = local_viscosity_args(args, I...) # local phase ratio - ratio_ijk = ratios_center[i, j, k] + ratio_ijk = ratios_center[I...] # compute second invariant of strain rate tensor - εij_normal = εxx[i, j, k], εyy[i, j, k], εzz[i, j, k] + εij_normal = εxx[I...], εyy[I...], εzz[I...] εij_normal = εij_normal .+ (εII_0, -εII_0 * 0.5, -εII_0 * 0.5) εij_shear = gather_yz(εyzv), gather_xz(εxzv), gather_xy(εxyv) εij = (εij_normal..., εij_shear...) @@ -159,8 +153,8 @@ end # update stress and effective viscosity ηi = compute_phase_viscosity_εII(rheology, ratio_ijk, εII, args_ijk) - ηi = continuation_log(ηi, η[i, j, k], ν) - η[i, j, k] = clamp(ηi, cutoff...) + ηi = continuation_log(ηi, η[I...], ν) + η[I...] = clamp(ηi, cutoff...) end return nothing @@ -180,6 +174,23 @@ end return local_args end +# @generated function compute_phase_viscosity_εII( +# rheology::NTuple{N,AbstractMaterialParamsStruct}, ratio, εII, args +# ) where {N} +# quote +# Base.@_inline_meta +# η = 0.0 +# Base.@nexprs $N i -> ( +# η += if iszero(ratio[i]) +# 0.0 +# else +# inv(compute_viscosity_εII(rheology[i].CompositeRheology[1], εII, args)) * ratio[i] +# end +# ) +# inv(η) +# end +# end + @generated function compute_phase_viscosity_εII( rheology::NTuple{N,AbstractMaterialParamsStruct}, ratio, εII, args ) where {N} @@ -190,9 +201,9 @@ end η += if iszero(ratio[i]) 0.0 else - inv(compute_viscosity_εII(rheology[i].CompositeRheology[1], εII, args)) * ratio[i] + compute_viscosity_εII(rheology[i].CompositeRheology[1], εII, args) * ratio[i] end ) - inv(η) + η end end diff --git a/src/stokes/PressureKernels.jl b/src/stokes/PressureKernels.jl index 5a0a6b68..2e86264d 100644 --- a/src/stokes/PressureKernels.jl +++ b/src/stokes/PressureKernels.jl @@ -40,21 +40,9 @@ function _compute_P!(P, ∇V, η, r, θ_dτ) return RP, P end -# function _compute_P!(P, P0, ∇V, η, K, dt, r, θ_dτ) -# _Kdt = inv(K * dt) -# RP = muladd(-(P - P0), _Kdt, -∇V) -# P += RP / (1.0 / (r / θ_dτ * η) + _Kdt) -# return RP, P -# end - function _compute_P!(P, P0, ∇V, η, K, dt, r, θ_dτ) _Kdt = inv(K * dt) - c = r / θ_dτ * η RP = muladd(-(P - P0), _Kdt, -∇V) - P += RP * c - - # P = (-∇V + P0 * _Kdt + P * inv(c)) * inv(inv(c) + _Kdt) - # RP = muladd(-(P - P0), _Kdt, -∇V) - + P += RP / (1.0 / (r / θ_dτ * η) + 1.0 * _Kdt) return RP, P end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 670a1bac..8ebf031d 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -45,6 +45,7 @@ using ParallelStencil.FiniteDifferences2D using GeoParams, LinearAlgebra, Printf import JustRelax: elastic_iter_params!, PTArray, Velocity, SymmetricTensor +import JustRelax: tensor_invariant!, compute_τ_nonlinear! import JustRelax: Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG import JustRelax: compute_maxloc!, solve! @@ -58,8 +59,7 @@ include("StressKernels.jl") export solve!, rotate_stress_particles_jaumann!, rotate_stress_particles_roation_matrix!, - compute_vorticity!, - tensor_invariant! + compute_vorticity! function update_τ_o!(stokes::StokesArrays{ViscoElastic,A,B,C,D,2}) where {A,B,C,D} τxx, τyy, τxy, τxy_c = stokes.τ.xx, stokes.τ.yy, stokes.τ.xy, stokes.τ.xy_c @@ -376,10 +376,10 @@ function JustRelax.solve!( update_halo!(ητ) @parallel (@idx ni) compute_τ_nonlinear!( - @tensor_center(stokes.τ)..., + @tensor_center(stokes.τ), stokes.τ.II, - @tensor(stokes.τ_o)..., - @strain(stokes)..., + @tensor(stokes.τ_o), + @strain(stokes), stokes.P, η, η_vep, @@ -496,6 +496,7 @@ function JustRelax.solve!( # solver loop @copy stokes.P0 stokes.P + @copy stokes.P0 stokes.P wtime0 = 0.0 θ = @zeros(ni...) λ = @zeros(ni...) @@ -515,7 +516,7 @@ function JustRelax.solve!( stokes.P0, stokes.R.RP, stokes.∇V, - η, + ητ, rheology, phase_ratios.center, dt, @@ -534,17 +535,18 @@ function JustRelax.solve!( if rem(iter, nout) == 0 @copy η0 η end - # if do_visc - ν = 1e-2 - compute_viscosity!( - η, - ν, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) + if do_visc + ν = 1e-2 + @parallel (@idx ni) compute_viscosity!( + η, + ν, + phase_ratios.center, + @strain(stokes)..., + args, + rheology, + viscosity_cutoff, + ) + end @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), @@ -686,12 +688,14 @@ function JustRelax.solve!( @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... ) - @parallel (@idx ni) compute_ρg!(ρg[2], ϕ, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel (@idx ni) compute_ρg!( + ρg[end], ϕ, rheology, (T=thermal.Tc, P=stokes.P) + ) @parallel (@idx ni) compute_τ_gp!( - @tensor_center(stokes.τ)..., + @tensor_center(stokes.τ), stokes.τ.II, - @tensor(stokes.τ_o)..., - @strain(stokes)..., + @tensor(stokes.τ_o), + @strain(stokes), η, η_vep, thermal.T, diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 61b16c7c..48b893f7 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -1,38 +1,3 @@ -## DIMENSION AGNOSTIC ELASTIC KERNELS - -@parallel function elastic_iter_params!( - dτ_Rho::AbstractArray, - Gdτ::AbstractArray, - ητ::AbstractArray, - Vpdτ::T, - G::T, - dt::M, - Re::T, - r::T, - max_li::T, -) where {T,M} - @all(dτ_Rho) = Vpdτ * max_li / Re / (one(T) / (one(T) / @all(ητ) + one(T) / (G * dt))) - @all(Gdτ) = Vpdτ^2 / @all(dτ_Rho) / (r + T(2.0)) - return nothing -end - -@parallel function elastic_iter_params!( - dτ_Rho::AbstractArray, - Gdτ::AbstractArray, - ητ::AbstractArray, - Vpdτ::T, - G::AbstractArray, - dt::M, - Re::T, - r::T, - max_li::T, -) where {T,M} - @all(dτ_Rho) = - Vpdτ * max_li / Re / (one(T) / (one(T) / @all(ητ) + one(T) / (@all(G) * dt))) - @all(Gdτ) = Vpdτ^2 / @all(dτ_Rho) / (r + T(2.0)) - return nothing -end - ## 3D ELASTICITY MODULE module Stokes3D @@ -46,18 +11,20 @@ using LinearAlgebra using Printf using GeoParams -import JustRelax: elastic_iter_params!, PTArray, Velocity, SymmetricTensor, pureshear_bc! +import JustRelax: PTArray, Velocity, SymmetricTensor, pureshear_bc! import JustRelax: Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG +import JustRelax: tensor_invariant!, compute_τ_nonlinear!, compute_τ_vertex!, compute_τ! import JustRelax: compute_maxloc!, solve! import JustRelax: mean_mpi, norm_mpi, minimum_mpi, maximum_mpi -export solve!, pureshear_bc! - +include("StressRotation.jl") include("PressureKernels.jl") -include("StressKernels.jl") include("VelocityKernels.jl") +export solve!, pureshear_bc! +rotate_stress_particles_jaumann!, +rotate_stress_particles_roation_matrix!, compute_vorticity!, @parallel function update_τ_o!( τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, τxx, τyy, τzz, τxy, τxz, τyz ) @@ -123,9 +90,6 @@ function JustRelax.solve!( ητ = deepcopy(η) compute_maxloc!(ητ, η) update_halo!(ητ) - # @parallel (1:size(ητ, 2), 1:size(ητ, 3)) free_slip_x!(ητ) - # @parallel (1:size(ητ, 1), 1:size(ητ, 3)) free_slip_y!(ητ) - # @parallel (1:size(ητ, 1), 1:size(ητ, 2)) free_slip_z!(ητ) # errors err = 2 * ϵ @@ -298,8 +262,8 @@ function JustRelax.solve!( stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... ) - # Update buoyancy - @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) + # # Update buoyancy + # @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) ν = 1e-3 @parallel (@idx ni) compute_viscosity!( @@ -319,9 +283,7 @@ function JustRelax.solve!( @strain(stokes), stokes.P, η, - η_vep, - λ, - rheology, # needs to be a tuple + @ones(ni...), dt, pt_stokes.θ_dτ, ) @@ -462,7 +424,7 @@ function JustRelax.solve!( stokes.P0, stokes.R.RP, stokes.∇V, - η, + ητ, rheology, phase_ratios.center, dt, @@ -475,9 +437,7 @@ function JustRelax.solve!( ) # Update buoyancy - @parallel (@idx ni) compute_ρg!( - ρg[3], phase_ratios.center, rheology, (T=thermal.T, P=stokes.P) - ) + @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) # Update viscosity ν = 1e-2 diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 673f01fd..3ac47590 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -70,6 +70,23 @@ end return nothing end +@parallel_indices (i, j) function compute_τ_vertex!( + τxy::AbstractArray{T,2}, εxy, η, θ_dτ +) where {T} + @inline av(A) = _av_a(A, i, j) + + # Shear components + if all((i, j) .< size(τxy) .- 1) + I = i + 1, j + 1 + av_η_ij = av(η) + denominator = inv(θ_dτ + 1.0) + + τxy[I...] += (-τxy[I...] + 2.0 * av_η_ij * εxy[I...]) * denominator + end + + return nothing +end + @parallel_indices (i, j, k) function compute_τ!( τxx, τyy, @@ -164,7 +181,7 @@ end end @parallel_indices (i, j, k) function compute_τ_vertex!( - τyz, τxz, τxy, τyz_o, τxz_o, τxy_o, εyz, εxz, εxy, η, G, dt, θ_dτ + τyz, τxz, τxy, εyz, εxz, εxy, ηvep, θ_dτ ) harm_xy(A) = _harm_xyi(A, i, j, k) harm_xz(A) = _harm_xzi(A, i, j, k) @@ -177,20 +194,26 @@ end @inbounds begin # Compute τ_xy if (1 < i < size(τxy, 1)) && (1 < j < size(τxy, 2)) && k ≤ size(τxy, 3) - η_ij = harm_xy(η) + η_ij = harm_xy(ηvep) + denominator = inv(θ_dτ + 1.0) + τxy[i, j, k] += (-get(τxy) + 2.0 * η_ij * get(εxy)) * denominator denominator = inv(θ_dτ + 1.0) τxy[i, j, k] += (-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)) - η_ij = harm_xz(η) + η_ij = harm_xz(ηvep) + denominator = inv(θ_dτ + 1.0) + τxz[i, j, k] += (-get(τxz) + 2.0 * η_ij * get(εxz)) * denominator denominator = inv(θ_dτ + 1.0) τxz[i, j, k] += (-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)) - η_ij = harm_yz(η) + η_ij = harm_yz(ηvep) + denominator = inv(θ_dτ + 1.0) + τyz[i, j, k] += (-get(τyz) + 2.0 * η_ij * get(εyz)) * denominator denominator = inv(θ_dτ + 1.0) τyz[i, j, k] += (-get(τyz) + 2.0 * η_ij * get(εyz)) * denominator end @@ -309,7 +332,6 @@ end end @parallel_indices (I...) function tensor_invariant!(II, xx, yy, xyv) - # convinience closure @inline gather(A) = _gather(A, I...)