From da5a3e8845020396369a2b3b477bf46185eb8f6d Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Tue, 2 Apr 2024 08:34:57 +0200 Subject: [PATCH 01/91] massive rework --- Project.toml | 2 +- .../stokes2D/shear_heating/Shearheating2D.jl | 86 +- .../shear_heating/Shearheating_rheology.jl | 4 +- src/Interpolations.jl | 29 +- src/JustRelax.jl | 55 +- src/JustRelax_CPU.jl | 45 + src/advection/weno5.jl | 10 +- src/boundaryconditions/BoundaryConditions.jl | 34 +- src/common.jl | 69 ++ src/phases/phases.jl | 117 ++- src/rheology/BuoyancyForces.jl | 20 +- src/rheology/Viscosity.jl | 58 +- src/stokes/Stokes2D.jl | 136 +-- src/stokes/Stokes2D_old.jl | 808 +++++++++++++++++ src/stokes/Stokes3D.jl | 84 +- src/stokes/Stokes3D_old.jl | 561 ++++++++++++ src/stokes/StressKernels.jl | 81 +- src/thermal_diffusion/DiffusionPT.jl | 819 +----------------- .../DiffusionPT_GeoParams.jl | 103 +++ .../DiffusionPT_coefficients.jl | 85 ++ src/thermal_diffusion/DiffusionPT_kernels.jl | 549 ++++++++++++ src/thermal_diffusion/DiffusionPT_solver.jl | 206 +++++ src/thermal_diffusion/Shearheating.jl | 31 +- src/types/heat_diffusion.jl | 95 ++ src/types/stokes.jl | 246 ++++++ src/types/traits.jl | 6 + test_types.jl | 67 ++ 27 files changed, 3149 insertions(+), 1257 deletions(-) create mode 100644 src/JustRelax_CPU.jl create mode 100644 src/common.jl create mode 100644 src/stokes/Stokes2D_old.jl create mode 100644 src/stokes/Stokes3D_old.jl create mode 100644 src/thermal_diffusion/DiffusionPT_GeoParams.jl create mode 100644 src/thermal_diffusion/DiffusionPT_coefficients.jl create mode 100644 src/thermal_diffusion/DiffusionPT_kernels.jl create mode 100644 src/thermal_diffusion/DiffusionPT_solver.jl create mode 100644 src/types/heat_diffusion.jl create mode 100644 src/types/stokes.jl create mode 100644 src/types/traits.jl create mode 100644 test_types.jl diff --git a/Project.toml b/Project.toml index 1227b3f7..dc2513e4 100644 --- a/Project.toml +++ b/Project.toml @@ -31,7 +31,7 @@ HDF5 = "0.17.1" ImplicitGlobalGrid = "0.14.0, 0.15.0" MPI = "0.20" MuladdMacro = "0.2" -ParallelStencil = "0.9, 0.10, 0.11" +ParallelStencil = "0.12" Reexport = "1.2.2" StaticArrays = "1" Statistics = "1" diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index b500b1e7..842b3fec 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -1,8 +1,9 @@ # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D +using JustRelax.DataIO +# import JustRelax.@cell +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -13,8 +14,8 @@ using JustPIC._2D const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +# model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) +# environment!(model) # Load script dependencies using Printf, LinearAlgebra, GeoParams, CellArrays @@ -23,8 +24,7 @@ using GLMakie # Load file with all the rheology configurations include("Shearheating_rheology.jl") - -# ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- function copyinn_x!(A, B) @@ -72,7 +72,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 40, 1 + nxcell, max_xcell, min_xcell = 20, 40, 10 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -89,12 +89,12 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) r_anomaly = 3e3 # radius of perturbation init_phases!(pPhases, particles, lx/2, yc_anomaly, r_anomaly) phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) # ---------------------------------------------------- @@ -102,51 +102,45 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) thermal = ThermalArrays(ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), - periodicity = (left = false, right = false, top = false, bot = false), ) # Initialize constant temperature @views thermal.T .= 273.0 + 400 thermal_bcs!(thermal.T, thermal_bc) - - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces - ρg = @zeros(ni...), @zeros(ni...) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + ρg = @zeros(ni...), @zeros(ni...) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) # 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, (-Inf, Inf) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - η_vep = copy(η) # PT coefficients for thermal diffusion - pt_thermal = PTThermalCoeffs( + pt_thermal = PTThermalCoeffs( rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1 ) # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - periodicity = (left = false, right = false, top = false, bot = false), + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), ) ## Compression and not extension - fix this - εbg = 5e-14 - stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) + εbg = 5e-14 + stokes.V.Vx .= PTArray()([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray()([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(stokes.V.Vx, stokes.V.Vy) # IO ----- ------------------------------------------- # if it does not exist, make folder where figures are stored if do_vtk - vtk_dir = figdir*"\\vtk" + vtk_dir = joinpath(figdir, "vtk") take(vtk_dir) end take(figdir) @@ -160,7 +154,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./1e3) - scatter!(ax2, Array(log10.(η[:])), Y./1e3) + scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) ylims!(ax2, minimum(xvi[2])./1e3, 0) hideydecorations!(ax2) @@ -185,10 +179,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) while it < 100 # 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, (-Inf, Inf) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -198,18 +192,18 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 75e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) + kwargs = (; + iterMax = 75e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -219,12 +213,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) @views thermal.T[2:end-1, :] .= T_buffer temperature2center!(thermal) - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, rheology, # needs to be a tuple dt, ) @@ -347,4 +339,6 @@ else igg end -main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) +# main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) + +foo(args...; x=1) = x \ No newline at end of file diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl index 74b5694f..40fb2c8a 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl @@ -63,7 +63,7 @@ function init_phases!(phases, particles, Lx, d, r) x = JustRelax.@cell px[ip, i, j] depth = -(JustRelax.@cell py[ip, i, j]) - @cell phases[ip, i, j] = 1.0 # matrix + JustRelax.@cell phases[ip, i, j] = 1.0 # matrix # thermal anomaly - circular if ((x - Lx)^2 + (depth - d)^2 ≤ r^2) @@ -73,5 +73,5 @@ function init_phases!(phases, particles, Lx, d, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, d) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, d) end diff --git a/src/Interpolations.jl b/src/Interpolations.jl index 1dc4f679..52fb268c 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -17,11 +17,11 @@ end # From cell vertices to cell center function temperature2center!(thermal::ThermalArrays) - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + @parallel (@idx size(thermal.Tc)...) temperature2center_kernel!(thermal.Tc, thermal.T) return nothing end -@parallel_indices (i, j) function temperature2center!( +@parallel_indices (i, j) function temperature2center_kernel!( T_center::T, T_vertex::T ) where {T<:AbstractArray{_T,2} where {_T<:Real}} T_center[i, j] = @@ -34,7 +34,7 @@ end return nothing end -@parallel_indices (i, j, k) function temperature2center!( +@parallel_indices (i, j, k) function temperature2center_kernel!( T_center::T, T_vertex::T ) where {T<:AbstractArray{_T,3} where {_T<:Real}} @inline av_T() = _av(T_vertex, i, j, k) @@ -44,17 +44,34 @@ end return nothing end -@parallel function vertex2center!(center, vertex) +function vertex2center!(center, vertex) + @parallel vertex2center_kernel!(center, vertex) + return nothing +end + +@parallel function vertex2center_kernel!(center, vertex) @all(center) = @av(vertex) return nothing end -@parallel function center2vertex!(vertex, center) +function center2vertex!(vertex, center) + @parallel center2vertex_kernel!(vertex, center) + return nothing +end + +@parallel function center2vertex_kernel!(vertex, center) @inn(vertex) = @av(center) return nothing end -@parallel_indices (i, j, k) function center2vertex!( +function center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) + @parallel center2vertex_kernel!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + return nothing +end + +@parallel_indices (i, j, k) function center2vertex_kernel!( vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy ) i1, j1, k1 = (i, j, k) .+ 1 diff --git a/src/JustRelax.jl b/src/JustRelax.jl index f1728e46..f4b6ade1 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -1,7 +1,6 @@ module JustRelax using Reexport -@reexport using ParallelStencil @reexport using ImplicitGlobalGrid using LinearAlgebra using Printf @@ -13,60 +12,28 @@ using StaticArrays function solve!() end +include("types/traits.jl") + include("topology/Topology.jl") export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g -include("MiniKernels.jl") -export _d_xa, - _d_ya, - _d_za, - _d_za, - _d_xi, - _d_yi, - _d_zi, - _d_zi, - _av, - _av_a, - _av_xa, - _av_ya, - _av_x, - _av_y, - _av_z, - _av_yz, - _av_xz, - _av_xy, - _av_yzi, - _av_xzi, - _av_xyi, - _gather, - _gather_yz, - _gather_xz, - _gather_xy, - _harm_x, - _harm_y, - _harm_z, - _harm_yz, - _harm_xz, - _harm_xy, - _harm_xyi, - _harm_xzi, - _harm_yzi, - _current +# include("MiniKernels.jl") include("phases/CellArrays.jl") export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! -include("rheology/StressUpdate.jl") -export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear! +# include("rheology/StressUpdate.jl") +# export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear! -include("MetaJustRelax.jl") +include("JustRelax_CPU.jl") +# include("MetaJustRelax.jl") -include("stokes/MetaStokes.jl") -export PS_Setup, environment!, ps_reset! +# include("stokes/MetaStokes.jl") +# export PS_Setup, environment!, ps_reset! -include("thermal_diffusion/MetaDiffusion.jl") +# include("thermal_diffusion/MetaDiffusion.jl") -include("thermal_diffusion/Rheology.jl") +# include("thermal_diffusion/Rheology.jl") include("IO/DataIO.jl") diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl new file mode 100644 index 00000000..b784fe47 --- /dev/null +++ b/src/JustRelax_CPU.jl @@ -0,0 +1,45 @@ +module JustRelax2D + +using ..JustRelax +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences2D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf + +import JustRelax: IGG, BackendTrait, CPUBackendTrait + +@init_parallel_stencil(Threads, Float64, 2) + +export PTArray + +PTArray() = Array + +include("common.jl") +include("stokes/Stokes2D.jl") +export solve! + +end + +module JustRelax3D + +using ..JustRelax +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences3D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf + +import JustRelax: IGG, BackendTrait, CPUBackendTrait + +@init_parallel_stencil(Threads, Float64, 3) + +export PTArray + +PTArray() = Array + +include("common.jl") +include("stokes/Stokes3D.jl") +export solve! + +end diff --git a/src/advection/weno5.jl b/src/advection/weno5.jl index ef3f8201..3df1e91f 100644 --- a/src/advection/weno5.jl +++ b/src/advection/weno5.jl @@ -47,11 +47,11 @@ The `WENO5` structure contains the parameters and temporary variables used in th # grid size ni::NTuple{N,Int64} # fluxes - ut::A = @zeros(ni...) - fL::A = @zeros(ni...) - fR::A = @zeros(ni...) - fB::A = @zeros(ni...) - fT::A = @zeros(ni...) + ut::A = PTArray(zeros(ni...)) + fL::A = PTArray(zeros(ni...)) + fR::A = PTArray(zeros(ni...)) + fB::A = PTArray(zeros(ni...)) + fT::A = PTArray(zeros(ni...)) # method method::M = Val{1} # 1:JS, 2:Z end diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 9a79a705..32e4c7ab 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -55,8 +55,6 @@ function thermal_bcs!(T, bcs::TemperatureBoundaryConditions) # no flux boundary conditions do_bc(bcs.no_flux) && (@parallel (@idx n) free_slip!(T, bcs.no_flux)) - # periodic conditions - do_bc(bcs.periodicity) && (@parallel (@idx n) periodic_boundaries!(T, bcs.periodicity)) return nothing end @@ -66,6 +64,11 @@ end Apply the prescribed flow boundary conditions `bc` on the `stokes` """ +flow_bcs!(stokes, bcs::FlowBoundaryConditions) = _flow_bcs!(bcs, @velocity(stokes)) +function flow_bcs!(bcs::FlowBoundaryConditions, V::Vararg{T,N}) where {T,N} + return _flow_bcs!(bcs, tuple(V...)) +end + function _flow_bcs!(bcs::FlowBoundaryConditions, V) n = bc_index(V) # no slip boundary conditions @@ -79,11 +82,6 @@ function _flow_bcs!(bcs::FlowBoundaryConditions, V) return nothing end -flow_bcs!(stokes, bcs::FlowBoundaryConditions) = _flow_bcs!(bcs, @velocity(stokes)) -function flow_bcs!(bcs::FlowBoundaryConditions, V::Vararg{T,N}) where {T,N} - return _flow_bcs!(bcs, tuple(V...)) -end - # BOUNDARY CONDITIONS KERNELS @parallel_indices (i) function no_slip!(Ax, Ay, bc) @@ -270,19 +268,19 @@ end function pureshear_bc!( stokes::StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg ) where {T} - stokes.V.Vx[:, 2:(end - 1)] .= PTArray([εbg * x for x in xvi[1], y in xci[2]]) - stokes.V.Vy[2:(end - 1), :] .= PTArray([-εbg * y for x in xci[1], y in xvi[2]]) + stokes.V.Vx[:, 2:(end - 1)] .= PTArray(([εbg * x for x in xvi[1], y in xci[2]])) + stokes.V.Vy[2:(end - 1), :] .= PTArray(([-εbg * y for x in xci[1], y in xvi[2]])) return nothing end -@parallel_indices (j) function free_slip_x!(A::AbstractArray{eltype(PTArray),2}) +@parallel_indices (j) function free_slip_x!(A::AbstractArray{T,2}) where {T} A[1, j] = A[2, j] A[end, j] = A[end - 1, j] return nothing end -@parallel_indices (i) function free_slip_y!(A::AbstractArray{eltype(PTArray),2}) +@parallel_indices (i) function free_slip_y!(A::AbstractArray{T,2}) where {T} A[i, 1] = A[i, 2] A[i, end] = A[i, end - 1] return nothing @@ -358,17 +356,3 @@ function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{3,T}}, Vx, Vy, Vz) w @parallel (1:size(Vy, 1), 1:size(Vy, 2)) free_slip_z!(Vy) end end - -function thermal_boundary_conditions!( - insulation::NamedTuple, T::AbstractArray{_T,3} -) where {_T} - # vertical boundaries - frontal, lateral = insulation - - nx, ny, nz = size(T) - - frontal && (@parallel (1:nx, 1:nz) free_slip_y!(T)) - lateral && (@parallel (1:ny, 1:nz) free_slip_x!(T)) - - return nothing -end diff --git a/src/common.jl b/src/common.jl new file mode 100644 index 00000000..1a7f38f2 --- /dev/null +++ b/src/common.jl @@ -0,0 +1,69 @@ +include("types/stokes.jl") +export StokesArrays, PTStokesCoeffs + +include("types/heat_diffusion.jl") +export ThermalArrays, PTThermalCoeffs + +include("Utils.jl") +export @allocate, + @add, + @idx, + @copy, + @velocity, + @strain, + @stress, + @tensor, + @shear, + @normal, + @stress_center, + @strain_center, + @tensor_center, + @qT, + @qT2, + @residuals, + compute_dt, + multi_copy!, + take + +include("boundaryconditions/BoundaryConditions.jl") +export FlowBoundaryConditions, + TemperatureBoundaryConditions, flow_bcs!, thermal_bcs!, pureshear_bc!, apply_free_slip! + +include("MiniKernels.jl") + +include("phases/phases.jl") +export PhaseRatio, fn_ratio, phase_ratios_center + +include("rheology/BuoyancyForces.jl") +export compute_ρg! + +include("rheology/Viscosity.jl") +export compute_viscosity! + +# include("thermal_diffusion/DiffusionExplicit.jl") +# export ThermalParameters + +include("particles/subgrid_diffusion.jl") +export subgrid_characteristic_time! + +include("Interpolations.jl") +export vertex2center!, center2vertex!, temperature2center!, velocity2vertex! + +include("advection/weno5.jl") +export WENO5, WENO_advection! + +# Stokes + +include("rheology/GeoParams.jl") +include("stokes/StressRotation.jl") +include("stokes/StressKernels.jl") +export tensor_invariant! + +include("stokes/PressureKernels.jl") +include("stokes/VelocityKernels.jl") + +# thermal diffusion + +include("thermal_diffusion/DiffusionPT.jl") +export heatdiffusion_PT! +export compute_shear_heating! diff --git a/src/phases/phases.jl b/src/phases/phases.jl index 9a7890a5..91aa0d73 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -34,58 +34,52 @@ end return N end -# ParallelStencil launch kernel for 2D -@parallel_indices (i, j) function phase_ratios_center(x, phases) - phase_ratios_center(x, phases, i, j) - return nothing -end - -# ParallelStencil launch kernel for 3D -@parallel_indices (i, j, k) function phase_ratios_center(x, phases) - phase_ratios_center(x, phases, i, j, k) - return nothing -end - -""" - phase_ratios_center(x::PhaseRatio, cell::Vararg{Int, N}) - -Compute the phase ratios at the center of the cell `cell` in `x::PhaseRatio`. -""" -function phase_ratios_center(x::PhaseRatio, phases, cell::Vararg{Int,N}) where {N} - return phase_ratios_center(x.center, phases, cell...) -end - -@inline function phase_ratios_center(x::CellArray, phases, cell::Vararg{Int,N}) where {N} - # total number of material phases - num_phases = nphases(x) - # number of active particles in this cell - n = 0 - for j in cellaxes(phases) - n += isinteger(@cell(phases[j, cell...])) && @cell(phases[j, cell...]) != 0 - end - _n = inv(n) - # compute phase ratios - ratios = _phase_ratios_center(phases, num_phases, _n, cell...) - for (i, ratio) in enumerate(ratios) - @cell x[i, cell...] = ratio - end -end - -@generated function _phase_ratios_center( - phases, ::Val{N1}, _n, cell::Vararg{Int,N2} -) where {N1,N2} - quote - Base.@_inline_meta - Base.@nexprs $N1 i -> reps_i = begin - c = 0 - for j in cellaxes(phases) - c += @cell(phases[j, cell...]) == i - end - c * _n - end - Base.@ncall $N1 tuple reps - end -end +# # ParallelStencil launch kernel +# @parallel_indices (I...) function phase_ratios_center(x, phases) +# phase_ratios_center(x, phases, I...) +# return nothing +# end + +# """ +# phase_ratios_center(x::PhaseRatio, cell::Vararg{Int, N}) + +# Compute the phase ratios at the center of the cell `cell` in `x::PhaseRatio`. +# """ +# function phase_ratios_center(x::PhaseRatio, phases, cell::Vararg{Int,N}) where {N} +# return phase_ratios_center(x.center, phases, cell...) +# end + +# @inline function phase_ratios_center(x::CellArray, phases, cell::Vararg{Int,N}) where {N} +# # total number of material phases +# num_phases = nphases(x) +# # number of active particles in this cell +# n = 0 +# for j in cellaxes(phases) +# n += isinteger(@cell(phases[j, cell...])) && @cell(phases[j, cell...]) != 0 +# end +# _n = inv(n) +# # compute phase ratios +# ratios = _phase_ratios_center(phases, num_phases, _n, cell...) +# for (i, ratio) in enumerate(ratios) +# @cell x[i, cell...] = ratio +# end +# end + +# @generated function _phase_ratios_center( +# phases, ::Val{N1}, _n, cell::Vararg{Int,N2} +# ) where {N1,N2} +# quote +# Base.@_inline_meta +# Base.@nexprs $N1 i -> reps_i = begin +# c = 0 +# for j in cellaxes(phases) +# c += @cell(phases[j, cell...]) == i +# end +# c * _n +# end +# Base.@ncall $N1 tuple reps +# end +# end """ fn_ratio(fn::F, rheology::NTuple{N, AbstractMaterialParamsStruct}, ratio) where {N, F} @@ -118,8 +112,15 @@ end end end -# ParallelStencil launch kernel for 2D -@parallel_indices (I...) function phase_ratios_center( +function phase_ratios_center(phase_ratios, particles, grid::Geometry, phases) + ni = size(phases) + @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( ratio_centers, pxi::NTuple{N,T1}, xci::NTuple{N,T2}, di::NTuple{N,T3}, phases ) where {N,T1,T2,T3} @@ -127,15 +128,11 @@ end cell_center = ntuple(i -> xci[i][I[i]], Val(N)) # phase ratios weights (∑w = 1.0) w = phase_ratio_weights( - getindex.(pxi, I...), - phases[I...], - cell_center, - di, - JustRelax.nphases(ratio_centers), + getindex.(pxi, I...), phases[I...], cell_center, di, nphases(ratio_centers) ) # update phase ratios array for k in 1:numphases(ratio_centers) - JustRelax.@cell ratio_centers[k, I...] = w[k] + @cell ratio_centers[k, I...] = w[k] end return nothing diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 59a84548..1b430b6f 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -3,9 +3,15 @@ Calculate the buoyance forces `ρg` for the given GeoParams.jl `rheology` object and correspondent arguments `args`. """ -@parallel_indices (I...) function compute_ρg!(ρg, rheology, args) # index arguments for the current cell cell center +function compute_ρg!(ρg, rheology, args) + ni = size(ρg) + @parallel (@idx ni) compute_ρg_kernel!(ρg, rheology, args) + return nothing +end + +@parallel_indices (I...) function compute_ρg_kernel!(ρg, rheology, args) args_ijk = ntuple_idx(args, I...) - ρg[I...] = JustRelax.compute_buoyancy(rheology, args_ijk) + ρg[I...] = compute_buoyancy(rheology, args_ijk) return nothing end @@ -15,9 +21,15 @@ end 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...) function compute_ρg!(ρg, phase_ratios, rheology, args) # index arguments for the current cell cell center +function compute_ρg!(ρg, phase_ratios::PhaseRatio, rheology, args) + ni = size(ρg) + @parallel (@idx ni) compute_ρg_kernel!(ρg, phase_ratios.center, rheology, args) + return nothing +end + +@parallel_indices (I...) function compute_ρg_kernel!(ρg, phase_ratios, rheology, args) args_ijk = ntuple_idx(args, I...) - ρg[I...] = JustRelax.compute_buoyancy(rheology, args_ijk, phase_ratios[I...]) + ρg[I...] = compute_buoyancy(rheology, args_ijk, phase_ratios[I...]) return nothing end diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index e07ca4c1..660e7f74 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,6 +1,13 @@ ## 2D KERNELS +function compute_viscosity!(stokes::StokesArrays, ν, args, rheology, cutoff) + ni = size(stokes.viscosity.η) + @parallel (@idx ni) compute_viscosity_kernel!( + stokes.viscosity.η, ν, @strain(stokes)..., args, rheology, cutoff + ) + return nothing +end -@parallel_indices (I...) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity_kernel!( η, ν, εxx, εyy, εxyv, args, rheology, cutoff ) @@ -30,7 +37,15 @@ return nothing end -@parallel_indices (I...) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) +function compute_viscosity!(η, ν, εII::AbstractArray, args, rheology, cutoff) + ni = size(stokes.viscosity.η) + @parallel (@idx ni) compute_viscosity_kernel!(η, ν, εII, args, rheology, cutoff) + return nothing +end + +@parallel_indices (I...) function compute_viscosity_kernel!( + η, ν, εII, args, rheology, cutoff +) @inbounds begin # argument fields at local index args_ij = local_viscosity_args(args, I...) @@ -47,7 +62,23 @@ end return nothing end -@parallel_indices (I...) function compute_viscosity!( +function compute_viscosity!( + stokes::StokesArrays, ν, phase_ratios::PhaseRatio, args, rheology, cutoff +) + ni = size(stokes.viscosity.η) + @parallel (@idx ni) compute_viscosity_kernel!( + stokes.viscosity.η, + ν, + phase_ratios.center, + @strain(stokes)..., + args, + rheology, + cutoff, + ) + return nothing +end + +@parallel_indices (I...) function compute_viscosity_kernel!( η, ν, ratios_center, εxx, εyy, εxyv, args, rheology, cutoff ) @@ -82,7 +113,7 @@ end ## 3D KERNELS -@parallel_indices (I...) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity_kernel!( η, ν, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, cutoff ) @@ -115,24 +146,7 @@ end return nothing end -@parallel_indices (I...) function compute_viscosity!(η, ν, εII, args, rheology, cutoff) - @inbounds begin - # # argument fields at local index - args_ijk = local_viscosity_args(args, I...) - - # compute second invariant of strain rate tensor - εII_ij = εII[I...] - - # update stress and effective viscosity - ηi = compute_viscosity_εII(rheology, εII_ij, args_ijk) - ηi = continuation_log(ηi, η[I...], ν) - η[I...] = clamp(ηi, cutoff...) - end - - return nothing -end - -@parallel_indices (I...) function compute_viscosity!( +@parallel_indices (I...) function compute_viscosity_kernel!( η, ν, ratios_center, εxx, εyy, εzz, εyzv, εxzv, εxyv, args, rheology, cutoff ) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index e2bb80f7..45868fff 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -1,88 +1,25 @@ -## DIMENSION AGNOSTIC 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 - ## 2D STOKES MODULE - -module Stokes2D - -using ImplicitGlobalGrid -using ..JustRelax -using CUDA, AMDGPU -using ParallelStencil -# using ParallelStencil.FiniteDifferences2D -using GeoParams, LinearAlgebra, Printf - -import JustRelax: elastic_iter_params!, PTArray, Velocity, SymmetricTensor -import JustRelax: - Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG -import JustRelax: compute_maxloc!, solve! -import JustRelax: mean_mpi, norm_mpi, maximum_mpi, minimum_mpi, backend - -@eval @init_parallel_stencil($backend, Float64, 2) - -include("../rheology/GeoParams.jl") -include("StressRotation.jl") -include("StressKernels.jl") -include("PressureKernels.jl") -include("VelocityKernels.jl") -include("StressKernels.jl") - -export solve! - -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 - τxx_o, τyy_o, τxy_o, τxy_o_c = stokes.τ_o.xx, - stokes.τ_o.yy, stokes.τ_o.xy, - stokes.τ_o.xy_c - +function update_τ_o!(stokes::StokesArrays) @parallel (@idx size(τxy)) multi_copy!( - (τxx_o, τyy_o, τxy_o, τxy_o_c), (τxx, τyy, τxy, τxy_c) + @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) ) return nothing end ## 2D VISCO-ELASTIC STOKES SOLVER -# viscous solver -function JustRelax.solve!( - stokes::StokesArrays{Viscous,A,B,C,D,2}, +# backend trait +function solve!(stokes::StokesArrays, args...; kwargs) + return solve!(backend(stokes), stokes, args...; kwargs...) +end + +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{2,T}, flow_bcs::FlowBoundaryConditions, ρg, - η, K, dt, igg::IGG; @@ -90,11 +27,13 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 1), verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # unpack _dx, _dy = inv.(di) (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η) = stokes.viscosity ni = size(stokes.P) # ~preconditioner @@ -203,13 +142,13 @@ function JustRelax.solve!( end # visco-elastic solver -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{2,T}, flow_bcs, ρg, - η, G, K, dt, @@ -218,11 +157,13 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 1), verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # unpack _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; ϵ, r, θ_dτ) = pt_stokes + (; η) = stokes.viscosity ni = size(stokes.P) # ~preconditioner @@ -318,14 +259,13 @@ end # GeoParams: general (visco-elasto-plastic) solver -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{2,T}, flow_bcs, ρg, - η, - η_vep, rheology::MaterialParams, args, dt, @@ -335,11 +275,13 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 0), verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # unpack _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; ϵ, r, θ_dτ) = pt_stokes + (; η, η_vep) = stokes.viscosity ni = size(stokes.P) # ~preconditioner @@ -480,14 +422,13 @@ end ## With phase ratios -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{2,T}, flow_bcs, ρg, - η, - η_vep, phase_ratios::PhaseRatio, rheology, args, @@ -501,12 +442,14 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 0), verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # unpack _di = inv.(di) (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity ni = size(stokes.P) # ~preconditioner @@ -693,16 +636,15 @@ function JustRelax.solve!( ) end -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, thermal::ThermalArrays, pt_stokes::PTStokesCoeffs, di::NTuple{2,T}, flow_bcs, ϕ, ρg, - η, - η_vep, phase_v, phase_c, args_η, @@ -713,12 +655,14 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 1), verbose=true, -) where {A,B,C,D,N,T} + kwargs..., +) where {N,T} # unpack _di = inv.(di) (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + (; η, η_vep) = stokes.viscosity ni = size(stokes.P) # ~preconditioner ητ = deepcopy(η) @@ -838,5 +782,3 @@ function JustRelax.solve!( norm_∇V=norm_∇V, ) end - -end # END OF MODULE diff --git a/src/stokes/Stokes2D_old.jl b/src/stokes/Stokes2D_old.jl new file mode 100644 index 00000000..cfe38ce3 --- /dev/null +++ b/src/stokes/Stokes2D_old.jl @@ -0,0 +1,808 @@ + +## 2D STOKES MODULE + +module Stokes2D + +using ImplicitGlobalGrid +using ..JustRelax +using CUDA, AMDGPU +using ParallelStencil +# using ParallelStencil.FiniteDifferences2D +using GeoParams, LinearAlgebra, Printf + +import JustRelax: elastic_iter_params!, PTArray, Velocity, SymmetricTensor +import JustRelax: + Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG +import JustRelax: compute_maxloc!, solve! +import JustRelax: mean_mpi, norm_mpi, maximum_mpi, minimum_mpi, backend + +@eval @init_parallel_stencil($backend, Float64, 2) + +include("../rheology/GeoParams.jl") +include("StressRotation.jl") +include("StressKernels.jl") +include("PressureKernels.jl") +include("VelocityKernels.jl") +include("StressKernels.jl") + +export solve! + +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 + τxx_o, τyy_o, τxy_o, τxy_o_c = stokes.τ_o.xx, + stokes.τ_o.yy, stokes.τ_o.xy, + stokes.τ_o.xy_c + + @parallel (@idx size(τxy)) multi_copy!( + (τxx_o, τyy_o, τxy_o, τxy_o_c), (τxx, τyy, τxy, τxy_c) + ) + return nothing +end + +## 2D VISCO-ELASTIC STOKES SOLVER + +# viscous solver +function JustRelax.solve!( + stokes::StokesArrays{Viscous,A,B,C,D,2}, + pt_stokes::PTStokesCoeffs, + di::NTuple{2,T}, + flow_bcs::FlowBoundaryConditions, + ρg, + η, + K, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 1), + verbose=true, +) where {A,B,C,D,T} + + # unpack + _dx, _dy = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + @parallel compute_P!( + stokes.P, stokes.P0, stokes.RP, stokes.∇V, η, K, dt, r, θ_dτ + ) + @parallel (@idx ni .+ 1) compute_τ!( + @stress(stokes)..., @strain(stokes)..., η, θ_dτ + ) + + @hide_communication b_width begin + @parallel compute_V!( + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + ητ, + _di..., + dt, + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(@velocity(stokes)...) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, + stokes.R.Ry, + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + ρg..., + _di..., + dt, + ) + Vmin, Vmax = extrema(stokes.V.Vx) + Pmin, Pmax = extrema(stokes.P) + push!( + norm_Rx, + norm_mpi(stokes.R.Rx) / (Pmax - Pmin) * lx / sqrt(length(stokes.R.Rx)), + ) + push!( + norm_Ry, + norm_mpi(stokes.R.Ry) / (Pmax - Pmin) * lx / sqrt(length(stokes.R.Ry)), + ) + push!( + norm_∇V, norm_mpi(stokes.∇V) / (Vmax - Vmin) * lx / sqrt(length(stokes.∇V)) + ) + err = maximum_mpi(norm_Rx[end], norm_Ry[end], norm_∇V[end]) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end + +# visco-elastic solver +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, + pt_stokes::PTStokesCoeffs, + di::NTuple{2,T}, + flow_bcs, + ρg, + η, + G, + K, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 1), + verbose=true, +) where {A,B,C,D,T} + + # unpack + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + @parallel (@idx ni) compute_∇V!(stokes.∇V, stokes.V.Vx, stokes.V.Vy, _di...) + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + @parallel compute_P!( + stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, η, K, dt, r, θ_dτ + ) + @parallel (@idx ni) compute_τ!( + @stress(stokes)..., + @tensor(stokes.τ_o)..., + @strain(stokes)..., + η, + G, + θ_dτ, + dt, + ) + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + ητ, + _di..., + ) + # free slip boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, stokes.R.Ry, stokes.P, @stress(stokes)..., ρg..., _di... + ) + errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum_mpi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end + +# GeoParams: general (visco-elasto-plastic) solver + +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, + pt_stokes::PTStokesCoeffs, + di::NTuple{2,T}, + flow_bcs, + ρg, + η, + η_vep, + rheology::MaterialParams, + args, + dt, + igg::IGG; + viscosity_cutoff=(1e16, 1e24), + iterMax=10e3, + nout=500, + b_width=(4, 4, 0), + verbose=true, +) where {A,B,C,D,T} + + # unpack + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + Kb = get_Kb(rheology) + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + + # solver loop + wtime0 = 0.0 + λ = @zeros(ni...) + θ = @zeros(ni...) + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + @parallel compute_P!( + stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, η, Kb, dt, r, θ_dτ + ) + + @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) + + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + ν = 1e-2 + @parallel (@idx ni) compute_viscosity!( + η, ν, @strain(stokes)..., args, rheology, viscosity_cutoff + ) + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + + @parallel (@idx ni) compute_τ_nonlinear!( + @tensor_center(stokes.τ), + stokes.τ.II, + @tensor(stokes.τ_o), + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + stokes.P, + θ, + η, + η_vep, + λ, + tupleize(rheology), # needs to be a tuple + dt, + θ_dτ, + ) + + @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) + update_halo!(stokes.τ.xy) + + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + θ, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + ητ, + _di..., + dt, + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, + stokes.R.Ry, + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + ρg..., + _di..., + dt, + ) + + errs = maximum.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + stokes.P .= θ + + # accumulate plastic strain tensor + @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end + +## With phase ratios + +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, + pt_stokes::PTStokesCoeffs, + di::NTuple{2,T}, + flow_bcs, + ρg, + η, + η_vep, + phase_ratios::PhaseRatio, + rheology, + args, + dt, + igg::IGG; + viscosity_cutoff=(1e16, 1e24), + iterMax=50e3, + iterMin=1e2, + viscosity_relaxation=1e-2, + free_surface=false, + nout=500, + b_width=(4, 4, 0), + verbose=true, +) where {A,B,C,D,T} + + # unpack + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + sizehint!(norm_Rx, Int(iterMax)) + sizehint!(norm_Ry, Int(iterMax)) + sizehint!(norm_∇V, Int(iterMax)) + sizehint!(err_evo1, Int(iterMax)) + sizehint!(err_evo2, Int(iterMax)) + + # solver loop + @copy stokes.P0 stokes.P + wtime0 = 0.0 + θ = @zeros(ni...) + λ = @zeros(ni...) + η0 = deepcopy(η) + do_visc = true + + for Aij in @tensor_center(stokes.ε_pl) + Aij .= 0.0 + end + Vx_on_Vy = @zeros(size(stokes.V.Vy)) + + while iter ≤ iterMax + iterMin < iter && err < ϵ && break + + wtime0 += @elapsed begin + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + + @parallel (@idx ni) compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + dt, + r, + θ_dτ, + ) + + if rem(iter, 5) == 0 + @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + end + + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + + if rem(iter, nout) == 0 + @copy η0 η + end + if do_visc + @parallel (@idx ni) compute_viscosity!( + η, + viscosity_relaxation, + phase_ratios.center, + @strain(stokes)..., + args, + rheology, + viscosity_cutoff, + ) + end + + @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, + θ, + η, + η_vep, + λ, + phase_ratios.center, + tupleize(rheology), # needs to be a tuple + dt, + θ_dτ, + ) + + @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) + update_halo!(stokes.τ.xy) + + @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( + Vx_on_Vy, stokes.V.Vx + ) + + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + Vx_on_Vy, + θ, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + ητ, + _di..., + dt * free_surface, + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy) + end + end + + iter += 1 + + if iter % nout == 0 && iter > 1 + er_η = norm_mpi(@.(log10(η) - log10(η0))) + er_η < 1e-3 && (do_visc = false) + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, + stokes.R.Ry, + @velocity(stokes)..., + Vx_on_Vy, + stokes.P, + @stress(stokes)..., + ρg..., + _di..., + dt * free_surface, + ) + # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + errs = ( + norm_mpi(stokes.R.Rx) / length(stokes.R.Rx), + norm_mpi(stokes.R.Ry) / length(stokes.R.Ry), + norm_mpi(stokes.R.RP) / length(stokes.R.RP), + ) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum_mpi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + + if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ && iter ≥ 20000 + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + stokes.P .= θ + + # accumulate plastic strain tensor + @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end + +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, + thermal::ThermalArrays, + pt_stokes::PTStokesCoeffs, + di::NTuple{2,T}, + flow_bcs, + ϕ, + ρg, + η, + η_vep, + phase_v, + phase_c, + args_η, + rheology::NTuple{N,MaterialParams}, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 1), + verbose=true, +) where {A,B,C,D,N,T} + + # unpack + + _di = inv.(di) + (; ϵ, r, θ_dτ, ηdτ) = pt_stokes + ni = size(stokes.P) + # ~preconditioner + ητ = deepcopy(η) + # @hide_communication b_width begin # communication/computation overlap + compute_maxloc!(ητ, η; window=(1, 1)) + update_halo!(ητ) + # end + + # errors + err = 2 * ϵ + iter = 0 + err_evo1 = Float64[] + err_evo2 = Float64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_∇V = Float64[] + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + @parallel (@idx ni) compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + η, + rheology, + phase_c, + dt, + r, + θ_dτ, + ) + @parallel (@idx ni .+ 1) compute_strain_rate!( + @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... + ) + @parallel (@idx ni) compute_ρg!( + ρg[end], ϕ, rheology, (T=thermal.Tc, P=stokes.P) + ) + @parallel (@idx ni) compute_τ_gp!( + @tensor_center(stokes.τ), + stokes.τ.II, + @tensor(stokes.τ_o), + @strain(stokes), + η, + η_vep, + thermal.T, + phase_v, + phase_c, + args_η, + rheology, # needs to be a tuple + dt, + θ_dτ, + ) + @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + pt_stokes.ηdτ, + ρg..., + ητ, + _di..., + dt, + ) + # apply boundary conditions boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + @parallel (@idx ni) compute_Res!( + stokes.R.Rx, + stokes.R.Ry, + @velocity(stokes)..., + stokes.P, + @stress(stokes)..., + ρg..., + _di..., + dt, + ) + errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) + push!(norm_Rx, errs[1]) + push!(norm_Ry, errs[2]) + push!(norm_∇V, errs[3]) + err = maximum_mpi(errs) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && (verbose || iter == iterMax) + @printf( + "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[end], + norm_Ry[end], + norm_∇V[end] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_∇V=norm_∇V, + ) +end + +end # END OF MODULE diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 2495b47d..fbaba8f7 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -1,33 +1,3 @@ -## 3D ELASTICITY MODULE - -module Stokes3D - -using ImplicitGlobalGrid -using ParallelStencil -using ParallelStencil.FiniteDifferences3D -using JustRelax -using CUDA, AMDGPU -using LinearAlgebra -using Printf -using GeoParams - -import JustRelax: PTArray, Velocity, SymmetricTensor, pureshear_bc! -import JustRelax: - Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG -import JustRelax: compute_maxloc!, solve! -import JustRelax: mean_mpi, norm_mpi, minimum_mpi, maximum_mpi, backend - -@eval @init_parallel_stencil($backend, Float64, 3) - -include("../rheology/GeoParams.jl") -include("StressRotation.jl") -include("StressKernels.jl") -include("PressureKernels.jl") -include("VelocityKernels.jl") -include("StressKernels.jl") - -export solve!, pureshear_bc! - @parallel function update_τ_o!( τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, τxx, τyy, τzz, τxy, τxz, τyz ) @@ -40,39 +10,20 @@ export solve!, pureshear_bc! return nothing end -function update_τ_o!(stokes::StokesArrays{ViscoElastic,A,B,C,D,3}) where {A,B,C,D} +function update_τ_o!(stokes::StokesArrays) @parallel update_τ_o!(@tensor(stokes.τ_o)..., @stress(stokes)...) end -## BOUNDARY CONDITIONS - -function JustRelax.pureshear_bc!( - stokes::StokesArrays, di::NTuple{3,T}, li::NTuple{3,T}, εbg -) where {T} - # unpack - Vx, _, Vz = stokes.V.Vx, stokes.V.Vy, stokes.V.Vz - dx, _, dz = di - lx, _, lz = li - # Velocity pure shear boundary conditions - stokes.V.Vx .= PTArray([ - -εbg * ((i - 1) * dx - 0.5 * lx) for i in 1:size(Vx, 1), j in 1:size(Vx, 2), - k in 1:size(Vx, 3) - ]) - return stokes.V.Vz .= PTArray([ - εbg * ((k - 1) * dz - 0.5 * lz) for i in 1:size(Vz, 1), j in 1:size(Vz, 2), - k in 1:size(Vz, 3) - ]) -end - ## 3D VISCO-ELASTIC STOKES SOLVER +solve!(stokes::StokesArrays, args...) = solve!(CPUBackendTrait(stokes), stokes, args...) -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bcs, ρg, - η, K, G, dt, @@ -81,13 +32,14 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 4), verbose=true, -) where {A,B,C,D,T} +) where {T} # solver related ϵ = pt_stokes.ϵ # geometry _di = @. 1 / di ni = size(stokes.P) + (; η) = stokes.viscosity # ~preconditioner ητ = deepcopy(η) @@ -199,14 +151,13 @@ end ## 3D VISCO-ELASTO-PLASTIC STOKES SOLVER WITH GeoParams.jl -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bcs::FlowBoundaryConditions, ρg, - η, - η_vep, rheology::MaterialParams, args, dt, @@ -215,13 +166,14 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 4), verbose=true, -) where {A,B,C,D,T} +) where {T} # solver related ϵ = pt_stokes.ϵ # geometry _di = @. 1 / di ni = size(stokes.P) + (; η, η_vep) = stokes.viscosity # ~preconditioner ητ = deepcopy(η) @@ -372,14 +324,13 @@ function JustRelax.solve!( end # GeoParams and multiple phases -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, +function solve!( + ::CPUBackendTrait, + stokes::StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bc::FlowBoundaryConditions, ρg, - η, - η_vep, phase_ratios::PhaseRatio, rheology::NTuple{N,AbstractMaterialParamsStruct}, args, @@ -390,7 +341,7 @@ function JustRelax.solve!( b_width=(4, 4, 4), verbose=true, viscosity_cutoff=(-Inf, Inf), -) where {A,B,C,D,T,N} +) where {T,N} ## UNPACK @@ -399,6 +350,7 @@ function JustRelax.solve!( # geometry _di = @. 1 / di ni = size(stokes.P) + (; η, η_vep) = stokes.viscosity # errors err = Inf @@ -557,5 +509,3 @@ function JustRelax.solve!( av_time=av_time, ) end - -end # END OF MODULE diff --git a/src/stokes/Stokes3D_old.jl b/src/stokes/Stokes3D_old.jl new file mode 100644 index 00000000..2495b47d --- /dev/null +++ b/src/stokes/Stokes3D_old.jl @@ -0,0 +1,561 @@ +## 3D ELASTICITY MODULE + +module Stokes3D + +using ImplicitGlobalGrid +using ParallelStencil +using ParallelStencil.FiniteDifferences3D +using JustRelax +using CUDA, AMDGPU +using LinearAlgebra +using Printf +using GeoParams + +import JustRelax: PTArray, Velocity, SymmetricTensor, pureshear_bc! +import JustRelax: + Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG +import JustRelax: compute_maxloc!, solve! +import JustRelax: mean_mpi, norm_mpi, minimum_mpi, maximum_mpi, backend + +@eval @init_parallel_stencil($backend, Float64, 3) + +include("../rheology/GeoParams.jl") +include("StressRotation.jl") +include("StressKernels.jl") +include("PressureKernels.jl") +include("VelocityKernels.jl") +include("StressKernels.jl") + +export solve!, pureshear_bc! + +@parallel function update_τ_o!( + τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, τxx, τyy, τzz, τxy, τxz, τyz +) + @all(τxx_o) = @all(τxx) + @all(τyy_o) = @all(τyy) + @all(τzz_o) = @all(τzz) + @all(τxy_o) = @all(τxy) + @all(τxz_o) = @all(τxz) + @all(τyz_o) = @all(τyz) + return nothing +end + +function update_τ_o!(stokes::StokesArrays{ViscoElastic,A,B,C,D,3}) where {A,B,C,D} + @parallel update_τ_o!(@tensor(stokes.τ_o)..., @stress(stokes)...) +end + +## BOUNDARY CONDITIONS + +function JustRelax.pureshear_bc!( + stokes::StokesArrays, di::NTuple{3,T}, li::NTuple{3,T}, εbg +) where {T} + # unpack + Vx, _, Vz = stokes.V.Vx, stokes.V.Vy, stokes.V.Vz + dx, _, dz = di + lx, _, lz = li + # Velocity pure shear boundary conditions + stokes.V.Vx .= PTArray([ + -εbg * ((i - 1) * dx - 0.5 * lx) for i in 1:size(Vx, 1), j in 1:size(Vx, 2), + k in 1:size(Vx, 3) + ]) + return stokes.V.Vz .= PTArray([ + εbg * ((k - 1) * dz - 0.5 * lz) for i in 1:size(Vz, 1), j in 1:size(Vz, 2), + k in 1:size(Vz, 3) + ]) +end + +## 3D VISCO-ELASTIC STOKES SOLVER + +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, + pt_stokes::PTStokesCoeffs, + di::NTuple{3,T}, + flow_bcs, + ρg, + η, + K, + G, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 4), + verbose=true, +) where {A,B,C,D,T} + + # solver related + ϵ = pt_stokes.ϵ + # geometry + _di = @. 1 / di + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + compute_maxloc!(ητ, η) + update_halo!(ητ) + + # errors + err = 2 * ϵ + iter = 0 + cont = 0 + err_evo1 = Float64[] + err_evo2 = Int64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_Rz = Float64[] + norm_∇V = Float64[] + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + @parallel compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + η, + K, + dt, + pt_stokes.r, + pt_stokes.θ_dτ, + ) + @parallel (@idx ni .+ 1) compute_strain_rate!( + stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... + ) + @parallel (@idx ni .+ 1) compute_τ!( + @stress(stokes)..., + @tensor(stokes.τ_o)..., + @strain(stokes)..., + η, + G, + dt, + pt_stokes.θ_dτ, + ) + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.R.Rz, + stokes.P, + ρg..., + @stress(stokes)..., + ητ, + pt_stokes.ηdτ, + _di..., + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + cont += 1 + push!(norm_Rx, maximum_mpi(abs.(stokes.R.Rx))) + push!(norm_Ry, maximum_mpi(abs.(stokes.R.Ry))) + push!(norm_Rz, maximum_mpi(abs.(stokes.R.Rz))) + push!(norm_∇V, maximum_mpi(abs.(stokes.R.RP))) + err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) + @printf( + "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[cont], + norm_Ry[cont], + norm_Rz[cont], + norm_∇V[cont] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + av_time = wtime0 / (iter - 1) # average time per iteration + update_τ_o!(stokes) # copy τ into τ_o + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_Rz=norm_Rz, + norm_∇V=norm_∇V, + time=wtime0, + av_time=av_time, + ) +end + +## 3D VISCO-ELASTO-PLASTIC STOKES SOLVER WITH GeoParams.jl + +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, + pt_stokes::PTStokesCoeffs, + di::NTuple{3,T}, + flow_bcs::FlowBoundaryConditions, + ρg, + η, + η_vep, + rheology::MaterialParams, + args, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 4), + verbose=true, +) where {A,B,C,D,T} + + # solver related + ϵ = pt_stokes.ϵ + # geometry + _di = @. 1 / di + ni = size(stokes.P) + + # ~preconditioner + ητ = deepcopy(η) + + # errors + err = 2 * ϵ + iter = 0 + cont = 0 + err_evo1 = Float64[] + err_evo2 = Int64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_Rz = Float64[] + norm_∇V = Float64[] + + Kb = get_Kb(rheology) + G = get_shear_modulus(rheology) + @copy stokes.P0 stokes.P + λ = @zeros(ni...) + θ = @zeros(ni...) + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + compute_maxloc!(ητ, η) + update_halo!(ητ) + + @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) + @parallel (@idx ni) compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + Kb, + dt, + pt_stokes.r, + pt_stokes.θ_dτ, + ) + @parallel (@idx ni) compute_strain_rate!( + stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... + ) + + # # Update buoyancy + # @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) + + ν = 1e-3 + @parallel (@idx ni) compute_viscosity!( + η, + 1.0, + phase_ratios.center, + @strain(stokes)..., + args, + rheology, + viscosity_cutoff, + ) + + @parallel (@idx ni) compute_τ_nonlinear!( + @tensor_center(stokes.τ), + stokes.τ.II, + @tensor(stokes.τ_o), + @strain(stokes), + @tensor_center(stokes.ε_pl), + stokes.EII_pl, + stokes.P, + θ, + η, + @ones(ni...), + λ, + tupleize(rheology), # needs to be a tuple + dt, + pt_stokes.θ_dτ, + ) + + @parallel (@idx ni .+ 1) compute_τ_vertex!( + @shear(stokes.τ)..., + @shear(stokes.τ_o)..., + @shear(stokes.ε)..., + η_vep, + G, + dt, + pt_stokes.θ_dτ, + ) + + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + stokes.R.Rx, + stokes.R.Ry, + stokes.R.Rz, + stokes.P, + ρg..., + @stress(stokes)..., + ητ, + pt_stokes.ηdτ, + _di..., + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bcs) + update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + end + end + + iter += 1 + if iter % nout == 0 && iter > 1 + cont += 1 + push!(norm_Rx, maximum_mpi(abs.(stokes.R.Rx))) + push!(norm_Ry, maximum_mpi(abs.(stokes.R.Ry))) + push!(norm_Rz, maximum_mpi(abs.(stokes.R.Rz))) + push!(norm_∇V, maximum_mpi(abs.(stokes.R.RP))) + err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && (verbose || iter == iterMax) + @printf( + "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[cont], + norm_Ry[cont], + norm_Rz[cont], + norm_∇V[cont] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + av_time = wtime0 / (iter - 1) # average time per iteration + update_τ_o!(stokes) # copy τ into τ_o + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_Rz=norm_Rz, + norm_∇V=norm_∇V, + time=wtime0, + av_time=av_time, + ) +end + +# GeoParams and multiple phases +function JustRelax.solve!( + stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, + pt_stokes::PTStokesCoeffs, + di::NTuple{3,T}, + flow_bc::FlowBoundaryConditions, + ρg, + η, + η_vep, + phase_ratios::PhaseRatio, + rheology::NTuple{N,AbstractMaterialParamsStruct}, + args, + dt, + igg::IGG; + iterMax=10e3, + nout=500, + b_width=(4, 4, 4), + verbose=true, + viscosity_cutoff=(-Inf, Inf), +) where {A,B,C,D,T,N} + + ## UNPACK + + # solver related + ϵ = pt_stokes.ϵ + # geometry + _di = @. 1 / di + ni = size(stokes.P) + + # errors + err = Inf + iter = 0 + cont = 0 + err_evo1 = Float64[] + err_evo2 = Int64[] + norm_Rx = Float64[] + norm_Ry = Float64[] + norm_Rz = Float64[] + norm_∇V = Float64[] + + @copy stokes.P0 stokes.P + λ = @zeros(ni...) + θ = @zeros(ni...) + + # solver loop + wtime0 = 0.0 + while iter < 2 || (err > ϵ && iter ≤ iterMax) + wtime0 += @elapsed begin + # ~preconditioner + ητ = deepcopy(η) + compute_maxloc!(ητ, η) + 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...) + @parallel (@idx ni) compute_P!( + stokes.P, + stokes.P0, + stokes.R.RP, + stokes.∇V, + ητ, + rheology, + phase_ratios.center, + dt, + pt_stokes.r, + pt_stokes.θ_dτ, + ) + + @parallel (@idx ni) compute_strain_rate!( + stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... + ) + + # Update buoyancy + @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + + # Update viscosity + ν = 1e-2 + @parallel (@idx ni) compute_viscosity!( + η, + ν, + phase_ratios.center, + @strain(stokes)..., + args, + rheology, + viscosity_cutoff, + ) + + @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, + θ, + η, + η_vep, + λ, + phase_ratios.center, + tupleize(rheology), # needs to be a tuple + dt, + pt_stokes.θ_dτ, + ) + + @parallel (@idx ni .+ 1) center2vertex!( + stokes.τ.yz, + stokes.τ.xz, + stokes.τ.xy, + stokes.τ.yz_c, + stokes.τ.xz_c, + stokes.τ.xy_c, + ) + update_halo!(stokes.τ.yz, stokes.τ.xz, stokes.τ.xy) + + # @parallel (@idx ni .+ 1) compute_τ_vertex!( + # @shear(stokes.τ)..., @shear(stokes.ε)..., η_vep, pt_stokes.θ_dτ + # ) + + @hide_communication b_width begin # communication/computation overlap + @parallel compute_V!( + @velocity(stokes)..., + @residuals(stokes.R)..., + θ, + ρg..., + @stress(stokes)..., + ητ, + pt_stokes.ηdτ, + _di..., + ) + # apply boundary conditions + flow_bcs!(stokes, flow_bc) + update_halo!(@velocity(stokes)...) + end + end + + stokes.P .= θ + + iter += 1 + if iter % nout == 0 && iter > 1 + cont += 1 + for (norm_Ri, Ri) in zip((norm_Rx, norm_Ry, norm_Rz), @residuals(stokes.R)) + push!(norm_Ri, maximum(abs.(Ri))) + end + push!(norm_∇V, maximum(abs.(stokes.R.RP))) + err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) + push!(err_evo1, err) + push!(err_evo2, iter) + if igg.me == 0 && (verbose || iter == iterMax) + @printf( + "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", + iter, + err, + norm_Rx[cont], + norm_Ry[cont], + norm_Rz[cont], + norm_∇V[cont] + ) + end + isnan(err) && error("NaN(s)") + end + + if igg.me == 0 && err ≤ ϵ + println("Pseudo-transient iterations converged in $iter iterations") + end + end + + av_time = wtime0 / (iter - 1) # average time per iteration + + stokes.P .= θ + + return ( + iter=iter, + err_evo1=err_evo1, + err_evo2=err_evo2, + norm_Rx=norm_Rx, + norm_Ry=norm_Ry, + norm_Rz=norm_Rz, + norm_∇V=norm_∇V, + time=wtime0, + av_time=av_time, + ) +end + +end # END OF MODULE diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index de79d5fc..58f01b07 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -70,25 +70,8 @@ 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, + τxx::AbstractArray{T,3}, τyy, τzz, τyz, @@ -110,7 +93,7 @@ end G, dt, θ_dτ, -) +) where {T} 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) @@ -197,8 +180,6 @@ end η_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 @@ -206,50 +187,12 @@ end η_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(η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 - end - return nothing -end - -@parallel_indices (i, j, k) function compute_τ_vertex!( - τyz, τxz, τxy, εyz, εxz, εxy, η, θ_dτ -) - I = i, j, k - harm_xy(A) = _harm_xyi(A, I...) - harm_xz(A) = _harm_xzi(A, I...) - harm_yz(A) = _harm_yzi(A, I...) - av_xy(A) = _av_xyi(A, I...) - av_xz(A) = _av_xzi(A, I...) - av_yz(A) = _av_yzi(A, I...) - - @inbounds begin - # Compute τ_xy - if (1 < i < size(τxy, 1)) && (1 < j < size(τxy, 2)) && k ≤ size(τxy, 3) - η_ij = av_xy(η) - denominator = inv(θ_dτ + 1.0) - τxy[I...] += (-τxy[I...] + 2.0 * η_ij * εxy[I...]) * denominator - end - # Compute τ_xz - if (1 < i < size(τxz, 1)) && j ≤ size(τxz, 2) && (1 < k < size(τxz, 3)) - η_ij = av_xz(η) - denominator = inv(θ_dτ + 1.0) - τxz[I...] += (-τxz[I...] + 2.0 * η_ij * εxz[I...]) * denominator - end - # Compute τ_yz - if i ≤ size(τyz, 1) && (1 < j < size(τyz, 2)) && (1 < k < size(τyz, 3)) - η_ij = av_yz(η) - denominator = inv(θ_dτ + 1.0) - τyz[I...] += (-τyz[I...] + 2.0 * η_ij * εyz[I...]) * denominator end end return nothing @@ -354,19 +297,33 @@ end return nothing end -@parallel_indices (I...) function tensor_invariant!(II, xx, yy, xyv) +""" + tensor_invariant!(A::SymmetricTensor) + +Compute the tensor invariant of the given symmetric tensor `A`. + +# Arguments +- `A::SymmetricTensor`: The input symmetric tensor. +""" +function tensor_invariant!(A::SymmetricTensor) + ni = size(A.II) + @parallel (@idx ni) tensor_invariant_kernel!(A.II, @tensor(A)...) + return nothing +end + +@parallel_indices (I...) function tensor_invariant_kernel!(II, xx, yy, xy) # convenience closure @inline gather(A) = _gather(A, I...) @inbounds begin - τ = xx[I...], yy[I...], gather(xyv) + τ = xx[I...], yy[I...], gather(xy) II[I...] = second_invariant_staggered(τ...) end return nothing end -@parallel_indices (I...) function tensor_invariant!(II, xx, yy, zz, yz, xz, xy) +@parallel_indices (I...) function tensor_invariant_kernel!(II, xx, yy, zz, yz, xz, xy) # convenience closures @inline gather_yz(A) = _gather_yz(A, I...) diff --git a/src/thermal_diffusion/DiffusionPT.jl b/src/thermal_diffusion/DiffusionPT.jl index 2101f7e4..cdb22e20 100644 --- a/src/thermal_diffusion/DiffusionPT.jl +++ b/src/thermal_diffusion/DiffusionPT.jl @@ -1,814 +1,5 @@ -## GeoParams - -# include("Rheology.jl") -include("../rheology/GeoParams.jl") - -@inline get_phase(x::PhaseRatio) = x.center -@inline get_phase(x) = x - -update_pt_thermal_arrays!(::Vararg{Any,N}) where {N} = nothing - -function update_pt_thermal_arrays!( - pt_thermal, phase_ratios::PhaseRatio, rheology, args, _dt -) - ni = size(phase_ratios.center) - - @parallel (@idx ni) compute_pt_thermal_arrays!( - pt_thermal.θr_dτ, - pt_thermal.dτ_ρ, - rheology, - phase_ratios.center, - args, - pt_thermal.max_lxyz, - pt_thermal.Vpdτ, - _dt, - ) - - return nothing -end - -@inline function compute_phase(fn::F, rheology, phase::Int, args) where {F} - return fn(rheology, phase, args) -end - -@inline function compute_phase(fn::F, rheology, phase::SVector, args) where {F} - return fn_ratio(fn, rheology, phase, args) -end - -@inline compute_phase(fn::F, rheology, ::Nothing, args) where {F} = fn(rheology, args) - -@inline Base.@propagate_inbounds function getindex_phase( - phase::AbstractArray, I::Vararg{Int,N} -) where {N} - return phase[I...] -end -@inline getindex_phase(::Nothing, I::Vararg{Int,N}) where {N} = nothing - -@inline function compute_ρCp(rheology, args) - return compute_heatcapacity(rheology, args) * compute_density(rheology, args) -end - -@inline function compute_ρCp(rheology, phase::Union{Nothing,Int}, args) - return compute_phase(compute_heatcapacity, rheology, phase, args) * - compute_phase(compute_density, rheology, phase, args) -end - -@inline function compute_ρCp(rheology, ρ, args) - return compute_heatcapacity(rheology, args) * ρ -end - -@inline function compute_ρCp(rheology, ρ, phase::Union{Nothing,Int}, args) - return compute_phase(compute_heatcapacity, rheology, phase, args) * ρ -end - -## 3D KERNELS - -@parallel_indices (i, j, k) function compute_flux!( - qTx::AbstractArray{_T,3}, qTy, qTz, qTx2, qTy2, qTz2, T, K, θr_dτ, _dx, _dy, _dz -) where {_T} - d_xi(A) = _d_xi(A, i, j, k, _dx) - d_yi(A) = _d_yi(A, i, j, k, _dy) - d_zi(A) = _d_zi(A, i, j, k, _dz) - av_xy(A) = _av_xy(A, i, j, k) - av_xz(A) = _av_xz(A, i, j, k) - av_yz(A) = _av_yz(A, i, j, k) - - I = i, j, k - - if all(I .≤ size(qTx)) - qx = qTx2[I...] = -av_yz(K) * d_xi(T) - qTx[I...] = (qTx[I...] * av_yz(θr_dτ) + qx) / (1.0 + av_yz(θr_dτ)) - end - - if all(I .≤ size(qTy)) - qy = qTy2[I...] = -av_xz(K) * d_yi(T) - qTy[I...] = (qTy[I...] * av_xz(θr_dτ) + qy) / (1.0 + av_xz(θr_dτ)) - end - - if all(I .≤ size(qTz)) - qz = qTz2[I...] = -av_xy(K) * d_zi(T) - qTz[I...] = (qTz[I...] * av_xy(θr_dτ) + qz) / (1.0 + av_xy(θr_dτ)) - end - - return nothing -end - -@parallel_indices (i, j, k) function compute_flux!( - qTx::AbstractArray{_T,3}, - qTy, - qTz, - qTx2, - qTy2, - qTz2, - T, - rheology, - phase, - θr_dτ, - _dx, - _dy, - _dz, - args, -) where {_T} - d_xi(A) = _d_xi(A, i, j, k, _dx) - d_yi(A) = _d_yi(A, i, j, k, _dy) - d_zi(A) = _d_zi(A, i, j, k, _dz) - av_xy(A) = _av_xy(A, i, j, k) - av_xz(A) = _av_xz(A, i, j, k) - av_yz(A) = _av_yz(A, i, j, k) - - get_K(idx, args) = compute_phase(compute_conductivity, rheology, idx, args) - - I = i, j, k - - @inbounds if all(I .≤ size(qTx)) - T_ijk = (T[(I .+ 1)...] + T[i, j + 1, k + 1]) * 0.5 - args_ijk = (; T=T_ijk, P=av_yz(args.P)) - K = - ( - get_K(getindex_phase(phase, i, j + 1, k), args_ijk) + - get_K(getindex_phase(phase, i, j, k + 1), args_ijk) + - get_K(getindex_phase(phase, i, j + 1, k + 1), args_ijk) + - get_K(getindex_phase(phase, i, j, k), args_ijk) - ) * 0.25 - - qx = qTx2[I...] = -K * d_xi(T) - qTx[I...] = (qTx[I...] * av_yz(θr_dτ) + qx) / (1.0 + av_yz(θr_dτ)) - end - - @inbounds if all(I .≤ size(qTy)) - T_ijk = (T[(I .+ 1)...] + T[i + 1, j, k + 1]) * 0.5 - args_ijk = (; T=T_ijk, P=av_xz(args.P)) - K = - ( - get_K(getindex_phase(phase, i + 1, j, k), args_ijk) + - get_K(getindex_phase(phase, i, j, k + 1), args_ijk) + - get_K(getindex_phase(phase, i + 1, j, k + 1), args_ijk) + - get_K(getindex_phase(phase, i, j, k), args_ijk) - ) * 0.25 - - qy = qTy2[I...] = -K * d_yi(T) - qTy[I...] = (qTy[I...] * av_xz(θr_dτ) + qy) / (1.0 + av_xz(θr_dτ)) - end - - @inbounds if all(I .≤ size(qTz)) - T_ijk = (T[(I .+ 1)...] + T[i + 1, j + 1, k]) * 0.5 - args_ijk = (; T=T_ijk, P=av_xy(args.P)) - K = - ( - get_K(getindex_phase(phase, i + 1, j + 1, k), args_ijk) + - get_K(getindex_phase(phase, i, j + 1, k), args_ijk) + - get_K(getindex_phase(phase, i + 1, j, k), args_ijk) + - get_K(getindex_phase(phase, i, j, k), args_ijk) - ) * 0.25 - - qz = qTz2[I...] = -K * d_zi(T) - qTz[I...] = (qTz[I...] * av_xy(θr_dτ) + qz) / (1.0 + av_xy(θr_dτ)) - end - - return nothing -end - -@parallel_indices (I...) function update_T!( - T::AbstractArray{_T,3}, - Told, - qTx, - qTy, - qTz, - H, - shear_heating, - ρCp, - dτ_ρ, - _dt, - _dx, - _dy, - _dz, -) where {_T} - av(A) = _av(A, I...) - d_xa(A) = _d_xa(A, I..., _dx) - d_ya(A) = _d_ya(A, I..., _dy) - d_za(A) = _d_za(A, I..., _dz) - - I1 = I .+ 1 - T[I1...] += - av(dτ_ρ) * ( - (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) - - av(ρCp) * (T[I1...] - Told[I1...]) * _dt - ) + - av(H) + - av(shear_heating) - - return nothing -end - -@parallel_indices (i, j, k) function update_T!( - T::AbstractArray{_T,3}, - Told, - qTx, - qTy, - qTz, - H, - shear_heating, - rheology, - phase, - dτ_ρ, - _dt, - _dx, - _dy, - _dz, - args, -) where {_T} - av(A) = _av(A, i, j, k) - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) - - I = i + 1, j + 1, k + 1 - - T_ijk = T[I...] - args_ijk = (; T=T_ijk, P=av(args.P)) - phase_ijk = getindex_phase(phase, i, j, k) - - T[I...] = - T_ijk + - av(dτ_ρ) * ( - (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) - - compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt + - av(H) + - av(shear_heating) - ) - return nothing -end - -@parallel_indices (i, j, k) function check_res!( - ResT::AbstractArray{_T,3}, - T, - Told, - qTx2, - qTy2, - qTz2, - H, - shear_heating, - ρCp, - _dt, - _dx, - _dy, - _dz, -) where {_T} - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) - av(A) = _av(A, i, j, k) - - I = i + 1, j + 1, k + 1 - - ResT[i, j, k] = - -av(ρCp) * (T[I...] - Told[I...]) * _dt - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + - av(H) + - av(shear_heating) - - return nothing -end - -@parallel_indices (i, j, k) function check_res!( - ResT::AbstractArray{_T,3}, - T, - Told, - qTx2, - qTy2, - qTz2, - H, - shear_heating, - rheology, - phase, - _dt, - _dx, - _dy, - _dz, - args, -) where {_T} - d_xa(A) = _d_xa(A, i, j, k, _dx) - d_ya(A) = _d_ya(A, i, j, k, _dy) - d_za(A) = _d_za(A, i, j, k, _dz) - av(A) = _av(A, i, j, k) - - I = i + 1, j + 1, k + 1 - T_ijk = T[I...] - args_ijk = (; T=T_ijk, P=av(args.P)) - phase_ijk = getindex_phase(phase, i, j, k) - - ResT[i, j, k] = - -compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt - - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + - av(H) + - av(shear_heating) - - return nothing -end - -## 2D KERNELS - -@parallel_indices (i, j) function compute_flux!( - qTx::AbstractArray{_T,2}, qTy, qTx2, qTy2, T, K, θr_dτ, _dx, _dy -) where {_T} - nx = size(θr_dτ, 1) - - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 - av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 - - @inbounds if all((i, j) .≤ size(qTx)) - qx = qTx2[i, j] = -av_xa(K) * d_xi(T) - qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) - end - - @inbounds if all((i, j) .≤ size(qTy)) - qy = qTy2[i, j] = -av_ya(K) * d_yi(T) - qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) - end - return nothing -end - -@parallel_indices (i, j) function compute_flux!( - qTx::AbstractArray{_T,2}, qTy, qTx2, qTy2, T, rheology, phase, θr_dτ, _dx, _dy, args -) where {_T} - nx = size(θr_dτ, 1) - - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 - av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 - - if all((i, j) .≤ size(qTx)) - ii, jj = clamp(i - 1, 1, nx), j + 1 - phase_ij = getindex_phase(phase, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K1 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) - - ii, jj = clamp(i - 1, 1, nx), j - phase_ij = getindex_phase(phase, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K2 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) - K = (K1 + K2) * 0.5 - - qx = qTx2[i, j] = -K * d_xi(T) - qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) - end - - if all((i, j) .≤ size(qTy)) - ii, jj = min(i, nx), j - phase_ij = getindex_phase(phase, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K1 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) - - ii, jj = max(i - 1, 1), j - phase_ij = getindex_phase(phase, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K2 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) - K = (K1 + K2) * 0.5 - - qy = qTy2[i, j] = -K * d_yi(T) - qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) - end - - return nothing -end - -@parallel_indices (i, j) function compute_flux!( - qTx::AbstractArray{_T,2}, - qTy, - qTx2, - qTy2, - T, - rheology::NTuple{N,AbstractMaterialParamsStruct}, - phase_ratios::CellArray{C1,C2,C3,C4}, - θr_dτ, - _dx, - _dy, - args, -) where {_T,N,C1,C2,C3,C4} - nx = size(θr_dτ, 1) - - d_xi(A) = _d_xi(A, i, j, _dx) - d_yi(A) = _d_yi(A, i, j, _dy) - av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 - av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 - compute_K(phase, args) = fn_ratio(compute_conductivity, rheology, phase, args) - - @inbounds if all((i, j) .≤ size(qTx)) - ii, jj = clamp(i - 1, 1, nx), j + 1 - phase_ij = getindex_phase(phase_ratios, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K1 = compute_K(phase_ij, args_ij) - - ii, jj = clamp(i - 1, 1, nx), j - phase_ij = getindex_phase(phase_ratios, ii, jj) - K2 = compute_K(phase_ij, args_ij) - K = (K1 + K2) * 0.5 - - qx = qTx2[i, j] = -K * d_xi(T) - qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) - end - - @inbounds if all((i, j) .≤ size(qTy)) - ii, jj = min(i, nx), j - phase_ij = getindex_phase(phase_ratios, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K1 = compute_K(phase_ij, args_ij) - - ii, jj = clamp(i - 1, 1, nx), j - phase_ij = getindex_phase(phase_ratios, ii, jj) - args_ij = ntuple_idx(args, ii, jj) - K2 = compute_K(phase_ij, args_ij) - K = (K1 + K2) * 0.5 - - qy = qTy2[i, j] = -K * d_yi(T) - qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) - end - - return nothing -end - -@parallel_indices (i, j) function update_T!( - T::AbstractArray{_T,2}, Told, qTx, qTy, H, shear_heating, ρCp, dτ_ρ, _dt, _dx, _dy -) where {_T} - nx, ny = size(ρCp) - - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) - #! format: off - function av(A) - ( - A[clamp(i - 1, 1, nx), clamp(j - 1, 1, ny)] + - A[clamp(i - 1, 1, nx), j] + - A[i, clamp(j - 1, 1, ny)] + - A[i, j] - ) * 0.25 - end - #! format: on - - T[i + 1, j + 1] += - av(dτ_ρ) * ( - (-(d_xa(qTx) + d_ya(qTy))) - - av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt + - av(H) + - av(shear_heating) - ) - return nothing -end - -@parallel_indices (i, j) function update_T!( - T::AbstractArray{_T,2}, - Told, - qTx, - qTy, - H, - shear_heating, - rheology, - phase, - dτ_ρ, - _dt, - _dx, - _dy, - args::NamedTuple, -) where {_T} - nx, ny = size(args.P) - - i0 = clamp(i - 1, 1, nx) - i1 = clamp(i, 1, nx) - j0 = clamp(j - 1, 1, ny) - j1 = clamp(j, 1, ny) - - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) - av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 - - T_ij = T[i + 1, j + 1] - args_ij = (; T=T_ij, P=av(args.P)) - - ρCp = - ( - compute_ρCp(rheology, getindex_phase(phase, i0, j0), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i0, j1), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i1, j0), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) - ) * 0.25 - - T[i + 1, j + 1] += - av(dτ_ρ) * ( - (-(d_xa(qTx) + d_ya(qTy))) - ρCp * (T_ij - Told[i + 1, j + 1]) * _dt + - av(H) + - av(shear_heating) - ) - return nothing -end - -@parallel_indices (i, j) function check_res!( - ResT::AbstractArray{_T,2}, T, Told, qTx2, qTy2, H, shear_heating, ρCp, _dt, _dx, _dy -) where {_T} - nx, ny = size(ρCp) - - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) - #! format: off - function av(A) - ( - A[clamp(i - 1, 1, nx), clamp(j - 1, 1, ny)] + - A[clamp(i - 1, 1, nx), clamp(j, 1, ny)] + - A[clamp(i, 1, nx), clamp(j - 1, 1, ny)] + - A[clamp(i, 1, nx), clamp(j, 1, ny)] - ) * 0.25 - end - #! format: on - - ResT[i, j] = - -av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - - (d_xa(qTx2) + d_ya(qTy2)) + - av(H) + - av(shear_heating) - return nothing -end - -@parallel_indices (i, j) function check_res!( - ResT::AbstractArray{_T,2}, - T, - Told, - qTx2, - qTy2, - H, - shear_heating, - rheology, - phase, - _dt, - _dx, - _dy, - args, -) where {_T} - nx, ny = size(args.P) - - i0 = clamp(i - 1, 1, nx) - i1 = clamp(i, 1, nx) - j0 = clamp(j - 1, 1, ny) - j1 = clamp(j, 1, ny) - - d_xa(A) = _d_xa(A, i, j, _dx) - d_ya(A) = _d_ya(A, i, j, _dy) - av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 - - T_ij = T[i + 1, j + 1] - args_ij = (; T=T_ij, P=av(args.P)) - - ρCp = - ( - compute_ρCp(rheology, getindex_phase(phase, i0, j0), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i0, j1), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i1, j0), args_ij) + - compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) - ) * 0.25 - - ResT[i, j] = - -ρCp * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) + - av(H) + - av(shear_heating) - return nothing -end - -@parallel_indices (I...) function update_ΔT!(ΔT, T, Told) - ΔT[I...] = T[I...] - Told[I...] - return nothing -end - -### SOLVERS COLLECTION BELOW - THEY SHOULD BE DIMENSION AGNOSTIC - -@inline flux_range(nx, ny) = @idx (nx + 3, ny + 1) -@inline flux_range(nx, ny, nz) = @idx (nx, ny, nz) - -@inline update_range(nx, ny) = @idx (nx + 1, ny - 1) -@inline update_range(nx, ny, nz) = residual_range(nx, ny, nz) - -@inline residual_range(nx, ny) = update_range(nx, ny) -@inline residual_range(nx, ny, nz) = @idx (nx - 1, ny - 1, nz - 1) - -function update_T(::Nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) - @parallel update_range(ni...) update_T!( - thermal.T, - thermal.Told, - @qT(thermal)..., - thermal.H, - thermal.shear_heating, - ρCp, - pt_thermal.dτ_ρ, - _dt, - _di..., - ) -end - -function update_T( - ::Nothing, b_width, thermal, rheology, phase, pt_thermal, _dt, _di, ni, args -) - @parallel update_range(ni...) update_T!( - thermal.T, - thermal.Told, - @qT(thermal)..., - thermal.H, - thermal.shear_heating, - rheology, - phase, - pt_thermal.dτ_ρ, - _dt, - _di..., - args, - ) -end - -""" - heatdiffusion_PT!(thermal, pt_thermal, K, ρCp, dt, di; iterMax, nout, verbose) - -Heat diffusion solver using Pseudo-Transient iterations. Both `K` and `ρCp` are n-dimensional arrays. -""" -function heatdiffusion_PT!( - thermal::ThermalArrays, - pt_thermal::PTThermalCoeffs, - thermal_bc::TemperatureBoundaryConditions, - K::AbstractArray, - ρCp::AbstractArray, - dt, - di; - igg=nothing, - b_width=(4, 4, 4), - iterMax=50e3, - nout=1e3, - verbose=true, -) - # Compute some constant stuff - _dt = inv(dt) - _di = inv.(di) - _sq_len_RT = inv(sqrt(length(thermal.ResT))) - ϵ = pt_thermal.ϵ - ni = size(thermal.Tc) - @copy thermal.Told thermal.T - - # errors - iter_count = Int64[] - norm_ResT = Float64[] - sizehint!(iter_count, Int(iterMax)) - sizehint!(norm_ResT, Int(iterMax)) - - # Pseudo-transient iteration - iter = 0 - wtime0 = 0e0 - err = 2 * ϵ - - println("\n ====================================\n") - println("Starting thermal diffusion solver...\n") - - while err > ϵ && iter < iterMax - wtime0 += @elapsed begin - @parallel flux_range(ni...) compute_flux!( - @qT(thermal)..., @qT2(thermal)..., thermal.T, K, pt_thermal.θr_dτ, _di... - ) - update_T(nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) - thermal_bcs!(thermal.T, thermal_bc) - update_halo!(thermal.T) - end - - iter += 1 - - if iter % nout == 0 - wtime0 += @elapsed begin - @parallel residual_range(ni...) check_res!( - thermal.ResT, - thermal.T, - thermal.Told, - @qT2(thermal)..., - thermal.H, - thermal.shear_heating, - ρCp, - _dt, - _di..., - ) - end - - err = norm(thermal.ResT) * _sq_len_RT - - push!(norm_ResT, err) - push!(iter_count, iter) - - if verbose - @printf("iter = %d, err = %1.3e \n", iter, err) - end - end - end - - println("\n ...solver finished in $(round(wtime0, sigdigits=5)) seconds \n") - println("====================================\n") - - @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) - - return nothing -end - -""" - heatdiffusion_PT!(thermal, pt_thermal, rheology, dt, di; iterMax, nout, verbose) - -Heat diffusion solver using Pseudo-Transient iterations. -""" -function heatdiffusion_PT!( - thermal::ThermalArrays, - pt_thermal::PTThermalCoeffs, - thermal_bc::TemperatureBoundaryConditions, - rheology, - args::NamedTuple, - dt, - di; - igg=nothing, - phase=nothing, - b_width=(4, 4, 4), - iterMax=50e3, - nout=1e3, - verbose=true, -) - phases = get_phase(phase) - - # Compute some constant stuff - _dt = inv(dt) - _di = inv.(di) - _sq_len_RT = inv(sqrt(length(thermal.ResT))) - ϵ = pt_thermal.ϵ - ni = size(thermal.Tc) - @copy thermal.Told thermal.T - update_pt_thermal_arrays!(pt_thermal, phase, rheology, args, _dt) - - # errors - iter_count = Int64[] - norm_ResT = Float64[] - sizehint!(iter_count, Int(iterMax)) - sizehint!(norm_ResT, Int(iterMax)) - - # Pseudo-transient iteration - iter = 0 - wtime0 = 0e0 - err = 2 * ϵ - - println("\n ====================================\n") - println("Starting thermal diffusion solver...\n") - - while err > ϵ && iter < iterMax - wtime0 += @elapsed begin - @parallel flux_range(ni...) compute_flux!( - @qT(thermal)..., - @qT2(thermal)..., - thermal.T, - rheology, - phases, - pt_thermal.θr_dτ, - _di..., - args, - ) - update_T( - nothing, b_width, thermal, rheology, phases, pt_thermal, _dt, _di, ni, args - ) - thermal_bcs!(thermal.T, thermal_bc) - update_halo!(thermal.T) - end - - iter += 1 - - if iter % nout == 0 - wtime0 += @elapsed begin - @parallel residual_range(ni...) check_res!( - thermal.ResT, - thermal.T, - thermal.Told, - @qT2(thermal)..., - thermal.H, - thermal.shear_heating, - rheology, - phases, - _dt, - _di..., - args, - ) - end - - err = norm(thermal.ResT) * _sq_len_RT - - push!(norm_ResT, err) - push!(iter_count, iter) - - if verbose - @printf("iter = %d, err = %1.3e \n", iter, err) - end - end - end - - println("\n ...solver finished in $(round(wtime0, sigdigits=5)) seconds \n") - println("====================================\n") - - @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) - - return nothing -end +include("DiffusionPT_GeoParams.jl") +include("DiffusionPT_kernels.jl") +include("DiffusionPT_coefficients.jl") +include("ShearHeating.jl") +include("DiffusionPT_solver.jl") diff --git a/src/thermal_diffusion/DiffusionPT_GeoParams.jl b/src/thermal_diffusion/DiffusionPT_GeoParams.jl new file mode 100644 index 00000000..67464c7a --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_GeoParams.jl @@ -0,0 +1,103 @@ +## Phases + +@inline get_phase(x::PhaseRatio) = x.center +@inline get_phase(x) = x + +# update_pt_thermal_arrays!(::Vararg{Any,N}) where {N} = nothing + +function update_pt_thermal_arrays!( + pt_thermal, phase_ratios::PhaseRatio, rheology, args, _dt +) + ni = size(phase_ratios.center) + + @parallel (@idx ni) compute_pt_thermal_arrays!( + pt_thermal.θr_dτ, + pt_thermal.dτ_ρ, + rheology, + phase_ratios.center, + args, + pt_thermal.max_lxyz, + pt_thermal.Vpdτ, + _dt, + ) + + return nothing +end + +@inline function compute_phase(fn::F, rheology, phase::Int, args) where {F} + return fn(rheology, phase, args) +end + +@inline function compute_phase(fn::F, rheology, phase::SVector, args) where {F} + return fn_ratio(fn, rheology, phase, args) +end + +@inline compute_phase(fn::F, rheology, ::Nothing, args) where {F} = fn(rheology, args) + +@inline Base.@propagate_inbounds function getindex_phase( + phase::AbstractArray, I::Vararg{Int,N} +) where {N} + return phase[I...] +end + +@inline getindex_phase(::Nothing, I::Vararg{Int,N}) where {N} = nothing + +# Diffusivity + +@inline function compute_diffusivity(rheology, args) + return compute_conductivity(rheology, args) * + inv(compute_heatcapacity(rheology, args) * compute_density(rheology, args)) +end + +@inline function compute_diffusivity(rheology, phase::Union{Nothing,Int}, args) + return compute_conductivity(rheology, phase, args) * inv( + compute_heatcapacity(rheology, phase, args) * compute_density(rheology, phase, args) + ) +end + +@inline function compute_diffusivity(rheology, ρ, args) + return compute_conductivity(rheology, args) * + inv(compute_heatcapacity(rheology, args) * ρ) +end + +@inline function compute_diffusivity(rheology, ρ, phase::Union{Nothing,Int}, args) + return compute_conductivity(rheology, phase, args) * + inv(compute_heatcapacity(rheology, phase, args) * ρ) +end + +@inline function compute_diffusivity( + rheology::NTuple{N,AbstractMaterialParamsStruct}, phase_ratios::SArray, args +) where {N} + ρ = compute_density_ratio(phase_ratios, rheology, args) + conductivity = fn_ratio(compute_conductivity, rheology, phase_ratios, args) + heatcapacity = fn_ratio(compute_heatcapacity, rheology, phase_ratios, args) + return conductivity * inv(heatcapacity * ρ) +end + +# ρ*Cp + +@inline function compute_ρCp(rheology, args) + return compute_heatcapacity(rheology, args) * compute_density(rheology, args) +end + +@inline function compute_ρCp(rheology, phase::Union{Nothing,Int}, args) + return compute_phase(compute_heatcapacity, rheology, phase, args) * + compute_phase(compute_density, rheology, phase, args) +end + +@inline function compute_ρCp(rheology, ρ, args) + return compute_heatcapacity(rheology, args) * ρ +end + +@inline function compute_ρCp(rheology, ρ, phase::Union{Nothing,Int}, args) + return compute_phase(compute_heatcapacity, rheology, phase, args) * ρ +end + +@inline function compute_ρCp(rheology, phase_ratios::SArray, args) + return fn_ratio(compute_heatcapacity, rheology, phase_ratios, args) * + fn_ratio(compute_density, rheology, phase_ratios, args) +end + +@inline function compute_ρCp(rheology, ρ, phase_ratios::SArray, args) + return fn_ratio(compute_heatcapacity, rheology, phase_ratios, args) * ρ +end diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl new file mode 100644 index 00000000..11d6ff33 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -0,0 +1,85 @@ +# with phase ratios +function PTThermalCoeffs( + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + Vpdτ = min(di...) * CFL + max_lxyz = max(li...) + θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) + + @parallel (@idx ni) compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, phase_ratios.center, args, max_lxyz, Vpdτ, inv(dt) + ) + + return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) +end + +# without phase ratios +function PTThermalCoeffs( + rheology, args, dt, ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 +) where {nDim,T} + Vpdτ = min(di...) * CFL + max_lxyz = max(li...) + θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) + + @parallel (@idx ni) compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, inv(dt) + ) + + return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) +end + +@parallel_indices (I...) function compute_pt_thermal_arrays!( + θr_dτ::AbstractArray, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt +) + _compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, I... + ) + + return nothing +end + +@parallel_indices (I...) function compute_pt_thermal_arrays!( + θr_dτ::AbstractArray, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt +) + _compute_pt_thermal_arrays!(θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, I...) + + return nothing +end + +function _compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} +) where {N} + args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) + phase_ij = phase[Idx...] + ρCp = compute_ρCp(rheology, phase_ij, args_ij) + _K = inv(fn_ratio(compute_conductivity, rheology, phase_ij, args_ij)) + + _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number + θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re + dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re + + return nothing +end + +function _compute_pt_thermal_arrays!( + θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} +) where {N} + args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) + + ρCp = compute_ρCp(rheology, args_ij) + _K = inv(compute_conductivity(rheology, args_ij)) + + _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number + θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re + dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re + + return nothing +end diff --git a/src/thermal_diffusion/DiffusionPT_kernels.jl b/src/thermal_diffusion/DiffusionPT_kernels.jl new file mode 100644 index 00000000..7ad5d8d7 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_kernels.jl @@ -0,0 +1,549 @@ + +## 3D KERNELS + +@parallel_indices (i, j, k) function compute_flux!( + qTx::AbstractArray{_T,3}, qTy, qTz, qTx2, qTy2, qTz2, T, K, θr_dτ, _dx, _dy, _dz +) where {_T} + d_xi(A) = _d_xi(A, i, j, k, _dx) + d_yi(A) = _d_yi(A, i, j, k, _dy) + d_zi(A) = _d_zi(A, i, j, k, _dz) + av_xy(A) = _av_xy(A, i, j, k) + av_xz(A) = _av_xz(A, i, j, k) + av_yz(A) = _av_yz(A, i, j, k) + + I = i, j, k + + if all(I .≤ size(qTx)) + qx = qTx2[I...] = -av_yz(K) * d_xi(T) + qTx[I...] = (qTx[I...] * av_yz(θr_dτ) + qx) / (1.0 + av_yz(θr_dτ)) + end + + if all(I .≤ size(qTy)) + qy = qTy2[I...] = -av_xz(K) * d_yi(T) + qTy[I...] = (qTy[I...] * av_xz(θr_dτ) + qy) / (1.0 + av_xz(θr_dτ)) + end + + if all(I .≤ size(qTz)) + qz = qTz2[I...] = -av_xy(K) * d_zi(T) + qTz[I...] = (qTz[I...] * av_xy(θr_dτ) + qz) / (1.0 + av_xy(θr_dτ)) + end + + return nothing +end + +@parallel_indices (i, j, k) function compute_flux!( + qTx::AbstractArray{_T,3}, + qTy, + qTz, + qTx2, + qTy2, + qTz2, + T, + rheology, + phase, + θr_dτ, + _dx, + _dy, + _dz, + args, +) where {_T} + d_xi(A) = _d_xi(A, i, j, k, _dx) + d_yi(A) = _d_yi(A, i, j, k, _dy) + d_zi(A) = _d_zi(A, i, j, k, _dz) + av_xy(A) = _av_xy(A, i, j, k) + av_xz(A) = _av_xz(A, i, j, k) + av_yz(A) = _av_yz(A, i, j, k) + + get_K(idx, args) = compute_phase(compute_conductivity, rheology, idx, args) + + I = i, j, k + + @inbounds if all(I .≤ size(qTx)) + T_ijk = (T[(I .+ 1)...] + T[i, j + 1, k + 1]) * 0.5 + args_ijk = (; T=T_ijk, P=av_yz(args.P)) + K = + ( + get_K(getindex_phase(phase, i, j + 1, k), args_ijk) + + get_K(getindex_phase(phase, i, j, k + 1), args_ijk) + + get_K(getindex_phase(phase, i, j + 1, k + 1), args_ijk) + + get_K(getindex_phase(phase, i, j, k), args_ijk) + ) * 0.25 + + qx = qTx2[I...] = -K * d_xi(T) + qTx[I...] = (qTx[I...] * av_yz(θr_dτ) + qx) / (1.0 + av_yz(θr_dτ)) + end + + @inbounds if all(I .≤ size(qTy)) + T_ijk = (T[(I .+ 1)...] + T[i + 1, j, k + 1]) * 0.5 + args_ijk = (; T=T_ijk, P=av_xz(args.P)) + K = + ( + get_K(getindex_phase(phase, i + 1, j, k), args_ijk) + + get_K(getindex_phase(phase, i, j, k + 1), args_ijk) + + get_K(getindex_phase(phase, i + 1, j, k + 1), args_ijk) + + get_K(getindex_phase(phase, i, j, k), args_ijk) + ) * 0.25 + + qy = qTy2[I...] = -K * d_yi(T) + qTy[I...] = (qTy[I...] * av_xz(θr_dτ) + qy) / (1.0 + av_xz(θr_dτ)) + end + + @inbounds if all(I .≤ size(qTz)) + T_ijk = (T[(I .+ 1)...] + T[i + 1, j + 1, k]) * 0.5 + args_ijk = (; T=T_ijk, P=av_xy(args.P)) + K = + ( + get_K(getindex_phase(phase, i + 1, j + 1, k), args_ijk) + + get_K(getindex_phase(phase, i, j + 1, k), args_ijk) + + get_K(getindex_phase(phase, i + 1, j, k), args_ijk) + + get_K(getindex_phase(phase, i, j, k), args_ijk) + ) * 0.25 + + qz = qTz2[I...] = -K * d_zi(T) + qTz[I...] = (qTz[I...] * av_xy(θr_dτ) + qz) / (1.0 + av_xy(θr_dτ)) + end + + return nothing +end + +@parallel_indices (I...) function update_T!( + T::AbstractArray{_T,3}, + Told, + qTx, + qTy, + qTz, + H, + shear_heating, + ρCp, + dτ_ρ, + _dt, + _dx, + _dy, + _dz, +) where {_T} + av(A) = _av(A, I...) + d_xa(A) = _d_xa(A, I..., _dx) + d_ya(A) = _d_ya(A, I..., _dy) + d_za(A) = _d_za(A, I..., _dz) + + I1 = I .+ 1 + T[I1...] += + av(dτ_ρ) * ( + (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) - + av(ρCp) * (T[I1...] - Told[I1...]) * _dt + ) + + av(H) + + av(shear_heating) + + return nothing +end + +@parallel_indices (i, j, k) function update_T!( + T::AbstractArray{_T,3}, + Told, + qTx, + qTy, + qTz, + H, + shear_heating, + rheology, + phase, + dτ_ρ, + _dt, + _dx, + _dy, + _dz, + args, +) where {_T} + av(A) = _av(A, i, j, k) + d_xa(A) = _d_xa(A, i, j, k, _dx) + d_ya(A) = _d_ya(A, i, j, k, _dy) + d_za(A) = _d_za(A, i, j, k, _dz) + + I = i + 1, j + 1, k + 1 + + T_ijk = T[I...] + args_ijk = (; T=T_ijk, P=av(args.P)) + phase_ijk = getindex_phase(phase, i, j, k) + + T[I...] = + T_ijk + + av(dτ_ρ) * ( + (-(d_xa(qTx) + d_ya(qTy) + d_za(qTz))) - + compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt + + av(H) + + av(shear_heating) + ) + return nothing +end + +@parallel_indices (i, j, k) function check_res!( + ResT::AbstractArray{_T,3}, + T, + Told, + qTx2, + qTy2, + qTz2, + H, + shear_heating, + ρCp, + _dt, + _dx, + _dy, + _dz, +) where {_T} + d_xa(A) = _d_xa(A, i, j, k, _dx) + d_ya(A) = _d_ya(A, i, j, k, _dy) + d_za(A) = _d_za(A, i, j, k, _dz) + av(A) = _av(A, i, j, k) + + I = i + 1, j + 1, k + 1 + + ResT[i, j, k] = + -av(ρCp) * (T[I...] - Told[I...]) * _dt - (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + + av(H) + + av(shear_heating) + + return nothing +end + +@parallel_indices (i, j, k) function check_res!( + ResT::AbstractArray{_T,3}, + T, + Told, + qTx2, + qTy2, + qTz2, + H, + shear_heating, + rheology, + phase, + _dt, + _dx, + _dy, + _dz, + args, +) where {_T} + d_xa(A) = _d_xa(A, i, j, k, _dx) + d_ya(A) = _d_ya(A, i, j, k, _dy) + d_za(A) = _d_za(A, i, j, k, _dz) + av(A) = _av(A, i, j, k) + + I = i + 1, j + 1, k + 1 + T_ijk = T[I...] + args_ijk = (; T=T_ijk, P=av(args.P)) + phase_ijk = getindex_phase(phase, i, j, k) + + ResT[i, j, k] = + -compute_ρCp(rheology, phase_ijk, args_ijk) * (T_ijk - Told[I...]) * _dt - + (d_xa(qTx2) + d_ya(qTy2) + d_za(qTz2)) + + av(H) + + av(shear_heating) + + return nothing +end + +## 2D KERNELS + +@parallel_indices (i, j) function compute_flux!( + qTx::AbstractArray{_T,2}, qTy, qTx2, qTy2, T, K, θr_dτ, _dx, _dy +) where {_T} + nx = size(θr_dτ, 1) + + d_xi(A) = _d_xi(A, i, j, _dx) + d_yi(A) = _d_yi(A, i, j, _dy) + av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 + av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 + + @inbounds if all((i, j) .≤ size(qTx)) + qx = qTx2[i, j] = -av_xa(K) * d_xi(T) + qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) + end + + @inbounds if all((i, j) .≤ size(qTy)) + qy = qTy2[i, j] = -av_ya(K) * d_yi(T) + qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) + end + return nothing +end + +@parallel_indices (i, j) function compute_flux!( + qTx::AbstractArray{_T,2}, qTy, qTx2, qTy2, T, rheology, phase, θr_dτ, _dx, _dy, args +) where {_T} + nx = size(θr_dτ, 1) + + d_xi(A) = _d_xi(A, i, j, _dx) + d_yi(A) = _d_yi(A, i, j, _dy) + av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 + av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 + + if all((i, j) .≤ size(qTx)) + ii, jj = clamp(i - 1, 1, nx), j + 1 + phase_ij = getindex_phase(phase, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K1 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) + + ii, jj = clamp(i - 1, 1, nx), j + phase_ij = getindex_phase(phase, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K2 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) + K = (K1 + K2) * 0.5 + + qx = qTx2[i, j] = -K * d_xi(T) + qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) + end + + if all((i, j) .≤ size(qTy)) + ii, jj = min(i, nx), j + phase_ij = getindex_phase(phase, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K1 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) + + ii, jj = max(i - 1, 1), j + phase_ij = getindex_phase(phase, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K2 = compute_phase(compute_conductivity, rheology, phase_ij, args_ij) + K = (K1 + K2) * 0.5 + + qy = qTy2[i, j] = -K * d_yi(T) + qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) + end + + return nothing +end + +@parallel_indices (i, j) function compute_flux!( + qTx::AbstractArray{_T,2}, + qTy, + qTx2, + qTy2, + T, + rheology::NTuple{N,AbstractMaterialParamsStruct}, + phase_ratios::CellArray{C1,C2,C3,C4}, + θr_dτ, + _dx, + _dy, + args, +) where {_T,N,C1,C2,C3,C4} + nx = size(θr_dτ, 1) + + d_xi(A) = _d_xi(A, i, j, _dx) + d_yi(A) = _d_yi(A, i, j, _dy) + av_xa(A) = (A[clamp(i - 1, 1, nx), j + 1] + A[clamp(i - 1, 1, nx), j]) * 0.5 + av_ya(A) = (A[clamp(i, 1, nx), j] + A[clamp(i - 1, 1, nx), j]) * 0.5 + compute_K(phase, args) = fn_ratio(compute_conductivity, rheology, phase, args) + + @inbounds if all((i, j) .≤ size(qTx)) + ii, jj = clamp(i - 1, 1, nx), j + 1 + phase_ij = getindex_phase(phase_ratios, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K1 = compute_K(phase_ij, args_ij) + + ii, jj = clamp(i - 1, 1, nx), j + phase_ij = getindex_phase(phase_ratios, ii, jj) + K2 = compute_K(phase_ij, args_ij) + K = (K1 + K2) * 0.5 + + qx = qTx2[i, j] = -K * d_xi(T) + qTx[i, j] = (qTx[i, j] * av_xa(θr_dτ) + qx) / (1.0 + av_xa(θr_dτ)) + end + + @inbounds if all((i, j) .≤ size(qTy)) + ii, jj = min(i, nx), j + phase_ij = getindex_phase(phase_ratios, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K1 = compute_K(phase_ij, args_ij) + + ii, jj = clamp(i - 1, 1, nx), j + phase_ij = getindex_phase(phase_ratios, ii, jj) + args_ij = ntuple_idx(args, ii, jj) + K2 = compute_K(phase_ij, args_ij) + K = (K1 + K2) * 0.5 + + qy = qTy2[i, j] = -K * d_yi(T) + qTy[i, j] = (qTy[i, j] * av_ya(θr_dτ) + qy) / (1.0 + av_ya(θr_dτ)) + end + + return nothing +end + +@parallel_indices (i, j) function update_T!( + T::AbstractArray{_T,2}, Told, qTx, qTy, H, shear_heating, ρCp, dτ_ρ, _dt, _dx, _dy +) where {_T} + nx, ny = size(ρCp) + + d_xa(A) = _d_xa(A, i, j, _dx) + d_ya(A) = _d_ya(A, i, j, _dy) + #! format: off + function av(A) + ( + A[clamp(i - 1, 1, nx), clamp(j - 1, 1, ny)] + + A[clamp(i - 1, 1, nx), j] + + A[i, clamp(j - 1, 1, ny)] + + A[i, j] + ) * 0.25 + end + #! format: on + + T[i + 1, j + 1] += + av(dτ_ρ) * ( + (-(d_xa(qTx) + d_ya(qTy))) - + av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt + + av(H) + + av(shear_heating) + ) + return nothing +end + +@parallel_indices (i, j) function update_T!( + T::AbstractArray{_T,2}, + Told, + qTx, + qTy, + H, + shear_heating, + rheology, + phase, + dτ_ρ, + _dt, + _dx, + _dy, + args::NamedTuple, +) where {_T} + nx, ny = size(args.P) + + i0 = clamp(i - 1, 1, nx) + i1 = clamp(i, 1, nx) + j0 = clamp(j - 1, 1, ny) + j1 = clamp(j, 1, ny) + + d_xa(A) = _d_xa(A, i, j, _dx) + d_ya(A) = _d_ya(A, i, j, _dy) + av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 + + T_ij = T[i + 1, j + 1] + args_ij = (; T=T_ij, P=av(args.P)) + + ρCp = + ( + compute_ρCp(rheology, getindex_phase(phase, i0, j0), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i0, j1), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i1, j0), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) + ) * 0.25 + + T[i + 1, j + 1] += + av(dτ_ρ) * ( + (-(d_xa(qTx) + d_ya(qTy))) - ρCp * (T_ij - Told[i + 1, j + 1]) * _dt + + av(H) + + av(shear_heating) + ) + return nothing +end + +@parallel_indices (i, j) function check_res!( + ResT::AbstractArray{_T,2}, T, Told, qTx2, qTy2, H, shear_heating, ρCp, _dt, _dx, _dy +) where {_T} + nx, ny = size(ρCp) + + d_xa(A) = _d_xa(A, i, j, _dx) + d_ya(A) = _d_ya(A, i, j, _dy) + #! format: off + function av(A) + ( + A[clamp(i - 1, 1, nx), clamp(j - 1, 1, ny)] + + A[clamp(i - 1, 1, nx), clamp(j, 1, ny)] + + A[clamp(i, 1, nx), clamp(j - 1, 1, ny)] + + A[clamp(i, 1, nx), clamp(j, 1, ny)] + ) * 0.25 + end + #! format: on + + ResT[i, j] = + -av(ρCp) * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - + (d_xa(qTx2) + d_ya(qTy2)) + + av(H) + + av(shear_heating) + return nothing +end + +@parallel_indices (i, j) function check_res!( + ResT::AbstractArray{_T,2}, + T, + Told, + qTx2, + qTy2, + H, + shear_heating, + rheology, + phase, + _dt, + _dx, + _dy, + args, +) where {_T} + nx, ny = size(args.P) + + i0 = clamp(i - 1, 1, nx) + i1 = clamp(i, 1, nx) + j0 = clamp(j - 1, 1, ny) + j1 = clamp(j, 1, ny) + + d_xa(A) = _d_xa(A, i, j, _dx) + d_ya(A) = _d_ya(A, i, j, _dy) + av(A) = (A[i0, j0] + A[i0, j1] + A[i1, j0] + A[i1, j1]) * 0.25 + + T_ij = T[i + 1, j + 1] + args_ij = (; T=T_ij, P=av(args.P)) + + ρCp = + ( + compute_ρCp(rheology, getindex_phase(phase, i0, j0), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i0, j1), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i1, j0), args_ij) + + compute_ρCp(rheology, getindex_phase(phase, i1, j1), args_ij) + ) * 0.25 + + ResT[i, j] = + -ρCp * (T[i + 1, j + 1] - Told[i + 1, j + 1]) * _dt - (d_xa(qTx2) + d_ya(qTy2)) + + av(H) + + av(shear_heating) + return nothing +end + +@parallel_indices (I...) function update_ΔT!(ΔT, T, Told) + ΔT[I...] = T[I...] - Told[I...] + return nothing +end + +function update_T(::Nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) + @parallel update_range(ni...) update_T!( + thermal.T, + thermal.Told, + @qT(thermal)..., + thermal.H, + thermal.shear_heating, + ρCp, + pt_thermal.dτ_ρ, + _dt, + _di..., + ) +end + +function update_T( + ::Nothing, b_width, thermal, rheology, phase, pt_thermal, _dt, _di, ni, args +) + @parallel update_range(ni...) update_T!( + thermal.T, + thermal.Told, + @qT(thermal)..., + thermal.H, + thermal.shear_heating, + rheology, + phase, + pt_thermal.dτ_ρ, + _dt, + _di..., + args, + ) +end diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl new file mode 100644 index 00000000..03349ad4 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -0,0 +1,206 @@ +function heatdiffusion_PT!(thermal::ThermalArrays, args...; kwargs) + return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs...) +end + +""" + heatdiffusion_PT!(thermal, pt_thermal, K, ρCp, dt, di; iterMax, nout, verbose) + +Heat diffusion solver using Pseudo-Transient iterations. Both `K` and `ρCp` are n-dimensional arrays. +""" +function heatdiffusion_PT!( + ::CPUBackendTrait, + thermal::ThermalArrays, + pt_thermal::PTThermalCoeffs, + thermal_bc::TemperatureBoundaryConditions, + K::AbstractArray, + ρCp::AbstractArray, + dt, + di; + igg=nothing, + b_width=(4, 4, 4), + iterMax=50e3, + nout=1e3, + verbose=true, +) + # Compute some constant stuff + _dt = inv(dt) + _di = inv.(di) + _sq_len_RT = inv(sqrt(length(thermal.ResT))) + ϵ = pt_thermal.ϵ + ni = size(thermal.Tc) + @copy thermal.Told thermal.T + + # errors + iter_count = Int64[] + norm_ResT = Float64[] + sizehint!(iter_count, Int(iterMax)) + sizehint!(norm_ResT, Int(iterMax)) + + # Pseudo-transient iteration + iter = 0 + wtime0 = 0e0 + err = 2 * ϵ + + println("\n ====================================\n") + println("Starting thermal diffusion solver...\n") + + while err > ϵ && iter < iterMax + wtime0 += @elapsed begin + @parallel flux_range(ni...) compute_flux!( + @qT(thermal)..., @qT2(thermal)..., thermal.T, K, pt_thermal.θr_dτ, _di... + ) + update_T(nothing, b_width, thermal, ρCp, pt_thermal, _dt, _di, ni) + thermal_bcs!(thermal.T, thermal_bc) + update_halo!(thermal.T) + end + + iter += 1 + + if iter % nout == 0 + wtime0 += @elapsed begin + @parallel residual_range(ni...) check_res!( + thermal.ResT, + thermal.T, + thermal.Told, + @qT2(thermal)..., + thermal.H, + thermal.shear_heating, + ρCp, + _dt, + _di..., + ) + end + + err = norm(thermal.ResT) * _sq_len_RT + + push!(norm_ResT, err) + push!(iter_count, iter) + + if verbose + @printf("iter = %d, err = %1.3e \n", iter, err) + end + end + end + + println("\n ...solver finished in $(round(wtime0, sigdigits=5)) seconds \n") + println("====================================\n") + + @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) + + @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + + return nothing +end + +""" + heatdiffusion_PT!(thermal, pt_thermal, rheology, dt, di; iterMax, nout, verbose) + +Heat diffusion solver using Pseudo-Transient iterations. +""" +function heatdiffusion_PT!( + ::CPUBackendTrait, + thermal::ThermalArrays, + pt_thermal::PTThermalCoeffs, + thermal_bc::TemperatureBoundaryConditions, + rheology, + args::NamedTuple, + dt, + di; + igg=nothing, + phase=nothing, + b_width=(4, 4, 4), + iterMax=50e3, + nout=1e3, + verbose=true, +) + phases = get_phase(phase) + + # Compute some constant stuff + _dt = inv(dt) + _di = inv.(di) + _sq_len_RT = inv(sqrt(length(thermal.ResT))) + ϵ = pt_thermal.ϵ + ni = size(thermal.Tc) + @copy thermal.Told thermal.T + update_pt_thermal_arrays!(pt_thermal, phase, rheology, args, _dt) + + # errors + iter_count = Int64[] + norm_ResT = Float64[] + sizehint!(iter_count, Int(iterMax)) + sizehint!(norm_ResT, Int(iterMax)) + + # Pseudo-transient iteration + iter = 0 + wtime0 = 0e0 + err = 2 * ϵ + + println("\n ====================================\n") + println("Starting thermal diffusion solver...\n") + + while err > ϵ && iter < iterMax + wtime0 += @elapsed begin + @parallel flux_range(ni...) compute_flux!( + @qT(thermal)..., + @qT2(thermal)..., + thermal.T, + rheology, + phases, + pt_thermal.θr_dτ, + _di..., + args, + ) + update_T( + nothing, b_width, thermal, rheology, phases, pt_thermal, _dt, _di, ni, args + ) + thermal_bcs!(thermal.T, thermal_bc) + update_halo!(thermal.T) + end + + iter += 1 + + if iter % nout == 0 + wtime0 += @elapsed begin + @parallel residual_range(ni...) check_res!( + thermal.ResT, + thermal.T, + thermal.Told, + @qT2(thermal)..., + thermal.H, + thermal.shear_heating, + rheology, + phases, + _dt, + _di..., + args, + ) + end + + err = norm(thermal.ResT) * _sq_len_RT + + push!(norm_ResT, err) + push!(iter_count, iter) + + if verbose + @printf("iter = %d, err = %1.3e \n", iter, err) + end + end + end + + println("\n ...solver finished in $(round(wtime0, sigdigits=5)) seconds \n") + println("====================================\n") + + @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) + @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + + return nothing +end + +@inline flux_range(nx, ny) = @idx (nx + 3, ny + 1) +@inline flux_range(nx, ny, nz) = @idx (nx, ny, nz) + +@inline update_range(nx, ny) = @idx (nx + 1, ny - 1) +@inline update_range(nx, ny, nz) = residual_range(nx, ny, nz) + +@inline residual_range(nx, ny) = update_range(nx, ny) +@inline residual_range(nx, ny, nz) = @idx (nx - 1, ny - 1, nz - 1) diff --git a/src/thermal_diffusion/Shearheating.jl b/src/thermal_diffusion/Shearheating.jl index 8fc1cb9c..062ae3b6 100644 --- a/src/thermal_diffusion/Shearheating.jl +++ b/src/thermal_diffusion/Shearheating.jl @@ -1,4 +1,17 @@ -@parallel_indices (I...) function compute_shear_heating!( +function compute_shear_heating!(thermal, stokes, rheology, dt) + ni = size(thermal.shear_heating) + @parallel (ni) compute_shear_heating_kernel!( + thermal.shear_heating, + @tensor_center(stokes.τ), + @tensor_center(stokes.τ_o), + @strain(stokes), + rheology, + dt, + ) + return nothing +end + +@parallel_indices (I...) function compute_shear_heating_kernel!( shear_heating, τ::NTuple{N,T}, τ_old::NTuple{N,T}, ε::NTuple{N,T}, rheology, dt ) where {N,T} _Gdt = inv(get_shear_modulus(rheology) * dt) @@ -8,7 +21,21 @@ return nothing end -@parallel_indices (I...) function compute_shear_heating!( +function compute_shear_heating!(thermal, stokes, phase_ratios::PhaseRatio, rheology, dt) + ni = size(thermal.shear_heating) + @parallel (ni) compute_shear_heating_kernel!( + thermal.shear_heating, + @tensor_center(stokes.τ), + @tensor_center(stokes.τ_o), + @strain(stokes), + phase_ratios.center, + rheology, + dt, + ) + return nothing +end + +@parallel_indices (I...) function compute_shear_heating_kernel!( shear_heating, τ::NTuple{N,T}, τ_old::NTuple{N,T}, diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl new file mode 100644 index 00000000..cd55b455 --- /dev/null +++ b/src/types/heat_diffusion.jl @@ -0,0 +1,95 @@ +struct ThermalArrays{_T} + T::_T # Temperature @ grid nodes + Tc::_T # Temperature @ cell centers + ΔT::_T + Told::_T + dT_dt::_T + qTx::_T + qTy::_T + qTz::Union{_T,Nothing} + qTx2::_T + qTy2::_T + qTz2::Union{_T,Nothing} + H::_T # source terms + shear_heating::_T # shear heating terms + ResT::_T + + function ThermalArrays(nx::Integer, ny::Integer) + T = @zeros(nx + 3, ny + 1) + ΔT = @zeros(nx + 3, ny + 1) + Told = @zeros(nx + 3, ny + 1) + Tc = @zeros(nx, ny) + H = @zeros(nx, ny) + shear_heating = @zeros(nx, ny) + dT_dt = @zeros(nx + 1, ny - 1) + qTx = @zeros(nx + 2, ny - 1) + qTy = @zeros(nx + 1, ny) + qTx2 = @zeros(nx + 2, ny - 1) + qTy2 = @zeros(nx + 1, ny) + ResT = @zeros(nx + 1, ny - 1) + return new{typeof(T)}( + T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTx2, qTy2, H, shear_heating, ResT + ) + end + + function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) + T = @zeros(nx + 1, ny + 1, nz + 1) + ΔT = @zeros(nx + 1, ny + 1, nz + 1) + Told = @zeros(nx + 1, ny + 1, nz + 1) + Tc = @zeros(nx, ny, nz) + H = @zeros(nx, ny, nz) + shear_heating = @zeros(nx, ny, nz) + dT_dt = @zeros(ni .- 1) + qTx = @zeros(nx, ny - 1, nz - 1) + qTy = @zeros(nx - 1, ny, nz - 1) + qTz = @zeros(nx - 1, ny - 1, nz) + qTx2 = @zeros(nx, ny - 1, nz - 1) + qTy2 = @zeros(nx - 1, ny, nz - 1) + qTz2 = @zeros(nx - 1, ny - 1, nz) + ResT = @zeros((ni .- 1)...) + return new{typeof(T)}( + T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT + ) + end +end + +ThermalArrays(ni::NTuple{N,Number}) where {N} = ThermalArrays(ni...) +function ThermalArrays(::Number, ::Number) + throw(ArgumentError("ThermalArrays dimensions must be given as integers")) +end +function ThermalArrays(::Number, ::Number, ::Number) + throw(ArgumentError("ThermalArrays dimensions must be given as integers")) +end + +# traits +@inline backend(::ThermalArrays{T}) where {T} = backend(T) + +## Thermal diffusion coefficients + +struct PTThermalCoeffs{T,M} + CFL::T + ϵ::T + max_lxyz::T + max_lxyz2::T + Vpdτ::T + θr_dτ::M + dτ_ρ::M + + function PTThermalCoeffs( + CFL::T, ϵ::T, max_lxyz::T, max_lxyz2::T, Vpdτ::T, θr_dτ::M, dτ_ρ::M + ) where {T,M} + return new{T,M}(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) + end +end + +function PTThermalCoeffs( + K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 +) where {nDim,T} + Vpdτ = min(di...) * CFL + max_lxyz = max(li...) + max_lxyz2 = max_lxyz^2 + Re = @. π + √(π * π + ρCp * max_lxyz2 / K / dt) # Numerical Reynolds number + θr_dτ = @. max_lxyz / Vpdτ / Re + dτ_ρ = @. Vpdτ * max_lxyz / K / Re + return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) +end diff --git a/src/types/stokes.jl b/src/types/stokes.jl new file mode 100644 index 00000000..ef54a975 --- /dev/null +++ b/src/types/stokes.jl @@ -0,0 +1,246 @@ +# 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 + +## Velocity type + +struct Velocity{T} + Vx::T + Vy::T + Vz::Union{T,Nothing} + + Velocity(Vx::T, Vy::T, Vz::Union{T,Nothing}) where {T} = new{T}(Vx, Vy, Vz) +end + +Velocity(Vx::T, Vy::T) where {T} = Velocity(Vx, Vy, nothing) + +function Velocity(nx::Integer, ny::Integer) + nVx = (nx + 1, ny + 2) + nVy = (nx + 2, ny + 1) + + Vx, Vy = @zeros(nVx...), @zeros(nVy) + return Velocity(Vx, Vy, nothing) +end + +function Velocity(nx::Integer, ny::Integer, nz::Integer) + nVx = (nx + 1, ny + 2, nz + 2) + nVy = (nx + 2, ny + 1, nz + 2) + nVz = (nx + 2, ny + 2, nz + 1) + + Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) + return Velocity{PTArray}(Vx, Vy, Vz) +end + +Velocity(ni::NTuple{N,Number}) where {N} = Velocity(ni...) +function Velocity(::Number, ::Number) + throw(ArgumentError("Velocity dimensions must be given as integers")) +end +function Velocity(::Number, ::Number, ::Number) + throw(ArgumentError("Velocity dimensions must be given as integers")) +end + +## Viscosity type + +struct Viscosity{T} + η::T # with no plasticity + η_vep::T # with plasticity + ητ::T # PT viscosi + + Viscosity(args::Vararg{T,N}) where {T,N} = new{T}(args...) +end + +Viscosity(args...) = Viscosity(promote(args...)...) + +function Viscosity(ni::NTuple{N,Integer}) where {N} + η = @zeros(ni...) + η_vep = @zeros(ni...) + ητ = @zeros(ni...) + return Viscosity(η, η_vep, ητ) +end + +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")) +end + +## SymmetricTensor type + +struct SymmetricTensor{T} + xx::T + yy::T + zz::Union{T,Nothing} + xy::T + yz::Union{T,Nothing} + xz::Union{T,Nothing} + xy_c::T + yz_c::Union{T,Nothing} + xz_c::Union{T,Nothing} + II::T + + function SymmetricTensor( + xx::T, + yy::T, + zz::Union{T,Nothing}, + xy::T, + yz::Union{T,Nothing}, + xz::Union{T,Nothing}, + xy_c::T, + yz_c::Union{T,Nothing}, + xz_c::Union{T,Nothing}, + II::T, + ) where {T} + return new{T}(xx, yy, zz, xy, yz, xz, xy_c, yz_c, xz_c, II) + end +end + +function SymmetricTensor(xx::T, yy::T, xy::T, xy_c::T, II::T) where {T} + return SymmetricTensor( + xx, yy, nothing, xy, nothing, nothing, xy_c, nothing, nothing, II + ) +end + +function SymmetricTensor(nx::Integer, ny::Integer) + return SymmetricTensor( + @zeros(nx, ny), # xx + @zeros(nx, ny), # yy + @zeros(nx + 1, ny + 1), # xy + @zeros(nx, ny), # xy @ cell center + @zeros(nx, ny) # II (second invariant) + ) +end + +function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) + return SymmetricTensor( + @zeros(nx, ny, nz), # xx + @zeros(nx + 1, ny + 1, nz), # xy + @zeros(nx, ny, nz), # yy + @zeros(nx + 1, ny, nz + 1), # xz + @zeros(nx, ny + 1, nz + 1), # yz + @zeros(nx, ny, nz), # zz + @zeros(nx, ny, nz), # yz @ cell center + @zeros(nx, ny, nz), # xz @ cell center + @zeros(nx, ny, nz), # xy @ cell center + @zeros(nx, ny, nz), # II (second invariant) + ) +end + +SymmetricTensor(ni::NTuple{N,Number}) where {N} = SymmetricTensor(ni...) +function SymmetricTensor(::Number, ::Number) + throw(ArgumentError("SymmetricTensor dimensions must be given as integers")) +end +function SymmetricTensor(::Number, ::Number, ::Number) + throw(ArgumentError("SymmetricTensor dimensions must be given as integers")) +end + +## Residual type + +struct Residual{T} + RP::T + Rx::T + Ry::T + Rz::Union{T,Nothing} + + Residual(RP::T, Rx::T, Ry::T, Rz::Union{T,Nothing}) where {T} = new{T}(RP, Rx, Ry, Rz) +end + +Residual(RP::T, Rx::T, Ry::T) where {T} = Residual(RP, Rx, Ry, nothing) + +function Residual(nx::Integer, ny::Integer) + Rx = @zeros(nx - 1, ny) + Ry = @zeros(nx, ny - 1) + RP = @zeros(nx, ny) + return Residual(RP, Rx, Ry) +end + +function Residual(nx::Integer, ny::Integer, nz::Integer) + Rx = @zeros(nx - 1, ny, nz) + Ry = @zeros(nx, ny - 1, nz) + Rz = @zeros(nx, ny, nz - 1) + RP = @zeros(nx, ny, nz) + return Residual(RP, Rx, Ry, Rz) +end + +Residual(ni::NTuple{N,Number}) where {N} = Residual(ni...) +function Residual(::Number, ::Number) + throw(ArgumentError("Residual dimensions must be given as integers")) +end +function Residual(::Number, ::Number, ::Number) + throw(ArgumentError("Residual dimensions must be given as integers")) +end + +## StokesArrays type + +struct StokesArrays{A,B,C,D,T} + P::T + P0::T + V::A + ∇V::T + τ::B + ε::B + ε_pl::B + EII_pl::T + viscosity::D + τ_o::Union{B,Nothing} + R::C + + function StokesArrays(ni::NTuple{N,Integer}) where {N} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(ni) + τ = SymmetricTensor(ni) + τ_o = SymmetricTensor(ni) + ε = SymmetricTensor(ni) + ε_pl = SymmetricTensor(ni) + EII_pl = @zeros(ni...) + viscosity = Viscosity(ni) + R = Residual(ni) + + return new{typeof(V),typeof(τ),typeof(R),typeof(viscosity),typeof(P)}( + P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R + ) + end +end + +StokesArrays(ni::Vararg{Integer,N}) where {N} = StokesArrays(tuple(ni...)) +function StokesArrays(::Number, ::Number) + throw(ArgumentError("StokesArrays dimensions must be given as integers")) +end +function StokesArrays(::Number, ::Number, ::Number) + throw(ArgumentError("StokesArrays dimensions must be given as integers")) +end + +# traits +@inline backend(x::StokesArrays) = backend(x.P) + +## PTStokesCoeffs type + +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; + ϵ::Float64=1e-8, + Re::Float64=3π, + CFL::Float64=(N == 2 ? 0.9 / √2.1 : 0.9 / √3.1), + r::Float64=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{Float64}(CFL, ϵ, Re, r, Vpdτ, θ_dτ, ηdτ) + end +end diff --git a/src/types/traits.jl b/src/types/traits.jl new file mode 100644 index 00000000..73bbab95 --- /dev/null +++ b/src/types/traits.jl @@ -0,0 +1,6 @@ +abstract type BackendTrait end +struct CPUBackendTrait <: BackendTrait end + +@inline backend(::Array) = CPUBackendTrait() +@inline backend(::Type{<:Array}) = CPUBackendTrait() +@inline backend(::T) where {T} = throw(ArgumentError("Backend $(T) not supported")) diff --git a/test_types.jl b/test_types.jl new file mode 100644 index 00000000..e5960886 --- /dev/null +++ b/test_types.jl @@ -0,0 +1,67 @@ +using Test +using ParallelStencil +@init_parallel_stencil(Threads,Float64,2) + +include("src/stokes/types/types.jl") + +T = Array + +ni = nx, ny = (10, 10) + +R = Residual(ni) +@test isnothing(R.Rz) +@test size(R.Rx) == (nx-1, ny) +@test size(R.Ry) == (nx, ny-1) +@test size(R.RP) == ni +@test R.Rx isa Array +@test R.Ry isa Array +@test R.RP isa Array + +@test Residual(ni) isa Residual +@test Residual(nx, ny) isa Residual +@test_throws ArgumentError Residual(10.0, 10.0) + +visc = Viscosity(ni) +@test size(visc.η) == ni +@test size(visc.η_vep) == ni +@test size(visc.ητ) == ni +@test visc.η isa Array +@test visc.η_vep isa Array +@test visc.ητ isa Array +@test_throws ArgumentError Viscosity(10.0, 10.0) + +v = Velocity(ni) +tensor = SymmetricTensor(ni) + +@test size(tensor.xx) == (nx, ny) +@test size(tensor.yy) == (nx, ny) +@test size(tensor.xy) == (nx + 1, ny + 1) +@test size(tensor.xy_c) == (nx, ny) +@test size(tensor.II) == (nx, ny) + +@test tensor.xx isa Array +@test tensor.yy isa Array +@test tensor.xy isa Array +@test tensor.xy_c isa Array +@test tensor.II isa Array + +stokes = StokesArrays(ni) + +@test size(stokes.P) == ni +@test size(stokes.P0) == ni +@test size(stokes.∇V) == ni +@test size(stokes.EII_pl) == ni + +@test stokes.P isa Array +@test stokes.P0 isa Array +@test stokes.∇V isa Array +@test stokes.V isa Velocity +@test stokes.τ isa SymmetricTensor +@test stokes.τ_o isa SymmetricTensor +@test stokes.ε isa SymmetricTensor +@test stokes.ε_pl isa SymmetricTensor +@test stokes.EII_pl isa Array +@test stokes.viscosity isa Viscosity +@test stokes.R isa Residual + +@test_throws ArgumentError StokesArrays(10.0, 10.0) From 32db87e5a5eafb7e14b032f25278816e0a7fc919 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 12:02:50 +0200 Subject: [PATCH 02/91] adapt shear heating minapp --- .../stokes2D/shear_heating/Shearheating2D.jl | 36 +++++++++---------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index 842b3fec..e002e937 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -83,12 +83,11 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) particle_args = (pT, pPhases) # Elliptical temperature anomaly - xc_anomaly = lx/2 # origin of thermal anomaly - yc_anomaly = 40e3 # origin of thermal anomaly - # yc_anomaly = 39e3 # origin of thermal anomaly - r_anomaly = 3e3 # radius of perturbation - init_phases!(pPhases, particles, lx/2, yc_anomaly, r_anomaly) + xc_anomaly = lx/2 # origin of thermal anomaly + yc_anomaly = 40e3 # origin of thermal anomaly + r_anomaly = 3e3 # radius of perturbation phase_ratios = PhaseRatio(ni, length(rheology)) + init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- @@ -120,10 +119,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) compute_viscosity!( stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1 + rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-1 / √2.1 ) # Boundary conditions @@ -230,11 +228,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) ) # ------------------------------ @@ -252,14 +252,14 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) inject = check_injection(particles) inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + 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) + checkpointing(figdir, stokes, thermal.T, stokes.viscosity.η, t) if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) @@ -276,7 +276,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) τyy = Array(stokes.τ.yy), εxx = Array(stokes.ε.xx), εyy = Array(stokes.ε.yy), - η = Array(η), + η = Array(stokes.viscosity.η), ) do_vtk( joinpath(vtk_dir, "vtk_" * lpad("$it", 6, "0")), @@ -308,7 +308,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # 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) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) @@ -339,6 +339,4 @@ else igg end -# main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) - -foo(args...; x=1) = x \ No newline at end of file +main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) From 91b1bb807c5ba85a0b63b4e3b9445ac6a41c863b Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 12:03:26 +0200 Subject: [PATCH 03/91] more fixes --- src/JustRelax.jl | 1 + src/JustRelax_CPU.jl | 6 ++++-- src/common.jl | 1 + src/stokes/Stokes2D.jl | 15 +++++---------- src/thermal_diffusion/DiffusionPT_solver.jl | 5 ++--- src/thermal_diffusion/Shearheating.jl | 2 +- src/types/heat_diffusion.jl | 2 +- src/types/stokes.jl | 4 ++-- 8 files changed, 17 insertions(+), 19 deletions(-) diff --git a/src/JustRelax.jl b/src/JustRelax.jl index f4b6ade1..787ffa65 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -13,6 +13,7 @@ using StaticArrays function solve!() end include("types/traits.jl") +export BackendTrait, CPUBackendTrait include("topology/Topology.jl") export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index b784fe47..a5ea5eab 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -6,8 +6,9 @@ using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences2D using ImplicitGlobalGrid using GeoParams, LinearAlgebra, Printf +using MPI -import JustRelax: IGG, BackendTrait, CPUBackendTrait +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend @init_parallel_stencil(Threads, Float64, 2) @@ -29,8 +30,9 @@ using CellArrays using ParallelStencil, ParallelStencil.FiniteDifferences3D using ImplicitGlobalGrid using GeoParams, LinearAlgebra, Printf +using MPI -import JustRelax: IGG, BackendTrait, CPUBackendTrait +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend @init_parallel_stencil(Threads, Float64, 3) diff --git a/src/common.jl b/src/common.jl index 1a7f38f2..9364141a 100644 --- a/src/common.jl +++ b/src/common.jl @@ -55,6 +55,7 @@ export WENO5, WENO_advection! # Stokes include("rheology/GeoParams.jl") +include("rheology/StressUpdate.jl") include("stokes/StressRotation.jl") include("stokes/StressKernels.jl") export tensor_invariant! diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 45868fff..47b9b00d 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -509,7 +509,8 @@ function solve!( ) if rem(iter, 5) == 0 - @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + # @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) end @parallel (@idx ni .+ 1) compute_strain_rate!( @@ -520,14 +521,8 @@ function solve!( @copy η0 η end if do_visc - @parallel (@idx ni) compute_viscosity!( - η, - viscosity_relaxation, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, + compute_viscosity!( + stokes,viscosity_relaxation, phase_ratios, args, rheology, viscosity_cutoff ) end @@ -549,7 +544,7 @@ function solve!( θ_dτ, ) - @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) + center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index 03349ad4..0e1b505f 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -86,8 +86,7 @@ function heatdiffusion_PT!( println("====================================\n") @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) return nothing end @@ -191,7 +190,7 @@ function heatdiffusion_PT!( println("====================================\n") @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) return nothing end diff --git a/src/thermal_diffusion/Shearheating.jl b/src/thermal_diffusion/Shearheating.jl index 062ae3b6..cb1a3755 100644 --- a/src/thermal_diffusion/Shearheating.jl +++ b/src/thermal_diffusion/Shearheating.jl @@ -23,7 +23,7 @@ end function compute_shear_heating!(thermal, stokes, phase_ratios::PhaseRatio, rheology, dt) ni = size(thermal.shear_heating) - @parallel (ni) compute_shear_heating_kernel!( + @parallel (@idx ni) compute_shear_heating_kernel!( thermal.shear_heating, @tensor_center(stokes.τ), @tensor_center(stokes.τ_o), diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index cd55b455..28fdb90a 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -28,7 +28,7 @@ struct ThermalArrays{_T} qTy2 = @zeros(nx + 1, ny) ResT = @zeros(nx + 1, ny - 1) return new{typeof(T)}( - T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTx2, qTy2, H, shear_heating, ResT + T, Tc, ΔT, Told, dT_dt, qTx, qTy, nothing, qTx2, qTy2, nothing, H, shear_heating, ResT ) end diff --git a/src/types/stokes.jl b/src/types/stokes.jl index ef54a975..c42ceffa 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -55,8 +55,8 @@ end Viscosity(args...) = Viscosity(promote(args...)...) function Viscosity(ni::NTuple{N,Integer}) where {N} - η = @zeros(ni...) - η_vep = @zeros(ni...) + η = @ones(ni...) + η_vep = @ones(ni...) ητ = @zeros(ni...) return Viscosity(η, η_vep, ητ) end From 6ac7cb35e6a2189b6d5396f6580a46ab1bcf5b82 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 12:18:11 +0200 Subject: [PATCH 04/91] fixes and cleanup --- .../benchmarks/stokes2D/RunStokesBench2D.jl | 6 +- .../elastic_buildup/Elastic_BuildUp.jl | 13 +- .../PlumeFreeSurface_2D.jl | 52 +- .../RayleighTaylor2D.jl | 7 +- .../stokes2D/shear_heating/Shearheating2D.jl | 14 +- src/MetaJustRelax.jl | 173 ---- src/stokes/MetaStokes.jl | 248 ------ src/stokes/Stokes2D.jl | 9 +- src/stokes/Stokes2D_old.jl | 808 ------------------ src/stokes/Stokes3D_old.jl | 561 ------------ src/thermal_diffusion/MetaDiffusion.jl | 216 ----- 11 files changed, 41 insertions(+), 2066 deletions(-) delete mode 100644 src/MetaJustRelax.jl delete mode 100644 src/stokes/MetaStokes.jl delete mode 100644 src/stokes/Stokes2D_old.jl delete mode 100644 src/stokes/Stokes3D_old.jl delete mode 100644 src/thermal_diffusion/MetaDiffusion.jl diff --git a/miniapps/benchmarks/stokes2D/RunStokesBench2D.jl b/miniapps/benchmarks/stokes2D/RunStokesBench2D.jl index 792b64cf..18ba84c2 100644 --- a/miniapps/benchmarks/stokes2D/RunStokesBench2D.jl +++ b/miniapps/benchmarks/stokes2D/RunStokesBench2D.jl @@ -1,14 +1,10 @@ -using JustRelax, Printf, LinearAlgebra +using JustRelax, JustRelax.JustRelax2D, Printf, LinearAlgebra using MPI: MPI using GLMakie using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) # :cpu | :CUDA | :AMDGPU -environment!(model) - # choose benchmark benchmark = :solcx diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl index 36650f17..123c69ce 100644 --- a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl +++ b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl @@ -50,7 +50,7 @@ function elastic_buildup(; ## Allocate arrays needed for every Stokes problem # general stokes arrays - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(ni) # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL=1 / √2.1) @@ -78,15 +78,16 @@ function elastic_buildup(; di, flow_bcs, ρg, - η, Gc, Kb, dt, igg; - iterMax=150e3, - nout=1000, - b_width=(4, 4, 1), - verbose=true, + kwargs = (; + iterMax=150e3, + nout=1000, + b_width=(4, 4, 1), + verbose=true, + ) ) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 7ea24726..4408b1a8 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -1,14 +1,8 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - +using JustRelax, JustRelax.JustRelax2D using JustPIC, JustPIC._2D const backend = CPUBackend -using ParallelStencil +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) # Load script dependencies @@ -44,10 +38,10 @@ function init_phases!(phases, particles) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, i, j]) == 0 && continue + @cell(index[ip, i, j]) == 0 && continue - x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) + x = @cell px[ip, i, j] + depth = -(@cell py[ip, i, j]) @cell phases[ip, i, j] = 2.0 if 0e0 ≤ depth ≤ 100e3 @@ -124,7 +118,7 @@ function main(igg, nx, ny) # Elliptical temperature anomaly init_phases!(pPhases, particles) phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -139,14 +133,10 @@ function main(igg, nx, ny) # Buoyancy forces & rheology ρg = @zeros(ni...), @zeros(ni...) - η = @ones(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - η_vep = copy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -160,7 +150,7 @@ function main(igg, nx, ny) ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") scatter!(ax1, Array(ρg[2][:]./9.81), Y./1e3) - scatter!(ax2, Array(log10.(η[:])), Y./1e3) + scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) # scatter!(ax2, Array(stokes.P[:]), Y./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) ylims!(ax2, minimum(xvi[2])./1e3, 0) @@ -181,12 +171,8 @@ function main(igg, nx, ny) # Stokes solver ---------------- args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) solve!( stokes, @@ -194,16 +180,16 @@ function main(igg, nx, ny) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 150e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) + kwargs = (; + iterMax = 150e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf) + ) ) dt = compute_dt(stokes, di) / 2 # ------------------------------ @@ -217,8 +203,8 @@ function main(igg, nx, ny) inject = check_injection(particles) inject && inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - + phase_ratios_center(phase_ratios, particles, grid, pPhases) + @show it += 1 t += dt @@ -227,7 +213,7 @@ function main(igg, nx, ny) nt = 5 fig = Figure(resolution = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") - heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)), colormap = :grayC) + heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) arrows!( ax, xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 6ec170e0..d8a2d849 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -1,6 +1,5 @@ # using CUDA using JustRelax, JustRelax.DataIO -import JustRelax.@cell # setup ParallelStencil.jl environment model = PS_Setup(:Threads, Float64, 2) @@ -43,10 +42,10 @@ function init_phases!(phases, particles, A) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, i, j]) == 0 && continue + @cell(index[ip, i, j]) == 0 && continue - x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) + x = @cell px[ip, i, j] + depth = -(@cell py[ip, i, j]) @cell phases[ip, i, j] = 2.0 if 0e0 ≤ depth ≤ 100e3 diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index e002e937..0b4a21ec 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -55,14 +55,14 @@ end function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Physical domain ------------------------------------ - ly = 40e3 # domain length in y - lx = 70e3 # domain length in x - ni = nx, ny # number of cells - li = lx, ly # domain length in x- and y- - di = @. li / ni # grid step in x- and -y - origin = 0.0, -ly # origin coordinates (15km f sticky air layer) + ly = 40e3 # domain length in y + lx = 70e3 # domain length in x + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / ni # grid step in x- and -y + origin = 0.0, -ly # origin coordinates (15km f sticky air layer) grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # Physical properties using GeoParams ---------------- diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl deleted file mode 100644 index 2bf96618..00000000 --- a/src/MetaJustRelax.jl +++ /dev/null @@ -1,173 +0,0 @@ -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 - ## 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! - - include(joinpath(@__DIR__, "phases/phases.jl")) - export PhaseRatio, fn_ratio, 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!, tensor_invariant! - - include(joinpath(@__DIR__, "stokes/Stokes3D.jl")) - export solve!, tensor_invariant! - - 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/stokes/MetaStokes.jl b/src/stokes/MetaStokes.jl deleted file mode 100644 index c768b9e3..00000000 --- a/src/stokes/MetaStokes.jl +++ /dev/null @@ -1,248 +0,0 @@ -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 - 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 - 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 - end -end - -function make_stokes_struct!() - @eval begin - struct StokesArrays{M<:AbstractStokesModel,A,B,C,T,nDim} - P::T - P0::T - V::A - ∇V::T - τ::B - ε::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) - ε_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(P),2}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),2}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),3}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),3}( - P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, deepcopy(τ), R - ) - 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/Stokes2D.jl b/src/stokes/Stokes2D.jl index 47b9b00d..cd11eb2d 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -317,15 +317,15 @@ function solve!( stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, η, Kb, dt, r, θ_dτ ) - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) + compute_ρg!(ρg[2], rheology, args) @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... ) ν = 1e-2 - @parallel (@idx ni) compute_viscosity!( - η, ν, @strain(stokes)..., args, rheology, viscosity_cutoff + compute_viscosity!( + stokes, ν, args, rheology, viscosity_cutoff ) compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) @@ -347,7 +347,7 @@ function solve!( θ_dτ, ) - @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) + center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) @hide_communication b_width begin # communication/computation overlap @@ -510,7 +510,6 @@ function solve!( if rem(iter, 5) == 0 compute_ρg!(ρg[2], phase_ratios, rheology, args) - # @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) end @parallel (@idx ni .+ 1) compute_strain_rate!( diff --git a/src/stokes/Stokes2D_old.jl b/src/stokes/Stokes2D_old.jl deleted file mode 100644 index cfe38ce3..00000000 --- a/src/stokes/Stokes2D_old.jl +++ /dev/null @@ -1,808 +0,0 @@ - -## 2D STOKES MODULE - -module Stokes2D - -using ImplicitGlobalGrid -using ..JustRelax -using CUDA, AMDGPU -using ParallelStencil -# using ParallelStencil.FiniteDifferences2D -using GeoParams, LinearAlgebra, Printf - -import JustRelax: elastic_iter_params!, PTArray, Velocity, SymmetricTensor -import JustRelax: - Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG -import JustRelax: compute_maxloc!, solve! -import JustRelax: mean_mpi, norm_mpi, maximum_mpi, minimum_mpi, backend - -@eval @init_parallel_stencil($backend, Float64, 2) - -include("../rheology/GeoParams.jl") -include("StressRotation.jl") -include("StressKernels.jl") -include("PressureKernels.jl") -include("VelocityKernels.jl") -include("StressKernels.jl") - -export solve! - -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 - τxx_o, τyy_o, τxy_o, τxy_o_c = stokes.τ_o.xx, - stokes.τ_o.yy, stokes.τ_o.xy, - stokes.τ_o.xy_c - - @parallel (@idx size(τxy)) multi_copy!( - (τxx_o, τyy_o, τxy_o, τxy_o_c), (τxx, τyy, τxy, τxy_c) - ) - return nothing -end - -## 2D VISCO-ELASTIC STOKES SOLVER - -# viscous solver -function JustRelax.solve!( - stokes::StokesArrays{Viscous,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, - di::NTuple{2,T}, - flow_bcs::FlowBoundaryConditions, - ρg, - η, - K, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 1), - verbose=true, -) where {A,B,C,D,T} - - # unpack - _dx, _dy = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - @parallel compute_P!( - stokes.P, stokes.P0, stokes.RP, stokes.∇V, η, K, dt, r, θ_dτ - ) - @parallel (@idx ni .+ 1) compute_τ!( - @stress(stokes)..., @strain(stokes)..., η, θ_dτ - ) - - @hide_communication b_width begin - @parallel compute_V!( - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - ητ, - _di..., - dt, - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(@velocity(stokes)...) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, - stokes.R.Ry, - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - ρg..., - _di..., - dt, - ) - Vmin, Vmax = extrema(stokes.V.Vx) - Pmin, Pmax = extrema(stokes.P) - push!( - norm_Rx, - norm_mpi(stokes.R.Rx) / (Pmax - Pmin) * lx / sqrt(length(stokes.R.Rx)), - ) - push!( - norm_Ry, - norm_mpi(stokes.R.Ry) / (Pmax - Pmin) * lx / sqrt(length(stokes.R.Ry)), - ) - push!( - norm_∇V, norm_mpi(stokes.∇V) / (Vmax - Vmin) * lx / sqrt(length(stokes.∇V)) - ) - err = maximum_mpi(norm_Rx[end], norm_Ry[end], norm_∇V[end]) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end - -# visco-elastic solver -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, - di::NTuple{2,T}, - flow_bcs, - ρg, - η, - G, - K, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 1), - verbose=true, -) where {A,B,C,D,T} - - # unpack - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - @parallel (@idx ni) compute_∇V!(stokes.∇V, stokes.V.Vx, stokes.V.Vy, _di...) - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - @parallel compute_P!( - stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, η, K, dt, r, θ_dτ - ) - @parallel (@idx ni) compute_τ!( - @stress(stokes)..., - @tensor(stokes.τ_o)..., - @strain(stokes)..., - η, - G, - θ_dτ, - dt, - ) - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - ητ, - _di..., - ) - # free slip boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, stokes.R.Ry, stokes.P, @stress(stokes)..., ρg..., _di... - ) - errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum_mpi(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end - -# GeoParams: general (visco-elasto-plastic) solver - -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, - di::NTuple{2,T}, - flow_bcs, - ρg, - η, - η_vep, - rheology::MaterialParams, - args, - dt, - igg::IGG; - viscosity_cutoff=(1e16, 1e24), - iterMax=10e3, - nout=500, - b_width=(4, 4, 0), - verbose=true, -) where {A,B,C,D,T} - - # unpack - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - Kb = get_Kb(rheology) - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - - # solver loop - wtime0 = 0.0 - λ = @zeros(ni...) - θ = @zeros(ni...) - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - @parallel compute_P!( - stokes.P, stokes.P0, stokes.R.RP, stokes.∇V, η, Kb, dt, r, θ_dτ - ) - - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) - - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - ν = 1e-2 - @parallel (@idx ni) compute_viscosity!( - η, ν, @strain(stokes)..., args, rheology, viscosity_cutoff - ) - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - - @parallel (@idx ni) compute_τ_nonlinear!( - @tensor_center(stokes.τ), - stokes.τ.II, - @tensor(stokes.τ_o), - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - stokes.P, - θ, - η, - η_vep, - λ, - tupleize(rheology), # needs to be a tuple - dt, - θ_dτ, - ) - - @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) - update_halo!(stokes.τ.xy) - - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - θ, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - ητ, - _di..., - dt, - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, - stokes.R.Ry, - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - ρg..., - _di..., - dt, - ) - - errs = maximum.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - stokes.P .= θ - - # accumulate plastic strain tensor - @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end - -## With phase ratios - -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, - di::NTuple{2,T}, - flow_bcs, - ρg, - η, - η_vep, - phase_ratios::PhaseRatio, - rheology, - args, - dt, - igg::IGG; - viscosity_cutoff=(1e16, 1e24), - iterMax=50e3, - iterMin=1e2, - viscosity_relaxation=1e-2, - free_surface=false, - nout=500, - b_width=(4, 4, 0), - verbose=true, -) where {A,B,C,D,T} - - # unpack - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - sizehint!(norm_Rx, Int(iterMax)) - sizehint!(norm_Ry, Int(iterMax)) - sizehint!(norm_∇V, Int(iterMax)) - sizehint!(err_evo1, Int(iterMax)) - sizehint!(err_evo2, Int(iterMax)) - - # solver loop - @copy stokes.P0 stokes.P - wtime0 = 0.0 - θ = @zeros(ni...) - λ = @zeros(ni...) - η0 = deepcopy(η) - do_visc = true - - for Aij in @tensor_center(stokes.ε_pl) - Aij .= 0.0 - end - Vx_on_Vy = @zeros(size(stokes.V.Vy)) - - while iter ≤ iterMax - iterMin < iter && err < ϵ && break - - wtime0 += @elapsed begin - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - - @parallel (@idx ni) compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - dt, - r, - θ_dτ, - ) - - if rem(iter, 5) == 0 - @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) - end - - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - - if rem(iter, nout) == 0 - @copy η0 η - end - if do_visc - @parallel (@idx ni) compute_viscosity!( - η, - viscosity_relaxation, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) - end - - @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, - θ, - η, - η_vep, - λ, - phase_ratios.center, - tupleize(rheology), # needs to be a tuple - dt, - θ_dτ, - ) - - @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) - update_halo!(stokes.τ.xy) - - @parallel (1:(size(stokes.V.Vy, 1) - 2), 1:size(stokes.V.Vy, 2)) interp_Vx_on_Vy!( - Vx_on_Vy, stokes.V.Vx - ) - - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - Vx_on_Vy, - θ, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - ητ, - _di..., - dt * free_surface, - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy) - end - end - - iter += 1 - - if iter % nout == 0 && iter > 1 - er_η = norm_mpi(@.(log10(η) - log10(η0))) - er_η < 1e-3 && (do_visc = false) - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, - stokes.R.Ry, - @velocity(stokes)..., - Vx_on_Vy, - stokes.P, - @stress(stokes)..., - ρg..., - _di..., - dt * free_surface, - ) - # errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - errs = ( - norm_mpi(stokes.R.Rx) / length(stokes.R.Rx), - norm_mpi(stokes.R.Ry) / length(stokes.R.Ry), - norm_mpi(stokes.R.RP) / length(stokes.R.RP), - ) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum_mpi(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - - if igg.me == 0 #&& ((verbose && err > ϵ) || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ && iter ≥ 20000 - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - stokes.P .= θ - - # accumulate plastic strain tensor - @parallel (@idx ni) accumulate_tensor!(stokes.EII_pl, @tensor_center(stokes.ε_pl), dt) - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end - -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - thermal::ThermalArrays, - pt_stokes::PTStokesCoeffs, - di::NTuple{2,T}, - flow_bcs, - ϕ, - ρg, - η, - η_vep, - phase_v, - phase_c, - args_η, - rheology::NTuple{N,MaterialParams}, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 1), - verbose=true, -) where {A,B,C,D,N,T} - - # unpack - - _di = inv.(di) - (; ϵ, r, θ_dτ, ηdτ) = pt_stokes - ni = size(stokes.P) - # ~preconditioner - ητ = deepcopy(η) - # @hide_communication b_width begin # communication/computation overlap - compute_maxloc!(ητ, η; window=(1, 1)) - update_halo!(ητ) - # end - - # errors - err = 2 * ϵ - iter = 0 - err_evo1 = Float64[] - err_evo2 = Float64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_∇V = Float64[] - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - @parallel (@idx ni) compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - η, - rheology, - phase_c, - dt, - r, - θ_dτ, - ) - @parallel (@idx ni .+ 1) compute_strain_rate!( - @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... - ) - @parallel (@idx ni) compute_ρg!( - ρg[end], ϕ, rheology, (T=thermal.Tc, P=stokes.P) - ) - @parallel (@idx ni) compute_τ_gp!( - @tensor_center(stokes.τ), - stokes.τ.II, - @tensor(stokes.τ_o), - @strain(stokes), - η, - η_vep, - thermal.T, - phase_v, - phase_c, - args_η, - rheology, # needs to be a tuple - dt, - θ_dτ, - ) - @parallel center2vertex!(stokes.τ.xy, stokes.τ.xy_c) - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - pt_stokes.ηdτ, - ρg..., - ητ, - _di..., - dt, - ) - # apply boundary conditions boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - @parallel (@idx ni) compute_Res!( - stokes.R.Rx, - stokes.R.Ry, - @velocity(stokes)..., - stokes.P, - @stress(stokes)..., - ρg..., - _di..., - dt, - ) - errs = maximum_mpi.((abs.(stokes.R.Rx), abs.(stokes.R.Ry), abs.(stokes.R.RP))) - push!(norm_Rx, errs[1]) - push!(norm_Ry, errs[2]) - push!(norm_∇V, errs[3]) - err = maximum_mpi(errs) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && (verbose || iter == iterMax) - @printf( - "Total steps = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[end], - norm_Ry[end], - norm_∇V[end] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_∇V=norm_∇V, - ) -end - -end # END OF MODULE diff --git a/src/stokes/Stokes3D_old.jl b/src/stokes/Stokes3D_old.jl deleted file mode 100644 index 2495b47d..00000000 --- a/src/stokes/Stokes3D_old.jl +++ /dev/null @@ -1,561 +0,0 @@ -## 3D ELASTICITY MODULE - -module Stokes3D - -using ImplicitGlobalGrid -using ParallelStencil -using ParallelStencil.FiniteDifferences3D -using JustRelax -using CUDA, AMDGPU -using LinearAlgebra -using Printf -using GeoParams - -import JustRelax: PTArray, Velocity, SymmetricTensor, pureshear_bc! -import JustRelax: - Residual, StokesArrays, PTStokesCoeffs, AbstractStokesModel, ViscoElastic, IGG -import JustRelax: compute_maxloc!, solve! -import JustRelax: mean_mpi, norm_mpi, minimum_mpi, maximum_mpi, backend - -@eval @init_parallel_stencil($backend, Float64, 3) - -include("../rheology/GeoParams.jl") -include("StressRotation.jl") -include("StressKernels.jl") -include("PressureKernels.jl") -include("VelocityKernels.jl") -include("StressKernels.jl") - -export solve!, pureshear_bc! - -@parallel function update_τ_o!( - τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, τxx, τyy, τzz, τxy, τxz, τyz -) - @all(τxx_o) = @all(τxx) - @all(τyy_o) = @all(τyy) - @all(τzz_o) = @all(τzz) - @all(τxy_o) = @all(τxy) - @all(τxz_o) = @all(τxz) - @all(τyz_o) = @all(τyz) - return nothing -end - -function update_τ_o!(stokes::StokesArrays{ViscoElastic,A,B,C,D,3}) where {A,B,C,D} - @parallel update_τ_o!(@tensor(stokes.τ_o)..., @stress(stokes)...) -end - -## BOUNDARY CONDITIONS - -function JustRelax.pureshear_bc!( - stokes::StokesArrays, di::NTuple{3,T}, li::NTuple{3,T}, εbg -) where {T} - # unpack - Vx, _, Vz = stokes.V.Vx, stokes.V.Vy, stokes.V.Vz - dx, _, dz = di - lx, _, lz = li - # Velocity pure shear boundary conditions - stokes.V.Vx .= PTArray([ - -εbg * ((i - 1) * dx - 0.5 * lx) for i in 1:size(Vx, 1), j in 1:size(Vx, 2), - k in 1:size(Vx, 3) - ]) - return stokes.V.Vz .= PTArray([ - εbg * ((k - 1) * dz - 0.5 * lz) for i in 1:size(Vz, 1), j in 1:size(Vz, 2), - k in 1:size(Vz, 3) - ]) -end - -## 3D VISCO-ELASTIC STOKES SOLVER - -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, - di::NTuple{3,T}, - flow_bcs, - ρg, - η, - K, - G, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 4), - verbose=true, -) where {A,B,C,D,T} - - # solver related - ϵ = pt_stokes.ϵ - # geometry - _di = @. 1 / di - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - compute_maxloc!(ητ, η) - update_halo!(ητ) - - # errors - err = 2 * ϵ - iter = 0 - cont = 0 - err_evo1 = Float64[] - err_evo2 = Int64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_Rz = Float64[] - norm_∇V = Float64[] - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - @parallel compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - η, - K, - dt, - pt_stokes.r, - pt_stokes.θ_dτ, - ) - @parallel (@idx ni .+ 1) compute_strain_rate!( - stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... - ) - @parallel (@idx ni .+ 1) compute_τ!( - @stress(stokes)..., - @tensor(stokes.τ_o)..., - @strain(stokes)..., - η, - G, - dt, - pt_stokes.θ_dτ, - ) - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.R.Rz, - stokes.P, - ρg..., - @stress(stokes)..., - ητ, - pt_stokes.ηdτ, - _di..., - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - cont += 1 - push!(norm_Rx, maximum_mpi(abs.(stokes.R.Rx))) - push!(norm_Ry, maximum_mpi(abs.(stokes.R.Ry))) - push!(norm_Rz, maximum_mpi(abs.(stokes.R.Rz))) - push!(norm_∇V, maximum_mpi(abs.(stokes.R.RP))) - err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && ((verbose && err > ϵ) || iter == iterMax) - @printf( - "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[cont], - norm_Ry[cont], - norm_Rz[cont], - norm_∇V[cont] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - av_time = wtime0 / (iter - 1) # average time per iteration - update_τ_o!(stokes) # copy τ into τ_o - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_Rz=norm_Rz, - norm_∇V=norm_∇V, - time=wtime0, - av_time=av_time, - ) -end - -## 3D VISCO-ELASTO-PLASTIC STOKES SOLVER WITH GeoParams.jl - -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, - di::NTuple{3,T}, - flow_bcs::FlowBoundaryConditions, - ρg, - η, - η_vep, - rheology::MaterialParams, - args, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 4), - verbose=true, -) where {A,B,C,D,T} - - # solver related - ϵ = pt_stokes.ϵ - # geometry - _di = @. 1 / di - ni = size(stokes.P) - - # ~preconditioner - ητ = deepcopy(η) - - # errors - err = 2 * ϵ - iter = 0 - cont = 0 - err_evo1 = Float64[] - err_evo2 = Int64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_Rz = Float64[] - norm_∇V = Float64[] - - Kb = get_Kb(rheology) - G = get_shear_modulus(rheology) - @copy stokes.P0 stokes.P - λ = @zeros(ni...) - θ = @zeros(ni...) - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - compute_maxloc!(ητ, η) - update_halo!(ητ) - - @parallel (@idx ni) compute_∇V!(stokes.∇V, @velocity(stokes)..., _di...) - @parallel (@idx ni) compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - Kb, - dt, - pt_stokes.r, - pt_stokes.θ_dτ, - ) - @parallel (@idx ni) compute_strain_rate!( - stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... - ) - - # # Update buoyancy - # @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) - - ν = 1e-3 - @parallel (@idx ni) compute_viscosity!( - η, - 1.0, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) - - @parallel (@idx ni) compute_τ_nonlinear!( - @tensor_center(stokes.τ), - stokes.τ.II, - @tensor(stokes.τ_o), - @strain(stokes), - @tensor_center(stokes.ε_pl), - stokes.EII_pl, - stokes.P, - θ, - η, - @ones(ni...), - λ, - tupleize(rheology), # needs to be a tuple - dt, - pt_stokes.θ_dτ, - ) - - @parallel (@idx ni .+ 1) compute_τ_vertex!( - @shear(stokes.τ)..., - @shear(stokes.τ_o)..., - @shear(stokes.ε)..., - η_vep, - G, - dt, - pt_stokes.θ_dτ, - ) - - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - stokes.R.Rx, - stokes.R.Ry, - stokes.R.Rz, - stokes.P, - ρg..., - @stress(stokes)..., - ητ, - pt_stokes.ηdτ, - _di..., - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bcs) - update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) - end - end - - iter += 1 - if iter % nout == 0 && iter > 1 - cont += 1 - push!(norm_Rx, maximum_mpi(abs.(stokes.R.Rx))) - push!(norm_Ry, maximum_mpi(abs.(stokes.R.Ry))) - push!(norm_Rz, maximum_mpi(abs.(stokes.R.Rz))) - push!(norm_∇V, maximum_mpi(abs.(stokes.R.RP))) - err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && (verbose || iter == iterMax) - @printf( - "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[cont], - norm_Ry[cont], - norm_Rz[cont], - norm_∇V[cont] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - av_time = wtime0 / (iter - 1) # average time per iteration - update_τ_o!(stokes) # copy τ into τ_o - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_Rz=norm_Rz, - norm_∇V=norm_∇V, - time=wtime0, - av_time=av_time, - ) -end - -# GeoParams and multiple phases -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, - di::NTuple{3,T}, - flow_bc::FlowBoundaryConditions, - ρg, - η, - η_vep, - phase_ratios::PhaseRatio, - rheology::NTuple{N,AbstractMaterialParamsStruct}, - args, - dt, - igg::IGG; - iterMax=10e3, - nout=500, - b_width=(4, 4, 4), - verbose=true, - viscosity_cutoff=(-Inf, Inf), -) where {A,B,C,D,T,N} - - ## UNPACK - - # solver related - ϵ = pt_stokes.ϵ - # geometry - _di = @. 1 / di - ni = size(stokes.P) - - # errors - err = Inf - iter = 0 - cont = 0 - err_evo1 = Float64[] - err_evo2 = Int64[] - norm_Rx = Float64[] - norm_Ry = Float64[] - norm_Rz = Float64[] - norm_∇V = Float64[] - - @copy stokes.P0 stokes.P - λ = @zeros(ni...) - θ = @zeros(ni...) - - # solver loop - wtime0 = 0.0 - while iter < 2 || (err > ϵ && iter ≤ iterMax) - wtime0 += @elapsed begin - # ~preconditioner - ητ = deepcopy(η) - compute_maxloc!(ητ, η) - 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...) - @parallel (@idx ni) compute_P!( - stokes.P, - stokes.P0, - stokes.R.RP, - stokes.∇V, - ητ, - rheology, - phase_ratios.center, - dt, - pt_stokes.r, - pt_stokes.θ_dτ, - ) - - @parallel (@idx ni) compute_strain_rate!( - stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... - ) - - # Update buoyancy - @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) - - # Update viscosity - ν = 1e-2 - @parallel (@idx ni) compute_viscosity!( - η, - ν, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) - - @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, - θ, - η, - η_vep, - λ, - phase_ratios.center, - tupleize(rheology), # needs to be a tuple - dt, - pt_stokes.θ_dτ, - ) - - @parallel (@idx ni .+ 1) center2vertex!( - stokes.τ.yz, - stokes.τ.xz, - stokes.τ.xy, - stokes.τ.yz_c, - stokes.τ.xz_c, - stokes.τ.xy_c, - ) - update_halo!(stokes.τ.yz, stokes.τ.xz, stokes.τ.xy) - - # @parallel (@idx ni .+ 1) compute_τ_vertex!( - # @shear(stokes.τ)..., @shear(stokes.ε)..., η_vep, pt_stokes.θ_dτ - # ) - - @hide_communication b_width begin # communication/computation overlap - @parallel compute_V!( - @velocity(stokes)..., - @residuals(stokes.R)..., - θ, - ρg..., - @stress(stokes)..., - ητ, - pt_stokes.ηdτ, - _di..., - ) - # apply boundary conditions - flow_bcs!(stokes, flow_bc) - update_halo!(@velocity(stokes)...) - end - end - - stokes.P .= θ - - iter += 1 - if iter % nout == 0 && iter > 1 - cont += 1 - for (norm_Ri, Ri) in zip((norm_Rx, norm_Ry, norm_Rz), @residuals(stokes.R)) - push!(norm_Ri, maximum(abs.(Ri))) - end - push!(norm_∇V, maximum(abs.(stokes.R.RP))) - err = max(norm_Rx[cont], norm_Ry[cont], norm_Rz[cont], norm_∇V[cont]) - push!(err_evo1, err) - push!(err_evo2, iter) - if igg.me == 0 && (verbose || iter == iterMax) - @printf( - "iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", - iter, - err, - norm_Rx[cont], - norm_Ry[cont], - norm_Rz[cont], - norm_∇V[cont] - ) - end - isnan(err) && error("NaN(s)") - end - - if igg.me == 0 && err ≤ ϵ - println("Pseudo-transient iterations converged in $iter iterations") - end - end - - av_time = wtime0 / (iter - 1) # average time per iteration - - stokes.P .= θ - - return ( - iter=iter, - err_evo1=err_evo1, - err_evo2=err_evo2, - norm_Rx=norm_Rx, - norm_Ry=norm_Ry, - norm_Rz=norm_Rz, - norm_∇V=norm_∇V, - time=wtime0, - av_time=av_time, - ) -end - -end # END OF MODULE diff --git a/src/thermal_diffusion/MetaDiffusion.jl b/src/thermal_diffusion/MetaDiffusion.jl deleted file mode 100644 index 412fe682..00000000 --- a/src/thermal_diffusion/MetaDiffusion.jl +++ /dev/null @@ -1,216 +0,0 @@ -function make_thermal_arrays!(ndim) - flux_sym1 = (:qTx, :qTy, :qTz) - flux_sym2 = (:qTx2, :qTy2, :qTz2) - flux1 = [:($(flux_sym1[i])::_T) for i in 1:ndim] - flux2 = [:($(flux_sym2[i])::_T) for i in 1:ndim] - - @eval begin - struct ThermalArrays{_T} - T::_T # Temperature @ grid nodes - Tc::_T # Temperature @ cell centers - ΔT::_T - Told::_T - dT_dt::_T - $(flux1...) - $(flux2...) - H::_T # source terms - shear_heating::_T # shear heating terms - ResT::_T - - function ThermalArrays(ni::NTuple{1,Integer}) - nx, = ni - T, ΔT, Told = @zeros(ni...), @zeros(ni...), @zeros(ni...) - dT_dt = @zeros(nx - 2) - qTx = @zeros(nx - 1) - qTx2 = @zeros(nx - 1) - ResT = @zeros(nx - 2) - return new{typeof(T)}(T, ΔT, Told, dT_dt, qTx, qTx2, ResT) - end - - function ThermalArrays(ni::NTuple{2,Integer}) - nx, ny = ni - T = @zeros(nx + 3, ny + 1) - ΔT = @zeros(nx + 3, ny + 1) - Told = @zeros(nx + 3, ny + 1) - Tc = @zeros(ni...) - H = @zeros(ni...) - shear_heating = @zeros(ni...) - dT_dt = @zeros(nx + 1, ny - 1) - qTx = @zeros(nx + 2, ny - 1) - qTy = @zeros(nx + 1, ny) - qTx2 = @zeros(nx + 2, ny - 1) - qTy2 = @zeros(nx + 1, ny) - ResT = @zeros(nx + 1, ny - 1) - return new{typeof(T)}( - T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTx2, qTy2, H, shear_heating, ResT - ) - end - - function ThermalArrays(ni::NTuple{3,Integer}) - nx, ny, nz = ni - T, ΔT, Told = @zeros(ni .+ 1...), @zeros(ni .+ 1...), @zeros(ni .+ 1...) - Tc = @zeros(ni...) - H = @zeros(ni...) - shear_heating = @zeros(ni...) - dT_dt = @zeros(ni .- 1) - qTx = @zeros(nx, ny - 1, nz - 1) - qTy = @zeros(nx - 1, ny, nz - 1) - qTz = @zeros(nx - 1, ny - 1, nz) - qTx2 = @zeros(nx, ny - 1, nz - 1) - qTy2 = @zeros(nx - 1, ny, nz - 1) - qTz2 = @zeros(nx - 1, ny - 1, nz) - ResT = @zeros((ni .- 1)...) - return new{typeof(T)}( - T, - Tc, - ΔT, - Told, - dT_dt, - qTx, - qTy, - qTz, - qTx2, - qTy2, - qTz2, - H, - shear_heating, - ResT, - ) - end - end - end -end - -function make_PTthermal_struct!() - @eval begin - struct PTThermalCoeffs{T,M,nDim} - CFL::T - ϵ::T - max_lxyz::T - max_lxyz2::T - Vpdτ::T - θr_dτ::M - dτ_ρ::M - - function PTThermalCoeffs( - K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - max_lxyz2 = max_lxyz^2 - Re = @. π + √(π * π + ρCp * max_lxyz2 / K / dt) # Numerical Reynolds number - θr_dτ = @. max_lxyz / Vpdτ / Re - dτ_ρ = @. Vpdτ * max_lxyz / K / Re - - return new{eltype(Vpdτ),typeof(dτ_ρ),nDim}( - CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ - ) - end - - function PTThermalCoeffs( - CFL::T, ϵ::T, max_lxyz::T, Vpdτ::T, θr_dτ::M, dτ_ρ::M - ) where {T,M} - nDim = length(size(θr_dτ)) - return new{T,M,nDim}(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) - end - end - - # with phase ratios - function PTThermalCoeffs( - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) - - idx = ntuple(i -> 1:ni[i], Val(nDim)) - @parallel (idx) compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase_ratios.center, args, max_lxyz, Vpdτ, inv(dt) - ) - - return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) - end - - # without phase ratios - function PTThermalCoeffs( - # function PTThermalCoeffs( - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) - - idx = ntuple(i -> 1:ni[i], Val(nDim)) - @parallel (idx) compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, inv(dt) - ) - - return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) - end - - @parallel_indices (I...) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt - ) - _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, I... - ) - - return nothing - end - - @parallel_indices (I...) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt - ) - _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, I... - ) - - return nothing - end - - function _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} - ) where {N} - args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) - phase_ij = phase[Idx...] - ρCp = compute_ρCp(rheology, phase_ij, args_ij) - _K = inv(fn_ratio(compute_conductivity, rheology, phase_ij, args_ij)) - - _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number - θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re - dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re - - return nothing - end - - function _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} - ) where {N} - args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) - - ρCp = compute_ρCp(rheology, args_ij) - _K = inv(compute_conductivity(rheology, args_ij)) - - _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number - θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re - dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re - - return nothing - end - end -end From a8e0dc7b26d000e8394b0e3172ab5c4d6aed80a3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 13:12:14 +0200 Subject: [PATCH 05/91] format --- src/JustRelax.jl | 13 ------------- src/stokes/Stokes2D.jl | 11 +++++++---- src/types/heat_diffusion.jl | 15 ++++++++++++++- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 787ffa65..a2a7f3b3 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -18,23 +18,10 @@ export BackendTrait, CPUBackendTrait include("topology/Topology.jl") export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g -# include("MiniKernels.jl") - include("phases/CellArrays.jl") export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! -# include("rheology/StressUpdate.jl") -# export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear! - include("JustRelax_CPU.jl") -# include("MetaJustRelax.jl") - -# include("stokes/MetaStokes.jl") -# export PS_Setup, environment!, ps_reset! - -# include("thermal_diffusion/MetaDiffusion.jl") - -# include("thermal_diffusion/Rheology.jl") include("IO/DataIO.jl") diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index cd11eb2d..5d9c3b9f 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -324,9 +324,7 @@ function solve!( ) ν = 1e-2 - compute_viscosity!( - stokes, ν, args, rheology, viscosity_cutoff - ) + compute_viscosity!(stokes, ν, args, rheology, viscosity_cutoff) compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) @@ -521,7 +519,12 @@ function solve!( end if do_visc compute_viscosity!( - stokes,viscosity_relaxation, phase_ratios, args, rheology, viscosity_cutoff + stokes, + viscosity_relaxation, + phase_ratios, + args, + rheology, + viscosity_cutoff, ) end diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index 28fdb90a..8340c6dd 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -28,7 +28,20 @@ struct ThermalArrays{_T} qTy2 = @zeros(nx + 1, ny) ResT = @zeros(nx + 1, ny - 1) return new{typeof(T)}( - T, Tc, ΔT, Told, dT_dt, qTx, qTy, nothing, qTx2, qTy2, nothing, H, shear_heating, ResT + T, + Tc, + ΔT, + Told, + dT_dt, + qTx, + qTy, + nothing, + qTx2, + qTy2, + nothing, + H, + shear_heating, + ResT, ) end From 915f35fdc60010d31a184f5ac0d93a5484bab0ac Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 18:44:40 +0200 Subject: [PATCH 06/91] CUDA extension --- Project.toml | 8 +- ext/JustRelaxCUDAExt.jl | 177 +++++++++++++++++++ src/Interpolations.jl | 7 +- src/JustRelax.jl | 11 +- src/JustRelax_CPU.jl | 12 +- src/boundaryconditions/BoundaryConditions.jl | 82 ++------- src/common.jl | 6 +- src/phases/phases.jl | 77 ++------ src/rheology/Viscosity.jl | 22 ++- src/stokes/Stokes2D.jl | 31 ++-- src/thermal_diffusion/DiffusionPT_solver.jl | 2 +- src/types/heat_diffusion.jl | 2 + src/types/phases.jl | 18 ++ src/types/stokes.jl | 11 +- src/types/traits.jl | 2 + 15 files changed, 294 insertions(+), 174 deletions(-) create mode 100644 ext/JustRelaxCUDAExt.jl create mode 100644 src/types/phases.jl diff --git a/Project.toml b/Project.toml index dc2513e4..9e0d7138 100644 --- a/Project.toml +++ b/Project.toml @@ -6,11 +6,11 @@ version = "0.1.1" [deps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0" +JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" @@ -21,6 +21,12 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" +[weakdeps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" + +[extensions] +JustRelaxCUDAExt = "CUDA" + [compat] AMDGPU = "0.6, 0.7, 0.8" Adapt = "3.7.2" diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl new file mode 100644 index 00000000..bb7d3fa0 --- /dev/null +++ b/ext/JustRelaxCUDAExt.jl @@ -0,0 +1,177 @@ +module JustRelaxCUDAExt + +using CUDA +using JustRelax: JustRelax +import JustRelax: PTArray + +JustRelax.PTArray(::Type{CUDABackend}) = CuArray + +module JustRelax2D + + using JustRelax: JustRelax + using CUDA + using StaticArrays + using CellArrays + using ParallelStencil, ParallelStencil.FiniteDifferences2D + using ImplicitGlobalGrid + using GeoParams, LinearAlgebra, Printf + using MPI + + import JustRelax: + IGG, + BackendTrait, + CPUBackendTrait, + backend, + CPUBackend, + Geometry, + @cell + + @init_parallel_stencil(CUDA, Float64, 2) + + include("../src/common.jl") + include("../src/stokes/Stokes2D.jl") + + # add CUDA traits + struct CUDABackendTrait <: BackendTrait end + + @inline backend(::CuArray) = CUDABackendTrait() + @inline backend(::Type{<:CuArray}) = CUDABackendTrait() + + # Types + function JustRelax.JustRelax2D.StokesArrays( + ::Type{CUDABackend}, ni::Vararg{Integer,N} + ) where {N} + return StokesArrays(tuple(ni...)) + end + + function JustRelax.JustRelax2D.StokesArrays( + ::Type{CUDABackend}, ni::NTuple{N,Integer} + ) where {N} + return StokesArrays(ni) + end + + function JustRelax.JustRelax2D.ThermalArrays( + ::Type{CUDABackend}, ni::NTuple{N,Number} + ) where {N} + return ThermalArrays(ni...) + end + + function JustRelax.JustRelax2D.ThermalArrays( + ::Type{CUDABackend}, ni::Vararg{Number,N} + ) where {N} + return ThermalArrays(ni...) + end + + function JustRelax.JustRelax2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) + end + + # Boundary conditions + function JustRelax.JustRelax2D.flow_bcs!( + ::CUDABackendTrait, stokes::StokesArrays, bcs + ) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function flow_bcs!( + ::CUDABackendTrait, stokes::StokesArrays, bcs + ) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function JustRelax.JustRelax2D.thermal_bcs!( + ::CUDABackendTrait, thermal::ThermalArrays, bcs + ) + return thermal_bcs!(thermal.T, bcs) + end + + function thermal_bcs!( + ::CUDABackendTrait, thermal::ThermalArrays, bcs + ) + return thermal_bcs!(thermal.T, bcs) + end + + # Phases + function JustRelax.JustRelax2D.phase_ratios_center( + ::CUDABackendTrait, phase_ratios::PhaseRatio, particles, grid::Geometry, phases + ) + return _phase_ratios_center(phase_ratios, particles, grid, phases) + end + + # Rheology + ## viscosity + function JustRelax.JustRelax2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function JustRelax.JustRelax2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function JustRelax.JustRelax2D.compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + ## Stress + JustRelax.JustRelax2D.tensor_invariant!(A::SymmetricTensor) = tensor_invariant!(A) + + ## Buoyancy forces + function JustRelax.JustRelax2D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) + end + function JustRelax.JustRelax2D.compute_ρg!( + ρg::CuArray, phase_ratios::PhaseRatio, rheology, args + ) + return compute_ρg!(ρg, phase_ratios, rheology, args) + end + + # Interpolations + function JustRelax.JustRelax2D.temperature2center!( + ::CUDABackendTrait, thermal::ThermalArrays + ) + return _temperature2center!(thermal) + end + function JustRelax.JustRelax2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) + end + function JustRelax.JustRelax2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + return center2vertex!(vertex, center) + end + + function JustRelax.JustRelax2D.center2vertex!( + vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T + ) where {T<:CuArray} + return center2vertex!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + end + + # Solvers + JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) + + # Utils + JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I::IGG) = compute_dt(S, di, dt_diff, I::IGG) + +end + +end diff --git a/src/Interpolations.jl b/src/Interpolations.jl index 52fb268c..733ad2ee 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -16,7 +16,12 @@ end # From cell vertices to cell center -function temperature2center!(thermal::ThermalArrays) +temperature2center!(thermal) = temperature2center!(backend(thermal), thermal) +function temperature2center!(::CPUBackendTrait, thermal::ThermalArrays) + return _temperature2center!(thermal) +end + +function _temperature2center!(thermal::ThermalArrays) @parallel (@idx size(thermal.Tc)...) temperature2center_kernel!(thermal.Tc, thermal.T) return nothing end diff --git a/src/JustRelax.jl b/src/JustRelax.jl index a2a7f3b3..ffb01d1b 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -12,6 +12,15 @@ using StaticArrays function solve!() end +struct CPUBackend end +struct AMDGPUBackend end + +PTArray() = Array +PTArray(::Type{CPUBackend}) = Array +PTArray(::T) where {T} = error(ArgumentError("Unknown backend $T")) + +export PTArray, CPUBackend, CUDABackend, AMDGPUBackend + include("types/traits.jl") export BackendTrait, CPUBackendTrait @@ -23,6 +32,6 @@ export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! include("JustRelax_CPU.jl") -include("IO/DataIO.jl") +# include("IO/DataIO.jl") end # module diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index a5ea5eab..0500e9e4 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -8,14 +8,10 @@ using ImplicitGlobalGrid using GeoParams, LinearAlgebra, Printf using MPI -import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend @init_parallel_stencil(Threads, Float64, 2) -export PTArray - -PTArray() = Array - include("common.jl") include("stokes/Stokes2D.jl") export solve! @@ -32,14 +28,10 @@ using ImplicitGlobalGrid using GeoParams, LinearAlgebra, Printf using MPI -import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend @init_parallel_stencil(Threads, Float64, 3) -export PTArray - -PTArray() = Array - include("common.jl") include("stokes/Stokes3D.jl") export solve! diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 32e4c7ab..99cf16d9 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -50,7 +50,14 @@ end Apply the prescribed heat boundary conditions `bc` on the `T` """ -function thermal_bcs!(T, bcs::TemperatureBoundaryConditions) +thermal_bcs!(thermal, bcs) = thermal_bcs!(backend(thermal), thermal, bcs) +function thermal_bcs!( + ::CPUBackendTrait, thermal::ThermalArrays, bcs::FlowBoundaryConditions +) + return thermal_bcs!(thermal.T, bcs) +end + +function thermal_bcs!(T::AbstractArray, bcs::TemperatureBoundaryConditions) n = bc_index(T) # no flux boundary conditions @@ -64,12 +71,15 @@ end Apply the prescribed flow boundary conditions `bc` on the `stokes` """ -flow_bcs!(stokes, bcs::FlowBoundaryConditions) = _flow_bcs!(bcs, @velocity(stokes)) -function flow_bcs!(bcs::FlowBoundaryConditions, V::Vararg{T,N}) where {T,N} +flow_bcs!(stokes, bcs) = flow_bcs!(backend(stokes), stokes, bcs) +function flow_bcs!(::CPUBackendTrait, stokes, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end +function flow_bcs!(bcs, V::Vararg{T,N}) where {T,N} return _flow_bcs!(bcs, tuple(V...)) end -function _flow_bcs!(bcs::FlowBoundaryConditions, V) +function _flow_bcs!(bcs, V) n = bc_index(V) # no slip boundary conditions do_bc(bcs.no_slip) && (@parallel (@idx n) no_slip!(V..., bcs.no_slip)) @@ -214,57 +224,6 @@ end return nothing end -@parallel_indices (i) function periodic_boundaries!(Ax, Ay, bc) - @inbounds begin - if i ≤ size(Ax, 1) - bc.bot && (Ax[i, 1] = Ax[i, end - 1]) - bc.top && (Ax[i, end] = Ax[i, 2]) - end - if i ≤ size(Ay, 2) - bc.left && (Ay[1, i] = Ay[end - 1, i]) - bc.right && (Ay[end, i] = Ay[2, i]) - end - end - return nothing -end - -@parallel_indices (i) function periodic_boundaries!( - T::_T, bc -) where {_T<:AbstractArray{<:Any,2}} - @inbounds begin - if i ≤ size(T, 1) - bc.bot && (T[i, 1] = T[i, end - 1]) - bc.top && (T[i, end] = T[i, 2]) - end - if i ≤ size(T, 2) - bc.left && (T[1, i] = T[end - 1, i]) - bc.right && (T[end, i] = T[2, i]) - end - end - return nothing -end - -@parallel_indices (i, j) function periodic_boundaries!( - T::_T, bc -) where {_T<:AbstractArray{<:Any,3}} - nx, ny, nz = size(T) - @inbounds begin - if i ≤ nx && j ≤ ny - bc.bot && (T[i, j, 1] = T[i, j, end - 1]) - bc.top && (T[i, j, end] = T[i, j, 2]) - end - if i ≤ ny && j ≤ nz - bc.left && (T[1, i, j] = T[end - 1, i, j]) - bc.right && (T[end, i, j] = T[2, i, j]) - end - if i ≤ nx && j ≤ nz - bc.front && (T[i, 1, j] = T[i, end - 1, j]) - bc.back && (T[i, end, j] = T[i, 2, j]) - end - end - return nothing -end - function pureshear_bc!( stokes::StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg ) where {T} @@ -307,19 +266,6 @@ function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{2,T}}, Vx, Vy) where return nothing end -function thermal_boundary_conditions!( - insulation::NamedTuple, T::AbstractArray{_T,2} -) where {_T} - insulation_x, insulation_y = insulation - - nx, ny = size(T) - - insulation_x && (@parallel (1:ny) free_slip_x!(T)) - insulation_y && (@parallel (1:nx) free_slip_y!(T)) - - return nothing -end - # 3D KERNELS @parallel_indices (j, k) function free_slip_x!(A::AbstractArray{T,3}) where {T} diff --git a/src/common.jl b/src/common.jl index 9364141a..1d961fb0 100644 --- a/src/common.jl +++ b/src/common.jl @@ -4,6 +4,9 @@ export StokesArrays, PTStokesCoeffs include("types/heat_diffusion.jl") export ThermalArrays, PTThermalCoeffs +include("types/phases.jl") +export PhaseRatio + include("Utils.jl") export @allocate, @add, @@ -32,7 +35,8 @@ export FlowBoundaryConditions, include("MiniKernels.jl") include("phases/phases.jl") -export PhaseRatio, fn_ratio, phase_ratios_center +export fn_ratio, phase_ratios_center +# export PhaseRatio, fn_ratio, phase_ratios_center include("rheology/BuoyancyForces.jl") export compute_ρg! diff --git a/src/phases/phases.jl b/src/phases/phases.jl index 91aa0d73..7a2fb037 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -1,21 +1,3 @@ - -struct Phases{T} - vertex::T - center::T -end - -struct PhaseRatio{T} - vertex::T - center::T - - function PhaseRatio(ni, num_phases) - center = @fill(0.0, ni..., celldims = (num_phases,)) - vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) - T = typeof(center) - return new{T}(vertex, center) - end -end - """ nphases(x::PhaseRatio) @@ -34,53 +16,6 @@ end return N end -# # ParallelStencil launch kernel -# @parallel_indices (I...) function phase_ratios_center(x, phases) -# phase_ratios_center(x, phases, I...) -# return nothing -# end - -# """ -# phase_ratios_center(x::PhaseRatio, cell::Vararg{Int, N}) - -# Compute the phase ratios at the center of the cell `cell` in `x::PhaseRatio`. -# """ -# function phase_ratios_center(x::PhaseRatio, phases, cell::Vararg{Int,N}) where {N} -# return phase_ratios_center(x.center, phases, cell...) -# end - -# @inline function phase_ratios_center(x::CellArray, phases, cell::Vararg{Int,N}) where {N} -# # total number of material phases -# num_phases = nphases(x) -# # number of active particles in this cell -# n = 0 -# for j in cellaxes(phases) -# n += isinteger(@cell(phases[j, cell...])) && @cell(phases[j, cell...]) != 0 -# end -# _n = inv(n) -# # compute phase ratios -# ratios = _phase_ratios_center(phases, num_phases, _n, cell...) -# for (i, ratio) in enumerate(ratios) -# @cell x[i, cell...] = ratio -# end -# end - -# @generated function _phase_ratios_center( -# phases, ::Val{N1}, _n, cell::Vararg{Int,N2} -# ) where {N1,N2} -# quote -# Base.@_inline_meta -# Base.@nexprs $N1 i -> reps_i = begin -# c = 0 -# for j in cellaxes(phases) -# c += @cell(phases[j, cell...]) == i -# end -# c * _n -# end -# Base.@ncall $N1 tuple reps -# end -# end - """ fn_ratio(fn::F, rheology::NTuple{N, AbstractMaterialParamsStruct}, ratio) where {N, F} @@ -112,7 +47,17 @@ end end end -function phase_ratios_center(phase_ratios, particles, grid::Geometry, phases) +function phase_ratios_center(phase_ratios, particles, grid, phases) + return phase_ratios_center(backend(phase_ratios), phase_ratios, particles, grid, phases) +end + +function phase_ratios_center( + ::CPUBackendTrait, phase_ratios::PhaseRatio, particles, grid::Geometry, phases +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +function _phase_ratios_center(phase_ratios::PhaseRatio, particles, grid::Geometry, phases) ni = size(phases) @parallel (@idx ni) phase_ratios_center_kernel( phase_ratios.center, particles.coords, grid.xci, grid.di, phases diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 660e7f74..4ec06176 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,5 +1,14 @@ ## 2D KERNELS -function compute_viscosity!(stokes::StokesArrays, ν, args, rheology, cutoff) +function compute_viscosity!(stokes, ν, args, rheology, cutoff) + return compute_viscosity!(backend(stokes), ν, args, rheology, cutoff) +end +function compute_viscosity!( + ::CPUBackendTrait, stokes, ν, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end + +function _compute_viscosity!(stokes::StokesArrays, ν, args, rheology, cutoff) ni = size(stokes.viscosity.η) @parallel (@idx ni) compute_viscosity_kernel!( stokes.viscosity.η, ν, @strain(stokes)..., args, rheology, cutoff @@ -62,7 +71,18 @@ end return nothing end +function compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return compute_viscosity!( + backend(stokes), stokes, ν, phase_ratios, args, rheology, cutoff + ) +end function compute_viscosity!( + ::CPUBackendTrait, stokes::StokesArrays, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end + +function _compute_viscosity!( stokes::StokesArrays, ν, phase_ratios::PhaseRatio, args, rheology, cutoff ) ni = size(stokes.viscosity.η) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 5d9c3b9f..577b9960 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -9,14 +9,15 @@ end ## 2D VISCO-ELASTIC STOKES SOLVER # backend trait -function solve!(stokes::StokesArrays, args...; kwargs) - return solve!(backend(stokes), stokes, args...; kwargs...) +function solve!(stokes, args...; kwargs) + return solve!(backend(stokes), stokes, args...; kwargs) end -function solve!( - ::CPUBackendTrait, +solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) + +function _solve!( stokes::StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{2,T}, flow_bcs::FlowBoundaryConditions, ρg, @@ -142,10 +143,9 @@ function solve!( end # visco-elastic solver -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, @@ -259,10 +259,9 @@ end # GeoParams: general (visco-elasto-plastic) solver -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, @@ -420,10 +419,9 @@ end ## With phase ratios -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, @@ -633,11 +631,10 @@ function solve!( ) end -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::StokesArrays, thermal::ThermalArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{2,T}, flow_bcs, ϕ, diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index 0e1b505f..72e36923 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -1,4 +1,4 @@ -function heatdiffusion_PT!(thermal::ThermalArrays, args...; kwargs) +function heatdiffusion_PT!(thermal, args...; kwargs) return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs...) end diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index 8340c6dd..21798b13 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -66,6 +66,8 @@ struct ThermalArrays{_T} end end +ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Number}) where {N} = ThermalArrays(ni...) +ThermalArrays(::Type{CPUBackend}, ni::Vararg{Number,N}) where {N} = ThermalArrays(ni...) ThermalArrays(ni::NTuple{N,Number}) where {N} = ThermalArrays(ni...) function ThermalArrays(::Number, ::Number) throw(ArgumentError("ThermalArrays dimensions must be given as integers")) diff --git a/src/types/phases.jl b/src/types/phases.jl new file mode 100644 index 00000000..52f2d53d --- /dev/null +++ b/src/types/phases.jl @@ -0,0 +1,18 @@ +struct PhaseRatio{T} + vertex::T + center::T + + PhaseRatio(vertex::T, center::T) where {T<:AbstractArray} = new{T}(vertex, center) +end + +function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} + center = @fill(0.0, ni..., celldims = (num_phases,)) + vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) + # T = typeof(center) + return PhaseRatio(vertex, center) +end + +@inline PhaseRatio(::Type{CPUBackend}, ni, num_phases) = PhaseRatio(ni, num_phases) + +# backend trait +@inline backend(x::PhaseRatio) = backend(x.center.data) diff --git a/src/types/stokes.jl b/src/types/stokes.jl index c42ceffa..74a7b622 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -1,10 +1,3 @@ -# 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 - ## Velocity type struct Velocity{T} @@ -206,6 +199,10 @@ struct StokesArrays{A,B,C,D,T} end end +function StokesArrays(::Type{CPUBackend}, ni::Vararg{Integer,N}) where {N} + return StokesArrays(tuple(ni...)) +end +StokesArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where {N} = StokesArrays(ni) StokesArrays(ni::Vararg{Integer,N}) where {N} = StokesArrays(tuple(ni...)) function StokesArrays(::Number, ::Number) throw(ArgumentError("StokesArrays dimensions must be given as integers")) diff --git a/src/types/traits.jl b/src/types/traits.jl index 73bbab95..04c1acb8 100644 --- a/src/types/traits.jl +++ b/src/types/traits.jl @@ -1,5 +1,7 @@ abstract type BackendTrait end struct CPUBackendTrait <: BackendTrait end +# struct CUDABackendTrait <: BackendTrait end +# struct AMDGPUBackendTrait <: BackendTrait end @inline backend(::Array) = CPUBackendTrait() @inline backend(::Type{<:Array}) = CPUBackendTrait() From 17fa2430590eb9765f595764cf0b3fdd00ccb914 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 18:45:56 +0200 Subject: [PATCH 07/91] ext fix --- ext/JustRelaxCUDAExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index bb7d3fa0..70aff836 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -170,7 +170,7 @@ module JustRelax2D JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) # Utils - JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I::IGG) = compute_dt(S, di, dt_diff, I::IGG) + JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) = compute_dt(S, di, dt_diff, I::IGG) end From fd10f2fba2beba2c8eaa3269fec167aacc2aab9c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 2 Apr 2024 18:52:09 +0200 Subject: [PATCH 08/91] ext fix --- ext/JustRelaxCUDAExt.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 70aff836..87b89f02 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -171,7 +171,9 @@ module JustRelax2D # Utils JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) = compute_dt(S, di, dt_diff, I::IGG) - + JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff) = compute_dt(S, di, dt_diff) + JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di) = compute_dt(S, di) + end end From fd5e89c4b3941a613a0500396f46e57fe8f1ed95 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Tue, 2 Apr 2024 22:05:08 +0200 Subject: [PATCH 09/91] traits tests --- test_traits.jl.jl | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 test_traits.jl.jl diff --git a/test_traits.jl.jl b/test_traits.jl.jl new file mode 100644 index 00000000..bbba4703 --- /dev/null +++ b/test_traits.jl.jl @@ -0,0 +1,7 @@ +@test backend(Array) == CPUBackendTrait() +@test backend(Matrix) == CPUBackendTrait() +@test backend(Vector) == CPUBackendTrait() +@test backend(rand(2)) == CPUBackendTrait() +@test backend(rand(2,2)) == CPUBackendTrait() +@test backend(rand(2,2,2)) == CPUBackendTrait() +@test_throws ArgumentError backend(rand()) From 6a5f0a91e152bca5db609c1e0c5f6fef3dfe1caf Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Wed, 3 Apr 2024 14:16:25 +0200 Subject: [PATCH 10/91] up miniapp --- .../stokes2D/shear_heating/Shearheating2D.jl | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index 0b4a21ec..f3e5c780 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -1,17 +1,20 @@ +using CUDA # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax2D -using JustRelax.DataIO +# using JustRelax.DataIO # import JustRelax.@cell using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +# @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +@init_parallel_stencil(CUDA, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +# const backend = JustPIC.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, 2) or (:AMDGPU, Float64, 2) @@ -27,7 +30,6 @@ include("Shearheating_rheology.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- function copyinn_x!(A, B) - @parallel function f_x(A, B) @all(A) = @inn_x(B) return nothing @@ -86,19 +88,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx/2 # origin of thermal anomaly yc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -130,8 +132,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) ) ## Compression and not extension - fix this εbg = 5e-14 - stokes.V.Vx .= PTArray()([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray()([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) + stokes.V.Vx .= PTArray(backend )([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([ (ly - abs(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) @@ -174,7 +176,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) end # Time loop t, it = 0.0, 0 - while it < 100 + while it < 1 # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @@ -339,4 +341,4 @@ else igg end -main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) +main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) \ No newline at end of file From c1463cd0e1c99d202075f1a37cd82e9f5b755c6b Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 15:59:00 +0200 Subject: [PATCH 11/91] rework types & format --- src/Interpolations.jl | 4 +- src/JustRelax.jl | 11 +- src/JustRelax_CPU.jl | 2 + src/Utils.jl | 72 +++++---- src/array_conversions.jl | 27 ---- src/boundaryconditions/BoundaryConditions.jl | 22 ++- src/common.jl | 9 +- src/particles/subgrid_diffusion.jl | 10 +- src/phases/phases.jl | 13 +- src/rheology/BuoyancyForces.jl | 2 +- src/rheology/Viscosity.jl | 21 ++- src/stokes/Stokes2D.jl | 18 +-- src/stokes/Stokes3D.jl | 14 +- src/stokes/StressKernels.jl | 6 +- src/thermal_diffusion/DiffusionExplicit.jl | 32 ++-- .../DiffusionPT_GeoParams.jl | 4 +- .../DiffusionPT_coefficients.jl | 4 +- src/thermal_diffusion/DiffusionPT_solver.jl | 8 +- src/thermal_diffusion/Shearheating.jl | 4 +- src/types/constructors/heat_diffusion.jl | 51 ++++++ src/types/constructors/phases.jl | 6 + src/types/constructors/stokes.jl | 88 ++++++++++ src/types/heat_diffusion.jl | 100 ++++++------ src/types/phases.jl | 14 +- src/types/stokes.jl | 150 +++++++++--------- src/types/traits.jl | 11 ++ 26 files changed, 437 insertions(+), 266 deletions(-) create mode 100644 src/types/constructors/heat_diffusion.jl create mode 100644 src/types/constructors/phases.jl create mode 100644 src/types/constructors/stokes.jl diff --git a/src/Interpolations.jl b/src/Interpolations.jl index 733ad2ee..f6b5a102 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -17,11 +17,11 @@ end # From cell vertices to cell center temperature2center!(thermal) = temperature2center!(backend(thermal), thermal) -function temperature2center!(::CPUBackendTrait, thermal::ThermalArrays) +function temperature2center!(::CPUBackendTrait, thermal::JustRelax.ThermalArrays) return _temperature2center!(thermal) end -function _temperature2center!(thermal::ThermalArrays) +function _temperature2center!(thermal::JustRelax.ThermalArrays) @parallel (@idx size(thermal.Tc)...) temperature2center_kernel!(thermal.Tc, thermal.T) return nothing end diff --git a/src/JustRelax.jl b/src/JustRelax.jl index ffb01d1b..36c06769 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -21,8 +21,17 @@ PTArray(::T) where {T} = error(ArgumentError("Unknown backend $T")) export PTArray, CPUBackend, CUDABackend, AMDGPUBackend +include("types/stokes.jl") +# export StokesArrays, PTStokesCoeffs + +include("types/heat_diffusion.jl") +# export ThermalArrays, PTThermalCoeffs + +include("types/phases.jl") +# export PhaseRatio + include("types/traits.jl") -export BackendTrait, CPUBackendTrait +export BackendTrait, CPUBackendTrait, NonCPUBackendTrait include("topology/Topology.jl") export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 0500e9e4..7caa2e2a 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -9,6 +9,7 @@ using GeoParams, LinearAlgebra, Printf using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend +import JustRelax: PTStokesCoeffs @init_parallel_stencil(Threads, Float64, 2) @@ -29,6 +30,7 @@ using GeoParams, LinearAlgebra, Printf using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend +import JustRelax: PTStokesCoeffs @init_parallel_stencil(Threads, Float64, 3) diff --git a/src/Utils.jl b/src/Utils.jl index e9c8eada..b13f5c3e 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -81,10 +81,11 @@ macro tuple(A) end end -@inline _tuple(V::Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy -@inline _tuple(V::Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz -@inline _tuple(A::SymmetricTensor{<:AbstractArray{T,2}}) where {T} = A.xx, A.yy, A.xy_c -@inline function _tuple(A::SymmetricTensor{<:AbstractArray{T,3}}) where {T} +@inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy +@inline _tuple(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz +@inline _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}}) where {T} = + A.xx, A.yy, A.xy_c +@inline function _tuple(A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}}) where {T} return A.xx, A.yy, A.zz, A.yz_c, A.xz_c, A.xy_c end @@ -99,8 +100,9 @@ macro velocity(A) end end -@inline unpack_velocity(V::Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy -@inline unpack_velocity(V::Velocity{<:AbstractArray{T,3}}) where {T} = V.Vx, V.Vy, V.Vz +@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,2}}) where {T} = V.Vx, V.Vy +@inline unpack_velocity(V::JustRelax.Velocity{<:AbstractArray{T,3}}) where {T} = + V.Vx, V.Vy, V.Vz """ @qT(V) @@ -113,8 +115,9 @@ macro qT(A) end end -@inline unpack_qT(A::ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx, A.qTy -@inline unpack_qT(A::ThermalArrays{<:AbstractArray{T,3}}) where {T} = A.qTx, A.qTy, A.qTz +@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx, A.qTy +@inline unpack_qT(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} = + A.qTx, A.qTy, A.qTz """ @qT2(V) @@ -127,8 +130,9 @@ macro qT2(A) end end -@inline unpack_qT2(A::ThermalArrays{<:AbstractArray{T,2}}) where {T} = A.qTx2, A.qTy2 -@inline function unpack_qT2(A::ThermalArrays{<:AbstractArray{T,3}}) where {T} +@inline unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,2}}) where {T} = + A.qTx2, A.qTy2 +@inline function unpack_qT2(A::JustRelax.ThermalArrays{<:AbstractArray{T,3}}) where {T} return A.qTx2, A.qTy2, A.qTz2 end @@ -180,10 +184,14 @@ macro tensor(A) end end -@inline function unpack_tensor_stag(A::SymmetricTensor{<:AbstractArray{T,2}}) where {T} +@inline function unpack_tensor_stag( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}} +) where {T} return A.xx, A.yy, A.xy end -@inline function unpack_tensor_stag(A::SymmetricTensor{<:AbstractArray{T,3}}) where {T} +@inline function unpack_tensor_stag( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}} +) where {T} return A.xx, A.yy, A.zz, A.yz, A.xz, A.xy end @@ -200,12 +208,12 @@ macro shear(A) end @inline function unpack_shear_components_stag( - A::SymmetricTensor{<:AbstractArray{T,2}} + A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}} ) where {T} return A.xy end @inline function unpack_shear_components_stag( - A::SymmetricTensor{<:AbstractArray{T,3}} + A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}} ) where {T} return A.yz, A.xz, A.xy end @@ -223,7 +231,7 @@ macro normal(A) end @generated function unpack_normal_components_stag( - A::SymmetricTensor{<:AbstractArray{T,N}} + A::JustRelax.SymmetricTensor{<:AbstractArray{T,N}} ) where {T,N} syms = (:xx, :yy, :zz) quote @@ -269,10 +277,14 @@ macro tensor_center(A) end end -@inline function unpack_tensor_center(A::SymmetricTensor{<:AbstractArray{T,2}}) where {T} +@inline function unpack_tensor_center( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}} +) where {T} return A.xx, A.yy, A.xy_c end -@inline function unpack_tensor_center(A::SymmetricTensor{<:AbstractArray{T,3}}) where {T} +@inline function unpack_tensor_center( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}} +) where {T} return A.xx, A.yy, A.zz, A.yz_c, A.xz_c, A.xy_c end @@ -287,10 +299,14 @@ macro residuals(A) end end -@inline function unpack_residuals(A::Residual{<:AbstractArray{T,2}}) where {T} +@inline function unpack_residuals( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}} +) where {T} return A.Rx, A.Ry end -@inline function unpack_residuals(A::Residual{<:AbstractArray{T,3}}) where {T} +@inline function unpack_residuals( + A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}} +) where {T} return A.Rx, A.Ry, A.Rz end @@ -388,19 +404,20 @@ macro unpack(x) end """ - compute_dt(S::StokesArrays, di) + compute_dt(S::JustRelax.StokesArrays, di) Compute the time step `dt` for the velocity field `S.V` for a regular grid with grid spacing `di`. """ -@inline compute_dt(S::StokesArrays, di) = compute_dt(@velocity(S), di, Inf) +@inline compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(@velocity(S), di, Inf) """ - compute_dt(S::StokesArrays, di, dt_diff) + compute_dt(S::JustRelax.StokesArrays, di, dt_diff) Compute the time step `dt` for the velocity field `S.V` and the diffusive maximum time step `dt_diff` for a regular gridwith grid spacing `di`. """ -@inline compute_dt(S::StokesArrays, di, dt_diff) = compute_dt(@velocity(S), di, dt_diff) +@inline compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = + compute_dt(@velocity(S), di, dt_diff) @inline function compute_dt(V::NTuple, di, dt_diff) n = inv(length(V) + 0.1) @@ -408,22 +425,23 @@ Compute the time step `dt` for the velocity field `S.V` and the diffusive maximu return min(dt_diff, dt_adv) end """ - compute_dt(S::StokesArrays, di, igg) + compute_dt(S::JustRelax.StokesArrays, di, igg) Compute the time step `dt` for the velocity field `S.V` for a regular gridwith grid spacing `di`. The implicit global grid variable `I` implies that the time step is calculated globally and not separately on each block. """ -@inline compute_dt(S::StokesArrays, di, I::IGG) = compute_dt(@velocity(S), di, Inf, I::IGG) +@inline compute_dt(S::JustRelax.StokesArrays, di, I::IGG) = + compute_dt(@velocity(S), di, Inf, I::IGG) """ - compute_dt(S::StokesArrays, di, dt_diff) + compute_dt(S::JustRelax.StokesArrays, di, dt_diff) Compute the time step `dt` for the velocity field `S.V` and the diffusive maximum time step `dt_diff` for a regular gridwith grid spacing `di`. The implicit global grid variable `I` implies that the time step is calculated globally and not separately on each block. """ -@inline function compute_dt(S::StokesArrays, di, dt_diff, I::IGG) +@inline function compute_dt(S::JustRelax.StokesArrays, di, dt_diff, I::IGG) return compute_dt(@velocity(S), di, dt_diff, I::IGG) end diff --git a/src/array_conversions.jl b/src/array_conversions.jl index 51a740f5..b09501d5 100644 --- a/src/array_conversions.jl +++ b/src/array_conversions.jl @@ -1,30 +1,3 @@ -# Device trait system - -abstract type DeviceTrait end -struct CPUDeviceTrait <: DeviceTrait end -struct NonCPUDeviceTrait <: DeviceTrait end - -@inline iscpu(::Array) = CPUDeviceTrait() -@inline iscpu(::AbstractArray) = NonCPUDeviceTrait() -@inline iscpu(::T) where {T} = throw(ArgumentError("Unknown device")) - -@inline iscpu(::Velocity{Array{T,N}}) where {T,N} = CPUDeviceTrait() -@inline iscpu(::Velocity{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait() - -@inline iscpu(::SymmetricTensor{Array{T,N}}) where {T,N} = CPUDeviceTrait() -@inline iscpu(::SymmetricTensor{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait() - -@inline iscpu(::Residual{Array{T,N}}) where {T,N} = CPUDeviceTrait() -@inline iscpu(::Residual{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait() - -@inline iscpu(::ThermalArrays{Array{T,N}}) where {T,N} = CPUDeviceTrait() -@inline iscpu(::ThermalArrays{AbstractArray{T,N}}) where {T,N} = NonCPUDeviceTrait() - -@inline iscpu(::StokesArrays{M,A,B,C,Array{T,N},nDim}) where {M,A,B,C,T,N,nDim} = - CPUDeviceTrait() -@inline iscpu(::StokesArrays{M,A,B,C,AbstractArray{T,N},nDim}) where {M,A,B,C,T,N,nDim} = - NonCPUDeviceTrait() - ## Conversion of structs to CPU @inline remove_parameters(::T) where {T} = Base.typename(T).wrapper diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 15964d85..5dd126ef 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -49,7 +49,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::ThermalArrays, bcs::FlowBoundaryConditions + ::CPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs::FlowBoundaryConditions ) return thermal_bcs!(thermal.T, bcs) end @@ -243,21 +243,19 @@ function free_surface_bcs!( end end -function free_surface_bcs!( - stokes::StokesArrays{A,B,C,D,E,2}, bcs::FlowBoundaryConditions -) where {A,B,C,D,E} +function free_surface_bcs!(stokes::JustRelax.StokesArrays, bcs::FlowBoundaryConditions) if bcs.free_surface @views stokes.τ.yy[:, end] .= 0.0 end end -function free_surface_bcs!( - stokes::StokesArrays{A,B,C,D,E,3}, bcs::FlowBoundaryConditions -) where {A,B,C,D,E} - if bcs.free_surface - @views stokes.τ.zz[:, :, end] .= 0.0 - end -end +# function free_surface_bcs!( +# stokes::JustRelax.StokesArrays{A,B,C,D,E,3}, bcs::FlowBoundaryConditions +# ) where {A,B,C,D,E} +# if bcs.free_surface +# @views stokes.τ.zz[:, :, end] .= 0.0 +# end +# end @parallel_indices (i) function FreeSurface_Vy!( Vx::AbstractArray{T,2}, @@ -319,7 +317,7 @@ end end function pureshear_bc!( - stokes::StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg + stokes::JustRelax.StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg ) where {T} stokes.V.Vx[:, 2:(end - 1)] .= PTArray(([εbg * x for x in xvi[1], y in xci[2]])) stokes.V.Vy[2:(end - 1), :] .= PTArray(([-εbg * y for x in xci[1], y in xvi[2]])) diff --git a/src/common.jl b/src/common.jl index 1d961fb0..c475a4e8 100644 --- a/src/common.jl +++ b/src/common.jl @@ -1,10 +1,10 @@ -include("types/stokes.jl") +include("types/constructors/stokes.jl") export StokesArrays, PTStokesCoeffs -include("types/heat_diffusion.jl") +include("types/constructors/heat_diffusion.jl") export ThermalArrays, PTThermalCoeffs -include("types/phases.jl") +include("types/constructors/phases.jl") export PhaseRatio include("Utils.jl") @@ -70,5 +70,4 @@ include("stokes/VelocityKernels.jl") # thermal diffusion include("thermal_diffusion/DiffusionPT.jl") -export heatdiffusion_PT! -export compute_shear_heating! +export PTThermalCoeffs, heatdiffusion_PT!, compute_shear_heating! diff --git a/src/particles/subgrid_diffusion.jl b/src/particles/subgrid_diffusion.jl index 8c970576..bc457642 100644 --- a/src/particles/subgrid_diffusion.jl +++ b/src/particles/subgrid_diffusion.jl @@ -4,10 +4,10 @@ function subgrid_characteristic_time!( subgrid_arrays, particles, dt₀, - phases::PhaseRatio, + phases::JustRelax.PhaseRatio, rheology, - thermal::ThermalArrays, - stokes::StokesArrays, + thermal::JustRelax.ThermalArrays, + stokes::JustRelax.StokesArrays, xci, di, ) @@ -24,8 +24,8 @@ function subgrid_characteristic_time!( dt₀, phases::AbstractArray{Int,N}, rheology, - thermal::ThermalArrays, - stokes::StokesArrays, + thermal::JustRelax.ThermalArrays, + stokes::JustRelax.StokesArrays, xci, di, ) where {N} diff --git a/src/phases/phases.jl b/src/phases/phases.jl index 7a2fb037..22452832 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -1,9 +1,10 @@ """ - nphases(x::PhaseRatio) + nphases(x::JustRelax.PhaseRatio) -Return the number of phases in `x::PhaseRatio`. +Return the number of phases in `x::JustRelax.PhaseRatio`. """ -@inline nphases(x::PhaseRatio) = nphases(x.center) +@inline nphases(x::JustRelax.PhaseRatio) = nphases(x.center) + @inline function nphases( ::CellArray{StaticArraysCore.SArray{Tuple{N},T,N1,N},N2,N3,T_Array} ) where {N,T,N1,N2,N3,T_Array} @@ -52,12 +53,14 @@ function phase_ratios_center(phase_ratios, particles, grid, phases) end function phase_ratios_center( - ::CPUBackendTrait, phase_ratios::PhaseRatio, particles, grid::Geometry, phases + ::CPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases ) return _phase_ratios_center(phase_ratios, particles, grid, phases) end -function _phase_ratios_center(phase_ratios::PhaseRatio, particles, grid::Geometry, phases) +function _phase_ratios_center( + phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases +) ni = size(phases) @parallel (@idx ni) phase_ratios_center_kernel( phase_ratios.center, particles.coords, grid.xci, grid.di, phases diff --git a/src/rheology/BuoyancyForces.jl b/src/rheology/BuoyancyForces.jl index 1b430b6f..553b5af6 100644 --- a/src/rheology/BuoyancyForces.jl +++ b/src/rheology/BuoyancyForces.jl @@ -21,7 +21,7 @@ end 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. """ -function compute_ρg!(ρg, phase_ratios::PhaseRatio, rheology, args) +function compute_ρg!(ρg, phase_ratios::JustRelax.PhaseRatio, rheology, args) ni = size(ρg) @parallel (@idx ni) compute_ρg_kernel!(ρg, phase_ratios.center, rheology, args) return nothing diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 4ec06176..5104ca8a 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -2,13 +2,11 @@ function compute_viscosity!(stokes, ν, args, rheology, cutoff) return compute_viscosity!(backend(stokes), ν, args, rheology, cutoff) end -function compute_viscosity!( - ::CPUBackendTrait, stokes, ν, args, rheology, cutoff -) +function compute_viscosity!(::CPUBackendTrait, stokes, ν, args, rheology, cutoff) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end -function _compute_viscosity!(stokes::StokesArrays, ν, args, rheology, cutoff) +function _compute_viscosity!(stokes::JustRelax.StokesArrays, ν, args, rheology, cutoff) ni = size(stokes.viscosity.η) @parallel (@idx ni) compute_viscosity_kernel!( stokes.viscosity.η, ν, @strain(stokes)..., args, rheology, cutoff @@ -77,13 +75,24 @@ function compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) ) end function compute_viscosity!( - ::CPUBackendTrait, stokes::StokesArrays, ν, phase_ratios, args, rheology, cutoff + ::CPUBackendTrait, + stokes::JustRelax.StokesArrays, + ν, + phase_ratios, + args, + rheology, + cutoff, ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end function _compute_viscosity!( - stokes::StokesArrays, ν, phase_ratios::PhaseRatio, args, rheology, cutoff + stokes::JustRelax.StokesArrays, + ν, + phase_ratios::JustRelax.PhaseRatio, + args, + rheology, + cutoff, ) ni = size(stokes.viscosity.η) @parallel (@idx ni) compute_viscosity_kernel!( diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 77cb2e74..b041aa6a 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -1,5 +1,5 @@ ## 2D STOKES MODULE -function update_τ_o!(stokes::StokesArrays) +function update_τ_o!(stokes::JustRelax.StokesArrays) @parallel (@idx size(τxy)) multi_copy!( @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) ) @@ -13,10 +13,10 @@ function solve!(stokes, args...; kwargs) return solve!(backend(stokes), stokes, args...; kwargs) end -solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) +solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) function _solve!( - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes, di::NTuple{2,T}, flow_bcs::FlowBoundaryConditions, @@ -144,7 +144,7 @@ end # visco-elastic solver function _solve!( - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes, di::NTuple{2,T}, flow_bcs, @@ -260,7 +260,7 @@ end # GeoParams: general (visco-elasto-plastic) solver function _solve!( - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes, di::NTuple{2,T}, flow_bcs, @@ -429,12 +429,12 @@ end ## With phase ratios function _solve!( - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, rheology, args, dt, @@ -650,8 +650,8 @@ function _solve!( end function _solve!( - stokes::StokesArrays, - thermal::ThermalArrays, + stokes::JustRelax.StokesArrays, + thermal::JustRelax.ThermalArrays, pt_stokes, di::NTuple{2,T}, flow_bcs, diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index f2f39573..e4c725f0 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -10,16 +10,18 @@ return nothing end -function update_τ_o!(stokes::StokesArrays) +function update_τ_o!(stokes::JustRelax.StokesArrays) @parallel update_τ_o!(@tensor(stokes.τ_o)..., @stress(stokes)...) end ## 3D VISCO-ELASTIC STOKES SOLVER -solve!(stokes::StokesArrays, args...) = solve!(CPUBackendTrait(stokes), stokes, args...) +function solve!(stokes::JustRelax.StokesArrays, args...) + return solve!(CPUBackendTrait(stokes), stokes, args...) +end function solve!( ::CPUBackendTrait, - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bcs, @@ -153,7 +155,7 @@ end function solve!( ::CPUBackendTrait, - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bcs::FlowBoundaryConditions, @@ -326,12 +328,12 @@ end # GeoParams and multiple phases function solve!( ::CPUBackendTrait, - stokes::StokesArrays, + stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, flow_bc::FlowBoundaryConditions, ρg, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, rheology::NTuple{N,AbstractMaterialParamsStruct}, args, dt, diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 58f01b07..1beba926 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -298,14 +298,14 @@ end end """ - tensor_invariant!(A::SymmetricTensor) + tensor_invariant!(A::JustRelax.SymmetricTensor) Compute the tensor invariant of the given symmetric tensor `A`. # Arguments -- `A::SymmetricTensor`: The input symmetric tensor. +- `A::JustRelax.SymmetricTensor`: The input symmetric tensor. """ -function tensor_invariant!(A::SymmetricTensor) +function tensor_invariant!(A::JustRelax.SymmetricTensor) ni = size(A.II) @parallel (@idx ni) tensor_invariant_kernel!(A.II, @tensor(A)...) return nothing diff --git a/src/thermal_diffusion/DiffusionExplicit.jl b/src/thermal_diffusion/DiffusionExplicit.jl index f2c4d52d..5241e235 100644 --- a/src/thermal_diffusion/DiffusionExplicit.jl +++ b/src/thermal_diffusion/DiffusionExplicit.jl @@ -54,8 +54,8 @@ end ## SOLVER function JustRelax.solve!( - thermal::ThermalArrays{M}, - pt_thermal::PTThermalCoeffs, + thermal::JustRelax.ThermalArrays{M}, + pt_thermal::JustRelax.PTThermalCoeffs, thermal_parameters::ThermalParameters{<:AbstractArray{_T,1}}, ni::NTuple{1,Integer}, di::NTuple{1,_T}, @@ -302,7 +302,7 @@ end ## SOLVER function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_parameters::ThermalParameters{<:AbstractArray{_T,2}}, thermal_bc::NamedTuple, di::NTuple{2,_T}, @@ -328,7 +328,7 @@ end # upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_parameters::ThermalParameters{<:AbstractArray{_T,2}}, stokes, thermal_bc::TemperatureBoundaryConditions, @@ -366,7 +366,7 @@ end # GEOPARAMS VERSION function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology, args::NamedTuple, @@ -398,10 +398,10 @@ end # with multiple material phases function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology::NTuple{N,AbstractMaterialParamsStruct}, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, args::NamedTuple, di::NTuple{2,_T}, dt, @@ -428,7 +428,7 @@ end # Upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, stokes, rheology, @@ -466,7 +466,7 @@ end # Upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, stokes, phases, @@ -763,7 +763,7 @@ end ## SOLVER function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_parameters::ThermalParameters{<:AbstractArray{_T,3}}, thermal_bc::NamedTuple, di::NTuple{3,_T}, @@ -791,7 +791,7 @@ end # upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_parameters::ThermalParameters{<:AbstractArray{_T,3}}, thermal_bc::NamedTuple, stokes, @@ -834,7 +834,7 @@ end # GEOPARAMS VERSION function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology, args::NamedTuple, @@ -871,10 +871,10 @@ end # with multiple material phases - no advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, rheology::NTuple{N,AbstractMaterialParamsStruct}, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, args::NamedTuple, di::NTuple{3,_T}, dt; @@ -916,7 +916,7 @@ end # upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, stokes, rheology, @@ -967,7 +967,7 @@ end # upwind advection function JustRelax.solve!( - thermal::ThermalArrays{M}, + thermal::JustRelax.ThermalArrays{M}, thermal_bc::TemperatureBoundaryConditions, stokes, phases, diff --git a/src/thermal_diffusion/DiffusionPT_GeoParams.jl b/src/thermal_diffusion/DiffusionPT_GeoParams.jl index 67464c7a..c356cb01 100644 --- a/src/thermal_diffusion/DiffusionPT_GeoParams.jl +++ b/src/thermal_diffusion/DiffusionPT_GeoParams.jl @@ -1,12 +1,12 @@ ## Phases -@inline get_phase(x::PhaseRatio) = x.center +@inline get_phase(x::JustRelax.PhaseRatio) = x.center @inline get_phase(x) = x # update_pt_thermal_arrays!(::Vararg{Any,N}) where {N} = nothing function update_pt_thermal_arrays!( - pt_thermal, phase_ratios::PhaseRatio, rheology, args, _dt + pt_thermal, phase_ratios::JustRelax.PhaseRatio, rheology, args, _dt ) ni = size(phase_ratios.center) diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index 11d6ff33..b8de9f8c 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -18,7 +18,7 @@ function PTThermalCoeffs( θr_dτ, dτ_ρ, rheology, phase_ratios.center, args, max_lxyz, Vpdτ, inv(dt) ) - return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) + return JustRelax.PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) end # without phase ratios @@ -33,7 +33,7 @@ function PTThermalCoeffs( θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, inv(dt) ) - return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) + return JustRelax.PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) end @parallel_indices (I...) function compute_pt_thermal_arrays!( diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index 72e36923..eb73cd9a 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -9,8 +9,8 @@ Heat diffusion solver using Pseudo-Transient iterations. Both `K` and `ρCp` are """ function heatdiffusion_PT!( ::CPUBackendTrait, - thermal::ThermalArrays, - pt_thermal::PTThermalCoeffs, + thermal::JustRelax.ThermalArrays, + pt_thermal::JustRelax.PTThermalCoeffs, thermal_bc::TemperatureBoundaryConditions, K::AbstractArray, ρCp::AbstractArray, @@ -98,8 +98,8 @@ Heat diffusion solver using Pseudo-Transient iterations. """ function heatdiffusion_PT!( ::CPUBackendTrait, - thermal::ThermalArrays, - pt_thermal::PTThermalCoeffs, + thermal::JustRelax.ThermalArrays, + pt_thermal::JustRelax.PTThermalCoeffs, thermal_bc::TemperatureBoundaryConditions, rheology, args::NamedTuple, diff --git a/src/thermal_diffusion/Shearheating.jl b/src/thermal_diffusion/Shearheating.jl index cb1a3755..cd1d559b 100644 --- a/src/thermal_diffusion/Shearheating.jl +++ b/src/thermal_diffusion/Shearheating.jl @@ -21,7 +21,9 @@ end return nothing end -function compute_shear_heating!(thermal, stokes, phase_ratios::PhaseRatio, rheology, dt) +function compute_shear_heating!( + thermal, stokes, phase_ratios::JustRelax.PhaseRatio, rheology, dt +) ni = size(thermal.shear_heating) @parallel (@idx ni) compute_shear_heating_kernel!( thermal.shear_heating, diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl new file mode 100644 index 00000000..54d0695b --- /dev/null +++ b/src/types/constructors/heat_diffusion.jl @@ -0,0 +1,51 @@ + +function ThermalArrays(nx::Integer, ny::Integer) + T = @zeros(nx + 3, ny + 1) + ΔT = @zeros(nx + 3, ny + 1) + Told = @zeros(nx + 3, ny + 1) + Tc = @zeros(nx, ny) + H = @zeros(nx, ny) + shear_heating = @zeros(nx, ny) + dT_dt = @zeros(nx + 1, ny - 1) + qTx = @zeros(nx + 2, ny - 1) + qTy = @zeros(nx + 1, ny) + qTx2 = @zeros(nx + 2, ny - 1) + qTy2 = @zeros(nx + 1, ny) + ResT = @zeros(nx + 1, ny - 1) + return JustRelax.ThermalArrays( + T, + Tc, + ΔT, + Told, + dT_dt, + qTx, + qTy, + nothing, + qTx2, + qTy2, + nothing, + H, + shear_heating, + ResT, + ) +end + +function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) + T = @zeros(nx + 1, ny + 1, nz + 1) + ΔT = @zeros(nx + 1, ny + 1, nz + 1) + Told = @zeros(nx + 1, ny + 1, nz + 1) + Tc = @zeros(nx, ny, nz) + H = @zeros(nx, ny, nz) + shear_heating = @zeros(nx, ny, nz) + dT_dt = @zeros(ni .- 1) + qTx = @zeros(nx, ny - 1, nz - 1) + qTy = @zeros(nx - 1, ny, nz - 1) + qTz = @zeros(nx - 1, ny - 1, nz) + qTx2 = @zeros(nx, ny - 1, nz - 1) + qTy2 = @zeros(nx - 1, ny, nz - 1) + qTz2 = @zeros(nx - 1, ny - 1, nz) + ResT = @zeros((ni .- 1)...) + return JustRelax.ThermalArrays( + T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT + ) +end diff --git a/src/types/constructors/phases.jl b/src/types/constructors/phases.jl new file mode 100644 index 00000000..f7f1d5a4 --- /dev/null +++ b/src/types/constructors/phases.jl @@ -0,0 +1,6 @@ +function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} + center = @fill(0.0, ni..., celldims = (num_phases,)) + vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) + # T = typeof(center) + return JustRelax.PhaseRatio(vertex, center) +end diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl new file mode 100644 index 00000000..fb01b005 --- /dev/null +++ b/src/types/constructors/stokes.jl @@ -0,0 +1,88 @@ +## Velocity type + +function Velocity(nx::Integer, ny::Integer) + nVx = (nx + 1, ny + 2) + nVy = (nx + 2, ny + 1) + + Vx, Vy = @zeros(nVx...), @zeros(nVy) + return JustRelax.Velocity(Vx, Vy, nothing) +end + +function Velocity(nx::Integer, ny::Integer, nz::Integer) + nVx = (nx + 1, ny + 2, nz + 2) + nVy = (nx + 2, ny + 1, nz + 2) + nVz = (nx + 2, ny + 2, nz + 1) + + Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) + return JustRelax.Velocity{PTArray}(Vx, Vy, Vz) +end + +## Viscosity type + +function Viscosity(ni::NTuple{N,Integer}) where {N} + η = @ones(ni...) + η_vep = @ones(ni...) + ητ = @zeros(ni...) + return JustRelax.Viscosity(η, η_vep, ητ) +end + +## SymmetricTensor type + +function SymmetricTensor(nx::Integer, ny::Integer) + return JustRelax.SymmetricTensor( + @zeros(nx, ny), # xx + @zeros(nx, ny), # yy + @zeros(nx + 1, ny + 1), # xy + @zeros(nx, ny), # xy @ cell center + @zeros(nx, ny) # II (second invariant) + ) +end + +function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) + return JustRelax.SymmetricTensor( + @zeros(nx, ny, nz), # xx + @zeros(nx + 1, ny + 1, nz), # xy + @zeros(nx, ny, nz), # yy + @zeros(nx + 1, ny, nz + 1), # xz + @zeros(nx, ny + 1, nz + 1), # yz + @zeros(nx, ny, nz), # zz + @zeros(nx, ny, nz), # yz @ cell center + @zeros(nx, ny, nz), # xz @ cell center + @zeros(nx, ny, nz), # xy @ cell center + @zeros(nx, ny, nz), # II (second invariant) + ) +end + +## Residual type + +function Residual(nx::Integer, ny::Integer) + Rx = @zeros(nx - 1, ny) + Ry = @zeros(nx, ny - 1) + RP = @zeros(nx, ny) + return JustRelax.Residual(RP, Rx, Ry) +end + +function Residual(nx::Integer, ny::Integer, nz::Integer) + Rx = @zeros(nx - 1, ny, nz) + Ry = @zeros(nx, ny - 1, nz) + Rz = @zeros(nx, ny, nz - 1) + RP = @zeros(nx, ny, nz) + return JustRelax.Residual(RP, Rx, Ry, Rz) +end + +## StokesArrays type +function StokesArrays(ni::NTuple{N,Integer}) where {N} + P = @zeros(ni...) + P0 = @zeros(ni...) + ∇V = @zeros(ni...) + V = Velocity(ni...) + τ = SymmetricTensor(ni...) + τ_o = SymmetricTensor(ni...) + ε = SymmetricTensor(ni...) + ε_pl = SymmetricTensor(ni...) + EII_pl = @zeros(ni...) + viscosity = Viscosity(ni) + R = Residual(ni...) + + return JustRelax.StokesArrays(P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R) +end diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index 21798b13..cf997d75 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -14,56 +14,56 @@ struct ThermalArrays{_T} shear_heating::_T # shear heating terms ResT::_T - function ThermalArrays(nx::Integer, ny::Integer) - T = @zeros(nx + 3, ny + 1) - ΔT = @zeros(nx + 3, ny + 1) - Told = @zeros(nx + 3, ny + 1) - Tc = @zeros(nx, ny) - H = @zeros(nx, ny) - shear_heating = @zeros(nx, ny) - dT_dt = @zeros(nx + 1, ny - 1) - qTx = @zeros(nx + 2, ny - 1) - qTy = @zeros(nx + 1, ny) - qTx2 = @zeros(nx + 2, ny - 1) - qTy2 = @zeros(nx + 1, ny) - ResT = @zeros(nx + 1, ny - 1) - return new{typeof(T)}( - T, - Tc, - ΔT, - Told, - dT_dt, - qTx, - qTy, - nothing, - qTx2, - qTy2, - nothing, - H, - shear_heating, - ResT, - ) - end + # function ThermalArrays(nx::Integer, ny::Integer) + # T = @zeros(nx + 3, ny + 1) + # ΔT = @zeros(nx + 3, ny + 1) + # Told = @zeros(nx + 3, ny + 1) + # Tc = @zeros(nx, ny) + # H = @zeros(nx, ny) + # shear_heating = @zeros(nx, ny) + # dT_dt = @zeros(nx + 1, ny - 1) + # qTx = @zeros(nx + 2, ny - 1) + # qTy = @zeros(nx + 1, ny) + # qTx2 = @zeros(nx + 2, ny - 1) + # qTy2 = @zeros(nx + 1, ny) + # ResT = @zeros(nx + 1, ny - 1) + # return new{typeof(T)}( + # T, + # Tc, + # ΔT, + # Told, + # dT_dt, + # qTx, + # qTy, + # nothing, + # qTx2, + # qTy2, + # nothing, + # H, + # shear_heating, + # ResT, + # ) + # end - function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) - T = @zeros(nx + 1, ny + 1, nz + 1) - ΔT = @zeros(nx + 1, ny + 1, nz + 1) - Told = @zeros(nx + 1, ny + 1, nz + 1) - Tc = @zeros(nx, ny, nz) - H = @zeros(nx, ny, nz) - shear_heating = @zeros(nx, ny, nz) - dT_dt = @zeros(ni .- 1) - qTx = @zeros(nx, ny - 1, nz - 1) - qTy = @zeros(nx - 1, ny, nz - 1) - qTz = @zeros(nx - 1, ny - 1, nz) - qTx2 = @zeros(nx, ny - 1, nz - 1) - qTy2 = @zeros(nx - 1, ny, nz - 1) - qTz2 = @zeros(nx - 1, ny - 1, nz) - ResT = @zeros((ni .- 1)...) - return new{typeof(T)}( - T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT - ) - end + # function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) + # T = @zeros(nx + 1, ny + 1, nz + 1) + # ΔT = @zeros(nx + 1, ny + 1, nz + 1) + # Told = @zeros(nx + 1, ny + 1, nz + 1) + # Tc = @zeros(nx, ny, nz) + # H = @zeros(nx, ny, nz) + # shear_heating = @zeros(nx, ny, nz) + # dT_dt = @zeros(ni .- 1) + # qTx = @zeros(nx, ny - 1, nz - 1) + # qTy = @zeros(nx - 1, ny, nz - 1) + # qTz = @zeros(nx - 1, ny - 1, nz) + # qTx2 = @zeros(nx, ny - 1, nz - 1) + # qTy2 = @zeros(nx - 1, ny, nz - 1) + # qTz2 = @zeros(nx - 1, ny - 1, nz) + # ResT = @zeros((ni .- 1)...) + # return new{typeof(T)}( + # T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT + # ) + # end end ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Number}) where {N} = ThermalArrays(ni...) @@ -77,7 +77,7 @@ function ThermalArrays(::Number, ::Number, ::Number) end # traits -@inline backend(::ThermalArrays{T}) where {T} = backend(T) +# @inline backend(::ThermalArrays{T}) where {T} = backend(T) ## Thermal diffusion coefficients diff --git a/src/types/phases.jl b/src/types/phases.jl index 52f2d53d..c82a27d4 100644 --- a/src/types/phases.jl +++ b/src/types/phases.jl @@ -5,14 +5,14 @@ struct PhaseRatio{T} PhaseRatio(vertex::T, center::T) where {T<:AbstractArray} = new{T}(vertex, center) end -function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} - center = @fill(0.0, ni..., celldims = (num_phases,)) - vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) - # T = typeof(center) - return PhaseRatio(vertex, center) -end +# function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} +# center = @fill(0.0, ni..., celldims = (num_phases,)) +# vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) +# # T = typeof(center) +# return PhaseRatio(vertex, center) +# end @inline PhaseRatio(::Type{CPUBackend}, ni, num_phases) = PhaseRatio(ni, num_phases) # backend trait -@inline backend(x::PhaseRatio) = backend(x.center.data) +# @inline backend(x::PhaseRatio) = backend(x.center.data) diff --git a/src/types/stokes.jl b/src/types/stokes.jl index 74a7b622..02450bcf 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -10,22 +10,22 @@ end Velocity(Vx::T, Vy::T) where {T} = Velocity(Vx, Vy, nothing) -function Velocity(nx::Integer, ny::Integer) - nVx = (nx + 1, ny + 2) - nVy = (nx + 2, ny + 1) +# function Velocity(nx::Integer, ny::Integer) +# nVx = (nx + 1, ny + 2) +# nVy = (nx + 2, ny + 1) - Vx, Vy = @zeros(nVx...), @zeros(nVy) - return Velocity(Vx, Vy, nothing) -end +# Vx, Vy = @zeros(nVx...), @zeros(nVy) +# return Velocity(Vx, Vy, nothing) +# end -function Velocity(nx::Integer, ny::Integer, nz::Integer) - nVx = (nx + 1, ny + 2, nz + 2) - nVy = (nx + 2, ny + 1, nz + 2) - nVz = (nx + 2, ny + 2, nz + 1) +# function Velocity(nx::Integer, ny::Integer, nz::Integer) +# nVx = (nx + 1, ny + 2, nz + 2) +# nVy = (nx + 2, ny + 1, nz + 2) +# nVz = (nx + 2, ny + 2, nz + 1) - Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) - return Velocity{PTArray}(Vx, Vy, Vz) -end +# Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) +# return Velocity{PTArray}(Vx, Vy, Vz) +# end Velocity(ni::NTuple{N,Number}) where {N} = Velocity(ni...) function Velocity(::Number, ::Number) @@ -47,12 +47,12 @@ end Viscosity(args...) = Viscosity(promote(args...)...) -function Viscosity(ni::NTuple{N,Integer}) where {N} - η = @ones(ni...) - η_vep = @ones(ni...) - ητ = @zeros(ni...) - return Viscosity(η, η_vep, ητ) -end +# function Viscosity(ni::NTuple{N,Integer}) where {N} +# η = @ones(ni...) +# η_vep = @ones(ni...) +# ητ = @zeros(ni...) +# return Viscosity(η, η_vep, ητ) +# end 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)) @@ -96,30 +96,30 @@ function SymmetricTensor(xx::T, yy::T, xy::T, xy_c::T, II::T) where {T} ) end -function SymmetricTensor(nx::Integer, ny::Integer) - return SymmetricTensor( - @zeros(nx, ny), # xx - @zeros(nx, ny), # yy - @zeros(nx + 1, ny + 1), # xy - @zeros(nx, ny), # xy @ cell center - @zeros(nx, ny) # II (second invariant) - ) -end - -function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) - return SymmetricTensor( - @zeros(nx, ny, nz), # xx - @zeros(nx + 1, ny + 1, nz), # xy - @zeros(nx, ny, nz), # yy - @zeros(nx + 1, ny, nz + 1), # xz - @zeros(nx, ny + 1, nz + 1), # yz - @zeros(nx, ny, nz), # zz - @zeros(nx, ny, nz), # yz @ cell center - @zeros(nx, ny, nz), # xz @ cell center - @zeros(nx, ny, nz), # xy @ cell center - @zeros(nx, ny, nz), # II (second invariant) - ) -end +# function SymmetricTensor(nx::Integer, ny::Integer) +# return SymmetricTensor( +# @zeros(nx, ny), # xx +# @zeros(nx, ny), # yy +# @zeros(nx + 1, ny + 1), # xy +# @zeros(nx, ny), # xy @ cell center +# @zeros(nx, ny) # II (second invariant) +# ) +# end + +# function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) +# return SymmetricTensor( +# @zeros(nx, ny, nz), # xx +# @zeros(nx + 1, ny + 1, nz), # xy +# @zeros(nx, ny, nz), # yy +# @zeros(nx + 1, ny, nz + 1), # xz +# @zeros(nx, ny + 1, nz + 1), # yz +# @zeros(nx, ny, nz), # zz +# @zeros(nx, ny, nz), # yz @ cell center +# @zeros(nx, ny, nz), # xz @ cell center +# @zeros(nx, ny, nz), # xy @ cell center +# @zeros(nx, ny, nz), # II (second invariant) +# ) +# end SymmetricTensor(ni::NTuple{N,Number}) where {N} = SymmetricTensor(ni...) function SymmetricTensor(::Number, ::Number) @@ -142,20 +142,20 @@ end Residual(RP::T, Rx::T, Ry::T) where {T} = Residual(RP, Rx, Ry, nothing) -function Residual(nx::Integer, ny::Integer) - Rx = @zeros(nx - 1, ny) - Ry = @zeros(nx, ny - 1) - RP = @zeros(nx, ny) - return Residual(RP, Rx, Ry) -end - -function Residual(nx::Integer, ny::Integer, nz::Integer) - Rx = @zeros(nx - 1, ny, nz) - Ry = @zeros(nx, ny - 1, nz) - Rz = @zeros(nx, ny, nz - 1) - RP = @zeros(nx, ny, nz) - return Residual(RP, Rx, Ry, Rz) -end +# function Residual(nx::Integer, ny::Integer) +# Rx = @zeros(nx - 1, ny) +# Ry = @zeros(nx, ny - 1) +# RP = @zeros(nx, ny) +# return Residual(RP, Rx, Ry) +# end + +# function Residual(nx::Integer, ny::Integer, nz::Integer) +# Rx = @zeros(nx - 1, ny, nz) +# Ry = @zeros(nx, ny - 1, nz) +# Rz = @zeros(nx, ny, nz - 1) +# RP = @zeros(nx, ny, nz) +# return Residual(RP, Rx, Ry, Rz) +# end Residual(ni::NTuple{N,Number}) where {N} = Residual(ni...) function Residual(::Number, ::Number) @@ -180,23 +180,23 @@ struct StokesArrays{A,B,C,D,T} τ_o::Union{B,Nothing} R::C - function StokesArrays(ni::NTuple{N,Integer}) where {N} - P = @zeros(ni...) - P0 = @zeros(ni...) - ∇V = @zeros(ni...) - V = Velocity(ni) - τ = SymmetricTensor(ni) - τ_o = SymmetricTensor(ni) - ε = SymmetricTensor(ni) - ε_pl = SymmetricTensor(ni) - EII_pl = @zeros(ni...) - viscosity = Viscosity(ni) - R = Residual(ni) - - return new{typeof(V),typeof(τ),typeof(R),typeof(viscosity),typeof(P)}( - P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R - ) - end + # function StokesArrays(ni::NTuple{N,Integer}) where {N} + # P = @zeros(ni...) + # P0 = @zeros(ni...) + # ∇V = @zeros(ni...) + # V = Velocity(ni) + # τ = SymmetricTensor(ni) + # τ_o = SymmetricTensor(ni) + # ε = SymmetricTensor(ni) + # ε_pl = SymmetricTensor(ni) + # EII_pl = @zeros(ni...) + # viscosity = Viscosity(ni) + # R = Residual(ni) + + # return new{typeof(V),typeof(τ),typeof(R),typeof(viscosity),typeof(P)}( + # P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R + # ) + # end end function StokesArrays(::Type{CPUBackend}, ni::Vararg{Integer,N}) where {N} @@ -212,7 +212,7 @@ function StokesArrays(::Number, ::Number, ::Number) end # traits -@inline backend(x::StokesArrays) = backend(x.P) +# @inline backend(x::StokesArrays) = backend(x.P) ## PTStokesCoeffs type diff --git a/src/types/traits.jl b/src/types/traits.jl index 04c1acb8..bcc0c8e8 100644 --- a/src/types/traits.jl +++ b/src/types/traits.jl @@ -1,8 +1,19 @@ abstract type BackendTrait end struct CPUBackendTrait <: BackendTrait end +struct NonCPUBackendTrait <: BackendTrait end # struct CUDABackendTrait <: BackendTrait end # struct AMDGPUBackendTrait <: BackendTrait end @inline backend(::Array) = CPUBackendTrait() @inline backend(::Type{<:Array}) = CPUBackendTrait() +@inline backend(::AbstractArray) = NonCPUDeviceTrait() +@inline backend(::Type{<:AbstractArray}) = NonCPUDeviceTrait() + @inline backend(::T) where {T} = throw(ArgumentError("Backend $(T) not supported")) + +@inline backend(::JustRelax.Velocity{T}) where {T} = backend(T) +@inline backend(::JustRelax.SymmetricTensor{T}) where {T} = backend(T) +@inline backend(::JustRelax.Residual{T}) where {T} = backend(T) +@inline backend(::JustRelax.ThermalArrays{T}) where {T} = backend(T) +@inline backend(x::JustRelax.StokesArrays) = backend(x.P) +@inline backend(x::JustRelax.PhaseRatio) = backend(x.center.data) From 4f0bfda40cd42522cf5e989da987cfae807509a4 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 16:33:25 +0200 Subject: [PATCH 12/91] acommodate constructors for extensions --- ext/JustRelaxCUDAExt.jl | 55 +++++++++++++++++-- src/JustRelax.jl | 5 +- .../DiffusionPT_coefficients.jl | 8 +-- src/types/constructors/heat_diffusion.jl | 7 +++ src/types/constructors/phases.jl | 4 ++ src/types/constructors/stokes.jl | 4 ++ 6 files changed, 71 insertions(+), 12 deletions(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 87b89f02..9056594b 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -38,12 +38,6 @@ module JustRelax2D @inline backend(::Type{<:CuArray}) = CUDABackendTrait() # Types - function JustRelax.JustRelax2D.StokesArrays( - ::Type{CUDABackend}, ni::Vararg{Integer,N} - ) where {N} - return StokesArrays(tuple(ni...)) - end - function JustRelax.JustRelax2D.StokesArrays( ::Type{CUDABackend}, ni::NTuple{N,Integer} ) where {N} @@ -66,6 +60,55 @@ module JustRelax2D return PhaseRatio(ni, num_phases) end + function PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ = 1e-8, + CFL = 0.9 / √3, + ) where {nDim,T} + return JustRelax.JustRelax2D.PTThermalCoeffs( + rheology, + phase_ratios, + args, + dt, + ni, + di, + li; + ϵ = ϵ, + CFL = CFL, + ) + end + + + function PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ = 1e-8, + CFL = 0.9 / √3, + ) where {nDim,T} + return JustRelax.JustRelax2D.PTThermalCoeffs( + rheology, + args, + dt, + ni, + di, + li; + ϵ = ϵ, + CFL = CFL, + ) + end + # Boundary conditions function JustRelax.JustRelax2D.flow_bcs!( ::CUDABackendTrait, stokes::StokesArrays, bcs diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 36c06769..49dd59a1 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -12,8 +12,9 @@ using StaticArrays function solve!() end -struct CPUBackend end -struct AMDGPUBackend end +abstract type AbstractBackend end +struct CPUBackend <: AbstractBackend end +struct AMDGPUBackend <: AbstractBackend end PTArray() = Array PTArray(::Type{CPUBackend}) = Array diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index b8de9f8c..df4388ce 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -1,5 +1,5 @@ -# with phase ratios function PTThermalCoeffs( + ::Type{CPUBackend}, rheology, phase_ratios, args, @@ -7,8 +7,8 @@ function PTThermalCoeffs( ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, + ϵ = 1e-8, + CFL = 0.9 / √3, ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) @@ -23,7 +23,7 @@ end # without phase ratios function PTThermalCoeffs( - rheology, args, dt, ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 + ::Type{CPUBackend}, rheology, args, dt, ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl index 54d0695b..e3904cee 100644 --- a/src/types/constructors/heat_diffusion.jl +++ b/src/types/constructors/heat_diffusion.jl @@ -1,3 +1,10 @@ +function ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where N + return ThermalArrays(ni...) +end + +function ThermalArrays(::Type{CPUBackend}, ni::Vararg{Integer, N}) where N + return ThermalArrays(ni...) +end function ThermalArrays(nx::Integer, ny::Integer) T = @zeros(nx + 3, ny + 1) diff --git a/src/types/constructors/phases.jl b/src/types/constructors/phases.jl index f7f1d5a4..6826523f 100644 --- a/src/types/constructors/phases.jl +++ b/src/types/constructors/phases.jl @@ -1,3 +1,7 @@ +function PhaseRatio(::Type{CPUBackend}, ni::NTuple{N,Integer}, num_phases::Integer) where N + return PhaseRatio(ni, num_phases) +end + function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} center = @fill(0.0, ni..., celldims = (num_phases,)) vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl index fb01b005..ba2f567e 100644 --- a/src/types/constructors/stokes.jl +++ b/src/types/constructors/stokes.jl @@ -71,6 +71,10 @@ function Residual(nx::Integer, ny::Integer, nz::Integer) end ## StokesArrays type +function StokesArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where N + return StokesArrays(ni) +end + function StokesArrays(ni::NTuple{N,Integer}) where {N} P = @zeros(ni...) P0 = @zeros(ni...) From 3589482b644398543886621b8511a06ce5a7e3c7 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:19:30 +0200 Subject: [PATCH 13/91] add Blakenbach example to the docs --- docs/Project.toml | 2 +- docs/make.jl | 3 + docs/src/Blankenbach.md | 391 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 395 insertions(+), 1 deletion(-) create mode 100644 docs/src/Blankenbach.md diff --git a/docs/Project.toml b/docs/Project.toml index 69325320..1814eb33 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -2,4 +2,4 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" [compat] -Documenter = "1.2.1" +Documenter = "1" diff --git a/docs/make.jl b/docs/make.jl index e3e4f116..92709052 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -10,6 +10,9 @@ makedocs(; warnonly = Documenter.except(:footnote), pages=[ "Home" => "index.md", + "Examples" => [ + "Blankenbach.md" + ], ], ) diff --git a/docs/src/Blankenbach.md b/docs/src/Blankenbach.md new file mode 100644 index 00000000..5886e32d --- /dev/null +++ b/docs/src/Blankenbach.md @@ -0,0 +1,391 @@ +# Blankenbach benchmark + +Thermal convection benchmark from [Blankenbach et al., 1989](https://academic.oup.com/gji/article/98/1/23/622167) + +# Initialize packages + +Load JustRelax necessary modules and define backend. +```julia +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend +``` + +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 = JustPIC.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 to define and compute physical properties of the materials: +```julia +using GeoParams +``` + +# Script + +## Model domain +```julia +nx = ny = 51 # number of cells per dimension +nit = 6e3 +igg = IGG( + init_global_grid(nx, ny, 1; init_MPI= true)... +) # initialize MPI grid +ly = 1.0 # domain length in y +lx = ly * ar # 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 = dt_diff = 0.9 * min(di...)^2 / 4.0 # diffusive CFL timestep limiter +``` + +## Rheology +```julia +rheology = ( + SetMaterialParams(; + Phase = 1, + Density = PT_Density(; ρ0 = 1, α = 1, β = 0.0), + HeatCapacity = ConstantHeatCapacity(; Cp = 1.0), + Conductivity = ConstantConductivity(; k = 1.0), + CompositeRheology = CompositeRheology((LinearViscous(; η = 1),)), + RadioactiveHeat = ConstantRadioactiveHeat(0.0), + Gravity = ConstantGravity(; g = 1e4), + ), +) +``` + +## Initialize particles +```julia +nxcell = 24 # initial number of perticles per cell +max_xcell = 35 # maximum number of perticles per cell +min_xcell = 12 # minimum number of perticles per cell +particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... +) # particles object +subgrid_arrays = SubgridDiffusionCellArrays(particles) # arrays needed for subgrid diffusion +# velocity grids +grid_vx, grid_vy = velocity_grids(xci, xvi, di) # staggered velocity grids +``` + +and we want to keep track of the temperature `pT`, temperature of the previous time step `pT0`, and material phase `pPhase`: + +```julia +pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) +particle_args = (pT, pT0, pPhases) +``` + +# Temperature anomaly +```julia +xc_anomaly = 0.0 # origin of thermal anomaly +yc_anomaly = 1 / 3 # origin of thermal anomaly +r_anomaly = 0.1 / 2 # radius of perturbation +``` + +Helper function to initialize material phases with `ParallelStencil.jl` +```julia +function init_phases!(phases, particles) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, index) + @inbounds for ip in JustRelax.cellaxes(phases) + # quick escape if the ip-th element of the [i,j]-th cell is empty + JustRelax.@cell(index[ip, i, j]) == 0 && continue + # all particles have phase number = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.index) +end + +init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly) +``` + +or we can use the alternative one-liners +```julia +@views pPhase.data[!isnan.(particles.index.data)] .= 1.0 +``` +or +```julia +map!(x -> isnan(x) ? NaN : 1.0, pPhase.data, particles.index.data) +``` + +and finally we need the phase ratios at the cell centers: +```julia +phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) +phase_ratios_center(phase_ratios, particles, grid, pPhases) +``` + +## Stokes and heat diffusion arrays + +Stokes arrays object +```julia +stokes = StokesArrays(backend_JR, ni) +``` + +and the correspondant heat diffusion one +```julia +thermal = ThermalArrays(backend_JR, ni) +``` + +## Initialize thermal profile and viscosity fields + +To initialize the thermal profile we use `ParallelStencil.jl` again +```julia +@parallel_indices (i, j) function init_T!(T, y) + T[i, j] = 1 - y[j] + return nothing +end + +@parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[2]) # cell vertices +@parallel (@idx size(thermal.Tc)) init_T!(thermal.Tc, xci[2]) # cell centers +``` + +and we define a rectangular thermal anomaly at $x \in [0, 0.05]$, $y \in [\frac{1}{3} - 0.05, \frac{1}{3} + 0.05]$ +```julia +function rectangular_perturbation!(T, xc, yc, r, xvi) + @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + T[i, j] += .2 + end + return nothing + end + ni = size(T) + @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) + return nothing +end + +xc_anomaly = 0.0 # center of the thermal anomaly +yc_anomaly = 1/3 # center of the thermal anomaly +r_anomaly = 0.1/2 # half-width of the thermal anomaly +rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) +``` + +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), +) +thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), +) +thermal_bcs!(thermal.T, thermal_bc) +thermal.Told .= thermal.T +``` + +## Pseuo-transient coefficients +```julia +pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) +pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 +) +``` + +## Just before solving the problem... +We need to allocate some arrays to be able to do the subgrid diffusion of the temperature field at the particles level: +```julia +T_buffer = @zeros(ni.+1) # without the ghost nodes at the x-direction +Told_buffer = similar(T_buffer) # without the ghost nodes at the x-direction +dt₀ = similar(stokes.P) # subgrid diffusion time scale +# copy temperature to buffer arrays +for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) +end +# interpolate temperatyre on the particles +grid2particle!(pT, xvi, T_buffer, particles) +pT0.data .= pT.data +``` +where +```julia +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 +``` + +In this benchmark we want to keep track of the time `trms`, the rms-velocity `Urms` + +$$ U_{rms} = \sqrt{ \int_{\Omega} (V_x^2 + V_y^2 ) d\Omega} $$ + +and the Nusselt number at the top of the model `Nu_top` + +$$ Nu_{top} = \int \frac{\partial T}{\partial x} dx $$ + +And we will store their time history in the vectors: +```julia +Urms = Float64[] +Nu_top = Float64[] +trms = Float64[] +``` + +We further need two buffer arrays where to interpolate the velocity field at the vertices of the grid cells +```julia +# Buffer arrays to compute velocity rms +Vx_v = @zeros(ni.+1...) +Vy_v = @zeros(ni.+1...) +``` + +## Advancing one time step + +1. Solve stokes +```julia +solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + Inf, + igg; + kwargs = (; + iterMax = 150e3, + nout = 200, + viscosity_cutoff = (-Inf, Inf), + verbose = true + ) +) +# calculate adaptive time step +dt = compute_dt(stokes, di, dt_diff) +``` +2. Heat diffusion solver +```julia +heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) +) +``` +3. Subgrid diffusion at the particle level +```julia +for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) +end +subgrid_characteristic_time!( + subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di +) +centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles) +subgrid_diffusion!( + pT, T_buffer, thermal.ΔT[2:end-1, :], subgrid_arrays, particles, xvi, di, dt +) +``` +4. Advect particles +```julia +# advect particles in space +advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), 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 +phase_ratios_center(phase_ratios, particles, grid, pPhases) +``` + +5. Interpolate `T` back to the grid +```julia +# interpolate fields from particle to grid vertices +particle2grid!(T_buffer, pT, xvi, particles) +@views T_buffer[:, end] .= 0.0 +@views T_buffer[:, 1] .= 1.0 +@views thermal.T[2:end-1, :] .= T_buffer +flow_bcs!(stokes, flow_bcs) # apply boundary conditions +temperature2center!(thermal) +``` + +6. Update buoyancy forces and viscosity +```julia +args = (; T = thermal.Tc, P = stokes.P, dt=Inf) +compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) +compute_ρg!(ρg[2], phase_ratios, rheology, args) +``` +7. Compute Nusselt number and rms-velocity +```julia +# Nusselt number, Nu = ∫ ∂T/∂z dx +Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) +push!(Nu_top, Nu_it) +# Compute U rms +# U₍ᵣₘₛ₎ = √ ∫∫ (vx²+vz²) dx dz +Urms_it = let + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) + @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude + sqrt(sum( Vx_v.^2 .* prod(di)) ) +end +push!(Urms, Urms_it) +push!(trms, t) +``` + +# Visualization +We will use `Makie.jl` to visualize the results +```julia +using GLMakie +``` + +## Fields +```julia +# Make particles plottable +p = particles.coords +ppx, ppy = p +pxv = ppx.data[:] +pyv = ppy.data[:] +clr = pT.data[:] +idxv = particles.index.data[:]; + +# Make Makie figure +fig = Figure(size = (900, 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 = "Vx [m/s]") +ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") +# grid temperature +h1 = heatmap!(ax1, xvi[1], xvi[2], Array(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(0, 1) ) +# y-velocity +h2 = heatmap!(ax2, xvi[1], xvi[2], Array(stokes.V.Vy) , colormap=:batlow) +# x-velocity +h3 = heatmap!(ax3, xvi[1], xvi[2], Array(stokes.V.Vx) , colormap=:batlow) +# particles temperature +h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) +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 +``` \ No newline at end of file From 8c49e8fd9b045e18129a42bd004a981d2063eaa2 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:22:33 +0200 Subject: [PATCH 14/91] change default viscosity cutoff --- src/stokes/Stokes2D.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index b041aa6a..2dd64c84 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -269,7 +269,7 @@ function _solve!( args, dt, igg::IGG; - viscosity_cutoff=(1e16, 1e24), + viscosity_cutoff=(-Inf, Inf), iterMax=10e3, nout=500, b_width=(4, 4, 0), @@ -439,7 +439,7 @@ function _solve!( args, dt, igg::IGG; - viscosity_cutoff=(1e16, 1e24), + viscosity_cutoff=(-Inf, Inf), iterMax=50e3, iterMin=1e2, viscosity_relaxation=1e-2, From 5e9cfa42cc13a52307aa955c0592b4e6275a78b7 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:22:54 +0200 Subject: [PATCH 15/91] add DataIO.jl back --- src/JustRelax.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 49dd59a1..456d0948 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -42,6 +42,6 @@ export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! include("JustRelax_CPU.jl") -# include("IO/DataIO.jl") +include("IO/DataIO.jl") end # module From 4427ea4de9061252fe609d46e2e1510b87f5df77 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:23:17 +0200 Subject: [PATCH 16/91] update Blankenbach --- .../Blankenbach2D/Benchmark2D_sgd_scaled.jl | 111 +++++++----------- .../Blankenbach_Rheology_scaled.jl | 6 +- 2 files changed, 43 insertions(+), 74 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl index a084e4a3..d89d3cb8 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -1,21 +1,17 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +const backend_JR = JustRelax.CPUBackend using JustPIC using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using Printf, LinearAlgebra, GeoParams, GLMakie # Load file with all the rheology configurations include("Blankenbach_Rheology_scaled.jl") @@ -23,27 +19,13 @@ include("Blankenbach_Rheology_scaled.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- function copyinn_x!(A, B) - @parallel function f_x(A, B) @all(A) = @inn_x(B) return nothing end - @parallel f_x(A, B) end -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - # Initial thermal profile @parallel_indices (i, j) function init_T!(T, y) T[i, j] = 1 - y[j] @@ -52,7 +34,6 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) - @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += .2 @@ -61,7 +42,6 @@ function rectangular_perturbation!(T, xc, yc, r, xvi) end ni = size(T) @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) - return nothing end ## END OF HELPER FUNCTION ------------------------------------------------------------ @@ -97,32 +77,32 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) particle_args = (pT, pT0, pPhases) - # Elliptical temperature anomaly - xc_anomaly = 0.0 # origin of thermal anomaly - yc_anomaly = 1/3 # origin of thermal anomaly - r_anomaly = 0.1/2 # radius of perturbation init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.85 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) # initialize thermal profile @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[2]) + # Elliptical temperature anomaly + xc_anomaly = 0.0 # origin of thermal anomaly + yc_anomaly = 1/3 # origin of thermal anomaly + r_anomaly = 0.1/2 # radius of perturbation rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) thermal_bcs!(thermal.T, thermal_bc) thermal.Told .= thermal.T - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Rayleigh number ------------------------------------ @@ -131,21 +111,17 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Buoyancy forces ------------------------------------ ρg = @zeros(ni...), @zeros(ni...) - for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - end + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) # 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, (1e19, 1e25) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - η_vep = copy(η) - + # PT coefficients for thermal diffusion ------------- pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 ) # Boundary conditions ------------------------------- @@ -205,14 +181,11 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk Vy_v = @zeros(ni.+1...) while it ≤ nit - @show it # 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, (-Inf, Inf) - ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -222,17 +195,17 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 150e3, - nout = 200, - viscosity_cutoff = (-Inf, Inf), - verbose = false + kwargs = (; + iterMax = 150e3, + nout = 200, + viscosity_cutoff = (-Inf, Inf), + verbose = true + ) ) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -246,12 +219,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 50, - verbose = false, + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) ) + # subgrid diffusion for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) end @@ -266,17 +242,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # clean_particles!(particles, xvi, particle_args) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # advect particles in memory move_particles!(particles, xvi, particle_args) - # clean_particles!(particles, xvi, particle_args) # check if we need to inject particles - inject = check_injection(particles) - # inject && break - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # Nusselt number, Nu = ∫ ∂T/∂z dx ---- Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) @@ -286,9 +258,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Compute U rms ----------------------------- # U₍ᵣₘₛ₎ = √ ∫∫ (vx²+vz²) dx dz Urms_it = let - JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude - sqrt( sum( Vx_v.^2 .* prod(di)) ) + sqrt(sum( Vx_v.^2 .* prod(di)) ) end push!(Urms, Urms_it) push!(trms, t) @@ -306,10 +278,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit || any(isnan.(thermal.T)) - checkpointing(figdir, stokes, thermal.T, η, t) if save_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), τxy = Array(stokes.τ.xy), diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl index 63a5c9c4..02605263 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl @@ -1,10 +1,8 @@ # from "A benchmark comparison for mantle convection codes"; Blankenbach et al., 1989 function init_rheologies() - # Define rheolgy struct rheology = ( - # Name = "UpperCrust", SetMaterialParams(; Phase = 1, Density = PT_Density(; ρ0=1, α = 1, β = 0.0), @@ -25,11 +23,11 @@ function init_phases!(phases, particles, Lx, d, r) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue - @cell phases[ip, i, j] = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 end return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) end From d506eee02bd8b20bbefdd67d7218bd3deaf927a7 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:25:43 +0200 Subject: [PATCH 17/91] typo --- docs/src/Blankenbach.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/Blankenbach.md b/docs/src/Blankenbach.md index 5886e32d..38b15379 100644 --- a/docs/src/Blankenbach.md +++ b/docs/src/Blankenbach.md @@ -131,7 +131,7 @@ Stokes arrays object stokes = StokesArrays(backend_JR, ni) ``` -and the correspondant heat diffusion one +and the correspondent heat diffusion one ```julia thermal = ThermalArrays(backend_JR, ni) ``` From 4777e40c44612ac9c4805e0b55b09204ff9d284d Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sat, 27 Apr 2024 22:25:54 +0200 Subject: [PATCH 18/91] format --- ext/JustRelaxCUDAExt.jl | 76 ++++++------------- .../DiffusionPT_coefficients.jl | 14 +++- src/types/constructors/heat_diffusion.jl | 4 +- src/types/constructors/phases.jl | 4 +- src/types/constructors/stokes.jl | 2 +- 5 files changed, 42 insertions(+), 58 deletions(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 9056594b..69820181 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -18,13 +18,7 @@ module JustRelax2D using MPI import JustRelax: - IGG, - BackendTrait, - CPUBackendTrait, - backend, - CPUBackend, - Geometry, - @cell + IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend, Geometry, @cell @init_parallel_stencil(CUDA, Float64, 2) @@ -69,23 +63,14 @@ module JustRelax2D ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; - ϵ = 1e-8, - CFL = 0.9 / √3, + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} return JustRelax.JustRelax2D.PTThermalCoeffs( - rheology, - phase_ratios, - args, - dt, - ni, - di, - li; - ϵ = ϵ, - CFL = CFL, + rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL ) end - - + function PTThermalCoeffs( ::Type{CUDABackend}, rheology, @@ -94,31 +79,20 @@ module JustRelax2D ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; - ϵ = 1e-8, - CFL = 0.9 / √3, + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} return JustRelax.JustRelax2D.PTThermalCoeffs( - rheology, - args, - dt, - ni, - di, - li; - ϵ = ϵ, - CFL = CFL, + rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL ) end # Boundary conditions - function JustRelax.JustRelax2D.flow_bcs!( - ::CUDABackendTrait, stokes::StokesArrays, bcs - ) + function JustRelax.JustRelax2D.flow_bcs!(::CUDABackendTrait, stokes::StokesArrays, bcs) return _flow_bcs!(bcs, @velocity(stokes)) end - function flow_bcs!( - ::CUDABackendTrait, stokes::StokesArrays, bcs - ) + function flow_bcs!(::CUDABackendTrait, stokes::StokesArrays, bcs) return _flow_bcs!(bcs, @velocity(stokes)) end @@ -127,10 +101,8 @@ module JustRelax2D ) return thermal_bcs!(thermal.T, bcs) end - - function thermal_bcs!( - ::CUDABackendTrait, thermal::ThermalArrays, bcs - ) + + function thermal_bcs!(::CUDABackendTrait, thermal::ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end @@ -159,9 +131,7 @@ module JustRelax2D return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end - function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, args, rheology, cutoff - ) + function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end function compute_viscosity!( @@ -169,9 +139,7 @@ module JustRelax2D ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end - function compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end @@ -210,13 +178,19 @@ module JustRelax2D end # Solvers - JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) - + function JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) + end + # Utils - JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) = compute_dt(S, di, dt_diff, I::IGG) - JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff) = compute_dt(S, di, dt_diff) + function JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) + return compute_dt(S, di, dt_diff, I::IGG) + end + function JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff) + return compute_dt(S, di, dt_diff) + end JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di) = compute_dt(S, di) - + end end diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index df4388ce..e39e5ab9 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -7,8 +7,8 @@ function PTThermalCoeffs( ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; - ϵ = 1e-8, - CFL = 0.9 / √3, + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) @@ -23,7 +23,15 @@ end # without phase ratios function PTThermalCoeffs( - ::Type{CPUBackend}, rheology, args, dt, ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 + ::Type{CPUBackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl index e3904cee..be01a4c8 100644 --- a/src/types/constructors/heat_diffusion.jl +++ b/src/types/constructors/heat_diffusion.jl @@ -1,8 +1,8 @@ -function ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where N +function ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where {N} return ThermalArrays(ni...) end -function ThermalArrays(::Type{CPUBackend}, ni::Vararg{Integer, N}) where N +function ThermalArrays(::Type{CPUBackend}, ni::Vararg{Integer,N}) where {N} return ThermalArrays(ni...) end diff --git a/src/types/constructors/phases.jl b/src/types/constructors/phases.jl index 6826523f..432a17cd 100644 --- a/src/types/constructors/phases.jl +++ b/src/types/constructors/phases.jl @@ -1,4 +1,6 @@ -function PhaseRatio(::Type{CPUBackend}, ni::NTuple{N,Integer}, num_phases::Integer) where N +function PhaseRatio( + ::Type{CPUBackend}, ni::NTuple{N,Integer}, num_phases::Integer +) where {N} return PhaseRatio(ni, num_phases) end diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl index ba2f567e..66825e0f 100644 --- a/src/types/constructors/stokes.jl +++ b/src/types/constructors/stokes.jl @@ -71,7 +71,7 @@ function Residual(nx::Integer, ny::Integer, nz::Integer) end ## StokesArrays type -function StokesArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where N +function StokesArrays(::Type{CPUBackend}, ni::NTuple{N,Integer}) where {N} return StokesArrays(ni) end From 9f4e04f8c7021403394125d417c16212778a93be Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:22:14 +0200 Subject: [PATCH 19/91] up docs --- docs/src/Blankenbach.md | 11 +++++++++-- docs/src/advection.md | 8 ++++++++ docs/src/backend.md | 28 ++++++++++++++++++++++++++++ docs/src/equations.md | 7 +++++++ docs/src/material_physics.md | 4 ++++ 5 files changed, 56 insertions(+), 2 deletions(-) create mode 100644 docs/src/advection.md create mode 100644 docs/src/backend.md create mode 100644 docs/src/equations.md create mode 100644 docs/src/material_physics.md diff --git a/docs/src/Blankenbach.md b/docs/src/Blankenbach.md index 38b15379..dbb6abe5 100644 --- a/docs/src/Blankenbach.md +++ b/docs/src/Blankenbach.md @@ -21,7 +21,7 @@ We will also use `ParallelStencil.jl` to write some device-agnostic helper funct using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) ``` -and will use GeoParams to define and compute physical properties of the materials: +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 ``` @@ -388,4 +388,11 @@ Colorbar(fig[2,4], h4) linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig -``` \ No newline at end of file +``` + +### Final model +Temperature field +![Temperatuere](../assets/Blankenbach/Temp.png) + +And time history of the rms-velocity and Nusselt number +![time series](../assets/Blankenbach/Time_Series_V_Nu.png) diff --git a/docs/src/advection.md b/docs/src/advection.md new file mode 100644 index 00000000..2ce1fc37 --- /dev/null +++ b/docs/src/advection.md @@ -0,0 +1,8 @@ +# Field advection + +## Particles-in-Cell +`JustRelax` relies on [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl) for advections of particles containing material information. + +## Upwind + +## WENO5 \ No newline at end of file diff --git a/docs/src/backend.md b/docs/src/backend.md new file mode 100644 index 00000000..39a75f2f --- /dev/null +++ b/docs/src/backend.md @@ -0,0 +1,28 @@ +# Selecting the backend + +JustRelax supports three backends: CPU, and CUDA and AMD GPU cards. To use the default CPU backend, simply load JustRelax: + +```julia +using JustRelax +``` + +The GPU backends are implemented as extensions, so it is enough to load the appropiate GPU Pkg before loading JustRelax. That is, to use CUDA cards: +```julia +using CUDA, JustRelax +``` +and for AMD cards: +```julia +using AMDGPU, JustRelax +``` + +Two and three dimensional solvers are implemented in different submodules, which also need to be loaded. To use the two-dimensional backend: + +```julia +using JustRelax.JustRelax2D +``` + +and for the three-dimensional backend: + +```julia +using JustRelax.JustRelax3D +``` \ No newline at end of file diff --git a/docs/src/equations.md b/docs/src/equations.md new file mode 100644 index 00000000..0154a877 --- /dev/null +++ b/docs/src/equations.md @@ -0,0 +1,7 @@ +# Pseudo-transient iterative method + +## Heat diffusion + +## Stokes equations + +## Constitutive equations diff --git a/docs/src/material_physics.md b/docs/src/material_physics.md new file mode 100644 index 00000000..7fb95b6e --- /dev/null +++ b/docs/src/material_physics.md @@ -0,0 +1,4 @@ +# Material physics + +`JustRelax` is fully integrated with [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl/tree/main) to perform all the material physics computations. + From df21b12910c06d82f229a2b5838b565701aaf833 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:22:44 +0200 Subject: [PATCH 20/91] remove CUDA and AMDGPU deps --- src/IO/DataIO.jl | 1 - src/IO/H5.jl | 8 ++------ 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/IO/DataIO.jl b/src/IO/DataIO.jl index 76302f52..c44a8119 100644 --- a/src/IO/DataIO.jl +++ b/src/IO/DataIO.jl @@ -3,7 +3,6 @@ module DataIO using WriteVTK using HDF5 using MPI -using CUDA, AMDGPU import ..JustRelax: Geometry diff --git a/src/IO/H5.jl b/src/IO/H5.jl index 5caae2d4..2c7ac491 100644 --- a/src/IO/H5.jl +++ b/src/IO/H5.jl @@ -4,14 +4,10 @@ macro namevar(x) name = split(string(x), ".")[end] return quote tmp = $(esc(x)) - $(esc(name)), _tocpu(tmp) + $(esc(name)), Array(tmp) end end -_tocpu(x) = x -_tocpu(x::T) where {T<:CuArray} = Array(x) -_tocpu(x::T) where {T<:ROCArray} = Array(x) - """ checkpointing(dst, stokes, T, η, time) @@ -26,7 +22,7 @@ function checkpointing(dst, stokes, T, η, time) write(file, @namevar(stokes.V.Vy)...) write(file, @namevar(stokes.P)...) write(file, @namevar(T)...) - write(file, "viscosity", _tocpu(η)) + write(file, "viscosity", Array(η)) end end From ed56bc96366f35ed311110c950d05fb42691aca6 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:43:52 +0200 Subject: [PATCH 21/91] typo --- docs/src/backend.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/backend.md b/docs/src/backend.md index 39a75f2f..98bd5105 100644 --- a/docs/src/backend.md +++ b/docs/src/backend.md @@ -6,7 +6,7 @@ JustRelax supports three backends: CPU, and CUDA and AMD GPU cards. To use the d using JustRelax ``` -The GPU backends are implemented as extensions, so it is enough to load the appropiate GPU Pkg before loading JustRelax. That is, to use CUDA cards: +The GPU backends are implemented as extensions, so it is enough to load the appropriate GPU Pkg before loading JustRelax. That is, to use CUDA cards: ```julia using CUDA, JustRelax ``` From ec1fb174de59f2746a9cef05dd30b62dea942ff0 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:48:10 +0200 Subject: [PATCH 22/91] Blankenbachs up to speed --- .../Blankenbach2D/Benchmark2D_WENO5.jl | 183 ++++++++---------- .../stokes2D/Blankenbach2D/Benchmark2D_sgd.jl | 127 +++++------- .../Blankenbach2D/Benchmark2D_sgd_scaled.jl | 53 ++--- .../Blankenbach2D/Blankenbach_Rheology.jl | 11 +- .../Blankenbach_Rheology_scaled.jl | 9 +- 5 files changed, 171 insertions(+), 212 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl index 385255ea..664c4579 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl @@ -1,6 +1,7 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -8,30 +9,23 @@ using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions, CellArrays +using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays # Load file with all the rheology configurations include("Blankenbach_Rheology.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- +function copyinn_x!(A, B) -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing + @parallel f_x(A, B) end # Initial thermal profile @@ -46,7 +40,6 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) - @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += 20.0 @@ -55,17 +48,16 @@ function rectangular_perturbation!(T, xc, yc, r, xvi) end ni = size(T) @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) - return nothing end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk =false) +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) # Physical domain ------------------------------------ ly = 1000e3 # domain length in y - lx = ly * ar # domain length in x + 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 @@ -83,70 +75,61 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Weno model ----------------------------------------- weno = WENO5(ni=(nx,ny).+1, method=Val{2}()) # ni.+1 for Temp # ---------------------------------------------------- - + # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 12, 24, 6 + nxcell, max_xcell, min_xcell = 24, 36, 12 particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) - subgrid_arrays = SubgridDiffusionCellArrays(particles) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) # temperature - pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) - particle_args = (pT, pT0, pPhases) - - # Elliptical temperature anomaly - xc_anomaly = 0.0 # origin of thermal anomaly - yc_anomaly = -600e3 # origin of thermal anomaly - r_anomaly = 100e3 # radius of perturbation - init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + pPhases, = init_cell_arrays(particles, Val(1)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.85 / √2.1) + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) # initialize thermal profile @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[2]) + # Elliptical temperature anomaly + xc_anomaly = 0.0 # origin of thermal anomaly + yc_anomaly = -600e3 # origin of thermal anomaly + r_anomaly = 100e3 # radius of perturbation rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) thermal_bcs!(thermal.T, thermal_bc) thermal.Told .= thermal.T - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Rayleigh number - ΔT = thermal.T[1,1] - thermal.T[1,end] - Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / - (κ * rheology[1].CompositeRheology[1].elements[1].η ) + ΔT = thermal.T[1,1] - thermal.T[1,end] + Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / + (κ * rheology[1].CompositeRheology[1].elements[1].η ) @show Ra - # Buoyancy forces ------------------------------------ + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + + # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) - for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - 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, (1e19, 1e25) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - η_vep = copy(η) # PT coefficients for thermal diffusion ------------- pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 0.5 / √2.1 ) # Boundary conditions ------------------------------- @@ -181,9 +164,6 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk fig end - # WENO arrays - T_WENO = @zeros(ni.+1) - local Vx_v, Vy_v if save_vtk Vx_v = @zeros(ni.+1...) @@ -199,15 +179,16 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) + # WENO arrays + T_WENO = @zeros(ni.+1) + while it ≤ nit @show it # 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, (-Inf, Inf) - ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -217,22 +198,22 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 150e3, - nout = 50, - viscosity_cutoff = (-Inf, Inf), - verbose = false + kwargs = (; + iterMax = 150e3, + nout = 200, + viscosity_cutoff = (-Inf, Inf), + verbose = true + ) ) - dt = compute_dt(stokes, di, dt_diff) + dt = compute_dt(stokes, di, dt_diff) # ------------------------------ - # Weno advection + # Thermal solver --------------- heatdiffusion_PT!( thermal, pt_thermal, @@ -241,32 +222,22 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = false, + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) ) - T_WENO .= thermal.T[2:end-1, :] - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + @views T_WENO .= thermal.T[2:end-1, :] + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) - thermal.T[2:end-1, :] .= T_WENO + @views thermal.T[2:end-1, :] .= T_WENO + @views thermal.T[2:end-1,end] .= 273.0 + @views thermal.T[2:end-1,1] .= 1273.0 + temperature2center!(thermal) # ------------------------------ - thermal.T[2:end-1,end] .= 273.0 - thermal.T[2:end-1,1] .= 1273.0 - - # Advection -------------------- - # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # clean_particles!(particles, xvi, particle_args) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # clean_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, pPhases) # Nusselt number, Nu = H/ΔT/L ∫ ∂T/∂z dx ---- Nu_it = (ly / (1000.0*lx)) * @@ -277,7 +248,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Compute U rms ----------------------------- # U₍ᵣₘₛ₎ = H*ρ₀*c₍ₚ₎/k * √ 1/H/L * ∫∫ (vx²+vz²) dx dz Urms_it = let - JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) + # JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * ((ly * rheology[1].Density[1].ρ0 * rheology[1].HeatCapacity[1].Cp) / rheology[1].Conductivity[1].k ) @@ -288,7 +259,6 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit - checkpointing(figdir, stokes, thermal.T, η, t) if save_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) @@ -315,29 +285,38 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk data_c ) end + + # # Make particles plottable + # p = particles.coords + # ppx, ppy = p + # pxv = ppx.data[:]./1e3 + # pyv = ppy.data[:]./1e3 + # clr = pT.data[:] #pPhases.data[:] + # idxv = particles.index.data[:]; + # Make Makie figure fig = Figure(size = (900, 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 = "Vx [m/s]") - #ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") + ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") # - h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(T_WENO) , colormap=:lajolla, colorrange=(273,1273) ) + h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(273,1273) ) # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) # h3 = heatmap!(ax3, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vx) , colormap=:batlow) # - #h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) - #h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)) , colormap=:batlow) + # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)) , 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) + Colorbar(fig[2,4], h4) + linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl index 0e093229..0d11d5a4 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl @@ -1,6 +1,7 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -8,11 +9,7 @@ using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays @@ -32,18 +29,6 @@ function copyinn_x!(A, B) @parallel f_x(A, B) end -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - # Initial thermal profile @parallel_indices (i, j) function init_T!(T, y) depth = -y[j] @@ -56,7 +41,6 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) - @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += 20.0 @@ -65,17 +49,16 @@ function rectangular_perturbation!(T, xc, yc, r, xvi) end ni = size(T) @parallel (@idx ni) _rectangular_perturbation!(T, xc, yc, r, xvi...) - return nothing end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk =false) +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) # Physical domain ------------------------------------ ly = 1000e3 # domain length in y - lx = ly * ar # domain length in x + 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 @@ -91,9 +74,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 12, 24, 6 + nxcell, max_xcell, min_xcell = 24, 36, 12 particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) subgrid_arrays = SubgridDiffusionCellArrays(particles) # velocity grids @@ -101,58 +84,53 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # temperature pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) particle_args = (pT, pT0, pPhases) - - # Elliptical temperature anomaly - xc_anomaly = 0.0 # origin of thermal anomaly - yc_anomaly = -600e3 # origin of thermal anomaly - r_anomaly = 100e3 # radius of perturbation - init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.85 / √2.1) + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) # initialize thermal profile @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[2]) + # Elliptical temperature anomaly + xc_anomaly = 0.0 # origin of thermal anomaly + yc_anomaly = -600e3 # origin of thermal anomaly + r_anomaly = 100e3 # radius of perturbation rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) thermal_bcs!(thermal.T, thermal_bc) thermal.Told .= thermal.T - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Rayleigh number - ΔT = thermal.T[1,1] - thermal.T[1,end] - Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / - (κ * rheology[1].CompositeRheology[1].elements[1].η ) + ΔT = thermal.T[1,1] - thermal.T[1,end] + Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / + (κ * rheology[1].CompositeRheology[1].elements[1].η ) @show Ra - # Buoyancy forces ------------------------------------ + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + + # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) - for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - 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, (1e19, 1e25) + compute_ρg!(ρg[2], phase_ratios, rheology, args) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - η_vep = copy(η) # PT coefficients for thermal diffusion ------------- pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 0.5 / √2.1 ) # Boundary conditions ------------------------------- @@ -216,10 +194,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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, (-Inf, Inf) - ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -229,19 +205,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 150e3, - nout = 200, - viscosity_cutoff = (-Inf, Inf), - verbose = false + kwargs = (; + iterMax = 150e3, + nout = 200, + viscosity_cutoff = (-Inf, Inf), + verbose = true + ) ) - dt = compute_dt(stokes, di, dt_diff) + dt = compute_dt(stokes, di, dt_diff) # ------------------------------ # Thermal solver --------------- @@ -253,11 +229,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 50, - verbose = true, + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) ) for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) @@ -273,17 +251,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # clean_particles!(particles, xvi, particle_args) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # advect particles in memory move_particles!(particles, xvi, particle_args) - # clean_particles!(particles, xvi, particle_args) # check if we need to inject particles - inject = check_injection(particles) - # inject && break - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer, ), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # Nusselt number, Nu = H/ΔT/L ∫ ∂T/∂z dx ---- Nu_it = (ly / (1000.0*lx)) * @@ -294,7 +268,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Compute U rms ----------------------------- # U₍ᵣₘₛ₎ = H*ρ₀*c₍ₚ₎/k * √ 1/H/L * ∫∫ (vx²+vz²) dx dz Urms_it = let - JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * ((ly * rheology[1].Density[1].ρ0 * rheology[1].HeatCapacity[1].Cp) / rheology[1].Conductivity[1].k ) @@ -313,7 +287,6 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit - checkpointing(figdir, stokes, thermal.T, η, t) if save_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl index d89d3cb8..12c327ed 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -1,7 +1,8 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend + using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -const backend_JR = JustRelax.CPUBackend using JustPIC using JustPIC._2D @@ -17,7 +18,6 @@ using Printf, LinearAlgebra, GeoParams, GLMakie include("Blankenbach_Rheology_scaled.jl") ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- - function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -47,11 +47,11 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk =false) +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) # Physical domain ------------------------------------ ly = 1.0 # domain length in y - lx = ly * ar # domain length in x + 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 @@ -66,9 +66,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 12, 24, 6 + nxcell, max_xcell, min_xcell = 24, 36, 12 particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + backend, nxcell, max_xcell, min_xcell, xvi, di, ni ) subgrid_arrays = SubgridDiffusionCellArrays(particles) # velocity grids @@ -76,16 +76,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # temperature pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) particle_args = (pT, pT0, pPhases) - - init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles) phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.85 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- @@ -109,19 +108,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk Ra = rheology[1].Gravity[1].g println("Ra = $Ra") - # Buoyancy forces ------------------------------------ + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + + # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - # Rheology ------------------------------------------ η = @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) ) - + # PT coefficients for thermal diffusion ------------- pt_thermal = PTThermalCoeffs( - backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 1e-1 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 0.5 / √2.1 ) # Boundary conditions ------------------------------- @@ -207,7 +206,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk verbose = true ) ) - dt = compute_dt(stokes, di, dt_diff) + dt = compute_dt(stokes, di, dt_diff) # ------------------------------ # Thermal solver --------------- @@ -326,15 +325,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # h3 = heatmap!(ax3, xvi[1], xvi[2], Array(stokes.V.Vx) , colormap=:batlow) # - h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), markersize=3) - # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) + # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), markersize=3) + h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) 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) + Colorbar(fig[2,4], h4) linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig @@ -345,6 +344,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk l1 = lines!(ax21,trms,(Urms)) l2 = lines!(ax22,trms,(Nu_top)) save(joinpath(figdir, "Time_Series_V_Nu.png"), fig2) + + cmap = ([:white, :white, :white, :white]) + fig3 = Figure(size = (900, 900), title = "t = $t") + ax = Axis(fig3[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + h1 = heatmap!(ax, xvi..., thermal.T[2:end-1,:], colormap=:lipari, colorrange=(0, 1)) + contour!(ax, xvi..., thermal.T[2:end-1,:], linewidth=5, levels= 0.2:0.2:0.8, colormap=cmap) + Colorbar(fig3[1,2], h1) + save(joinpath(figdir, "Temp.png"), fig3) + fig3 end it += 1 t += dt @@ -385,10 +393,10 @@ end figdir = "Blankenbach_subgrid_scaled" save_vtk = false # set to true to generate VTK files for ParaView ar = 1 # aspect ratio -n = 51 +n = 64 nx = n ny = n -nit = 6e3 +nit = 2e3#6e3 igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) else @@ -397,3 +405,4 @@ end # run main script main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, save_vtk = save_vtk); + diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl index 712df554..77da2cfd 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology.jl @@ -17,19 +17,18 @@ function init_rheologies() ) end -function init_phases!(phases, particles, Lx, d, r) +function init_phases!(phases, particles) ni = size(phases) - @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) + @parallel_indices (i, j) function init_phases!(phases, index) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue - - @cell phases[ip, i, j] = 1.0 - + JustRelax.@cell phases[ip, i, j] = 1.0 end return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) + @parallel (@idx ni) init_phases!(phases, particles.index) + return nothing end \ No newline at end of file diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl index 02605263..15954fa3 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Blankenbach_Rheology_scaled.jl @@ -15,19 +15,18 @@ function init_rheologies() ) end -function init_phases!(phases, particles, Lx, d, r) +function init_phases!(phases, particles) ni = size(phases) - @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx) + @parallel_indices (i, j) function init_phases!(phases, index) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue - JustRelax.@cell phases[ip, i, j] = 1.0 - end return nothing end - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) + @parallel (@idx ni) init_phases!(phases, particles.index) + return nothing end From 79ad82a8a0da42d0358a2dc110c6405202013407 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:48:21 +0200 Subject: [PATCH 23/91] fix WENO --- src/advection/weno5.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/advection/weno5.jl b/src/advection/weno5.jl index 3df1e91f..e7efa256 100644 --- a/src/advection/weno5.jl +++ b/src/advection/weno5.jl @@ -47,11 +47,11 @@ The `WENO5` structure contains the parameters and temporary variables used in th # grid size ni::NTuple{N,Int64} # fluxes - ut::A = PTArray(zeros(ni...)) - fL::A = PTArray(zeros(ni...)) - fR::A = PTArray(zeros(ni...)) - fB::A = PTArray(zeros(ni...)) - fT::A = PTArray(zeros(ni...)) + ut::A = zeros(ni...) + fL::A = zeros(ni...) + fR::A = zeros(ni...) + fB::A = zeros(ni...) + fT::A = zeros(ni...) # method method::M = Val{1} # 1:JS, 2:Z end From 24e05484a3de8ce387986d824230eaa690932566 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 09:48:32 +0200 Subject: [PATCH 24/91] docs --- docs/make.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 92709052..5c42274e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -9,8 +9,11 @@ makedocs(; format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true"), # easier local build warnonly = Documenter.except(:footnote), pages=[ - "Home" => "index.md", - "Examples" => [ + "Home" => "index.md", + "Backend" => "backend.md", + "Equations" => "equations.md", + "Advection" => "advection.md", + "Examples" => [ "Blankenbach.md" ], ], From 0f1ea9db0a1993897fe84f4daa1ec372b7ff1f8c Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 18:16:58 +0200 Subject: [PATCH 25/91] fixes and polish 3D solvers --- src/Utils.jl | 4 +- src/boundaryconditions/BoundaryConditions.jl | 10 +-- src/common.jl | 1 - src/stokes/Stokes2D.jl | 5 +- src/stokes/Stokes3D.jl | 44 +++++----- src/types/constructors/heat_diffusion.jl | 11 ++- src/types/constructors/stokes.jl | 8 +- src/types/heat_diffusion.jl | 59 +------------- src/types/stokes.jl | 86 -------------------- 9 files changed, 47 insertions(+), 181 deletions(-) diff --git a/src/Utils.jl b/src/Utils.jl index b13f5c3e..bc204aa5 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -300,12 +300,12 @@ macro residuals(A) end @inline function unpack_residuals( - A::JustRelax.SymmetricTensor{<:AbstractArray{T,2}} + A::JustRelax.Residual{<:AbstractArray{T,2}} ) where {T} return A.Rx, A.Ry end @inline function unpack_residuals( - A::JustRelax.SymmetricTensor{<:AbstractArray{T,3}} + A::JustRelax.Residual{<:AbstractArray{T,3}} ) where {T} return A.Rx, A.Ry, A.Rz end diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 5dd126ef..6a73b505 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -243,11 +243,11 @@ function free_surface_bcs!( end end -function free_surface_bcs!(stokes::JustRelax.StokesArrays, bcs::FlowBoundaryConditions) - if bcs.free_surface - @views stokes.τ.yy[:, end] .= 0.0 - end -end +# function free_surface_bcs!(stokes::JustRelax.StokesArrays, bcs::FlowBoundaryConditions) +# if bcs.free_surface +# @views stokes.τ.yy[:, end] .= 0.0 +# end +# end # function free_surface_bcs!( # stokes::JustRelax.StokesArrays{A,B,C,D,E,3}, bcs::FlowBoundaryConditions diff --git a/src/common.jl b/src/common.jl index c475a4e8..8c4d3d80 100644 --- a/src/common.jl +++ b/src/common.jl @@ -36,7 +36,6 @@ include("MiniKernels.jl") include("phases/phases.jl") export fn_ratio, phase_ratios_center -# export PhaseRatio, fn_ratio, phase_ratios_center include("rheology/BuoyancyForces.jl") export compute_ρg! diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 2dd64c84..688b7aae 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -9,10 +9,11 @@ end ## 2D VISCO-ELASTIC STOKES SOLVER # backend trait -function solve!(stokes, args...; kwargs) +function solve!(stokes::JustRelax.StokesArrays, args...; kwargs) return solve!(backend(stokes), stokes, args...; kwargs) end +# entry point for extensions solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) function _solve!( @@ -558,7 +559,7 @@ function _solve!( dt, θ_dτ, ) - free_surface_bcs!(stokes, flow_bcs) + # free_surface_bcs!(stokes, flow_bcs) center2vertex!(stokes.τ.xy, stokes.τ.xy_c) update_halo!(stokes.τ.xy) diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index e4c725f0..6fc7b8ff 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -15,12 +15,14 @@ function update_τ_o!(stokes::JustRelax.StokesArrays) end ## 3D VISCO-ELASTIC STOKES SOLVER -function solve!(stokes::JustRelax.StokesArrays, args...) - return solve!(CPUBackendTrait(stokes), stokes, args...) +function solve!(stokes::JustRelax.StokesArrays, args...; kwargs) + return solve!(backend(stokes), stokes, args...; kwargs) end -function solve!( - ::CPUBackendTrait, +# entry point for extensions +solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) + +function _solve!( stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, @@ -34,6 +36,7 @@ function solve!( nout=500, b_width=(4, 4, 4), verbose=true, + kwargs..., ) where {T} # solver related @@ -153,8 +156,7 @@ end ## 3D VISCO-ELASTO-PLASTIC STOKES SOLVER WITH GeoParams.jl -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, @@ -168,6 +170,7 @@ function solve!( nout=500, b_width=(4, 4, 4), verbose=true, + kwargs..., ) where {T} # solver related @@ -326,8 +329,7 @@ function solve!( end # GeoParams and multiple phases -function solve!( - ::CPUBackendTrait, +function _solve!( stokes::JustRelax.StokesArrays, pt_stokes::PTStokesCoeffs, di::NTuple{3,T}, @@ -342,7 +344,9 @@ function solve!( nout=500, b_width=(4, 4, 4), verbose=true, + viscosity_relaxation=1e-2, viscosity_cutoff=(-Inf, Inf), + kwargs..., ) where {T,N} ## UNPACK @@ -399,19 +403,17 @@ function solve!( ) # Update buoyancy - @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + compute_ρg!(ρg[end], phase_ratios, rheology, args) # Update viscosity - ν = 1e-2 - @parallel (@idx ni) compute_viscosity!( - η, - ν, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) + compute_viscosity!( + stokes, + viscosity_relaxation, + phase_ratios, + args, + rheology, + viscosity_cutoff, + ) @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), @@ -430,9 +432,9 @@ function solve!( dt, pt_stokes.θ_dτ, ) - free_surface_bcs!(stokes, flow_bc) + # free_surface_bcs!(stokes, flow_bc) - @parallel (@idx ni .+ 1) center2vertex!( + center2vertex!( stokes.τ.yz, stokes.τ.xz, stokes.τ.xy, diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl index be01a4c8..d193aabc 100644 --- a/src/types/constructors/heat_diffusion.jl +++ b/src/types/constructors/heat_diffusion.jl @@ -9,6 +9,7 @@ end function ThermalArrays(nx::Integer, ny::Integer) T = @zeros(nx + 3, ny + 1) ΔT = @zeros(nx + 3, ny + 1) + ΔTc = @zeros(nx, ny) Told = @zeros(nx + 3, ny + 1) Tc = @zeros(nx, ny) H = @zeros(nx, ny) @@ -22,8 +23,9 @@ function ThermalArrays(nx::Integer, ny::Integer) return JustRelax.ThermalArrays( T, Tc, - ΔT, Told, + ΔT, + ΔTc, dT_dt, qTx, qTy, @@ -40,19 +42,20 @@ end function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) T = @zeros(nx + 1, ny + 1, nz + 1) ΔT = @zeros(nx + 1, ny + 1, nz + 1) + ΔTc = @zeros(nx, ny, ny) Told = @zeros(nx + 1, ny + 1, nz + 1) Tc = @zeros(nx, ny, nz) H = @zeros(nx, ny, nz) shear_heating = @zeros(nx, ny, nz) - dT_dt = @zeros(ni .- 1) + dT_dt = @zeros(nx -1, ny - 1, nz - 1) qTx = @zeros(nx, ny - 1, nz - 1) qTy = @zeros(nx - 1, ny, nz - 1) qTz = @zeros(nx - 1, ny - 1, nz) qTx2 = @zeros(nx, ny - 1, nz - 1) qTy2 = @zeros(nx - 1, ny, nz - 1) qTz2 = @zeros(nx - 1, ny - 1, nz) - ResT = @zeros((ni .- 1)...) + ResT = @zeros(nx -1, ny - 1, nz - 1) return JustRelax.ThermalArrays( - T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT + T, Tc, Told, ΔT,ΔTc, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT ) end diff --git a/src/types/constructors/stokes.jl b/src/types/constructors/stokes.jl index 66825e0f..db12d424 100644 --- a/src/types/constructors/stokes.jl +++ b/src/types/constructors/stokes.jl @@ -14,7 +14,7 @@ function Velocity(nx::Integer, ny::Integer, nz::Integer) nVz = (nx + 2, ny + 2, nz + 1) Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) - return JustRelax.Velocity{PTArray}(Vx, Vy, Vz) + return JustRelax.Velocity(Vx, Vy, Vz) end ## Viscosity type @@ -41,11 +41,11 @@ end function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) return JustRelax.SymmetricTensor( @zeros(nx, ny, nz), # xx - @zeros(nx + 1, ny + 1, nz), # xy @zeros(nx, ny, nz), # yy - @zeros(nx + 1, ny, nz + 1), # xz - @zeros(nx, ny + 1, nz + 1), # yz @zeros(nx, ny, nz), # zz + @zeros(nx + 1, ny + 1, nz), # xy + @zeros(nx, ny + 1, nz + 1), # yz + @zeros(nx + 1, ny, nz + 1), # xz @zeros(nx, ny, nz), # yz @ cell center @zeros(nx, ny, nz), # xz @ cell center @zeros(nx, ny, nz), # xy @ cell center diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index cf997d75..2f8f475e 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -1,8 +1,9 @@ struct ThermalArrays{_T} - T::_T # Temperature @ grid nodes + T::_T # Temperature @ grid nodes Tc::_T # Temperature @ cell centers - ΔT::_T Told::_T + ΔT::_T + ΔTc::_T dT_dt::_T qTx::_T qTy::_T @@ -13,57 +14,6 @@ struct ThermalArrays{_T} H::_T # source terms shear_heating::_T # shear heating terms ResT::_T - - # function ThermalArrays(nx::Integer, ny::Integer) - # T = @zeros(nx + 3, ny + 1) - # ΔT = @zeros(nx + 3, ny + 1) - # Told = @zeros(nx + 3, ny + 1) - # Tc = @zeros(nx, ny) - # H = @zeros(nx, ny) - # shear_heating = @zeros(nx, ny) - # dT_dt = @zeros(nx + 1, ny - 1) - # qTx = @zeros(nx + 2, ny - 1) - # qTy = @zeros(nx + 1, ny) - # qTx2 = @zeros(nx + 2, ny - 1) - # qTy2 = @zeros(nx + 1, ny) - # ResT = @zeros(nx + 1, ny - 1) - # return new{typeof(T)}( - # T, - # Tc, - # ΔT, - # Told, - # dT_dt, - # qTx, - # qTy, - # nothing, - # qTx2, - # qTy2, - # nothing, - # H, - # shear_heating, - # ResT, - # ) - # end - - # function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) - # T = @zeros(nx + 1, ny + 1, nz + 1) - # ΔT = @zeros(nx + 1, ny + 1, nz + 1) - # Told = @zeros(nx + 1, ny + 1, nz + 1) - # Tc = @zeros(nx, ny, nz) - # H = @zeros(nx, ny, nz) - # shear_heating = @zeros(nx, ny, nz) - # dT_dt = @zeros(ni .- 1) - # qTx = @zeros(nx, ny - 1, nz - 1) - # qTy = @zeros(nx - 1, ny, nz - 1) - # qTz = @zeros(nx - 1, ny - 1, nz) - # qTx2 = @zeros(nx, ny - 1, nz - 1) - # qTy2 = @zeros(nx - 1, ny, nz - 1) - # qTz2 = @zeros(nx - 1, ny - 1, nz) - # ResT = @zeros((ni .- 1)...) - # return new{typeof(T)}( - # T, Tc, ΔT, Told, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT - # ) - # end end ThermalArrays(::Type{CPUBackend}, ni::NTuple{N,Number}) where {N} = ThermalArrays(ni...) @@ -76,9 +26,6 @@ function ThermalArrays(::Number, ::Number, ::Number) throw(ArgumentError("ThermalArrays dimensions must be given as integers")) end -# traits -# @inline backend(::ThermalArrays{T}) where {T} = backend(T) - ## Thermal diffusion coefficients struct PTThermalCoeffs{T,M} diff --git a/src/types/stokes.jl b/src/types/stokes.jl index 02450bcf..e35743aa 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -10,22 +10,6 @@ end Velocity(Vx::T, Vy::T) where {T} = Velocity(Vx, Vy, nothing) -# function Velocity(nx::Integer, ny::Integer) -# nVx = (nx + 1, ny + 2) -# nVy = (nx + 2, ny + 1) - -# Vx, Vy = @zeros(nVx...), @zeros(nVy) -# return Velocity(Vx, Vy, nothing) -# end - -# function Velocity(nx::Integer, ny::Integer, nz::Integer) -# nVx = (nx + 1, ny + 2, nz + 2) -# nVy = (nx + 2, ny + 1, nz + 2) -# nVz = (nx + 2, ny + 2, nz + 1) - -# Vx, Vy, Vz = @zeros(nVx...), @zeros(nVy), @zeros(nVz) -# return Velocity{PTArray}(Vx, Vy, Vz) -# end Velocity(ni::NTuple{N,Number}) where {N} = Velocity(ni...) function Velocity(::Number, ::Number) @@ -46,14 +30,6 @@ struct Viscosity{T} end Viscosity(args...) = Viscosity(promote(args...)...) - -# function Viscosity(ni::NTuple{N,Integer}) where {N} -# η = @ones(ni...) -# η_vep = @ones(ni...) -# ητ = @zeros(ni...) -# return Viscosity(η, η_vep, ητ) -# end - 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} @@ -96,31 +72,6 @@ function SymmetricTensor(xx::T, yy::T, xy::T, xy_c::T, II::T) where {T} ) end -# function SymmetricTensor(nx::Integer, ny::Integer) -# return SymmetricTensor( -# @zeros(nx, ny), # xx -# @zeros(nx, ny), # yy -# @zeros(nx + 1, ny + 1), # xy -# @zeros(nx, ny), # xy @ cell center -# @zeros(nx, ny) # II (second invariant) -# ) -# end - -# function SymmetricTensor(nx::Integer, ny::Integer, nz::Integer) -# return SymmetricTensor( -# @zeros(nx, ny, nz), # xx -# @zeros(nx + 1, ny + 1, nz), # xy -# @zeros(nx, ny, nz), # yy -# @zeros(nx + 1, ny, nz + 1), # xz -# @zeros(nx, ny + 1, nz + 1), # yz -# @zeros(nx, ny, nz), # zz -# @zeros(nx, ny, nz), # yz @ cell center -# @zeros(nx, ny, nz), # xz @ cell center -# @zeros(nx, ny, nz), # xy @ cell center -# @zeros(nx, ny, nz), # II (second invariant) -# ) -# end - SymmetricTensor(ni::NTuple{N,Number}) where {N} = SymmetricTensor(ni...) function SymmetricTensor(::Number, ::Number) throw(ArgumentError("SymmetricTensor dimensions must be given as integers")) @@ -141,22 +92,6 @@ struct Residual{T} end Residual(RP::T, Rx::T, Ry::T) where {T} = Residual(RP, Rx, Ry, nothing) - -# function Residual(nx::Integer, ny::Integer) -# Rx = @zeros(nx - 1, ny) -# Ry = @zeros(nx, ny - 1) -# RP = @zeros(nx, ny) -# return Residual(RP, Rx, Ry) -# end - -# function Residual(nx::Integer, ny::Integer, nz::Integer) -# Rx = @zeros(nx - 1, ny, nz) -# Ry = @zeros(nx, ny - 1, nz) -# Rz = @zeros(nx, ny, nz - 1) -# RP = @zeros(nx, ny, nz) -# return Residual(RP, Rx, Ry, Rz) -# end - Residual(ni::NTuple{N,Number}) where {N} = Residual(ni...) function Residual(::Number, ::Number) throw(ArgumentError("Residual dimensions must be given as integers")) @@ -179,24 +114,6 @@ struct StokesArrays{A,B,C,D,T} viscosity::D τ_o::Union{B,Nothing} R::C - - # function StokesArrays(ni::NTuple{N,Integer}) where {N} - # P = @zeros(ni...) - # P0 = @zeros(ni...) - # ∇V = @zeros(ni...) - # V = Velocity(ni) - # τ = SymmetricTensor(ni) - # τ_o = SymmetricTensor(ni) - # ε = SymmetricTensor(ni) - # ε_pl = SymmetricTensor(ni) - # EII_pl = @zeros(ni...) - # viscosity = Viscosity(ni) - # R = Residual(ni) - - # return new{typeof(V),typeof(τ),typeof(R),typeof(viscosity),typeof(P)}( - # P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, viscosity, τ_o, R - # ) - # end end function StokesArrays(::Type{CPUBackend}, ni::Vararg{Integer,N}) where {N} @@ -211,9 +128,6 @@ function StokesArrays(::Number, ::Number, ::Number) throw(ArgumentError("StokesArrays dimensions must be given as integers")) end -# traits -# @inline backend(x::StokesArrays) = backend(x.P) - ## PTStokesCoeffs type struct PTStokesCoeffs{T} From 253645ec33cbc0d60ae6b2776ef654b14e4a74be Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 18:17:15 +0200 Subject: [PATCH 26/91] format & clean thermal stress miniapps --- .../Thermal_Stress_Magma_Chamber_nondim.jl | 423 ++++++++---------- .../Thermal_Stress_Magma_Chamber_nondim3D.jl | 215 +++------ 2 files changed, 239 insertions(+), 399 deletions(-) diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl index 75fb53f8..1bd4b732 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl @@ -1,8 +1,8 @@ -# using CUDA -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -using ParallelStencil -# @init_parallel_stencil(CUDA, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -10,14 +10,9 @@ using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -# model = PS_Setup(:CUDA, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -using Printf, Statistics, LinearAlgebra, GeoParams, CairoMakie, CellArrays +using Printf, Statistics, LinearAlgebra, GeoParams, GLMakie, CellArrays using StaticArrays using ImplicitGlobalGrid using MPI: MPI @@ -25,7 +20,6 @@ using WriteVTK # ----------------------------------------------------------------------------------------- ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- - function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -41,16 +35,6 @@ macro all_j(A) return esc(:($A[$idx_j])) end -# Initial pressure profile - not accurate -# @parallel function init_P!(P, ρg, z, sticky_air) -# @all(P) = (abs(@all(ρg) * (@all_j(z))) - (@all(ρg) * sticky_air)) * <((@all_j(z)), 0.0) -# return nothing -# end -# @parallel function init_P!(P, ρg, z) -# @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) -# return nothing -# end - @parallel_indices (i, j) function init_P!(P, ρg, z, dz) P[i, j] = sum(abs(ρg[i, jj] * z[jj]) for jj in j:size(P, 2)) * dz return nothing @@ -84,7 +68,7 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, r_anomaly, stic return nothing end - @parallel (JustRelax.@idx ni) init_phases!( + @parallel (@idx ni) init_phases!( phases, particles.coords..., particles.index, @@ -133,108 +117,82 @@ function circular_perturbation!(T, δT, xc_anomaly, yc_anomaly, r_anomaly, xvi, ) end - -function phase_change!(phases, particles) - ni = size(phases) - @parallel_indices (I...) function _phase_change!(phases, px, py, index) - @inbounds for ip in JustRelax.cellaxes(phases) - #quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue - - x = JustRelax.@cell px[ip, I...] - y = (JustRelax.@cell py[ip, I...]) - phase_ij = @cell phases[ip, I...] - if y > 0.0 && (phase_ij == 2.0 || phase_ij == 3.0) - @cell phases[ip, I...] = 4.0 - end - end - return nothing - end - - @parallel (JustRelax.@idx ni) _phase_change!( - phases, particles.coords..., particles.index - ) -end - function init_rheology(CharDim; is_compressible = false, steady_state=true) # plasticity setup - do_DP = true # do_DP=false: Von Mises, do_DP=true: Drucker-Prager (friction angle) - η_reg = 1.0e16Pa*s # regularisation "viscosity" for Drucker-Prager - Coh = 10.0MPa # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) - ϕ = 30.0 * do_DP # friction angle - G0 = 6.0e11Pa # elastic shear modulus - G_magma = 6.0e11Pa # elastic shear modulus perturbation - - # soft_C = LinearSoftening((ustrip(Coh)/2, ustrip(Coh)), (0e0, 1e-1)) # softening law + do_DP = true # do_DP=false: Von Mises, do_DP=true: Drucker-Prager (friction angle) + η_reg = 1.0e16Pa * s # regularisation "viscosity" for Drucker-Prager + Coh = 10.0MPa # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) + ϕ = 30.0 * do_DP # friction angle + G0 = 6.0e11Pa # elastic shear modulus + G_magma = 6.0e11Pa # elastic shear modulus perturbation + soft_C = NonLinearSoftening(; ξ₀=ustrip(Coh), Δ=ustrip(Coh) / 2) # softening law pl = DruckerPrager_regularised(; C=Coh, ϕ=ϕ, η_vp=η_reg, Ψ=0.0, softening_C = soft_C) # plasticity if is_compressible == true - el = SetConstantElasticity(; G=G0, ν=0.25) # elastic spring + el = SetConstantElasticity(; G=G0, ν=0.25) # elastic spring el_magma = SetConstantElasticity(; G=G_magma, ν=0.25)# elastic spring - β_rock = 6.0e-11 - β_magma = 6.0e-11 + β_rock = 6.0e-11 + β_magma = 6.0e-11 else - el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring + el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring el_magma = SetConstantElasticity(; G=G_magma, ν=0.5) # elastic spring - β_rock = inv(get_Kb(el)) - β_magma = inv(get_Kb(el_magma)) + β_rock = inv(get_Kb(el)) + β_magma = inv(get_Kb(el_magma)) end if steady_state == true - creep_rock = LinearViscous(; η=1e23 * Pa * s) + creep_rock = LinearViscous(; η=1e23 * Pa * s) creep_magma = LinearViscous(; η=1e18 * Pa * s) - creep_air = LinearViscous(; η=1e18 * Pa * s) + creep_air = LinearViscous(; η=1e18 * Pa * s) else - creep_rock = DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) + creep_rock = DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) creep_magma = DislocationCreep(; A=1.67e-24, n=3.5, E=1.87e5, V=6e-6, r=0.0, R=8.3145) - creep_air = LinearViscous(; η=1e18 * Pa * s) - β_rock = 6.0e-11 - β_magma = 6.0e-11 + creep_air = LinearViscous(; η=1e18 * Pa * s) + β_rock = 6.0e-11 + β_magma = 6.0e-11 end - g=9.81m/s^2 + g = 9.81m/s^2 rheology = ( #Name="UpperCrust" SetMaterialParams(; - Phase=1, - 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, el, pl)), - # CompositeRheology=CompositeRheology((creep_rock,)), - Melting=MeltingParam_Caricchi(), - Gravity = ConstantGravity(; g=g), - Elasticity=el, - CharDim=CharDim, + Phase = 1, + 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, el, pl)), + 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), - 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, )), - CompositeRheology=CompositeRheology((creep_magma, el_magma)), - Melting=MeltingParam_Caricchi(), - Gravity = ConstantGravity(; g=g), - Elasticity=el_magma, - CharDim=CharDim, + Phase = 2, + Density = PT_Density(; ρ0=2650kg / m^3, T0=0.0C, β=β_magma / Pa), + HeatCapacity = ConstantHeatCapacity(; Cp=1050J / kg / K), + Conductivity = ConstantConductivity(; k=1.5Watt / K / m), + LatentHeat = ConstantLatentHeat(; Q_L=350e3J / kg), + ShearHeat = ConstantShearheating(0.0NoUnits), + CompositeRheology = CompositeRheology((creep_magma, el_magma)), + Melting = MeltingParam_Caricchi(), + 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, + 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, ), ) @@ -243,52 +201,52 @@ end function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) - CharDim = GEO_units(;length=12.5km, viscosity=1e21, temperature = 1e3C) + # Characteristic lengths + CharDim = GEO_units(;length=12.5km, viscosity=1e21, temperature = 1e3C) + #-------JustRelax parameters------------------------------------------------------------- # Domain setup for JustRelax - sticky_air = nondimensionalize(1.5km, CharDim) # thickness of the sticky air layer - ly = nondimensionalize(12.5km,CharDim) + sticky_air # domain length in y-direction - lx = nondimensionalize(15.5km, CharDim) # domain length in x-direction - li = lx, ly # domain length in x- and y-direction - ni = nx, ny # number of grid points in x- and y-direction - di = @. li / ni # grid step in x- and y-direction - origin = nondimensionalize(0.0km,CharDim), -ly # 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 + sticky_air = nondimensionalize(1.5km, CharDim) # thickness of the sticky air layer + ly = nondimensionalize(12.5km,CharDim) + sticky_air # domain length in y-direction + lx = nondimensionalize(15.5km, CharDim) # domain length in x-direction + li = lx, ly # domain length in x- and y-direction + ni = nx, ny # number of grid points in x- and y-direction + di = @. li / ni # grid step in x- and y-direction + origin = nondimensionalize(0.0km,CharDim), -ly # origin coordinates of the domain + grid = Geometry(ni, li; origin=origin) + εbg = nondimensionalize(0.0 / s,CharDim) # background strain rate + (; xci, xvi) = grid # nodes at the center and vertices of the cells #--------------------------------------------------------------------------------------- # Physical Parameters - rheology = init_rheology(CharDim; is_compressible=true, steady_state=false) - cutoff_visc = nondimensionalize((1e16Pa*s, 1e24Pa*s),CharDim) - κ = (4 / (rheology[1].HeatCapacity[1].Cp * rheology[1].Density[1].ρ0)) + rheology = init_rheology(CharDim; is_compressible=true, steady_state=false) + cutoff_visc = nondimensionalize((1e16Pa*s, 1e24Pa*s),CharDim) + κ = (4 / (rheology[1].HeatCapacity[1].Cp * rheology[1].Density[1].ρ0)) dt = dt_diff = (0.5 * min(di...)^2 / κ / 2.01) # diffusive CFL timestep limiter # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 20, 40, 15 - particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni...) - + particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni...) subgrid_arrays = SubgridDiffusionCellArrays(particles) - # velocity grids grid_vx, grid_vy = velocity_grids(xci, xvi, di) # temperature - pT, pPhases = init_cell_arrays(particles, Val(3)) - particle_args = (pT, pPhases) - - # Circular temperature anomaly-------------------------- - x_anomaly = lx * 0.5 - y_anomaly = nondimensionalize(-5km,CharDim) # origin of the small thermal anomaly - r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation - anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) + pT, pPhases = init_cell_arrays(particles, Val(2)) + particle_args = (pT, pPhases) + + # Circular temperature anomaly ----------------------- + x_anomaly = lx * 0.5 + y_anomaly = nondimensionalize(-5km,CharDim) # origin of the small thermal anomaly + r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation + anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) init_phases!(pPhases, particles, x_anomaly, y_anomaly, r_anomaly, sticky_air, nondimensionalize(0.0km,CharDim), nondimensionalize(20km,CharDim)) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # Initialisation of thermal profile - thermal = ThermalArrays(ni) # initialise thermal arrays and boundary conditions - thermal_bc = TemperatureBoundaryConditions(; - no_flux=(left=true, right=true, top=false, bot=false), + thermal = ThermalArrays(backend_JR, ni) # initialise thermal arrays and boundary conditions + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left=true, right=true, top=false, bot=false), ) @parallel (@idx ni .+ 1) init_T!( thermal.T, xvi[2], @@ -302,54 +260,48 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) thermal.T, anomaly, x_anomaly, y_anomaly, r_anomaly, xvi, sticky_air ) thermal_bcs!(thermal.T, thermal_bc) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!( - thermal.Tc, thermal.T - ) + temperature2center!(thermal) + # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) # initialise stokes arrays with the defined regime - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL=0.9 / √2.1) + stokes = StokesArrays(backend_JR, ni) # initialise stokes arrays with the defined regime + pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-4, CFL = 1 / √2.1) # ---------------------------------------------------- args = (; T=thermal.Tc, P=stokes.P, dt=dt) - pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.8 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.8 / √2.1 ) - # Boundary conditions of the flow - stokes.V.Vx .= PTArray([ + + # Pure shear far-field boundary conditions + stokes.V.Vx .= PTArray(backend_JR)([ εbg * (x - lx * 0.5) / (lx / 2) / 2 for x in xvi[1], _ in 1:(ny + 2) ]) - stokes.V.Vy .= PTArray([ + stokes.V.Vy .= PTArray(backend_JR)([ (abs(y) - sticky_air) * εbg * (abs(y) > sticky_air) for _ in 1:(nx + 2), y in xvi[2] ]) flow_bcs = FlowBoundaryConditions(; - free_slip=(left=true, right=true, top=true, bot=true), - free_surface =true, + free_slip = (left=true, right=true, top=true, bot=true), + free_surface = true, ) flow_bcs!(stokes, flow_bcs) update_halo!(stokes.V.Vx, stokes.V.Vy) - η = @ones(ni...) # initialise viscosity - ϕ = similar(η) # melt fraction center + η = @ones(ni...) # initialise viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, cutoff_visc - ) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) η_vep = copy(η) # Buoyancy force - ρg = @zeros(ni...), @zeros(ni...) # ρg[1] is the buoyancy force in the x direction, ρg[2] is the buoyancy force in the y direction - for _ in 1:10 - @parallel (JustRelax.@idx ni) compute_ρg!( - ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P) - ) + ρg = @zeros(ni...), @zeros(ni...) # ρg[1] is the buoyancy force in the x direction, ρg[2] is the buoyancy force in the y direction + for _ in 1:5 + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2], di[2]) end # Arguments for functions - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) + args = (; T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) @copy thermal.Told thermal.T # IO ------------------------------------------------ @@ -386,7 +338,6 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) # Time loop t, it = 0.0, 0 - local iters local Vx_v, Vy_v if do_vtk Vx_v = @zeros(ni .+ 1...) @@ -402,20 +353,17 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) grid2particle!(pT, xvi, T_buffer, particles) @copy stokes.P0 stokes.P thermal.Told .= thermal.T - P_init = deepcopy(stokes.P) - Tsurf = thermal.T[1, end] - Tbot = thermal.T[1, 1] + P_init = deepcopy(stokes.P) + Tsurf = thermal.T[1, end] + Tbot = thermal.T[1, 1] - while it < 25 #nt + while it < 25 # Update buoyancy and viscosity - - args = (; ϕ=ϕ, T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, cutoff_visc - ) - @parallel (@idx ni) compute_ρg!( - ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P) - ) + args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + # Stokes solver ----------------- solve!( stokes, @@ -423,21 +371,20 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax=100e3, - free_surface=true, - nout=5e3, - viscosity_cutoff=cutoff_visc, - ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!( - stokes.ε.II, @strain(stokes)... - ) + kwargs = (; + iterMax = 100e3, + free_surface = true, + nout = 5e3, + viscosity_cutoff = cutoff_visc, + ) + ) + tensor_invariant!(stokes.ε) + @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @parallel (@idx ni) multi_copy!( @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) @@ -445,15 +392,14 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) dt = compute_dt(stokes, di, dt_diff, igg) # -------------------------------- - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, - rheology, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, + rheology, # needs to be a tuple dt, ) + # Thermal solver --------------- heatdiffusion_PT!( thermal, @@ -463,13 +409,15 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) args, dt, di; - igg=igg, - phase=phase_ratios, - iterMax=150e3, - nout=1e3, - verbose=true, + kwargs =(; + igg = igg, + phase = phase_ratios, + iterMax = 150e3, + nout = 1e3, + verbose = true, + ) + ) - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) end @@ -484,40 +432,39 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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) + 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) + phase_ratios_center(phase_ratios, particles, grid, pPhases) particle2grid!(T_buffer, pT, xvi, particles) - @views T_buffer[:, end] .= Tsurf - @views T_buffer[:, 1] .= Tbot + @views T_buffer[:, end] .= Tsurf + @views T_buffer[:, 1] .= Tbot @views thermal.T[2:end - 1, :] .= T_buffer thermal_bcs!(thermal.T, thermal_bc) temperature2center!(thermal) thermal.ΔT .= thermal.T .- thermal.Told - @parallel (@idx size(thermal.ΔTc)...) temperature2center!(thermal.ΔTc, thermal.ΔT) + vertex2center!(thermal.ΔTc, thermal.ΔT) @show it += 1 t += dt # # # Plotting ------------------------------------------------------- if it == 1 || rem(it, 1) == 0 - checkpointing(figdir, stokes, thermal.T, η, t) + # checkpointing(figdir, stokes, thermal.T, η, t) if igg.me == 0 if do_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; - T=Array(ustrip.(dimensionalize(thermal.T[2:(end - 1), :], C, CharDim))), - τxy= Array(ustrip.(dimensionalize(stokes.τ.xy, s^-1, CharDim))), - εxy= Array(ustrip.(dimensionalize(stokes.ε.xy, s^-1, CharDim))), - Vx = Array(ustrip.(dimensionalize(Vx_v,cm/yr,CharDim))), - Vy = Array(ustrip.(dimensionalize(Vy_v, cm/yr, CharDim))), + T = Array(ustrip.(dimensionalize(thermal.T[2:(end - 1), :], C, CharDim))), + τxy = Array(ustrip.(dimensionalize(stokes.τ.xy, s^-1, CharDim))), + εxy = Array(ustrip.(dimensionalize(stokes.ε.xy, s^-1, CharDim))), + Vx = Array(ustrip.(dimensionalize(Vx_v,cm/yr,CharDim))), + Vy = Array(ustrip.(dimensionalize(Vy_v, cm/yr, CharDim))), ) data_c = (; P = Array(ustrip.(dimensionalize(stokes.P,MPa,CharDim))), @@ -538,18 +485,11 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) ) end - # Make particles plottable - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:] ./ 1e3 - pyv = ppy.data[:] ./ 1e3 - clr = pPhases.data[:] - idxv = particles.index.data[:] - t_dim = (dimensionalize(t, yr, CharDim).val / 1e3) + t_Kyrs = t_dim * 1e3 # Make Makie figure fig = Figure(; size=(2000, 1800), createmissing=true) - ar = li[1] / li[2] + ar = li[1] / li[2] ax0 = Axis( fig[1, 1:2]; @@ -558,24 +498,24 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) titlesize=50, height=0.0, ) - ax0.ylabelvisible = false - ax0.xlabelvisible = false - ax0.xgridvisible = false - ax0.ygridvisible = false - ax0.xticksvisible = false - ax0.yticksvisible = false + ax0.ylabelvisible = false + ax0.xlabelvisible = false + ax0.xgridvisible = false + ax0.ygridvisible = false + ax0.xticksvisible = false + ax0.yticksvisible = false ax0.yminorticksvisible = false ax0.xminorticksvisible = false - ax0.xgridcolor = :white - ax0.ygridcolor = :white - ax0.ytickcolor = :white - ax0.xtickcolor = :white - ax0.yticklabelcolor = :white - ax0.xticklabelcolor = :white - ax0.yticklabelsize = 0 - ax0.xticklabelsize = 0 - ax0.xlabelcolor = :white - ax0.ylabelcolor = :white + ax0.xgridcolor = :white + ax0.ygridcolor = :white + ax0.ytickcolor = :white + ax0.xtickcolor = :white + ax0.yticklabelcolor = :white + ax0.xticklabelcolor = :white + ax0.yticklabelsize = 0 + ax0.xticklabelsize = 0 + ax0.xlabelcolor = :white + ax0.ylabelcolor = :white ax1 = Axis( fig[2, 1][1, 1]; @@ -744,16 +684,17 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) end end end - # finalize_global_grid() - + + finalize_global_grid() + + return nothing end figdir = "Thermal_stresses_around_cooling_magma" do_vtk = true # set to true to generate VTK files for ParaView -ar = 1 # aspect ratio -n = 128 -nx = n * ar - 2 -ny = n - 2 +n = 64 +nx = n * ar - 2 +ny = n - 2 igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid IGG(init_global_grid(nx, ny, 1; init_MPI=true)...) else @@ -761,18 +702,4 @@ else end # run main script -main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=do_vtk); - -function plot_particles(particles, pPhases) - p = particles.coords - ppx, ppy = p - pxv = ppx.data[:] - pyv = ppy.data[:] - clr = pPhases.data[:] - idxv = particles.index.data[:] - f, ax, h = scatter( - Array(pxv[idxv]), Array(pyv[idxv]); color=Array(clr[idxv]), colormap=:roma - ) - Colorbar(f[1, 2], h) - return f -end +main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=do_vtk); \ No newline at end of file diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl index 948c60f1..4bcb849f 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl @@ -1,22 +1,16 @@ -using CUDA -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell -using ParallelStencil -@init_parallel_stencil(CUDA, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -# @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences3D +@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) 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 = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:CUDA, Float64, 3) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -# model = PS_Setup(:cpu, Float64, 3) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using Printf, Statistics, LinearAlgebra, GeoParams, GLMakie, CellArrays using StaticArrays @@ -32,27 +26,12 @@ macro all_j(A) return esc(:($A[$idx_j])) end -# Initial pressure profile - not accurate -# @parallel function init_P!(P, ρg, z, sticky_air) -# @all(P) = (abs(@all(ρg) * (@all_j(z))) - (@all(ρg) * sticky_air)) * <((@all_j(z)), 0.0) -# return nothing -# end -# @parallel function init_P!(P, ρg, z) -# @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) -# return nothing -# end - @parallel function init_P!(P, ρg, z, sticky_air) - @all(P) = abs(@all(ρg) * (@all_j(z) + sticky_air)) * <((@all_j(z) + sticky_air), 0.0) + @all(P) = @all(ρg) + # @all(P) = abs(@all(ρg) * (@all_j(z) + sticky_air)) * <((@all_j(z) + sticky_air), 0.0) return nothing end -# @parallel_indices (i, j) function init_P!(P, ρg, z, dz) -# P[i, j] = sum(abs(ρg[i, jj] * z[jj]) for jj in j:size(P, 2)) * dz -# return nothing -# end - - function init_phases!(phases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, sticky_air,top, bottom) ni = size(phases) @@ -76,7 +55,7 @@ function init_phases!(phases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_a end if z < top - @cell phases[ip, I...] = 3.0 + JustRelax.@cell phases[ip, I...] = 3.0 end end return nothing @@ -180,7 +159,6 @@ function init_rheology(CharDim; is_compressible = false, steady_state=true) LatentHeat = ConstantLatentHeat(; Q_L=350e3J / kg), ShearHeat = ConstantShearheating(1.0NoUnits), CompositeRheology = CompositeRheology((creep_rock, el, pl)), - # CompositeRheology=CompositeRheology((creep_rock,)), Melting = MeltingParam_Caricchi(), Gravity = ConstantGravity(; g=g), Elasticity = el, @@ -195,7 +173,6 @@ function init_rheology(CharDim; is_compressible = false, steady_state=true) Conductivity = ConstantConductivity(; k=1.5Watt / K / m), LatentHeat = ConstantLatentHeat(; Q_L=350e3J / kg), ShearHeat = ConstantShearheating(0.0NoUnits), - # CompositeRheology=CompositeRheology((creep_magma, )), CompositeRheology = CompositeRheology((creep_magma, el_magma)), Melting = MeltingParam_Caricchi(), Gravity = ConstantGravity(; g=g), @@ -205,13 +182,13 @@ function init_rheology(CharDim; is_compressible = false, steady_state=true) #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,)), + 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, ), @@ -229,14 +206,14 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # Domain setup for JustRelax sticky_air = nondimensionalize(1.5km, CharDim) # thickness of the sticky air layer lz = nondimensionalize(12.5km,CharDim) + sticky_air # domain length in y-direction - lx = ly = nondimensionalize(15.5km, CharDim) # domain length in x-direction + lx = ly = nondimensionalize(15.5km, CharDim) # 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 + εbg = nondimensionalize(0.0 / s, CharDim) # background strain rate #--------------------------------------------------------------------------------------- # Physical Parameters @@ -263,12 +240,12 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals z_anomaly = nondimensionalize(-5km,CharDim) # origin of the small thermal anomaly r_anomaly = nondimensionalize(1.5km, CharDim) # radius of perturbation anomaly = nondimensionalize((750 + 273)K, CharDim) # thermal perturbation (in K) - phase_ratios = PhaseRatio(ni, length(rheology)) + 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), nondimensionalize(20km,CharDim)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # Initialisation of thermal profile - thermal = ThermalArrays(ni) # initialise thermal arrays and boundary conditions + 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), ) @@ -285,30 +262,18 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals thermal.T, anomaly, x_anomaly, y_anomaly, z_anomaly, r_anomaly, xvi, sticky_air ) thermal_bcs!(thermal.T, thermal_bc) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!( - thermal.Tc, thermal.T - ) + temperature2center!(thermal) + # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) # initialise stokes arrays with the defined regime + stokes = StokesArrays(backend_JR, ni) # initialise stokes arrays with the defined regime pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL=0.9 / √3.1) # ---------------------------------------------------- - args = (; T=thermal.Tc, P=stokes.P, dt=dt) - + args = (; T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.8 / √3.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.8 / √3.1 ) - # Boundary conditions of the flow - # stokes.V.Vx .= PTArray([ - # εbg * (x - lx * 0.5) for x in xvi[1], _ in 1:(ny + 2), _ in 1:(nz + 2) - # ]) - # stokes.V.Vy .= PTArray([ - # εbg * (y - ly * 0.5) for _ in 1:(nx + 2), y in xvi[2], _ in 1:(nz + 2) - # ]) - # stokes.V.Vz .= PTArray([ - # (abs(z) + sticky_air) * εbg * ((z + sticky_air) > 0) for _ in 1:(nx + 2), _ in 1:(ny + 2), z in 1:(nz + 2) - # ]) flow_bcs = FlowBoundaryConditions(; free_slip = (left=true, right=true, front=true, back=true, top=true, bot=true), @@ -318,26 +283,16 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals flow_bcs!(stokes, flow_bcs) update_halo!(@velocity(stokes)...) - η = @ones(ni...) # initialise viscosity - η_vep = similar(η) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, cutoff_visc - ) - - # Buoyancy force + # Buoyancy force & viscosity ρ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 - for _ in 1:2 - @parallel (JustRelax.@idx ni) compute_ρg!( - ρg[3], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P) - ) + for _ in 1:5 + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[3], xci[3], sticky_air) - # @parallel init_P!(stokes.P, ρg[2], xci[2], di[2]) - # @parallel init_P!(stokes.P, ρg[2], xci[2]) end + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) # Arguments for functions - args = (; T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) @copy thermal.Told thermal.T # IO ------------------------------------------------ @@ -380,7 +335,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals Vz_v = @zeros(ni .+ 1...) end - dt₀ = similar(stokes.P) + dt₀ = similar(stokes.P) grid2particle!(pT, xvi, thermal.T, particles) @copy stokes.P0 stokes.P @@ -388,16 +343,13 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals Tsurf = thermal.T[1, 1, end] Tbot = thermal.T[1, 1, 1] - while it < 25 #nt + while it < 25 # Update buoyancy and viscosity - - args = (; T=thermal.Tc, P=stokes.P, dt=dt, ΔTc=thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, cutoff_visc - ) - @parallel (@idx ni) compute_ρg!( - ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P) - ) + args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + # Stokes solver ----------------- solve!( stokes, @@ -405,20 +357,19 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 100e3, - nout = 5e3, - viscosity_cutoff = cutoff_visc, - ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!( - stokes.ε.II, @strain(stokes)... + kwargs = (; + iterMax = 150e3, + free_surface = false, + nout = 5e3, + viscosity_cutoff = cutoff_visc, + ) ) + tensor_invariant!(stokes.ε) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @parallel (@idx ni) multi_copy!( @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) @@ -426,15 +377,14 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals dt = compute_dt(stokes, di, dt_diff, igg) # -------------------------------- - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, - rheology, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, + rheology, # needs to be a tuple dt, ) + # Thermal solver --------------- heatdiffusion_PT!( thermal, @@ -444,12 +394,15 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 150e3, - nout = 1e3, - verbose = true, + 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 ) @@ -461,14 +414,13 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vxi..., dt, 2 / 3) + 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 = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), xvi) + 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) + phase_ratios_center(phase_ratios, particles, grid, pPhases) particle2grid!(thermal.T, pT, xvi, particles) @views thermal.T[:, :, end] .= Tsurf @@ -476,18 +428,19 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals thermal_bcs!(thermal.T, thermal_bc) temperature2center!(thermal) thermal.ΔT .= thermal.T .- thermal.Told - @parallel (@idx size(thermal.ΔTc)...) temperature2center!(thermal.ΔTc, thermal.ΔT) + vertex2center!(thermal.ΔTc, thermal.ΔT) @show it += 1 t += dt # # # Plotting ------------------------------------------------------- if it == 1 || rem(it, 1) == 0 - checkpointing(figdir, stokes, thermal.T, η, t) + (; η) = stokes.viscosity + # checkpointing(figdir, stokes, thermal.T, η, t) if igg.me == 0 if do_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + 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))), @@ -543,7 +496,6 @@ end figdir = "Thermal_stresses_around_cooling_magma_3D" do_vtk = true # set to true to generate VTK files for ParaView -ar = 1 # aspect ratio n = 64 nx = n ny = n @@ -556,42 +508,3 @@ end # run main script main3D(igg; figdir=figdir, nx=nx, ny=ny, nz=nz, do_vtk = do_vtk); - - -# 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(ρg[2][:]), - # Array(ustrip.(dimensionalize(ρg[2][:], kg/m^3 * m / s^2, CharDim)))./9.81, - Array(ustrip.(dimensionalize(stokes.P[:], MPa, CharDim))), - ustrip.(dimensionalize(Z, km, CharDim)), - ) - hideydecorations!(ax2) - save(joinpath(figdir, "initial_profile.png"), fig) - fig -end - -# function plot_particles(particles, pPhases) -# p = particles.coords -# ppx, ppy = p -# pxv = ppx.data[:] -# pyv = ppy.data[:] -# clr = pPhases.data[:] -# idxv = particles.index.data[:] -# f, ax, h = scatter( -# Array(pxv[idxv]), Array(pyv[idxv]); color=Array(clr[idxv]), colormap=:roma -# ) -# Colorbar(f[1, 2], h) -# return f -# end From 67eac4ab5c5f2f6867aef974839234c985be84b5 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 22:53:57 +0200 Subject: [PATCH 27/91] adapt a bunch of tests --- .../elastic_buildup/Elastic_BuildUp.jl | 15 +-- miniapps/benchmarks/stokes2D/solcx/SolCx.jl | 38 +++--- miniapps/benchmarks/stokes2D/solkz/SolKz.jl | 24 ++-- test/test_arrays_conversions.jl | 53 +++++---- test/test_shearband2D.jl | 41 +++---- test/test_shearband2D_softening.jl | 40 +++---- test/test_shearheating2D.jl | 108 ++++++++---------- test/test_sinking_block.jl | 56 ++++----- test/test_stokes_elastic_buildup.jl | 5 +- test/test_stokes_solcx.jl | 5 +- test/test_stokes_solkz.jl | 5 +- 11 files changed, 185 insertions(+), 205 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl index 123c69ce..162673d5 100644 --- a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl +++ b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl @@ -42,18 +42,19 @@ function elastic_buildup(; kyr = 1e3 * yr ttot = endtime * kyr # total simulation time - ## Setup-specific parameters and fields - η = fill(η0, nx, ny) - g = 0.0 # gravity - Gc = @fill(G, ni...) - Kb = @fill(Inf, ni...) - ## Allocate arrays needed for every Stokes problem # general stokes arrays - stokes = StokesArrays(ni) + stokes = StokesArrays(backend, ni) # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL=1 / √2.1) + ## Setup-specific parameters and fields + (; η) = stokes.viscosity + η .= fill(η0, nx, ny) + g = 0.0 # gravity + Gc = @fill(G, ni...) + Kb = @fill(Inf, ni...) + ## Boundary conditions pureshear_bc!(stokes, xci, xvi, εbg) flow_bcs = FlowBoundaryConditions(; diff --git a/miniapps/benchmarks/stokes2D/solcx/SolCx.jl b/miniapps/benchmarks/stokes2D/solcx/SolCx.jl index 016516e2..970efda3 100644 --- a/miniapps/benchmarks/stokes2D/solcx/SolCx.jl +++ b/miniapps/benchmarks/stokes2D/solcx/SolCx.jl @@ -13,7 +13,7 @@ end function solCx_viscosity(xci, ni, di; Δη=1e6) xc, yc = xci # make grid array (will be eaten by GC) - x = PTArray([xci for xci in xc, _ in yc]) + x = PTArray(backend)([xci for xci in xc, _ in yc]) η = @zeros(ni...) _viscosity(x, Δη) = ifelse(x ≤ 0.5, 1e0, Δη) @@ -33,9 +33,9 @@ end function solCx_density(xci, ni, di) xc, yc = xci # make grid array (will be eaten by GC) - x = PTArray([xci for xci in xc, _ in yc]) - y = PTArray([yci for _ in xc, yci in yc]) - ρ = PTArray(zeros(ni)) + x = PTArray(backend)([xci for xci in xc, _ in yc]) + y = PTArray(backend)([yci for _ in xc, yci in yc]) + ρ = PTArray(backend)(zeros(ni)) _density(x, y) = -sin(π * y) * cos(π * x) @@ -80,18 +80,19 @@ function solCx( ## Allocate arrays needed for every Stokes problem # general stokes arrays - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; CFL = 1 / √2.1, ϵ = 1e-8) ## Setup-specific parameters and fields - η = solCx_viscosity(xci, ni, di; Δη = Δη) # viscosity field - ρ = solCx_density(xci, ni, di) - fy = ρ .* g - ρg = @zeros(ni...), fy - dt = 1 - G = @fill(Inf, ni...) - K = @fill(Inf, ni...) + (; η) = stokes.viscosity + η .= solCx_viscosity(xci, ni, di; Δη = Δη) # viscosity field + ρ = solCx_density(xci, ni, di) + fy = ρ .* g + ρg = @zeros(ni...), fy + dt = 1 + G = @fill(Inf, ni...) + K = @fill(Inf, ni...) # smooth viscosity jump (otherwise no convergence for Δη > ~15) η2 = deepcopy(η) @@ -100,7 +101,7 @@ function solCx( @parallel smooth!(η2, η, 1.0) update_halo!(η2, η) end - @parallel (1:size(η2, 1)) free_slip_y!(η2) + @parallel (1:size(η2, 1)) JustRelax.JustRelax2D.free_slip_y!(η2) η, η2 = η2, η end @@ -121,15 +122,16 @@ function solCx( di, flow_bcs, ρg, - η, G, K, 0.1, igg; - iterMax = 500e3, - nout = 5e3, - b_width = (4, 4, 1), - verbose = false, + kwargs = ( + iterMax = 500e3, + nout = 5e3, + b_width = (4, 4, 1), + verbose = true, + ) ) t += Δt end diff --git a/miniapps/benchmarks/stokes2D/solkz/SolKz.jl b/miniapps/benchmarks/stokes2D/solkz/SolKz.jl index 7629a973..9047d85a 100644 --- a/miniapps/benchmarks/stokes2D/solkz/SolKz.jl +++ b/miniapps/benchmarks/stokes2D/solkz/SolKz.jl @@ -7,7 +7,7 @@ function solKz_viscosity(xci, ni, di; B=log(1e6)) xc, yc = xci # make grid array (will be eaten by GC) y = @zeros(ni...) - y = PTArray([yci for _ in xc, yci in yc]) + y = PTArray(backend)([yci for _ in xc, yci in yc]) η = @zeros(ni...) _viscosity(y, B) = exp(B * y) @@ -27,8 +27,8 @@ end function solKz_density(xci, ni, di) xc, yc = xci # make grid array (will be eaten by GC) - x = PTArray([xci for xci in xc, _ in yc]) - y = PTArray([yci for _ in xc, yci in yc]) + x = PTArray(backend)([xci for xci in xc, _ in yc]) + y = PTArray(backend)([yci for _ in xc, yci in yc]) ρ = @zeros(ni...) _density(x, y) = -sin(2 * y) * cos(3 * π * x) @@ -68,14 +68,15 @@ function solKz(; ## Allocate arrays needed for every Stokes problem # general stokes arrays - stokes = StokesArrays(ni, ViscoElastic) + (; η) = stokes.viscosity + stokes = StokesArrays(backend, ni, ViscoElastic) # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; Re = 5π, CFL = 1 / √2.1) ## Setup-specific parameters and fields - η = solKz_viscosity(xci, ni, di; B = log(Δη)) # viscosity field + η .= solKz_viscosity(xci, ni, di; B = log(Δη)) # viscosity field ρ = solKz_density(xci, ni, di) - fy = ρ * g + fy = ρ .* g ρg = @zeros(ni...), fy dt = 0.1 G = @fill(Inf, ni...) @@ -98,15 +99,16 @@ function solKz(; di, flow_bcs, ρg, - η, G, Kb, dt, igg; - iterMax = 150e3, - nout = 1e3, - b_width = (4, 4, 1), - verbose = false, + kwargs = ( + iterMax = 150e3, + nout = 1e3, + b_width = (4, 4, 1), + verbose = false + ), ) t += Δt end diff --git a/test/test_arrays_conversions.jl b/test/test_arrays_conversions.jl index 88d42fa3..1973517a 100644 --- a/test/test_arrays_conversions.jl +++ b/test/test_arrays_conversions.jl @@ -1,24 +1,35 @@ -using JustRelax, Test -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using JustRelax, JustRelax.JustRelax2D, Test +const bk = JustRelax.backend -@testset "Array conversions" begin - ni = 10, 10 - stokes = StokesArrays(ni, ViscoElastic) - thermal = ThermalArrays(ni) +# @testset "Array conversions" begin + ni = 2, 2 + stokes = StokesArrays(CPUBackend, ni) + thermal = ThermalArrays(CPUBackend, ni) - @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} - @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} - @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} - @test Array(stokes.P) isa Array{T, N} where {T, N} - @test Array(stokes) isa StokesArrays - @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} + # @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} + # @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} + # @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} + # @test Array(stokes.P) isa Array{T, N} where {T, N} + # @test Array(stokes) isa StokesArrays + # @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} - @test JustRelax.iscpu(stokes.V) isa JustRelax.CPUDeviceTrait - @test JustRelax.iscpu(stokes.τ) isa JustRelax.CPUDeviceTrait - @test JustRelax.iscpu(stokes.R) isa JustRelax.CPUDeviceTrait - @test JustRelax.iscpu(stokes.P) isa JustRelax.CPUDeviceTrait - @test JustRelax.iscpu(stokes) isa JustRelax.CPUDeviceTrait - @test JustRelax.iscpu(thermal) isa JustRelax.CPUDeviceTrait - @test_throws ArgumentError("Unknown device") JustRelax.iscpu("potato") -end + + # test generic arrays + @test bk(Array) === CPUBackendTrait() + @test bk(Matrix) === CPUBackendTrait() + @test bk(Vector) === CPUBackendTrait() + @test bk(rand(2)) === CPUBackendTrait() + @test bk(rand(2,2)) === CPUBackendTrait() + @test bk(rand(2,2,2)) === CPUBackendTrait() + @test_throws ArgumentError backend(rand()) + + # test JR structs + @test bk(stokes.V) === CPUBackendTrait() + @test bk(stokes.τ) === CPUBackendTrait() + @test bk(stokes.R) === CPUBackendTrait() + @test bk(stokes.P) === CPUBackendTrait() + @test bk(stokes) === CPUBackendTrait() + @test bk(thermal) === CPUBackendTrait() + @test_throws ArgumentError bk("potato") +# end + diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index e07b03d1..ed11cd21 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -2,13 +2,11 @@ push!(LOAD_PATH, "..") using Test, Suppressor using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +const backend = CPUBackend # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -31,16 +29,16 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) end # MAIN SCRIPT -------------------------------------------------------------------- function ShearBand2D() - n = 32 - nx = n - ny = n + n = 32 + nx = n + ny = n init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ------------------------------------ ly = 1e0 # domain length in y @@ -98,20 +96,17 @@ function ShearBand2D() # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right = true, top = true, bot = true), @@ -139,19 +134,19 @@ function ShearBand2D() di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose = false, - iterMax = 50e3, - nout = 1e2, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + verbose = false, + iterMax = 50e3, + nout = 1e2, + viscosity_cutoff = (-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl index 224d1f08..ac4c919a 100644 --- a/test/test_shearband2D_softening.jl +++ b/test/test_shearband2D_softening.jl @@ -2,13 +2,12 @@ push!(LOAD_PATH, "..") using Test, Suppressor using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +const backend = CPUBackend + # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -31,16 +30,16 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) end # MAIN SCRIPT -------------------------------------------------------------------- function ShearBand2D() - n = 32 - nx = n - ny = n + n = 32 + nx = n + ny = n init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ------------------------------------ ly = 1e0 # domain length in y @@ -63,7 +62,7 @@ function ShearBand2D() εbg = 1.0 # background strain-rate η_reg = 8e-3 # regularisation "viscosity" dt = η0/G0/4.0 # assumes Maxwell time of 4 - dt /= 5 + dt /= 5 el_bg = ConstantElasticity(; G=G0, Kb=4) el_inc = ConstantElasticity(; G=Gi, Kb=4) visc = LinearViscous(; η=η0) @@ -102,7 +101,7 @@ function ShearBand2D() # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces @@ -112,8 +111,8 @@ function ShearBand2D() # Rheology η = @ones(ni...) η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions @@ -143,19 +142,19 @@ function ShearBand2D() di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose = false, - iterMax = 50e3, - nout = 1e2, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + verbose = false, + iterMax = 50e3, + nout = 1e2, + viscosity_cutoff = (-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @@ -177,7 +176,6 @@ function ShearBand2D() return iters, τII, sol - end @testset "NonLinearSoftening_ShearBand2D" begin diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 0a0c95d0..04a199e6 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -1,11 +1,9 @@ -push!(LOAD_PATH, "..") -using Test, Suppressor - # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -13,14 +11,11 @@ using JustPIC._2D # Threads is the default backend, # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays +using Printf, LinearAlgebra, GeoParams # Load file with all the rheology configurations include("../miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl") @@ -49,17 +44,15 @@ end @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) return nothing end - - ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function Shearheating2D() - n = 32 - nx = n - ny = n + n = 32 + nx = n + ny = n init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ------------------------------------ ly = 40e3 # domain length in y @@ -79,7 +72,7 @@ function Shearheating2D() # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 40, 1 + nxcell, max_xcell, min_xcell = 20, 32, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -90,18 +83,17 @@ function Shearheating2D() particle_args = (pT, pPhases) # Elliptical temperature anomaly - xc_anomaly = lx/2 # origin of thermal anomaly - yc_anomaly = 40e3 # origin of thermal anomaly - # yc_anomaly = 39e3 # origin of thermal anomaly + xc_anomaly = lx / 2 # origin of thermal anomaly + yc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - init_phases!(pPhases, particles, lx/2, yc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) # ---------------------------------------------------- @@ -114,27 +106,21 @@ function Shearheating2D() # Initialize constant temperature @views thermal.T .= 273.0 + 400 thermal_bcs!(thermal.T, thermal_bc) - - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) # Rheology - η = @ones(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - η_vep = copy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1 ) # Boundary conditions @@ -143,8 +129,8 @@ function Shearheating2D() ) ## Compression and not extension - fix this εbg = 5e-14 - stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) + stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(stokes.V.Vx, stokes.V.Vy) @@ -160,33 +146,32 @@ function Shearheating2D() local iters, thermal while it < 10 # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt=dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- - iters = solve!( + solve!( stokes, pt_stokes, di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 75e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf), - verbose=false, + kwargs = (; + iterMax = 75e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf) + ) ) - @parallel (@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -215,17 +200,19 @@ function Shearheating2D() args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = false, + kwargs = (; + 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, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # advect particles in memory move_particles!(particles, xvi, particle_args) # interpolate fields from grid vertices to particles @@ -234,11 +221,10 @@ function Shearheating2D() end grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - + phase_ratios_center(phase_ratios, particles, grid, pPhases) + @show it += 1 t += dt @@ -247,8 +233,6 @@ function Shearheating2D() finalize_global_grid(; finalize_MPI = true) return iters, thermal - - end @testset "Shearheating2D" begin diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index dce3131f..a5d11804 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -1,20 +1,14 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, 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_JR = CPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using Printf, LinearAlgebra, GeoParams, CellArrays @@ -58,7 +52,7 @@ function init_phases!(phases, particles, xc, yc, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, xc, yc, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, xc, yc, r) end import ParallelStencil.INDICES @@ -76,12 +70,12 @@ end # BEGIN MAIN SCRIPT # -------------------------------------------------------------------------------- function Sinking_Block2D() - ar = 1 - n = 32 - nx = n - ny = n + ar = 1 + n = 32 + nx = n + ny = n init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ------------------------------------ ly = 500e3 @@ -117,7 +111,7 @@ function Sinking_Block2D() # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 20, 1 + nxcell, max_xcell, min_xcell = 20, 40, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -128,28 +122,24 @@ function Sinking_Block2D() xc_anomaly = 250e3 # origin of thermal anomaly yc_anomaly = -(ly-400e3) # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, xc_anomaly, abs(yc_anomaly), r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, CFL = 0.95 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=@ones(ni...), P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=@ones(ni...), P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) # ---------------------------------------------------- # Viscosity - η = @ones(ni...) args = (; dt = dt, ΔTc = @zeros(ni...)) η_cutoff = -Inf, Inf - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, η_cutoff - ) - η_vep = deepcopy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions @@ -168,24 +158,24 @@ function Sinking_Block2D() di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax=150e3, - nout=1e3, - viscosity_cutoff = η_cutoff, - verbose = false, + kwargs = ( + iterMax=150e3, + nout=1e3, + viscosity_cutoff = η_cutoff, + verbose = false, + ) ); dt = compute_dt(stokes, di, igg) # ------------------------------ Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) velocity = @. √(Vx_v^2 + Vy_v^2 ) finalize_global_grid(; finalize_MPI = true) diff --git a/test/test_stokes_elastic_buildup.jl b/test/test_stokes_elastic_buildup.jl index 3ac5b741..7ae7e950 100644 --- a/test/test_stokes_elastic_buildup.jl +++ b/test/test_stokes_elastic_buildup.jl @@ -1,12 +1,11 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +const backend = CPUBackend include("../miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl") diff --git a/test/test_stokes_solcx.jl b/test/test_stokes_solcx.jl index 88a81a0d..a4034172 100644 --- a/test/test_stokes_solcx.jl +++ b/test/test_stokes_solcx.jl @@ -1,12 +1,11 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +const backend = CPUBackend include("../miniapps/benchmarks/stokes2D/solcx/SolCx.jl") diff --git a/test/test_stokes_solkz.jl b/test/test_stokes_solkz.jl index 0cb7a78d..9c0246c1 100644 --- a/test/test_stokes_solkz.jl +++ b/test/test_stokes_solkz.jl @@ -1,12 +1,11 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +const backend = CPUBackend include("../miniapps/benchmarks/stokes2D/solkz/SolKz.jl") From 8f3390704d708c41388ec338a259a1cd5c2412df Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 22:54:26 +0200 Subject: [PATCH 28/91] new test suits --- test/runtests.jl | 9 +++ test/test_traits.jl.jl | 44 ++++++++++++++ test/test_types.jl | 134 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+) create mode 100644 test/test_traits.jl.jl create mode 100644 test/test_types.jl diff --git a/test/runtests.jl b/test/runtests.jl index 6a2032dd..3b2d6d5e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,6 +12,15 @@ function runtests() nfail = 0 printstyled("Testing package JustRelax.jl\n"; bold=true, color=:white) + f0 = "test_traits.jl", "test_types.jl" + for f in f0 + try + include(f) + catch ex + nfail += 1 + end + end + for f in testfiles occursin("burstedde", f) && continue diff --git a/test/test_traits.jl.jl b/test/test_traits.jl.jl new file mode 100644 index 00000000..033e8f39 --- /dev/null +++ b/test/test_traits.jl.jl @@ -0,0 +1,44 @@ +using JustRelax, Test +import JustRelax.JustRelax2D as JR2 +import JustRelax.JustRelax3D as JR3 + +const bk = JustRelax.backend + +@testset "Traits" begin + # test generic arrays + @test bk(Array) === CPUBackendTrait() + @test bk(Matrix) === CPUBackendTrait() + @test bk(Vector) === CPUBackendTrait() + @test bk(rand(2)) === CPUBackendTrait() + @test bk(rand(2,2)) === CPUBackendTrait() + @test bk(rand(2,2,2)) === CPUBackendTrait() + + # test error handling + @test_throws ArgumentError backend(rand()) + @test_throws ArgumentError bk("potato") + + # test JR structs + ## 2D + ni = 2, 2 + stokes2 = JR2.StokesArrays(CPUBackend, ni) + thermal2 = JR2.ThermalArrays(CPUBackend, ni) + + @test bk(stokes2.V) === CPUBackendTrait() + @test bk(stokes2.τ) === CPUBackendTrait() + @test bk(stokes2.R) === CPUBackendTrait() + @test bk(stokes2.P) === CPUBackendTrait() + @test bk(stokes2) === CPUBackendTrait() + @test bk(thermal2) === CPUBackendTrait() + + ## 3D + ni = 2, 2, 2 + stokes3 = JR3.StokesArrays(CPUBackend, ni) + thermal3 = JR3.ThermalArrays(CPUBackend, ni) + + @test bk(stokes3.V) === CPUBackendTrait() + @test bk(stokes3.τ) === CPUBackendTrait() + @test bk(stokes3.R) === CPUBackendTrait() + @test bk(stokes3.P) === CPUBackendTrait() + @test bk(stokes3) === CPUBackendTrait() + @test bk(thermal3) === CPUBackendTrait() +end \ No newline at end of file diff --git a/test/test_types.jl b/test/test_types.jl new file mode 100644 index 00000000..7c57a577 --- /dev/null +++ b/test/test_types.jl @@ -0,0 +1,134 @@ +using JustRelax, Test +import JustRelax.JustRelax2D as JR2 +import JustRelax.JustRelax3D as JR3 + +const backend = CPUBackend + +@testset "2D allocators" begin + ni = nx, ny = (2, 2) + + R = JR2.Residual(ni...) + @test R isa JustRelax.Residual + @test isnothing(R.Rz) + @test size(R.Rx) == (nx-1, ny) + @test size(R.Ry) == (nx, ny-1) + @test size(R.RP) == ni + @test R.Rx isa Array + @test R.Ry isa Array + @test R.RP isa Array + @test_throws MethodError JR2.Residual(10.0, 10.0) + + visc = JR2.Viscosity(ni) + @test size(visc.η) == ni + @test size(visc.η_vep) == ni + @test size(visc.ητ) == ni + @test visc.η isa Array + @test visc.η_vep isa Array + @test visc.ητ isa Array + @test_throws MethodError JR2.Viscosity(10.0, 10.0) + + v = JR2.Velocity(ni...) + tensor = JR2.SymmetricTensor(ni...) + + @test size(tensor.xx) == (nx, ny) + @test size(tensor.yy) == (nx, ny) + @test size(tensor.xy) == (nx + 1, ny + 1) + @test size(tensor.xy_c) == (nx, ny) + @test size(tensor.II) == (nx, ny) + + @test tensor.xx isa Array + @test tensor.yy isa Array + @test tensor.xy isa Array + @test tensor.xy_c isa Array + @test tensor.II isa Array + + stokes = JR2.StokesArrays(backend, ni) + + @test size(stokes.P) == ni + @test size(stokes.P0) == ni + @test size(stokes.∇V) == ni + @test size(stokes.EII_pl) == ni + + @test stokes.P isa Array + @test stokes.P0 isa Array + @test stokes.∇V isa Array + @test stokes.V isa JustRelax.Velocity + @test stokes.τ isa JustRelax.SymmetricTensor + @test stokes.τ_o isa JustRelax.SymmetricTensor + @test stokes.ε isa JustRelax.SymmetricTensor + @test stokes.ε_pl isa JustRelax.SymmetricTensor + @test stokes.EII_pl isa Array + @test stokes.viscosity isa JustRelax.Viscosity + @test stokes.R isa JustRelax.Residual + + @test_throws MethodError JR2.StokesArrays(backend, 10.0, 10.0) +end + +@testset "3D allocators" begin + ni = nx, ny, nz = (2, 2, 2) + + R = JR3.Residual(ni...) + @test R isa JustRelax.Residual + @test size(R.Rx) == (nx-1, ny, nz) + @test size(R.Ry) == (nx, ny-1, nz) + @test size(R.Rz) == (nx, ny, nz-1) + @test size(R.RP) == ni + @test R.Rx isa Array + @test R.Ry isa Array + @test R.Rz isa Array + @test R.RP isa Array + @test_throws MethodError JR3.Residual(1.0, 1.0, 1.0) + + visc = JR3.Viscosity(ni) + @test size(visc.η) == ni + @test size(visc.η_vep) == ni + @test size(visc.ητ) == ni + @test visc.η isa Array + @test visc.η_vep isa Array + @test visc.ητ isa Array + @test_throws MethodError JR3.Viscosity(1.0, 1.0, 1.0) + + v = JR3.Velocity(ni...) + tensor = JR3.SymmetricTensor(ni...) + + @test size(tensor.xx) == ni + @test size(tensor.yy) == ni + @test size(tensor.xy) == (nx + 1, ny + 1, nz ) + @test size(tensor.yz) == (nx , ny + 1, nz + 1) + @test size(tensor.xz) == (nx + 1, ny , nz + 1) + @test size(tensor.xy_c) == ni + @test size(tensor.yz_c) == ni + @test size(tensor.xz_c) == ni + @test size(tensor.II) == ni + + @test tensor.xx isa Array + @test tensor.yy isa Array + @test tensor.xy isa Array + @test tensor.yz isa Array + @test tensor.xz isa Array + @test tensor.xy_c isa Array + @test tensor.yz_c isa Array + @test tensor.xz_c isa Array + @test tensor.II isa Array + + stokes = JR3.StokesArrays(backend, ni) + + @test size(stokes.P) == ni + @test size(stokes.P0) == ni + @test size(stokes.∇V) == ni + @test size(stokes.EII_pl) == ni + + @test stokes.P isa Array + @test stokes.P0 isa Array + @test stokes.∇V isa Array + @test stokes.V isa JustRelax.Velocity + @test stokes.τ isa JustRelax.SymmetricTensor + @test stokes.τ_o isa JustRelax.SymmetricTensor + @test stokes.ε isa JustRelax.SymmetricTensor + @test stokes.ε_pl isa JustRelax.SymmetricTensor + @test stokes.EII_pl isa Array + @test stokes.viscosity isa JustRelax.Viscosity + @test stokes.R isa JustRelax.Residual + + @test_throws MethodError JR3.StokesArrays(backend, 10.0, 10.0) +end \ No newline at end of file From 10d4e4a33be9497b47a91fedc3a6f1f6282fcde8 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 22:54:37 +0200 Subject: [PATCH 29/91] remove dead files --- test_traits.jl.jl | 7 ----- test_types.jl | 67 ----------------------------------------------- 2 files changed, 74 deletions(-) delete mode 100644 test_traits.jl.jl delete mode 100644 test_types.jl diff --git a/test_traits.jl.jl b/test_traits.jl.jl deleted file mode 100644 index bbba4703..00000000 --- a/test_traits.jl.jl +++ /dev/null @@ -1,7 +0,0 @@ -@test backend(Array) == CPUBackendTrait() -@test backend(Matrix) == CPUBackendTrait() -@test backend(Vector) == CPUBackendTrait() -@test backend(rand(2)) == CPUBackendTrait() -@test backend(rand(2,2)) == CPUBackendTrait() -@test backend(rand(2,2,2)) == CPUBackendTrait() -@test_throws ArgumentError backend(rand()) diff --git a/test_types.jl b/test_types.jl deleted file mode 100644 index e5960886..00000000 --- a/test_types.jl +++ /dev/null @@ -1,67 +0,0 @@ -using Test -using ParallelStencil -@init_parallel_stencil(Threads,Float64,2) - -include("src/stokes/types/types.jl") - -T = Array - -ni = nx, ny = (10, 10) - -R = Residual(ni) -@test isnothing(R.Rz) -@test size(R.Rx) == (nx-1, ny) -@test size(R.Ry) == (nx, ny-1) -@test size(R.RP) == ni -@test R.Rx isa Array -@test R.Ry isa Array -@test R.RP isa Array - -@test Residual(ni) isa Residual -@test Residual(nx, ny) isa Residual -@test_throws ArgumentError Residual(10.0, 10.0) - -visc = Viscosity(ni) -@test size(visc.η) == ni -@test size(visc.η_vep) == ni -@test size(visc.ητ) == ni -@test visc.η isa Array -@test visc.η_vep isa Array -@test visc.ητ isa Array -@test_throws ArgumentError Viscosity(10.0, 10.0) - -v = Velocity(ni) -tensor = SymmetricTensor(ni) - -@test size(tensor.xx) == (nx, ny) -@test size(tensor.yy) == (nx, ny) -@test size(tensor.xy) == (nx + 1, ny + 1) -@test size(tensor.xy_c) == (nx, ny) -@test size(tensor.II) == (nx, ny) - -@test tensor.xx isa Array -@test tensor.yy isa Array -@test tensor.xy isa Array -@test tensor.xy_c isa Array -@test tensor.II isa Array - -stokes = StokesArrays(ni) - -@test size(stokes.P) == ni -@test size(stokes.P0) == ni -@test size(stokes.∇V) == ni -@test size(stokes.EII_pl) == ni - -@test stokes.P isa Array -@test stokes.P0 isa Array -@test stokes.∇V isa Array -@test stokes.V isa Velocity -@test stokes.τ isa SymmetricTensor -@test stokes.τ_o isa SymmetricTensor -@test stokes.ε isa SymmetricTensor -@test stokes.ε_pl isa SymmetricTensor -@test stokes.EII_pl isa Array -@test stokes.viscosity isa Viscosity -@test stokes.R isa Residual - -@test_throws ArgumentError StokesArrays(10.0, 10.0) From 4e1c1ee741e0eae4b038e32513e659bc696faa36 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Sun, 28 Apr 2024 22:54:50 +0200 Subject: [PATCH 30/91] format --- src/Utils.jl | 8 ++------ src/stokes/Stokes3D.jl | 9 ++------- src/types/constructors/heat_diffusion.jl | 6 +++--- src/types/stokes.jl | 1 - 4 files changed, 7 insertions(+), 17 deletions(-) diff --git a/src/Utils.jl b/src/Utils.jl index bc204aa5..6acdb73e 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -299,14 +299,10 @@ macro residuals(A) end end -@inline function unpack_residuals( - A::JustRelax.Residual{<:AbstractArray{T,2}} -) where {T} +@inline function unpack_residuals(A::JustRelax.Residual{<:AbstractArray{T,2}}) where {T} return A.Rx, A.Ry end -@inline function unpack_residuals( - A::JustRelax.Residual{<:AbstractArray{T,3}} -) where {T} +@inline function unpack_residuals(A::JustRelax.Residual{<:AbstractArray{T,3}}) where {T} return A.Rx, A.Ry, A.Rz end diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 6fc7b8ff..58582845 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -407,13 +407,8 @@ function _solve!( # Update viscosity compute_viscosity!( - stokes, - viscosity_relaxation, - phase_ratios, - args, - rheology, - viscosity_cutoff, - ) + stokes, viscosity_relaxation, phase_ratios, args, rheology, viscosity_cutoff + ) @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl index d193aabc..90772dde 100644 --- a/src/types/constructors/heat_diffusion.jl +++ b/src/types/constructors/heat_diffusion.jl @@ -47,15 +47,15 @@ function ThermalArrays(nx::Integer, ny::Integer, nz::Integer) Tc = @zeros(nx, ny, nz) H = @zeros(nx, ny, nz) shear_heating = @zeros(nx, ny, nz) - dT_dt = @zeros(nx -1, ny - 1, nz - 1) + dT_dt = @zeros(nx - 1, ny - 1, nz - 1) qTx = @zeros(nx, ny - 1, nz - 1) qTy = @zeros(nx - 1, ny, nz - 1) qTz = @zeros(nx - 1, ny - 1, nz) qTx2 = @zeros(nx, ny - 1, nz - 1) qTy2 = @zeros(nx - 1, ny, nz - 1) qTz2 = @zeros(nx - 1, ny - 1, nz) - ResT = @zeros(nx -1, ny - 1, nz - 1) + ResT = @zeros(nx - 1, ny - 1, nz - 1) return JustRelax.ThermalArrays( - T, Tc, Told, ΔT,ΔTc, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT + T, Tc, Told, ΔT, ΔTc, dT_dt, qTx, qTy, qTz, qTx2, qTy2, qTz2, H, shear_heating, ResT ) end diff --git a/src/types/stokes.jl b/src/types/stokes.jl index e35743aa..e7b7cd48 100644 --- a/src/types/stokes.jl +++ b/src/types/stokes.jl @@ -10,7 +10,6 @@ end 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")) From af511fb38b4f44df98a4ceb3a82651b3c50dbea5 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 10:05:26 +0200 Subject: [PATCH 31/91] fix edge-case in heat diffusion --- .../DiffusionPT_coefficients.jl | 68 +++++++++++++++++++ src/thermal_diffusion/DiffusionPT_solver.jl | 16 +++-- src/types/heat_diffusion.jl | 12 ---- 3 files changed, 77 insertions(+), 19 deletions(-) diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index e39e5ab9..c64725b2 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -1,3 +1,25 @@ +function PTThermalCoeffs( + ::Type{CPUBackend}, K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 +) where {nDim,T} + + PTThermalCoeffs( + K, ρCp, dt, di, li; ϵ=ϵ, CFL=CFL + ) +end + +function PTThermalCoeffs( + K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 +) where {nDim,T} + Vpdτ = min(di...) * CFL + max_lxyz = max(li...) + max_lxyz2 = max_lxyz^2 + Re = @. π + √(π * π + ρCp * max_lxyz2 / K / dt) # Numerical Reynolds number + θr_dτ = @. max_lxyz / Vpdτ / Re + dτ_ρ = @. Vpdτ * max_lxyz / K / Re + return JustRelax.PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) +end + + function PTThermalCoeffs( ::Type{CPUBackend}, rheology, @@ -9,6 +31,30 @@ function PTThermalCoeffs( li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3, +) where {nDim,T} + PTThermalCoeffs( + rheology, + phase_ratios, + args, + dt, + ni, + di, + li; + ϵ=ϵ, + CFL=CFL, + ) +end + +function PTThermalCoeffs( + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) @@ -32,6 +78,28 @@ function PTThermalCoeffs( li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3, +) where {nDim,T} + PTThermalCoeffs( + rheology, + args, + dt, + ni, + di, + li; + ϵ=ϵ, + CFL=CFL, + ) +end + +function PTThermalCoeffs( + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index eb73cd9a..8b5170fe 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -1,14 +1,15 @@ function heatdiffusion_PT!(thermal, args...; kwargs) - return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs...) + return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs=kwargs) end +heatdiffusion_PT!(::CPUBackendTrait, thermal, args...; kwargs) = _heatdiffusion_PT!(thermal, args...; kwargs...) + """ heatdiffusion_PT!(thermal, pt_thermal, K, ρCp, dt, di; iterMax, nout, verbose) Heat diffusion solver using Pseudo-Transient iterations. Both `K` and `ρCp` are n-dimensional arrays. """ -function heatdiffusion_PT!( - ::CPUBackendTrait, +function _heatdiffusion_PT!( thermal::JustRelax.ThermalArrays, pt_thermal::JustRelax.PTThermalCoeffs, thermal_bc::TemperatureBoundaryConditions, @@ -17,10 +18,11 @@ function heatdiffusion_PT!( dt, di; igg=nothing, - b_width=(4, 4, 4), + b_width=(4, 4, 1), iterMax=50e3, nout=1e3, verbose=true, + kwargs... ) # Compute some constant stuff _dt = inv(dt) @@ -96,8 +98,7 @@ end Heat diffusion solver using Pseudo-Transient iterations. """ -function heatdiffusion_PT!( - ::CPUBackendTrait, +function _heatdiffusion_PT!( thermal::JustRelax.ThermalArrays, pt_thermal::JustRelax.PTThermalCoeffs, thermal_bc::TemperatureBoundaryConditions, @@ -111,6 +112,7 @@ function heatdiffusion_PT!( iterMax=50e3, nout=1e3, verbose=true, + kwargs... ) phases = get_phase(phase) @@ -121,7 +123,7 @@ function heatdiffusion_PT!( ϵ = pt_thermal.ϵ ni = size(thermal.Tc) @copy thermal.Told thermal.T - update_pt_thermal_arrays!(pt_thermal, phase, rheology, args, _dt) + !isnothing(phase) && update_pt_thermal_arrays!(pt_thermal, phase, rheology, args, _dt) # errors iter_count = Int64[] diff --git a/src/types/heat_diffusion.jl b/src/types/heat_diffusion.jl index 2f8f475e..6d855057 100644 --- a/src/types/heat_diffusion.jl +++ b/src/types/heat_diffusion.jl @@ -43,15 +43,3 @@ struct PTThermalCoeffs{T,M} return new{T,M}(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) end end - -function PTThermalCoeffs( - K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 -) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - max_lxyz2 = max_lxyz^2 - Re = @. π + √(π * π + ρCp * max_lxyz2 / K / dt) # Numerical Reynolds number - θr_dτ = @. max_lxyz / Vpdτ / Re - dτ_ρ = @. Vpdτ * max_lxyz / K / Re - return PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) -end From 4c0da5a0cf175592d704cf6b9721003efd133533 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 10:06:12 +0200 Subject: [PATCH 32/91] adapt more tests --- test/test_VanKeken.jl | 73 ++++++++++------------------- test/test_diffusion2D.jl | 26 +++++----- test/test_diffusion2D_multiphase.jl | 50 ++++++++++---------- 3 files changed, 64 insertions(+), 85 deletions(-) diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index 74d20538..2a78e8ee 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -5,41 +5,27 @@ using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) using Printf, LinearAlgebra, GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D import JustRelax.@cell -using JustPIC -using JustPIC._2D +const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + + +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # x-length of the domain const λ = 0.9142 # HELPER FUNCTIONS --------------------------------------------------------------- -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -# Initial pressure guess -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - # Initialize phases on the particles function init_phases!(phases, particles) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index) - @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 @@ -56,7 +42,7 @@ function init_phases!(phases, particles) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) end # END OF HELPER FUNCTIONS -------------------------------------------------------- @@ -98,7 +84,7 @@ function VanKeken2D(ny=32, nx=32) ) # Initialize particles ------------------------------- - nxcell, max_p, min_p = 40, 40, 15 + nxcell, max_p, min_p = 40, 80, 20 particles = init_particles( backend, nxcell, max_p, min_p, xvi..., di..., nx, ny ) @@ -107,27 +93,22 @@ function VanKeken2D(ny=32, nx=32) # temperature pPhases, = init_cell_arrays(particles, Val(1)) particle_args = (pPhases, ) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, 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) + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; r=1e0, ϵ=1e-8, CFL = 1 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) - @parallel (JustRelax.@idx ni) JustRelax.compute_ρg!(ρg[2], phase_ratios.center, rheology, args) - @parallel init_P!(stokes.P, ρg[2], xci[2]) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -154,7 +135,7 @@ function VanKeken2D(ny=32, nx=32) while it < nt # Update buoyancy - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -164,23 +145,23 @@ function VanKeken2D(ny=32, nx=32) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 10e3, - nout = 50, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + iterMax = 10e3, + nout = 50, + viscosity_cutoff = (-Inf, Inf) + ) ) dt = compute_dt(stokes, di) / 10 # ------------------------------ # Compute U rms --------------- Urms_it = let - JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) + velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude sum(Vx_v.^2) * prod(di) |> sqrt end @@ -189,15 +170,13 @@ function VanKeken2D(ny=32, nx=32) # ------------------------------ # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # # advect particles in memory move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - @show inject = check_injection(particles) # inject && break - inject && inject_particles_phase!(particles, pPhases, (), (), xvi) + inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt diff --git a/test/test_diffusion2D.jl b/test/test_diffusion2D.jl index 724ed5eb..760d485f 100644 --- a/test/test_diffusion2D.jl +++ b/test/test_diffusion2D.jl @@ -1,14 +1,14 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +import JustRelax.@cell +using GeoParams # HELPER FUNCTIONS --------------------------------------------------------------- @parallel_indices (i, j) function init_T!(T, z) @@ -33,15 +33,15 @@ function elliptical_perturbation!(T, δT, xc, yc, r, xvi) @parallel _elliptical_perturbation!(T, δT, xc, yc, r, xvi...) end + # MAIN SCRIPT -------------------------------------------------------------------- function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, K0=3.0) kyr = 1e3 * 3600 * 24 * 365.25 Myr = 1e3 * kyr ttot = 1 * Myr # total simulation time dt = 50 * kyr # physical time step - init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ni = (nx, ny) @@ -63,7 +63,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, args = (; P=P) ## Allocate arrays needed for every Thermal Diffusion - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal.H .= 1e-6 # radiogenic heat production # physical parameters ρ = @fill(ρ0, ni...) @@ -71,7 +71,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, K = @fill(K0, ni...) ρCp = @. Cp * ρ - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li) + pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -98,7 +98,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, args, dt, di; - verbose=false, + kwargs = (; verbose=false), ) t += dt @@ -113,7 +113,7 @@ end nx=32; ny=32; thermal = diffusion_2D(; nx = nx, ny = ny) - @test thermal.T[Int(ceil(size(thermal.T)[1]/2)), Int(ceil(size(thermal.T)[2]/2))] ≈ 1823.6076461523571 atol=1e-1 - @test thermal.Tc[Int(ceil(size(thermal.Tc)[1]/2)), Int(ceil(size(thermal.Tc)[2]/2))] ≈ 1828.3169386441218 atol=1e-1 + @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1823.6076461523571 atol=1e-1 + @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1828.3169386441218 atol=1e-1 end end diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl index d4cb8d6a..89346d78 100644 --- a/test/test_diffusion2D_multiphase.jl +++ b/test/test_diffusion2D_multiphase.jl @@ -1,25 +1,22 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using Printf, LinearAlgebra, GeoParams, CellArrays, StaticArrays -using JustRelax +using GeoParams +using JustRelax, JustRelax.JustRelax2D + +const backend_JR = CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using JustPIC -using JustPIC._2D +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend import JustRelax.@cell - distance(p1, p2) = mapreduce(x->(x[1]-x[2])^2, +, zip(p1, p2)) |> sqrt @parallel_indices (i, j) function init_T!(T, z) @@ -68,7 +65,7 @@ function init_phases!(phases, particles, xc, yc, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) end @parallel_indices (I...) function compute_temperature_source_terms!(H, rheology, phase_ratios, args) @@ -80,13 +77,14 @@ end end function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) + kyr = 1e3 * 3600 * 24 * 365.25 Myr = 1e3 * kyr ttot = 1 * Myr # total simulation time dt = 50 * kyr # physical time step init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ni = (nx, ny) @@ -118,7 +116,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) args = (; P=P) ## Allocate arrays needed for every Thermal Diffusion - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -129,7 +127,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) r = 10e3 # thermal perturbation radius center_perturbation = lx/2, -ly/2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 40, 40, 1 @@ -139,16 +137,16 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) # temperature pPhases, = init_cell_arrays(particles, Val(1)) init_phases!(pPhases, particles, center_perturbation..., r) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) JustRelax.phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- @parallel (@idx ni) compute_temperature_source_terms!(thermal.H, rheology, phase_ratios.center, args) # PT coefficients for thermal diffusion args = (; P=P, T=thermal.Tc) - pt_thermal = JustRelax.PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.65 / √2 + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.65 / √2 ) # Time loop @@ -164,10 +162,12 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) args, dt, di; - phase = phase_ratios, - iterMax = 1e3, - nout = 10, - verbose=false, + kwargs = ( + phase = phase_ratios, + iterMax = 1e3, + nout = 10, + verbose=false, + ) ) it += 1 @@ -182,7 +182,7 @@ end nx=32; ny=32; thermal = diffusion_2D(; nx = nx, ny = ny) - @test thermal.T[Int(ceil(size(thermal.T)[1]/2)), Int(ceil(size(thermal.T)[2]/2))] ≈ 1819.2297931741878 atol=1e-1 - @test thermal.Tc[Int(ceil(size(thermal.Tc)[1]/2)), Int(ceil(size(thermal.Tc)[2]/2))] ≈ 1824.3532934301472 atol=1e-1 + @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1819.2297931741878 atol=1e-1 + @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1824.3532934301472 atol=1e-1 end -end +end \ No newline at end of file From f49544b6aec06bec84c563e4383a228f15740a95 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 12:17:23 +0200 Subject: [PATCH 33/91] more tests --- .../shear_heating/Shearheating_rheology.jl | 2 +- miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl | 24 +++- test/test_shearheating2D.jl | 20 ++-- test/test_shearheating3D.jl | 112 ++++++++---------- test/test_stokes_solvi3D.jl | 8 +- 5 files changed, 81 insertions(+), 85 deletions(-) diff --git a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl index 5041a86a..81ac67dd 100644 --- a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl +++ b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl @@ -74,5 +74,5 @@ function init_phases!(phases, particles, Lx, Ly, d, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, Ly, d) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, Ly, d) end diff --git a/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl b/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl index 7d6ba67c..e3ce39f0 100644 --- a/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl +++ b/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl @@ -4,7 +4,7 @@ using ParallelStencil.FiniteDifferences3D # D. W. Schmid and Y. Y. Podladchikov. Analytical solutions for deformable elliptical inclusions in # general shear. Geophysical Journal International, 155(1):269–288, 2003. -include("vizSolVi3D.jl") +# include("vizSolVi3D.jl") @parallel function smooth!(A2::AbstractArray{T,3}, A::AbstractArray{T,3}, fact::T) where {T} @inn(A2) = @inn(A) + one(T) / 6.1 / fact * (@d2_xi(A) + @d2_yi(A) + @d2_zi(A)) @@ -73,7 +73,8 @@ function solVi3D(; ## Allocate arrays needed for every Stokes problem # general stokes arrays - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) + (; η) = stokes.viscosity # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; CFL=1 / √3) @@ -84,12 +85,12 @@ function solVi3D(; G = 1.0 # elastic shear modulus # dt = η0 / (G * ξ) dt = Inf - η = viscosity(ni, di, li, rc, η0, ηi) + η .= viscosity(ni, di, li, rc, η0, ηi) Gc = @fill(G, ni...) Kb = @fill(Inf, ni...) ## Boundary conditions - pureshear_bc!(stokes, di, li, εbg) + pureshear_bc!(stokes, xci, xvi, εbg) flow_bcs = FlowBoundaryConditions(; free_slip = (left=false, right=false, top=false, bot=false, back=false, front=false), no_slip = (left=false, right=false, top=false, bot=false, back=false, front=false), @@ -105,7 +106,20 @@ function solVi3D(; local iters while t < ttot iters = solve!( - stokes, pt_stokes, di, flow_bcs, ρg, η, Kb, Gc, dt, igg; iterMax=5000, nout=100, verbose=false, + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + Kb, + Gc, + dt, + igg; + kwargs = (; + iterMax=5000, + nout=100, + verbose=false + ), ) t += Δt end diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 04a199e6..0192c9c6 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -1,13 +1,12 @@ # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D const backend_JR = JustRelax.CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using JustPIC -using JustPIC._2D +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. @@ -15,14 +14,13 @@ const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBac # const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams +using Printf, GeoParams # Load file with all the rheology configurations include("../miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl") # ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- - function copyinn_x!(A, B) @parallel function f_x(A, B) @@ -115,7 +113,7 @@ function Shearheating2D() @parallel init_P!(stokes.P, ρg[2], xci[2]) # Rheology - args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion @@ -181,12 +179,10 @@ function Shearheating2D() @views thermal.T[2:end-1, :] .= T_buffer temperature2center!(thermal) - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, rheology, # needs to be a tuple dt, ) diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index 4c622f6e..9be61145 100644 --- a/test/test_shearheating3D.jl +++ b/test/test_shearheating3D.jl @@ -3,26 +3,23 @@ using Test, Suppressor # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D +const backend_JR = JustRelax.CPUBackend import JustRelax.@cell -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) -using JustPIC -using JustPIC._3D +using ParallelStencil, ParallelStencil.FiniteDifferences3D +@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) + +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 3) #or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays +using Printf, GeoParams -# Load file with all the rheology configurations # Load file with all the rheology configurations include("../miniapps/benchmarks/stokes3D/shear_heating/Shearheating_rheology.jl") @@ -41,11 +38,14 @@ end end ## END OF HELPER FUNCTION ------------------------------------------------------------ +nx,ny,nz = 32,32,32 + ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- function Shearheating3D(nx=16, ny=16, nz=16) init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, nz; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, nz; init_MPI = init_mpi)...) + # Physical domain ------------------------------------ lx = 70e3 # domain length in x ly = 70e3 # domain length in y @@ -65,7 +65,7 @@ function Shearheating3D(nx=16, ny=16, nz=16) # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 40, 1 + nxcell, max_xcell, min_xcell = 20, 40, 10 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -80,19 +80,19 @@ function Shearheating3D(nx=16, ny=16, nz=16) yc_anomaly = ly/2 # origin of thermal anomaly zc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √3.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) @@ -100,27 +100,21 @@ function Shearheating3D(nx=16, ny=16, nz=16) # Initialize constant temperature @views thermal.T .= 273.0 + 400 thermal_bcs!(thermal.T, thermal_bc) - - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(3)) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + 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 = dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - η_vep = deepcopy(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=5e-2 / √3 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=5e-2 / √3 ) # Boundary conditions @@ -130,9 +124,9 @@ function Shearheating3D(nx=16, ny=16, nz=16) ) ## Compression and not extension - fix this εbg = 5e-14 - stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2]) - stokes.V.Vy .= PTArray([ -(y - ly/2) * εbg for _ in 1:nx+2, y in xvi[2], _ in 1:nz+2]) - stokes.V.Vz .= PTArray([ (lz - abs(z)) * εbg for _ in 1:nx+2, _ in 1:ny+2, z in xvi[3]]) + stokes.V.Vx .= PTArray(backend_JR)([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2]) + stokes.V.Vy .= PTArray(backend_JR)([ -(y - ly/2) * εbg for _ in 1:nx+2, y in xvi[2], _ in 1:nz+2]) + stokes.V.Vz .= PTArray(backend_JR)([ (lz - abs(z)) * εbg for _ in 1:nx+2, _ in 1:ny+2, z in xvi[3]]) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) @@ -143,11 +137,9 @@ function Shearheating3D(nx=16, ny=16, nz=16) local iters while it < 5 # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_ρg!(ρg[end], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -157,19 +149,19 @@ function Shearheating3D(nx=16, ny=16, nz=16) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 100e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf), - verbose=false, + kwargs = ( + iterMax = 100e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf), + verbose=false, + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -177,12 +169,10 @@ function Shearheating3D(nx=16, ny=16, nz=16) particle2grid!(thermal.T, pT, xvi, particles) temperature2center!(thermal) - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, rheology, # needs to be a tuple dt, ) @@ -196,38 +186,36 @@ function Shearheating3D(nx=16, ny=16, nz=16) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = false, + kwargs =( + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = false, + ) ) # ------------------------------ # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy, grid_vz), dt) # advect particles in memory move_particles!(particles, xvi, particle_args) # interpolate fields from grid vertices to particles grid2particle_flip!(pT, xvi, thermal.T, thermal.Told, particles) # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), xvi) + 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) + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt # ------------------------------ - end finalize_global_grid(; finalize_MPI = true) return iters, thermal - - end @testset "Shearheating3D" begin diff --git a/test/test_stokes_solvi3D.jl b/test/test_stokes_solvi3D.jl index a06cc519..23e7860f 100644 --- a/test/test_stokes_solvi3D.jl +++ b/test/test_stokes_solvi3D.jl @@ -1,14 +1,12 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax +using JustRelax, JustRelax.JustRelax3D +const backend = CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 3) -environment!(model) - include("../miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl") function check_convergence_case1() From a3e6b4176302ea31c1053e87d53730a2ad80cad3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 12:17:50 +0200 Subject: [PATCH 34/91] fix pure shear BC kernel --- src/boundaryconditions/BoundaryConditions.jl | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 6a73b505..ede956e0 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -319,8 +319,23 @@ end function pureshear_bc!( stokes::JustRelax.StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg ) where {T} - stokes.V.Vx[:, 2:(end - 1)] .= PTArray(([εbg * x for x in xvi[1], y in xci[2]])) - stokes.V.Vy[2:(end - 1), :] .= PTArray(([-εbg * y for x in xci[1], y in xvi[2]])) + _T = typeof(stokes.V.Vx) + stokes.V.Vx[:, 2:(end - 1)] .= _T([ εbg * x for x in xvi[1], y in xci[2]]) + stokes.V.Vy[2:(end - 1), :] .= _T([-εbg * y for x in xci[1], y in xvi[2]]) + + return nothing +end + +function pureshear_bc!( + stokes::JustRelax.StokesArrays, xci::NTuple{3,T}, xvi::NTuple{3,T}, εbg +) where {T} + xv, yv, zv = xvi + xc, yc, zc = xci + _T = typeof(stokes.V.Vx) + + stokes.V.Vx[:, 2:(end - 1), 2:(end - 1)] .= _T([ εbg * x for x in xv, y in yc, z in zc]) + stokes.V.Vy[2:(end - 1), :, 2:(end - 1)] .= _T([ εbg * y for x in xc, y in xv, z in zc]) + stokes.V.Vz[2:(end - 1), 2:(end - 1), :] .= _T([-εbg * z for x in xc, y in xc, z in zv]) return nothing end From 61fff6999ceeae652833cd5f63dd0d84073802de Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 14:32:36 +0200 Subject: [PATCH 35/91] more test fixes --- test/runtests.jl | 7 +--- test/test_CellArrays2D.jl | 21 +++++------ test/test_CellArrays3D.jl | 11 +++--- test/test_boundary_conditions2D.jl | 57 ++++++++++++++--------------- test/test_diffusion2D.jl | 2 + test/test_diffusion2D_multiphase.jl | 6 ++- test/test_diffusion3D.jl | 28 +++++++------- test/test_diffusion3D_multiphase.jl | 22 +++++------ test/test_grid2D.jl | 5 +-- test/test_grid3D.jl | 5 +-- test/test_traits.jl.jl | 2 +- 11 files changed, 77 insertions(+), 89 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 3b2d6d5e..a78f0b5a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -14,11 +14,7 @@ function runtests() f0 = "test_traits.jl", "test_types.jl" for f in f0 - try - include(f) - catch ex - nfail += 1 - end + include(f) end for f in testfiles @@ -32,6 +28,7 @@ function runtests() nfail += 1 end end + return nfail end diff --git a/test/test_CellArrays2D.jl b/test/test_CellArrays2D.jl index 2e10e86c..f5d9f30b 100644 --- a/test/test_CellArrays2D.jl +++ b/test/test_CellArrays2D.jl @@ -1,27 +1,26 @@ -using Test, Suppressor, StaticArrays, AllocCheck, JustRelax +using Test, StaticArrays, AllocCheck +using Suppressor +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - @testset "CellArrays 2D" begin - # @suppress begin + @suppress begin ni = 5, 5 - A = JustRelax.@fill(false, ni..., celldims=(2,), eltype=Bool) + A = @fill(false, ni..., celldims=(2,), eltype=Bool) - @test cellaxes(A) === Base.OneTo(2) - @test cellnum(A) == 2 + @test cellaxes(A) === Base.OneTo(2) + @test cellnum(A) == 2 @test new_empty_cell(A) === SA[false, false] @test @cell(A[1, 1, 1]) === false @test (@allocated @cell A[1, 1, 1]) === 0 @cell A[1, 1, 1] = true - @test @cell(A[1, 1, 1]) === true + @test @cell(A[1, 1, 1]) === true @test (@allocated @cell A[1, 1, 1] = true) === 0 - # end - # @test A[1, 1] === SA[true, false] + @test A[1, 1] === SA[true, false] + end # allocs = check_allocs(getindex, (typeof(A), Int64, Int64)) # @test isempty(allocs) end diff --git a/test/test_CellArrays3D.jl b/test/test_CellArrays3D.jl index 35822c68..c14fa7a1 100644 --- a/test/test_CellArrays3D.jl +++ b/test/test_CellArrays3D.jl @@ -1,15 +1,14 @@ -using Test, Suppressor, StaticArrays, AllocCheck, JustRelax +using Test, StaticArrays, AllocCheck +using Suppressor +using JustRelax, JustRelax.JustRelax3D using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) -model = PS_Setup(:Threads, Float64, 3) -environment!(model) - @testset "CellArrays 3D" begin @suppress begin ni = 5, 5, 5 - A = JustRelax.@fill(false, ni..., celldims=(2,), eltype=Bool) + A = @fill(false, ni..., celldims=(2,), eltype=Bool) @test @cell(A[1, 1, 1, 1]) === false @test (@allocated @cell A[1, 1, 1, 1]) == 0 @@ -18,7 +17,7 @@ environment!(model) @test @cell(A[1, 1, 1, 1]) === true @test (@allocated @cell A[1, 1, 1, 1] = true) == 0 - # @test A[1, 1, 1] == SA[true, false] + @test A[1, 1, 1] == SA[true, false] # allocs = check_allocs(getindex, (typeof(A), Int64, Int64, Int64)) # @test isempty(allocs) end diff --git a/test/test_boundary_conditions2D.jl b/test/test_boundary_conditions2D.jl index e0d9dd0b..86326672 100644 --- a/test/test_boundary_conditions2D.jl +++ b/test/test_boundary_conditions2D.jl @@ -1,29 +1,28 @@ -using JustRelax -using Test, Suppressor -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +using JustRelax, JustRelax.JustRelax2D +using Test#, Suppressor +const backend = CPUBackend @testset "Boundary Conditions" begin @suppress begin - n = 5 # number of elements - Vx, Vy = PTArray(rand(n + 1, n + 2)), PTArray(rand(n + 2, n + 1)) + n = 5 # number of elements + Vx, Vy = PTArray(backend)(rand(n + 1, n + 2)), PTArray(backend)(rand(n + 2, n + 1)) # free-slip - bcs = FlowBoundaryConditions(; - no_slip=(left=false, right=false, top=false, bot=false), - free_slip=(left=true, right=true, top=true, bot=true), + bcs = FlowBoundaryConditions(; + no_slip = (left=false, right=false, top=false, bot=false), + free_slip = (left=true, right=true, top=true, bot=true), ) flow_bcs!(bcs, Vx, Vy) - @test @views Vx[:, 1] == Vx[:, 2] - @test @views Vx[:, end] == Vx[:, end - 1] - @test @views Vy[1, :] == Vy[2, :] - @test @views Vy[end, :] == Vy[end - 1, :] + @test @views Vx[: , 1] == Vx[:, 2] + @test @views Vx[: , end] == Vx[:, end - 1] + @test @views Vy[1 , :] == Vy[2, :] + @test @views Vy[end, :] == Vy[end - 1, :] # no-slip - Vx, Vy = PTArray(rand(n + 1, n + 2)), PTArray(rand(n + 2, n + 1)) - bcs = FlowBoundaryConditions(; - no_slip=(left=true, right=true, top=true, bot=true), - free_slip=(left=false, right=false, top=false, bot=false), + Vx, Vy = PTArray(backend)(rand(n + 1, n + 2)), PTArray(backend)(rand(n + 2, n + 1)) + bcs = FlowBoundaryConditions(; + no_slip = (left=true, right=true, top=true, bot=true), + free_slip = (left=false, right=false, top=false, bot=false), ) flow_bcs!(bcs, Vx, Vy) @test sum(!iszero(Vx[1 , i]) for i in axes(Vx,2)) == 0 @@ -36,26 +35,26 @@ environment!(model) @test @views Vx[: , end] == -Vx[: , end - 1] # test with StokesArrays - ni = 5, 5 - stokes = StokesArrays(ni, ViscoElastic) - stokes.V.Vx .= PTArray(rand(n + 1, n + 2)) - stokes.V.Vy .= PTArray(rand(n + 2, n + 1)) + ni = 5, 5 + stokes = StokesArrays(backend, ni) + stokes.V.Vx .= PTArray(backend)(rand(n + 1, n + 2)) + stokes.V.Vy .= PTArray(backend)(rand(n + 2, n + 1)) # free-slip - flow_bcs = FlowBoundaryConditions(; - no_slip=(left=false, right=false, top=false, bot=false), - free_slip=(left=true, right=true, top=true, bot=true), + flow_bcs = FlowBoundaryConditions(; + no_slip = (left=false, right=false, top=false, bot=false), + free_slip = (left=true, right=true, top=true, bot=true), ) flow_bcs!(stokes, flow_bcs) - @test @views stokes.V.Vx[:, 1] == stokes.V.Vx[:, 2] + @test @views stokes.V.Vx[:, 1] == stokes.V.Vx[:, 2] @test @views stokes.V.Vx[:, end] == stokes.V.Vx[:, end - 1] - @test @views stokes.V.Vy[1, :] == stokes.V.Vy[2, :] + @test @views stokes.V.Vy[1, :] == stokes.V.Vy[2, :] @test @views stokes.V.Vy[end, :] == stokes.V.Vy[end - 1, :] # no-slip - flow_bcs = FlowBoundaryConditions(; - no_slip=(left=true, right=true, top=true, bot=true), - free_slip=(left=false, right=false, top=false, bot=false), + flow_bcs = FlowBoundaryConditions(; + no_slip = (left=true, right=true, top=true, bot=true), + free_slip = (left=false, right=false, top=false, bot=false), ) flow_bcs!(stokes, flow_bcs) diff --git a/test/test_diffusion2D.jl b/test/test_diffusion2D.jl index 760d485f..39689add 100644 --- a/test/test_diffusion2D.jl +++ b/test/test_diffusion2D.jl @@ -113,6 +113,8 @@ end nx=32; ny=32; thermal = diffusion_2D(; nx = nx, ny = ny) + + nx_T, ny_T = size(thermal.T) @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1823.6076461523571 atol=1e-1 @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1828.3169386441218 atol=1e-1 end diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl index 89346d78..03306e8a 100644 --- a/test/test_diffusion2D_multiphase.jl +++ b/test/test_diffusion2D_multiphase.jl @@ -182,7 +182,9 @@ end nx=32; ny=32; thermal = diffusion_2D(; nx = nx, ny = ny) - @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1819.2297931741878 atol=1e-1 - @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1824.3532934301472 atol=1e-1 + + nx_T, ny_T = size(thermal.T) + @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1819.2297931741878 atol=1e-1 + @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1824.3532934301472 atol=1e-1 end end \ No newline at end of file diff --git a/test/test_diffusion3D.jl b/test/test_diffusion3D.jl index 19661f01..c64500b2 100644 --- a/test/test_diffusion3D.jl +++ b/test/test_diffusion3D.jl @@ -2,13 +2,11 @@ push!(LOAD_PATH, "..") using Test, Suppressor using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) -environment!(model) +const backend = CPUBackend # HELPER FUNCTIONS --------------------------------------------------------------- @parallel_indices (i, j, k) function init_T!(T, z) @@ -48,10 +46,10 @@ function diffusion_3D(; finalize_MPI = false, ) - kyr = 1e3 * 3600 * 24 * 365.25 - Myr = 1e6 * 3600 * 24 * 365.25 - ttot = 1 * Myr # total simulation time - dt = 50 * kyr # physical time step + kyr = 1e3 * 3600 * 24 * 365.25 + Myr = 1e6 * 3600 * 24 * 365.25 + ttot = 1 * Myr # total simulation time + dt = 50 * kyr # physical time step # Physical domain ni = (nx, ny, nz) @@ -76,7 +74,7 @@ function diffusion_3D(; ## Allocate arrays needed for every Thermal Diffusion # general thermal arrays - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal.H .= 1e-6 # physical parameters ρ = @fill(ρ0, ni...) @@ -85,7 +83,7 @@ function diffusion_3D(; ρCp = @. Cp * ρ # Boundary conditions - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li; CFL = 0.75 / √3.1) + pt_thermal = PTThermalCoeffs(backend, K, ρCp, dt, di, li; CFL = 0.75 / √3.1) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) @@ -95,7 +93,7 @@ function diffusion_3D(; # Add thermal perturbation δT = 100e0 # thermal perturbation r = 10e3 # thermal perturbation radius - center_perturbation = lx/2, ly/2, -lz/2 + center_perturbation = lx / 2, ly / 2, -lz / 2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) t = 0.0 @@ -111,9 +109,11 @@ function diffusion_3D(; rheology, args, dt, - di,; - igg, - verbose=false, + di; + kwargs = (; + igg, + verbose=false + ), ) t += dt diff --git a/test/test_diffusion3D_multiphase.jl b/test/test_diffusion3D_multiphase.jl index 24b3bca5..33ca4354 100644 --- a/test/test_diffusion3D_multiphase.jl +++ b/test/test_diffusion3D_multiphase.jl @@ -1,7 +1,7 @@ push!(LOAD_PATH, "..") -using Test, Suppressor -using Printf, LinearAlgebra, GeoParams, CellArrays, StaticArrays +using Test, Suppressor, GeoParams +using JustRelax, JustRelax.JustRelax3D using JustRelax using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -13,10 +13,6 @@ using JustPIC._3D # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - import JustRelax.@cell @parallel_indices (i, j, k) function init_T!(T, z) @@ -49,24 +45,24 @@ function init_phases!(phases, particles, xc, yc, zc, r) @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, center, r) @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) # quick escape - JustRelax.@cell(index[ip, I...]) == 0 && continue + @cell(index[ip, I...]) == 0 && continue - x = JustRelax.@cell px[ip, I...] - y = JustRelax.@cell py[ip, I...] - z = JustRelax.@cell pz[ip, I...] + x = @cell px[ip, I...] + y = @cell py[ip, I...] + z = @cell pz[ip, I...] # plume - rectangular if (((x - center[1]))^2 + ((y - center[2]))^2 + ((z - center[3]))^2) ≤ r^2 - JustRelax.@cell phases[ip, I...] = 2.0 + @cell phases[ip, I...] = 2.0 else - JustRelax.@cell phases[ip, I...] = 1.0 + @cell phases[ip, I...] = 1.0 end end return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) end function diffusion_3D(; diff --git a/test/test_grid2D.jl b/test/test_grid2D.jl index 092a7b6c..314420cd 100644 --- a/test/test_grid2D.jl +++ b/test/test_grid2D.jl @@ -1,7 +1,4 @@ -using Test, Suppressor, JustRelax - -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +using Test, Suppressor, JustRelax, JustRelax.JustRelax2D @testset "Grid2D" begin @suppress begin diff --git a/test/test_grid3D.jl b/test/test_grid3D.jl index adb8bdf3..044f0ace 100644 --- a/test/test_grid3D.jl +++ b/test/test_grid3D.jl @@ -1,7 +1,4 @@ -using Test, Suppressor, JustRelax - -model = PS_Setup(:cpu, Float64, 3) -environment!(model) +using Test, Suppressor, JustRelax, JustRelax.JustRelax3D @testset "Grid3D" begin @suppress begin diff --git a/test/test_traits.jl.jl b/test/test_traits.jl.jl index 033e8f39..bc942702 100644 --- a/test/test_traits.jl.jl +++ b/test/test_traits.jl.jl @@ -14,7 +14,7 @@ const bk = JustRelax.backend @test bk(rand(2,2,2)) === CPUBackendTrait() # test error handling - @test_throws ArgumentError backend(rand()) + @test_throws ArgumentError bk(rand()) @test_throws ArgumentError bk("potato") # test JR structs From f01ed46979273c981f644aa128a7315b6ce7af5d Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 29 Apr 2024 15:48:50 +0200 Subject: [PATCH 36/91] fix file typo, add docs --- docs/make.jl | 20 +-- docs/src/index.md | 35 +++++- docs/src/{ => man}/Blankenbach.md | 0 docs/src/{ => man}/advection.md | 0 docs/src/{ => man}/backend.md | 0 docs/src/{ => man}/equations.md | 0 docs/src/man/installation.md | 118 ++++++++++++++++++ docs/src/man/listfunctions.md | 6 + docs/src/{ => man}/material_physics.md | 0 .../{Shearheating.jl => ShearHeating.jl} | 0 10 files changed, 171 insertions(+), 8 deletions(-) rename docs/src/{ => man}/Blankenbach.md (100%) rename docs/src/{ => man}/advection.md (100%) rename docs/src/{ => man}/backend.md (100%) rename docs/src/{ => man}/equations.md (100%) create mode 100644 docs/src/man/installation.md create mode 100644 docs/src/man/listfunctions.md rename docs/src/{ => man}/material_physics.md (100%) rename src/thermal_diffusion/{Shearheating.jl => ShearHeating.jl} (100%) diff --git a/docs/make.jl b/docs/make.jl index 5c42274e..85b05772 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -6,16 +6,22 @@ makedocs(; sitename="JustRelax.jl", authors="Albert de Montserrat and contributors", modules=[JustRelax], - format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true"), # easier local build + format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true", + size_threshold_ignore = ["man/listfunctions.md"]), # easier local build + warnonly = Documenter.except(:footnote), pages=[ - "Home" => "index.md", - "Backend" => "backend.md", - "Equations" => "equations.md", - "Advection" => "advection.md", - "Examples" => [ - "Blankenbach.md" + "Home" => "man/index.md", + "User guide"=> Any[ + "Installation" => "man/installation.md", + ], + "Backend" => "man/backend.md", + "Equations" => "man/equations.md", + "Advection" => "man/advection.md", + "Examples" => Any[ + "Blankenbach" => "man/Blankenbach.md", ], + "List of functions" => "man/listfunctions.md", ], ) diff --git a/docs/src/index.md b/docs/src/index.md index 86514583..263c471c 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,3 +1,36 @@ +```@meta +CurrentModule = JustRelax +``` + # JustRelax.jl -Multi-XPU stokes solver for geodynamic modelling. \ No newline at end of file +Documentation for [JustRelax.jl](https://github.com/PTsolvers/JustRelax.jl) + +Need to solve a very large multi-physics problem on many GPUs in parallel? Just Relax! + +`JustRelax.jl` is a collection of accelerated iterative pseudo-transient solvers using MPI and multiple CPU or GPU backends. It's part of the [PTSolvers organisation](https://ptsolvers.github.io) and +developed within the [GPU4GEO project](https://www.pasc-ch.org/projects/2021-2024/gpu4geo/). Current publications, outreach and news can be found on the [GPU4GEO website](https://ptsolvers.github.io/GPU4GEO/). + +The package relies on other packages as building blocks and parallelisation tools: + +* [ParallelStencil.jl](https://github.com/omlins/ParallelStencil.jl) +* [ImplicitGlobalGrid.jl](https://github.com/omlins/ImplicitGlobalGrid.jl) +* [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl) +* [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl) + + +The package serves several purposes: + + * It provides a collection of solvers to be used in quickly developing new applications + * It provides some standardization so that application codes can + + - more easily handle local material properties through the use of [GeoParams.jl]((https://github.com/JuliaGeodynamics/GeoParams.jl)) + - more easily switch between a pseudo-transient solver and another solvers (e.g. an explicit thermal solvers) + + * It provides a natural repository for contributions of new solvers for use by the larger community + +We provide several miniapps, each designed to solve a well-specified benchmark problem, in order to provide + + - examples of usage in high-performance computing + - basis on which to build more full-featured application codes + - cases for reference and performance tests diff --git a/docs/src/Blankenbach.md b/docs/src/man/Blankenbach.md similarity index 100% rename from docs/src/Blankenbach.md rename to docs/src/man/Blankenbach.md diff --git a/docs/src/advection.md b/docs/src/man/advection.md similarity index 100% rename from docs/src/advection.md rename to docs/src/man/advection.md diff --git a/docs/src/backend.md b/docs/src/man/backend.md similarity index 100% rename from docs/src/backend.md rename to docs/src/man/backend.md diff --git a/docs/src/equations.md b/docs/src/man/equations.md similarity index 100% rename from docs/src/equations.md rename to docs/src/man/equations.md diff --git a/docs/src/man/installation.md b/docs/src/man/installation.md new file mode 100644 index 00000000..374e10f4 --- /dev/null +++ b/docs/src/man/installation.md @@ -0,0 +1,118 @@ +# Installation instructions + +JustRelax.jl is written in the [julia](https://julialang.org) programming language, which is an extremely powerful, modern, scientific computing language. Julia works on all major operating systems, is free, fast, and has a very active user basis (with many useful packages). In case you haven't heard about julia yet, you are not alone. Yet, perhaps a look at [this](https://www.nature.com/articles/d41586-019-02310-3) or [this](https://thenextweb.com/news/watch-out-python-julia-programming-coding-language-coming-for-crown-syndication) article, which explains nicely why it has an enormous potential for computational geosciences as well. + +### 1. Install julia +In order to use then package you need to install julia. We recommend downloading from the [julia](https://julialang.org) webpage and follow the instructions . + + +### 2. Install Visual Studio Code +The julia files itself are text files (just like matlab scripts). You may want to edit or modify them at some stage, for which you can use any text editor for that. We prefer to use the freely available [Visual Studio Code](https://code.visualstudio.com) as it has a build-in terminal and is the comes with the (official) julia debugger (install the Julia extension for that). + +### 3. Getting started with julia +You start julia on the command line with: +``` +usr$ julia +``` +This will start the command-line interface of julia: +```julia + _ + _ _ _(_)_ | Documentation: https://docs.julialang.org + (_) | (_) (_) | + _ _ _| |_ __ _ | Type "?" for help, "]?" for Pkg help. + | | | | | | |/ _` | | + | | |_| | | | (_| | | Version 1.10.2 (2024-03-01) + _/ |\__'_|_|_|\__'_| | Official https://julialang.org/ release +|__/ | + +julia> + +``` + +From the julia prompt, you start the package manager by typing `]`: +```julia +(@v1.10) pkg> +``` +And you return to the command line with a backspace. + +Also useful is that julia has a build-in terminal, which you can reach by typing `;` on the command line: +```julia +julia>; +shell> +``` +In the shell, you can use the normal commands like listing the content of a directory, or the current path: +```julia +shell> ls +LICENSE Manifest.toml Project.toml README.md docs src test tutorial +shell> pwd +/home/usr/Documents/JustRelax.jl +``` +As before, return to the main command line (called `REPL`) with a backspace. + +If you want to see help information for any julia function, type `?` followed by the command. +An example for `Array` is: +```julia +help?> Array +search: Array SubArray BitArray DenseArray StridedArray PermutedDimsArray AbstractArray AbstractRange AbstractIrrational + + Array{T,N} <: AbstractArray{T,N} + + + N-dimensional dense array with elements of type T. + + ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── + + Array{T}(undef, dims) + Array{T,N}(undef, dims) + + + Construct an uninitialized N-dimensional Array containing elements of type T. N can either be supplied explicitly, as + in Array{T,N}(undef, dims), or be determined by the length or number of dims. dims may be a tuple or a series of + integer arguments corresponding to the lengths in each dimension. If the rank N is supplied explicitly, then it must + match the length or number of dims. Here undef is the UndefInitializer. + + Examples + ≡≡≡≡≡≡≡≡ + + julia> A = Array{Float64, 2}(undef, 2, 3) # N given explicitly + 2×3 Matrix{Float64}: + 6.90198e-310 6.90198e-310 6.90198e-310 + 6.90198e-310 6.90198e-310 0.0 + + [...] +``` + +If you are in a directory that has a julia file (which have the extension `*.jl`), you can open that file with Visual Studio Code: +```julia +shell> code runtests.jl +``` +Execute the file with: +```julia +julia> include("runtests") +``` +Note that you do not include the `*.jl` extension. Another option is to run the code line by line through the `Julia REPL` via `Shift+Enter`, this is especially useful for debugging purposes. + + +### 4. Install JustRelax.jl + +`JustRelax.jl` is a registered package and can be added as follows: + +```julia +using Pkg; Pkg.add("JustRelax") +``` +or + +```julia +julia> ] +(@v1.10) pkg> add JustRelax +``` + +However, as the API is changing and not every feature leads to a new release, one can also do `add JustRelax#main` which will clone the main branch of the repository. +After installation, you can test the package by running the following commands: + +```julia +using JustRelax +julia> ] + pkg> test JustRelax +``` +The test will take a while, so grab a :coffee: or :tea: diff --git a/docs/src/man/listfunctions.md b/docs/src/man/listfunctions.md new file mode 100644 index 00000000..181a8a49 --- /dev/null +++ b/docs/src/man/listfunctions.md @@ -0,0 +1,6 @@ +# List of all functions + +Here an overview of all functions: +```@autodocs +Modules = [JustRelax] +``` diff --git a/docs/src/material_physics.md b/docs/src/man/material_physics.md similarity index 100% rename from docs/src/material_physics.md rename to docs/src/man/material_physics.md diff --git a/src/thermal_diffusion/Shearheating.jl b/src/thermal_diffusion/ShearHeating.jl similarity index 100% rename from src/thermal_diffusion/Shearheating.jl rename to src/thermal_diffusion/ShearHeating.jl From 90f185679ac6bf12562e8e3ff20c34adde761ff2 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 15:54:04 +0200 Subject: [PATCH 37/91] remove Julia installation from the docs --- docs/src/man/installation.md | 97 +----------------------------------- 1 file changed, 1 insertion(+), 96 deletions(-) diff --git a/docs/src/man/installation.md b/docs/src/man/installation.md index 374e10f4..330e9090 100644 --- a/docs/src/man/installation.md +++ b/docs/src/man/installation.md @@ -1,99 +1,4 @@ -# Installation instructions - -JustRelax.jl is written in the [julia](https://julialang.org) programming language, which is an extremely powerful, modern, scientific computing language. Julia works on all major operating systems, is free, fast, and has a very active user basis (with many useful packages). In case you haven't heard about julia yet, you are not alone. Yet, perhaps a look at [this](https://www.nature.com/articles/d41586-019-02310-3) or [this](https://thenextweb.com/news/watch-out-python-julia-programming-coding-language-coming-for-crown-syndication) article, which explains nicely why it has an enormous potential for computational geosciences as well. - -### 1. Install julia -In order to use then package you need to install julia. We recommend downloading from the [julia](https://julialang.org) webpage and follow the instructions . - - -### 2. Install Visual Studio Code -The julia files itself are text files (just like matlab scripts). You may want to edit or modify them at some stage, for which you can use any text editor for that. We prefer to use the freely available [Visual Studio Code](https://code.visualstudio.com) as it has a build-in terminal and is the comes with the (official) julia debugger (install the Julia extension for that). - -### 3. Getting started with julia -You start julia on the command line with: -``` -usr$ julia -``` -This will start the command-line interface of julia: -```julia - _ - _ _ _(_)_ | Documentation: https://docs.julialang.org - (_) | (_) (_) | - _ _ _| |_ __ _ | Type "?" for help, "]?" for Pkg help. - | | | | | | |/ _` | | - | | |_| | | | (_| | | Version 1.10.2 (2024-03-01) - _/ |\__'_|_|_|\__'_| | Official https://julialang.org/ release -|__/ | - -julia> - -``` - -From the julia prompt, you start the package manager by typing `]`: -```julia -(@v1.10) pkg> -``` -And you return to the command line with a backspace. - -Also useful is that julia has a build-in terminal, which you can reach by typing `;` on the command line: -```julia -julia>; -shell> -``` -In the shell, you can use the normal commands like listing the content of a directory, or the current path: -```julia -shell> ls -LICENSE Manifest.toml Project.toml README.md docs src test tutorial -shell> pwd -/home/usr/Documents/JustRelax.jl -``` -As before, return to the main command line (called `REPL`) with a backspace. - -If you want to see help information for any julia function, type `?` followed by the command. -An example for `Array` is: -```julia -help?> Array -search: Array SubArray BitArray DenseArray StridedArray PermutedDimsArray AbstractArray AbstractRange AbstractIrrational - - Array{T,N} <: AbstractArray{T,N} - - - N-dimensional dense array with elements of type T. - - ─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── - - Array{T}(undef, dims) - Array{T,N}(undef, dims) - - - Construct an uninitialized N-dimensional Array containing elements of type T. N can either be supplied explicitly, as - in Array{T,N}(undef, dims), or be determined by the length or number of dims. dims may be a tuple or a series of - integer arguments corresponding to the lengths in each dimension. If the rank N is supplied explicitly, then it must - match the length or number of dims. Here undef is the UndefInitializer. - - Examples - ≡≡≡≡≡≡≡≡ - - julia> A = Array{Float64, 2}(undef, 2, 3) # N given explicitly - 2×3 Matrix{Float64}: - 6.90198e-310 6.90198e-310 6.90198e-310 - 6.90198e-310 6.90198e-310 0.0 - - [...] -``` - -If you are in a directory that has a julia file (which have the extension `*.jl`), you can open that file with Visual Studio Code: -```julia -shell> code runtests.jl -``` -Execute the file with: -```julia -julia> include("runtests") -``` -Note that you do not include the `*.jl` extension. Another option is to run the code line by line through the `Julia REPL` via `Shift+Enter`, this is especially useful for debugging purposes. - - -### 4. Install JustRelax.jl +# Installation `JustRelax.jl` is a registered package and can be added as follows: From 7c8e5624d935e9372fce9587dd2f765e0a9ef6ba Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 15:55:40 +0200 Subject: [PATCH 38/91] remove JustPIC from dependency check --- .github/workflows/Dependency.yml | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/.github/workflows/Dependency.yml b/.github/workflows/Dependency.yml index 3ef7599e..3d20a637 100644 --- a/.github/workflows/Dependency.yml +++ b/.github/workflows/Dependency.yml @@ -15,8 +15,4 @@ jobs: echo "GLMakie dependency found, failing the test." exit 1 fi - if grep -q "JustPIC" ./Project.toml; then - echo "JustPIC dependency found, failing the test." - exit 1 - fi - echo "Neither GLMakie nor JustPIC dependencies found." + echo "Neither GLMakie dependencies found." From eb1a05ef66e0ff7c531733b1f36dd64966f21cd0 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 15:58:31 +0200 Subject: [PATCH 39/91] JP compat --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 6bf9d11d..48d96b3f 100644 --- a/Project.toml +++ b/Project.toml @@ -35,6 +35,7 @@ CellArrays = "0.2" GeoParams = "0.5" HDF5 = "0.17.1" ImplicitGlobalGrid = "0.15.0" +JustPIC = "0.3.2" MPI = "0.20" MuladdMacro = "0.2" ParallelStencil = "0.12" From eea044ea4532d16184930d45e39b187d2f0d5d2d Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 16:02:44 +0200 Subject: [PATCH 40/91] bump JP test compat --- test/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index c1444d4f..7509cf6e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -10,7 +10,7 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] -JustPIC = "0.2" +JustPIC = "0.3.2" ParallelStencil = "0.11, 0.12" CellArrays = "0.2" StaticArrays = "1" From 7795a2a18bf882fd544ce93c3115b7276fbc3c8a Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 29 Apr 2024 16:32:11 +0200 Subject: [PATCH 41/91] switch to `julia-actions/setup-julia@v2` and switch to julia 1.10 --- .github/workflows/Documenter.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/Documenter.yml b/.github/workflows/Documenter.yml index 82f4f36b..dc10a7ea 100644 --- a/.github/workflows/Documenter.yml +++ b/.github/workflows/Documenter.yml @@ -34,7 +34,7 @@ jobs: - uses: julia-actions/setup-julia@v1 with: version: '1.9' - - uses: julia-actions/cache@v1 + - uses: julia-actions/cache@v2 - name: Install dependencies run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' - name: Build and deploy From 5ff65d1fb655990319df07a9cceaa4e0ce3115c1 Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 29 Apr 2024 16:33:03 +0200 Subject: [PATCH 42/91] fix wrong commit --- .github/workflows/Documenter.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/Documenter.yml b/.github/workflows/Documenter.yml index dc10a7ea..80eca36f 100644 --- a/.github/workflows/Documenter.yml +++ b/.github/workflows/Documenter.yml @@ -31,10 +31,10 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v1 + - uses: julia-actions/setup-julia@v2 with: - version: '1.9' - - uses: julia-actions/cache@v2 + version: '1.10' + - uses: julia-actions/cache@v1 - name: Install dependencies run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' - name: Build and deploy From edacd6f673a60bf1c8538f3c908ea5e145d3945a Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 17:00:32 +0200 Subject: [PATCH 43/91] fully working test suit --- .../benchmarks/stokes2D/solcx/vizSolCx.jl | 12 ++-- .../benchmarks/stokes2D/solkz/vizSolKz.jl | 12 ++-- miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl | 6 +- .../benchmarks/stokes3D/solvi/vizSolVi3D.jl | 2 +- test/Project.toml | 6 +- test/test_arrays_conversions.jl | 58 +++++++++---------- test/test_boundary_conditions2D.jl | 2 +- test/test_diffusion3D_multiphase.jl | 25 ++++---- test/test_shearband2D.jl | 4 +- test/test_shearband2D_softening.jl | 4 +- test/test_shearheating2D.jl | 4 +- test/test_sinking_block.jl | 2 +- test/test_stokes_solvi3D.jl | 2 +- test/{test_traits.jl.jl => test_traits.jl} | 0 14 files changed, 74 insertions(+), 65 deletions(-) rename test/{test_traits.jl.jl => test_traits.jl} (100%) diff --git a/miniapps/benchmarks/stokes2D/solcx/vizSolCx.jl b/miniapps/benchmarks/stokes2D/solcx/vizSolCx.jl index 7e8583c2..a3951e08 100644 --- a/miniapps/benchmarks/stokes2D/solcx/vizSolCx.jl +++ b/miniapps/benchmarks/stokes2D/solcx/vizSolCx.jl @@ -30,19 +30,19 @@ function solCx_solution(geometry; η_left=1, η_right=1e6) return (vx=vxs, vy=vys, p=ps) end -function solcx_error(geometry, stokes::StokesArrays; order=2) +function solcx_error(geometry, stokes; order=2) Li(A, B; order=2) = norm(A .- B, order) solk = solCx_solution(geometry) gridsize = reduce(*, geometry.di) - L2_vx = Li(stokes.V.Vx[:, 2:(end - 1)], PTArray(solk.vx); order=order) * gridsize - L2_vy = Li(stokes.V.Vy[2:(end - 1), :], PTArray(solk.vy); order=order) * gridsize - L2_p = Li(stokes.P, PTArray(solk.p); order=order) * gridsize + L2_vx = Li(stokes.V.Vx[:, 2:(end - 1)], PTArray(backend)(solk.vx); order=order) * gridsize + L2_vy = Li(stokes.V.Vy[2:(end - 1), :], PTArray(backend)(solk.vy); order=order) * gridsize + L2_p = Li(stokes.P, PTArray(backend)(solk.p); order=order) * gridsize return L2_vx, L2_vy, L2_p end -function plot_solCx(geometry, stokes::StokesArrays, ρ; cmap=:vik, fun=heatmap!) +function plot_solCx(geometry, stokes, ρ; cmap=:vik, fun=heatmap!) f = Figure(; size=(3000, 1800), fontsize=28) #Density @@ -79,7 +79,7 @@ function plot_solCx(geometry, stokes::StokesArrays, ρ; cmap=:vik, fun=heatmap!) return f end -function plot_solCx_error(geometry, stokes::StokesArrays, Δη; cmap=:vik) +function plot_solCx_error(geometry, stokes, Δη; cmap=:vik) solc = solCx_solution(geometry; η_right=Δη) # Plot diff --git a/miniapps/benchmarks/stokes2D/solkz/vizSolKz.jl b/miniapps/benchmarks/stokes2D/solkz/vizSolKz.jl index 65e03f6e..a53a7f65 100644 --- a/miniapps/benchmarks/stokes2D/solkz/vizSolKz.jl +++ b/miniapps/benchmarks/stokes2D/solkz/vizSolKz.jl @@ -32,20 +32,20 @@ function solkz_solution(geometry) return (vx=(vxs), vy=(vys), p=(ps)) end -function Li_error(geometry, stokes::StokesArrays; order=2) +function Li_error(geometry, stokes; order=2) solk = solkz_solution(geometry) gridsize = reduce(*, geometry.di) Li(A, B; order=2) = norm(A .- B, order) - L2_vx = Li(stokes.V.Vx[:, 2:(end - 1)], PTArray(solk.vx); order=order) * gridsize - L2_vy = Li(stokes.V.Vy[2:(end - 1), :], PTArray(solk.vy); order=order) * gridsize - L2_p = Li(stokes.P, PTArray(solk.p); order=order) * gridsize + L2_vx = Li(stokes.V.Vx[:, 2:(end - 1)], PTArray(backend)(solk.vx); order=order) * gridsize + L2_vy = Li(stokes.V.Vy[2:(end - 1), :], PTArray(backend)(solk.vy); order=order) * gridsize + L2_p = Li(stokes.P, PTArray(backend)(solk.p); order=order) * gridsize return L2_vx, L2_vy, L2_p end -function plot_solkz(geometry, ρ, stokes::StokesArrays; cmap=:vik) +function plot_solkz(geometry, ρ, stokes; cmap=:vik) f = Figure(; size=(3000, 1800), fontsize=28) #Ddensity @@ -86,7 +86,7 @@ function plot_solkz(geometry, ρ, stokes::StokesArrays; cmap=:vik) return f end -function plot_solKz_error(geometry, stokes::StokesArrays; cmap=:vik) +function plot_solKz_error(geometry, stokes; cmap=:vik) solk = solkz_solution(geometry) # Plot diff --git a/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl b/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl index e3ce39f0..501cd989 100644 --- a/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl +++ b/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl @@ -4,7 +4,7 @@ using ParallelStencil.FiniteDifferences3D # D. W. Schmid and Y. Y. Podladchikov. Analytical solutions for deformable elliptical inclusions in # general shear. Geophysical Journal International, 155(1):269–288, 2003. -# include("vizSolVi3D.jl") +include("vizSolVi3D.jl") @parallel function smooth!(A2::AbstractArray{T,3}, A::AbstractArray{T,3}, fact::T) where {T} @inn(A2) = @inn(A) + one(T) / 6.1 / fact * (@d2_xi(A) + @d2_yi(A) + @d2_zi(A)) @@ -92,7 +92,7 @@ function solVi3D(; ## Boundary conditions pureshear_bc!(stokes, xci, xvi, εbg) flow_bcs = FlowBoundaryConditions(; - free_slip = (left=false, right=false, top=false, bot=false, back=false, front=false), + free_slip = (left=true, right=true, top=true, bot=true, back=true, front=true), no_slip = (left=false, right=false, top=false, bot=false, back=false, front=false), ) flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -127,4 +127,4 @@ function solVi3D(; finalize_global_grid(; finalize_MPI=finalize_MPI) return (ni=ni, xci=xci, xvi=xvi, li=li, di=di), stokes, iters -end +end \ No newline at end of file diff --git a/miniapps/benchmarks/stokes3D/solvi/vizSolVi3D.jl b/miniapps/benchmarks/stokes3D/solvi/vizSolVi3D.jl index 1e33c1d9..ad1b8f68 100644 --- a/miniapps/benchmarks/stokes3D/solvi/vizSolVi3D.jl +++ b/miniapps/benchmarks/stokes3D/solvi/vizSolVi3D.jl @@ -1,4 +1,4 @@ -function plot(stokes::StokesArrays, geometry, rc; cmap=:vik) +function plot(stokes, geometry, rc; cmap=:vik) xci, xvi = geometry.xci, geometry.xvi cx, cy = (geometry.xvi[1][end] - geometry.xvi[1][1]) / 2, diff --git a/test/Project.toml b/test/Project.toml index 7509cf6e..5bba385d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,8 +1,9 @@ [deps] AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" +CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" +GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" ParallelStencil = "94395366-693c-11ea-3b26-d9b7aac5d958" -CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" @@ -10,9 +11,10 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] +CellArrays = "0.2" +GeoParams = "0.5.7" JustPIC = "0.3.2" ParallelStencil = "0.11, 0.12" -CellArrays = "0.2" StaticArrays = "1" Statistics = "1" julia = "1.9, 1.10" diff --git a/test/test_arrays_conversions.jl b/test/test_arrays_conversions.jl index 1973517a..012e2eaf 100644 --- a/test/test_arrays_conversions.jl +++ b/test/test_arrays_conversions.jl @@ -1,35 +1,35 @@ -using JustRelax, JustRelax.JustRelax2D, Test -const bk = JustRelax.backend +# using JustRelax, JustRelax.JustRelax2D, Test +# const bk = JustRelax.backend -# @testset "Array conversions" begin - ni = 2, 2 - stokes = StokesArrays(CPUBackend, ni) - thermal = ThermalArrays(CPUBackend, ni) +# # @testset "Array conversions" begin +# ni = 2, 2 +# stokes = StokesArrays(CPUBackend, ni) +# thermal = ThermalArrays(CPUBackend, ni) - # @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} - # @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} - # @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} - # @test Array(stokes.P) isa Array{T, N} where {T, N} - # @test Array(stokes) isa StokesArrays - # @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} +# # @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} +# # @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} +# # @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} +# # @test Array(stokes.P) isa Array{T, N} where {T, N} +# # @test Array(stokes) isa StokesArrays +# # @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} - # test generic arrays - @test bk(Array) === CPUBackendTrait() - @test bk(Matrix) === CPUBackendTrait() - @test bk(Vector) === CPUBackendTrait() - @test bk(rand(2)) === CPUBackendTrait() - @test bk(rand(2,2)) === CPUBackendTrait() - @test bk(rand(2,2,2)) === CPUBackendTrait() - @test_throws ArgumentError backend(rand()) +# # test generic arrays +# @test bk(Array) === CPUBackendTrait() +# @test bk(Matrix) === CPUBackendTrait() +# @test bk(Vector) === CPUBackendTrait() +# @test bk(rand(2)) === CPUBackendTrait() +# @test bk(rand(2,2)) === CPUBackendTrait() +# @test bk(rand(2,2,2)) === CPUBackendTrait() +# @test_throws ArgumentError backend(rand()) - # test JR structs - @test bk(stokes.V) === CPUBackendTrait() - @test bk(stokes.τ) === CPUBackendTrait() - @test bk(stokes.R) === CPUBackendTrait() - @test bk(stokes.P) === CPUBackendTrait() - @test bk(stokes) === CPUBackendTrait() - @test bk(thermal) === CPUBackendTrait() - @test_throws ArgumentError bk("potato") -# end +# # test JR structs +# @test bk(stokes.V) === CPUBackendTrait() +# @test bk(stokes.τ) === CPUBackendTrait() +# @test bk(stokes.R) === CPUBackendTrait() +# @test bk(stokes.P) === CPUBackendTrait() +# @test bk(stokes) === CPUBackendTrait() +# @test bk(thermal) === CPUBackendTrait() +# @test_throws ArgumentError bk("potato") +# # end diff --git a/test/test_boundary_conditions2D.jl b/test/test_boundary_conditions2D.jl index 86326672..14e9d93d 100644 --- a/test/test_boundary_conditions2D.jl +++ b/test/test_boundary_conditions2D.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D -using Test#, Suppressor +using Test, Suppressor const backend = CPUBackend @testset "Boundary Conditions" begin diff --git a/test/test_diffusion3D_multiphase.jl b/test/test_diffusion3D_multiphase.jl index 33ca4354..c31f9ebd 100644 --- a/test/test_diffusion3D_multiphase.jl +++ b/test/test_diffusion3D_multiphase.jl @@ -6,12 +6,15 @@ using JustRelax using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + + 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend import JustRelax.@cell @@ -117,7 +120,7 @@ function diffusion_3D(; ## Allocate arrays needed for every Thermal Diffusion # general thermal arrays - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal.H .= 1e-6 # physical parameters ρ = @fill(ρ0, ni...) @@ -145,14 +148,14 @@ function diffusion_3D(; ) # temperature pPhases, = init_cell_arrays(particles, Val(1)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, center_perturbation..., r) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) JustRelax.phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # PT coefficients for thermal diffusion args = (; P=P, T=thermal.Tc) - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li; CFL = 0.75 / √3.1) + pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li; CFL = 0.75 / √3.1) t = 0.0 it = 0 @@ -168,11 +171,13 @@ function diffusion_3D(; args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = false, + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = false, + ) ) t += dt diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index ed11cd21..56879450 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -112,8 +112,8 @@ function ShearBand2D() 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]]) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-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) diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl index ac4c919a..0b6390a4 100644 --- a/test/test_shearband2D_softening.jl +++ b/test/test_shearband2D_softening.jl @@ -120,8 +120,8 @@ function ShearBand2D() 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]]) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-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) diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 0192c9c6..9d60518b 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -1,3 +1,5 @@ +using Test, Suppressor + # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax2D @@ -14,7 +16,7 @@ const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBac # const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, GeoParams +using GeoParams # Load file with all the rheology configurations include("../miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl") diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index a5d11804..bc8bce16 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -2,7 +2,7 @@ push!(LOAD_PATH, "..") using Test, Suppressor using JustRelax, JustRelax.JustRelax2D -using ParallelStencil +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) const backend_JR = CPUBackend diff --git a/test/test_stokes_solvi3D.jl b/test/test_stokes_solvi3D.jl index 23e7860f..170eda32 100644 --- a/test/test_stokes_solvi3D.jl +++ b/test/test_stokes_solvi3D.jl @@ -36,7 +36,7 @@ function check_convergence_case1() ) tol = 1e-8 - passed = iters.err_evo1[end] < tol + passed = iters.norm_Rx[end] < tol return passed end diff --git a/test/test_traits.jl.jl b/test/test_traits.jl similarity index 100% rename from test/test_traits.jl.jl rename to test/test_traits.jl From 4ebff5a3bbd221770356df120d235c6bdd2e1dd1 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Mon, 29 Apr 2024 17:02:42 +0200 Subject: [PATCH 44/91] last touch --- miniapps/benchmarks/stokes2D/solkz/SolKz.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/solkz/SolKz.jl b/miniapps/benchmarks/stokes2D/solkz/SolKz.jl index 9047d85a..c10a2044 100644 --- a/miniapps/benchmarks/stokes2D/solkz/SolKz.jl +++ b/miniapps/benchmarks/stokes2D/solkz/SolKz.jl @@ -68,8 +68,8 @@ function solKz(; ## Allocate arrays needed for every Stokes problem # general stokes arrays + stokes = StokesArrays(backend, ni) (; η) = stokes.viscosity - stokes = StokesArrays(backend, ni, ViscoElastic) # general numerical coeffs for PT stokes pt_stokes = PTStokesCoeffs(li, di; Re = 5π, CFL = 1 / √2.1) @@ -107,7 +107,7 @@ function solKz(; iterMax = 150e3, nout = 1e3, b_width = (4, 4, 1), - verbose = false + verbose = true ), ) t += Δt From 9a233067ec2ed91bd6863da02060ee594a151176 Mon Sep 17 00:00:00 2001 From: aelligp Date: Mon, 29 Apr 2024 17:27:28 +0200 Subject: [PATCH 45/91] testing deployment --- docs/Project.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/Project.toml b/docs/Project.toml index 1814eb33..21a73358 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,10 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" +JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" +JustRelax = "34418575-392d-4e26-8c6d-96b0910afa06" [compat] Documenter = "1" +JustPIC = "0.3.2" +JustRelax = "0.1.2" From ccde489224fd21e8d289dad4c7c2d3a6e58cf8d5 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 22:51:03 +0200 Subject: [PATCH 46/91] array conversions & unit test --- src/JustRelax.jl | 4 +++ src/array_conversions.jl | 14 +++++----- src/types/phases.jl | 10 -------- test/runtests.jl | 2 +- test/test_arrays_conversions.jl | 45 ++++++++++----------------------- 5 files changed, 26 insertions(+), 49 deletions(-) diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 456d0948..0f0e6f23 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -44,4 +44,8 @@ include("JustRelax_CPU.jl") include("IO/DataIO.jl") +include("array_conversions.jl") +export Array + + end # module diff --git a/src/array_conversions.jl b/src/array_conversions.jl index b09501d5..cf97a8d8 100644 --- a/src/array_conversions.jl +++ b/src/array_conversions.jl @@ -1,18 +1,20 @@ +import Base: Array + ## Conversion of structs to CPU @inline remove_parameters(::T) where {T} = Base.typename(T).wrapper function Array( x::T -) where {T<:Union{StokesArrays,SymmetricTensor,ThermalArrays,Velocity,Residual}} - return Array(iscpu(x), x) +) where {T<:Union{JustRelax.StokesArrays,JustRelax.SymmetricTensor,JustRelax.ThermalArrays,JustRelax.Velocity,JustRelax.Residual}} + return Array(backend(x), x) end -Array(::CPUDeviceTrait, x) = x +Array(::CPUBackendTrait, x) = x function Array( - ::NonCPUDeviceTrait, x::T -) where {T<:Union{SymmetricTensor,ThermalArrays,Velocity,Residual}} + ::NonCPUBackendTrait, x::T +) where {T<:Union{JustRelax.SymmetricTensor,JustRelax.ThermalArrays,JustRelax.Velocity,JustRelax.Residual}} nfields = fieldcount(T) cpu_fields = ntuple(Val(nfields)) do i Base.@_inline_meta @@ -22,7 +24,7 @@ function Array( return T_clean(cpu_fields...) end -function Array(::NonCPUDeviceTrait, x::StokesArrays{T,A,B,C,M,nDim}) where {T,A,B,C,M,nDim} +function Array(::NonCPUBackendTrait, x::JustRelax.StokesArrays) nfields = fieldcount(StokesArrays) cpu_fields = ntuple(Val(nfields)) do i Base.@_inline_meta diff --git a/src/types/phases.jl b/src/types/phases.jl index c82a27d4..8e7bd9dd 100644 --- a/src/types/phases.jl +++ b/src/types/phases.jl @@ -5,14 +5,4 @@ struct PhaseRatio{T} PhaseRatio(vertex::T, center::T) where {T<:AbstractArray} = new{T}(vertex, center) end -# function PhaseRatio(ni::NTuple{N,Integer}, num_phases::Integer) where {N} -# center = @fill(0.0, ni..., celldims = (num_phases,)) -# vertex = @fill(0.0, ni .+ 1..., celldims = (num_phases,)) -# # T = typeof(center) -# return PhaseRatio(vertex, center) -# end - @inline PhaseRatio(::Type{CPUBackend}, ni, num_phases) = PhaseRatio(ni, num_phases) - -# backend trait -# @inline backend(x::PhaseRatio) = backend(x.center.data) diff --git a/test/runtests.jl b/test/runtests.jl index a78f0b5a..f7bf478e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,7 +12,7 @@ function runtests() nfail = 0 printstyled("Testing package JustRelax.jl\n"; bold=true, color=:white) - f0 = "test_traits.jl", "test_types.jl" + f0 = ("test_traits.jl", "test_types.jl", "test_arrays_conversions.jl") for f in f0 include(f) end diff --git a/test/test_arrays_conversions.jl b/test/test_arrays_conversions.jl index 012e2eaf..375bbb60 100644 --- a/test/test_arrays_conversions.jl +++ b/test/test_arrays_conversions.jl @@ -1,35 +1,16 @@ -# using JustRelax, JustRelax.JustRelax2D, Test -# const bk = JustRelax.backend +using JustRelax, JustRelax.JustRelax2D, Test +const bk = JustRelax.backend -# # @testset "Array conversions" begin -# ni = 2, 2 -# stokes = StokesArrays(CPUBackend, ni) -# thermal = ThermalArrays(CPUBackend, ni) - -# # @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} -# # @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} -# # @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} -# # @test Array(stokes.P) isa Array{T, N} where {T, N} -# # @test Array(stokes) isa StokesArrays -# # @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} +@testset "Array conversions" begin + ni = 2, 2 + stokes = StokesArrays(CPUBackend, ni) + thermal = ThermalArrays(CPUBackend, ni) - -# # test generic arrays -# @test bk(Array) === CPUBackendTrait() -# @test bk(Matrix) === CPUBackendTrait() -# @test bk(Vector) === CPUBackendTrait() -# @test bk(rand(2)) === CPUBackendTrait() -# @test bk(rand(2,2)) === CPUBackendTrait() -# @test bk(rand(2,2,2)) === CPUBackendTrait() -# @test_throws ArgumentError backend(rand()) - -# # test JR structs -# @test bk(stokes.V) === CPUBackendTrait() -# @test bk(stokes.τ) === CPUBackendTrait() -# @test bk(stokes.R) === CPUBackendTrait() -# @test bk(stokes.P) === CPUBackendTrait() -# @test bk(stokes) === CPUBackendTrait() -# @test bk(thermal) === CPUBackendTrait() -# @test_throws ArgumentError bk("potato") -# # end + @test Array(stokes.V) isa JustRelax.Velocity{Array{T, N}} where {T, N} + @test Array(stokes.τ) isa JustRelax.SymmetricTensor{Array{T, N}} where {T, N} + @test Array(stokes.R) isa JustRelax.Residual{Array{T, N}} where {T, N} + @test Array(stokes.P) isa Array{T, N} where {T, N} + @test Array(stokes) isa JustRelax.StokesArrays + @test Array(thermal) isa JustRelax.ThermalArrays{Array{T, N}} where {T, N} +end From 51437bb10e2e0d49041f96f32633e2817703aa7a Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 22:51:46 +0200 Subject: [PATCH 47/91] format --- src/JustRelax.jl | 1 - src/array_conversions.jl | 19 +++++++- src/boundaryconditions/BoundaryConditions.jl | 6 +-- .../DiffusionPT_coefficients.jl | 47 +++++-------------- src/thermal_diffusion/DiffusionPT_solver.jl | 8 ++-- 5 files changed, 37 insertions(+), 44 deletions(-) diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 0f0e6f23..174f1941 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -47,5 +47,4 @@ include("IO/DataIO.jl") include("array_conversions.jl") export Array - end # module diff --git a/src/array_conversions.jl b/src/array_conversions.jl index cf97a8d8..febb4b12 100644 --- a/src/array_conversions.jl +++ b/src/array_conversions.jl @@ -6,7 +6,15 @@ import Base: Array function Array( x::T -) where {T<:Union{JustRelax.StokesArrays,JustRelax.SymmetricTensor,JustRelax.ThermalArrays,JustRelax.Velocity,JustRelax.Residual}} +) where { + T<:Union{ + JustRelax.StokesArrays, + JustRelax.SymmetricTensor, + JustRelax.ThermalArrays, + JustRelax.Velocity, + JustRelax.Residual, + }, +} return Array(backend(x), x) end @@ -14,7 +22,14 @@ Array(::CPUBackendTrait, x) = x function Array( ::NonCPUBackendTrait, x::T -) where {T<:Union{JustRelax.SymmetricTensor,JustRelax.ThermalArrays,JustRelax.Velocity,JustRelax.Residual}} +) where { + T<:Union{ + JustRelax.SymmetricTensor, + JustRelax.ThermalArrays, + JustRelax.Velocity, + JustRelax.Residual, + }, +} nfields = fieldcount(T) cpu_fields = ntuple(Val(nfields)) do i Base.@_inline_meta diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index ede956e0..88a483ef 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -320,7 +320,7 @@ function pureshear_bc!( stokes::JustRelax.StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg ) where {T} _T = typeof(stokes.V.Vx) - stokes.V.Vx[:, 2:(end - 1)] .= _T([ εbg * x for x in xvi[1], y in xci[2]]) + stokes.V.Vx[:, 2:(end - 1)] .= _T([εbg * x for x in xvi[1], y in xci[2]]) stokes.V.Vy[2:(end - 1), :] .= _T([-εbg * y for x in xci[1], y in xvi[2]]) return nothing @@ -333,8 +333,8 @@ function pureshear_bc!( xc, yc, zc = xci _T = typeof(stokes.V.Vx) - stokes.V.Vx[:, 2:(end - 1), 2:(end - 1)] .= _T([ εbg * x for x in xv, y in yc, z in zc]) - stokes.V.Vy[2:(end - 1), :, 2:(end - 1)] .= _T([ εbg * y for x in xc, y in xv, z in zc]) + stokes.V.Vx[:, 2:(end - 1), 2:(end - 1)] .= _T([εbg * x for x in xv, y in yc, z in zc]) + stokes.V.Vy[2:(end - 1), :, 2:(end - 1)] .= _T([εbg * y for x in xc, y in xv, z in zc]) stokes.V.Vz[2:(end - 1), 2:(end - 1), :] .= _T([-εbg * z for x in xc, y in xc, z in zv]) return nothing diff --git a/src/thermal_diffusion/DiffusionPT_coefficients.jl b/src/thermal_diffusion/DiffusionPT_coefficients.jl index c64725b2..a566ec17 100644 --- a/src/thermal_diffusion/DiffusionPT_coefficients.jl +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -1,10 +1,14 @@ function PTThermalCoeffs( - ::Type{CPUBackend}, K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 + ::Type{CPUBackend}, + K, + ρCp, + dt, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, ) where {nDim,T} - - PTThermalCoeffs( - K, ρCp, dt, di, li; ϵ=ϵ, CFL=CFL - ) + return PTThermalCoeffs(K, ρCp, dt, di, li; ϵ=ϵ, CFL=CFL) end function PTThermalCoeffs( @@ -19,7 +23,6 @@ function PTThermalCoeffs( return JustRelax.PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ) end - function PTThermalCoeffs( ::Type{CPUBackend}, rheology, @@ -32,17 +35,7 @@ function PTThermalCoeffs( ϵ=1e-8, CFL=0.9 / √3, ) where {nDim,T} - PTThermalCoeffs( - rheology, - phase_ratios, - args, - dt, - ni, - di, - li; - ϵ=ϵ, - CFL=CFL, - ) + return PTThermalCoeffs(rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) end function PTThermalCoeffs( @@ -79,27 +72,11 @@ function PTThermalCoeffs( ϵ=1e-8, CFL=0.9 / √3, ) where {nDim,T} - PTThermalCoeffs( - rheology, - args, - dt, - ni, - di, - li; - ϵ=ϵ, - CFL=CFL, - ) + return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) end function PTThermalCoeffs( - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, + rheology, args, dt, ni, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 ) where {nDim,T} Vpdτ = min(di...) * CFL max_lxyz = max(li...) diff --git a/src/thermal_diffusion/DiffusionPT_solver.jl b/src/thermal_diffusion/DiffusionPT_solver.jl index 8b5170fe..bdfbebd5 100644 --- a/src/thermal_diffusion/DiffusionPT_solver.jl +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -2,7 +2,9 @@ function heatdiffusion_PT!(thermal, args...; kwargs) return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs=kwargs) end -heatdiffusion_PT!(::CPUBackendTrait, thermal, args...; kwargs) = _heatdiffusion_PT!(thermal, args...; kwargs...) +function heatdiffusion_PT!(::CPUBackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) +end """ heatdiffusion_PT!(thermal, pt_thermal, K, ρCp, dt, di; iterMax, nout, verbose) @@ -22,7 +24,7 @@ function _heatdiffusion_PT!( iterMax=50e3, nout=1e3, verbose=true, - kwargs... + kwargs..., ) # Compute some constant stuff _dt = inv(dt) @@ -112,7 +114,7 @@ function _heatdiffusion_PT!( iterMax=50e3, nout=1e3, verbose=true, - kwargs... + kwargs..., ) phases = get_phase(phase) From 2d93ec4bcbb905e6fcb3f188b861434ddd4acbc2 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:02:49 +0200 Subject: [PATCH 48/91] Merge branch 'main' into adm/extensions --- Shearheating2D.jl | 271 +++++++++++++++++++++++++++++++++++++ Shearheating_rheology.jl | 77 +++++++++++ test_arrays_conversions.jl | 110 +++++++++++++++ 3 files changed, 458 insertions(+) create mode 100644 Shearheating2D.jl create mode 100644 Shearheating_rheology.jl create mode 100644 test_arrays_conversions.jl diff --git a/Shearheating2D.jl b/Shearheating2D.jl new file mode 100644 index 00000000..73cf8078 --- /dev/null +++ b/Shearheating2D.jl @@ -0,0 +1,271 @@ +# Benchmark of Duretz et al. 2014 +# http://dx.doi.org/10.1002/2014GL060438 +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) + +using JustPIC +using JustPIC._2D +# Threads is the default backend, +# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, +# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +# Load script dependencies +using Printf, LinearAlgebra, GeoParams, CellArrays +# using GLMakie + +# Load file with all the rheology configurations +include("Shearheating_rheology.jl") + +## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- + +function copyinn_x!(A, B) + @parallel function f_x(A, B) + @all(A) = @inn_x(B) + return nothing + end + + @parallel f_x(A, B) +end + +import ParallelStencil.INDICES +const idx_j = INDICES[2] +macro all_j(A) + esc(:($A[$idx_j])) +end + +# Initial pressure profile - not accurate +@parallel function init_P!(P, ρg, z) + @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) + return nothing +end + + +## END OF HELPER FUNCTION ------------------------------------------------------------ + +## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- +function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) + + # Physical domain ------------------------------------ + ly = 40e3 # domain length in y + lx = 70e3 # domain length in x + ni = nx, ny # number of cells + li = lx, ly # domain length in x- and y- + di = @. li / ni # grid step in x- and -y + origin = 0.0, -ly # origin coordinates (15km f sticky air layer) + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells + # ---------------------------------------------------- + + # Physical properties using GeoParams ---------------- + rheology = init_rheologies(; is_TP_Conductivity=false) + κ = (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, max_xcell, min_xcell = 20, 40, 10 + particles = init_particles( + backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... + ) + # velocity grids + grid_vx, grid_vy = velocity_grids(xci, xvi, di) + # temperature + pT, pPhases = init_cell_arrays(particles, Val(3)) + particle_args = (pT, pPhases) + + # Elliptical temperature anomaly + xc_anomaly = lx/2 # origin of thermal anomaly + yc_anomaly = 40e3 # origin of thermal anomaly + r_anomaly = 3e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) + phase_ratios_center(phase_ratios, particles, grid, pPhases) + # ---------------------------------------------------- + + # STOKES --------------------------------------------- + # Allocate arrays needed for every Stokes problem + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) + # ---------------------------------------------------- + + # TEMPERATURE PROFILE -------------------------------- + thermal = ThermalArrays(backend_JR, ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true, right = true, top = false, bot = false), + ) + + # Initialize constant temperature + @views thermal.T .= 273.0 + 400 + thermal_bcs!(thermal.T, thermal_bc) + temperature2center!(thermal) + # ---------------------------------------------------- + + # Buoyancy forces + ρg = @zeros(ni...), @zeros(ni...) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[2], xci[2]) + + # Rheology + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + + # PT coefficients for thermal diffusion + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-1 / √2.1 + ) + + # Boundary conditions + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), + ) + ## Compression and not extension - fix this + εbg = 5e-14 + stokes.V.Vx .= PTArray(backend_JR)([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend_JR)([ (ly - abs(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) + + # # 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) + # # ---------------------------------------------------- + + # # 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) + # 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 + + T_buffer = @zeros(ni.+1) + Told_buffer = similar(T_buffer) + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle!(pT, xvi, T_buffer, particles) + + local Vx_v, Vy_v + if do_vtk + Vx_v = @zeros(ni.+1...) + Vy_v = @zeros(ni.+1...) + end + # Time loop + t, it = 0.0, 0 + while it < 1 + # interpolate fields from particle to grid vertices + particle2grid!(T_buffer, pT, xvi, particles) + @views T_buffer[:, end] .= 273.0 + 400 + @views thermal.T[2:end-1, :] .= T_buffer + temperature2center!(thermal) + + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + ) + # ------------------------------ + + # Stokes solver ---------------- + solve!( + stokes, + pt_stokes, + di, + flow_bcs, + ρg, + phase_ratios, + rheology, + args, + dt, + igg; + kwargs = (; + iterMax = 75e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf) + ) + ) + tensor_invariant!(stokes.ε) + dt = compute_dt(stokes, di, dt_diff) + # ------------------------------ + + compute_shear_heating!( + thermal, + stokes, + phase_ratios, + rheology, # needs to be a tuple + dt, + ) + + # Thermal solver --------------- + heatdiffusion_PT!( + thermal, + pt_thermal, + thermal_bc, + rheology, + args, + dt, + di; + kwargs = (; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) + ) + # ------------------------------ + + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # interpolate fields from grid vertices to particles + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) + # update phase ratios + phase_ratios_center(phase_ratios, particles, grid, pPhases) + + @show it += 1 + # t += dt + end + + return nothing +end + +figdir = "Benchmark_Duretz_etal_2014" +do_vtk = false # set to true to generate VTK files for ParaView +ar = 1 # aspect ratio +n = 64 +nx = n*ar - 2 +ny = n - 2 +igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) +else + igg +end + +@show @time main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) +@show @time main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) \ No newline at end of file diff --git a/Shearheating_rheology.jl b/Shearheating_rheology.jl new file mode 100644 index 00000000..40fb2c8a --- /dev/null +++ b/Shearheating_rheology.jl @@ -0,0 +1,77 @@ +# from Duretz et al. 2014 - http://dx.doi.org/10.1002/2014GL060438 + +function init_rheologies(; is_TP_Conductivity=true) + + # Dislocation creep + Matrix = DislocationCreep(A=3.20e-20, n=3.0, E=276e3, V=0e0, r=0.0, R=8.3145) + Inclusion = DislocationCreep(A=3.16e-26, n=3.3, E=186e3, V=0e0, r=0.0, R=8.3145) + + K_Matrix = if is_TP_Conductivity + # crust + TP_Conductivity(; + a = 1.72, + b = 807e0, + c = 350, + d = 0.0, + ) + else + ConstantConductivity(k = 2.5) + end + + K_Inclusion = if is_TP_Conductivity + TP_Conductivity(; + a = 1.72, + b = 807e0, + c = 350, + d = 0.0, + ) + else + ConstantConductivity(k = 2.5) + end + + # Define rheolgy struct + rheology = ( + # Name = "Matrix", + SetMaterialParams(; + Phase = 1, + Density = ConstantDensity(ρ = 2700), + HeatCapacity = ConstantHeatCapacity(; Cp=1050.0), + Conductivity = K_Matrix, + ShearHeat = ConstantShearheating(1.0NoUnits), + CompositeRheology = CompositeRheology((Matrix, )), + Gravity = ConstantGravity(; g=9.81), + ), + # Name = "LowerCrust", + SetMaterialParams(; + Phase = 2, + Density = ConstantDensity(ρ = 2700), + HeatCapacity = ConstantHeatCapacity(; Cp=1050.0), + Conductivity = K_Inclusion, + ShearHeat = ConstantShearheating(1.0NoUnits), + CompositeRheology = CompositeRheology((Inclusion, )), + ), + ) +end + +function init_phases!(phases, particles, Lx, d, r) + ni = size(phases) + + @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx, d) + @inbounds for ip in JustRelax.cellaxes(phases) + # quick escape + JustRelax.@cell(index[ip, i, j]) == 0 && continue + + x = JustRelax.@cell px[ip, i, j] + depth = -(JustRelax.@cell py[ip, i, j]) + JustRelax.@cell phases[ip, i, j] = 1.0 # matrix + + # thermal anomaly - circular + if ((x - Lx)^2 + (depth - d)^2 ≤ r^2) + JustRelax.@cell phases[ip, i, j] = 2.0 + end + end + return nothing + end + + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, d) +end diff --git a/test_arrays_conversions.jl b/test_arrays_conversions.jl new file mode 100644 index 00000000..60b396e4 --- /dev/null +++ b/test_arrays_conversions.jl @@ -0,0 +1,110 @@ +using JustRelax, Test +model = PS_Setup(:Threads, Float64, 2) +environment!(model) + +# @testset "Array conversions" begin + ni = 10, 10 + stokes = StokesArrays(ni, ViscoElastic) + thermal = ThermalArrays(ni) + + @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} + @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} + @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} + @test Array(stokes.P) isa Array{T, N} where {T, N} + @test Array(stokes) isa StokesArrays + @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} +# end + +# abstract type DeviceTrait end +# struct CPUDeviceTrait <: DeviceTrait end +# struct NonCPUDeviceTrait <: DeviceTrait end + +# @inline iscpu(::Array) = CPUDeviceTrait() +# @inline iscpu(::Type{Array}) = CPUDeviceTrait() +# @inline iscpu(::AbstractArray) = NonCPUDeviceTrait() +# @inline iscpu(::T) where T = throw(ArgumentError("Unkown device")) + +# @inline iscpu(::Velocity{Array{T, N}}) where {T, N} = CPUDeviceTrait() +# @inline iscpu(::Velocity{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() + +# @inline iscpu(::SymmetricTensor{Array{T, N}}) where {T, N} = CPUDeviceTrait() +# @inline iscpu(::SymmetricTensor{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() + +# @inline iscpu(::Residual{Array{T, N}}) where {T, N} = CPUDeviceTrait() +# @inline iscpu(::Residual{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() + +# @inline iscpu(::ThermalArrays{Array{T, N}}) where {T, N} = CPUDeviceTrait() +# @inline iscpu(::ThermalArrays{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() + +# @inline iscpu(::StokesArrays{M,A,B,C,Array{T, N},nDim}) where {M,A,B,C,T,N,nDim} =CPUDeviceTrait() +# @inline iscpu(::StokesArrays{M,A,B,C,AbstractArray{T, N},nDim}) where {M,A,B,C,T,N,nDim} =NonCPUDeviceTrait() + + +@test JustRelax.iscpu(stokes.V) isa JustRelax.CPUDeviceTrait +@test JustRelax.iscpu(stokes.τ) isa JustRelax.CPUDeviceTrait +@test JustRelax.iscpu(stokes.R) isa JustRelax.CPUDeviceTrait +@test JustRelax.iscpu(stokes.P) isa JustRelax.CPUDeviceTrait +@test JustRelax.iscpu(stokes) isa JustRelax.CPUDeviceTrait +@test JustRelax.iscpu(thermal) isa JustRelax.CPUDeviceTrait +@test_throws ArgumentError("Unkown device") JustRelax.iscpu("potato") + +# import Base.Array + +# @inline remove_parameters(::T) where T = Base.typename(T).wrapper + +# function Array(x::T) where T<:Union{SymmetricTensor, ThermalArrays, Velocity, Residual} +# Array(iscpu(x), x) +# end + +# Array(::CPUDeviceTrait, x) = x + +# function Array(::NonCPUDeviceTrait, x::T) where T<:Union{SymmetricTensor, ThermalArrays, Velocity, Residual} +# nfields = fieldcount(T) +# cpu_fields = ntuple(Val(nfields)) do i +# Base.@_inline_meta +# Array(getfield(x, i)) +# end +# T_clean = remove_parameters(x) +# return T_clean(cpu_fields...) +# end + +# function Array(::NonCPUDeviceTrait, x::StokesArrays{T,A,B,C,M,nDim}) where {T,A,B,C,M,nDim} +# nfields = fieldcount(StokesArrays) +# cpu_fields = ntuple(Val(nfields)) do i +# Base.@_inline_meta +# Array(getfield(x, i)) +# end +# T_clean = remove_parameters(x) +# return T_clean(cpu_fields...) +# end + +@edit Array(stokes) + +using Chairmarks + +abstract type AbstractTrait end +struct ConcreteTrait <: AbstractTrait end +foo(::Array) = ConcreteTrait() +foo(::Type{Array{T,N}}) where {T,N} = ConcreteTrait() + +foo1(::Type{Array{Float64,2}}) = ConcreteTrait() + + +T=Array{Float64,2} + +foo2(::Type{T}) = ConcreteTrait() +@b foo2($T) + +foo(::Type{Array{T,N}}) where {T,N} = ConcreteTrait() +@b foo($T) + +foo3(::Type{T}) where T<:Array = ConcreteTrait() + +x=rand(2,2) +T=Array{Float64,2} + +@b foo($T) +@btime foo($T) +@b foo1($T) +@b foo2($T) +@b foo3($T) \ No newline at end of file From b50ebf23d621ffe050dd6b0fe7a9542ea11da3e2 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:03:34 +0200 Subject: [PATCH 49/91] remove files --- Shearheating2D.jl | 271 ------------------------------------- Shearheating_rheology.jl | 77 ----------- test_arrays_conversions.jl | 110 --------------- 3 files changed, 458 deletions(-) delete mode 100644 Shearheating2D.jl delete mode 100644 Shearheating_rheology.jl delete mode 100644 test_arrays_conversions.jl diff --git a/Shearheating2D.jl b/Shearheating2D.jl deleted file mode 100644 index 73cf8078..00000000 --- a/Shearheating2D.jl +++ /dev/null @@ -1,271 +0,0 @@ -# Benchmark of Duretz et al. 2014 -# http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend - -using ParallelStencil, ParallelStencil.FiniteDifferences2D -@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) - -using JustPIC -using JustPIC._2D -# Threads is the default backend, -# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, -# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays -# using GLMakie - -# Load file with all the rheology configurations -include("Shearheating_rheology.jl") - -## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- - -function copyinn_x!(A, B) - @parallel function f_x(A, B) - @all(A) = @inn_x(B) - return nothing - end - - @parallel f_x(A, B) -end - -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -# Initial pressure profile - not accurate -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - - -## END OF HELPER FUNCTION ------------------------------------------------------------ - -## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) - - # Physical domain ------------------------------------ - ly = 40e3 # domain length in y - lx = 70e3 # domain length in x - ni = nx, ny # number of cells - li = lx, ly # domain length in x- and y- - di = @. li / ni # grid step in x- and -y - origin = 0.0, -ly # origin coordinates (15km f sticky air layer) - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Physical properties using GeoParams ---------------- - rheology = init_rheologies(; is_TP_Conductivity=false) - κ = (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, max_xcell, min_xcell = 20, 40, 10 - particles = init_particles( - backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... - ) - # velocity grids - grid_vx, grid_vy = velocity_grids(xci, xvi, di) - # temperature - pT, pPhases = init_cell_arrays(particles, Val(3)) - particle_args = (pT, pPhases) - - # Elliptical temperature anomaly - xc_anomaly = lx/2 # origin of thermal anomaly - yc_anomaly = 40e3 # origin of thermal anomaly - r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) - init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) - phase_ratios_center(phase_ratios, particles, grid, pPhases) - # ---------------------------------------------------- - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend_JR, ni) - thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true, right = true, top = false, bot = false), - ) - - # Initialize constant temperature - @views thermal.T .= 273.0 + 400 - thermal_bcs!(thermal.T, thermal_bc) - temperature2center!(thermal) - # ---------------------------------------------------- - - # Buoyancy forces - ρg = @zeros(ni...), @zeros(ni...) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - - # Rheology - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) - - # PT coefficients for thermal diffusion - pt_thermal = PTThermalCoeffs( - backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-1 / √2.1 - ) - - # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - ) - ## Compression and not extension - fix this - εbg = 5e-14 - stokes.V.Vx .= PTArray(backend_JR)([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray(backend_JR)([ (ly - abs(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) - - # # 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) - # # ---------------------------------------------------- - - # # 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) - # 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 - - T_buffer = @zeros(ni.+1) - Told_buffer = similar(T_buffer) - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - grid2particle!(pT, xvi, T_buffer, particles) - - local Vx_v, Vy_v - if do_vtk - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - end - # Time loop - t, it = 0.0, 0 - while it < 1 - # interpolate fields from particle to grid vertices - particle2grid!(T_buffer, pT, xvi, particles) - @views T_buffer[:, end] .= 273.0 + 400 - @views thermal.T[2:end-1, :] .= T_buffer - temperature2center!(thermal) - - # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) - ) - # ------------------------------ - - # Stokes solver ---------------- - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - phase_ratios, - rheology, - args, - dt, - igg; - kwargs = (; - iterMax = 75e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) - ) - ) - tensor_invariant!(stokes.ε) - dt = compute_dt(stokes, di, dt_diff) - # ------------------------------ - - compute_shear_heating!( - thermal, - stokes, - phase_ratios, - rheology, # needs to be a tuple - dt, - ) - - # Thermal solver --------------- - heatdiffusion_PT!( - thermal, - pt_thermal, - thermal_bc, - rheology, - args, - dt, - di; - kwargs = (; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, - ) - ) - # ------------------------------ - - # Advection -------------------- - # advect particles in space - advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # interpolate fields from grid vertices to particles - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) - # check if we need to inject particles - inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) - # update phase ratios - phase_ratios_center(phase_ratios, particles, grid, pPhases) - - @show it += 1 - # t += dt - end - - return nothing -end - -figdir = "Benchmark_Duretz_etal_2014" -do_vtk = false # set to true to generate VTK files for ParaView -ar = 1 # aspect ratio -n = 64 -nx = n*ar - 2 -ny = n - 2 -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) -else - igg -end - -@show @time main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) -@show @time main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) \ No newline at end of file diff --git a/Shearheating_rheology.jl b/Shearheating_rheology.jl deleted file mode 100644 index 40fb2c8a..00000000 --- a/Shearheating_rheology.jl +++ /dev/null @@ -1,77 +0,0 @@ -# from Duretz et al. 2014 - http://dx.doi.org/10.1002/2014GL060438 - -function init_rheologies(; is_TP_Conductivity=true) - - # Dislocation creep - Matrix = DislocationCreep(A=3.20e-20, n=3.0, E=276e3, V=0e0, r=0.0, R=8.3145) - Inclusion = DislocationCreep(A=3.16e-26, n=3.3, E=186e3, V=0e0, r=0.0, R=8.3145) - - K_Matrix = if is_TP_Conductivity - # crust - TP_Conductivity(; - a = 1.72, - b = 807e0, - c = 350, - d = 0.0, - ) - else - ConstantConductivity(k = 2.5) - end - - K_Inclusion = if is_TP_Conductivity - TP_Conductivity(; - a = 1.72, - b = 807e0, - c = 350, - d = 0.0, - ) - else - ConstantConductivity(k = 2.5) - end - - # Define rheolgy struct - rheology = ( - # Name = "Matrix", - SetMaterialParams(; - Phase = 1, - Density = ConstantDensity(ρ = 2700), - HeatCapacity = ConstantHeatCapacity(; Cp=1050.0), - Conductivity = K_Matrix, - ShearHeat = ConstantShearheating(1.0NoUnits), - CompositeRheology = CompositeRheology((Matrix, )), - Gravity = ConstantGravity(; g=9.81), - ), - # Name = "LowerCrust", - SetMaterialParams(; - Phase = 2, - Density = ConstantDensity(ρ = 2700), - HeatCapacity = ConstantHeatCapacity(; Cp=1050.0), - Conductivity = K_Inclusion, - ShearHeat = ConstantShearheating(1.0NoUnits), - CompositeRheology = CompositeRheology((Inclusion, )), - ), - ) -end - -function init_phases!(phases, particles, Lx, d, r) - ni = size(phases) - - @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx, d) - @inbounds for ip in JustRelax.cellaxes(phases) - # quick escape - JustRelax.@cell(index[ip, i, j]) == 0 && continue - - x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) - JustRelax.@cell phases[ip, i, j] = 1.0 # matrix - - # thermal anomaly - circular - if ((x - Lx)^2 + (depth - d)^2 ≤ r^2) - JustRelax.@cell phases[ip, i, j] = 2.0 - end - end - return nothing - end - - @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, d) -end diff --git a/test_arrays_conversions.jl b/test_arrays_conversions.jl deleted file mode 100644 index 60b396e4..00000000 --- a/test_arrays_conversions.jl +++ /dev/null @@ -1,110 +0,0 @@ -using JustRelax, Test -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - -# @testset "Array conversions" begin - ni = 10, 10 - stokes = StokesArrays(ni, ViscoElastic) - thermal = ThermalArrays(ni) - - @test Array(stokes.V) isa Velocity{Array{T, N}} where {T, N} - @test Array(stokes.τ) isa SymmetricTensor{Array{T, N}} where {T, N} - @test Array(stokes.R) isa Residual{Array{T, N}} where {T, N} - @test Array(stokes.P) isa Array{T, N} where {T, N} - @test Array(stokes) isa StokesArrays - @test Array(thermal) isa ThermalArrays{Array{T, N}} where {T, N} -# end - -# abstract type DeviceTrait end -# struct CPUDeviceTrait <: DeviceTrait end -# struct NonCPUDeviceTrait <: DeviceTrait end - -# @inline iscpu(::Array) = CPUDeviceTrait() -# @inline iscpu(::Type{Array}) = CPUDeviceTrait() -# @inline iscpu(::AbstractArray) = NonCPUDeviceTrait() -# @inline iscpu(::T) where T = throw(ArgumentError("Unkown device")) - -# @inline iscpu(::Velocity{Array{T, N}}) where {T, N} = CPUDeviceTrait() -# @inline iscpu(::Velocity{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() - -# @inline iscpu(::SymmetricTensor{Array{T, N}}) where {T, N} = CPUDeviceTrait() -# @inline iscpu(::SymmetricTensor{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() - -# @inline iscpu(::Residual{Array{T, N}}) where {T, N} = CPUDeviceTrait() -# @inline iscpu(::Residual{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() - -# @inline iscpu(::ThermalArrays{Array{T, N}}) where {T, N} = CPUDeviceTrait() -# @inline iscpu(::ThermalArrays{AbstractArray{T, N}}) where {T, N} = NonCPUDeviceTrait() - -# @inline iscpu(::StokesArrays{M,A,B,C,Array{T, N},nDim}) where {M,A,B,C,T,N,nDim} =CPUDeviceTrait() -# @inline iscpu(::StokesArrays{M,A,B,C,AbstractArray{T, N},nDim}) where {M,A,B,C,T,N,nDim} =NonCPUDeviceTrait() - - -@test JustRelax.iscpu(stokes.V) isa JustRelax.CPUDeviceTrait -@test JustRelax.iscpu(stokes.τ) isa JustRelax.CPUDeviceTrait -@test JustRelax.iscpu(stokes.R) isa JustRelax.CPUDeviceTrait -@test JustRelax.iscpu(stokes.P) isa JustRelax.CPUDeviceTrait -@test JustRelax.iscpu(stokes) isa JustRelax.CPUDeviceTrait -@test JustRelax.iscpu(thermal) isa JustRelax.CPUDeviceTrait -@test_throws ArgumentError("Unkown device") JustRelax.iscpu("potato") - -# import Base.Array - -# @inline remove_parameters(::T) where T = Base.typename(T).wrapper - -# function Array(x::T) where T<:Union{SymmetricTensor, ThermalArrays, Velocity, Residual} -# Array(iscpu(x), x) -# end - -# Array(::CPUDeviceTrait, x) = x - -# function Array(::NonCPUDeviceTrait, x::T) where T<:Union{SymmetricTensor, ThermalArrays, Velocity, Residual} -# nfields = fieldcount(T) -# cpu_fields = ntuple(Val(nfields)) do i -# Base.@_inline_meta -# Array(getfield(x, i)) -# end -# T_clean = remove_parameters(x) -# return T_clean(cpu_fields...) -# end - -# function Array(::NonCPUDeviceTrait, x::StokesArrays{T,A,B,C,M,nDim}) where {T,A,B,C,M,nDim} -# nfields = fieldcount(StokesArrays) -# cpu_fields = ntuple(Val(nfields)) do i -# Base.@_inline_meta -# Array(getfield(x, i)) -# end -# T_clean = remove_parameters(x) -# return T_clean(cpu_fields...) -# end - -@edit Array(stokes) - -using Chairmarks - -abstract type AbstractTrait end -struct ConcreteTrait <: AbstractTrait end -foo(::Array) = ConcreteTrait() -foo(::Type{Array{T,N}}) where {T,N} = ConcreteTrait() - -foo1(::Type{Array{Float64,2}}) = ConcreteTrait() - - -T=Array{Float64,2} - -foo2(::Type{T}) = ConcreteTrait() -@b foo2($T) - -foo(::Type{Array{T,N}}) where {T,N} = ConcreteTrait() -@b foo($T) - -foo3(::Type{T}) where T<:Array = ConcreteTrait() - -x=rand(2,2) -T=Array{Float64,2} - -@b foo($T) -@btime foo($T) -@b foo1($T) -@b foo2($T) -@b foo3($T) \ No newline at end of file From f687aab859db864cd1f6303f808b366d941430e2 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:07:54 +0200 Subject: [PATCH 50/91] delete files --- src/stokes/MetaStokes.jl | 267 ------------------------- src/thermal_diffusion/MetaDiffusion.jl | 236 ---------------------- 2 files changed, 503 deletions(-) delete mode 100644 src/stokes/MetaStokes.jl delete mode 100644 src/thermal_diffusion/MetaDiffusion.jl diff --git a/src/stokes/MetaStokes.jl b/src/stokes/MetaStokes.jl deleted file mode 100644 index 609cac0d..00000000 --- a/src/stokes/MetaStokes.jl +++ /dev/null @@ -1,267 +0,0 @@ -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 - - $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...) - end - 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 - - $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...) - end - 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 - - $(name)(args::Vararg{T,N}) where {T<:AbstractArray,N} = new{T}(args...) - end - end -end - -function make_stokes_struct!(; name::Symbol=:StokesArrays) - @eval begin - struct StokesArrays{M<:AbstractStokesModel,A,B,C,T,nDim} - P::T - P0::T - V::A - ∇V::T - τ::B - ε::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) - ε_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(P),2}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),2}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),3}( - P, P0, V, ∇V, τ, ε, ε_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) - ε_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(P),3}( - P, P0, V, ∇V, τ, ε, ε_pl, EII_pl, deepcopy(τ), R - ) - end - - function $(name)(args::Vararg{Any,N}) where {N} - return new{ - ViscoElastic, - typeof(args[3]), - typeof(args[5]), - typeof(args[end]), - 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/thermal_diffusion/MetaDiffusion.jl b/src/thermal_diffusion/MetaDiffusion.jl deleted file mode 100644 index 1c392202..00000000 --- a/src/thermal_diffusion/MetaDiffusion.jl +++ /dev/null @@ -1,236 +0,0 @@ -function make_thermal_arrays!(ndim) - flux_sym1 = (:qTx, :qTy, :qTz) - flux_sym2 = (:qTx2, :qTy2, :qTz2) - flux1 = [:($(flux_sym1[i])::_T) for i in 1:ndim] - flux2 = [:($(flux_sym2[i])::_T) for i in 1:ndim] - - @eval begin - struct ThermalArrays{_T} - T::_T # Temperature @ grid nodes - Tc::_T # Temperature @ cell centers - ΔT::_T - ΔTc::_T - Told::_T - dT_dt::_T - $(flux1...) - $(flux2...) - H::_T # source terms - shear_heating::_T # shear heating terms - ResT::_T - - function ThermalArrays(args::Vararg{T,N}) where {T<:AbstractArray,N} - return new{T}(args...) - end - - function ThermalArrays(ni::NTuple{1,Integer}) - nx, = ni - T, ΔT, Told = @zeros(ni...), @zeros(ni...), @zeros(ni...) - dT_dt = @zeros(nx - 2) - qTx = @zeros(nx - 1) - qTx2 = @zeros(nx - 1) - ResT = @zeros(nx - 2) - return new{typeof(T)}(T, ΔT, Told, dT_dt, qTx, qTx2, ResT) - end - - function ThermalArrays(ni::NTuple{2,Integer}) - nx, ny = ni - T = @zeros(nx + 3, ny + 1) - ΔT = @zeros(nx + 3, ny + 1) - Told = @zeros(nx + 3, ny + 1) - Tc = @zeros(ni...) - ΔTc = @zeros(ni...) - H = @zeros(ni...) - shear_heating = @zeros(ni...) - dT_dt = @zeros(nx + 1, ny - 1) - qTx = @zeros(nx + 2, ny - 1) - qTy = @zeros(nx + 1, ny) - qTx2 = @zeros(nx + 2, ny - 1) - qTy2 = @zeros(nx + 1, ny) - ResT = @zeros(nx + 1, ny - 1) - return new{typeof(T)}( - T, - Tc, - ΔT, - ΔTc, - Told, - dT_dt, - qTx, - qTy, - qTx2, - qTy2, - H, - shear_heating, - ResT, - ) - end - - function ThermalArrays(ni::NTuple{3,Integer}) - nx, ny, nz = ni - T, ΔT, Told = @zeros(ni .+ 1...), @zeros(ni .+ 1...), @zeros(ni .+ 1...) - Tc = @zeros(ni...) - ΔTc = @zeros(ni...) - H = @zeros(ni...) - shear_heating = @zeros(ni...) - dT_dt = @zeros(ni .- 1) - qTx = @zeros(nx, ny - 1, nz - 1) - qTy = @zeros(nx - 1, ny, nz - 1) - qTz = @zeros(nx - 1, ny - 1, nz) - qTx2 = @zeros(nx, ny - 1, nz - 1) - qTy2 = @zeros(nx - 1, ny, nz - 1) - qTz2 = @zeros(nx - 1, ny - 1, nz) - ResT = @zeros((ni .- 1)...) - return new{typeof(T)}( - T, - Tc, - ΔT, - ΔTc, - Told, - dT_dt, - qTx, - qTy, - qTz, - qTx2, - qTy2, - qTz2, - H, - shear_heating, - ResT, - ) - end - end - end -end - -function make_PTthermal_struct!() - @eval begin - struct PTThermalCoeffs{T,M,nDim} - CFL::T - ϵ::T - max_lxyz::T - max_lxyz2::T - Vpdτ::T - θr_dτ::M - dτ_ρ::M - - function PTThermalCoeffs( - K, ρCp, dt, di::NTuple{nDim,T}, li::NTuple{nDim,Any}; ϵ=1e-8, CFL=0.9 / √3 - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - max_lxyz2 = max_lxyz^2 - Re = @. π + √(π * π + ρCp * max_lxyz2 / K / dt) # Numerical Reynolds number - θr_dτ = @. max_lxyz / Vpdτ / Re - dτ_ρ = @. Vpdτ * max_lxyz / K / Re - - return new{eltype(Vpdτ),typeof(dτ_ρ),nDim}( - CFL, ϵ, max_lxyz, max_lxyz2, Vpdτ, θr_dτ, dτ_ρ - ) - end - - function PTThermalCoeffs( - CFL::T, ϵ::T, max_lxyz::T, Vpdτ::T, θr_dτ::M, dτ_ρ::M - ) where {T,M} - nDim = length(size(θr_dτ)) - return new{T,M,nDim}(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) - end - end - - # with phase ratios - function PTThermalCoeffs( - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) - - idx = ntuple(i -> 1:ni[i], Val(nDim)) - @parallel (idx) compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase_ratios.center, args, max_lxyz, Vpdτ, inv(dt) - ) - - return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) - end - - # without phase ratios - function PTThermalCoeffs( - # function PTThermalCoeffs( - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - Vpdτ = min(di...) * CFL - max_lxyz = max(li...) - θr_dτ, dτ_ρ = @zeros(ni...), @zeros(ni...) - - idx = ntuple(i -> 1:ni[i], Val(nDim)) - @parallel (idx) compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, inv(dt) - ) - - return PTThermalCoeffs(CFL, ϵ, max_lxyz, Vpdτ, θr_dτ, dτ_ρ) - end - - @parallel_indices (I...) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt - ) - _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, I... - ) - - return nothing - end - - @parallel_indices (I...) function compute_pt_thermal_arrays!( - θr_dτ::AbstractArray, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt - ) - _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, I... - ) - - return nothing - end - - function _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, phase, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} - ) where {N} - args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) - phase_ij = phase[Idx...] - ρCp = compute_ρCp(rheology, phase_ij, args_ij) - _K = inv(fn_ratio(compute_conductivity, rheology, phase_ij, args_ij)) - - _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number - θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re - dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re - - return nothing - end - - function _compute_pt_thermal_arrays!( - θr_dτ, dτ_ρ, rheology, args, max_lxyz, Vpdτ, _dt, Idx::Vararg{Int,N} - ) where {N} - args_ij = (; T=args.T[Idx...], P=args.P[Idx...]) - - ρCp = compute_ρCp(rheology, args_ij) - _K = inv(compute_conductivity(rheology, args_ij)) - - _Re = inv(π + √(π * π + ρCp * max_lxyz^2 * _K * _dt)) # Numerical Reynolds number - θr_dτ[Idx...] = max_lxyz / Vpdτ * _Re - dτ_ρ[Idx...] = Vpdτ * max_lxyz * _K * _Re - - return nothing - end - end -end From 334c1002051e3d42cf1054f0b81d53437cda2989 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:10:26 +0200 Subject: [PATCH 51/91] remove JP --- Project.toml | 2 -- docs/Project.toml | 2 -- 2 files changed, 4 deletions(-) diff --git a/Project.toml b/Project.toml index 48d96b3f..ca1516b5 100644 --- a/Project.toml +++ b/Project.toml @@ -10,7 +10,6 @@ CellArrays = "d35fcfd7-7af4-4c67-b1aa-d78070614af4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" ImplicitGlobalGrid = "4d7a3746-15be-11ea-1130-334b0c4f5fa0" -JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" @@ -35,7 +34,6 @@ CellArrays = "0.2" GeoParams = "0.5" HDF5 = "0.17.1" ImplicitGlobalGrid = "0.15.0" -JustPIC = "0.3.2" MPI = "0.20" MuladdMacro = "0.2" ParallelStencil = "0.12" diff --git a/docs/Project.toml b/docs/Project.toml index 21a73358..bdf780a3 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,10 +1,8 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" -JustPIC = "10dc771f-8528-4cd9-9d3b-b21b2e693339" JustRelax = "34418575-392d-4e26-8c6d-96b0910afa06" [compat] Documenter = "1" -JustPIC = "0.3.2" JustRelax = "0.1.2" From b246d7a48ef41c4c5ca38b5a8f1ae2222fff144c Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:20:48 +0200 Subject: [PATCH 52/91] move index.md to man/ --- docs/src/{ => man}/index.md | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename docs/src/{ => man}/index.md (100%) diff --git a/docs/src/index.md b/docs/src/man/index.md similarity index 100% rename from docs/src/index.md rename to docs/src/man/index.md From fb1b81a72a6738ee49be04c83e996ef417deff08 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:36:20 +0200 Subject: [PATCH 53/91] fix syntax --- test/test_shearheating2D.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 9d60518b..4ee419ce 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -98,7 +98,7 @@ function Shearheating2D() # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) From 2d0030e276b75539990cb80d244f0105c674db54 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:44:21 +0200 Subject: [PATCH 54/91] ext: fix type syntax --- ext/JustRelaxCUDAExt.jl | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 69820181..6eaae952 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -38,16 +38,16 @@ module JustRelax2D return StokesArrays(ni) end - function JustRelax.JustRelax2D.ThermalArrays( + function JustRelax.JustRelax2D.::JustRelax.ThermalArrays( ::Type{CUDABackend}, ni::NTuple{N,Number} ) where {N} - return ThermalArrays(ni...) + return ::JustRelax.ThermalArrays(ni...) end - function JustRelax.JustRelax2D.ThermalArrays( + function JustRelax.JustRelax2D.::JustRelax.ThermalArrays( ::Type{CUDABackend}, ni::Vararg{Number,N} ) where {N} - return ThermalArrays(ni...) + return ::JustRelax.ThermalArrays(ni...) end function JustRelax.JustRelax2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) @@ -88,27 +88,27 @@ module JustRelax2D end # Boundary conditions - function JustRelax.JustRelax2D.flow_bcs!(::CUDABackendTrait, stokes::StokesArrays, bcs) + function JustRelax.JustRelax2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) return _flow_bcs!(bcs, @velocity(stokes)) end - function flow_bcs!(::CUDABackendTrait, stokes::StokesArrays, bcs) + function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) return _flow_bcs!(bcs, @velocity(stokes)) end function JustRelax.JustRelax2D.thermal_bcs!( - ::CUDABackendTrait, thermal::ThermalArrays, bcs + ::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs ) return thermal_bcs!(thermal.T, bcs) end - function thermal_bcs!(::CUDABackendTrait, thermal::ThermalArrays, bcs) + function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end # Phases function JustRelax.JustRelax2D.phase_ratios_center( - ::CUDABackendTrait, phase_ratios::PhaseRatio, particles, grid::Geometry, phases + ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases ) return _phase_ratios_center(phase_ratios, particles, grid, phases) end @@ -151,14 +151,14 @@ module JustRelax2D return compute_ρg!(ρg, rheology, args) end function JustRelax.JustRelax2D.compute_ρg!( - ρg::CuArray, phase_ratios::PhaseRatio, rheology, args + ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args ) return compute_ρg!(ρg, phase_ratios, rheology, args) end # Interpolations function JustRelax.JustRelax2D.temperature2center!( - ::CUDABackendTrait, thermal::ThermalArrays + ::CUDABackendTrait, thermal::JustRelax.ThermalArrays ) return _temperature2center!(thermal) end @@ -183,13 +183,13 @@ module JustRelax2D end # Utils - function JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff, I) + function JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di, dt_diff, I) return compute_dt(S, di, dt_diff, I::IGG) end - function JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di, dt_diff) + function JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di, dt_diff) return compute_dt(S, di, dt_diff) end - JustRelax.JustRelax2D.compute_dt(S::StokesArrays, di) = compute_dt(S, di) + JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(S, di) end From 1ecccb74c48ab367caa51dba994f8a13882263da Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Mon, 29 Apr 2024 23:46:50 +0200 Subject: [PATCH 55/91] ext: CUDA trait --- ext/JustRelaxCUDAExt.jl | 7 ++----- src/types/traits.jl | 4 ++-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 6eaae952..2783786e 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -18,16 +18,13 @@ module JustRelax2D using MPI import JustRelax: - IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend, Geometry, @cell + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell @init_parallel_stencil(CUDA, Float64, 2) include("../src/common.jl") include("../src/stokes/Stokes2D.jl") - # add CUDA traits - struct CUDABackendTrait <: BackendTrait end - @inline backend(::CuArray) = CUDABackendTrait() @inline backend(::Type{<:CuArray}) = CUDABackendTrait() @@ -144,7 +141,7 @@ module JustRelax2D end ## Stress - JustRelax.JustRelax2D.tensor_invariant!(A::SymmetricTensor) = tensor_invariant!(A) + JustRelax.JustRelax2D.tensor_invariant!(A::JustRelax.SymmetricTensor) = tensor_invariant!(A) ## Buoyancy forces function JustRelax.JustRelax2D.compute_ρg!(ρg::CuArray, rheology, args) diff --git a/src/types/traits.jl b/src/types/traits.jl index bcc0c8e8..20431a17 100644 --- a/src/types/traits.jl +++ b/src/types/traits.jl @@ -1,8 +1,8 @@ abstract type BackendTrait end struct CPUBackendTrait <: BackendTrait end struct NonCPUBackendTrait <: BackendTrait end -# struct CUDABackendTrait <: BackendTrait end -# struct AMDGPUBackendTrait <: BackendTrait end +struct CUDABackendTrait <: BackendTrait end +struct AMDGPUBackendTrait <: BackendTrait end @inline backend(::Array) = CPUBackendTrait() @inline backend(::Type{<:Array}) = CPUBackendTrait() From 15545554daa050c8fa78aebd6eff890cc2a34c99 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 09:10:04 +0200 Subject: [PATCH 56/91] last update --- test/test_shearheating2D.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 4ee419ce..86a0f1c1 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -129,8 +129,8 @@ function Shearheating2D() ) ## Compression and not extension - fix this εbg = 5e-14 - stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) + stokes.V.Vx .= PTArray(backend_JR)([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend_JR)([ (ly - abs(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) @@ -154,7 +154,7 @@ function Shearheating2D() # ------------------------------ # Stokes solver ---------------- - solve!( + iters = solve!( stokes, pt_stokes, di, From c229059309f2c63ef8d127b65f922b1d84f9d106 Mon Sep 17 00:00:00 2001 From: aelligp Date: Tue, 30 Apr 2024 10:51:11 +0200 Subject: [PATCH 57/91] rename `save_vtk` to avoid error --- .../Blankenbach2D/Benchmark2D_WENO5.jl | 48 ++++++++-------- .../stokes2D/Blankenbach2D/Benchmark2D_sgd.jl | 52 +++++++++--------- .../Blankenbach2D/Benchmark2D_sgd_scaled.jl | 55 +++++++++---------- 3 files changed, 77 insertions(+), 78 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl index 664c4579..abab0486 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl @@ -41,7 +41,7 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) - @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += 20.0 end return nothing @@ -53,7 +53,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =false) # Physical domain ------------------------------------ ly = 1000e3 # domain length in y @@ -75,7 +75,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # Weno model ----------------------------------------- weno = WENO5(ni=(nx,ny).+1, method=Val{2}()) # ni.+1 for Temp # ---------------------------------------------------- - + # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 24, 36, 12 particles = init_particles( @@ -113,11 +113,11 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # Rayleigh number ΔT = thermal.T[1,1] - thermal.T[1,end] - Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / + Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / (κ * rheology[1].CompositeRheology[1].elements[1].η ) @show Ra - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) @@ -141,8 +141,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # IO ------------------------------------------------ # if it does not exist, make folder where figures are stored - if save_vtk - vtk_dir = figdir*"\\vtk" + if do_vtk + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -165,7 +165,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f end local Vx_v, Vy_v - if save_vtk + if do_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) end @@ -184,7 +184,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f while it ≤ nit @show it - + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) @@ -240,8 +240,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # ------------------------------ # Nusselt number, Nu = H/ΔT/L ∫ ∂T/∂z dx ---- - Nu_it = (ly / (1000.0*lx)) * - sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) + Nu_it = (ly / (1000.0*lx)) * + sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) push!(Nu_top, Nu_it) # ------------------------------------------- @@ -250,7 +250,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f Urms_it = let # JustRelax.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude - sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * + sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * ((ly * rheology[1].Density[1].ρ0 * rheology[1].HeatCapacity[1].Cp) / rheology[1].Conductivity[1].k ) end push!(Urms, Urms_it) @@ -260,7 +260,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit - if save_vtk + if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), @@ -302,12 +302,12 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") # h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(273,1273) ) - # + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) - # + # h3 = heatmap!(ax3, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vx) , colormap=:batlow) - # - # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) + # + # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)) , colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) @@ -319,7 +319,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig - + fig2 = Figure(size = (900, 1200), title = "Time Series") ax21 = Axis(fig2[1,1], aspect = 3, title = "V_{RMS}") ax22 = Axis(fig2[2,1], aspect = 3, title = "Nu_{top}") @@ -332,21 +332,21 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # ------------------------------ end - # Horizontally averaged depth profile + # Horizontally averaged depth profile Tmean = @zeros(ny+1) Emean = @zeros(ny) let for j = 1:(ny+1) - Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) - end + Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) + end for j = 1:ny Emean[j] = sum(η[:,j])/nx end 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(η)⟩") + ax2 = Axis(fig[1,2], aspect = 2/3, title = "⟨log10(η)⟩") lines!(ax1, Tmean, xvi[2]./1e3) lines!(ax2, log10.(Emean), xci[2]./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) @@ -364,7 +364,7 @@ end # (Path)/folder where output data and figures are stored figdir = "Blankenbach_WENO" -save_vtk = false # set to true to generate VTK files for ParaView +do_vtk = false # set to true to generate VTK files for ParaView ar = 1 # aspect ratio n = 51 nx = n @@ -377,4 +377,4 @@ else end # run main script -main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, save_vtk = save_vtk); +main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, do_vtk = do_vtk); diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl index 0d11d5a4..bba6a999 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl @@ -12,7 +12,7 @@ using JustPIC._2D const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using Printf, LinearAlgebra, GeoParams, CairoMakie, CellArrays # Load file with all the rheology configurations include("Blankenbach_Rheology.jl") @@ -42,7 +42,7 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) - @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += 20.0 end return nothing @@ -54,7 +54,7 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =false) # Physical domain ------------------------------------ ly = 1000e3 # domain length in y @@ -114,11 +114,11 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # Rayleigh number ΔT = thermal.T[1,1] - thermal.T[1,end] - Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / + Ra = (rheology[1].Density[1].ρ0 * rheology[1].Gravity[1].g * rheology[1].Density[1].α * ΔT * ly^3.0 ) / (κ * rheology[1].CompositeRheology[1].elements[1].η ) @show Ra - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) @@ -142,8 +142,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # IO ------------------------------------------------ # if it does not exist, make folder where figures are stored - if save_vtk - vtk_dir = figdir*"\\vtk" + if do_vtk + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -175,7 +175,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f pT0.data .= pT.data local Vx_v, Vy_v - if save_vtk + if do_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) end @@ -191,7 +191,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f while it ≤ nit @show it - + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) @@ -260,8 +260,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f phase_ratios_center(phase_ratios, particles, grid, pPhases) # Nusselt number, Nu = H/ΔT/L ∫ ∂T/∂z dx ---- - Nu_it = (ly / (1000.0*lx)) * - sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) + Nu_it = (ly / (1000.0*lx)) * + sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) push!(Nu_top, Nu_it) # ------------------------------------------- @@ -270,7 +270,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f Urms_it = let JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @. Vx_v .= hypot.(Vx_v, Vy_v) # we reuse Vx_v to store the velocity magnitude - sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * + sqrt( sum( Vx_v.^2 .* prod(di)) / lx /ly ) * ((ly * rheology[1].Density[1].ρ0 * rheology[1].HeatCapacity[1].Cp) / rheology[1].Conductivity[1].k ) end push!(Urms, Urms_it) @@ -279,7 +279,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) - @views T_buffer[:, end] .= 273.0 + @views T_buffer[:, end] .= 273.0 @views T_buffer[:, 1] .= 1273.0 @views thermal.T[2:end-1, :] .= T_buffer flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -288,7 +288,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit - if save_vtk + if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), @@ -330,12 +330,12 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") # h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(273,1273) ) - # + # h2 = heatmap!(ax2, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy) , colormap=:batlow) - # + # h3 = heatmap!(ax3, xvi[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vx) , colormap=:batlow) - # - h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) + # + h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(273,1273), markersize=3) #h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η)) , colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) @@ -347,7 +347,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig - + fig2 = Figure(size = (900, 1200), title = "Time Series") ax21 = Axis(fig2[1,1], aspect = 3, title = "V_{RMS}") ax22 = Axis(fig2[2,1], aspect = 3, title = "Nu_{top}") @@ -360,21 +360,21 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # ------------------------------ end - # Horizontally averaged depth profile + # Horizontally averaged depth profile Tmean = @zeros(ny+1) Emean = @zeros(ny) let for j = 1:(ny+1) - Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) - end + Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) + end for j = 1:ny Emean[j] = sum(η[:,j])/nx end 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(η)⟩") + ax2 = Axis(fig[1,2], aspect = 2/3, title = "⟨log10(η)⟩") lines!(ax1, Tmean, xvi[2]./1e3) lines!(ax2, log10.(Emean), xci[2]./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) @@ -392,9 +392,9 @@ end # (Path)/folder where output data and figures are stored figdir = "Blankenbach_subgrid" -save_vtk = false # set to true to generate VTK files for ParaView +do_vtk = false # set to true to generate VTK files for ParaView ar = 1 # aspect ratio -n = 51 +n = 128 nx = n ny = n nit = 6e3 @@ -405,4 +405,4 @@ else end # run main script -main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, save_vtk = save_vtk); +main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, do_vtk = do_vtk); diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl index 12c327ed..677ee7f6 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -12,7 +12,7 @@ using JustPIC._2D const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie +using Printf, LinearAlgebra, GeoParams, CairoMakie # Load file with all the rheology configurations include("Blankenbach_Rheology_scaled.jl") @@ -27,7 +27,7 @@ function copyinn_x!(A, B) end # Initial thermal profile -@parallel_indices (i, j) function init_T!(T, y) +@parallel_indices (i, j) function init_T!(T, y) T[i, j] = 1 - y[j] return nothing end @@ -35,7 +35,7 @@ end # Thermal rectangular perturbation function rectangular_perturbation!(T, xc, yc, r, xvi) @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) - @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += .2 end return nothing @@ -47,8 +47,8 @@ end ## END OF HELPER FUNCTION ------------------------------------------------------------ ## BEGIN OF MAIN SCRIPT -------------------------------------------------------------- -function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =false) - +function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =false) + # Physical domain ------------------------------------ ly = 1.0 # domain length in y lx = ly # domain length in x @@ -108,8 +108,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f Ra = rheology[1].Gravity[1].g println("Ra = $Ra") - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + # Buoyancy forces & viscosity ---------------------- ρg = @zeros(ni...), @zeros(ni...) η = @ones(ni...) @@ -117,7 +117,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f compute_viscosity!( stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - + # PT coefficients for thermal diffusion ------------- pt_thermal = PTThermalCoeffs( backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL = 0.5 / √2.1 @@ -132,8 +132,8 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # IO ------------------------------------------------ # if it does not exist, make folder where figures are stored - if save_vtk - vtk_dir = figdir*"\\vtk" + if do_vtk + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -165,7 +165,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f pT0.data .= pT.data local Vx_v, Vy_v - if save_vtk + if do_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) end @@ -180,7 +180,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f Vy_v = @zeros(ni.+1...) while it ≤ nit - + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) @@ -250,7 +250,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f phase_ratios_center(phase_ratios, particles, grid, pPhases) # Nusselt number, Nu = ∫ ∂T/∂z dx ---- - Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) + Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) push!(Nu_top, Nu_it) # ------------------------------------------- @@ -267,18 +267,18 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) - @views T_buffer[:, end] .= 0.0 + @views T_buffer[:, end] .= 0.0 @views T_buffer[:, 1] .= 1.0 @views thermal.T[2:end-1, :] .= T_buffer flow_bcs!(stokes, flow_bcs) # apply boundary conditions temperature2center!(thermal) @show extrema(thermal.T) any(isnan.(thermal.T)) && break - + # Data I/O and plotting --------------------- if it == 1 || rem(it, 200) == 0 || it == nit || any(isnan.(thermal.T)) - if save_vtk + if do_vtk JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), @@ -320,12 +320,12 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f ax4 = Axis(fig[2,3], aspect = ar, title = "T [K]") # h1 = heatmap!(ax1, xvi[1], xvi[2], Array(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(0, 1) ) - # + # h2 = heatmap!(ax2, xvi[1], xvi[2], Array(stokes.V.Vy) , colormap=:batlow) - # + # h3 = heatmap!(ax3, xvi[1], xvi[2], Array(stokes.V.Vx) , colormap=:batlow) - # - # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), markersize=3) + # + # h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), markersize=3) h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) hidexdecorations!(ax1) hidexdecorations!(ax2) @@ -337,7 +337,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f linkaxes!(ax1, ax2, ax3, ax4) save(joinpath(figdir, "$(it).png"), fig) fig - + fig2 = Figure(size = (900, 1200), title = "Time Series") ax21 = Axis(fig2[1,1], aspect = 3, title = "(V_{RMS})") ax22 = Axis(fig2[2,1], aspect = 3, title = "(Nu_{top})") @@ -359,21 +359,21 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", save_vtk =f # ------------------------------ end - # Horizontally averaged depth profile + # Horizontally averaged depth profile Tmean = @zeros(ny+1) Emean = @zeros(ny) let for j = 1:(ny+1) - Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) - end + Tmean[j] = sum(thermal.T[2:end-1,j])/(nx+1) + end for j = 1:ny Emean[j] = sum(η[:,j])/nx end 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(η)⟩") + ax2 = Axis(fig[1,2], aspect = 2/3, title = "⟨log10(η)⟩") lines!(ax1, Tmean, (1 .- xvi[2])) lines!(ax2, log10.(Emean), (1 .- xci[2])) ylims!(ax1, maximum(xvi[2]), 0) @@ -391,7 +391,7 @@ end # (Path)/folder where output data and figures are stored figdir = "Blankenbach_subgrid_scaled" -save_vtk = false # set to true to generate VTK files for ParaView +do_vtk = false # set to true to generate VTK files for ParaView ar = 1 # aspect ratio n = 64 nx = n @@ -404,5 +404,4 @@ else end # run main script -main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, save_vtk = save_vtk); - +main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nit = nit, do_vtk = do_vtk); From 3474566b963248a5e3be36397610ebcb85ec951e Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 10:52:07 +0200 Subject: [PATCH 58/91] working CUDA ext --- ext/JustRelaxCUDAExt.jl | 75 +++++++++++--------- src/JustRelax.jl | 3 + src/JustRelax_CPU.jl | 2 + src/Utils.jl | 28 +++++--- src/boundaryconditions/BoundaryConditions.jl | 29 -------- src/boundaryconditions/types.jl | 28 ++++++++ src/stokes/StressKernels.jl | 9 +++ 7 files changed, 102 insertions(+), 72 deletions(-) create mode 100644 src/boundaryconditions/types.jl diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 2783786e..8da7fa2d 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -17,8 +17,11 @@ module JustRelax2D using GeoParams, LinearAlgebra, Printf using MPI + import JustRelax.JustRelax2D as JR2D + import JustRelax: IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions @init_parallel_stencil(CUDA, Float64, 2) @@ -29,29 +32,29 @@ module JustRelax2D @inline backend(::Type{<:CuArray}) = CUDABackendTrait() # Types - function JustRelax.JustRelax2D.StokesArrays( + function JR2D.StokesArrays( ::Type{CUDABackend}, ni::NTuple{N,Integer} ) where {N} return StokesArrays(ni) end - function JustRelax.JustRelax2D.::JustRelax.ThermalArrays( + function JR2D.ThermalArrays( ::Type{CUDABackend}, ni::NTuple{N,Number} ) where {N} - return ::JustRelax.ThermalArrays(ni...) + return ThermalArrays(ni...) end - function JustRelax.JustRelax2D.::JustRelax.ThermalArrays( + function JR2D.ThermalArrays( ::Type{CUDABackend}, ni::Vararg{Number,N} ) where {N} - return ::JustRelax.ThermalArrays(ni...) + return ThermalArrays(ni...) end - function JustRelax.JustRelax2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) return PhaseRatio(ni, num_phases) end - function PTThermalCoeffs( + function JR2D.PTThermalCoeffs( ::Type{CUDABackend}, rheology, phase_ratios, @@ -63,12 +66,12 @@ module JustRelax2D ϵ=1e-8, CFL=0.9 / √3, ) where {nDim,T} - return JustRelax.JustRelax2D.PTThermalCoeffs( + return PTThermalCoeffs( rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL ) end - function PTThermalCoeffs( + function JR2D.PTThermalCoeffs( ::Type{CUDABackend}, rheology, args, @@ -79,13 +82,13 @@ module JustRelax2D ϵ=1e-8, CFL=0.9 / √3, ) where {nDim,T} - return JustRelax.JustRelax2D.PTThermalCoeffs( + return PTThermalCoeffs( rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL ) end # Boundary conditions - function JustRelax.JustRelax2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) return _flow_bcs!(bcs, @velocity(stokes)) end @@ -93,18 +96,16 @@ module JustRelax2D return _flow_bcs!(bcs, @velocity(stokes)) end - function JustRelax.JustRelax2D.thermal_bcs!( - ::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs - ) + function JR2D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end - + function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) return thermal_bcs!(thermal.T, bcs) end # Phases - function JustRelax.JustRelax2D.phase_ratios_center( + function JR2D.phase_ratios_center( ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases ) return _phase_ratios_center(phase_ratios, particles, grid, phases) @@ -112,17 +113,17 @@ module JustRelax2D # Rheology ## viscosity - function JustRelax.JustRelax2D.compute_viscosity!( + function JR2D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end - function JustRelax.JustRelax2D.compute_viscosity!( + function JR2D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end - function JustRelax.JustRelax2D.compute_viscosity!( + function JR2D.compute_viscosity!( η, ν, εII::CuArray, args, rheology, cutoff ) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) @@ -141,32 +142,40 @@ module JustRelax2D end ## Stress - JustRelax.JustRelax2D.tensor_invariant!(A::JustRelax.SymmetricTensor) = tensor_invariant!(A) + function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) + end ## Buoyancy forces - function JustRelax.JustRelax2D.compute_ρg!(ρg::CuArray, rheology, args) + function JR2D.compute_ρg!(ρg::CuArray, rheology, args) return compute_ρg!(ρg, rheology, args) end - function JustRelax.JustRelax2D.compute_ρg!( + function JR2D.compute_ρg!( ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args ) return compute_ρg!(ρg, phase_ratios, rheology, args) end # Interpolations - function JustRelax.JustRelax2D.temperature2center!( + function JR2D.temperature2center!( ::CUDABackendTrait, thermal::JustRelax.ThermalArrays ) return _temperature2center!(thermal) end - function JustRelax.JustRelax2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + + function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) + end + + function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} return vertex2center!(center, vertex) end - function JustRelax.JustRelax2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + + function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} return center2vertex!(vertex, center) end - function JustRelax.JustRelax2D.center2vertex!( + function JR2D.center2vertex!( vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T ) where {T<:CuArray} return center2vertex!( @@ -175,18 +184,18 @@ module JustRelax2D end # Solvers - function JustRelax.JustRelax2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) return _solve!(stokes, args...; kwargs...) end - # Utils - function JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di, dt_diff, I) - return compute_dt(S, di, dt_diff, I::IGG) + function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) end - function JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di, dt_diff) - return compute_dt(S, di, dt_diff) + + # Utils + function JR2D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) end - JustRelax.JustRelax2D.compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(S, di) end diff --git a/src/JustRelax.jl b/src/JustRelax.jl index 174f1941..0322619f 100644 --- a/src/JustRelax.jl +++ b/src/JustRelax.jl @@ -31,6 +31,9 @@ include("types/heat_diffusion.jl") include("types/phases.jl") # export PhaseRatio +include("boundaryconditions/types.jl") +export TemperatureBoundaryConditions, FlowBoundaryConditions + include("types/traits.jl") export BackendTrait, CPUBackendTrait, NonCPUBackendTrait diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index 7caa2e2a..a859501e 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -10,6 +10,7 @@ using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend import JustRelax: PTStokesCoeffs +import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions @init_parallel_stencil(Threads, Float64, 2) @@ -31,6 +32,7 @@ using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend import JustRelax: PTStokesCoeffs +import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions @init_parallel_stencil(Threads, Float64, 3) diff --git a/src/Utils.jl b/src/Utils.jl index 6acdb73e..9c6c65a2 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -399,12 +399,20 @@ macro unpack(x) end end +function compute_dt(S::JustRelax.StokesArrays, args...) + return compute_dt(backend(S), S, args...) +end + +function compute_dt(::CPUBackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) +end + """ compute_dt(S::JustRelax.StokesArrays, di) Compute the time step `dt` for the velocity field `S.V` for a regular grid with grid spacing `di`. """ -@inline compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(@velocity(S), di, Inf) +@inline _compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(@velocity(S), di, Inf) """ compute_dt(S::JustRelax.StokesArrays, di, dt_diff) @@ -412,12 +420,12 @@ Compute the time step `dt` for the velocity field `S.V` for a regular grid with Compute the time step `dt` for the velocity field `S.V` and the diffusive maximum time step `dt_diff` for a regular gridwith grid spacing `di`. """ -@inline compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = - compute_dt(@velocity(S), di, dt_diff) +@inline _compute_dt(S::JustRelax.StokesArrays, di, dt_diff) = + _compute_dt(@velocity(S), di, dt_diff) -@inline function compute_dt(V::NTuple, di, dt_diff) +@inline function _compute_dt(V::NTuple, di, dt_diff) n = inv(length(V) + 0.1) - dt_adv = mapreduce(x -> x[1] * inv(maximum_mpi(abs.(x[2]))), min, zip(di, V)) * n + dt_adv = mapreduce(x -> x[1] * inv(maximum(abs.(x[2]))), min, zip(di, V)) * n return min(dt_diff, dt_adv) end """ @@ -427,8 +435,8 @@ Compute the time step `dt` for the velocity field `S.V` for a regular gridwith g The implicit global grid variable `I` implies that the time step is calculated globally and not separately on each block. """ -@inline compute_dt(S::JustRelax.StokesArrays, di, I::IGG) = - compute_dt(@velocity(S), di, Inf, I::IGG) +@inline _compute_dt(S::JustRelax.StokesArrays, di, I::IGG) = + _compute_dt(@velocity(S), di, Inf, I::IGG) """ compute_dt(S::JustRelax.StokesArrays, di, dt_diff) @@ -437,11 +445,11 @@ Compute the time step `dt` for the velocity field `S.V` and the diffusive maximu `dt_diff` for a regular gridwith grid spacing `di`. The implicit global grid variable `I` implies that the time step is calculated globally and not separately on each block. """ -@inline function compute_dt(S::JustRelax.StokesArrays, di, dt_diff, I::IGG) - return compute_dt(@velocity(S), di, dt_diff, I::IGG) +@inline function _compute_dt(S::JustRelax.StokesArrays, di, dt_diff, I::IGG) + return _compute_dt(@velocity(S), di, dt_diff, I::IGG) end -@inline function compute_dt(V::NTuple, di, dt_diff, I::IGG) +@inline function _compute_dt(V::NTuple, di, dt_diff, I::IGG) n = inv(length(V) + 0.1) dt_adv = mapreduce(x -> x[1] * inv(maximum_mpi(abs.(x[2]))), max, zip(di, V)) * n return min(dt_diff, dt_adv) diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 88a483ef..3ac377d1 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -1,32 +1,3 @@ -abstract type AbstractBoundaryConditions end - -struct TemperatureBoundaryConditions{T,nD} <: AbstractBoundaryConditions - no_flux::T - - function TemperatureBoundaryConditions(; - no_flux::T=(left=true, right=false, top=false, bot=false) - ) where {T} - nD = length(no_flux) == 4 ? 2 : 3 - return new{T,nD}(no_flux) - end -end - -struct FlowBoundaryConditions{T,nD} <: AbstractBoundaryConditions - no_slip::T - free_slip::T - free_surface::Bool - - function FlowBoundaryConditions(; - no_slip::T=(left=false, right=false, top=false, bot=false), - free_slip::T=(left=true, right=true, top=true, bot=true), - free_surface::Bool=false, - ) where {T} - @assert length(no_slip) === length(free_slip) - nD = length(no_slip) == 4 ? 2 : 3 - return new{T,nD}(no_slip, free_slip, free_surface) - end -end - @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)...) diff --git a/src/boundaryconditions/types.jl b/src/boundaryconditions/types.jl new file mode 100644 index 00000000..768b8dd8 --- /dev/null +++ b/src/boundaryconditions/types.jl @@ -0,0 +1,28 @@ +abstract type AbstractBoundaryConditions end + +struct TemperatureBoundaryConditions{T,nD} <: AbstractBoundaryConditions + no_flux::T + + function TemperatureBoundaryConditions(; + no_flux::T=(left=true, right=false, top=false, bot=false) + ) where {T} + nD = length(no_flux) == 4 ? 2 : 3 + return new{T,nD}(no_flux) + end +end + +struct FlowBoundaryConditions{T,nD} <: AbstractBoundaryConditions + no_slip::T + free_slip::T + free_surface::Bool + + function FlowBoundaryConditions(; + no_slip::T=(left=false, right=false, top=false, bot=false), + free_slip::T=(left=true, right=true, top=true, bot=true), + free_surface::Bool=false, + ) where {T} + @assert length(no_slip) === length(free_slip) + nD = length(no_slip) == 4 ? 2 : 3 + return new{T,nD}(no_slip, free_slip, free_surface) + end +end \ No newline at end of file diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 1beba926..2aaab4c3 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -306,6 +306,15 @@ Compute the tensor invariant of the given symmetric tensor `A`. - `A::JustRelax.SymmetricTensor`: The input symmetric tensor. """ function tensor_invariant!(A::JustRelax.SymmetricTensor) + tensor_invariant!(backend(A), A) + return nothing +end + +function tensor_invariant!(::CPUBackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) +end + +function _tensor_invariant!(A::JustRelax.SymmetricTensor) ni = size(A.II) @parallel (@idx ni) tensor_invariant_kernel!(A.II, @tensor(A)...) return nothing From a4e60bb20b975045547eb31d10763d2edfcb2f78 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 11:04:17 +0200 Subject: [PATCH 59/91] 3D extensions --- ext/JustRelaxCUDAExt.jl | 195 +--------------------------------------- src/ext/CUDA/2D.jl | 189 ++++++++++++++++++++++++++++++++++++++ src/ext/CUDA/3D.jl | 189 ++++++++++++++++++++++++++++++++++++++ src/stokes/Stokes3D.jl | 6 +- 4 files changed, 385 insertions(+), 194 deletions(-) create mode 100644 src/ext/CUDA/2D.jl create mode 100644 src/ext/CUDA/3D.jl diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 8da7fa2d..978b8041 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -6,197 +6,10 @@ import JustRelax: PTArray JustRelax.PTArray(::Type{CUDABackend}) = CuArray -module JustRelax2D +@inline backend(::CuArray) = CUDABackendTrait() +@inline backend(::Type{<:CuArray}) = CUDABackendTrait() - using JustRelax: JustRelax - using CUDA - using StaticArrays - using CellArrays - using ParallelStencil, ParallelStencil.FiniteDifferences2D - using ImplicitGlobalGrid - using GeoParams, LinearAlgebra, Printf - using MPI - - import JustRelax.JustRelax2D as JR2D - - import JustRelax: - IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell - import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions - - @init_parallel_stencil(CUDA, Float64, 2) - - include("../src/common.jl") - include("../src/stokes/Stokes2D.jl") - - @inline backend(::CuArray) = CUDABackendTrait() - @inline backend(::Type{<:CuArray}) = CUDABackendTrait() - - # Types - function JR2D.StokesArrays( - ::Type{CUDABackend}, ni::NTuple{N,Integer} - ) where {N} - return StokesArrays(ni) - end - - function JR2D.ThermalArrays( - ::Type{CUDABackend}, ni::NTuple{N,Number} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.ThermalArrays( - ::Type{CUDABackend}, ni::Vararg{Number,N} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) - end - - function JR2D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - function JR2D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - # Boundary conditions - function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function JR2D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - # Phases - function JR2D.phase_ratios_center( - ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases - ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) - end - - # Rheology - ## viscosity - function JR2D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - ## Stress - function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) - end - - ## Buoyancy forces - function JR2D.compute_ρg!(ρg::CuArray, rheology, args) - return compute_ρg!(ρg, rheology, args) - end - function JR2D.compute_ρg!( - ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args - ) - return compute_ρg!(ρg, phase_ratios, rheology, args) - end - - # Interpolations - function JR2D.temperature2center!( - ::CUDABackendTrait, thermal::JustRelax.ThermalArrays - ) - return _temperature2center!(thermal) - end - - function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) - return _temperature2center!(thermal) - end - - function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} - return vertex2center!(center, vertex) - end - - function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} - return center2vertex!(vertex, center) - end - - function JR2D.center2vertex!( - vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T - ) where {T<:CuArray} - return center2vertex!( - vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy - ) - end - - # Solvers - function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) - return _solve!(stokes, args...; kwargs...) - end - - function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) - return _heatdiffusion_PT!(thermal, args...; kwargs...) - end - - # Utils - function JR2D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) - end - -end +include("../src/ext/CUDA/2D.jl") +include("../src/ext/CUDA/3D.jl") end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl new file mode 100644 index 00000000..3ad6845b --- /dev/null +++ b/src/ext/CUDA/2D.jl @@ -0,0 +1,189 @@ +module JustRelax2D + + using JustRelax: JustRelax + using CUDA + using StaticArrays + using CellArrays + using ParallelStencil, ParallelStencil.FiniteDifferences2D + using ImplicitGlobalGrid + using GeoParams, LinearAlgebra, Printf + using MPI + + import JustRelax.JustRelax2D as JR2D + + import JustRelax: + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + + @init_parallel_stencil(CUDA, Float64, 2) + + include("../../common.jl") + include("../../stokes/Stokes2D.jl") + + # Types + function JR2D.StokesArrays( + ::Type{CUDABackend}, ni::NTuple{N,Integer} + ) where {N} + return StokesArrays(ni) + end + + function JR2D.ThermalArrays( + ::Type{CUDABackend}, ni::NTuple{N,Number} + ) where {N} + return ThermalArrays(ni...) + end + + function JR2D.ThermalArrays( + ::Type{CUDABackend}, ni::Vararg{Number,N} + ) where {N} + return ThermalArrays(ni...) + end + + function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) + end + + function JR2D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + function JR2D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + # Boundary conditions + function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function JR2D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + # Phases + function JR2D.phase_ratios_center( + ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases + ) + return _phase_ratios_center(phase_ratios, particles, grid, phases) + end + + # Rheology + ## viscosity + function JR2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function JR2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function JR2D.compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + ## Stress + function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) + end + + ## Buoyancy forces + function JR2D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) + end + function JR2D.compute_ρg!( + ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ) + return compute_ρg!(ρg, phase_ratios, rheology, args) + end + + # Interpolations + function JR2D.temperature2center!( + ::CUDABackendTrait, thermal::JustRelax.ThermalArrays + ) + return _temperature2center!(thermal) + end + + function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) + end + + function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) + end + + function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + return center2vertex!(vertex, center) + end + + function JR2D.center2vertex!( + vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T + ) where {T<:CuArray} + return center2vertex!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + end + + # Solvers + function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) + end + + function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) + end + + # Utils + function JR2D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) + end + +end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl new file mode 100644 index 00000000..22e6cc94 --- /dev/null +++ b/src/ext/CUDA/3D.jl @@ -0,0 +1,189 @@ +module JustRelax3D + + using JustRelax: JustRelax + using CUDA + using StaticArrays + using CellArrays + using ParallelStencil, ParallelStencil.FiniteDifferences3D + using ImplicitGlobalGrid + using GeoParams, LinearAlgebra, Printf + using MPI + + import JustRelax.JustRelax3D as JR3D + + import JustRelax: + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + + @init_parallel_stencil(CUDA, Float64, 3) + + include("../../common.jl") + include("../../stokes/Stokes3D.jl") + + # Types + function JR3D.StokesArrays( + ::Type{CUDABackend}, ni::NTuple{N,Integer} + ) where {N} + return StokesArrays(ni) + end + + function JR3D.ThermalArrays( + ::Type{CUDABackend}, ni::NTuple{N,Number} + ) where {N} + return ThermalArrays(ni...) + end + + function JR3D.ThermalArrays( + ::Type{CUDABackend}, ni::Vararg{Number,N} + ) where {N} + return ThermalArrays(ni...) + end + + function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) + end + + function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + # Boundary conditions + function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + # Phases + function JR3D.phase_ratios_center( + ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases + ) + return _phase_ratios_center(phase_ratios, particles, grid, phases) + end + + # Rheology + ## viscosity + function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function JR3D.compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + ## Stress + function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) + end + + ## Buoyancy forces + function JR3D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) + end + function JR3D.compute_ρg!( + ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ) + return compute_ρg!(ρg, phase_ratios, rheology, args) + end + + # Interpolations + function JR3D.temperature2center!( + ::CUDABackendTrait, thermal::JustRelax.ThermalArrays + ) + return _temperature2center!(thermal) + end + + function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) + end + + function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) + end + + 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<:CuArray} + return center2vertex!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + end + + # Solvers + function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) + end + + function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) + end + + # Utils + function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) + end + +end diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 58582845..74f2c4bc 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -24,7 +24,7 @@ solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kw function _solve!( stokes::JustRelax.StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{3,T}, flow_bcs, ρg, @@ -158,7 +158,7 @@ end function _solve!( stokes::JustRelax.StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{3,T}, flow_bcs::FlowBoundaryConditions, ρg, @@ -331,7 +331,7 @@ end # GeoParams and multiple phases function _solve!( stokes::JustRelax.StokesArrays, - pt_stokes::PTStokesCoeffs, + pt_stokes, di::NTuple{3,T}, flow_bc::FlowBoundaryConditions, ρg, From 01b7c090b25fcfc1591ee75c667618504bd4a555 Mon Sep 17 00:00:00 2001 From: aelligp Date: Tue, 30 Apr 2024 11:07:36 +0200 Subject: [PATCH 60/91] update `README.md` --- README.md | 134 ++++++++++++++++++++++++++++-------------------------- 1 file changed, 69 insertions(+), 65 deletions(-) diff --git a/README.md b/README.md index 8b94ef17..d82c541e 100644 --- a/README.md +++ b/README.md @@ -58,19 +58,17 @@ The test will take a while, so grab a :coffee: or :tea: ## Example: shear band localisation (2D) -![ShearBand2D](miniapps/benchmarks/stokes2D/shear_band/movies/DP_nx2058_2D.gif) +![ShearBand2D](docs/src/assets/movies/DP_nx2058_2D.gif) This example displays how the package can be used to simulate shear band localisation. The example is based on the [ShearBands2D.jl](miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl). ```julia -using GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO +using GeoParams, CellArrays, GLMakie +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +const backend = CPUBackend # HELPER FUNCTIONS --------------------------------------------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -101,30 +99,33 @@ function init_phases!(phase_ratios, xci, radius) end ``` -JustRelax allows to setup a model environment `PS_Setup` (and interplay with the underlying ParallelStencil package) to specify the dimension of the problem (2D or 3D) and the backend (CPU or GPU). The `PS_setup` functions takes `device`, `precision` and `dimensions` as argument: +# Initialize packages +Load JustRelax necessary modules and define backend. ```julia - model = PS_Setup(:Threads, Float64, 2) #running on the CPU in 2D - environment!(model) - - model = PS_Setup(:CUDA, Float64, 2) #running on an NVIDIA GPU in 2D - environment!(model) - - model = PS_Setup(:AMDGPU, Float64, 2) #running on an AMD GPU in 2D - environment!(model) +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend ``` -If you therefore want to run a 3D code, change the `dimensions` to 3 in the commands above. For this specific example we use particles to define the material phases, for which we rely on [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl). As in `JustRelax.jl`, we need to set up the environment of `JustPIC.jl`. This is done by running/including the following commands: - ```julia using JustPIC using JustPIC._2D - const backend = CPUBackend # Threads is the default backend - const backend = CUDABackend # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script - const backend = AMDGPUBackend # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. + const backend = JustPIC.CPUBackend # Threads is the default backend + const backend = JustPIC.CUDABackend # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script + const backend = JustPIC.AMDGPUBackend # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. +``` + +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 ``` For the initial setup, you will need to specify the number of nodes in x- and y- direction `nx` and `ny` as well as the directory where the figures are stored (`figdir`). The initialisation of the global grid and MPI environment is done with `igg = IGG(init_global_grid(nx, ny, 0; init_MPI = true)...)`: @@ -138,18 +139,19 @@ figdir = "ShearBands2D" igg = IGG(init_global_grid(nx, ny, 0; init_MPI = true)...) ``` -Initialisation of the physical domain and the grid. As `JustRelax.jl` relies on [ImplicitGlobalGrid.jl](https://github.com/omlins/ImplicitGlobalGrid.jl), the grid can be `MPIAWARE` through setting the grid steps in x- and y- direction to ` di = @. li / (nx_g(),ny_g())`. This makes it a global grid and the grid steps are automatically distributed over the MPI processes. +Initialisation of the physical domain and the grid. As `JustRelax.jl` relies on [ImplicitGlobalGrid.jl](https://github.com/omlins/ImplicitGlobalGrid.jl), the grid can be `MPIAWARE`. This makes it a global grid and the grid steps are automatically distributed over the MPI processes. ```julia # Physical domain ------------------------------------ -ly = 1e0 # 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 -xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells -dt = Inf +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 ``` Initialisation of the rheology with [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl). The rheology can be tailored to the specific problem with different creep laws and material parameters (see [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl)) or the miniapps in the [convection folder](miniapps/convection). @@ -196,22 +198,19 @@ Initialisation of the Stokes arrays and the necessary allocations. The rheology ```julia # Initialize phase ratios ------------------------------- radius = 0.1 -phase_ratios = PhaseRatio(ni, length(rheology)) +phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem -stokes = StokesArrays(ni, ViscoElastic) +stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # PT coefficients after Räss, L., Utkin, I., Duretz, T., Omlin, S., and Podladchikov, Y. Y.: Assessing the robustness and scalability of the accelerated pseudo-transient method, Geosci. Model Dev., 15, 5757–5786, https://doi.org/10.5194/gmd-15-5757-2022, 2022. # Buoyancy forces -ρg = @zeros(ni...), @zeros(ni...) -args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) -# Viscosity -η = @ones(ni...) -η_vep = similar(η) # effective visco-elasto-plastic viscosity -@parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) -) +ρ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)) ``` Define pure shear velocity boundary conditions @@ -238,26 +237,26 @@ t, it, tmax = 0.0, 0, 3.5 while t < tmax # Stokes solver ---------------- - solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - η, - η_vep, - phase_ratios, - rheology, - args, - dt, - igg; - verbose = false, - iterMax = 500e3, - nout = 1e3, - viscosity_cutoff = (-Inf, Inf) - ) + 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 + ) + ) # Compute second invariant of the strain rate tensor - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) # Update old stresses @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @@ -267,22 +266,24 @@ while t < tmax it += 1 t += dt + push!(sol, solution(εbg, t, G0, η0)) push!(ttot, t) println("it = $it; t = $t \n") - # visualisation + # 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 = "τII") - ax2 = Axis(fig[2,1], aspect = 1, title = "η_vep") - ax3 = Axis(fig[1,2], aspect = 1, title = "log10(εII)") + 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.(η_vep)) , 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) @@ -290,21 +291,24 @@ while t < tmax hidexdecorations!(ax1) hidexdecorations!(ax3) save(joinpath(figdir, "$(it).png"), fig) + fig end ``` ## Miniapps -Currently there are 3 convection miniapps with particles and 3 corresponding miniapps without. The miniapps with particles are: +Currently there are 4 convection miniapps with particles and 4 corresponding miniapps without. The miniapps with particles are: * [Layered_convection2D.jl](miniapps/convection/Particles2D/Layered_convection2D.jl) + * [Layered_convection2D_nonDim.jl](miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl) * [Layered_convection3D.jl](miniapps/convection/Particles3D/Layered_convection3D.jl) * [WENO_convection2D.jl](miniapps/convection/WENO5/WENO_convection2D.jl) The miniapps without particles are: * [GlobalConvection2D_Upwind.jl](miniapps/convection/GlobalConvection2D_Upwind.jl) - * [GlobalConvection3D_Upwind.jl](miniapps/convection/GlobalConvection3D_Upwind.jl) * [GlobalConvection2D_WENO5.jl](miniapps/convection/GlobalConvection2D_WENO5.jl) + * [GlobalConvection2D_WENO5_MPI.jl](miniapps/convection/GlobalConvection2D_WENO5_MPI.jl) + * [GlobalConvection3D_Upwind.jl](miniapps/convection/GlobalConvection3D_Upwind.jl) ## Benchmarks From 26e29d473c8462a1865348d818c643ca9a94540e Mon Sep 17 00:00:00 2001 From: aelligp Date: Tue, 30 Apr 2024 11:08:19 +0200 Subject: [PATCH 61/91] add ShearBands example to Doc; move shearband movie --- docs/make.jl | 9 +- .../src/assets}/movies/DP_nx2058_2D.gif | Bin docs/src/man/ShearBands.md | 238 ++++++++++++++++++ 3 files changed, 243 insertions(+), 4 deletions(-) rename {miniapps/benchmarks/stokes2D/shear_band => docs/src/assets}/movies/DP_nx2058_2D.gif (100%) create mode 100644 docs/src/man/ShearBands.md diff --git a/docs/make.jl b/docs/make.jl index 85b05772..9159dcc7 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -14,12 +14,13 @@ makedocs(; "Home" => "man/index.md", "User guide"=> Any[ "Installation" => "man/installation.md", - ], - "Backend" => "man/backend.md", - "Equations" => "man/equations.md", - "Advection" => "man/advection.md", + "Backend" => "man/backend.md", + "Equations" => "man/equations.md", + "Advection" => "man/advection.md", + ], "Examples" => Any[ "Blankenbach" => "man/Blankenbach.md", + "Shear Bands" => "man/ShearBands.md", ], "List of functions" => "man/listfunctions.md", ], diff --git a/miniapps/benchmarks/stokes2D/shear_band/movies/DP_nx2058_2D.gif b/docs/src/assets/movies/DP_nx2058_2D.gif similarity index 100% rename from miniapps/benchmarks/stokes2D/shear_band/movies/DP_nx2058_2D.gif rename to docs/src/assets/movies/DP_nx2058_2D.gif diff --git a/docs/src/man/ShearBands.md b/docs/src/man/ShearBands.md new file mode 100644 index 00000000..a2eddd3c --- /dev/null +++ b/docs/src/man/ShearBands.md @@ -0,0 +1,238 @@ +# ShearBand benchmark + +Shear Band benchmark to test the visco-elasto-plastic rheology implementation in JustRelax.jl + +# Initialize packages + +Load JustRelax necessary modules and define backend. +```julia +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = JustRelax.CPUBackend +``` + + +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 (JustRelax.@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 +Shear Bands evolution in a 2D visco-elasto-plastic rheology model +![Shearbands](../assets/movies/DP_nx2058_2D.gif) From 57023a72437ec683d7bfee5c015317daac65d343 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 11:37:14 +0200 Subject: [PATCH 62/91] AMDGPU ext & some polishing --- ext/JustRelaxAMDGPUExt.jl | 15 +++ ext/JustRelaxCUDAExt.jl | 4 +- src/ext/AMDGPU/2D.jl | 189 ++++++++++++++++++++++++++++++++++++++ src/ext/AMDGPU/3D.jl | 189 ++++++++++++++++++++++++++++++++++++++ src/types/traits.jl | 11 ++- 5 files changed, 402 insertions(+), 6 deletions(-) create mode 100644 ext/JustRelaxAMDGPUExt.jl create mode 100644 src/ext/AMDGPU/2D.jl create mode 100644 src/ext/AMDGPU/3D.jl diff --git a/ext/JustRelaxAMDGPUExt.jl b/ext/JustRelaxAMDGPUExt.jl new file mode 100644 index 00000000..517cfdef --- /dev/null +++ b/ext/JustRelaxAMDGPUExt.jl @@ -0,0 +1,15 @@ +module JustRelaxAMDGPUExt + +using AMDGPU +using JustRelax: JustRelax +import JustRelax: PTArray, backend, AMDGPUBackendTrait + +PTArray(::Type{AMDGPUBackend}) = RocArray + +@inline backend(::CuArray) = AMDGPUBackendTrait() +@inline backend(::Type{<:CuArray}) = AMDGPUBackendTrait() + +include("../src/ext/AMDGPU/2D.jl") +include("../src/ext/AMDGPU/3D.jl") + +end diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl index 978b8041..6b0e5bcd 100644 --- a/ext/JustRelaxCUDAExt.jl +++ b/ext/JustRelaxCUDAExt.jl @@ -2,9 +2,9 @@ module JustRelaxCUDAExt using CUDA using JustRelax: JustRelax -import JustRelax: PTArray +import JustRelax: PTArray, backend, CUDABackendTrait -JustRelax.PTArray(::Type{CUDABackend}) = CuArray +PTArray(::Type{CUDABackend}) = CuArray @inline backend(::CuArray) = CUDABackendTrait() @inline backend(::Type{<:CuArray}) = CUDABackendTrait() diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl new file mode 100644 index 00000000..daa103a0 --- /dev/null +++ b/src/ext/AMDGPU/2D.jl @@ -0,0 +1,189 @@ +module JustRelax2D + + using JustRelax: JustRelax + using AMDGPU + using StaticArrays + using CellArrays + using ParallelStencil, ParallelStencil.FiniteDifferences2D + using ImplicitGlobalGrid + using GeoParams, LinearAlgebra, Printf + using MPI + + import JustRelax.JustRelax2D as JR2D + + import JustRelax: + IGG, BackendTrait, CPUBackendTrait, AMDGPUBackendTrait, backend, CPUBackend, Geometry, @cell + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + + @init_parallel_stencil(AMDGPU, Float64, 2) + + include("../../common.jl") + include("../../stokes/Stokes2D.jl") + + # Types + function JR2D.StokesArrays( + ::Type{AMDGPUBackend}, ni::NTuple{N,Integer} + ) where {N} + return StokesArrays(ni) + end + + function JR2D.ThermalArrays( + ::Type{AMDGPUBackend}, ni::NTuple{N,Number} + ) where {N} + return ThermalArrays(ni...) + end + + function JR2D.ThermalArrays( + ::Type{AMDGPUBackend}, ni::Vararg{Number,N} + ) where {N} + return ThermalArrays(ni...) + end + + function JR2D.PhaseRatio(::Type{AMDGPUBackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) + end + + function JR2D.PTThermalCoeffs( + ::Type{AMDGPUBackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + function JR2D.PTThermalCoeffs( + ::Type{AMDGPUBackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + # Boundary conditions + function JR2D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function JR2D.thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + function thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + # Phases + function JR2D.phase_ratios_center( + ::AMDGPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases + ) + return _phase_ratios_center(phase_ratios, particles, grid, phases) + end + + # Rheology + ## viscosity + function JR2D.compute_viscosity!( + ::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function JR2D.compute_viscosity!( + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function JR2D.compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function compute_viscosity!( + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + ## Stress + function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) + end + + ## Buoyancy forces + function JR2D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) + end + function JR2D.compute_ρg!( + ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ) + return compute_ρg!(ρg, phase_ratios, rheology, args) + end + + # Interpolations + function JR2D.temperature2center!( + ::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays + ) + return _temperature2center!(thermal) + end + + function temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) + end + + function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) + end + + function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + return center2vertex!(vertex, center) + end + + function JR2D.center2vertex!( + vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T + ) where {T<:CuArray} + return center2vertex!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + end + + # Solvers + function JR2D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) + end + + function JR2D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) + end + + # Utils + function JR2D.compute_dt(::AMDGPUBackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) + end + +end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl new file mode 100644 index 00000000..22e6cc94 --- /dev/null +++ b/src/ext/AMDGPU/3D.jl @@ -0,0 +1,189 @@ +module JustRelax3D + + using JustRelax: JustRelax + using CUDA + using StaticArrays + using CellArrays + using ParallelStencil, ParallelStencil.FiniteDifferences3D + using ImplicitGlobalGrid + using GeoParams, LinearAlgebra, Printf + using MPI + + import JustRelax.JustRelax3D as JR3D + + import JustRelax: + IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell + import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + + @init_parallel_stencil(CUDA, Float64, 3) + + include("../../common.jl") + include("../../stokes/Stokes3D.jl") + + # Types + function JR3D.StokesArrays( + ::Type{CUDABackend}, ni::NTuple{N,Integer} + ) where {N} + return StokesArrays(ni) + end + + function JR3D.ThermalArrays( + ::Type{CUDABackend}, ni::NTuple{N,Number} + ) where {N} + return ThermalArrays(ni...) + end + + function JR3D.ThermalArrays( + ::Type{CUDABackend}, ni::Vararg{Number,N} + ) where {N} + return ThermalArrays(ni...) + end + + function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) + end + + function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, + ) where {nDim,T} + return PTThermalCoeffs( + rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL + ) + end + + # Boundary conditions + function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) + end + + function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) + end + + # Phases + function JR3D.phase_ratios_center( + ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases + ) + return _phase_ratios_center(phase_ratios, particles, grid, phases) + end + + # Rheology + ## viscosity + function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function JR3D.compute_viscosity!( + η, ν, εII::CuArray, args, rheology, cutoff + ) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) + end + function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) + end + + ## Stress + function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + _tensor_invariant!(A) + end + + ## Buoyancy forces + function JR3D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) + end + function JR3D.compute_ρg!( + ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args + ) + return compute_ρg!(ρg, phase_ratios, rheology, args) + end + + # Interpolations + function JR3D.temperature2center!( + ::CUDABackendTrait, thermal::JustRelax.ThermalArrays + ) + return _temperature2center!(thermal) + end + + function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) + end + + function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) + end + + 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<:CuArray} + return center2vertex!( + vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy + ) + end + + # Solvers + function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) + end + + function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) + end + + # Utils + function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + _compute_dt(S, args...) + end + +end diff --git a/src/types/traits.jl b/src/types/traits.jl index 20431a17..3eaf37de 100644 --- a/src/types/traits.jl +++ b/src/types/traits.jl @@ -4,16 +4,19 @@ struct NonCPUBackendTrait <: BackendTrait end struct CUDABackendTrait <: BackendTrait end struct AMDGPUBackendTrait <: BackendTrait end +# AbstractArray's @inline backend(::Array) = CPUBackendTrait() @inline backend(::Type{<:Array}) = CPUBackendTrait() -@inline backend(::AbstractArray) = NonCPUDeviceTrait() -@inline backend(::Type{<:AbstractArray}) = NonCPUDeviceTrait() - -@inline backend(::T) where {T} = throw(ArgumentError("Backend $(T) not supported")) +@inline backend(::AbstractArray) = NonCPUBackendTrait() +@inline backend(::Type{<:AbstractArray}) = NonCPUBackendTrait() +# Custom struct's @inline backend(::JustRelax.Velocity{T}) where {T} = backend(T) @inline backend(::JustRelax.SymmetricTensor{T}) where {T} = backend(T) @inline backend(::JustRelax.Residual{T}) where {T} = backend(T) @inline backend(::JustRelax.ThermalArrays{T}) where {T} = backend(T) @inline backend(x::JustRelax.StokesArrays) = backend(x.P) @inline backend(x::JustRelax.PhaseRatio) = backend(x.center.data) + +# Error handling +@inline backend(::T) where {T} = throw(ArgumentError("$(T) is not a supported backend")) From accf00c905b98d4feb71a99c60032a299d389b91 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 11:44:07 +0200 Subject: [PATCH 63/91] docs --- docs/src/man/ShearBands.md | 4 ++-- docs/src/man/index.md | 16 +++++----------- docs/src/man/installation.md | 2 +- 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/docs/src/man/ShearBands.md b/docs/src/man/ShearBands.md index a2eddd3c..c16cce35 100644 --- a/docs/src/man/ShearBands.md +++ b/docs/src/man/ShearBands.md @@ -40,7 +40,7 @@ grid = Geometry(ni, li; origin = origin) dt = Inf ``` -## Physical properties using GeoParams ---------------- +## Physical properties using GeoParams ```julia τ_y = 1.6 # yield stress. If do_DP=true, τ_y stand for the cohesion: c*cos(ϕ) ϕ = 30 # friction angle @@ -104,7 +104,7 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) end ``` diff --git a/docs/src/man/index.md b/docs/src/man/index.md index 263c471c..95a5ed48 100644 --- a/docs/src/man/index.md +++ b/docs/src/man/index.md @@ -1,22 +1,16 @@ -```@meta -CurrentModule = JustRelax -``` - # JustRelax.jl -Documentation for [JustRelax.jl](https://github.com/PTsolvers/JustRelax.jl) - Need to solve a very large multi-physics problem on many GPUs in parallel? Just Relax! -`JustRelax.jl` is a collection of accelerated iterative pseudo-transient solvers using MPI and multiple CPU or GPU backends. It's part of the [PTSolvers organisation](https://ptsolvers.github.io) and +`JustRelax` is a collection of accelerated iterative pseudo-transient solvers using MPI and multiple CPU or GPU backends. It's part of the [PTSolvers organisation](https://ptsolvers.github.io) and developed within the [GPU4GEO project](https://www.pasc-ch.org/projects/2021-2024/gpu4geo/). Current publications, outreach and news can be found on the [GPU4GEO website](https://ptsolvers.github.io/GPU4GEO/). The package relies on other packages as building blocks and parallelisation tools: -* [ParallelStencil.jl](https://github.com/omlins/ParallelStencil.jl) -* [ImplicitGlobalGrid.jl](https://github.com/omlins/ImplicitGlobalGrid.jl) -* [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl) -* [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl) +* [ParallelStencil.jl](https://github.com/omlins/ParallelStencil.jl) - device agnostic parallel kernels +* [ImplicitGlobalGrid.jl](https://github.com/omlins/ImplicitGlobalGrid.jl) - (CUDA-aware) MPI communication +* [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl) - Material physics +* [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl) - Particle-in-cell advection The package serves several purposes: diff --git a/docs/src/man/installation.md b/docs/src/man/installation.md index 330e9090..2d62ff47 100644 --- a/docs/src/man/installation.md +++ b/docs/src/man/installation.md @@ -1,6 +1,6 @@ # Installation -`JustRelax.jl` is a registered package and can be added as follows: +`JustRelax` is a registered package and can be added as follows: ```julia using Pkg; Pkg.add("JustRelax") From 27f74806d3e7c4e2ac9afff89e82cad650d668d0 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 11:49:35 +0200 Subject: [PATCH 64/91] formatting --- src/JustRelax_CPU.jl | 6 +- src/Utils.jl | 2 +- src/boundaryconditions/types.jl | 2 +- src/ext/AMDGPU/2D.jl | 362 ++++++++++++++++---------------- src/ext/AMDGPU/3D.jl | 362 ++++++++++++++++---------------- src/ext/CUDA/2D.jl | 362 ++++++++++++++++---------------- src/ext/CUDA/3D.jl | 362 ++++++++++++++++---------------- src/stokes/StressKernels.jl | 2 +- 8 files changed, 715 insertions(+), 745 deletions(-) diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl index a859501e..a75d0970 100644 --- a/src/JustRelax_CPU.jl +++ b/src/JustRelax_CPU.jl @@ -10,7 +10,8 @@ using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend import JustRelax: PTStokesCoeffs -import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions @init_parallel_stencil(Threads, Float64, 2) @@ -32,7 +33,8 @@ using MPI import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend import JustRelax: PTStokesCoeffs -import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions @init_parallel_stencil(Threads, Float64, 3) diff --git a/src/Utils.jl b/src/Utils.jl index 9c6c65a2..6859d1de 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -404,7 +404,7 @@ function compute_dt(S::JustRelax.StokesArrays, args...) end function compute_dt(::CPUBackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) + return _compute_dt(S, args...) end """ diff --git a/src/boundaryconditions/types.jl b/src/boundaryconditions/types.jl index 768b8dd8..7eec4792 100644 --- a/src/boundaryconditions/types.jl +++ b/src/boundaryconditions/types.jl @@ -25,4 +25,4 @@ struct FlowBoundaryConditions{T,nD} <: AbstractBoundaryConditions nD = length(no_slip) == 4 ? 2 : 3 return new{T,nD}(no_slip, free_slip, free_surface) end -end \ No newline at end of file +end diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index daa103a0..a85433b2 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -1,189 +1,181 @@ module JustRelax2D - using JustRelax: JustRelax - using AMDGPU - using StaticArrays - using CellArrays - using ParallelStencil, ParallelStencil.FiniteDifferences2D - using ImplicitGlobalGrid - using GeoParams, LinearAlgebra, Printf - using MPI - - import JustRelax.JustRelax2D as JR2D - - import JustRelax: - IGG, BackendTrait, CPUBackendTrait, AMDGPUBackendTrait, backend, CPUBackend, Geometry, @cell - import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions - - @init_parallel_stencil(AMDGPU, Float64, 2) - - include("../../common.jl") - include("../../stokes/Stokes2D.jl") - - # Types - function JR2D.StokesArrays( - ::Type{AMDGPUBackend}, ni::NTuple{N,Integer} - ) where {N} - return StokesArrays(ni) - end - - function JR2D.ThermalArrays( - ::Type{AMDGPUBackend}, ni::NTuple{N,Number} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.ThermalArrays( - ::Type{AMDGPUBackend}, ni::Vararg{Number,N} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.PhaseRatio(::Type{AMDGPUBackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) - end - - function JR2D.PTThermalCoeffs( - ::Type{AMDGPUBackend}, - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - function JR2D.PTThermalCoeffs( - ::Type{AMDGPUBackend}, - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - # Boundary conditions - function JR2D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function JR2D.thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - function thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - # Phases - function JR2D.phase_ratios_center( - ::AMDGPUBackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases - ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) - end - - # Rheology - ## viscosity - function JR2D.compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - ## Stress - function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) - end - - ## Buoyancy forces - function JR2D.compute_ρg!(ρg::CuArray, rheology, args) - return compute_ρg!(ρg, rheology, args) - end - function JR2D.compute_ρg!( - ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args - ) - return compute_ρg!(ρg, phase_ratios, rheology, args) - end - - # Interpolations - function JR2D.temperature2center!( - ::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays - ) - return _temperature2center!(thermal) - end - - function temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays) - return _temperature2center!(thermal) - end - - function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} - return vertex2center!(center, vertex) - end - - function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} - return center2vertex!(vertex, center) - end - - function JR2D.center2vertex!( - vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T - ) where {T<:CuArray} - return center2vertex!( - vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy - ) - end - - # Solvers - function JR2D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs) - return _solve!(stokes, args...; kwargs...) - end - - function JR2D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs) - return _heatdiffusion_PT!(thermal, args...; kwargs...) - end - - # Utils - function JR2D.compute_dt(::AMDGPUBackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) - end +using JustRelax: JustRelax +using AMDGPU +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences2D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf +using MPI + +import JustRelax.JustRelax2D as JR2D + +import JustRelax: + IGG, + BackendTrait, + CPUBackendTrait, + AMDGPUBackendTrait, + backend, + CPUBackend, + Geometry, + @cell +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(AMDGPU, Float64, 2) + +include("../../common.jl") +include("../../stokes/Stokes2D.jl") + +# Types +function JR2D.StokesArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Integer}) where {N} + return StokesArrays(ni) +end + +function JR2D.ThermalArrays(::Type{AMDGPUBackend}, ni::NTuple{N,Number}) where {N} + return ThermalArrays(ni...) +end + +function JR2D.ThermalArrays(::Type{AMDGPUBackend}, ni::Vararg{Number,N}) where {N} + return ThermalArrays(ni...) +end + +function JR2D.PhaseRatio(::Type{AMDGPUBackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) +end + +function JR2D.PTThermalCoeffs( + ::Type{AMDGPUBackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +function JR2D.PTThermalCoeffs( + ::Type{AMDGPUBackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +# Boundary conditions +function JR2D.flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function flow_bcs!(::AMDGPUBackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function JR2D.thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +function thermal_bcs!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +# Phases +function JR2D.phase_ratios_center( + ::AMDGPUBackendTrait, + phase_ratios::JustRelax.PhaseRatio, + particles, + grid::Geometry, + phases, +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +# Rheology +## viscosity +function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function JR2D.compute_viscosity!( + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function compute_viscosity!( + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +## Stress +function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTensor) + return _tensor_invariant!(A) +end + +## Buoyancy forces +function JR2D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end + +# Interpolations +function JR2D.temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) +end + +function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + return center2vertex!(vertex, center) +end + +function JR2D.center2vertex!( + vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T +) where {T<:CuArray} + return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) +end + +# Solvers +function JR2D.solve!(::AMDGPUBackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) +end + +function JR2D.heatdiffusion_PT!(::AMDGPUBackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) +end + +# Utils +function JR2D.compute_dt(::AMDGPUBackendTrait, S::JustRelax.StokesArrays, args...) + return _compute_dt(S, args...) +end end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 22e6cc94..eb4cb4f3 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -1,189 +1,181 @@ module JustRelax3D - using JustRelax: JustRelax - using CUDA - using StaticArrays - using CellArrays - using ParallelStencil, ParallelStencil.FiniteDifferences3D - using ImplicitGlobalGrid - using GeoParams, LinearAlgebra, Printf - using MPI - - import JustRelax.JustRelax3D as JR3D - - import JustRelax: - IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell - import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions - - @init_parallel_stencil(CUDA, Float64, 3) - - include("../../common.jl") - include("../../stokes/Stokes3D.jl") - - # Types - function JR3D.StokesArrays( - ::Type{CUDABackend}, ni::NTuple{N,Integer} - ) where {N} - return StokesArrays(ni) - end - - function JR3D.ThermalArrays( - ::Type{CUDABackend}, ni::NTuple{N,Number} - ) where {N} - return ThermalArrays(ni...) - end - - function JR3D.ThermalArrays( - ::Type{CUDABackend}, ni::Vararg{Number,N} - ) where {N} - return ThermalArrays(ni...) - end - - function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) - end - - function JR3D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - function JR3D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - # Boundary conditions - function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - # Phases - function JR3D.phase_ratios_center( - ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases - ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) - end - - # Rheology - ## viscosity - function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function JR3D.compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - ## Stress - function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) - end - - ## Buoyancy forces - function JR3D.compute_ρg!(ρg::CuArray, rheology, args) - return compute_ρg!(ρg, rheology, args) - end - function JR3D.compute_ρg!( - ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args - ) - return compute_ρg!(ρg, phase_ratios, rheology, args) - end - - # Interpolations - function JR3D.temperature2center!( - ::CUDABackendTrait, thermal::JustRelax.ThermalArrays - ) - return _temperature2center!(thermal) - end - - function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) - return _temperature2center!(thermal) - end - - function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} - return vertex2center!(center, vertex) - end - - 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<:CuArray} - return center2vertex!( - vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy - ) - end - - # Solvers - function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) - return _solve!(stokes, args...; kwargs...) - end - - function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) - return _heatdiffusion_PT!(thermal, args...; kwargs...) - end - - # Utils - function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) - end +using JustRelax: JustRelax +using CUDA +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences3D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf +using MPI + +import JustRelax.JustRelax3D as JR3D + +import JustRelax: + IGG, + BackendTrait, + CPUBackendTrait, + CUDABackendTrait, + backend, + CPUBackend, + Geometry, + @cell +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(CUDA, Float64, 3) + +include("../../common.jl") +include("../../stokes/Stokes3D.jl") + +# Types +function JR3D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} + return StokesArrays(ni) +end + +function JR3D.ThermalArrays(::Type{CUDABackend}, ni::NTuple{N,Number}) where {N} + return ThermalArrays(ni...) +end + +function JR3D.ThermalArrays(::Type{CUDABackend}, ni::Vararg{Number,N}) where {N} + return ThermalArrays(ni...) +end + +function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) +end + +function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +# Boundary conditions +function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +# Phases +function JR3D.phase_ratios_center( + ::CUDABackendTrait, + phase_ratios::JustRelax.PhaseRatio, + particles, + grid::Geometry, + phases, +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +# Rheology +## viscosity +function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +## Stress +function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + return _tensor_invariant!(A) +end + +## Buoyancy forces +function JR3D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end + +# Interpolations +function JR3D.temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) +end + +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<:CuArray} + return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) +end + +# Solvers +function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) +end + +function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) +end + +# Utils +function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + return _compute_dt(S, args...) +end end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 3ad6845b..0dacf8b6 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -1,189 +1,181 @@ module JustRelax2D - using JustRelax: JustRelax - using CUDA - using StaticArrays - using CellArrays - using ParallelStencil, ParallelStencil.FiniteDifferences2D - using ImplicitGlobalGrid - using GeoParams, LinearAlgebra, Printf - using MPI - - import JustRelax.JustRelax2D as JR2D - - import JustRelax: - IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell - import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions - - @init_parallel_stencil(CUDA, Float64, 2) - - include("../../common.jl") - include("../../stokes/Stokes2D.jl") - - # Types - function JR2D.StokesArrays( - ::Type{CUDABackend}, ni::NTuple{N,Integer} - ) where {N} - return StokesArrays(ni) - end - - function JR2D.ThermalArrays( - ::Type{CUDABackend}, ni::NTuple{N,Number} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.ThermalArrays( - ::Type{CUDABackend}, ni::Vararg{Number,N} - ) where {N} - return ThermalArrays(ni...) - end - - function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) - end - - function JR2D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - function JR2D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - # Boundary conditions - function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function JR2D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - # Phases - function JR2D.phase_ratios_center( - ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases - ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) - end - - # Rheology - ## viscosity - function JR2D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function JR2D.compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - ## Stress - function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) - end - - ## Buoyancy forces - function JR2D.compute_ρg!(ρg::CuArray, rheology, args) - return compute_ρg!(ρg, rheology, args) - end - function JR2D.compute_ρg!( - ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args - ) - return compute_ρg!(ρg, phase_ratios, rheology, args) - end - - # Interpolations - function JR2D.temperature2center!( - ::CUDABackendTrait, thermal::JustRelax.ThermalArrays - ) - return _temperature2center!(thermal) - end - - function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) - return _temperature2center!(thermal) - end - - function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} - return vertex2center!(center, vertex) - end - - function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} - return center2vertex!(vertex, center) - end - - function JR2D.center2vertex!( - vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T - ) where {T<:CuArray} - return center2vertex!( - vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy - ) - end - - # Solvers - function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) - return _solve!(stokes, args...; kwargs...) - end - - function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) - return _heatdiffusion_PT!(thermal, args...; kwargs...) - end - - # Utils - function JR2D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) - end +using JustRelax: JustRelax +using CUDA +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences2D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf +using MPI + +import JustRelax.JustRelax2D as JR2D + +import JustRelax: + IGG, + BackendTrait, + CPUBackendTrait, + CUDABackendTrait, + backend, + CPUBackend, + Geometry, + @cell +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(CUDA, Float64, 2) + +include("../../common.jl") +include("../../stokes/Stokes2D.jl") + +# Types +function JR2D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} + return StokesArrays(ni) +end + +function JR2D.ThermalArrays(::Type{CUDABackend}, ni::NTuple{N,Number}) where {N} + return ThermalArrays(ni...) +end + +function JR2D.ThermalArrays(::Type{CUDABackend}, ni::Vararg{Number,N}) where {N} + return ThermalArrays(ni...) +end + +function JR2D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) +end + +function JR2D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +function JR2D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +# Boundary conditions +function JR2D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function JR2D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +# Phases +function JR2D.phase_ratios_center( + ::CUDABackendTrait, + phase_ratios::JustRelax.PhaseRatio, + particles, + grid::Geometry, + phases, +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +# Rheology +## viscosity +function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function JR2D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +## Stress +function JR2D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + return _tensor_invariant!(A) +end + +## Buoyancy forces +function JR2D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end + +# Interpolations +function JR2D.temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) +end + +function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} + return center2vertex!(vertex, center) +end + +function JR2D.center2vertex!( + vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T +) where {T<:CuArray} + return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) +end + +# Solvers +function JR2D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) +end + +function JR2D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) +end + +# Utils +function JR2D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + return _compute_dt(S, args...) +end end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 22e6cc94..eb4cb4f3 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -1,189 +1,181 @@ module JustRelax3D - using JustRelax: JustRelax - using CUDA - using StaticArrays - using CellArrays - using ParallelStencil, ParallelStencil.FiniteDifferences3D - using ImplicitGlobalGrid - using GeoParams, LinearAlgebra, Printf - using MPI - - import JustRelax.JustRelax3D as JR3D - - import JustRelax: - IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell - import JustRelax: AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions - - @init_parallel_stencil(CUDA, Float64, 3) - - include("../../common.jl") - include("../../stokes/Stokes3D.jl") - - # Types - function JR3D.StokesArrays( - ::Type{CUDABackend}, ni::NTuple{N,Integer} - ) where {N} - return StokesArrays(ni) - end - - function JR3D.ThermalArrays( - ::Type{CUDABackend}, ni::NTuple{N,Number} - ) where {N} - return ThermalArrays(ni...) - end - - function JR3D.ThermalArrays( - ::Type{CUDABackend}, ni::Vararg{Number,N} - ) where {N} - return ThermalArrays(ni...) - end - - function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) - return PhaseRatio(ni, num_phases) - end - - function JR3D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - phase_ratios, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - function JR3D.PTThermalCoeffs( - ::Type{CUDABackend}, - rheology, - args, - dt, - ni, - di::NTuple{nDim,T}, - li::NTuple{nDim,Any}; - ϵ=1e-8, - CFL=0.9 / √3, - ) where {nDim,T} - return PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL - ) - end - - # Boundary conditions - function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) - return _flow_bcs!(bcs, @velocity(stokes)) - end - - function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) - return thermal_bcs!(thermal.T, bcs) - end - - # Phases - function JR3D.phase_ratios_center( - ::CUDABackendTrait, phase_ratios::JustRelax.PhaseRatio, particles, grid::Geometry, phases - ) - return _phase_ratios_center(phase_ratios, particles, grid, phases) - end - - # Rheology - ## viscosity - function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function JR3D.compute_viscosity!( - η, ν, εII::CuArray, args, rheology, cutoff - ) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) - end - function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff - ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) - end - function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) - return compute_viscosity!(η, ν, εII, args, rheology, cutoff) - end - - ## Stress - function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) - end - - ## Buoyancy forces - function JR3D.compute_ρg!(ρg::CuArray, rheology, args) - return compute_ρg!(ρg, rheology, args) - end - function JR3D.compute_ρg!( - ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args - ) - return compute_ρg!(ρg, phase_ratios, rheology, args) - end - - # Interpolations - function JR3D.temperature2center!( - ::CUDABackendTrait, thermal::JustRelax.ThermalArrays - ) - return _temperature2center!(thermal) - end - - function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) - return _temperature2center!(thermal) - end - - function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} - return vertex2center!(center, vertex) - end - - 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<:CuArray} - return center2vertex!( - vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy - ) - end - - # Solvers - function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) - return _solve!(stokes, args...; kwargs...) - end - - function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) - return _heatdiffusion_PT!(thermal, args...; kwargs...) - end - - # Utils - function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) - _compute_dt(S, args...) - end +using JustRelax: JustRelax +using CUDA +using StaticArrays +using CellArrays +using ParallelStencil, ParallelStencil.FiniteDifferences3D +using ImplicitGlobalGrid +using GeoParams, LinearAlgebra, Printf +using MPI + +import JustRelax.JustRelax3D as JR3D + +import JustRelax: + IGG, + BackendTrait, + CPUBackendTrait, + CUDABackendTrait, + backend, + CPUBackend, + Geometry, + @cell +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(CUDA, Float64, 3) + +include("../../common.jl") +include("../../stokes/Stokes3D.jl") + +# Types +function JR3D.StokesArrays(::Type{CUDABackend}, ni::NTuple{N,Integer}) where {N} + return StokesArrays(ni) +end + +function JR3D.ThermalArrays(::Type{CUDABackend}, ni::NTuple{N,Number}) where {N} + return ThermalArrays(ni...) +end + +function JR3D.ThermalArrays(::Type{CUDABackend}, ni::Vararg{Number,N}) where {N} + return ThermalArrays(ni...) +end + +function JR3D.PhaseRatio(::Type{CUDABackend}, ni, num_phases) + return PhaseRatio(ni, num_phases) +end + +function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + phase_ratios, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, phase_ratios, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +function JR3D.PTThermalCoeffs( + ::Type{CUDABackend}, + rheology, + args, + dt, + ni, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return PTThermalCoeffs(rheology, args, dt, ni, di, li; ϵ=ϵ, CFL=CFL) +end + +# Boundary conditions +function JR3D.flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function flow_bcs!(::CUDABackendTrait, stokes::JustRelax.StokesArrays, bcs) + return _flow_bcs!(bcs, @velocity(stokes)) +end + +function JR3D.thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +function thermal_bcs!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays, bcs) + return thermal_bcs!(thermal.T, bcs) +end + +# Phases +function JR3D.phase_ratios_center( + ::CUDABackendTrait, + phase_ratios::JustRelax.PhaseRatio, + particles, + grid::Geometry, + phases, +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +# Rheology +## viscosity +function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end +function compute_viscosity!( + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end +function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) + return compute_viscosity!(η, ν, εII, args, rheology, cutoff) +end + +## Stress +function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor) + return _tensor_invariant!(A) +end + +## Buoyancy forces +function JR3D.compute_ρg!(ρg::CuArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) + return compute_ρg!(ρg, phase_ratios, rheology, args) +end + +# Interpolations +function JR3D.temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} + return vertex2center!(center, vertex) +end + +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<:CuArray} + return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) +end + +# Solvers +function JR3D.solve!(::CUDABackendTrait, stokes, args...; kwargs) + return _solve!(stokes, args...; kwargs...) +end + +function JR3D.heatdiffusion_PT!(::CUDABackendTrait, thermal, args...; kwargs) + return _heatdiffusion_PT!(thermal, args...; kwargs...) +end + +# Utils +function JR3D.compute_dt(::CUDABackendTrait, S::JustRelax.StokesArrays, args...) + return _compute_dt(S, args...) +end end diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index 2aaab4c3..e48b7ca3 100644 --- a/src/stokes/StressKernels.jl +++ b/src/stokes/StressKernels.jl @@ -311,7 +311,7 @@ function tensor_invariant!(A::JustRelax.SymmetricTensor) end function tensor_invariant!(::CPUBackendTrait, A::JustRelax.SymmetricTensor) - _tensor_invariant!(A) + return _tensor_invariant!(A) end function _tensor_invariant!(A::JustRelax.SymmetricTensor) From 0ada5ed0800dc108e5e9c0f0965f37191ab3f998 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 12:08:02 +0200 Subject: [PATCH 65/91] missing "_" --- test/test_VanKeken.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index 2a78e8ee..d8edc087 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -9,7 +9,6 @@ using JustRelax, JustRelax.JustRelax2D import JustRelax.@cell const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - 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, From 3fb6dceed1fcea681e5344c6abd8ba2ac3049ee2 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 12:15:49 +0200 Subject: [PATCH 66/91] typos --- ext/JustRelaxAMDGPUExt.jl | 4 ++-- src/Utils.jl | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/JustRelaxAMDGPUExt.jl b/ext/JustRelaxAMDGPUExt.jl index 517cfdef..2dac9360 100644 --- a/ext/JustRelaxAMDGPUExt.jl +++ b/ext/JustRelaxAMDGPUExt.jl @@ -6,8 +6,8 @@ import JustRelax: PTArray, backend, AMDGPUBackendTrait PTArray(::Type{AMDGPUBackend}) = RocArray -@inline backend(::CuArray) = AMDGPUBackendTrait() -@inline backend(::Type{<:CuArray}) = AMDGPUBackendTrait() +@inline backend(::RocArray) = AMDGPUBackendTrait() +@inline backend(::Type{<:RocArray}) = AMDGPUBackendTrait() include("../src/ext/AMDGPU/2D.jl") include("../src/ext/AMDGPU/3D.jl") diff --git a/src/Utils.jl b/src/Utils.jl index 6859d1de..2be50507 100644 --- a/src/Utils.jl +++ b/src/Utils.jl @@ -412,7 +412,7 @@ end Compute the time step `dt` for the velocity field `S.V` for a regular grid with grid spacing `di`. """ -@inline _compute_dt(S::JustRelax.StokesArrays, di) = compute_dt(@velocity(S), di, Inf) +@inline _compute_dt(S::JustRelax.StokesArrays, di) = _compute_dt(@velocity(S), di, Inf) """ compute_dt(S::JustRelax.StokesArrays, di, dt_diff) From e14bd11ad88d779a50db3a9929307752d6e93b1a Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 15:29:58 +0200 Subject: [PATCH 67/91] update miniapps --- .../stokes2D/VanKeken.jl/VanKeken.jl | 75 ++-- .../Elastic_BuildUp_phases_incompressible.jl | 58 ++- .../PlumeFreeSurface_2D.jl | 17 +- .../RayleighTaylor2D.jl | 66 ++-- .../stokes2D/shear_band/ShearBand2D.jl | 52 ++- .../stokes2D/shear_band/ShearBand2D_MPI.jl | 42 +-- .../shear_band/ShearBand2D_softening.jl | 33 +- .../stokes2D/shear_heating/Shearheating2D.jl | 103 +++--- .../stokes2D/sinking_block/SinkingBlock2D.jl | 51 ++- .../diffusion/diffusion2D.jl | 26 +- .../diffusion/diffusion2D_MPI.jl | 28 +- .../diffusion/diffusion2D_multiphase.jl | 41 +-- .../diffusion/diffusion3D.jl | 34 +- .../diffusion/diffusion3D_MPI.jl | 41 +-- .../diffusion/diffusion3D_multiphase.jl | 43 ++- .../convection/GlobalConvection2D_WENO5.jl | 119 +++--- .../GlobalConvection2D_WENO5_MPI.jl | 342 ------------------ .../Particles2D/Layered_convection2D.jl | 93 +++-- .../Particles2D/Layered_rheology.jl | 12 +- .../Layered_convection2D.jl | 73 ++-- .../Particles2D_nonDim/Layered_rheology.jl | 12 +- .../Particles3D/Layered_convection3D.jl | 83 ++--- .../Particles3D/Layered_rheology.jl | 12 +- 23 files changed, 505 insertions(+), 951 deletions(-) delete mode 100644 miniapps/convection/GlobalConvection2D_WENO5_MPI.jl diff --git a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl index c0300c06..628b7ad1 100644 --- a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl +++ b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl @@ -1,42 +1,27 @@ using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO#, CSV, DataFrames +using Printf, LinearAlgebra, GeoParams, CellArrays +using JustRelax, JustRelax.JustRelax2D import JustRelax.@cell -using JustPIC -using JustPIC._2D +const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # x-length of the domain const λ = 0.9142 # HELPER FUNCTIONS --------------------------------------------------------------- -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -# Initial pressure guess -@parallel function init_P!(P, ρg, z) - @all(P) = abs(@all(ρg) * @all_j(z)) * <(@all_j(z), 0.0) - return nothing -end - # Initialize phases on the particles function init_phases!(phases, particles) ni = size(phases) @parallel_indices (i, j) function init_phases!(phases, px, py, index) - @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 @@ -53,11 +38,10 @@ function init_phases!(phases, particles) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index) end # END OF HELPER FUNCTIONS -------------------------------------------------------- - # MAIN SCRIPT -------------------------------------------------------------------- function main2D(igg; ny=64, nx=64, figdir="model_figs") @@ -91,7 +75,7 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") ) # Initialize particles ------------------------------- - nxcell, max_p, min_p = 40, 60, 20 + nxcell, max_p, min_p = 40, 80, 20 particles = init_particles( backend, nxcell, max_p, min_p, xvi..., di..., nx, ny ) @@ -100,27 +84,22 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") # temperature pPhases, = init_cell_arrays(particles, Val(1)) particle_args = (pPhases, ) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, 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) + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; r=1e0, ϵ=1e-8, CFL = 1 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) - @parallel (JustRelax.@idx ni) JustRelax.compute_ρg!(ρg[2], phase_ratios.center, rheology, args) - @parallel init_P!(stokes.P, ρg[2], xci[2]) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -150,7 +129,7 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") while t < tmax # Update buoyancy - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -160,16 +139,16 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, - dt, ## dt needs to be a number! not 0 or Inf + dt, igg; - iterMax = 10e3, - nout = 50, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + iterMax = 10e3, + nout = 50, + viscosity_cutoff = (-Inf, Inf) + ) ) dt = compute_dt(stokes, di) / 10 # ------------------------------ @@ -185,15 +164,13 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") # ------------------------------ # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # # advect particles in memory move_particles!(particles, xvi, particle_args) - # check if we need to inject particles - inject = check_injection(particles) # inject && break - inject && inject_particles_phase!(particles, pPhases, (), (), xvi) + inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl index 6ad2ea3a..8c9df15f 100644 --- a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl +++ b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl @@ -1,28 +1,25 @@ using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -using GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = CPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using GeoParams, GLMakie, CellArrays # HELPER FUNCTIONS --------------------------------------------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) # Initialize phases on the particles -function init_phases!(phase_ratios, xci, radius) +function init_phases!(phase_ratios) ni = size(phase_ratios.center) - origin = 0.5, 0.5 - @parallel_indices (i, j) function init_phases!(phases, xc, yc, o_x, o_y) + @parallel_indices (i, j) function init_phases!(phases) JustRelax.@cell phases[1, i, j] = 1.0 return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin...) + @parallel (@idx ni) init_phases!(phase_ratios.center) end # MAIN SCRIPT -------------------------------------------------------------------- @@ -59,12 +56,12 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(ni, length(rheology)) - init_phases!(phase_ratios, xci, radius) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) + init_phases!(phase_ratios) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL=0.75 / √2.1) # Buoyancy forces @@ -73,29 +70,18 @@ function main(igg; nx=64, ny=64, figdir="model_figs") args = (; T=@zeros(ni...), P=stokes.P, dt=dt) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, - 1.0, - phase_ratios.center, - stokes.ε.xx, - stokes.ε.yy, - stokes.ε.xy, - args, - rheology, - (-Inf, Inf), - ) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; - free_slip=(left=true, right=true, top=true, bot=true), - no_slip=(left=false, right=false, top=false, bot=false), + 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]]) + stokes.V.Vx .= PTArray(backend_JR)([x * εbg for x in xvi[1], _ in 1:(ny + 2)]) + stokes.V.Vy .= PTArray(backend_JR)([-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) + # IO ------------------------------------------------ # if it does not exist, make folder where figures are stored !isdir(figdir) && mkpath(figdir) @@ -118,19 +104,19 @@ function main(igg; nx=64, ny=64, figdir="model_figs") di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose=false, - iterMax=500e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf), + kwargs = ( + verbose=false, + iterMax=500e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf), + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) if !isinf(dt) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 009c421c..793dd4f0 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -1,12 +1,14 @@ using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + using JustPIC, JustPIC._2D -const backend = CPUBackend +const backend = JustPIC.CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using LinearAlgebra, GeoParams, GLMakie # Velocity helper grids for the particle advection function copyinn_x!(A, B) @@ -117,18 +119,18 @@ function main(igg, nx, ny) # Elliptical temperature anomaly init_phases!(pPhases, particles) - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.95 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) # ---------------------------------------------------- # Buoyancy forces & rheology @@ -196,12 +198,11 @@ function main(igg, nx, ny) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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, (), (), xvi) + inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios phase_ratios_center(phase_ratios, particles, grid, pPhases) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 56ed5504..891f7350 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -1,19 +1,17 @@ # using CUDA -using JustRelax, JustRelax.DataIO - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend using JustPIC, JustPIC._2D -const backend = CPUBackend +const backend = JustPIC.CPUBackend -using ParallelStencil +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using LinearAlgebra, GeoParams, GLMakie +## START OF HELPER FUNCTION ---------------------------------------------------------- function copyinn_x!(A, B) @parallel function f_x(A, B) @all(A) = @inn_x(B) @@ -38,7 +36,7 @@ function init_phases!(phases, particles, A) @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) - f(x, A, λ) = A * sin(π*x/λ) + f(x, A, λ) = A * sin(π * x / λ) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape @@ -117,31 +115,27 @@ function RT_2D(igg, nx, ny) # Elliptical temperature anomaly A = 5e3 # Amplitude of the anomaly + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, A) - 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) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-8, CFL = 0.95 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) # ---------------------------------------------------- # Buoyancy forces & rheology ρg = @zeros(ni...), @zeros(ni...) - η = @ones(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) - @parallel (@idx ni) compute_viscosity!( - η, 0.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e19, 1e24) - ) - η_vep = copy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -176,13 +170,11 @@ function RT_2D(igg, nx, ny) dt_max = 50e3 * (3600 * 24 * 365.25) while it < 500 # run only for 5 Myrs + 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)) + # Stokes solver ---------------- - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) solve!( stokes, pt_stokes, @@ -196,12 +188,14 @@ function RT_2D(igg, nx, ny) args, dt, igg; - iterMax = 150e3, - iterMin = 5e3, - viscosity_relaxation = 1e-2, - nout = 5e3, - free_surface = true, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + iterMax = 150e3, + iterMin = 5e3, + viscosity_relaxation = 1e-2, + nout = 5e3, + free_surface = true, + viscosity_cutoff = (-Inf, Inf) + ) ) dt = if it ≤ 10 min(compute_dt(stokes, di), 1e3 * (3600 * 24 * 365.25)) @@ -216,21 +210,20 @@ function RT_2D(igg, nx, ny) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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, (), (), xvi) + inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios # @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) - + phase_ratios_center(phase_ratios, particles, grid, pPhases) + @show it += 1 t += dt if it == 1 || rem(it, 5) == 0 - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) nt = 2 p = particles.coords @@ -256,7 +249,6 @@ function RT_2D(igg, nx, ny) ) fig save(joinpath(figdir, "$(it).png"), fig) - end end diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 4b47d48c..12fb2839 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -1,11 +1,9 @@ -using GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO +using GeoParams +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +const backend = CPUBackend # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -28,7 +26,7 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) end # MAIN SCRIPT -------------------------------------------------------------------- @@ -82,7 +80,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") Elasticity = el_inc, ), ) - + # Initialize phase ratios ------------------------------- radius = 0.1 phase_ratios = PhaseRatio(ni, length(rheology)) @@ -90,38 +88,30 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - # Boundary conditions 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]]) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-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) - - # IO ------------------------------------------------ - # if it does not exist, make folder where figures are stored - take(figdir) - # ---------------------------------------------------- - + # Time loop t, it = 0.0, 0 - tmax = 5 + tmax = 3.5 τII = Float64[] sol = Float64[] ttot = Float64[] @@ -129,25 +119,25 @@ function main(igg; nx=64, ny=64, figdir="model_figs") while t < tmax # Stokes solver ---------------- - solve!( + iters = solve!( stokes, pt_stokes, di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose = false, - iterMax = 50e3, - nout = 1e2, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + verbose = false, + iterMax = 50e3, + nout = 1e2, + viscosity_cutoff = (-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl index 334398ad..67381d19 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl @@ -1,15 +1,9 @@ -using CUDA -using GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO +using GeoParams +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -dimension = 2 # 2 | 3 -device = :cpu # :cpu | :CUDA | :AMDGPU -precision = Float64 -model = PS_Setup(device, precision, dimension) -environment!(model) +const backend = CPUBackend # HELPER FUNCTIONS --------------------------------------------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -32,7 +26,7 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin..., radius) end # MAIN SCRIPT -------------------------------------------------------------------- @@ -89,23 +83,21 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) init_phases!(phase_ratios, xci, radius) - # STOKES --------------------------------------------- + # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions @@ -151,19 +143,19 @@ function main(igg; nx=64, ny=64, figdir="model_figs") di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose = false, - iterMax = 50e3, - nout = 1e3, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + verbose = false, + iterMax = 50e3, + nout = 1e2, + viscosity_cutoff = (-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) @parallel (@idx ni .+ 1) multi_copy!( diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl index e64e6273..43452a09 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl @@ -88,32 +88,29 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) init_phases!(phase_ratios, xci, radius) - # STOKES --------------------------------------------- + # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-6, CFL = 0.75 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @zeros(ni...), P = stokes.P, dt = dt) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) ) - # Boundary conditions 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]]) + stokes.V.Vx .= PTArray(backend)([ x*εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend)([-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) @@ -138,19 +135,19 @@ function main(igg; nx=64, ny=64, figdir="model_figs") di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - verbose = false, - iterMax = 50e3, - nout = 1e2, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + verbose = false, + iterMax = 50e3, + nout = 1e2, + viscosity_cutoff = (-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index f3e5c780..80bc5d27 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -1,28 +1,19 @@ -using CUDA # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax2D -# using JustRelax.DataIO -# import JustRelax.@cell +const backend_JR = CPUBackend + using ParallelStencil, ParallelStencil.FiniteDifferences2D -# @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -@init_parallel_stencil(CUDA, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +@init_parallel_stencil(Threads, Float64, 2) -using JustPIC -using JustPIC._2D +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 = JustPIC.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, 2) or (:AMDGPU, Float64, 2) -# environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays -using GLMakie +using GeoParams # Load file with all the rheology configurations include("Shearheating_rheology.jl") @@ -74,7 +65,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 40, 10 + nxcell, max_xcell, min_xcell = 20, 32, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -85,22 +76,22 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) particle_args = (pT, pPhases) # Elliptical temperature anomaly - xc_anomaly = lx/2 # origin of thermal anomaly - yc_anomaly = 40e3 # origin of thermal anomaly - r_anomaly = 3e3 # radius of perturbation - phase_ratios = PhaseRatio(backend, ni, length(rheology)) + xc_anomaly = lx / 2 # origin of thermal anomaly + yc_anomaly = 40e3 # origin of thermal anomaly + r_anomaly = 3e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, r_anomaly) phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(backend, ni) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(backend, ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -112,31 +103,38 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # ---------------------------------------------------- # Buoyancy forces - ρg = @zeros(ni...), @zeros(ni...) + ρg = @zeros(ni...), @zeros(ni...) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) # Rheology - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) - ) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + # PT coefficients for thermal diffusion - pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-1 / √2.1 + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1 ) # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), + flow_bcs = FlowBoundaryConditions(; + free_slip = (left = true, right=true, top=true, bot=true), ) ## Compression and not extension - fix this - εbg = 5e-14 - stokes.V.Vx .= PTArray(backend )([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) - stokes.V.Vy .= PTArray(backend)([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]]) + εbg = 5e-14 + stokes.V.Vx .= PTArray(backend_JR)([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2]) + stokes.V.Vy .= PTArray(backend_JR)([ (ly - abs(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) + T_buffer = @zeros(ni.+1) + Told_buffer = similar(T_buffer) + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle!(pT, xvi, T_buffer, particles) + + # IO ----- ------------------------------------------- # if it does not exist, make folder where figures are stored if do_vtk @@ -162,18 +160,12 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) fig end - T_buffer = @zeros(ni.+1) - Told_buffer = similar(T_buffer) - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - grid2particle!(pT, xvi, T_buffer, particles) - local Vx_v, Vy_v if do_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) end + # Time loop t, it = 0.0, 0 while it < 1 @@ -240,21 +232,20 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) ) # ------------------------------ - # Advection -------------------- - # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) - # advect particles in memory - move_particles!(particles, xvi, particle_args) - # interpolate fields from grid vertices to particles - for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) - copyinn_x!(dst, src) - end - grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) - # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) - # update phase ratios - phase_ratios_center(phase_ratios, particles, grid, pPhases) + # Advection -------------------- + # advect particles in space + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) + # advect particles in memory + move_particles!(particles, xvi, particle_args) + # interpolate fields from grid vertices to particles + for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) + copyinn_x!(dst, src) + end + grid2particle_flip!(pT, xvi, T_buffer, Told_buffer, particles) + # check if we need to inject particles + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) + # update phase ratios + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl index e855a003..2c04b64c 100644 --- a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl +++ b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl @@ -1,19 +1,13 @@ -using JustRelax -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend -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 +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using JustPIC, JustPIC._2D +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using GeoParams ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- @@ -55,7 +49,7 @@ function init_phases!(phases, particles, xc, yc, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, xc, yc, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, xc, yc, r) end import ParallelStencil.INDICES @@ -108,7 +102,7 @@ function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_per # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 20, 20, 1 + nxcell, max_xcell, min_xcell = 20, 40, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -119,28 +113,24 @@ function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_per xc_anomaly = 250e3 # origin of thermal anomaly yc_anomaly = -(ly-400e3) # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, xc_anomaly, abs(yc_anomaly), r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-5, CFL = 0.95 / √2.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=@ones(ni...), P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=@ones(ni...), P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) # ---------------------------------------------------- # Viscosity - η = @ones(ni...) args = (; dt = dt, ΔTc = @zeros(ni...)) η_cutoff = -Inf, Inf - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, η_cutoff - ) - η_vep = deepcopy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions @@ -152,22 +142,23 @@ function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_per # Stokes solver ---------------- args = (; T = @ones(ni...), P = stokes.P, dt=dτ, ΔTc = @zeros(ni...)) - solve!( + olve!( stokes, pt_stokes, di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax=150e3, - nout=1e3, - viscosity_cutoff = η_cutoff + kwargs = ( + iterMax=150e3, + nout=1e3, + viscosity_cutoff = η_cutoff, + verbose = false, + ) ); dt = compute_dt(stokes, di, igg) # ------------------------------ diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl index d782ce71..9fd1a11a 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl @@ -1,3 +1,17 @@ +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) + +using Printf, LinearAlgebra, GeoParams, CellArrays +using JustRelax, JustRelax.JustRelax2D +import JustRelax.@cell +const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + +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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + @parallel_indices (i, j) function init_T!(T, z) if z[j] == maximum(z) T[i, j] = 300.0 @@ -26,9 +40,8 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, Myr = 1e3 * kyr ttot = 1 * Myr # total simulation time dt = 50 * kyr # physical time step - init_mpi = JustRelax.MPI.Initialized() ? false : true - igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = init_mpi)...) # Physical domain ni = (nx, ny) @@ -50,7 +63,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, args = (; P=P) ## Allocate arrays needed for every Thermal Diffusion - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal.H .= 1e-6 # radiogenic heat production # physical parameters ρ = @fill(ρ0, ni...) @@ -58,7 +71,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, K = @fill(K0, ni...) ρCp = @. Cp * ρ - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li) + pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -69,11 +82,13 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, r = 10e3 # thermal perturbation radius center_perturbation = lx/2, -ly/2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) + temperature2center!(thermal) # Time loop t = 0.0 it = 0 nt = Int(ceil(ttot / dt)) + while it < nt heatdiffusion_PT!( thermal, @@ -82,7 +97,8 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, rheology, args, dt, - di, + di; + kwargs = (; verbose=false), ) t += dt diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl index add48909..85f244fd 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_MPI.jl @@ -1,15 +1,14 @@ # using CairoMakie -using JustRelax, GeoParams -using GLMakie +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + using ParallelStencil -@init_parallel_stencil(Threads, Float64, 3) +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -# setup ParallelStencil.jl environment -dimension = 2 # 2 | 3 -device = :cpu # :cpu | :CUDA | :AMDGPU -precision = Float64 -model = PS_Setup(device, precision, dimension) -environment!(model) +import JustRelax.@cell +using GeoParams + +using GLMakie @parallel_indices (i, j) function init_T!(T, z, lz) if z[j] ≥ 0.0 @@ -70,7 +69,7 @@ function diffusion_2D(; args = (; P=P) ## Allocate arrays needed for every Thermal Diffusion - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal.H .= 1e-6 # radiogenic heat production # physical parameters ρ = @fill(ρ0, ni...) @@ -78,7 +77,7 @@ function diffusion_2D(; K = @fill(K0, ni...) ρCp = @. Cp * ρ - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li) + pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -89,6 +88,7 @@ function diffusion_2D(; r = 10e3 # thermal perturbation radius center_perturbation = lx/2, -ly/2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) + temperature2center!(thermal) # global array nx_v = ((nx + 2) - 2) * igg.dims[1] @@ -108,8 +108,10 @@ function diffusion_2D(; args, dt, di; - igg=igg, - b_width=(4, 4, 1), + kwargs = (; + igg = igg, + b_width = (4, 4, 1), + ) ) @views T_nohalo .= Array(thermal.T[2:end-2, 2:end-1]) # Copy data to CPU removing the halo diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl index 53fbfdb7..3d12b48d 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D_multiphase.jl @@ -1,21 +1,11 @@ -using Printf, LinearAlgebra, GeoParams, SpecialFunctions, CellArrays, StaticArrays -using JustRelax +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using JustPIC -using JustPIC._2D -# Threads is the default backend, -# to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script, -# and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - import JustRelax.@cell - +using GeoParams distance(p1, p2) = mapreduce(x->(x[1]-x[2])^2, +, zip(p1, p2)) |> sqrt @@ -65,7 +55,7 @@ function init_phases!(phases, particles, xc, yc, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) end @parallel_indices (I...) function compute_temperature_source_terms!(H, rheology, phase_ratios, args) @@ -115,7 +105,7 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) args = (; P=P) ## Allocate arrays needed for every Thermal Diffusion - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -126,26 +116,25 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) r = 10e3 # thermal perturbation radius center_perturbation = lx/2, -ly/2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # Initialize particles ------------------------------- nxcell, max_xcell, min_xcell = 40, 40, 1 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) - # temperature pPhases, = init_cell_arrays(particles, Val(1)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, center_perturbation..., r) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) JustRelax.phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- @parallel (@idx ni) compute_temperature_source_terms!(thermal.H, rheology, phase_ratios.center, args) # PT coefficients for thermal diffusion args = (; P=P, T=thermal.Tc) - pt_thermal = JustRelax.PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.65 / √2 + pt_thermal = PTThermalCoeffs( + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=0.65 / √2 ) # Time loop @@ -161,9 +150,11 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, Cp0=1.2e3, K0=3.0) args, dt, di; - phase = phase_ratios, - iterMax = 1e3, - nout = 10, + kwargs = (; + phase = phase_ratios, + iterMax = 1e3, + nout = 10, + ) ) it += 1 diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D.jl index 4d33b3af..b04647b2 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D.jl @@ -1,14 +1,9 @@ -using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using GeoParams +using JustRelax, JustRelax.JustRelax3D using ParallelStencil -@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) +@init_parallel_stencil(Threads, Float64, 3) -# setup ParallelStencil.jl environment -dimension = 3 # 2 | 3 -device = :cpu # :cpu | :CUDA | :AMDGPU -precision = Float64 -model = PS_Setup(device, precision, dimension) -environment!(model) +const backend = CPUBackend @parallel_indices (i, j, k) function init_T!(T, z) if z[k] == maximum(z) @@ -47,10 +42,10 @@ function diffusion_3D(; finalize_MPI = false, ) - kyr = 1e3 * 3600 * 24 * 365.25 - Myr = 1e6 * 3600 * 24 * 365.25 - ttot = 1 * Myr # total simulation time - dt = 50 * kyr # physical time step + kyr = 1e3 * 3600 * 24 * 365.25 + Myr = 1e6 * 3600 * 24 * 365.25 + ttot = 1 * Myr # total simulation time + dt = 50 * kyr # physical time step # Physical domain ni = (nx, ny, nz) @@ -75,7 +70,7 @@ function diffusion_3D(; ## Allocate arrays needed for every Thermal Diffusion # general thermal arrays - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal.H .= 1e-6 # physical parameters ρ = @fill(ρ0, ni...) @@ -84,7 +79,7 @@ function diffusion_3D(; ρCp = @. Cp * ρ # Boundary conditions - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li; CFL = 0.75 / √3.1) + pt_thermal = PTThermalCoeffs(backend, K, ρCp, dt, di, li; CFL = 0.75 / √3.1) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) @@ -94,7 +89,7 @@ function diffusion_3D(; # Add thermal perturbation δT = 100e0 # thermal perturbation r = 10e3 # thermal perturbation radius - center_perturbation = lx/2, ly/2, -lz/2 + center_perturbation = lx / 2, ly / 2, -lz / 2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) t = 0.0 @@ -110,8 +105,11 @@ function diffusion_3D(; rheology, args, dt, - di,; - igg + di; + kwargs = (; + igg, + verbose=false + ), ) t += dt diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_MPI.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_MPI.jl index 4d151cf8..839507c5 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_MPI.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_MPI.jl @@ -1,14 +1,10 @@ -using JustRelax, GeoParams -using GLMakie +using JustRelax, JustRelax.JustRelax3D using ParallelStencil -@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) +@init_parallel_stencil(Threads, Float64, 3) -# setup ParallelStencil.jl environment -dimension = 3 # 2 | 3 -device = :cpu # :cpu | :CUDA | :AMDGPU -precision = Float64 -model = PS_Setup(device, precision, dimension) -environment!(model) +const backend = CPUBackend + +using GLMakie, GeoParams @parallel_indices (i, j, k) function init_T!(T, z, lz) if z[k] ≥ 0.0 @@ -62,13 +58,12 @@ function diffusion_3D(; grid = Geometry(ni, li; origin = origin) (; xci, xvi) = grid # nodes at the center and vertices of the cells - # Define the thermal parameters with GeoParams rheology = SetMaterialParams(; - Phase = 1, - Density = PT_Density(; ρ0=3.1e3, β=0.0, T0=0.0, α = 1.5e-5), - HeatCapacity = ConstantHeatCapacity(; Cp=Cp0), - Conductivity = ConstantConductivity(; k=K0), + Phase = 1, + Density = PT_Density(; ρ0=3.1e3, β=0.0, T0=0.0, α = 1.5e-5), + HeatCapacity = ConstantHeatCapacity(; Cp=Cp0), + Conductivity = ConstantConductivity(; k=K0), ) # fields needed to compute density on the fly @@ -77,7 +72,7 @@ function diffusion_3D(; ## Allocate arrays needed for every Thermal Diffusion # general thermal arrays - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal.H .= 1e-6 # physical parameters ρ = @fill(ρ0, ni...) @@ -86,19 +81,17 @@ function diffusion_3D(; ρCp = @. Cp * ρ # Boundary conditions - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li; CFL = 0.75 / √3.1) + pt_thermal = PTThermalCoeffs(backend, K, ρCp, dt, di, li; CFL = 0.75 / √3.1) thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), + no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) - @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[3], lz) - - # return nothing + @parallel (@idx size(thermal.T)) init_T!(thermal.T, xvi[3]) # Add thermal perturbation δT = 100e0 # thermal perturbation r = 10e3 # thermal perturbation radius - center_perturbation = lx/2, ly/2, -lz/2 + center_perturbation = lx / 2, ly / 2, -lz / 2 elliptical_perturbation!(thermal.T, δT, center_perturbation..., r, xvi) # Visualization global arrays @@ -120,8 +113,10 @@ function diffusion_3D(; rheology, args, dt, - di,; - igg + di; + kwargs = (; + igg, + ) ) @views T_nohalo .= Array(thermal.T[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl index 011de511..2c4ba5b6 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl @@ -1,20 +1,18 @@ -using Printf, LinearAlgebra, GeoParams, CellArrays, StaticArrays -using JustRelax +using JustRelax, JustRelax.JustRelax3D +import JustRelax.@cell + +const backend_JR = CPUBackend + using ParallelStencil -@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +@init_parallel_stencil(Threads, Float64, 3) -using JustPIC -using JustPIC._3D +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - -import JustRelax.@cell +using GeoParams @parallel_indices (i, j, k) function init_T!(T, z) if z[k] == maximum(z) @@ -63,7 +61,7 @@ function init_phases!(phases, particles, xc, yc, zc, r) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, center, r) end function diffusion_3D(; @@ -118,7 +116,7 @@ function diffusion_3D(; ## Allocate arrays needed for every Thermal Diffusion # general thermal arrays - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal.H .= 1e-6 # physical parameters ρ = @fill(ρ0, ni...) @@ -144,16 +142,15 @@ function diffusion_3D(; particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni ) - # temperature pPhases, = init_cell_arrays(particles, Val(1)) + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, center_perturbation..., r) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) JustRelax.phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # PT coefficients for thermal diffusion args = (; P=P, T=thermal.Tc) - pt_thermal = PTThermalCoeffs(K, ρCp, dt, di, li; CFL = 0.75 / √3.1) + pt_thermal = PTThermalCoeffs(backend_JR, K, ρCp, dt, di, li; CFL = 0.75 / √3.1) t = 0.0 it = 0 @@ -169,11 +166,13 @@ function diffusion_3D(; args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs =(; + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true, + ) ) t += dt diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index 5827a849..c9ef42e7 100644 --- a/miniapps/convection/GlobalConvection2D_WENO5.jl +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -1,12 +1,12 @@ -using JustRelax -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +import JustRelax.@cell + +const backend = JustRelax.CPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions +using GeoParams, GLMakie, SpecialFunctions # function to compute strain rate (compulsory) @inline function custom_εII(a::CustomRheology, TauII; args...) @@ -29,7 +29,6 @@ end end # HELPER FUNCTIONS --------------------------------------------------------------- - import ParallelStencil.INDICES const idx_j = INDICES[2] macro all_j(A) @@ -37,7 +36,7 @@ macro all_j(A) end @parallel function init_P!(P, ρg, z) - @all(P) = @all(ρg)*abs(@all_j(z)) + @all(P) = @all(ρg) * abs(@all_j(z)) return nothing end @@ -95,7 +94,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # ---------------------------------------------------- # Weno model ----------------------------------------- - weno = WENO5(ni= ni .+ 1, method=Val{2}()) # ni.+1 for Temp + weno = WENO5(ni = ni .+ 1, method = Val(2)) # ni.+1 for Temp # ---------------------------------------------------- # create rheology struct @@ -137,7 +136,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -161,64 +160,45 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma end @views thermal.T[:, 1] .= Tmax @views thermal.T[:, end] .= Tmin + update_halo!(thermal.T) @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni, ViscoElastic) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) + # Buoyancy forces - ρg = @zeros(ni...), @zeros(ni...) - for _ in 1:2 - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + ρg = @zeros(ni...), @zeros(ni...) + for _ in 1:1 + compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end - + # Rheology - η = @ones(ni...) - depth = PTArray([y for x in xci[1], y in xci[2]]) - args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = dt, ΔTc = thermal.ΔTc) - η_cutoff = 1e18, 1e23 - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, η_cutoff - ) - η_vep = deepcopy(η) + viscosity_cutoff = (1e16, 1e24) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 + backend, rheology, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 ) # Boundary conditions flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), + free_slip = (left = true , right = true , top = true , bot = true), ) flow_bcs!(stokes, flow_bcs) # apply boundary conditions - update_halo!(stokes.V.Vx, stokes.V.Vy) + update_halo!(@velocity(stokes)...) # ---------------------------------------------------- # IO ------------------------------------------------- # if it does not exist, make folder where figures are stored - !isdir(figdir) && mkpath(figdir) + take(figdir) # ---------------------------------------------------- - # Plot initial T and η profiles - fig0 = let - Yv = [y for x in xvi[1], y in xvi[2]][:] - Y = [y for x in xci[1], y in xci[2]][:] - fig = Figure(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[2:end-1,:][:]), Yv./1e3) - lines!(ax2, Array(log10.(η[:])), Y./1e3) - ylims!(ax1, -2890, 0) - ylims!(ax2, -2890, 0) - hideydecorations!(ax2) - save(joinpath(figdir, "initial_profile.png"), fig) - fig - end - # WENO arrays T_WENO = @zeros(ni.+1) Vx_v = @zeros(ni.+1...) @@ -227,34 +207,40 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # Time loop t, it = 0.0, 0 local iters + igg.me == 0 && println("Starting model") while (t / (1e6 * 3600 * 24 * 365.25)) < 4.5e3 - # Stokes solver ---------------- - args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt=dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, η_cutoff + + # Update buoyancy and viscosity - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff ) + # ------------------------------ + igg.me == 0 && println("Starting stokes solver...") iters = solve!( stokes, pt_stokes, di, flow_bcs, ρg, - η, - η_vep, rheology, args, dt, igg; - iterMax=50e3, - nout=1e3, - viscosity_cutoff = η_cutoff + kwargs = (; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) ); dt = compute_dt(stokes, di, dt_diff, igg) + println("Rank $(igg.me) ...stokes solver finished") # ------------------------------ # Thermal solver --------------- + igg.me == 0 && println("Starting thermal solver...") heatdiffusion_PT!( thermal, pt_thermal, @@ -263,14 +249,17 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma args, dt, di; - igg = igg, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs = ( + igg = igg, + iterMax = 10e3, + nout = 1e2, + verbose = true + ), ) + igg.me == 0 && println("...thermal solver finished") # Weno advection T_WENO .= @views thermal.T[2:end-1, :] - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) @views thermal.T[2:end-1, :] .= T_WENO # ------------------------------ @@ -308,18 +297,20 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma end - return (ni=ni, xci=xci, li=li, di=di), thermal + finalize_global_grid() + + return nothing end figdir = "figs2D_test_weno" ar = 8 # aspect ratio -n = 128 +n = 32 nx = n*ar - 2 ny = n - 2 +thermal_perturbation = :circular igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) + IGG(init_global_grid(nx, ny, 1; init_MPI= true, select_device=false)...) else igg end - -thermal_convection2D(igg; figdir=figdir, ar=ar,nx=nx, ny=ny); +thermal_convection2D(igg; figdir=figdir, ar=ar,nx=nx, ny=ny, thermal_perturbation=thermal_perturbation); diff --git a/miniapps/convection/GlobalConvection2D_WENO5_MPI.jl b/miniapps/convection/GlobalConvection2D_WENO5_MPI.jl deleted file mode 100644 index d1976d24..00000000 --- a/miniapps/convection/GlobalConvection2D_WENO5_MPI.jl +++ /dev/null @@ -1,342 +0,0 @@ -using JustRelax -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - -using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions - -# function to compute strain rate (compulsory) -@inline function custom_εII(a::CustomRheology, TauII; args...) - η = custom_viscosity(a; args...) - return TauII / η * 0.5 -end - -# function to compute deviatoric stress (compulsory) -@inline function custom_τII(a::CustomRheology, EpsII; args...) - η = custom_viscosity(a; args...) - return 2.0 * η * EpsII -end - -# helper function (optional) -@inline function custom_viscosity(a::CustomRheology; P=0.0, T=273.0, depth=0.0, kwargs...) - (; η0, Ea, Va, T0, R, cutoff) = a.args - η = η0 * exp((Ea + P * Va) / (R * T) - Ea / (R * T0)) - correction = (depth ≤ 660e3) + (2740e3 ≥ depth > 660e3) * 1e1 + (depth > 2700e3) * 1e-1 - η = clamp(η * correction, cutoff...) -end - -# HELPER FUNCTIONS --------------------------------------------------------------- - -import ParallelStencil.INDICES -const idx_j = INDICES[2] -macro all_j(A) - esc(:($A[$idx_j])) -end - -@parallel function init_P!(P, ρg, z) - @all(P) = @all(ρg) * abs(@all_j(z)) - return nothing -end - -# Half-space-cooling model -@parallel_indices (i, j) function init_T!(T, z, κ, Tm, Tp, Tmin, Tmax) - yr = 3600*24*365.25 - dTdz = (Tm-Tp)/2890e3 - zᵢ = abs(z[j]) - Tᵢ = Tp + dTdz*(zᵢ) - time = 100e6 * yr - Ths = Tmin + (Tm -Tmin) * erf((zᵢ)*0.5/(κ*time)^0.5) - T[i, j] = min(Tᵢ, Ths) - return -end - -function circular_perturbation!(T, δT, xc, yc, r, xvi) - - @parallel_indices (i, j) function _circular_perturbation!(T, δT, xc, yc, r, x, y) - @inbounds if (((x[i] - xc))^2 + ((y[j] - yc))^2) ≤ r^2 - T[i, j] *= δT / 100 + 1 - end - return nothing - end - - @parallel _circular_perturbation!(T, δT, xc, yc, r, xvi...) -end - -function random_perturbation!(T, δT, xbox, ybox, xvi) - - @parallel_indices (i, j) function _random_perturbation!(T, δT, xbox, ybox, x, y) - @inbounds if (xbox[1] ≤ x[i] ≤ xbox[2]) && (abs(ybox[1]) ≤ abs(y[j]) ≤ abs(ybox[2])) - δTi = δT * (rand() - 0.5) # random perturbation within ±δT [%] - T[i, j] *= δTi / 100 + 1 - end - return nothing - end - - @parallel (@idx size(T)) _random_perturbation!(T, δT, xbox, ybox, xvi...) -end - -# -------------------------------------------------------------------------------- -# BEGIN MAIN SCRIPT -# -------------------------------------------------------------------------------- -function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) - - # Physical domain ------------------------------------ - ly = 2890e3 - lx = ly * ar - origin = 0.0, -ly # origin coordinates - ni = nx, ny # number of cells - li = lx, ly # domain length in x- and y- - di = @. li / (nx_g(), ny_g()) # grid step in x- and -y - grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells - # ---------------------------------------------------- - - # Weno model ----------------------------------------- - weno = WENO5(ni = ni .+ 1, method = Val(2)) # ni.+1 for Temp - # ---------------------------------------------------- - - # create rheology struct - v_args = (; η0=5e20, Ea=200e3, Va=2.6e-6, T0=1.6e3, R=8.3145, cutoff=(1e16, 1e25)) - creep = CustomRheology(custom_εII, custom_τII, v_args) - - # Physical properties using GeoParams ---------------- - η_reg = 1e16 - G0 = 70e9 # shear modulus - cohesion = 30e6 - friction = asind(0.01) - pl = DruckerPrager_regularised(; C = cohesion, ϕ=friction, η_vp=η_reg, Ψ=0.0) # non-regularized plasticity - el = SetConstantElasticity(; G=G0, ν=0.5) # elastic spring - β = inv(get_Kb(el)) - - rheology = SetMaterialParams(; - Name = "Mantle", - Phase = 1, - Density = PT_Density(; ρ0=3.1e3, β=β, T0=0.0, α = 1.5e-5), - HeatCapacity = ConstantHeatCapacity(; Cp=1.2e3), - Conductivity = ConstantConductivity(; k=3.0), - CompositeRheology = CompositeRheology((creep, el, )), - Elasticity = el, - Gravity = ConstantGravity(; g=9.81), - ) - rheology_plastic = SetMaterialParams(; - Name = "Mantle", - Phase = 1, - Density = PT_Density(; ρ0=3.5e3, β=β, T0=0.0, α = 1.5e-5), - HeatCapacity = ConstantHeatCapacity(; Cp=1.2e3), - Conductivity = ConstantConductivity(; k=3.0), - CompositeRheology = CompositeRheology((creep, el, pl)), - Elasticity = el, - Gravity = ConstantGravity(; g=9.81), - ) - # heat diffusivity - κ = (rheology.Conductivity[1].k / (rheology.HeatCapacity[1].Cp * rheology.Density[1].ρ0)).val - dt = dt_diff = 0.5 * min(di...)^2 / κ / 2.01 # diffusive CFL timestep limiter - # ---------------------------------------------------- - - # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) - thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true, right = true, top = false, bot = false), - ) - # initialize thermal profile - Half space cooling - adiabat = 0.3 # adiabatic gradient - Tp = 1900 - Tm = Tp + adiabat * 2890 - Tmin, Tmax = 300.0, 3.5e3 - @parallel init_T!(thermal.T, xvi[2], κ, Tm, Tp, Tmin, Tmax) - thermal_bcs!(thermal.T, thermal_bc) - # Temperature anomaly - if thermal_perturbation == :random - δT = 5.0 # thermal perturbation (in %) - random_perturbation!(thermal.T, δT, (lx*1/8, lx*7/8), (-2000e3, -2600e3), xvi) - - elseif thermal_perturbation == :circular - δT = 10.0 # thermal perturbation (in %) - xc, yc = 0.5*lx, -0.75*ly # center of the thermal anomaly - r = 150e3 # radius of perturbation - circular_perturbation!(thermal.T, δT, xc, yc, r, xvi) - end - @views thermal.T[:, 1] .= Tmax - @views thermal.T[:, end] .= Tmin - update_halo!(thermal.T) - @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) - # ---------------------------------------------------- - - # STOKES --------------------------------------------- - # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) - # Buoyancy forces - ρg = @zeros(ni...), @zeros(ni...) - for _ in 1:2 - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) - @parallel init_P!(stokes.P, ρg[2], xci[2]) - end - - # Rheology - η = @ones(ni...) - depth = PTArray([y for x in xci[1], y in xci[2]]) - args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = dt, ΔTc = thermal.ΔTc) - η_cutoff = 1e18, 1e23 - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, η_cutoff - ) - η_vep = deepcopy(η) - - # PT coefficients for thermal diffusion - pt_thermal = PTThermalCoeffs( - rheology, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1 - ) - - # Boundary conditions - flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true , right = true , top = true , bot = true), - ) - flow_bcs!(stokes, flow_bcs) # apply boundary conditions - update_halo!(@velocity(stokes)...) - # ---------------------------------------------------- - - # IO ------------------------------------------------- - # if it does not exist, make folder where figures are stored - take(figdir) - # ---------------------------------------------------- - - # # Plot initial T and η profiles - # fig0 = let - # Yv = [y for x in xvi[1], y in xvi[2]][:] - # Y = [y for x in xci[1], y in xci[2]][:] - # fig = Figure(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[2:end-1,:][:]), Yv./1e3) - # lines!(ax2, Array(log10.(η[:])), Y./1e3) - # ylims!(ax1, -2890, 0) - # ylims!(ax2, -2890, 0) - # hideydecorations!(ax2) - # save(joinpath(figdir, "initial_profile.png"), fig) - # fig - # end - - # WENO arrays - T_WENO = @zeros(ni.+1) - Vx_v = @zeros(ni.+1...) - Vy_v = @zeros(ni.+1...) - - # global array - nx_v = ((nx + 2) - 2) * igg.dims[1] - ny_v = ((ny + 1) - 2) * igg.dims[2] - T_v = zeros(nx_v, ny_v) - T_nohalo = zeros((nx + 2)-2, (ny + 1)-2) - - # Time loop - t, it = 0.0, 0 - local iters - igg.me == 0 && println("Starting model") - while (t / (1e6 * 3600 * 24 * 365.25)) < 4.5e3 - # Stokes solver ---------------- - args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt=dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, args) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, η_cutoff - ) - - igg.me == 0 && println("Starting stokes solver...") - iters = solve!( - stokes, - pt_stokes, - di, - flow_bcs, - ρg, - η, - η_vep, - rheology, - args, - dt, - igg; - iterMax=50e3, - nout=1e3, - viscosity_cutoff = η_cutoff - ); - dt = compute_dt(stokes, di, dt_diff, igg) - println("Rank $(igg.me) ...stokes solver finished") - # ------------------------------ - - # Thermal solver --------------- - igg.me == 0 && println("Starting thermal solver...") - heatdiffusion_PT!( - thermal, - pt_thermal, - thermal_bc, - rheology, - args, - dt, - di; - igg = igg, - iterMax = 10e3, - nout = 1e2, - verbose = true, - ) - igg.me == 0 && println("...thermal solver finished") - # Weno advection - T_WENO .= @views thermal.T[2:end-1, :] - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) - update_halo!(Vx_v, Vy_v) - WENO_advection!(T_WENO, (Vx_v, Vy_v), weno, di, dt) - update_halo!(T_WENO) - @views thermal.T[2:end-1, :] .= T_WENO - update_halo!(thermal.T) - # ------------------------------ - - return - it += 1 - t += dt - - if igg.me == 0 - println("\n") - println("Time step number $it") - println(" time = $(t/(1e6 * 3600 * 24 *365.25)) Myrs, dt = $(dt/(1e6 * 3600 * 24 *365.25)) Myrs") - println("\n") - end - @views T_nohalo .= Array(thermal.T[2:end-2, 2:end-1]) # Copy data to CPU removing the halo - gather!(T_nohalo, T_v) - - if igg.me == 0 - if it == 1 || rem(it, 10) == 0 - println("Saving figure...") - xv_global = LinRange(0, lx / 1e3, size(T_v, 1)) - yv_global = LinRange(-ly / 1e3, 0, size(T_v, 2)) - fig = Figure(size = (1000, 1000), title = "t = $t") - ax = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") - h1 = heatmap!(ax, xv_global,yv_global,T_v, colormap=:batlow) - Colorbar(fig[1,2], h1, height=100) - fig, = heatmap(T_v, colorrange=(1500,2000)) - save( joinpath(figdir, "$(it).png"), fig) - println("...saving figure") - end - end - - end - - finalize_global_grid() - - return nothing -end - -figdir = "figs2D_test_weno_mpi" -ar = 8 # aspect ratio -n = 32 -nx = n*ar - 2 -ny = n - 2 -thermal_perturbation = :circular -igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 1; init_MPI= true, select_device=false)...) -else - igg -end -# thermal_convection2D(igg; figdir=figdir, ar=ar,nx=nx, ny=ny, thermal_perturbation=thermal_perturbation); - - -# mpiexec.exe -np 4 julia --startup-file=no --project=. .\miniapps\convection\GlobalConvection2D_Upwind_MPI.jl diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index 167855da..342f1939 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -1,26 +1,23 @@ -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -using ParallelStencil + +const backend_JR = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using JustPIC -using JustPIC._2D +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using GeoParams, GLMakie # Load file with all the rheology configurations include("Layered_rheology.jl") - ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- function copyinn_x!(A, B) @@ -114,7 +111,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # ---------------------------------------------------- # Initialize particles ------------------------------- - nxcell, max_xcell, min_xcell = 25, 30, 8 + nxcell, max_xcell, min_xcell = 25, 30, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -129,19 +126,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx/2 # origin of thermal anomaly yc_anomaly = -610e3 # origin of thermal anomaly r_anomaly = 25e3 # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly, thick_air) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.75 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -150,35 +147,34 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) thermal_bcs!(thermal.T, thermal_bc) Tbot = thermal.T[1, 1] rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi, thick_air) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) ρg = @zeros(ni...), @zeros(ni...) for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end + # Rheology - η = @ones(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e16, 1e24) - ) - η_vep = copy(η) + viscosity_cutoff = (1e16, 1e24) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) + (; η, η_vep) = stokes.viscosity # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-2 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-2 / √2.1 ) # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right=true, top=true, bot=true), - periodicity = (left = false, right = false, top = false, bot = false), ) flow_bcs!(stokes, flow_bcs) # apply boundary conditions - update_halo!(stokes.V.Vx, stokes.V.Vy) + update_halo!(@velocity(stokes)...) + # IO ----- ------------------------------------------- # if it does not exist, make folder where figures are stored if do_vtk @@ -217,6 +213,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) Vx_v = @zeros(ni.+1...) Vy_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 @@ -229,11 +226,11 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) temperature2center!(thermal) # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt=dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -243,18 +240,18 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 150e3, - nout=1e3, - viscosity_cutoff=(1e18, 1e24) + kwargs = (; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -267,11 +264,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true + ), ) for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) @@ -287,21 +286,20 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt # Data I/O and plotting --------------------- if it == 1 || rem(it, 1) == 0 - checkpointing(figdir, stokes, thermal.T, η, t) + # checkpointing(figdir, stokes, thermal.T, η, t) if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) @@ -385,5 +383,4 @@ else end # run main script - main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, do_vtk = do_vtk); diff --git a/miniapps/convection/Particles2D/Layered_rheology.jl b/miniapps/convection/Particles2D/Layered_rheology.jl index 3c3f1912..502dc3cb 100644 --- a/miniapps/convection/Particles2D/Layered_rheology.jl +++ b/miniapps/convection/Particles2D/Layered_rheology.jl @@ -121,19 +121,19 @@ function init_phases!(phases, particles, Lx, d, r, thick_air) x = JustRelax.@cell px[ip, i, j] depth = -(JustRelax.@cell py[ip, i, j]) - thick_air if 0e0 ≤ depth ≤ 21e3 - @cell phases[ip, i, j] = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 elseif 35e3 ≥ depth > 21e3 - @cell phases[ip, i, j] = 2.0 + JustRelax.@cell phases[ip, i, j] = 2.0 elseif 90e3 ≥ depth > 35e3 - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 elseif depth > 90e3 - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 elseif depth < 0e0 - @cell phases[ip, i, j] = 5.0 + JustRelax.@cell phases[ip, i, j] = 5.0 end @@ -145,5 +145,5 @@ function init_phases!(phases, particles, Lx, d, r, thick_air) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx) end diff --git a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl index 61a0afdb..200192c1 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl @@ -1,21 +1,19 @@ -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell + +const backend_JR = JustRelax.CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) -using JustPIC -using JustPIC._2D +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using GeoParams, GLMakie # Load file with all the rheology configurations include("Layered_rheology.jl") @@ -135,19 +133,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) xc_anomaly = lx/2 # origin of thermal anomaly yc_anomaly = nondimensionalize(-610km, CharDim) # origin of thermal anomaly r_anomaly = nondimensionalize(25km, CharDim) # radius of perturbation + phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, lx, yc_anomaly, r_anomaly, thick_air, CharDim) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-7, CFL = 0.9 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -157,27 +155,23 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) Tbot = thermal.T[1, 1] Ttop = thermal.T[1, end] rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi, thick_air, CharDim) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end + # Rheology - η = @ones(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) viscosity_cutoff = nondimensionalize((1e16Pa*s, 1e24Pa*s), CharDim) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, viscosity_cutoff - ) - η_vep = copy(η) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-6, CFL= 1e-3 / √2.1 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-6, CFL= 1e-3 / √2.1 ) # Boundary conditions @@ -239,10 +233,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # 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, viscosity_cutoff + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -252,18 +246,18 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 150e3, - nout = 1e3, - viscosity_cutoff = viscosity_cutoff + kwargs = (; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -276,11 +270,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 50e3, - nout = 5e2, - verbose = true, + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true + ), ) for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) @@ -296,14 +292,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt diff --git a/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl b/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl index c1b091fb..c23d3385 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_rheology.jl @@ -169,19 +169,19 @@ function init_phases!(phases, particles, Lx, d, r, thick_air, CharDim) x = JustRelax.@cell px[ip, i, j] depth = -(JustRelax.@cell py[ip, i, j]) - nondimensionalize(thick_air * km, CharDim) if nondimensionalize(0e0km, CharDim) ≤ depth ≤ nondimensionalize(21km, CharDim) - @cell phases[ip, i, j] = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 elseif nondimensionalize(35km, CharDim) ≥ depth > nondimensionalize(21km, CharDim) - @cell phases[ip, i, j] = 2.0 + JustRelax.@cell phases[ip, i, j] = 2.0 elseif nondimensionalize(90km, CharDim) ≥ depth > nondimensionalize(35km, CharDim) - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 elseif depth > nondimensionalize(90km, CharDim) - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 elseif depth < nondimensionalize(0e0km, CharDim) - @cell phases[ip, i, j] = 5.0 + JustRelax.@cell phases[ip, i, j] = 5.0 end @@ -193,6 +193,6 @@ function init_phases!(phases, particles, Lx, d, r, thick_air, CharDim) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, CharDim) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx, CharDim) end diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl index 8efc2da1..ff7c1912 100644 --- a/miniapps/convection/Particles3D/Layered_convection3D.jl +++ b/miniapps/convection/Particles3D/Layered_convection3D.jl @@ -1,21 +1,19 @@ -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell + +const backend_JR = CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) -using JustPIC -using JustPIC._3D +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend - -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) # or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) -environment!(model) +const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays +using Printf, GeoParams, GLMakie, GeoParams # Load file with all the rheology configurations include("Layered_rheology.jl") @@ -116,18 +114,18 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) zc_anomaly = -610e3 # origin of thermal anomaly r_anomaly = 50e3 # radius of perturbation init_phases!(pPhases, particles, lx, ly; d=abs(zc_anomaly), r=r_anomaly) - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend_JR, 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) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.5 / √3.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) @@ -135,26 +133,21 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) @parallel init_T!(thermal.T, xvi[3]) rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly, xvi) thermal_bcs!(thermal.T, thermal_bc) - @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # 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 + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + @parallel init_P!(stokes.P, ρg[end], xci[end]) + # Rheology - η = @ones(ni...) - args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) - ) - η_vep = deepcopy(η) + viscosity_cutoff = (1e18, 1e24) + compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( - rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 + backend_JR, rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √3 ) # Boundary conditions @@ -208,11 +201,12 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) temperature2center!(thermal) # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt=dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) + compute_ρg!(ρg[end], phase_ratios, rheology, args) + compute_viscosity!( + stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff ) - @parallel (@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + # ------------------------------ # Stokes solver ---------------- solve!( @@ -221,18 +215,18 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax = 100e3, - nout = 1e3, - viscosity_cutoff = (1e18, 1e24) + kwargs =(; + iterMax = 100e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) ); - @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) / 2 # ------------------------------ @@ -245,11 +239,13 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kawargs = (; + 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 @@ -262,15 +258,14 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy, grid_vz), dt) # 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, ), (thermal.T,), xvi) + 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) - + phase_ratios_center(phase_ratios, particles, grid, pPhases) + @show it += 1 t += dt diff --git a/miniapps/convection/Particles3D/Layered_rheology.jl b/miniapps/convection/Particles3D/Layered_rheology.jl index bff86b86..10c9b840 100644 --- a/miniapps/convection/Particles3D/Layered_rheology.jl +++ b/miniapps/convection/Particles3D/Layered_rheology.jl @@ -121,7 +121,7 @@ function init_phases!(phases, particles, Lx, Ly; d=650e3, r=50e3) @parallel_indices (I...) function init_phases!(phases, px, py, pz, index, r, Lx, Ly) - @inbounds for ip in JustRelax.JustRelax.cellaxes(phases) + @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, I...]) == 0 && continue @@ -130,19 +130,19 @@ function init_phases!(phases, particles, Lx, Ly; d=650e3, r=50e3) depth = -(JustRelax.@cell pz[ip, I...]) if 0e0 ≤ depth ≤ 21e3 - @cell phases[ip, I...] = 1.0 + JustRelax.@cell phases[ip, I...] = 1.0 elseif 35e3 ≥ depth > 21e3 - @cell phases[ip, I...] = 2.0 + JustRelax.@cell phases[ip, I...] = 2.0 elseif 90e3 ≥ depth > 35e3 - @cell phases[ip, I...] = 3.0 + JustRelax.@cell phases[ip, I...] = 3.0 elseif depth > 90e3 - @cell phases[ip, I...] = 3.0 + JustRelax.@cell phases[ip, I...] = 3.0 elseif 0e0 > depth - @cell phases[ip, I...] = 5.0 + JustRelax.@cell phases[ip, I...] = 5.0 end From e7cde5f50f5ce5a496a7bf7d755053d6738eb5eb Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 15:30:18 +0200 Subject: [PATCH 68/91] update tests --- test/test_diffusion2D.jl | 4 +++- test/test_diffusion3D.jl | 2 +- test/test_shearband2D.jl | 4 ++-- test/test_shearheating2D.jl | 1 + test/test_sinking_block.jl | 5 +---- 5 files changed, 8 insertions(+), 8 deletions(-) diff --git a/test/test_diffusion2D.jl b/test/test_diffusion2D.jl index 39689add..21e3276c 100644 --- a/test/test_diffusion2D.jl +++ b/test/test_diffusion2D.jl @@ -98,7 +98,9 @@ function diffusion_2D(; nx=32, ny=32, lx=100e3, ly=100e3, ρ0=3.3e3, Cp0=1.2e3, args, dt, di; - kwargs = (; verbose=false), + kwargs = (; + verbose = false + ), ) t += dt diff --git a/test/test_diffusion3D.jl b/test/test_diffusion3D.jl index c64500b2..92126ed7 100644 --- a/test/test_diffusion3D.jl +++ b/test/test_diffusion3D.jl @@ -1,7 +1,7 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using GeoParams, CellArrays +using GeoParams using JustRelax, JustRelax.JustRelax3D using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index 56879450..2400e824 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -1,7 +1,7 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using GeoParams, CellArrays +using GeoParams using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) @@ -91,7 +91,7 @@ function ShearBand2D() # Initialize phase ratios ------------------------------- radius = 0.1 - phase_ratios = PhaseRatio(ni, length(rheology)) + phase_ratios = PhaseRatio(backend, ni, length(rheology)) init_phases!(phase_ratios, xci, radius) # STOKES --------------------------------------------- diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 86a0f1c1..7444d7c3 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -145,6 +145,7 @@ function Shearheating2D() t, it = 0.0, 0 local iters, thermal while it < 10 + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index bc8bce16..e07e837d 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -10,7 +10,7 @@ const backend_JR = CPUBackend using JustPIC, JustPIC._2D const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -using Printf, LinearAlgebra, GeoParams, CellArrays +using GeoParams ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- @@ -151,7 +151,6 @@ function Sinking_Block2D() # Stokes solver ---------------- args = (; T = @ones(ni...), P = stokes.P, dt=dt, ΔTc = @zeros(ni...)) - local iters, velocity iters = solve!( stokes, pt_stokes, @@ -181,8 +180,6 @@ function Sinking_Block2D() finalize_global_grid(; finalize_MPI = true) return iters, velocity - - end @testset "Sinking_Block2D" begin From 827175b4618b1caea3d7c0a09720a8e6cf47298d Mon Sep 17 00:00:00 2001 From: aelligp Date: Tue, 30 Apr 2024 15:50:20 +0200 Subject: [PATCH 69/91] add DocPreviewCleanup --- .github/workflows/DocPreviewCleanup.yml | 34 +++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 .github/workflows/DocPreviewCleanup.yml diff --git a/.github/workflows/DocPreviewCleanup.yml b/.github/workflows/DocPreviewCleanup.yml new file mode 100644 index 00000000..17b135ff --- /dev/null +++ b/.github/workflows/DocPreviewCleanup.yml @@ -0,0 +1,34 @@ +name: Doc Preview Cleanup + +on: + pull_request: + types: [closed] + +jobs: + doc-preview-cleanup: + # Do not run on forks to avoid authorization errors + # Source: https://github.community/t/have-github-action-only-run-on-master-repo-and-not-on-forks/140840/18 + # Note: This does not always work as intended - but you can just ignore + # the failed CI runs after merging a PR + if: github.repository_owner == 'PTSolvers' + runs-on: ubuntu-latest + steps: + - name: Checkout gh-pages branch + uses: actions/checkout@v4 + with: + ref: gh-pages + + - name: Delete preview and history + shell: bash + run: | + git config user.name "Documenter.jl" + git config user.email "documenter@juliadocs.github.io" + git rm -rf --ignore-unmatch "previews/PR$PRNUM" + git commit -m "delete preview" --allow-empty + git branch gh-pages-new $(echo "delete history" | git commit-tree HEAD^{tree}) + env: + PRNUM: ${{ github.event.number }} + + - name: Push changes + run: | + git push --force origin gh-pages-new:gh-pages From cbf9d170cb84d611d2ef4d5007546113b5cd9f32 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 17:02:59 +0200 Subject: [PATCH 70/91] miniapp fixes --- .../PlumeFreeSurface_2D.jl | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 793dd4f0..ed3831ff 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -40,20 +40,20 @@ function init_phases!(phases, particles) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape - @cell(index[ip, i, j]) == 0 && continue + JustRelax.@cell(index[ip, i, j]) == 0 && continue - x = @cell px[ip, i, j] - depth = -(@cell py[ip, i, j]) - @cell phases[ip, i, j] = 2.0 + x = JustRelax.@cell px[ip, i, j] + depth = -(JustRelax.@cell py[ip, i, j]) + JustRelax.@cell phases[ip, i, j] = 2.0 if 0e0 ≤ depth ≤ 100e3 - @cell phases[ip, i, j] = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 else - @cell phases[ip, i, j] = 2.0 + JustRelax.@cell phases[ip, i, j] = 2.0 if ((x - 250e3)^2 + (depth - 250e3)^2 ≤ r^2) - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 end end @@ -143,6 +143,7 @@ function main(igg, nx, ny) # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right=true, top=true, bot=true), + # free_surface = true ) # Plot initial T and η profiles @@ -188,9 +189,10 @@ function main(igg, nx, ny) dt, igg; kwargs = (; - iterMax = 150e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = (-Inf, Inf), + free_surface = true ) ) dt = compute_dt(stokes, di) / 2 @@ -210,7 +212,7 @@ function main(igg, nx, ny) t += dt if it == 1 || rem(it, 5) == 0 - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) nt = 5 fig = Figure(resolution = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") From cd256d71ed11a59b3dd4be43a6a671d3398e1116 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Tue, 30 Apr 2024 17:06:59 +0200 Subject: [PATCH 71/91] up miniapp --- .../RayleighTaylor2D.jl | 33 +++++-------------- 1 file changed, 9 insertions(+), 24 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 891f7350..c3e11063 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -40,24 +40,25 @@ function init_phases!(phases, particles, A) @inbounds for ip in JustRelax.cellaxes(phases) # quick escape - @cell(index[ip, i, j]) == 0 && continue + JustRelax.@cell(index[ip, i, j]) == 0 && continue - x = @cell px[ip, i, j] - depth = -(@cell py[ip, i, j]) - @cell phases[ip, i, j] = 2.0 + x = JustRelax.@cell px[ip, i, j] + depth = -(JustRelax.@cell py[ip, i, j]) + JustRelax.@cell phases[ip, i, j] = 2.0 if 0e0 ≤ depth ≤ 100e3 - @cell phases[ip, i, j] = 1.0 + JustRelax.@cell phases[ip, i, j] = 1.0 elseif depth > (-f(x, A, 500e3) + (200e3 - A)) - @cell phases[ip, i, j] = 3.0 + JustRelax.@cell phases[ip, i, j] = 3.0 + end end return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, A) + @parallel (@idx ni) init_phases!(phases, particles.coords..., particles.index, A) end ## END OF HELPER FUNCTION ------------------------------------------------------------ @@ -117,7 +118,7 @@ function RT_2D(igg, nx, ny) A = 5e3 # Amplitude of the anomaly phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, A) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- @@ -144,20 +145,6 @@ function RT_2D(igg, nx, ny) free_surface = true, ) - # Plot initial T and η profiles - let - Y = [y for x in xci[1], y in xci[2]][:] - fig = Figure(size = (1200, 900)) - ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") - ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") - scatter!(ax1, Array(ρg[2][:]./9.81), Y./1e3) - scatter!(ax2, Array(log10.(η[:])), Y./1e3) - ylims!(ax1, minimum(xvi[2])./1e3, 0) - ylims!(ax2, minimum(xvi[2])./1e3, 0) - hideydecorations!(ax2) - fig - end - Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) @@ -181,8 +168,6 @@ function RT_2D(igg, nx, ny) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, From 015e8848fff99e4c018a86f57ba5b1792abf4fa1 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 1 May 2024 14:28:56 +0200 Subject: [PATCH 72/91] up miniapp --- .../PlumeFreeSurface_2D.jl | 30 +++++-------------- 1 file changed, 7 insertions(+), 23 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index ed3831ff..79d996c0 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -142,25 +142,9 @@ function main(igg, nx, ny) # Boundary conditions flow_bcs = FlowBoundaryConditions(; - free_slip = (left = true, right=true, top=true, bot=true), - # free_surface = true + free_slip = (left = true, right = true, top = true, bot = true), ) - # Plot initial T and η profiles - let - Y = [y for x in xci[1], y in xci[2]][:] - fig = Figure(size = (1200, 900)) - ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") - ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") - scatter!(ax1, Array(ρg[2][:]./9.81), Y./1e3) - scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) - # scatter!(ax2, Array(stokes.P[:]), Y./1e3) - ylims!(ax1, minimum(xvi[2])./1e3, 0) - ylims!(ax2, minimum(xvi[2])./1e3, 0) - hideydecorations!(ax2) - fig - end - Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) @@ -170,13 +154,13 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 dt = 1e3 * (3600 * 24 *365.25) - while it < 10 # run only for 5 Myrs + while it < 15 # run only for 5 Myrs # Stokes solver ---------------- args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) - + solve!( stokes, pt_stokes, @@ -189,7 +173,7 @@ function main(igg, nx, ny) dt, igg; kwargs = (; - iterMax = 150e3, + iterMax = 50e3, nout = 1e3, viscosity_cutoff = (-Inf, Inf), free_surface = true @@ -207,14 +191,14 @@ function main(igg, nx, ny) inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios phase_ratios_center(phase_ratios, particles, grid, pPhases) - + @show it += 1 t += dt - if it == 1 || rem(it, 5) == 0 + if it == 1 || rem(it, 1) == 0 velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) nt = 5 - fig = Figure(resolution = (900, 900), title = "t = $t") + fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") heatmap!(ax, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η)), colormap = :grayC) arrows!( From bfd4fad1f356ab6259a1c5c618cb97e987e1587c Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 1 May 2024 19:05:01 +0200 Subject: [PATCH 73/91] fix free surface miniapps --- .../PlumeFreeSurface_2D.jl | 3 ++- .../RayleighTaylor2D.jl | 10 ++++---- src/boundaryconditions/BoundaryConditions.jl | 25 +++++-------------- 3 files changed, 13 insertions(+), 25 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 79d996c0..2f8ac79d 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -143,6 +143,7 @@ function main(igg, nx, ny) # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right = true, top = true, bot = true), + free_surface = true ) Vx_v = @zeros(ni.+1...) @@ -153,7 +154,7 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 *365.25) + dt = 1e3 * (3600 * 24 * 365.25) while it < 15 # run only for 5 Myrs # Stokes solver ---------------- diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index c3e11063..15fe60b4 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -124,7 +124,7 @@ function RT_2D(igg, nx, ny) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend_JR, ni) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-8, CFL = 0.95 / √2.1) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- @@ -142,6 +142,7 @@ function RT_2D(igg, nx, ny) flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right = true, top = true, bot = false), no_slip = (left = false, right = false, top = false, bot = true), + # free_slip = (left = true, right = true, top = true, bot = true), free_surface = true, ) @@ -174,9 +175,9 @@ function RT_2D(igg, nx, ny) dt, igg; kwargs = ( - iterMax = 150e3, - iterMin = 5e3, - viscosity_relaxation = 1e-2, + iterMax = 50e3, + iterMin = 1e3, + # viscosity_relaxation = 1e-2, nout = 5e3, free_surface = true, viscosity_cutoff = (-Inf, Inf) @@ -201,7 +202,6 @@ function RT_2D(igg, nx, ny) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios - # @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 3ac377d1..4d4ea76c 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -214,20 +214,6 @@ function free_surface_bcs!( end end -# function free_surface_bcs!(stokes::JustRelax.StokesArrays, bcs::FlowBoundaryConditions) -# if bcs.free_surface -# @views stokes.τ.yy[:, end] .= 0.0 -# end -# end - -# function free_surface_bcs!( -# stokes::JustRelax.StokesArrays{A,B,C,D,E,3}, bcs::FlowBoundaryConditions -# ) where {A,B,C,D,E} -# if bcs.free_surface -# @views stokes.τ.zz[:, :, end] .= 0.0 -# end -# end - @parallel_indices (i) function FreeSurface_Vy!( Vx::AbstractArray{T,2}, Vy::AbstractArray{T,2}, @@ -243,15 +229,16 @@ end ) where {T} phase = @inbounds phase_ratios[i, end] Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt + ν = 1e-2 Vy[i + 1, end] = - Vy[i + 1, end - 1] + - 3.0 / 2.0 * + ν * (Vy[i + 1, end - 1] + + (3 / 2) * ( - P[i, end] / (2.0 * η[i, end]) - - (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + + P[i, end] / (2.0 * η[i, end]) + #- + # (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + inv(3.0) * (Vx[i + 1, end - 1] - Vx[i, end - 1]) * inv(dx) ) * - dy + dy) + (1 - ν) * Vy[i + 1, end] return nothing end From c4f4df4350c0d1027e16a85f0b537cc933024bde Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 1 May 2024 19:07:17 +0200 Subject: [PATCH 74/91] format --- src/boundaryconditions/BoundaryConditions.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 4d4ea76c..763fd38a 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -231,14 +231,16 @@ end Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt ν = 1e-2 Vy[i + 1, end] = - ν * (Vy[i + 1, end - 1] + - (3 / 2) * - ( - P[i, end] / (2.0 * η[i, end]) + #- - # (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + - inv(3.0) * (Vx[i + 1, end - 1] - Vx[i, end - 1]) * inv(dx) - ) * - dy) + (1 - ν) * Vy[i + 1, end] + ν * ( + Vy[i + 1, end - 1] + + (3 / 2) * + ( + P[i, end] / (2.0 * η[i, end]) + #- + (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + + inv(3.0) * (Vx[i + 1, end - 1] - Vx[i, end - 1]) * inv(dx) + ) * + dy + ) + (1 - ν) * Vy[i + 1, end] return nothing end From 3da47fa287b303c653fb9373247b1683f9b780d9 Mon Sep 17 00:00:00 2001 From: Albert de Montserrat Date: Wed, 1 May 2024 22:21:28 +0200 Subject: [PATCH 75/91] some polishing --- src/boundaryconditions/BoundaryConditions.jl | 315 +------------------ src/boundaryconditions/free_slip.jl | 175 +++++++++++ src/boundaryconditions/free_surface.jl | 86 +++++ src/boundaryconditions/no_slip.jl | 25 ++ src/boundaryconditions/pure_shear.jl | 23 ++ src/stokes/Stokes2D.jl | 4 +- 6 files changed, 316 insertions(+), 312 deletions(-) create mode 100644 src/boundaryconditions/free_slip.jl create mode 100644 src/boundaryconditions/free_surface.jl create mode 100644 src/boundaryconditions/no_slip.jl create mode 100644 src/boundaryconditions/pure_shear.jl diff --git a/src/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 763fd38a..bdc9efeb 100644 --- a/src/boundaryconditions/BoundaryConditions.jl +++ b/src/boundaryconditions/BoundaryConditions.jl @@ -58,314 +58,7 @@ function _flow_bcs!(bcs, V) end # BOUNDARY CONDITIONS KERNELS - -@parallel_indices (i) function no_slip!(Ax, Ay, bc) - @inbounds begin - if bc.left - (i ≤ size(Ax, 2)) && (Ax[1, i] = 0.0) - (1 < i < size(Ay, 2)) && (Ay[1, i] = -Ay[2, i]) - end - if bc.right - (i ≤ size(Ax, 2)) && (Ax[end, i] = 0.0) - (1 < i < size(Ay, 2)) && (Ay[end, i] = -Ay[end - 1, i]) - end - if bc.bot - (i ≤ size(Ay, 1)) && (Ay[i, 1] = 0.0) - (1 < i < size(Ax, 1)) && (Ax[i, 1] = -Ax[i, 2]) - end - if bc.top - (i ≤ size(Ay, 1)) && (Ay[i, end] = 0.0) - (1 < i < size(Ax, 1)) && (Ax[i, end] = -Ax[i, end - 1]) - end - # corners - # bc.bot && (Ax[1, 1] = 0.0; Ax[1, 1] = 0.0) - # bc.left && bc.bot && (Ax[1, 1] = 0.0) - # bc.right && bc.top && (Ay[end, end] = 0.0) - end - return nothing -end - -@parallel_indices (i) function free_slip!(Ax, Ay, bc) - @inbounds begin - if i ≤ size(Ax, 1) - bc.bot && (Ax[i, 1] = Ax[i, 2]) - bc.top && (Ax[i, end] = Ax[i, end - 1]) - end - if i ≤ size(Ay, 2) - bc.left && (Ay[1, i] = Ay[2, i]) - bc.right && (Ay[end, i] = Ay[end - 1, i]) - end - end - return nothing -end - -@parallel_indices (i, j) function free_slip!(Ax, Ay, Az, bc) - @inbounds begin - # free slip in the front and back XZ planes - if bc.front - if i ≤ size(Ax, 1) && j ≤ size(Ax, 3) - Ax[i, 1, j] = Ax[i, 2, j] - end - if i ≤ size(Az, 1) && j ≤ size(Az, 3) - Az[i, 1, j] = Az[i, 2, j] - end - end - if bc.back - if i ≤ size(Ax, 1) && j ≤ size(Ax, 3) - Ax[i, end, j] = Ax[i, end - 1, j] - end - if i ≤ size(Az, 1) && j ≤ size(Az, 3) - Az[i, end, j] = Az[i, end - 1, j] - end - end - # free slip in the front and back XY planes - if bc.top - if i ≤ size(Ax, 1) && j ≤ size(Ax, 2) - Ax[i, j, 1] = Ax[i, j, 2] - end - if i ≤ size(Ay, 1) && j ≤ size(Ay, 2) - Ay[i, j, 1] = Ay[i, j, 2] - end - end - if bc.bot - if i ≤ size(Ax, 1) && j ≤ size(Ax, 2) - Ax[i, j, end] = Ax[i, j, end - 1] - end - if i ≤ size(Ay, 1) && j ≤ size(Ay, 2) - Ay[i, j, end] = Ay[i, j, end - 1] - end - end - # free slip in the front and back YZ planes - if bc.left - if i ≤ size(Ay, 2) && j ≤ size(Ay, 3) - Ay[1, i, j] = Ay[2, i, j] - end - if i ≤ size(Az, 2) && j ≤ size(Az, 3) - Az[1, i, j] = Az[2, i, j] - end - end - if bc.right - if i ≤ size(Ay, 2) && j ≤ size(Ay, 3) - Ay[end, i, j] = Ay[end - 1, i, j] - end - if i ≤ size(Az, 2) && j ≤ size(Az, 3) - Az[end, i, j] = Az[end - 1, i, j] - end - end - end - return nothing -end - -@parallel_indices (i) function free_slip!(T::_T, bc) where {_T<:AbstractArray{<:Any,2}} - @inbounds begin - if i ≤ size(T, 1) - bc.bot && (T[i, 1] = T[i, 2]) - bc.top && (T[i, end] = T[i, end - 1]) - end - if i ≤ size(T, 2) - bc.left && (T[1, i] = T[2, i]) - bc.right && (T[end, i] = T[end - 1, i]) - end - end - return nothing -end - -@parallel_indices (i, j) function free_slip!(T::_T, bc) where {_T<:AbstractArray{<:Any,3}} - nx, ny, nz = size(T) - @inbounds begin - if i ≤ nx && j ≤ ny - bc.bot && (T[i, j, 1] = T[i, j, 2]) - bc.top && (T[i, j, end] = T[i, j, end - 1]) - end - if i ≤ ny && j ≤ nz - bc.left && (T[1, i, j] = T[2, i, j]) - bc.right && (T[end, i, j] = T[end - 1, i, j]) - end - if i ≤ nx && j ≤ nz - bc.front && (T[i, 1, j] = T[i, 2, j]) - bc.back && (T[i, end, j] = T[i, end - 1, j]) - end - end - return nothing -end - -function free_surface_bcs!( - stokes, bcs::FlowBoundaryConditions, η, rheology, phase_ratios, dt, di -) - indices_range(::Any, Vy) = @idx (size(Vy, 1) - 2) - indices_range(::Any, ::Any, Vz) = @idx (size(Vz, 1) - 2, size(Vz, 2) - 2) - - V = @velocity(stokes) - n = indices_range(V...) - - if bcs.free_surface - # apply boundary conditions - @parallel n FreeSurface_Vy!( - V..., - stokes.P, - stokes.P0, - stokes.τ_o.yy, - η, - rheology, - phase_ratios.center, - dt, - di..., - ) - end -end - -@parallel_indices (i) function FreeSurface_Vy!( - Vx::AbstractArray{T,2}, - Vy::AbstractArray{T,2}, - P::AbstractArray{T,2}, - P_old::AbstractArray{T,2}, - τyy_old::AbstractArray{T,2}, - η::AbstractArray{T,2}, - rheology, - phase_ratios, - dt::T, - dx::T, - dy::T, -) where {T} - phase = @inbounds phase_ratios[i, end] - Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt - ν = 1e-2 - Vy[i + 1, end] = - ν * ( - Vy[i + 1, end - 1] + - (3 / 2) * - ( - P[i, end] / (2.0 * η[i, end]) + #- - (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + - inv(3.0) * (Vx[i + 1, end - 1] - Vx[i, end - 1]) * inv(dx) - ) * - dy - ) + (1 - ν) * Vy[i + 1, end] - return nothing -end - -@parallel_indices (i, j) function FreeSurface_Vy!( - Vx::AbstractArray{T,3}, - Vy::AbstractArray{T,3}, - Vz::AbstractArray{T,3}, - P::AbstractArray{T,3}, - P_old::AbstractArray{T,3}, - τyy_old::AbstractArray{T,3}, - η::AbstractArray{T,3}, - rheology, - phase_ratios, - dt::T, - dx::T, - dy::T, - dz::T, -) where {T} - phase = @inbounds phase_ratios[i, j, end] - Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt - Vz[i + 1, j + 1, end] = - Vz[i + 1, j + 1, end - 1] + - 3.0 / 2.0 * - ( - P[i, j, end] / (2.0 * η[i, j, end]) - - (τyy_old[i, j, end] + P_old[i, j, end]) / (2.0 * Gdt) + - inv(3.0) * ( - (Vx[i + 1, j + 1, end - 1] - Vx[i, j + 1, end - 1]) * inv(dx) + - (Vy[i + 1, j + 1, end - 1] - Vy[i + 1, j, end - 1]) * inv(dy) - ) - ) * - dz - return nothing -end - -function pureshear_bc!( - stokes::JustRelax.StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg -) where {T} - _T = typeof(stokes.V.Vx) - stokes.V.Vx[:, 2:(end - 1)] .= _T([εbg * x for x in xvi[1], y in xci[2]]) - stokes.V.Vy[2:(end - 1), :] .= _T([-εbg * y for x in xci[1], y in xvi[2]]) - - return nothing -end - -function pureshear_bc!( - stokes::JustRelax.StokesArrays, xci::NTuple{3,T}, xvi::NTuple{3,T}, εbg -) where {T} - xv, yv, zv = xvi - xc, yc, zc = xci - _T = typeof(stokes.V.Vx) - - stokes.V.Vx[:, 2:(end - 1), 2:(end - 1)] .= _T([εbg * x for x in xv, y in yc, z in zc]) - stokes.V.Vy[2:(end - 1), :, 2:(end - 1)] .= _T([εbg * y for x in xc, y in xv, z in zc]) - stokes.V.Vz[2:(end - 1), 2:(end - 1), :] .= _T([-εbg * z for x in xc, y in xc, z in zv]) - - return nothing -end - -@parallel_indices (j) function free_slip_x!(A::AbstractArray{T,2}) where {T} - A[1, j] = A[2, j] - A[end, j] = A[end - 1, j] - return nothing -end - -@parallel_indices (i) function free_slip_y!(A::AbstractArray{T,2}) where {T} - A[i, 1] = A[i, 2] - A[i, end] = A[i, end - 1] - return nothing -end - -@inbounds @parallel_indices (i) function _apply_free_slip!(Ax, Ay, freeslip_x, freeslip_y) - if freeslip_x && i ≤ size(Ax, 1) - Ax[i, 1] = Ax[i, 2] - Ax[i, end] = Ax[i, end - 1] - end - if freeslip_y && i ≤ size(Ay, 2) - Ay[1, i] = Ay[2, i] - Ay[end, i] = Ay[end - 1, i] - end - return nothing -end - -function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{2,T}}, Vx, Vy) where {T} - freeslip_x, freeslip_y = freeslip - n = max(size(Vx, 1), size(Vy, 2)) - # free slip boundary conditions - @parallel (1:n) _apply_free_slip!(Vx, Vy, freeslip_x, freeslip_y) - - return nothing -end - -# 3D KERNELS - -@parallel_indices (j, k) function free_slip_x!(A::AbstractArray{T,3}) where {T} - A[1, j, k] = A[2, j, k] - A[end, j, k] = A[end - 1, j, k] - return nothing -end - -@parallel_indices (i, k) function free_slip_y!(A::AbstractArray{T,3}) where {T} - A[i, 1, k] = A[i, 2, k] - A[i, end, k] = A[i, end - 1, k] - return nothing -end - -@parallel_indices (i, j) function free_slip_z!(A::AbstractArray{T,3}) where {T} - A[i, j, 1] = A[i, j, 2] - A[i, j, end] = A[i, j, end - 1] - return nothing -end - -function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{3,T}}, Vx, Vy, Vz) where {T} - freeslip_x, freeslip_y, freeslip_z = freeslip - # free slip boundary conditions - if freeslip_x - @parallel (1:size(Vy, 2), 1:size(Vy, 3)) free_slip_x!(Vy) - @parallel (1:size(Vz, 2), 1:size(Vz, 3)) free_slip_x!(Vz) - end - if freeslip_y - @parallel (1:size(Vx, 1), 1:size(Vx, 3)) free_slip_y!(Vx) - @parallel (1:size(Vz, 1), 1:size(Vz, 3)) free_slip_y!(Vz) - end - if freeslip_z - @parallel (1:size(Vx, 1), 1:size(Vx, 2)) free_slip_z!(Vx) - @parallel (1:size(Vy, 1), 1:size(Vy, 2)) free_slip_z!(Vy) - end -end +include("free_slip.jl") +include("free_surface.jl") +include("no_slip.jl") +include("pure_shear.jl") diff --git a/src/boundaryconditions/free_slip.jl b/src/boundaryconditions/free_slip.jl new file mode 100644 index 00000000..a9348db0 --- /dev/null +++ b/src/boundaryconditions/free_slip.jl @@ -0,0 +1,175 @@ +@parallel_indices (i) function free_slip!(Ax, Ay, bc) + @inbounds begin + if i ≤ size(Ax, 1) + bc.bot && (Ax[i, 1] = Ax[i, 2]) + bc.top && (Ax[i, end] = Ax[i, end - 1]) + end + if i ≤ size(Ay, 2) + bc.left && (Ay[1, i] = Ay[2, i]) + bc.right && (Ay[end, i] = Ay[end - 1, i]) + end + end + return nothing +end + +@parallel_indices (i, j) function free_slip!(Ax, Ay, Az, bc) + @inbounds begin + # free slip in the front and back XZ planes + if bc.front + if i ≤ size(Ax, 1) && j ≤ size(Ax, 3) + Ax[i, 1, j] = Ax[i, 2, j] + end + if i ≤ size(Az, 1) && j ≤ size(Az, 3) + Az[i, 1, j] = Az[i, 2, j] + end + end + if bc.back + if i ≤ size(Ax, 1) && j ≤ size(Ax, 3) + Ax[i, end, j] = Ax[i, end - 1, j] + end + if i ≤ size(Az, 1) && j ≤ size(Az, 3) + Az[i, end, j] = Az[i, end - 1, j] + end + end + # free slip in the front and back XY planes + if bc.top + if i ≤ size(Ax, 1) && j ≤ size(Ax, 2) + Ax[i, j, 1] = Ax[i, j, 2] + end + if i ≤ size(Ay, 1) && j ≤ size(Ay, 2) + Ay[i, j, 1] = Ay[i, j, 2] + end + end + if bc.bot + if i ≤ size(Ax, 1) && j ≤ size(Ax, 2) + Ax[i, j, end] = Ax[i, j, end - 1] + end + if i ≤ size(Ay, 1) && j ≤ size(Ay, 2) + Ay[i, j, end] = Ay[i, j, end - 1] + end + end + # free slip in the front and back YZ planes + if bc.left + if i ≤ size(Ay, 2) && j ≤ size(Ay, 3) + Ay[1, i, j] = Ay[2, i, j] + end + if i ≤ size(Az, 2) && j ≤ size(Az, 3) + Az[1, i, j] = Az[2, i, j] + end + end + if bc.right + if i ≤ size(Ay, 2) && j ≤ size(Ay, 3) + Ay[end, i, j] = Ay[end - 1, i, j] + end + if i ≤ size(Az, 2) && j ≤ size(Az, 3) + Az[end, i, j] = Az[end - 1, i, j] + end + end + end + return nothing +end + +@parallel_indices (i) function free_slip!(T::_T, bc) where {_T<:AbstractArray{<:Any,2}} + @inbounds begin + if i ≤ size(T, 1) + bc.bot && (T[i, 1] = T[i, 2]) + bc.top && (T[i, end] = T[i, end - 1]) + end + if i ≤ size(T, 2) + bc.left && (T[1, i] = T[2, i]) + bc.right && (T[end, i] = T[end - 1, i]) + end + end + return nothing +end + +@parallel_indices (i, j) function free_slip!(T::_T, bc) where {_T<:AbstractArray{<:Any,3}} + nx, ny, nz = size(T) + @inbounds begin + if i ≤ nx && j ≤ ny + bc.bot && (T[i, j, 1] = T[i, j, 2]) + bc.top && (T[i, j, end] = T[i, j, end - 1]) + end + if i ≤ ny && j ≤ nz + bc.left && (T[1, i, j] = T[2, i, j]) + bc.right && (T[end, i, j] = T[end - 1, i, j]) + end + if i ≤ nx && j ≤ nz + bc.front && (T[i, 1, j] = T[i, 2, j]) + bc.back && (T[i, end, j] = T[i, end - 1, j]) + end + end + return nothing +end + + + +@parallel_indices (j) function free_slip_x!(A::AbstractArray{T,2}) where {T} + A[1, j] = A[2, j] + A[end, j] = A[end - 1, j] + return nothing +end + +@parallel_indices (i) function free_slip_y!(A::AbstractArray{T,2}) where {T} + A[i, 1] = A[i, 2] + A[i, end] = A[i, end - 1] + return nothing +end + +@inbounds @parallel_indices (i) function _apply_free_slip!(Ax, Ay, freeslip_x, freeslip_y) + if freeslip_x && i ≤ size(Ax, 1) + Ax[i, 1] = Ax[i, 2] + Ax[i, end] = Ax[i, end - 1] + end + if freeslip_y && i ≤ size(Ay, 2) + Ay[1, i] = Ay[2, i] + Ay[end, i] = Ay[end - 1, i] + end + return nothing +end + +function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{2,T}}, Vx, Vy) where {T} + freeslip_x, freeslip_y = freeslip + n = max(size(Vx, 1), size(Vy, 2)) + # free slip boundary conditions + @parallel (1:n) _apply_free_slip!(Vx, Vy, freeslip_x, freeslip_y) + + return nothing +end + +# 3D KERNELS + +@parallel_indices (j, k) function free_slip_x!(A::AbstractArray{T,3}) where {T} + A[1, j, k] = A[2, j, k] + A[end, j, k] = A[end - 1, j, k] + return nothing +end + +@parallel_indices (i, k) function free_slip_y!(A::AbstractArray{T,3}) where {T} + A[i, 1, k] = A[i, 2, k] + A[i, end, k] = A[i, end - 1, k] + return nothing +end + +@parallel_indices (i, j) function free_slip_z!(A::AbstractArray{T,3}) where {T} + A[i, j, 1] = A[i, j, 2] + A[i, j, end] = A[i, j, end - 1] + return nothing +end + +function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{3,T}}, Vx, Vy, Vz) where {T} + freeslip_x, freeslip_y, freeslip_z = freeslip + # free slip boundary conditions + if freeslip_x + @parallel (1:size(Vy, 2), 1:size(Vy, 3)) free_slip_x!(Vy) + @parallel (1:size(Vz, 2), 1:size(Vz, 3)) free_slip_x!(Vz) + end + if freeslip_y + @parallel (1:size(Vx, 1), 1:size(Vx, 3)) free_slip_y!(Vx) + @parallel (1:size(Vz, 1), 1:size(Vz, 3)) free_slip_y!(Vz) + end + if freeslip_z + @parallel (1:size(Vx, 1), 1:size(Vx, 2)) free_slip_z!(Vx) + @parallel (1:size(Vy, 1), 1:size(Vy, 2)) free_slip_z!(Vy) + end +end \ No newline at end of file diff --git a/src/boundaryconditions/free_surface.jl b/src/boundaryconditions/free_surface.jl new file mode 100644 index 00000000..d130b331 --- /dev/null +++ b/src/boundaryconditions/free_surface.jl @@ -0,0 +1,86 @@ +function free_surface_bcs!( + stokes, bcs::FlowBoundaryConditions, η, rheology, phase_ratios, dt, di +) + indices_range(::Any, Vy) = @idx (size(Vy, 1) - 2) + indices_range(::Any, ::Any, Vz) = @idx (size(Vz, 1) - 2, size(Vz, 2) - 2) + + V = @velocity(stokes) + n = indices_range(V...) + + if bcs.free_surface + # apply boundary conditions + @parallel n FreeSurface_Vy!( + V..., + stokes.P, + stokes.P0, + stokes.τ_o.yy, + η, + rheology, + phase_ratios.center, + dt, + di..., + ) + end +end + +@parallel_indices (i) function FreeSurface_Vy!( + Vx::AbstractArray{T,2}, + Vy::AbstractArray{T,2}, + P::AbstractArray{T,2}, + P_old::AbstractArray{T,2}, + τyy_old::AbstractArray{T,2}, + η::AbstractArray{T,2}, + rheology, + phase_ratios, + dt::T, + dx::T, + dy::T, +) where {T} + phase = @inbounds phase_ratios[i, end] + Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt + ν = 1e-2 + Vy[i + 1, end] = + ν * ( + Vy[i + 1, end - 1] + + (3 / 2) * + ( + P[i, end] / (2.0 * η[i, end]) + #- + (τyy_old[i, end] + P_old[i, end]) / (2.0 * Gdt) + + inv(3.0) * (Vx[i + 1, end - 1] - Vx[i, end - 1]) * inv(dx) + ) * + dy + ) + (1 - ν) * Vy[i + 1, end] + return nothing +end + +@parallel_indices (i, j) function FreeSurface_Vy!( + Vx::AbstractArray{T,3}, + Vy::AbstractArray{T,3}, + Vz::AbstractArray{T,3}, + P::AbstractArray{T,3}, + P_old::AbstractArray{T,3}, + τyy_old::AbstractArray{T,3}, + η::AbstractArray{T,3}, + rheology, + phase_ratios, + dt::T, + dx::T, + dy::T, + dz::T, +) where {T} + phase = @inbounds phase_ratios[i, j, end] + Gdt = fn_ratio(get_shear_modulus, rheology, phase) * dt + Vz[i + 1, j + 1, end] = + Vz[i + 1, j + 1, end - 1] + + 3.0 / 2.0 * + ( + P[i, j, end] / (2.0 * η[i, j, end]) - + (τyy_old[i, j, end] + P_old[i, j, end]) / (2.0 * Gdt) + + inv(3.0) * ( + (Vx[i + 1, j + 1, end - 1] - Vx[i, j + 1, end - 1]) * inv(dx) + + (Vy[i + 1, j + 1, end - 1] - Vy[i + 1, j, end - 1]) * inv(dy) + ) + ) * + dz + return nothing +end \ No newline at end of file diff --git a/src/boundaryconditions/no_slip.jl b/src/boundaryconditions/no_slip.jl new file mode 100644 index 00000000..bbae5065 --- /dev/null +++ b/src/boundaryconditions/no_slip.jl @@ -0,0 +1,25 @@ +@parallel_indices (i) function no_slip!(Ax, Ay, bc) + @inbounds begin + if bc.left + (i ≤ size(Ax, 2)) && (Ax[1, i] = 0.0) + (1 < i < size(Ay, 2)) && (Ay[1, i] = -Ay[2, i]) + end + if bc.right + (i ≤ size(Ax, 2)) && (Ax[end, i] = 0.0) + (1 < i < size(Ay, 2)) && (Ay[end, i] = -Ay[end - 1, i]) + end + if bc.bot + (i ≤ size(Ay, 1)) && (Ay[i, 1] = 0.0) + (1 < i < size(Ax, 1)) && (Ax[i, 1] = -Ax[i, 2]) + end + if bc.top + (i ≤ size(Ay, 1)) && (Ay[i, end] = 0.0) + (1 < i < size(Ax, 1)) && (Ax[i, end] = -Ax[i, end - 1]) + end + # corners + # bc.bot && (Ax[1, 1] = 0.0; Ax[1, 1] = 0.0) + # bc.left && bc.bot && (Ax[1, 1] = 0.0) + # bc.right && bc.top && (Ay[end, end] = 0.0) + end + return nothing +end \ No newline at end of file diff --git a/src/boundaryconditions/pure_shear.jl b/src/boundaryconditions/pure_shear.jl new file mode 100644 index 00000000..ebc83714 --- /dev/null +++ b/src/boundaryconditions/pure_shear.jl @@ -0,0 +1,23 @@ +function pureshear_bc!( + stokes::JustRelax.StokesArrays, xci::NTuple{2,T}, xvi::NTuple{2,T}, εbg +) where {T} + _T = typeof(stokes.V.Vx) + stokes.V.Vx[:, 2:(end - 1)] .= _T([εbg * x for x in xvi[1], y in xci[2]]) + stokes.V.Vy[2:(end - 1), :] .= _T([-εbg * y for x in xci[1], y in xvi[2]]) + + return nothing +end + +function pureshear_bc!( + stokes::JustRelax.StokesArrays, xci::NTuple{3,T}, xvi::NTuple{3,T}, εbg +) where {T} + xv, yv, zv = xvi + xc, yc, zc = xci + _T = typeof(stokes.V.Vx) + + stokes.V.Vx[:, 2:(end - 1), 2:(end - 1)] .= _T([εbg * x for x in xv, y in yc, z in zc]) + stokes.V.Vy[2:(end - 1), :, 2:(end - 1)] .= _T([εbg * y for x in xc, y in xv, z in zc]) + stokes.V.Vz[2:(end - 1), 2:(end - 1), :] .= _T([-εbg * z for x in xc, y in xc, z in zv]) + + return nothing +end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 688b7aae..c4c84127 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -569,6 +569,8 @@ function _solve!( ) @hide_communication b_width begin # communication/computation overlap + flow_bcs!(stokes, flow_bcs) + @parallel compute_V!( @velocity(stokes)..., Vx_on_Vy, @@ -582,7 +584,7 @@ function _solve!( ) # apply boundary conditions free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - flow_bcs!(stokes, flow_bcs) + # flow_bcs!(stokes, flow_bcs) update_halo!(@velocity(stokes)...) end end From 4fbf6720b04267b372c5bffc1b2f0e7a0af3e8a9 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 10:33:46 +0200 Subject: [PATCH 76/91] format --- src/boundaryconditions/free_slip.jl | 4 +--- src/boundaryconditions/free_surface.jl | 2 +- src/boundaryconditions/no_slip.jl | 2 +- src/stokes/Stokes2D.jl | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/boundaryconditions/free_slip.jl b/src/boundaryconditions/free_slip.jl index a9348db0..f054ad0a 100644 --- a/src/boundaryconditions/free_slip.jl +++ b/src/boundaryconditions/free_slip.jl @@ -102,8 +102,6 @@ end return nothing end - - @parallel_indices (j) function free_slip_x!(A::AbstractArray{T,2}) where {T} A[1, j] = A[2, j] A[end, j] = A[end - 1, j] @@ -172,4 +170,4 @@ function apply_free_slip!(freeslip::NamedTuple{<:Any,NTuple{3,T}}, Vx, Vy, Vz) w @parallel (1:size(Vx, 1), 1:size(Vx, 2)) free_slip_z!(Vx) @parallel (1:size(Vy, 1), 1:size(Vy, 2)) free_slip_z!(Vy) end -end \ No newline at end of file +end diff --git a/src/boundaryconditions/free_surface.jl b/src/boundaryconditions/free_surface.jl index d130b331..1e4ff780 100644 --- a/src/boundaryconditions/free_surface.jl +++ b/src/boundaryconditions/free_surface.jl @@ -83,4 +83,4 @@ end ) * dz return nothing -end \ No newline at end of file +end diff --git a/src/boundaryconditions/no_slip.jl b/src/boundaryconditions/no_slip.jl index bbae5065..2a677277 100644 --- a/src/boundaryconditions/no_slip.jl +++ b/src/boundaryconditions/no_slip.jl @@ -22,4 +22,4 @@ # bc.right && bc.top && (Ay[end, end] = 0.0) end return nothing -end \ No newline at end of file +end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index c4c84127..a0501f51 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -570,7 +570,7 @@ function _solve!( @hide_communication b_width begin # communication/computation overlap flow_bcs!(stokes, flow_bcs) - + @parallel compute_V!( @velocity(stokes)..., Vx_on_Vy, From 5a9af0eacb4775425bf1efba9b5667e3e0de7ca3 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 11:28:02 +0200 Subject: [PATCH 77/91] reworked public viscosity kernel --- .../Blankenbach2D/Benchmark2D_WENO5.jl | 4 +- .../stokes2D/Blankenbach2D/Benchmark2D_sgd.jl | 4 +- .../Blankenbach2D/Benchmark2D_sgd_scaled.jl | 4 +- .../stokes2D/VanKeken.jl/VanKeken.jl | 2 +- .../Elastic_BuildUp_phases_incompressible.jl | 2 +- .../PlumeFreeSurface_2D.jl | 15 +-- .../RayleighTaylor2D.jl | 23 ++-- .../stokes2D/shear_band/ShearBand2D.jl | 2 +- .../stokes2D/shear_band/ShearBand2D_MPI.jl | 2 +- .../shear_band/ShearBand2D_softening.jl | 2 +- .../stokes2D/shear_heating/Shearheating2D.jl | 4 +- .../stokes2D/sinking_block/SinkingBlock2D.jl | 2 +- .../stokes3D/shear_band/ShearBand3D.jl | 35 +++--- .../stokes3D/shear_heating/Shearheating3D.jl | 103 ++++++++---------- .../Thermal_Stress_Magma_Chamber_nondim.jl | 4 +- .../Thermal_Stress_Magma_Chamber_nondim3D.jl | 4 +- .../convection/GlobalConvection2D_WENO5.jl | 4 +- .../Particles2D/Layered_convection2D.jl | 4 +- .../Layered_convection2D.jl | 4 +- .../Particles3D/Layered_convection3D.jl | 4 +- src/ext/AMDGPU/2D.jl | 16 +-- src/ext/AMDGPU/3D.jl | 18 +-- src/ext/CUDA/2D.jl | 12 +- src/ext/CUDA/3D.jl | 12 +- src/rheology/Viscosity.jl | 47 ++++---- src/stokes/Stokes2D.jl | 18 ++- src/stokes/Stokes3D.jl | 19 +--- test/test_VanKeken.jl | 6 +- test/test_shearband2D.jl | 2 +- test/test_shearband2D_softening.jl | 4 +- test/test_shearheating2D.jl | 10 +- test/test_shearheating3D.jl | 9 +- test/test_sinking_block.jl | 5 +- 33 files changed, 175 insertions(+), 231 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl index abab0486..a1eadfd0 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl @@ -124,7 +124,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal η = @ones(ni...) compute_ρg!(ρg[2], phase_ratios, rheology, args) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # PT coefficients for thermal diffusion ------------- @@ -187,7 +187,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl index bba6a999..83edd811 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl @@ -125,7 +125,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal η = @ones(ni...) compute_ρg!(ρg[2], phase_ratios, rheology, args) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # PT coefficients for thermal diffusion ------------- @@ -194,7 +194,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl index 677ee7f6..5cc1e812 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -115,7 +115,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal η = @ones(ni...) compute_ρg!(ρg[2], phase_ratios, rheology, args) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # PT coefficients for thermal diffusion ------------- @@ -183,7 +183,7 @@ function main2D(igg; ar=1, nx=32, ny=32, nit = 1e1, figdir="figs2D", do_vtk =fal # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ diff --git a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl index 628b7ad1..b1d74002 100644 --- a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl +++ b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl @@ -99,7 +99,7 @@ function main2D(igg; ny=64, nx=64, figdir="model_figs") compute_ρg!(ρg[2], phase_ratios, rheology, args) # Rheology - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl index 8c9df15f..b28722f8 100644 --- a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl +++ b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp_phases_incompressible.jl @@ -70,7 +70,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") args = (; T=@zeros(ni...), P=stokes.P, dt=dt) # Rheology - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 2f8ac79d..7360820c 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -138,7 +138,7 @@ function main(igg, nx, ny) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -155,13 +155,8 @@ function main(igg, nx, ny) # Time loop t, it = 0.0, 0 dt = 1e3 * (3600 * 24 * 365.25) - while it < 15 # run only for 5 Myrs + while it < 15 - # Stokes solver ---------------- - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) - solve!( stokes, pt_stokes, @@ -174,10 +169,10 @@ function main(igg, nx, ny) dt, igg; kwargs = (; - iterMax = 50e3, - nout = 1e3, + iterMax = 50e3, + nout = 1e3, viscosity_cutoff = (-Inf, Inf), - free_surface = true + free_surface = true ) ) dt = compute_dt(stokes, di) / 2 diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index 15fe60b4..b416047e 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -136,13 +136,12 @@ function RT_2D(igg, nx, ny) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right = true, top = true, bot = false), no_slip = (left = false, right = false, top = false, bot = true), - # free_slip = (left = true, right = true, top = true, bot = true), free_surface = true, ) @@ -153,14 +152,10 @@ function RT_2D(igg, nx, ny) take(figdir) # Time loop - t, it = 0.0, 0 - dt = 1e3 * (3600 * 24 * 365.25) - dt_max = 50e3 * (3600 * 24 * 365.25) - while it < 500 # run only for 5 Myrs - - 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)) + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + dt_max = 50e3 * (3600 * 24 * 365.25) + while it < 500 # Stokes solver ---------------- solve!( @@ -175,9 +170,9 @@ function RT_2D(igg, nx, ny) dt, igg; kwargs = ( - iterMax = 50e3, - iterMin = 1e3, - # viscosity_relaxation = 1e-2, + iterMax = 50e3, + iterMin = 1e3, + viscosity_relaxation = 1e-2, nout = 5e3, free_surface = true, viscosity_cutoff = (-Inf, Inf) @@ -251,4 +246,4 @@ else igg end -RT_2D(igg, nx, ny) +# RT_2D(igg, nx, ny) diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 12fb2839..6fd8a182 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -97,7 +97,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Rheology compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions flow_bcs = FlowBoundaryConditions(; diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl index 67381d19..ca2baccf 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl @@ -97,7 +97,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Rheology compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl index 43452a09..b1275fdf 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl @@ -102,7 +102,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") # Rheology compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions flow_bcs = FlowBoundaryConditions(; diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index 80bc5d27..16506b9d 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -109,7 +109,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Rheology args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -173,7 +173,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # ------------------------------ diff --git a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl index 2c04b64c..96ed16dd 100644 --- a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl +++ b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl @@ -130,7 +130,7 @@ function sinking_block2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_per # Viscosity args = (; dt = dt, ΔTc = @zeros(ni...)) η_cutoff = -Inf, Inf - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl index c148538e..cd075758 100644 --- a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl @@ -1,12 +1,12 @@ +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO +import JustRelax.@cell +const backend_JR = JustRelax.CPUBackend + using Printf, GeoParams, GLMakie, CellArrays, CSV, DataFrames -using JustRelax, JustRelax.DataIO + using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) -environment!(model) - # HELPER FUNCTIONS --------------------------------------------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) @@ -29,7 +29,7 @@ function init_phases!(phase_ratios, xci, radius) return nothing end - @parallel (JustRelax.@idx ni) init_phases!(phase_ratios.center, xci..., origin...) + @parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin...) end # MAIN SCRIPT -------------------------------------------------------------------- @@ -52,7 +52,6 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs") 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" η_reg = 1.25e-2 # regularisation "viscosity" dt = η0/G0/4.0 # assumes Maxwell time of 4 el_bg = ConstantElasticity(; G=G0, ν=0.5) @@ -90,17 +89,13 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs") # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-4, CFL = 0.05 / √3.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...), @zeros(ni...) - args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @zeros(ni...), P = stokes.P, dt = Inf) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -132,16 +127,16 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs") di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 150e3, - nout = 1e3, - viscosity_cutoff = (-Inf, Inf) + kwargs = (; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = (-Inf, Inf) + ) ) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @@ -149,7 +144,7 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs") @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε.II, @strain(stokes)...) push!(τII, maximum(stokes.τ.xx)) it += 1 diff --git a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl index c3180b7b..b5c28801 100644 --- a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl @@ -1,7 +1,9 @@ # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell +const backend_JR = JustRelax.CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) @@ -12,10 +14,6 @@ using JustPIC._3D # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 3) #or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) -environment!(model) - # Load script dependencies using Printf, LinearAlgebra, GeoParams, CellArrays using GLMakie @@ -42,15 +40,15 @@ end function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =false) # Physical domain ------------------------------------ - lx = 70e3 # domain length in x - ly = 70e3 # domain length in y - lz = 40e3 # domain length in y - ni = nx, ny, nz # number of cells - li = lx, ly, lz # domain length in x- and y- - di = @. li / (nx_g(),ny_g(),nz_g()) # grid step in x- and -y - origin = 0.0, 0.0, -lz # origin coordinates (15km f sticky air layer) + lx = 70e3 # domain length in x + ly = 70e3 # domain length in y + lz = 40e3 # domain length in y + ni = nx, ny, nz # number of cells + li = lx, ly, lz # domain length in x- and y- + di = @. li / (nx_g(),ny_g(),nz_g()) # grid step in x- and -y + origin = 0.0, 0.0, -lz # origin coordinates (15km f sticky air layer) grid = Geometry(ni, li; origin = origin) - (; xci, xvi) = grid # nodes at the center and vertices of the cells + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # Physical properties using GeoParams ---------------- @@ -75,43 +73,37 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal yc_anomaly = ly/2 # origin of thermal anomaly zc_anomaly = 40e3 # origin of thermal anomaly r_anomaly = 3e3 # radius of perturbation - init_phases!(pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) + init_phases!(backend_JR, pPhases, particles, xc_anomaly, yc_anomaly, zc_anomaly, r_anomaly) phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) + phase_ratios_center(phase_ratios, particles, grid, pPhases) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √3.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) - thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), + thermal = ThermalArrays(backend_JR, ni) + thermal_bc = TemperatureBoundaryConditions(; + no_flux = (left = true , right = true , top = false, bot = false, front = true , back = true), ) # Initialize constant temperature @views thermal.T .= 273.0 + 400 thermal_bcs!(thermal.T, thermal_bc) - - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces ρg = ntuple(_ -> @zeros(ni...), Val(3)) - - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + 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 = dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - η_vep = deepcopy(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -167,11 +159,9 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal t, it = 0.0, 0 while it < 10 # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf) - ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[3], phase_ratios.center, rheology, args) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) + compute_ρg!(ρg[3], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -181,18 +171,18 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 100e3, - nout=1e3, - viscosity_cutoff=(-Inf, Inf) + kwargs = (; + iterMax = 100e3, + nout=1e3, + viscosity_cutoff=(-Inf, Inf) + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -200,12 +190,10 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal particle2grid!(thermal.T, pT, xvi, particles) temperature2center!(thermal) - @parallel (@idx ni) compute_shear_heating!( - thermal.shear_heating, - @tensor_center(stokes.τ), - @tensor_center(stokes.τ_o), - @strain(stokes), - phase_ratios.center, + compute_shear_heating!( + thermal, + stokes, + phase_ratios, rheology, # needs to be a tuple dt, ) @@ -219,36 +207,37 @@ function main3D(igg; ar=8, ny=16, nx=ny*8, nz=ny*8, figdir="figs3D", do_vtk =fal args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs = (; + 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) + advection!(particles, RungeKutta2(), @velocity(stokes), ( grid_vx, grid_vy, grid_vz), dt) # advect particles in memory move_particles!(particles, xvi, particle_args) # interpolate fields from grid vertices to particles grid2particle_flip!(pT, xvi, thermal.T, thermal.Told, particles) # check if we need to inject particles - inject = check_injection(particles) - inject && inject_particles_phase!(particles, pPhases, (pT, ), (thermal.T,), xvi) + 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) + 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) + # checkpointing(figdir, stokes, thermal.T, η, t) if do_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) + velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T), τxy = Array(stokes.τ.xy), diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl index 1bd4b732..a1e1d0f9 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl @@ -290,7 +290,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) η = @ones(ni...) # initialise viscosity - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) η_vep = copy(η) # Buoyancy force @@ -362,7 +362,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) # Update buoyancy and viscosity - args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) # Stokes solver ----------------- solve!( diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl index 4bcb849f..977373a0 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl @@ -290,7 +290,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[3], xci[3], sticky_air) end - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) # Arguments for functions @copy thermal.Told thermal.T @@ -348,7 +348,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals # Update buoyancy and viscosity - args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, cutoff_visc) + compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) # Stokes solver ----------------- solve!( diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index c9ef42e7..3f32bf13 100644 --- a/miniapps/convection/GlobalConvection2D_WENO5.jl +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -179,7 +179,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # Rheology viscosity_cutoff = (1e16, 1e24) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -214,7 +214,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff + stokes, phase_ratios, args, rheology, viscosity_cutoff ) # ------------------------------ diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index 342f1939..de8f6184 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -160,7 +160,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Rheology viscosity_cutoff = (1e16, 1e24) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) (; η, η_vep) = stokes.viscosity # PT coefficients for thermal diffusion @@ -229,7 +229,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff + stokes, phase_ratios, args, rheology, viscosity_cutoff ) # ------------------------------ diff --git a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl index 200192c1..bb596a5e 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl @@ -167,7 +167,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Rheology viscosity_cutoff = nondimensionalize((1e16Pa*s, 1e24Pa*s), CharDim) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -235,7 +235,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff + stokes, phase_ratios, args, rheology, viscosity_cutoff ) # ------------------------------ diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl index ff7c1912..1b793bef 100644 --- a/miniapps/convection/Particles3D/Layered_convection3D.jl +++ b/miniapps/convection/Particles3D/Layered_convection3D.jl @@ -143,7 +143,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) # Rheology viscosity_cutoff = (1e18, 1e24) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -204,7 +204,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, args) compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, viscosity_cutoff + stokes, phase_ratios, args, rheology, viscosity_cutoff ) # ------------------------------ diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index a85433b2..3afbfc6f 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -104,27 +104,27 @@ end # Rheology ## viscosity -function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR2D.compute_viscosity!( ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end -function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) +function JR2D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end -function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) +function compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index eb4cb4f3..2ebd22dd 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -104,27 +104,27 @@ end # Rheology ## viscosity -function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function JR3D.compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end -function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) +function JR3D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end -function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) +function compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index 0dacf8b6..bca19c23 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -104,8 +104,8 @@ end # Rheology ## viscosity -function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR2D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff @@ -116,13 +116,13 @@ function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index eb4cb4f3..14db8455 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -104,8 +104,8 @@ end # Rheology ## viscosity -function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR3D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff @@ -116,13 +116,13 @@ function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) - return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +function compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) + return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff + ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 ) - return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index 5104ca8a..f6ac86ef 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,7 +1,8 @@ ## 2D KERNELS -function compute_viscosity!(stokes, ν, args, rheology, cutoff) - return compute_viscosity!(backend(stokes), ν, args, rheology, cutoff) +function compute_viscosity!(stokes::JustRelax.StokesArrays, args, rheology, cutoff; relaxation = 1e0) + return compute_viscosity!(backend(stokes), stokes, relaxation, args, rheology, cutoff) end + function compute_viscosity!(::CPUBackendTrait, stokes, ν, args, rheology, cutoff) return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end @@ -69,9 +70,9 @@ end return nothing end -function compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +function compute_viscosity!(stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; relaxation = 1e0) return compute_viscosity!( - backend(stokes), stokes, ν, phase_ratios, args, rheology, cutoff + backend(stokes), stokes, relaxation, phase_ratios, args, rheology, cutoff ) end function compute_viscosity!( @@ -225,23 +226,6 @@ 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} @@ -252,9 +236,26 @@ end η += if iszero(ratio[i]) 0.0 else - compute_viscosity_εII(rheology[i].CompositeRheology[1], εII, args) * ratio[i] + 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} +# quote +# Base.@_inline_meta +# η = 0.0 +# Base.@nexprs $N i -> ( +# η += if iszero(ratio[i]) +# 0.0 +# else +# compute_viscosity_εII(rheology[i].CompositeRheology[1], εII, args) * ratio[i] +# end +# ) +# η +# end +# end diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index a0501f51..e78660d9 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -271,6 +271,7 @@ function _solve!( dt, igg::IGG; viscosity_cutoff=(-Inf, Inf), + viscosity_relaxation=1e-2, iterMax=10e3, nout=500, b_width=(4, 4, 0), @@ -326,8 +327,7 @@ function _solve!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... ) - ν = 1e-2 - compute_viscosity!(stokes, ν, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation) compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) @@ -441,9 +441,9 @@ function _solve!( dt, igg::IGG; viscosity_cutoff=(-Inf, Inf), + viscosity_relaxation=1e-2, iterMax=50e3, iterMin=1e2, - viscosity_relaxation=1e-2, free_surface=false, nout=500, b_width=(4, 4, 0), @@ -520,9 +520,7 @@ function _solve!( # stokes.P[1, end] = stokes.P[2, end] # stokes.P[end, end] = stokes.P[end - 1, end] - if rem(iter, 5) == 0 - compute_ρg!(ρg[2], phase_ratios, rheology, args) - end + compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel (@idx ni .+ 1) compute_strain_rate!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... @@ -534,11 +532,11 @@ function _solve!( if do_visc compute_viscosity!( stokes, - viscosity_relaxation, phase_ratios, args, rheology, - viscosity_cutoff, + viscosity_cutoff; + relaxation = viscosity_relaxation, ) end @@ -569,8 +567,6 @@ function _solve!( ) @hide_communication b_width begin # communication/computation overlap - flow_bcs!(stokes, flow_bcs) - @parallel compute_V!( @velocity(stokes)..., Vx_on_Vy, @@ -584,7 +580,7 @@ function _solve!( ) # apply boundary conditions free_surface_bcs!(stokes, flow_bcs, η, rheology, phase_ratios, dt, di) - # flow_bcs!(stokes, flow_bcs) + flow_bcs!(stokes, flow_bcs) update_halo!(@velocity(stokes)...) end end diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 74f2c4bc..d7cb7fc7 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -32,6 +32,7 @@ function _solve!( G, dt, igg::IGG; + viscosity_relaxation=1e-2, iterMax=10e3, nout=500, b_width=(4, 4, 4), @@ -169,6 +170,7 @@ function _solve!( iterMax=10e3, nout=500, b_width=(4, 4, 4), + viscosity_relaxation = 1e-2, verbose=true, kwargs..., ) where {T} @@ -223,19 +225,10 @@ function _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!( - η, - 1.0, - phase_ratios.center, - @strain(stokes)..., - args, - rheology, - viscosity_cutoff, - ) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation) @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), @@ -407,7 +400,7 @@ function _solve!( # Update viscosity compute_viscosity!( - stokes, viscosity_relaxation, phase_ratios, args, rheology, viscosity_cutoff + stokes, phase_ratios, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation ) @parallel (@idx ni) compute_τ_nonlinear!( diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index d8edc087..27356d61 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -107,7 +107,7 @@ function VanKeken2D(ny=32, nx=32) compute_ρg!(ρg[2], phase_ratios, rheology, args) # Rheology - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -133,10 +133,6 @@ function VanKeken2D(ny=32, nx=32) while it < nt - # Update buoyancy - compute_ρg!(ρg[2], phase_ratios, rheology, args) - # ------------------------------ - # Stokes solver ---------------- iters = solve!( stokes, diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl index 2400e824..79aaeac5 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -105,7 +105,7 @@ function ShearBand2D() # Rheology compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions flow_bcs = FlowBoundaryConditions(; diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl index 0b6390a4..a8380598 100644 --- a/test/test_shearband2D_softening.jl +++ b/test/test_shearband2D_softening.jl @@ -109,10 +109,8 @@ function ShearBand2D() args = (; T = @zeros(ni...), P = stokes.P, dt = dt, ΔTc = @zeros(ni...)) # Rheology - η = @ones(ni...) - η_vep = similar(η) # effective visco-elasto-plastic viscosity compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index 7444d7c3..f3bac853 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -116,7 +116,7 @@ function Shearheating2D() # Rheology args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -145,14 +145,6 @@ function Shearheating2D() t, it = 0.0, 0 local iters, thermal while it < 10 - - # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) - compute_viscosity!( - stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf) - ) - # ------------------------------ # Stokes solver ---------------- iters = solve!( diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index 9be61145..eae56d94 100644 --- a/test/test_shearheating3D.jl +++ b/test/test_shearheating3D.jl @@ -110,7 +110,7 @@ function Shearheating3D(nx=16, ny=16, nz=16) # Rheology args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -136,12 +136,7 @@ function Shearheating3D(nx=16, ny=16, nz=16) t, it = 0.0, 0 local iters while it < 5 - # Update buoyancy and viscosity - - args = (; T = thermal.Tc, P = stokes.P, dt = Inf) - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) - compute_ρg!(ρg[end], phase_ratios, rheology, args) - # ------------------------------ - + # Stokes solver ---------------- iters = solve!( stokes, diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index e07e837d..50c6f58e 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -137,9 +137,9 @@ function Sinking_Block2D() # ---------------------------------------------------- # Viscosity - args = (; dt = dt, ΔTc = @zeros(ni...)) + args = (; T = @ones(ni...), P = stokes.P, dt=Inf) η_cutoff = -Inf, Inf - compute_viscosity!(stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions @@ -150,7 +150,6 @@ function Sinking_Block2D() update_halo!(stokes.V.Vx, stokes.V.Vy) # Stokes solver ---------------- - args = (; T = @ones(ni...), P = stokes.P, dt=dt, ΔTc = @zeros(ni...)) iters = solve!( stokes, pt_stokes, From c7b94c651a0f9dcfc9eac1e123645a84c764aa13 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 11:29:24 +0200 Subject: [PATCH 78/91] format --- src/ext/AMDGPU/2D.jl | 10 +++++++--- src/ext/AMDGPU/3D.jl | 10 +++++++--- src/ext/CUDA/2D.jl | 10 +++++++--- src/ext/CUDA/3D.jl | 10 +++++++--- src/rheology/Viscosity.jl | 8 ++++++-- src/stokes/Stokes2D.jl | 6 ++++-- src/stokes/Stokes3D.jl | 18 +++++++++++++++--- 7 files changed, 53 insertions(+), 19 deletions(-) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index 3afbfc6f..bc53f702 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -104,7 +104,9 @@ end # Rheology ## viscosity -function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function JR2D.compute_viscosity!( + ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR2D.compute_viscosity!( @@ -116,11 +118,13 @@ function JR2D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 + ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 ) return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index 2ebd22dd..a0a1d9ba 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -104,7 +104,9 @@ end # Rheology ## viscosity -function JR3D.compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function JR3D.compute_viscosity!( + ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR3D.compute_viscosity!( @@ -116,11 +118,13 @@ function JR3D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 + ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 ) return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index bca19c23..bf0fbe81 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -104,7 +104,9 @@ end # Rheology ## viscosity -function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function JR2D.compute_viscosity!( + ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR2D.compute_viscosity!( @@ -116,11 +118,13 @@ function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 + ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 ) return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 14db8455..52c68d3c 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -104,7 +104,9 @@ end # Rheology ## viscosity -function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function JR3D.compute_viscosity!( + ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function JR3D.compute_viscosity!( @@ -116,11 +118,13 @@ function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!(::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 +) return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) end function compute_viscosity!( - ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation = 1e0 + ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 ) return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) end diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index f6ac86ef..e9effe38 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,5 +1,7 @@ ## 2D KERNELS -function compute_viscosity!(stokes::JustRelax.StokesArrays, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + stokes::JustRelax.StokesArrays, args, rheology, cutoff; relaxation=1e0 +) return compute_viscosity!(backend(stokes), stokes, relaxation, args, rheology, cutoff) end @@ -70,7 +72,9 @@ end return nothing end -function compute_viscosity!(stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; relaxation = 1e0) +function compute_viscosity!( + stokes::JustRelax.StokesArrays, phase_ratios, args, rheology, cutoff; relaxation=1e0 +) return compute_viscosity!( backend(stokes), stokes, relaxation, phase_ratios, args, rheology, cutoff ) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index e78660d9..3a4137dd 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -327,7 +327,9 @@ function _solve!( @strain(stokes)..., stokes.∇V, @velocity(stokes)..., _di... ) - compute_viscosity!(stokes, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation) + compute_viscosity!( + stokes, args, rheology, viscosity_cutoff; relaxation=viscosity_relaxation + ) compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) @@ -536,7 +538,7 @@ function _solve!( args, rheology, viscosity_cutoff; - relaxation = viscosity_relaxation, + relaxation=viscosity_relaxation, ) end diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index d7cb7fc7..3fc23485 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -170,7 +170,7 @@ function _solve!( iterMax=10e3, nout=500, b_width=(4, 4, 4), - viscosity_relaxation = 1e-2, + viscosity_relaxation=1e-2, verbose=true, kwargs..., ) where {T} @@ -228,7 +228,14 @@ function _solve!( # Update buoyancy @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation) + compute_viscosity!( + stokes, + phase_ratios, + args, + rheology, + viscosity_cutoff; + relaxation=viscosity_relaxation, + ) @parallel (@idx ni) compute_τ_nonlinear!( @tensor_center(stokes.τ), @@ -400,7 +407,12 @@ function _solve!( # Update viscosity compute_viscosity!( - stokes, phase_ratios, args, rheology, viscosity_cutoff; relaxation = viscosity_relaxation + stokes, + phase_ratios, + args, + rheology, + viscosity_cutoff; + relaxation=viscosity_relaxation, ) @parallel (@idx ni) compute_τ_nonlinear!( From daa3bc6453f33b6ab6453e779a6f6cbad9c2a417 Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 11:59:09 +0200 Subject: [PATCH 79/91] fix viscosity kernel in ext --- src/ext/AMDGPU/2D.jl | 31 ++++++++++++++++--------------- src/ext/AMDGPU/3D.jl | 33 +++++++++++++++++---------------- src/ext/CUDA/2D.jl | 20 ++++++++++---------- src/ext/CUDA/3D.jl | 20 ++++++++++---------- src/rheology/Viscosity.jl | 1 + 5 files changed, 54 insertions(+), 51 deletions(-) diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl index bc53f702..dcab6325 100644 --- a/src/ext/AMDGPU/2D.jl +++ b/src/ext/AMDGPU/2D.jl @@ -103,31 +103,32 @@ function JR2D.phase_ratios_center( end # Rheology + ## viscosity -function JR2D.compute_viscosity!( - ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function JR2D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function JR2D.compute_viscosity!( ::AMDGPUBackendTrait, 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) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) - return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end @@ -138,10 +139,10 @@ function JR2D.tensor_invariant!(::AMDGPUBackendTrait, A::JustRelax.SymmetricTens end ## Buoyancy forces -function JR2D.compute_ρg!(ρg::CuArray, rheology, args) +function JR2D.compute_ρg!(ρg::RocArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR2D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR2D.compute_ρg!(ρg::RocArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -154,17 +155,17 @@ function temperature2center!(::AMDGPUBackendTrait, thermal::JustRelax.ThermalArr return _temperature2center!(thermal) end -function JR2D.vertex2center!(center::T, vertex::T) where {T<:CuArray} +function JR2D.vertex2center!(center::T, vertex::T) where {T<:RocArray} return vertex2center!(center, vertex) end -function JR2D.center2vertex!(vertex::T, center::T) where {T<:CuArray} +function JR2D.center2vertex!(vertex::T, center::T) where {T<:RocArray} return center2vertex!(vertex, center) end function JR2D.center2vertex!( vertex_yz::T, vertex_xz::T, vertex_xy::T, center_yz::T, center_xz::T, center_xy::T -) where {T<:CuArray} +) where {T<:RocArray} return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) end diff --git a/src/ext/AMDGPU/3D.jl b/src/ext/AMDGPU/3D.jl index a0a1d9ba..472905e1 100644 --- a/src/ext/AMDGPU/3D.jl +++ b/src/ext/AMDGPU/3D.jl @@ -103,31 +103,32 @@ function JR3D.phase_ratios_center( end # Rheology + ## viscosity -function JR3D.compute_viscosity!( - ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function JR3D.compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function JR3D.compute_viscosity!( ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end -function JR3D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) + +function JR2D.compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function compute_viscosity!(::AMDGPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function compute_viscosity!( - ::AMDGPUBackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 + ::AMDGPUBackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) - return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function compute_viscosity!(η, ν, εII::RocArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end @@ -138,10 +139,10 @@ function JR3D.tensor_invariant!(::CUDABackendTrait, A::JustRelax.SymmetricTensor end ## Buoyancy forces -function JR3D.compute_ρg!(ρg::CuArray, rheology, args) +function JR3D.compute_ρg!(ρg::RocArray, rheology, args) return compute_ρg!(ρg, rheology, args) end -function JR3D.compute_ρg!(ρg::CuArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) +function JR3D.compute_ρg!(ρg::RocArray, phase_ratios::JustRelax.PhaseRatio, rheology, args) return compute_ρg!(ρg, phase_ratios, rheology, args) end @@ -154,17 +155,17 @@ function temperature2center!(::CUDABackendTrait, thermal::JustRelax.ThermalArray return _temperature2center!(thermal) end -function JR3D.vertex2center!(center::T, vertex::T) where {T<:CuArray} +function JR3D.vertex2center!(center::T, vertex::T) where {T<:RocArray} return vertex2center!(center, vertex) end -function JR3D.center2vertex!(vertex::T, center::T) where {T<:CuArray} +function JR3D.center2vertex!(vertex::T, center::T) where {T<:RocArray} 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<:CuArray} +) where {T<:RocArray} return center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy) end diff --git a/src/ext/CUDA/2D.jl b/src/ext/CUDA/2D.jl index bf0fbe81..e4701169 100644 --- a/src/ext/CUDA/2D.jl +++ b/src/ext/CUDA/2D.jl @@ -104,30 +104,30 @@ end # Rheology ## viscosity -function JR2D.compute_viscosity!( - ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function JR2D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function JR2D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function JR2D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!( - ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function compute_viscosity!( - ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) - return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end diff --git a/src/ext/CUDA/3D.jl b/src/ext/CUDA/3D.jl index 52c68d3c..73158c3a 100644 --- a/src/ext/CUDA/3D.jl +++ b/src/ext/CUDA/3D.jl @@ -104,30 +104,30 @@ end # Rheology ## viscosity -function JR3D.compute_viscosity!( - ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function JR3D.compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function JR3D.compute_viscosity!( ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function JR3D.compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end -function compute_viscosity!( - ::CUDABackendTrait, stokes, args, rheology, cutoff; relaxation=1e0 -) - return _compute_viscosity!(stokes, relaxation, args, rheology, cutoff) +function compute_viscosity!(::CUDABackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) end + function compute_viscosity!( - ::CUDABackendTrait, stokes, phase_ratios, args, rheology, cutoff; relaxation=1e0 + ::CUDABackendTrait, stokes, ν, phase_ratios, args, rheology, cutoff ) - return _compute_viscosity!(stokes, relaxation, phase_ratios, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) end + function compute_viscosity!(η, ν, εII::CuArray, args, rheology, cutoff) return compute_viscosity!(η, ν, εII, args, rheology, cutoff) end diff --git a/src/rheology/Viscosity.jl b/src/rheology/Viscosity.jl index e9effe38..f051d501 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -79,6 +79,7 @@ function compute_viscosity!( backend(stokes), stokes, relaxation, phase_ratios, args, rheology, cutoff ) end + function compute_viscosity!( ::CPUBackendTrait, stokes::JustRelax.StokesArrays, From fedcf752717dec8679d8432ab577d3124f6a1c6b Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 13:28:25 +0200 Subject: [PATCH 80/91] move elastic stress copy back into the solver --- .../elastic_buildup/Elastic_BuildUp.jl | 5 ----- .../stokes2D/shear_band/ShearBand2D.jl | 7 +------ .../stokes2D/shear_band/ShearBand2D_MPI.jl | 9 -------- .../shear_band/ShearBand2D_softening.jl | 5 ----- src/stokes/Stokes2D.jl | 21 ++++++++++++++++++- src/stokes/Stokes3D.jl | 17 ++++++++++++++- 6 files changed, 37 insertions(+), 27 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl index 162673d5..35e88bba 100644 --- a/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl +++ b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl @@ -91,11 +91,6 @@ function elastic_buildup(; ) ) - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) - ) - t += dt it += 1 println("Iteration $it => t = $(t/kyr) kyrs") diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index 6fd8a182..ba1fc9bd 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -139,12 +139,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") ) 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 diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl index ca2baccf..9f5f399e 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_MPI.jl @@ -158,15 +158,6 @@ function main(igg; nx=64, ny=64, figdir="model_figs") 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 diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl index b1275fdf..26aa0efd 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl @@ -150,11 +150,6 @@ function main(igg; nx=64, ny=64, figdir="model_figs") 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 diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 3a4137dd..0bc2b742 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -133,6 +133,11 @@ function _solve!( 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.τ) + ) + return ( iter=iter, err_evo1=err_evo1, @@ -248,6 +253,11 @@ function _solve!( 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.τ) + ) + return ( iter=iter, err_evo1=err_evo1, @@ -415,6 +425,11 @@ function _solve!( end 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) @@ -635,7 +650,11 @@ function _solve!( end stokes.P .= θ - # @views stokes.P .-= stokes.P[:, end] + + @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) diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 3fc23485..9f40f5d0 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -313,7 +313,14 @@ function _solve!( end av_time = wtime0 / (iter - 1) # average time per iteration - update_τ_o!(stokes) # copy τ into τ_o + + @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) return ( iter=iter, @@ -500,6 +507,14 @@ function _solve!( av_time = wtime0 / (iter - 1) # average time per iteration 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) return ( iter=iter, From ca715a13103ae0453f398ba46177a4b96dfed59c Mon Sep 17 00:00:00 2001 From: albert-de-montserrat Date: Thu, 2 May 2024 13:35:25 +0200 Subject: [PATCH 81/91] update 3D global updwind convection --- .../convection/GlobalConvection3D_Upwind.jl | 79 ++++++++----------- 1 file changed, 35 insertions(+), 44 deletions(-) diff --git a/miniapps/convection/GlobalConvection3D_Upwind.jl b/miniapps/convection/GlobalConvection3D_Upwind.jl index 3c1cc602..2adef7d4 100644 --- a/miniapps/convection/GlobalConvection3D_Upwind.jl +++ b/miniapps/convection/GlobalConvection3D_Upwind.jl @@ -1,11 +1,11 @@ -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO +import JustRelax.@cell + +const backend_JR = CPUBackend + using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 3) #or (:CUDA, Float64, 3) or (:AMDGPU, Float64, 3) -environment!(model) - using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions # function to compute strain rate (compulsory) @@ -24,7 +24,6 @@ end @inline function custom_viscosity(a::CustomRheology; P=0.0, T=273.0, depth=0.0, kwargs...) (; η0, Ea, Va, T0, R, cutoff) = a.args η = η0 * exp((Ea + P * Va) / (R * T) - Ea / (R * T0)) - # correction = (depth ≤ 660e3) + (2740e3 ≥ depth > 660e3) * 1e1 + (depth > 2740e3) * 1e-2 correction = (depth ≤ 660e3) + (2740e3 ≥ depth > 660e3) * 1e1 + (depth > 2700e3) * 1e-1 η = clamp(η * correction, cutoff...) end @@ -90,18 +89,19 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th igg = IGG(init_global_grid(nx, ny, nz; init_MPI = JustRelax.MPI.Initialized() ? false : true)...) # Physical domain ------------------------------------ - lz = 2890e3 - lx = ly = lz * ar - origin = 0.0, 0.0, -lz # origin coordinates - ni = nx, ny, nz # number of cells - li = lx, ly, lz # domain length in x- and y- - di = @. li / (nx_g(), ny_g(), nz_g()) # grid step in x- and -y - xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + lz = 2890e3 + lx = ly = lz * ar + origin = 0.0, 0.0, -lz # origin coordinates + ni = nx, ny, nz # number of cells + li = lx, ly, lz # domain length in x- and y- + di = @. li / (nx_g(), ny_g(), nz_g()) # grid step in x- and -y + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # create rheology struct v_args = (; η0=5e20, Ea=200e3, Va=2.6e-6, T0=1.6e3, R=8.3145, cutoff=(1e16, 1e25)) - creep = CustomRheology(custom_εII, custom_τII, v_args) + creep = CustomRheology(custom_εII, custom_τII, v_args) # Physical properties using GeoParams ---------------- η_reg = 1e18 @@ -139,9 +139,9 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend, ni) thermal_bc = TemperatureBoundaryConditions(; - no_flux = (left = true, right = true, top = false, bot = false, front=true, back=true), + no_flux = (left = true, right = true, top = false, bot = false, front=true, back=true), ) # initialize thermal profile - Half space cooling adiabat = 0.3 # adiabatic gradient @@ -166,34 +166,30 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th @views thermal.T[:, :, 1] .= Tmax @views thermal.T[:, :, end] .= Tmin - @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 1.0 / √3.1) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...), @zeros(ni...) for _ in 1:2 - @parallel (@idx ni) compute_ρg!(ρg[3], rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[3], rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[3], xci[3]) end # Rheology - η = @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 = dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, (1e18, 1e24) - ) - η_vep = deepcopy(η) + args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf) + compute_viscosity!(stokes, args, rheology, (1e18, 1e24)) + # 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), ) flow_bcs!(stokes, flow_bcs) # apply boundary conditions - update_halo!(stokes.V.Vx, stokes.V.Vy, stokes.V.Vz) + update_halo!(@velocity(stokes)...) # ---------------------------------------------------- # IO ------------------------------------------------- @@ -226,7 +222,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th # Update arguments needed to compute several physical properties # e.g. density, viscosity, etc - - args = (; T=thermal.Tc, P=stokes.P, depth=depth, dt=dt, ΔTc = thermal.ΔTc) + args = (; T=thermal.Tc, P=stokes.P, depth=depth, dt=Inf) # Stokes solver ---------------- solve!( @@ -235,14 +231,14 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th di, flow_bcs, ρg, - η, - η_vep, rheology, args, Inf, igg; - iterMax=150e3, - nout=2e3, + kwargs = (; + iterMax=150e3, + nout=2e3, + ) ); println("starting non linear iterations") @@ -303,15 +299,10 @@ 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() - figdir = "figs3D_test" - ar = 3 # aspect ratio - n = 32 - nx = n*ar - 2 - ny = nx - nz = n - 2 - -# thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz); -# end - -# run() +figdir = "figs3D_test" +ar = 3 # aspect ratio +n = 32 +nx = n*ar - 2 +ny = nx +nz = n - 2 +thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz); From 898685b2d53be31e87464cbe004d825e54e28570 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:51:13 +0200 Subject: [PATCH 82/91] rm JR/JR from backend calls --- docs/src/man/Blankenbach.md | 42 +++++++++---------- docs/src/man/ShearBands.md | 2 +- .../Blankenbach2D/Benchmark2D_WENO5.jl | 4 +- .../stokes2D/Blankenbach2D/Benchmark2D_sgd.jl | 4 +- .../stokes2D/VanKeken.jl/VanKeken.jl | 2 +- .../PlumeFreeSurface_2D.jl | 6 +-- .../RayleighTaylor2D.jl | 30 ++++++------- .../stokes2D/shear_heating/Shearheating2D.jl | 4 +- .../stokes2D/sinking_block/SinkingBlock2D.jl | 2 +- .../stokes3D/shear_band/ShearBand3D.jl | 2 +- .../stokes3D/shear_heating/Shearheating3D.jl | 2 +- .../diffusion/diffusion2D.jl | 2 +- .../diffusion/diffusion3D_multiphase.jl | 2 +- .../Thermal_Stress_Magma_Chamber_nondim.jl | 16 +++---- .../Thermal_Stress_Magma_Chamber_nondim3D.jl | 6 +-- .../Particles2D/Layered_convection2D.jl | 6 +-- .../Layered_convection2D.jl | 24 +++++------ .../Particles3D/Layered_convection3D.jl | 6 +-- test/test_VanKeken.jl | 4 +- test/test_diffusion2D_multiphase.jl | 6 +-- test/test_diffusion3D_multiphase.jl | 2 +- test/test_shearheating2D.jl | 6 +-- test/test_shearheating3D.jl | 8 ++-- test/test_sinking_block.jl | 2 +- 24 files changed, 95 insertions(+), 95 deletions(-) diff --git a/docs/src/man/Blankenbach.md b/docs/src/man/Blankenbach.md index dbb6abe5..eb358fd3 100644 --- a/docs/src/man/Blankenbach.md +++ b/docs/src/man/Blankenbach.md @@ -7,13 +7,13 @@ Thermal convection benchmark from [Blankenbach et al., 1989](https://academic.o Load JustRelax necessary modules and define backend. ```julia using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend ``` 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend ``` We will also use `ParallelStencil.jl` to write some device-agnostic helper functions: @@ -57,14 +57,14 @@ rheology = ( CompositeRheology = CompositeRheology((LinearViscous(; η = 1),)), RadioactiveHeat = ConstantRadioactiveHeat(0.0), Gravity = ConstantGravity(; g = 1e4), - ), + ), ) ``` ## Initialize particles ```julia nxcell = 24 # initial number of perticles per cell -max_xcell = 35 # maximum number of perticles per cell +max_xcell = 35 # maximum number of perticles per cell min_xcell = 12 # minimum number of perticles per cell particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... @@ -77,7 +77,7 @@ grid_vx, grid_vy = velocity_grids(xci, xvi, di) # staggered velocity grids and we want to keep track of the temperature `pT`, temperature of the previous time step `pT0`, and material phase `pPhase`: ```julia -pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) +pT, pT0, pPhases = init_cell_arrays(particles, Val(3)) particle_args = (pT, pT0, pPhases) ``` @@ -113,7 +113,7 @@ or we can use the alternative one-liners ```julia @views pPhase.data[!isnan.(particles.index.data)] .= 1.0 ``` -or +or ```julia map!(x -> isnan(x) ? NaN : 1.0, pPhase.data, particles.index.data) ``` @@ -127,7 +127,7 @@ phase_ratios_center(phase_ratios, particles, grid, pPhases) ## Stokes and heat diffusion arrays Stokes arrays object -```julia +```julia stokes = StokesArrays(backend_JR, ni) ``` @@ -140,7 +140,7 @@ thermal = ThermalArrays(backend_JR, ni) To initialize the thermal profile we use `ParallelStencil.jl` again ```julia -@parallel_indices (i, j) function init_T!(T, y) +@parallel_indices (i, j) function init_T!(T, y) T[i, j] = 1 - y[j] return nothing end @@ -153,7 +153,7 @@ and we define a rectangular thermal anomaly at $x \in [0, 0.05]$, $y \in [\frac{ ```julia function rectangular_perturbation!(T, xc, yc, r, xvi) @parallel_indices (i, j) function _rectangular_perturbation!(T, xc, yc, r, x, y) - @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) + @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc)^2 ≤ r^2) T[i, j] += .2 end return nothing @@ -173,14 +173,14 @@ We initialize the buoyancy forces and viscosity ```julia ρg = @zeros(ni...), @zeros(ni...) η = @ones(ni...) -args = (; T = thermal.Tc, P = stokes.P, dt = Inf) +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 +```julia flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right=true, top=true, bot=true), ) @@ -213,7 +213,7 @@ end grid2particle!(pT, xvi, T_buffer, particles) pT0.data .= pT.data ``` -where +where ```julia function copyinn_x!(A, B) @parallel function f_x(A, B) @@ -270,8 +270,8 @@ solve!( ) # calculate adaptive time step dt = compute_dt(stokes, di, dt_diff) -``` -2. Heat diffusion solver +``` +2. Heat diffusion solver ```julia heatdiffusion_PT!( thermal, @@ -319,7 +319,7 @@ phase_ratios_center(phase_ratios, particles, grid, pPhases) ```julia # interpolate fields from particle to grid vertices particle2grid!(T_buffer, pT, xvi, particles) -@views T_buffer[:, end] .= 0.0 +@views T_buffer[:, end] .= 0.0 @views T_buffer[:, 1] .= 1.0 @views thermal.T[2:end-1, :] .= T_buffer flow_bcs!(stokes, flow_bcs) # apply boundary conditions @@ -334,10 +334,10 @@ compute_ρg!(ρg[2], phase_ratios, rheology, args) ``` 7. Compute Nusselt number and rms-velocity ```julia -# Nusselt number, Nu = ∫ ∂T/∂z dx -Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) +# Nusselt number, Nu = ∫ ∂T/∂z dx +Nu_it = sum( ((abs.(thermal.T[2:end-1,end] - thermal.T[2:end-1,end-1])) ./ di[2]) .*di[1]) push!(Nu_top, Nu_it) -# Compute U rms +# Compute U rms # U₍ᵣₘₛ₎ = √ ∫∫ (vx²+vz²) dx dz Urms_it = let JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, stokes.V.Vx, stokes.V.Vy; ghost_nodes=true) @@ -350,7 +350,7 @@ push!(trms, t) # Visualization We will use `Makie.jl` to visualize the results -```julia +```julia using GLMakie ``` @@ -377,7 +377,7 @@ h2 = heatmap!(ax2, xvi[1], xvi[2], Array(stokes.V.Vy) , colormap=:batlow) # x-velocity h3 = heatmap!(ax3, xvi[1], xvi[2], Array(stokes.V.Vx) , colormap=:batlow) # particles temperature -h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) +h4 = scatter!(ax4, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), colormap=:lajolla, colorrange=(0, 1), markersize=3) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) diff --git a/docs/src/man/ShearBands.md b/docs/src/man/ShearBands.md index c16cce35..33273296 100644 --- a/docs/src/man/ShearBands.md +++ b/docs/src/man/ShearBands.md @@ -7,7 +7,7 @@ Shear Band benchmark to test the visco-elasto-plastic rheology implementation in Load JustRelax necessary modules and define backend. ```julia using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend ``` diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl index a1eadfd0..3bf7c5a9 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -9,7 +9,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl index 83edd811..7f3868d8 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -9,7 +9,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, LinearAlgebra, GeoParams, CairoMakie, CellArrays diff --git a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl index b1d74002..e7a0dc0b 100644 --- a/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl +++ b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # x-length of the domain const λ = 0.9142 diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl index 7360820c..bfe0d97f 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -2,7 +2,7 @@ using JustRelax, JustRelax.JustRelax2D const backend_JR = CPUBackend using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend +const backend = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) @@ -43,7 +43,7 @@ function init_phases!(phases, particles) JustRelax.@cell(index[ip, i, j]) == 0 && continue x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) + depth = -(JustRelax.@cell py[ip, i, j]) JustRelax.@cell phases[ip, i, j] = 2.0 if 0e0 ≤ depth ≤ 100e3 @@ -187,7 +187,7 @@ function main(igg, nx, ny) inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios phase_ratios_center(phase_ratios, particles, grid, pPhases) - + @show it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl index b416047e..f864a07f 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -3,7 +3,7 @@ using JustRelax, JustRelax.JustRelax2D const backend_JR = CPUBackend using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend +const backend = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) @@ -33,19 +33,19 @@ end function init_phases!(phases, particles, A) ni = size(phases) - + @parallel_indices (i, j) function init_phases!(phases, px, py, index, A) - + f(x, A, λ) = A * sin(π * x / λ) - + @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) + depth = -(JustRelax.@cell py[ip, i, j]) JustRelax.@cell phases[ip, i, j] = 2.0 - + if 0e0 ≤ depth ≤ 100e3 JustRelax.@cell phases[ip, i, j] = 1.0 @@ -114,7 +114,7 @@ function RT_2D(igg, nx, ny) pT, pPhases = init_cell_arrays(particles, Val(2)) particle_args = (pT, pPhases) - # Elliptical temperature anomaly + # Elliptical temperature anomaly A = 5e3 # Amplitude of the anomaly phase_ratios = PhaseRatio(backend_JR, ni, length(rheology)) init_phases!(pPhases, particles, A) @@ -130,7 +130,7 @@ function RT_2D(igg, nx, ny) # TEMPERATURE PROFILE -------------------------------- thermal = ThermalArrays(backend_JR, ni) # ---------------------------------------------------- - + # Buoyancy forces & rheology ρg = @zeros(ni...), @zeros(ni...) args = (; T = thermal.Tc, P = stokes.P, dt = Inf) @@ -139,7 +139,7 @@ function RT_2D(igg, nx, ny) compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions - flow_bcs = FlowBoundaryConditions(; + flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right = true, top = true, bot = false), no_slip = (left = false, right = false, top = false, bot = true), free_surface = true, @@ -193,7 +193,7 @@ function RT_2D(igg, nx, ny) # advect particles in space advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # advect particles in memory - move_particles!(particles, xvi, particle_args) + move_particles!(particles, xvi, particle_args) # check if we need to inject particles inject_particles_phase!(particles, pPhases, (), (), xvi) # update phase ratios @@ -215,15 +215,15 @@ function RT_2D(igg, nx, ny) fig = Figure(size = (900, 900), title = "t = $t") ax = Axis(fig[1,1], aspect = 1, title = " t=$(t/(1e3 * 3600 * 24 *365.25)) Kyrs") scatter!( - ax, - pxv, pyv, - color=clr, + ax, + pxv, pyv, + color=clr, colormap = :lajolla, markersize = 3 ) arrows!( ax, - xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., + xvi[1][1:nt:end-1]./1e3, xvi[2][1:nt:end-1]./1e3, Array.((Vx_v[1:nt:end-1, 1:nt:end-1], Vy_v[1:nt:end-1, 1:nt:end-1]))..., lengthscale = 25 / max(maximum(Vx_v), maximum(Vy_v)), color = :darkblue, ) @@ -232,7 +232,7 @@ function RT_2D(igg, nx, ny) end end - return + return end ## END OF MAIN SCRIPT ---------------------------------------------------------------- diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index 16506b9d..fd83a6b4 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using GeoParams @@ -332,4 +332,4 @@ else igg end -main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) \ No newline at end of file +main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk) diff --git a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl index 96ed16dd..a9799b9b 100644 --- a/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl +++ b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl @@ -5,7 +5,7 @@ using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using GeoParams diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl index cd075758..21777c0b 100644 --- a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl @@ -1,6 +1,6 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using Printf, GeoParams, GLMakie, CellArrays, CSV, DataFrames diff --git a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl index b5c28801..e8f7a0be 100644 --- a/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl +++ b/miniapps/benchmarks/stokes3D/shear_heating/Shearheating3D.jl @@ -2,7 +2,7 @@ # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 3) or (AMDGPU, Float64, 3) diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl index 9fd1a11a..caeacec1 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend @parallel_indices (i, j) function init_T!(T, z) if z[j] == maximum(z) diff --git a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl index 2c4ba5b6..4f5387e8 100644 --- a/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl +++ b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion3D_multiphase.jl @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using GeoParams diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl index a1e1d0f9..94132853 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl @@ -1,6 +1,6 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using Printf, Statistics, LinearAlgebra, GeoParams, GLMakie, CellArrays using StaticArrays @@ -363,7 +363,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) - + # Stokes solver ----------------- solve!( stokes, @@ -382,7 +382,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) nout = 5e3, viscosity_cutoff = cutoff_visc, ) - ) + ) tensor_invariant!(stokes.ε) @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) @@ -416,7 +416,7 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) nout = 1e3, verbose = true, ) - + ) for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told)) copyinn_x!(dst, src) @@ -684,9 +684,9 @@ function main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=false) end end end - + finalize_global_grid() - + return nothing end @@ -702,4 +702,4 @@ else end # run main script -main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=do_vtk); \ No newline at end of file +main2D(igg; figdir=figdir, nx=nx, ny=ny, do_vtk=do_vtk); diff --git a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl index 977373a0..3d38c660 100644 --- a/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl +++ b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim3D.jl @@ -1,6 +1,6 @@ using JustRelax, JustRelax.JustRelax3D, JustRelax.DataIO import JustRelax.@cell -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences3D @init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using Printf, Statistics, LinearAlgebra, GeoParams, GLMakie, CellArrays using StaticArrays @@ -349,7 +349,7 @@ function main3D(igg; figdir = "output", nx = 64, ny = 64, nz = 64, do_vtk = fals args = (; T=thermal.Tc, P=stokes.P, dt=Inf, ΔTc=thermal.ΔTc) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!(stokes, phase_ratios, args, rheology, cutoff_visc) - + # Stokes solver ----------------- solve!( stokes, diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index de8f6184..dd42d543 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -3,14 +3,14 @@ import JustRelax.@cell const backend_JR = CPUBackend -using ParallelStencil, ParallelStencil.FiniteDifferences2D +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using GeoParams, GLMakie @@ -157,7 +157,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end - + # Rheology viscosity_cutoff = (1e16, 1e24) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) diff --git a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl index bb596a5e..560fd1ea 100644 --- a/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D_nonDim/Layered_convection2D.jl @@ -1,7 +1,7 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -49,20 +49,20 @@ end if depth < nondimensionalize(0e0km, CharDim) T[i + 1, j] = nondimensionalize(273e0K, CharDim) - elseif nondimensionalize(0e0km, CharDim) ≤ (depth) < nondimensionalize(35km, CharDim) + elseif nondimensionalize(0e0km, CharDim) ≤ (depth) < nondimensionalize(35km, CharDim) dTdZ = nondimensionalize((923-273)/35 * K/km, CharDim) - offset = nondimensionalize(273e0K, CharDim) + offset = nondimensionalize(273e0K, CharDim) T[i + 1, j] = (depth) * dTdZ + offset elseif nondimensionalize(110km, CharDim) > (depth) ≥ nondimensionalize(35km, CharDim) dTdZ = nondimensionalize((1492-923)/75 * K/km, CharDim) - offset = nondimensionalize(923K, CharDim) + offset = nondimensionalize(923K, CharDim) T[i + 1, j] = (depth - nondimensionalize(35km, CharDim)) * dTdZ + offset - elseif (depth) ≥ nondimensionalize(110km, CharDim) + elseif (depth) ≥ nondimensionalize(110km, CharDim) dTdZ = nondimensionalize((1837 - 1492)/590 * K/km, CharDim) - offset = nondimensionalize(1492e0K, CharDim) + offset = nondimensionalize(1492e0K, CharDim) T[i + 1, j] = (depth - nondimensionalize(110km, CharDim)) * dTdZ + offset end @@ -77,7 +77,7 @@ function rectangular_perturbation!(T, xc, yc, r, xvi, thick_air, CharDim) @inbounds if ((x[i]-xc)^2 ≤ r^2) && ((y[j] - yc - thick_air)^2 ≤ r^2) depth = -y[j] - thick_air dTdZ = nondimensionalize((2047 - 2017)K / 50km, CharDim) - offset = nondimensionalize(2017e0K, CharDim) + offset = nondimensionalize(2017e0K, CharDim) T[i + 1, j] = (depth - nondimensionalize(585km, CharDim)) * dTdZ + offset end return nothing @@ -96,7 +96,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) thickness = 700 * km η0 = 1e20 - CharDim = GEO_units(; + CharDim = GEO_units(; length = thickness, viscosity = η0, temperature = 1e3K ) # Physical domain ------------------------------------ @@ -157,14 +157,14 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi, thick_air, CharDim) temperature2center!(thermal) # ---------------------------------------------------- - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) for _ in 1:1 compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end - + # Rheology viscosity_cutoff = nondimensionalize((1e16Pa*s, 1e24Pa*s), CharDim) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) @@ -180,7 +180,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) ) flow_bcs!(stokes, flow_bcs) # apply boundary conditions update_halo!(stokes.V.Vx, stokes.V.Vy) - + # IO ------------------------------------------------ # if it does not exist, make folder where figures are stored if do_vtk @@ -230,7 +230,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) @views thermal.T[:, 1] .= Tbot thermal_bcs!(thermal.T, thermal_bc) temperature2center!(thermal) - + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl index 1b793bef..b3878993 100644 --- a/miniapps/convection/Particles3D/Layered_convection3D.jl +++ b/miniapps/convection/Particles3D/Layered_convection3D.jl @@ -10,7 +10,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, GeoParams, GLMakie, GeoParams @@ -140,7 +140,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) ρg = ntuple(_ -> @zeros(ni...), Val(3)) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[end], xci[end]) - + # Rheology viscosity_cutoff = (1e18, 1e24) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) @@ -265,7 +265,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false) inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios phase_ratios_center(phase_ratios, particles, grid, pPhases) - + @show it += 1 t += dt diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl index 27356d61..48d1daaa 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -13,7 +13,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # x-length of the domain const λ = 0.9142 @@ -184,6 +184,6 @@ end @suppress begin iters, Urms = VanKeken2D() @test passed = iters.err_evo1[end] < 1e-4 - @test all(<(1e-2), Urms) + @test all(<(1e-2), Urms) end end diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl index 03306e8a..e02ca525 100644 --- a/test/test_diffusion2D_multiphase.jl +++ b/test/test_diffusion2D_multiphase.jl @@ -13,7 +13,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend import JustRelax.@cell @@ -182,9 +182,9 @@ end nx=32; ny=32; thermal = diffusion_2D(; nx = nx, ny = ny) - + nx_T, ny_T = size(thermal.T) @test thermal.T[nx_T >>> 1 + 1, ny_T >>> 1 + 1] ≈ 1819.2297931741878 atol=1e-1 @test thermal.Tc[ nx >>> 1 , nx >>> 1 ] ≈ 1824.3532934301472 atol=1e-1 end -end \ No newline at end of file +end diff --git a/test/test_diffusion3D_multiphase.jl b/test/test_diffusion3D_multiphase.jl index c31f9ebd..6d68a9e2 100644 --- a/test/test_diffusion3D_multiphase.jl +++ b/test/test_diffusion3D_multiphase.jl @@ -14,7 +14,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend import JustRelax.@cell diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl index f3bac853..23516ba2 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -3,7 +3,7 @@ using Test, Suppressor # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax2D -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -12,7 +12,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies @@ -215,7 +215,7 @@ function Shearheating2D() inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios phase_ratios_center(phase_ratios, particles, grid, pPhases) - + @show it += 1 t += dt diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index eae56d94..c6ed8764 100644 --- a/test/test_shearheating3D.jl +++ b/test/test_shearheating3D.jl @@ -4,7 +4,7 @@ using Test, Suppressor # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 using JustRelax, JustRelax.JustRelax3D -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend import JustRelax.@cell using ParallelStencil, ParallelStencil.FiniteDifferences3D @@ -14,7 +14,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies @@ -136,7 +136,7 @@ function Shearheating3D(nx=16, ny=16, nz=16) t, it = 0.0, 0 local iters while it < 5 - + # Stokes solver ---------------- iters = solve!( stokes, @@ -149,7 +149,7 @@ function Shearheating3D(nx=16, ny=16, nz=16) args, Inf, igg; - kwargs = ( + kwargs = ( iterMax = 100e3, nout=1e3, viscosity_cutoff=(-Inf, Inf), diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl index 50c6f58e..5f44969c 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -8,7 +8,7 @@ using ParallelStencil, ParallelStencil.FiniteDifferences2D const backend_JR = CPUBackend using JustPIC, JustPIC._2D -const backend = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend using GeoParams From 6fa522dedcb9167aaa5c4a34ab408a1789346775 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:51:28 +0200 Subject: [PATCH 83/91] fix plotting --- miniapps/convection/GlobalConvection3D_Upwind.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miniapps/convection/GlobalConvection3D_Upwind.jl b/miniapps/convection/GlobalConvection3D_Upwind.jl index 2adef7d4..6c3d6838 100644 --- a/miniapps/convection/GlobalConvection3D_Upwind.jl +++ b/miniapps/convection/GlobalConvection3D_Upwind.jl @@ -207,7 +207,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th 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[:]), Zv./1e3) - scatter!(ax2, Array(log10.(η[:])), Z./1e3 ) + scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Z./1e3 ) ylims!(ax1, minimum(xvi[3])./1e3, 0) ylims!(ax2, minimum(xvi[3])./1e3, 0) hideydecorations!(ax2) @@ -277,7 +277,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[3].*1e-3, Array(thermal.T[:,slice_j,:]) , colormap=:batlow) h2 = heatmap!(ax2, xci[1].*1e-3, xvi[3].*1e-3, Array(stokes.V.Vz[:,slice_j,:]) , colormap=:batlow) h3 = heatmap!(ax3, xci[1].*1e-3, xci[3].*1e-3, Array(stokes.τ.II[:,slice_j,:].*1e-6) , colormap=:batlow) - h4 = heatmap!(ax4, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(η_vep[:,slice_j,:])) , colormap=:batlow) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(stokes.viscosity.η_vep[:,slice_j,:])) , colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) From 6553a41d8c54f9410454dee67ee656e48c8ec918 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:52:56 +0200 Subject: [PATCH 84/91] bring Globalconv2D_upwind up to speed --- .../convection/GlobalConvection2D_Upwind.jl | 60 +++++++++---------- 1 file changed, 29 insertions(+), 31 deletions(-) diff --git a/miniapps/convection/GlobalConvection2D_Upwind.jl b/miniapps/convection/GlobalConvection2D_Upwind.jl index 377adbe0..0ffe02f3 100644 --- a/miniapps/convection/GlobalConvection2D_Upwind.jl +++ b/miniapps/convection/GlobalConvection2D_Upwind.jl @@ -1,10 +1,8 @@ -using JustRelax -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = CPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:threads, Float64, 2) -environment!(model) +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions @@ -84,7 +82,7 @@ end function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_perturbation = :circular) # initialize MPI - igg = IGG(init_global_grid(nx, ny, 0; init_MPI = JustRelax.MPI.Initialized() ? false : true)...) + igg = IGG(init_global_grid(nx, ny, 1; init_MPI = JustRelax.MPI.Initialized() ? false : true)...) # Physical domain ------------------------------------ ly = 2890e3 @@ -136,7 +134,7 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -160,29 +158,28 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p end @views thermal.T[:, 1] .= Tmax @views thermal.T[:, end] .= Tmin - @parallel (@idx ni) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) - pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) + stokes = StokesArrays(backend_JR, ni) + pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-7, CFL = 0.9 / √2.1) + + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) for _ in 1:2 - @parallel (@idx ni) compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) end # Rheology - η = @ones(ni...) - depth = PTArray([y for x in xci[1], y in xci[2]]) + depth = PTArray(backend_JR)([y for x in xci[1], y in xci[2]]) args = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = dt, ΔTc = thermal.ΔTc) viscosity_cutoff = 1e18, 1e23 - @parallel (@idx ni) compute_viscosity!( - η, 1.0, @strain(stokes)..., args, rheology, viscosity_cutoff - ) - η_vep = deepcopy(η) + compute_viscosity!(stokes, args, rheology, viscosity_cutoff) + # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right=true, top=true, bot=true), @@ -193,7 +190,7 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p # IO ------------------------------------------------- # if it does not exist, make folder where figures are stored - !isdir(figdir) && mkpath(figdir) + take(figdir) # ---------------------------------------------------- # Plot initial T and η profiles @@ -204,7 +201,7 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p ax1 = Axis(fig[1,1], aspect = 2/3, title = "T") ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)") lines!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3) - lines!(ax2, Array(log10.(η[:])), Y./1e3) + lines!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) ylims!(ax1, -2890, 0) ylims!(ax2, -2890, 0) hideydecorations!(ax2) @@ -224,15 +221,16 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p di, flow_bcs, ρg, - η, - η_vep, rheology, args, dt, igg; - iterMax=50e3, - nout=1e3, - viscosity_cutoff = viscosity_cutoff + kwargs = ( + viscosity_cutoff = viscosity_cutoff, + iterMax = 10e3, + nout = 1e2, + verbose = true + ), ); dt = compute_dt(stokes, di, dt_diff, igg) # ------------------------------ @@ -261,14 +259,14 @@ function thermal_convection2D(; ar=8, ny=16, nx=ny*8, figdir="figs2D", thermal_p # Plotting --------------------- if it == 1 || rem(it, 10) == 0 fig = Figure(size = (1000, 1000), title = "t = $t") - ax1 = Axis(fig[1,1], aspect = ar, title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") - ax2 = Axis(fig[2,1], aspect = ar, title = "Vy [m/s]") - ax3 = Axis(fig[3,1], aspect = ar, title = "τII [MPa]") - ax4 = Axis(fig[4,1], aspect = ar, title = "log10(η)") + ax1 = Axis(fig[1,1], aspect = DataAspect(), title = "T [K] (t=$(t/(1e6 * 3600 * 24 *365.25)) Myrs)") + ax2 = Axis(fig[2,1], aspect = DataAspect(), title = "Vy [m/s]") + ax3 = Axis(fig[3,1], aspect = DataAspect(), title = "τII [MPa]") + ax4 = Axis(fig[4,1], aspect = DataAspect(), title = "log10(η)") h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(thermal.T) , colormap=:batlow) h2 = heatmap!(ax2, xci[1].*1e-3, xvi[2].*1e-3, Array(stokes.V.Vy[2:end-1,:]) , colormap=:batlow) h3 = heatmap!(ax3, xci[1].*1e-3, xci[2].*1e-3, Array(stokes.τ.II.*1e-6) , colormap=:batlow) - h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(η_vep)) , colormap=:batlow) + h4 = heatmap!(ax4, xci[1].*1e-3, xci[2].*1e-3, Array(log10.(stokes.viscosity.η_vep)) , colormap=:batlow) hidexdecorations!(ax1) hidexdecorations!(ax2) hidexdecorations!(ax3) @@ -288,7 +286,7 @@ end function run() figdir = "figs2D_test" - ar = 8 # aspect ratio + ar = 2 # aspect ratio n = 128 nx = n*ar - 2 ny = n - 2 From c6985871ce7adc35f5b98cb782bb5b9b2643501a Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:53:17 +0200 Subject: [PATCH 85/91] rm JR from backend --- miniapps/convection/GlobalConvection2D_WENO5.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index 3f32bf13..9ff9c5d1 100644 --- a/miniapps/convection/GlobalConvection2D_WENO5.jl +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -1,7 +1,7 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -const backend = JustRelax.CPUBackend +const backend = CPUBackend using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -168,7 +168,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # Allocate arrays needed for every Stokes problem stokes = StokesArrays(backend, ni, ViscoElastic) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.8 / √2.1) - + # Buoyancy forces args = (; T = thermal.Tc, P = stokes.P, dt = Inf) ρg = @zeros(ni...), @zeros(ni...) @@ -176,7 +176,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma compute_ρg!(ρg[2], phase_ratios, rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end - + # Rheology viscosity_cutoff = (1e16, 1e24) compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) @@ -209,7 +209,7 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma local iters igg.me == 0 && println("Starting model") while (t / (1e6 * 3600 * 24 * 365.25)) < 4.5e3 - + # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) From 0331b80b53cbeb10d1031c6011114e9fa92bb275 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:55:39 +0200 Subject: [PATCH 86/91] rm phase_ratios from WENO5 --- miniapps/convection/GlobalConvection2D_WENO5.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index 9ff9c5d1..5e674a03 100644 --- a/miniapps/convection/GlobalConvection2D_WENO5.jl +++ b/miniapps/convection/GlobalConvection2D_WENO5.jl @@ -173,13 +173,13 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma args = (; T = thermal.Tc, P = stokes.P, dt = Inf) ρg = @zeros(ni...), @zeros(ni...) for _ in 1:1 - compute_ρg!(ρg[2], phase_ratios, rheology, args) + compute_ρg!(ρg[2], rheology, args) @parallel init_P!(stokes.P, ρg[2], xci[2]) end # Rheology viscosity_cutoff = (1e16, 1e24) - compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + compute_viscosity!(stokes, args, rheology, viscosity_cutoff) # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -212,9 +212,9 @@ function thermal_convection2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", therma # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=Inf) - compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[end], rheology, (T=thermal.Tc, P=stokes.P)) compute_viscosity!( - stokes, phase_ratios, args, rheology, viscosity_cutoff + stokes, args, rheology, viscosity_cutoff ) # ------------------------------ From 51ea886d4e537bd2b7a6f87d4485aeae55d24850 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 13:56:53 +0200 Subject: [PATCH 87/91] format --- src/stokes/Stokes2D.jl | 20 ++++++-------------- src/stokes/Stokes3D.jl | 12 ++++-------- 2 files changed, 10 insertions(+), 22 deletions(-) diff --git a/src/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index 0bc2b742..b792b520 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -134,9 +134,7 @@ function _solve!( end @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) - ) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) return ( iter=iter, @@ -254,9 +252,7 @@ function _solve!( end @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) - ) + @parallel (@idx ni) multi_copy!(@tensor_center(stokes.τ_o), @tensor_center(stokes.τ)) return ( iter=iter, @@ -425,11 +421,9 @@ function _solve!( end 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.τ) - ) + @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) @@ -650,11 +644,9 @@ function _solve!( end 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.τ) - ) + @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) diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 9f40f5d0..2e8fd9e2 100644 --- a/src/stokes/Stokes3D.jl +++ b/src/stokes/Stokes3D.jl @@ -313,11 +313,9 @@ function _solve!( end 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.τ) - ) + @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) @@ -507,11 +505,9 @@ function _solve!( av_time = wtime0 / (iter - 1) # average time per iteration 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.τ) - ) + @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) From 38cefac34c235456c0e33508319fc3233fadf479 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 14:00:15 +0200 Subject: [PATCH 88/91] shearband examples --- miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl | 8 ++++---- .../stokes2D/shear_band/ShearBand2D_softening.jl | 7 ++----- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl index ba1fc9bd..9c7460c7 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl @@ -1,4 +1,4 @@ -using GeoParams +using GeoParams, GLMakie, CellArrays using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) @@ -80,7 +80,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") Elasticity = el_inc, ), ) - + # Initialize phase ratios ------------------------------- radius = 0.1 phase_ratios = PhaseRatio(ni, length(rheology)) @@ -108,7 +108,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") stokes.V.Vy .= PTArray(backend)([-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) - + # Time loop t, it = 0.0, 0 tmax = 3.5 @@ -139,7 +139,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs") ) tensor_invariant!(stokes.ε) push!(τII, maximum(stokes.τ.xx)) - + it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl index 26aa0efd..786f87e6 100644 --- a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl +++ b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D_softening.jl @@ -1,12 +1,9 @@ using GeoParams, GLMakie, CellArrays -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D using ParallelStencil @init_parallel_stencil(Threads, Float64, 2) -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - +const backend = CPUBackend # HELPER FUNCTIONS ----------------------------------- ---------------------------- solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η)) From 5c5bb92dd40d7945fad87c807e6504fccdaf0b6e Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 14:08:52 +0200 Subject: [PATCH 89/91] up convection --- .../Particles2D/Layered_convection2D.jl | 4 +- .../convection/WENO5/WENO_convection2D.jl | 93 +++++++++---------- 2 files changed, 47 insertions(+), 50 deletions(-) diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index dd42d543..441c402e 100644 --- a/miniapps/convection/Particles2D/Layered_convection2D.jl +++ b/miniapps/convection/Particles2D/Layered_convection2D.jl @@ -178,7 +178,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # IO ----- ------------------------------------------- # if it does not exist, make folder where figures are stored if do_vtk - vtk_dir = figdir*"\\vtk" + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -192,7 +192,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./1e3) - scatter!(ax2, Array(log10.(η[:])), Y./1e3) + scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) ylims!(ax2, minimum(xvi[2])./1e3, 0) hideydecorations!(ax2) diff --git a/miniapps/convection/WENO5/WENO_convection2D.jl b/miniapps/convection/WENO5/WENO_convection2D.jl index c2754da1..74f2073e 100644 --- a/miniapps/convection/WENO5/WENO_convection2D.jl +++ b/miniapps/convection/WENO5/WENO_convection2D.jl @@ -1,19 +1,16 @@ -using CUDA -CUDA.allowscalar(false) # for safety - -using JustRelax, JustRelax.DataIO +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO import JustRelax.@cell -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) -## NOTE: need to run one of the lines below if one wishes to switch from one backend to another -const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -using JustPIC -using JustPIC._2D +const backend_JR = CPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) + +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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions, CellArrays @@ -92,7 +89,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) li = lx, ly # domain length in x- and y- di = @. li / ni # grid step in x- and -y origin = 0.0, -ly # origin coordinates (15km f sticky air layer) - xci, xvi = lazy_grid(di, li, ni; origin=origin) # nodes at the center and vertices of the cells + grid = Geometry(ni, li; origin = origin) + (; xci, xvi) = grid # nodes at the center and vertices of the cells # ---------------------------------------------------- # Physical properties using GeoParams ---------------- @@ -127,37 +125,35 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # STOKES --------------------------------------------- # Allocate arrays needed for every Stokes problem - stokes = StokesArrays(ni, ViscoElastic) + stokes = StokesArrays(backend_JR, ni) pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.1 / √2.1) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), - periodicity = (left = false, right = false, top = false, bot = false), ) # initialize thermal profile - Half space cooling @parallel (@idx ni .+ 1) init_T!(thermal.T, xvi[2]) thermal_bcs!(thermal.T, thermal_bc) rectangular_perturbation!(thermal.T, xc_anomaly, yc_anomaly, r_anomaly, xvi) - @parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) + temperature2center!(thermal) # ---------------------------------------------------- # Buoyancy forces ρg = @zeros(ni...), @zeros(ni...) for _ in 1:1 - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P)) + compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) @parallel init_P!(stokes.P, ρg[2], xci[2]) end - # Rheology - η = @zeros(ni...) + args = (; T = thermal.Tc, P = stokes.P, dt = dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e16, 1e24) - ) - η_vep = copy(η) + # Rheology + viscosity_cutoff = (1e16, 1e24) + compute_viscosity!(stokes, phase_ratios, args, rheology, viscosity_cutoff) + (; η, η_vep) = stokes.viscosity # PT coefficients for thermal diffusion pt_thermal = PTThermalCoeffs( @@ -167,13 +163,14 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Boundary conditions flow_bcs = FlowBoundaryConditions(; free_slip = (left = true, right=true, top=true, bot=true), - periodicity = (left = false, right = false, top = false, bot = false), ) + flow_bcs!(stokes, flow_bcs) # apply boundary conditions + update_halo!(@velocity(stokes)...) # IO ----- ------------------------------------------- # if it does not exist, make folder where figures are stored if do_vtk - vtk_dir = figdir*"\\vtk" + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -187,7 +184,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./1e3) - scatter!(ax2, Array(log10.(η[:])), Y./1e3) + scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3) ylims!(ax1, minimum(xvi[2])./1e3, 0) ylims!(ax2, minimum(xvi[2])./1e3, 0) hideydecorations!(ax2) @@ -205,10 +202,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Update buoyancy and viscosity - args = (; T = thermal.Tc, P = stokes.P, dt=dt, ΔTc = thermal.ΔTc) - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24) + compute_ρg!(ρg[end], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, phase_ratios, args, rheology, viscosity_cutoff ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -218,18 +215,18 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, Inf, igg; - iterMax=50e3, - nout=1e3, - viscosity_cutoff=(1e18, 1e24) + kwargs = (; + iterMax = 150e3, + nout = 1e3, + viscosity_cutoff = viscosity_cutoff + ) ) - @parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...) + tensor_invariant!(stokes.ε) dt = compute_dt(stokes, di, dt_diff) # ------------------------------ @@ -242,11 +239,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) args, dt, di; - igg = igg, - phase = phase_ratios, - iterMax = 10e3, - nout = 1e2, - verbose = true, + kwargs = ( + igg = igg, + phase = phase_ratios, + iterMax = 10e3, + nout = 1e2, + verbose = true + ), ) T_WENO .= thermal.T[2:end-1, :] JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) @@ -256,15 +255,13 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false) # Advection -------------------- # advect particles in space - advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, dt, 2 / 3) + advection!(particles, RungeKutta2(), @velocity(stokes), (grid_vx, grid_vy), dt) # 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) + inject_particles_phase!(particles, pPhases, (pT, ), (T_buffer,), xvi) # update phase ratios - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases) - + phase_ratios_center(phase_ratios, particles, grid, pPhases) @show it += 1 t += dt @@ -294,7 +291,7 @@ n = 256 nx = n*ar - 2 ny = n - 2 igg = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid - IGG(init_global_grid(nx, ny, 0; init_MPI= true)...) + IGG(init_global_grid(nx, ny, 1; init_MPI= true)...) else igg end From 241c8f10609b4ca33bfd3d15e612af9a97d152d1 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 14:09:31 +0200 Subject: [PATCH 90/91] up readme --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index d82c541e..e25e1f5c 100644 --- a/README.md +++ b/README.md @@ -104,7 +104,7 @@ end Load JustRelax necessary modules and define backend. ```julia using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend ``` For this specific example we use particles to define the material phases, for which we rely on [JustPIC.jl](https://github.com/JuliaGeodynamics/JustPIC.jl). As in `JustRelax.jl`, we need to set up the environment of `JustPIC.jl`. This is done by running/including the following commands: @@ -113,9 +113,9 @@ For this specific example we use particles to define the material phases, for wh using JustPIC using JustPIC._2D - const backend = JustPIC.CPUBackend # Threads is the default backend - const backend = JustPIC.CUDABackend # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script - const backend = JustPIC.AMDGPUBackend # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. + const backend = CPUBackend # Threads is the default backend + const backend = CUDABackend # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script + const backend = AMDGPUBackend # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script. ``` We will also use `ParallelStencil.jl` to write some device-agnostic helper functions: From 3cbcc62ad7d0dc39b763f7a3eb08f63397999a63 Mon Sep 17 00:00:00 2001 From: aelligp Date: Thu, 2 May 2024 14:10:10 +0200 Subject: [PATCH 91/91] up Blankenbach --- .../stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl index 5cc1e812..4233b273 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -1,5 +1,5 @@ using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO -const backend_JR = JustRelax.CPUBackend +const backend_JR = CPUBackend using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) @@ -9,7 +9,7 @@ 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 = JustPIC.CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies using Printf, LinearAlgebra, GeoParams, CairoMakie