diff --git a/miniapps/subduction/2D/Subduction2D.jl b/miniapps/subduction/2D/Subduction2D.jl index 0d24ca06..65b328d8 100644 --- a/miniapps/subduction/2D/Subduction2D.jl +++ b/miniapps/subduction/2D/Subduction2D.jl @@ -1,5 +1,5 @@ -const isCUDA = false -# const isCUDA = true +# const isCUDA = false +const isCUDA = true @static if isCUDA using CUDA @@ -7,7 +7,6 @@ end using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO - const backend = @static if isCUDA CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend else @@ -92,7 +91,13 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk grid_vxi = velocity_grids(xci, xvi, di) # material phase & temperature pPhases, pT = init_cell_arrays(particles, Val(2)) - particle_args = (pT, pPhases) + + # particle fields for the stress rotation + pτ = pτxx, pτyy, pτxy = init_cell_arrays(particles, Val(3)) # stress + # pτ_o = pτxx_o, pτyy_o, pτxy_o = init_cell_arrays(particles, Val(3)) # old stress + pω = pωxy, = init_cell_arrays(particles, Val(1)) # vorticity + particle_args = (pT, pPhases, pτ..., pω...) + particle_args_reduced = (pT, pτ..., pω...) # Assign particles phases anomaly phases_device = PTArray(backend)(phases_GMG) @@ -104,7 +109,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re=3π, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, Re = 3e0, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 + # pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, Re = 2π√2, r=0.7, CFL = 0.9 / √2.1) # Re=3π, r=0.7 # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- @@ -133,7 +139,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - backend, rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-5, CFL=0.95 / √3 + backend, rheology, phase_ratios, args0, dt, ni, di, li; ϵ=1e-8, CFL=0.95 / √2 ) # Boundary conditions @@ -159,26 +165,6 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk Vy_v = @zeros(ni.+1...) end - # # Smooth out thermal field --------------------------- - # for _ in 1:10 - # heatdiffusion_PT!( - # thermal, - # pt_thermal, - # thermal_bc, - # rheology, - # 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) @@ -187,9 +173,16 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk end grid2particle!(pT, xvi, T_buffer, particles) + τxx_v = @zeros(ni.+1...) + τyy_v = @zeros(ni.+1...) + # Time loop t, it = 0.0, 0 + fig_iters = Figure(size=(1200, 800)) + ax_iters1 = Axis(fig_iters[1,1], aspect = 1, title = "error") + ax_iters2 = Axis(fig_iters[1,2], aspect = 1, title = "num iters / ny") + while it < 1000 # run only for 5 Myrs # interpolate fields from particle to grid vertices @@ -202,6 +195,10 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + particle2centroid!(stokes.τ.xx, pτxx, xci, particles) + particle2centroid!(stokes.τ.yy, pτyy, xci, particles) + particle2grid!(stokes.τ.xy, pτxy, xvi, particles) + # Stokes solver ---------------- t_stokes = @elapsed begin out = solve!( @@ -216,14 +213,30 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk dt, igg; kwargs = ( - iterMax = 50e3, + iterMax = 100e3, nout = 2e3, viscosity_cutoff = viscosity_cutoff, free_surface = false, viscosity_relaxation = 1e-2 ) ); + + scatter!(ax_iters1, [it], [log10(out.err_evo1[end])], markersize = 10, color=:blue) + scatter!(ax_iters2, [it], [out.iter/ny], markersize = 10, color=:blue) + fig_iters + + if it == 1 || rem(it, 10) == 0 + save(joinpath(figdir, "errors.png"), fig_iters) + end end + + center2vertex!(τxx_v, stokes.τ.xx) + center2vertex!(τyy_v, stokes.τ.yy) + centroid2particle!(pτxx , xci, stokes.τ.xx, particles) + centroid2particle!(pτyy , xci, stokes.τ.yy, particles) + grid2particle!(pτxy, xvi, stokes.τ.xy, particles) + rotate_stress_particles!(pτ, pω, particles, dt) + println("Stokes solver time ") println(" Total time: $t_stokes s") println(" Time/iteration: $(t_stokes / out.iter) s") @@ -262,9 +275,16 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk 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) + # inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + inject_particles_phase!( + particles, + pPhases, + particle_args_reduced, + (T_buffer, τxx_v, τyy_v, stokes.τ.xy, stokes.ω.xy), + xvi + ) + # update phase ratios update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases) @@ -323,7 +343,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk # Plot particles phase h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1) # Plot 2nd invariant of strain rate - h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + # h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.ε.II)) , colormap=:batlow) + h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array((stokes.τ.II)) , colormap=:batlow) # Plot effective viscosity h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) hidexdecorations!(ax1) @@ -348,7 +369,7 @@ end do_vtk = true # set to true to generate VTK files for ParaView figdir = "Subduction2D" n = 128 -nx, ny = 128, 64 +nx, ny = 256, 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)...) @@ -356,4 +377,4 @@ else igg end -main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); +main(li, origin, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk); \ No newline at end of file diff --git a/miniapps/subduction/2D/Subduction2D_rheology.jl b/miniapps/subduction/2D/Subduction2D_rheology.jl index 66545156..7fd01521 100644 --- a/miniapps/subduction/2D/Subduction2D_rheology.jl +++ b/miniapps/subduction/2D/Subduction2D_rheology.jl @@ -7,7 +7,9 @@ function init_rheology_nonNewtonian() # diffusion laws diff_wet_olivine = SetDiffusionCreep(Diffusion.wet_olivine_Hirth_2003) - lithosphere_rheology = CompositeRheology( (disl_wet_olivine, diff_wet_olivine)) + el = ConstantElasticity(; G = 40e9) + + lithosphere_rheology = CompositeRheology( (el, disl_wet_olivine, diff_wet_olivine)) init_rheologies(lithosphere_rheology) end @@ -19,10 +21,11 @@ function init_rheology_nonNewtonian_plastic() # plasticity ϕ_wet_olivine = asind(0.1) C_wet_olivine = 1e6 - η_reg = 1e19 - + η_reg = 1e16 + el = ConstantElasticity(; G = 40e9, ν = 0.45) lithosphere_rheology = CompositeRheology( ( + el, disl_wet_olivine, diff_wet_olivine, DruckerPrager_regularised(; C = C_wet_olivine, ϕ = ϕ_wet_olivine, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity @@ -32,15 +35,16 @@ function init_rheology_nonNewtonian_plastic() end function init_rheology_linear() - lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e22),)) + el = ConstantElasticity(; G = 40e9, ν = 0.45) + # lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), )) + lithosphere_rheology = CompositeRheology( (LinearViscous(; η=1e23), el)) init_rheologies(lithosphere_rheology) end function init_rheologies(lithosphere_rheology) # common physical properties α = 2.4e-5 # 1 / K - Cp = 750 # J / kg K - + Cp = 750 # J / kg K # Define rheolgy struct rheology = ( # Name = "Asthenoshpere", diff --git a/src/common.jl b/src/common.jl index 349a793b..c88da13b 100644 --- a/src/common.jl +++ b/src/common.jl @@ -57,9 +57,6 @@ export compute_viscosity! include("rheology/Melting.jl") export compute_melt_fraction! -# include("thermal_diffusion/DiffusionExplicit.jl") -# export ThermalParameters - include("particles/subgrid_diffusion.jl") export subgrid_characteristic_time! @@ -72,12 +69,17 @@ export WENO_advection! # Stokes include("rheology/GeoParams.jl") + include("rheology/StressUpdate.jl") + include("stokes/StressRotation.jl") + include("stokes/StressKernels.jl") export tensor_invariant! include("stokes/PressureKernels.jl") +export rotate_stress_particles! + include("stokes/VelocityKernels.jl") # thermal diffusion diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 0b5b9804..b48fcb19 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -368,4 +368,27 @@ function JR2D.WENO_advection!(u::ROCArrayArray, Vxi::NTuple, weno, di, dt) return WENO_advection!(u, Vxi, weno, di, dt) end +# stress rotation on particles + +function JR2D.rotate_stress_particles!( + τ::NTuple, + ω::NTuple, + particles::Particles{JustPIC.AMDGPUBackend}, + dt; + method::Symbol=:matrix, +) + fn = if method === :matrix + rotate_stress_particles_rotation_matrix! + + elseif method === :jaumann + rotate_stress_particles_jaumann! + + else + error("Unknown method: $method. Valid methods are :matrix and :jaumann") + end + @parallel (@idx size(particles.index)) fn(τ..., ω..., particles.index, dt) + + return nothing +end + end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index a40c754c..76fce83d 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -386,4 +386,25 @@ function JR3D.WENO_advection!(u::ROCArray, Vxi::NTuple, weno, di, dt) return WENO_advection!(u, Vxi, weno, di, dt) end +function JR3D.rotate_stress_particles!( + τ::NTuple, + ω::NTuple, + particles::Particles{JustPIC.AMDGPUBackend}, + dt; + method::Symbol=:matrix, +) + fn = if method === :matrix + rotate_stress_particles_rotation_matrix! + + elseif method === :jaumann + rotate_stress_particles_jaumann! + + else + error("Unknown method: $method. Valid methods are :matrix and :jaumann") + end + @parallel (@idx size(particles.index)) fn(τ..., ω..., particles.index, dt) + + return nothing +end + end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 5f448607..c4315698 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -353,4 +353,23 @@ function JR2D.WENO_advection!(u::CuArray, Vxi::NTuple, weno, di, dt) return WENO_advection!(u, Vxi, weno, di, dt) end +# stress rotation on particles + +function JR2D.rotate_stress_particles!( + τ::NTuple, ω::NTuple, particles::Particles{CUDABackend}, dt; method::Symbol=:matrix +) + fn = if method === :matrix + rotate_stress_particles_rotation_matrix! + + elseif method === :jaumann + rotate_stress_particles_jaumann! + + else + error("Unknown method: $method. Valid methods are :matrix and :jaumann") + end + @parallel (@idx size(particles.index)) fn(τ..., ω..., particles.index, dt) + + return nothing +end + end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index da702989..4de3b352 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -385,4 +385,23 @@ function JR3D.WENO_advection!(u::CuArray, Vxi::NTuple, weno, di, dt) return WENO_advection!(u, Vxi, weno, di, dt) end +# stress rotation on particles + +function JR3D.rotate_stress_particles!( + τ::NTuple, ω::NTuple, particles::Particles{CUDABackend}, dt; method::Symbol=:matrix +) + fn = if method === :matrix + rotate_stress_particles_rotation_matrix! + + elseif method === :jaumann + rotate_stress_particles_jaumann! + + else + error("Unknown method: $method. Valid methods are :matrix and :jaumann") + end + @parallel (@idx size(particles.index)) fn(τ..., ω..., particles.index, dt) + + return nothing +end + end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 274b51b4..ce65dd0b 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -335,6 +335,9 @@ function _solve!( # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + while iter < 2 || (err > ϵ && iter ≤ iterMax) wtime0 += @elapsed begin @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) @@ -438,8 +441,10 @@ function _solve!( stokes.P .= θ # θ = P + plastic_overpressure - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) @@ -526,6 +531,9 @@ function _solve!( compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) displacement2velocity!(stokes, dt, flow_bcs) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + while iter ≤ iterMax iterMin < iter && err < ϵ && break @@ -661,13 +669,15 @@ function _solve!( isnan(err) && error("NaN(s)") end - if igg.me == 0 && err ≤ ϵ && iter ≥ 20000 + if igg.me == 0 && err ≤ ϵ println("Pseudo-transient iterations converged in $iter iterations") end end - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index d42527ea..d4617339 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -214,6 +214,9 @@ function _solve!( # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + # solver loop wtime0 = 0.0 while iter < 2 || (err > ϵ && iter ≤ iterMax) @@ -345,8 +348,10 @@ function _solve!( av_time = wtime0 / (iter - 1) # average time per iteration - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.yz, stokes.ω.xz, stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) @@ -423,6 +428,9 @@ function _solve!( # convert displacement to velocity displacement2velocity!(stokes, dt, flow_bcs) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + while iter < 2 || (err > ϵ && iter ≤ iterMax) wtime0 += @elapsed begin # ~preconditioner @@ -541,8 +549,10 @@ function _solve!( av_time = wtime0 / (iter - 1) # average time per iteration - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) + # compute vorticity + @parallel (@idx ni .+ 1) compute_vorticity!( + stokes.ω.yz, stokes.ω.xz, stokes.ω.xy, @velocity(stokes)..., inv.(di)... + ) # accumulate plastic strain tensor @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) diff --git a/src/stokes/StressRotation.jl b/src/stokes/StressRotation.jl index 3ab7635b..860ca311 100644 --- a/src/stokes/StressRotation.jl +++ b/src/stokes/StressRotation.jl @@ -1,50 +1,122 @@ using StaticArrays +# Vorticity tensor + +@parallel_indices (I...) function compute_vorticity!(ωxy, Vx, Vy, _dx, _dy) + dx(A) = _d_xa(A, I..., _dx) + dy(A) = _d_ya(A, I..., _dy) + + ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) + + return nothing +end + +@parallel_indices (I...) function compute_vorticity!( + ωyz, ωxz, ωxy, Vx, Vy, Vz, _dx, _dy, _dz +) + dx(A) = _d_xa(A, I..., _dx) + dy(A) = _d_ya(A, I..., _dy) + dz(A) = _d_za(A, I..., _dz) + + ωyz[I...] = 0.5 * (dy(Vz) - dz(Vy)) + ωxz[I...] = 0.5 * (dz(Vx) - dx(Vz)) + ωxy[I...] = 0.5 * (dx(Vy) - dy(Vx)) + + return nothing +end + ## 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) +function rotate_stress_particles!( + τ::NTuple, ω::NTuple, particles::Particles, dt; method::Symbol=:matrix +) + @parallel (@idx size(particles.index)) rotate_stress_particles_GeoParams!( + τ..., ω..., particles.index, dt + ) + return nothing +end - vorticity[i, j] = 0.5 * (dx(Vy) - dy(Vx)) +@parallel_indices (I...) function rotate_stress_particles_GeoParams!( + xx, yy, xy, ω, index, dt +) + for ip in cellaxes(index) + @index(index[ip, I...]) || continue # no particle in this location + + ω_xy = @index ω[ip, I...] + τ_xx = @index xx[ip, I...] + τ_yy = @index yy[ip, I...] + τ_xy = @index xy[ip, I...] + + τ_rotated = GeoParams.rotate_elastic_stress2D(ω_xy, (τ_xx, τ_yy, τ_xy), dt) + + @index xx[ip, I...] = τ_rotated[1] + @index yy[ip, I...] = τ_rotated[2] + @index xy[ip, I...] = τ_rotated[3] + end return nothing end -@parallel_indices (i, j) function rotate_stress_particles_jaumann!(xx, yy, xy, ω, index, dt) - cell = i, j +@parallel_indices (I...) function rotate_stress_particles_GeoParams!( + xx, yy, zz, yz, xz, xy, ωyz, ωxz, ωxy, index, dt +) + for ip in cellaxes(index) + @index(index[ip, I...]) || continue # no particle in this location + + ω_yz = @index ωyz[ip, I...] + ω_xz = @index ωxz[ip, I...] + ω_xy = @index ωxy[ip, I...] + τ_xx = @index xx[ip, I...] + τ_yy = @index yy[ip, I...] + τ_yz = @index yz[ip, I...] + τ_xz = @index xz[ip, I...] + τ_xy = @index xy[ip, I...] + + τ_rotated = GeoParams.rotate_elastic_stress3D( + (ω_yz, ω_xz, ω_xy), (τ_xx, τ_yy, τ_xy, τ_yz, τ_xz, τ_xy), dt + ) + + @index xx[ip, I...] = τ_rotated[1] + @index yy[ip, I...] = τ_rotated[2] + @index zz[ip, I...] = τ_rotated[3] + @index yz[ip, I...] = τ_rotated[4] + @index xz[ip, I...] = τ_rotated[5] + @index xy[ip, I...] = τ_rotated[6] + end - for ip in JustRelax.cellaxes(index) - !@index(index[ip, cell...]) && continue # no particle in this location + return nothing +end + +@parallel_indices (I) function rotate_stress_particles_jaumann!(xx, yy, xy, ω, index, dt) + for ip in cellaxes(index) + !@index(index[ip, I...]) && continue # no particle in this location - ω_xy = @index ω[ip, cell...] - τ_xx = @index xx[ip, cell...] - τ_yy = @index yy[ip, cell...] - τ_xy = @index xy[ip, cell...] + ω_xy = @index ω[ip, I...] + τ_xx = @index xx[ip, I...] + τ_yy = @index yy[ip, I...] + τ_xy = @index xy[ip, I...] tmp = τ_xy * ω_xy * 2.0 - @index xx[ip, cell...] = fma(dt, cte, τ_xx) - @index yy[ip, cell...] = fma(dt, cte, τ_yy) - @index xy[ip, cell...] = fma(dt, (τ_xx - τ_yy) * ω_xy, τ_xy) + @index xx[ip, I...] = fma(dt, cte, τ_xx) + @index yy[ip, I...] = fma(dt, cte, τ_yy) + @index xy[ip, I...] = fma(dt, (τ_xx - τ_yy) * ω_xy, τ_xy) end return nothing end -@parallel_indices (i, j) function rotate_stress_particles_rotation_matrix!( +@parallel_indices (I...) function rotate_stress_particles_rotation_matrix!( xx, yy, xy, ω, index, dt ) - cell = i, j - - for ip in JustRelax.cellaxes(index) - !@index(index[ip, cell...]) && continue # no particle in this location + for ip in cellaxes(index) + !@index(index[ip, I...]) && continue # no particle in this location - θ = dt * @index ω[ip, cell...] + θ = dt * @index ω[ip, I...] sinθ, cosθ = sincos(θ) - τ_xx = @index xx[ip, cell...] - τ_yy = @index yy[ip, cell...] - τ_xy = @index xy[ip, cell...] + τ_xx = @index xx[ip, I...] + τ_yy = @index yy[ip, I...] + τ_xy = @index xy[ip, I...] R = @SMatrix [ cosθ -sinθ @@ -59,9 +131,9 @@ end # this could be fully unrolled in 2D τr = R * τ * R' - @index xx[ip, cell...] = τr[1, 1] - @index yy[ip, cell...] = τr[2, 2] - @index xy[ip, cell...] = τr[1, 2] + @index xx[ip, I...] = τr[1, 1] + @index yy[ip, I...] = τr[2, 2] + @index xy[ip, I...] = τr[1, 2] end return nothing diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl index 374ad3d8..eb107986 100644 --- a/src/types/constructors/stokes.jl +++ b/src/types/constructors/stokes.jl @@ -39,15 +39,15 @@ end ## Vorticity type function Vorticity(nx::Integer, ny::Integer) - xy = @zeros(nx, ny) + xy = @zeros(nx + 1, ny + 1) 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) + yz = @zeros(nx, ny + 1, nz + 1) + xz = @zeros(nx + 1, ny, nz + 1) + xy = @zeros(nx + 1, ny + 1, nz) return JustRelax.Vorticity(yz, xz, xy) end