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." 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 diff --git a/.github/workflows/Documenter.yml b/.github/workflows/Documenter.yml index 82f4f36b..80eca36f 100644 --- a/.github/workflows/Documenter.yml +++ b/.github/workflows/Documenter.yml @@ -31,9 +31,9 @@ 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' + 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()' diff --git a/Project.toml b/Project.toml index 5a2cd9ad..ca1516b5 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ version = "0.1.2" [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" @@ -21,6 +20,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" diff --git a/README.md b/README.md index 8b94ef17..e25e1f5c 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,23 +99,16 @@ 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 = 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 @@ -127,6 +118,16 @@ For this specific example we use particles to define the material phases, for wh 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: +```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)...)`: ```julia @@ -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 diff --git a/docs/Project.toml b/docs/Project.toml index 69325320..bdf780a3 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,8 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38" +JustRelax = "34418575-392d-4e26-8c6d-96b0910afa06" [compat] -Documenter = "1.2.1" +Documenter = "1" +JustRelax = "0.1.2" diff --git a/docs/make.jl b/docs/make.jl index e3e4f116..9159dcc7 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -6,10 +6,23 @@ 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", + "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", + "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/index.md b/docs/src/index.md deleted file mode 100644 index 86514583..00000000 --- a/docs/src/index.md +++ /dev/null @@ -1,3 +0,0 @@ -# JustRelax.jl - -Multi-XPU stokes solver for geodynamic modelling. \ No newline at end of file diff --git a/docs/src/man/Blankenbach.md b/docs/src/man/Blankenbach.md new file mode 100644 index 00000000..eb358fd3 --- /dev/null +++ b/docs/src/man/Blankenbach.md @@ -0,0 +1,398 @@ +# 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 = 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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend +``` + +We will also use `ParallelStencil.jl` to write some device-agnostic helper functions: +```julia +using ParallelStencil +@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +``` +and will use [GeoParams.jl](https://github.com/JuliaGeodynamics/GeoParams.jl/tree/main) to define and compute physical properties of the materials: +```julia +using GeoParams +``` + +# Script + +## Model domain +```julia +nx = ny = 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 correspondent 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 +``` + +### 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/man/ShearBands.md b/docs/src/man/ShearBands.md new file mode 100644 index 00000000..33273296 --- /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 = 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 (@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) diff --git a/docs/src/man/advection.md b/docs/src/man/advection.md new file mode 100644 index 00000000..2ce1fc37 --- /dev/null +++ b/docs/src/man/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/man/backend.md b/docs/src/man/backend.md new file mode 100644 index 00000000..98bd5105 --- /dev/null +++ b/docs/src/man/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 appropriate 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/man/equations.md b/docs/src/man/equations.md new file mode 100644 index 00000000..0154a877 --- /dev/null +++ b/docs/src/man/equations.md @@ -0,0 +1,7 @@ +# Pseudo-transient iterative method + +## Heat diffusion + +## Stokes equations + +## Constitutive equations diff --git a/docs/src/man/index.md b/docs/src/man/index.md new file mode 100644 index 00000000..95a5ed48 --- /dev/null +++ b/docs/src/man/index.md @@ -0,0 +1,30 @@ +# JustRelax.jl + +Need to solve a very large multi-physics problem on many GPUs in parallel? Just Relax! + +`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) - 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: + + * 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/man/installation.md b/docs/src/man/installation.md new file mode 100644 index 00000000..2d62ff47 --- /dev/null +++ b/docs/src/man/installation.md @@ -0,0 +1,23 @@ +# Installation + +`JustRelax` 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/man/material_physics.md b/docs/src/man/material_physics.md new file mode 100644 index 00000000..7fb95b6e --- /dev/null +++ b/docs/src/man/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. + diff --git a/ext/JustRelaxAMDGPUExt.jl b/ext/JustRelaxAMDGPUExt.jl new file mode 100644 index 00000000..2dac9360 --- /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(::RocArray) = AMDGPUBackendTrait() +@inline backend(::Type{<:RocArray}) = AMDGPUBackendTrait() + +include("../src/ext/AMDGPU/2D.jl") +include("../src/ext/AMDGPU/3D.jl") + +end diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl new file mode 100644 index 00000000..6b0e5bcd --- /dev/null +++ b/ext/JustRelaxCUDAExt.jl @@ -0,0 +1,15 @@ +module JustRelaxCUDAExt + +using CUDA +using JustRelax: JustRelax +import JustRelax: PTArray, backend, CUDABackendTrait + +PTArray(::Type{CUDABackend}) = CuArray + +@inline backend(::CuArray) = CUDABackendTrait() +@inline backend(::Type{<:CuArray}) = CUDABackendTrait() + +include("../src/ext/CUDA/2D.jl") +include("../src/ext/CUDA/3D.jl") + +end diff --git a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_WENO5.jl index 385255ea..3bf7c5a9 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 = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -10,28 +11,21 @@ using JustPIC._2D # 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, 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,26 +40,24 @@ 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 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", do_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 @@ -85,68 +77,59 @@ 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 - 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, 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 ------------------------------- @@ -158,8 +141,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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) @@ -181,11 +164,8 @@ 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 + if do_vtk Vx_v = @zeros(ni.+1...) Vy_v = @zeros(ni.+1...) end @@ -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, 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,45 +222,35 @@ 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)) * - 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) # ------------------------------------------- # 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 ) * + 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) @@ -288,9 +259,8 @@ 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 + if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), @@ -315,32 +285,41 @@ 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(thermal.T[2:end-1,:]) , colormap=:lajolla, colorrange=(273,1273) ) # - h1 = heatmap!(ax1, xvi[1].*1e-3, xvi[2].*1e-3, Array(T_WENO) , 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 - + 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}") @@ -353,21 +332,21 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # ------------------------------ 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) @@ -385,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 @@ -398,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 0e093229..7f3868d8 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 = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -10,12 +11,8 @@ using JustPIC._2D # 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 Printf, LinearAlgebra, GeoParams, CairoMakie, CellArrays # Load file with all the rheology configurations include("Blankenbach_Rheology.jl") @@ -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,26 +41,24 @@ 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 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", do_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, 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 ------------------------------- @@ -164,8 +142,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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) @@ -197,7 +175,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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 @@ -213,13 +191,11 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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, 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,30 +251,26 @@ 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)) * - 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) # ------------------------------------------- # 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 ) * + 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) @@ -305,7 +279,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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 @@ -313,9 +287,8 @@ 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 + if do_vtk JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), @@ -357,12 +330,12 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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) @@ -374,7 +347,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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}") @@ -387,21 +360,21 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # ------------------------------ 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) @@ -419,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 @@ -432,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 a084e4a3..4233b273 100755 --- a/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl +++ b/miniapps/benchmarks/stokes2D/Blankenbach2D/Benchmark2D_sgd_scaled.jl @@ -1,6 +1,7 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D, JustRelax.DataIO +const backend_JR = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -10,68 +11,47 @@ using JustPIC._2D # 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 Printf, LinearAlgebra, GeoParams, CairoMakie # Load file with all the rheology configurations 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) +@parallel_indices (i, j) function init_T!(T, y) T[i, j] = 1 - y[j] return nothing 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 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", do_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 @@ -86,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 @@ -96,56 +76,51 @@ 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 = 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)) + 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 = 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 ------------------------------------ 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...) - 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, 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 ------------------------------- @@ -157,8 +132,8 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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) @@ -190,7 +165,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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 @@ -205,14 +180,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, phase_ratios, args, rheology, (-Inf, Inf)) + compute_ρg!(ρg[2], phase_ratios, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -222,19 +194,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 --------------- @@ -246,12 +218,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,29 +241,25 @@ 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]) + 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.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) @@ -296,20 +267,19 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk # 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)) - checkpointing(figdir, stokes, thermal.T, η, t) - if save_vtk - JustRelax.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) + if do_vtk + JustRelax.JustRelax2D.velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...) data_v = (; T = Array(thermal.T[2:end-1, :]), τxy = Array(stokes.τ.xy), @@ -350,51 +320,60 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, nit = 1e1, figdir="figs2D", save_vtk 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]), 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 - + 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})") 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 # ------------------------------ 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) @@ -412,12 +391,12 @@ 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 = 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 @@ -425,4 +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); 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 63a5c9c4..15954fa3 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), @@ -17,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 - - @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 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/VanKeken.jl/VanKeken.jl b/miniapps/benchmarks/stokes2D/VanKeken.jl/VanKeken.jl index c0300c06..e7a0dc0b 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) - # 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, 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.jl b/miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl index 36650f17..35e88bba 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, ViscoElastic) + 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(; @@ -78,20 +79,16 @@ function elastic_buildup(; di, flow_bcs, ρg, - η, Gc, Kb, dt, igg; - iterMax=150e3, - nout=1000, - b_width=(4, 4, 1), - verbose=true, - ) - - @parallel (@idx ni .+ 1) multi_copy!(@tensor(stokes.τ_o), @tensor(stokes.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) + kwargs = (; + iterMax=150e3, + nout=1000, + b_width=(4, 4, 1), + verbose=true, + ) ) 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..b28722f8 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, 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 78f6b10c..bfe0d97f 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/PlumeFreeSurface_2D.jl @@ -1,18 +1,14 @@ -using JustRelax, JustRelax.DataIO -import JustRelax.@cell - -# 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 -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 # Velocity helper grids for the particle advection function copyinn_x!(A, B) @@ -48,16 +44,16 @@ function init_phases!(phases, particles) x = JustRelax.@cell px[ip, i, j] depth = -(JustRelax.@cell py[ip, i, j]) - @cell phases[ip, i, j] = 2.0 + 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 @@ -123,51 +119,33 @@ 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 = 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 ρ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, phase_ratios, args, rheology, (-Inf, Inf)) # 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), + 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) - # 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...) @@ -176,17 +154,8 @@ 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 - - # 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) - ) + dt = 1e3 * (3600 * 24 * 365.25) + while it < 15 solve!( stokes, @@ -194,40 +163,40 @@ 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 = 50e3, + nout = 1e3, + viscosity_cutoff = (-Inf, Inf), + free_surface = true + ) ) dt = compute_dt(stokes, di) / 2 # ------------------------------ # 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) + 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)...) + 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.(η)), 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 3a13d3fb..f864a07f 100644 --- a/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl +++ b/miniapps/benchmarks/stokes2D/free_surface_stabilization/RayleighTaylor2D.jl @@ -1,20 +1,17 @@ # using CUDA -using JustRelax, JustRelax.DataIO -import JustRelax.@cell - -# 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 -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) @@ -36,31 +33,32 @@ 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/λ) - + + f(x, A, λ) = A * sin(π * x / λ) + @inbounds for ip in JustRelax.cellaxes(phases) # quick escape JustRelax.@cell(index[ip, i, j]) == 0 && continue x = JustRelax.@cell px[ip, i, j] - depth = -(JustRelax.@cell py[ip, i, j]) - @cell phases[ip, i, j] = 2.0 - + 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 ------------------------------------------------------------ @@ -116,55 +114,37 @@ 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) - phase_ratios = PhaseRatio(ni, length(rheology)) - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, 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-8, CFL = 0.95 / √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) # ---------------------------------------------------- - + # 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, 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, ) - # 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...) @@ -172,37 +152,31 @@ 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 + t, it = 0.0, 0 + dt = 1e3 * (3600 * 24 * 365.25) + dt_max = 50e3 * (3600 * 24 * 365.25) + while it < 500 # 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, di, flow_bcs, ρg, - η, - η_vep, phase_ratios, rheology, args, dt, igg; - iterMax = 150e3, - iterMin = 5e3, - viscosity_relaxation = 1e-2, - nout = 5e3, - free_surface = true, - viscosity_cutoff = (-Inf, Inf) + kwargs = ( + iterMax = 50e3, + iterMin = 1e3, + 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)) @@ -217,21 +191,19 @@ 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) + 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 @@ -243,25 +215,24 @@ 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, ) fig save(joinpath(figdir, "$(it).png"), fig) - end end - return + return end ## END OF MAIN SCRIPT ---------------------------------------------------------------- @@ -275,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 4b47d48c..9c7460c7 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 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 -------------------------------------------------------------------- @@ -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, 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,32 +119,27 @@ 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.τ)) - @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 334398ad..9f5f399e 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, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions @@ -151,30 +143,21 @@ 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!( - @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 e64e6273..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 / η)) @@ -88,32 +85,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, 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,26 +132,21 @@ 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.τ)) - @parallel (@idx ni) multi_copy!( - @tensor_center(stokes.τ_o), @tensor_center(stokes.τ) - ) - it += 1 t += dt diff --git a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl index cf05964f..fd83a6b4 100644 --- a/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl +++ b/miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl @@ -1,33 +1,26 @@ # Benchmark of Duretz et al. 2014 # http://dx.doi.org/10.1002/2014GL060438 -using JustRelax, JustRelax.DataIO -import JustRelax.@cell -using ParallelStencil -@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend -using JustPIC -using JustPIC._2D +using ParallelStencil, ParallelStencil.FiniteDifferences2D +@init_parallel_stencil(Threads, 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 -# setup ParallelStencil.jl environment -model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) - # Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays -using GLMakie +using GeoParams # 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) - @parallel function f_x(A, B) @all(A) = @inn_x(B) return nothing @@ -55,14 +48,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 ---------------- @@ -72,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, 1 + nxcell, max_xcell, min_xcell = 20, 32, 12 particles = init_particles( backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni... ) @@ -83,23 +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 - # 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) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -107,27 +99,21 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =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)) + 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(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 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 @@ -136,15 +122,23 @@ 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_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 - vtk_dir = figdir*"\\vtk" + vtk_dir = joinpath(figdir, "vtk") take(vtk_dir) end take(figdir) @@ -158,7 +152,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) @@ -166,27 +160,21 @@ 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 < 100 + while it < 1 # 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[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) # ------------------------------ # Stokes solver ---------------- @@ -196,18 +184,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) # ------------------------------ @@ -217,12 +205,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, ) @@ -236,36 +222,37 @@ 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, + ) ) # ------------------------------ - # 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 - @parallel (@idx ni) phase_ratios_center(phase_ratios.center, 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 # 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)...) @@ -282,7 +269,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")), @@ -314,7 +301,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) 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/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl b/miniapps/benchmarks/stokes2D/sinking_block/SinkingBlock2D.jl index e855a003..a9799b9b 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 +using JustRelax, JustRelax.JustRelax2D +const backend_JR = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @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. +using JustPIC, JustPIC._2D const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -# setup ParallelStencil.jl environment -model = PS_Setup(:Threads, Float64, 2) -environment!(model) - -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, 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/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/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/SolKz.jl b/miniapps/benchmarks/stokes2D/solkz/SolKz.jl index 7629a973..c10a2044 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 = StokesArrays(backend, ni) + (; η) = stokes.viscosity # 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 = true + ), ) t += Δt end 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/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl index c148538e..21777c0b 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 = 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..e8f7a0be 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 = 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/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..501cd989 100644 --- a/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl +++ b/miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl @@ -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,14 +85,14 @@ 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), + 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 @@ -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 @@ -113,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/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl b/miniapps/benchmarks/thermal_diffusion/diffusion/diffusion2D.jl index d782ce71..caeacec1 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 = 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..4f5387e8 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 -# 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/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl b/miniapps/benchmarks/thermal_stress/Thermal_Stress_Magma_Chamber_nondim.jl index 75fb53f8..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,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 = CPUBackend + +using ParallelStencil, ParallelStencil.FiniteDifferences2D @init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2) using JustPIC @@ -12,12 +12,7 @@ using JustPIC._2D # 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) - -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, 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, 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 @@ -762,17 +703,3 @@ 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 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..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,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 = 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 = 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, 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, 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 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 diff --git a/miniapps/convection/GlobalConvection2D_WENO5.jl b/miniapps/convection/GlobalConvection2D_WENO5.jl index 5827a849..5e674a03 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 = 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], 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, 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], rheology, (T=thermal.Tc, P=stokes.P)) + compute_viscosity!( + stokes, 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/GlobalConvection3D_Upwind.jl b/miniapps/convection/GlobalConvection3D_Upwind.jl index 3c1cc602..6c3d6838 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 ------------------------------------------------- @@ -211,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) @@ -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") @@ -281,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) @@ -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); diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl index 167855da..441c402e 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) - # 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,39 +147,38 @@ 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, 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 - vtk_dir = figdir*"\\vtk" + vtk_dir = joinpath(figdir,"vtk") take(vtk_dir) end take(figdir) @@ -196,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) @@ -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, 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..560fd1ea 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 = 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") @@ -51,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 @@ -79,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 @@ -98,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 ------------------------------------ @@ -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) # ---------------------------------------------------- - + args = (; T = thermal.Tc, P = stokes.P, dt=Inf) # 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, 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 @@ -186,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 @@ -236,13 +230,13 @@ 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) - @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, 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..b3878993 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) - # 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, 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, 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,14 +258,13 @@ 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 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 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 diff --git a/src/Interpolations.jl b/src/Interpolations.jl index 1dc4f679..f6b5a102 100644 --- a/src/Interpolations.jl +++ b/src/Interpolations.jl @@ -16,12 +16,17 @@ end # From cell vertices to cell center -function temperature2center!(thermal::ThermalArrays) - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) +temperature2center!(thermal) = temperature2center!(backend(thermal), thermal) +function temperature2center!(::CPUBackendTrait, thermal::JustRelax.ThermalArrays) + return _temperature2center!(thermal) +end + +function _temperature2center!(thermal::JustRelax.ThermalArrays) + @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 +39,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 +49,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..0322619f 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,61 +12,42 @@ using StaticArrays function solve!() end -include("topology/Topology.jl") -export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g +abstract type AbstractBackend end +struct CPUBackend <: AbstractBackend end +struct AMDGPUBackend <: AbstractBackend end -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 +PTArray() = Array +PTArray(::Type{CPUBackend}) = Array +PTArray(::T) where {T} = error(ArgumentError("Unknown backend $T")) -include("phases/CellArrays.jl") -export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! +export PTArray, CPUBackend, CUDABackend, AMDGPUBackend + +include("types/stokes.jl") +# export StokesArrays, PTStokesCoeffs + +include("types/heat_diffusion.jl") +# export ThermalArrays, PTThermalCoeffs -include("rheology/StressUpdate.jl") -export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear! +include("types/phases.jl") +# export PhaseRatio -include("MetaJustRelax.jl") +include("boundaryconditions/types.jl") +export TemperatureBoundaryConditions, FlowBoundaryConditions -include("stokes/MetaStokes.jl") -export PS_Setup, environment!, ps_reset! +include("types/traits.jl") +export BackendTrait, CPUBackendTrait, NonCPUBackendTrait -include("thermal_diffusion/MetaDiffusion.jl") +include("topology/Topology.jl") +export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g + +include("phases/CellArrays.jl") +export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex! -include("thermal_diffusion/Rheology.jl") +include("JustRelax_CPU.jl") include("IO/DataIO.jl") +include("array_conversions.jl") +export Array + end # module diff --git a/src/JustRelax_CPU.jl b/src/JustRelax_CPU.jl new file mode 100644 index 00000000..a75d0970 --- /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 +using MPI + +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend +import JustRelax: PTStokesCoeffs +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(Threads, Float64, 2) + +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 +using MPI + +import JustRelax: IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend +import JustRelax: PTStokesCoeffs +import JustRelax: + AbstractBoundaryConditions, TemperatureBoundaryConditions, FlowBoundaryConditions + +@init_parallel_stencil(Threads, Float64, 3) + +include("common.jl") +include("stokes/Stokes3D.jl") +export solve! + +end diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl deleted file mode 100644 index 21092997..00000000 --- a/src/MetaJustRelax.jl +++ /dev/null @@ -1,177 +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__, "array_conversions.jl")) - export Array - - include(joinpath(@__DIR__, "Utils.jl")) - export @allocate, @add, @idx, @copy - export @velocity, - @strain, - @stress, - @tensor, - @shear, - @normal, - @stress_center, - @strain_center, - @tensor_center, - @qT, - @qT2, - @residuals, - compute_dt, - assign!, - tupleize, - compute_maxloc!, - continuation_log, - continuation_linear, - mean_mpi, - norm_mpi, - minimum_mpi, - maximum_mpi, - multi_copy!, - take - - include(joinpath(@__DIR__, "boundaryconditions/BoundaryConditions.jl")) - export pureshear_bc!, - FlowBoundaryConditions, - flow_bcs!, - TemperatureBoundaryConditions, - thermal_boundary_conditions!, - thermal_bcs!, - free_slip_x!, - free_slip_y!, - free_slip_z!, - apply_free_slip!, - free_surface_bcs! - - include(joinpath(@__DIR__, "phases/phases.jl")) - export PhaseRatio, fn_ratio, phase_ratios_center - - 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/Utils.jl b/src/Utils.jl index e9c8eada..2be50507 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,10 @@ macro residuals(A) end end -@inline function unpack_residuals(A::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::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 @@ -387,47 +399,57 @@ 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...) + return _compute_dt(S, args...) +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) +@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 """ - 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) - 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/advection/weno5.jl b/src/advection/weno5.jl index ef3f8201..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 = @zeros(ni...) - fL::A = @zeros(ni...) - fR::A = @zeros(ni...) - fB::A = @zeros(ni...) - fT::A = @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 diff --git a/src/array_conversions.jl b/src/array_conversions.jl index 237c6430..febb4b12 100644 --- a/src/array_conversions.jl +++ b/src/array_conversions.jl @@ -1,29 +1,4 @@ -# Device trait system - -abstract type DeviceTrait end -struct CPUDeviceTrait <: DeviceTrait end -struct NonCPUDeviceTrait <: DeviceTrait end - -@inline iscpu(::Array) = CPUDeviceTrait() -@inline iscpu(::T) where {T<: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() +import Base: Array ## Conversion of structs to CPU @@ -31,15 +6,30 @@ struct NonCPUDeviceTrait <: DeviceTrait end 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 @@ -49,7 +39,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/boundaryconditions/BoundaryConditions.jl b/src/boundaryconditions/BoundaryConditions.jl index 594da411..bdc9efeb 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)...) @@ -47,7 +18,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::JustRelax.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 @@ -61,13 +39,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)) @@ -78,339 +58,7 @@ function _flow_bcs!(bcs::FlowBoundaryConditions, 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 - -function free_surface_bcs!( - stokes::StokesArrays{A,B,C,D,E,2}, bcs::FlowBoundaryConditions -) where {A,B,C,D,E} - 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 - -@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 - Vy[i + 1, end] = - Vy[i + 1, end - 1] + - 3.0 / 2.0 * - ( - 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 - 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::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]]) - - return nothing -end - -@parallel_indices (j) function free_slip_x!(A::AbstractArray{eltype(PTArray),2}) - 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}) - 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 - -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} - 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 - -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 +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..f054ad0a --- /dev/null +++ b/src/boundaryconditions/free_slip.jl @@ -0,0 +1,173 @@ +@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 diff --git a/src/boundaryconditions/free_surface.jl b/src/boundaryconditions/free_surface.jl new file mode 100644 index 00000000..1e4ff780 --- /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 diff --git a/src/boundaryconditions/no_slip.jl b/src/boundaryconditions/no_slip.jl new file mode 100644 index 00000000..2a677277 --- /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 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/boundaryconditions/types.jl b/src/boundaryconditions/types.jl new file mode 100644 index 00000000..7eec4792 --- /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 diff --git a/src/common.jl b/src/common.jl new file mode 100644 index 00000000..8c4d3d80 --- /dev/null +++ b/src/common.jl @@ -0,0 +1,72 @@ +include("types/constructors/stokes.jl") +export StokesArrays, PTStokesCoeffs + +include("types/constructors/heat_diffusion.jl") +export ThermalArrays, PTThermalCoeffs + +include("types/constructors/phases.jl") +export PhaseRatio + +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 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("rheology/StressUpdate.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 PTThermalCoeffs, heatdiffusion_PT!, compute_shear_heating! diff --git a/src/ext/AMDGPU/2D.jl b/src/ext/AMDGPU/2D.jl new file mode 100644 index 00000000..dcab6325 --- /dev/null +++ b/src/ext/AMDGPU/2D.jl @@ -0,0 +1,186 @@ +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::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) +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::RocArray, 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::RocArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR2D.compute_ρg!(ρg::RocArray, 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<:RocArray} + return vertex2center!(center, vertex) +end + +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<:RocArray} + 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 new file mode 100644 index 00000000..472905e1 --- /dev/null +++ b/src/ext/AMDGPU/3D.jl @@ -0,0 +1,186 @@ +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!(::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 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) +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::RocArray, 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::RocArray, rheology, args) + return compute_ρg!(ρg, rheology, args) +end +function JR3D.compute_ρg!(ρg::RocArray, 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<:RocArray} + return vertex2center!(center, vertex) +end + +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<:RocArray} + 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 new file mode 100644 index 00000000..e4701169 --- /dev/null +++ b/src/ext/CUDA/2D.jl @@ -0,0 +1,185 @@ +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) + 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 new file mode 100644 index 00000000..73158c3a --- /dev/null +++ b/src/ext/CUDA/3D.jl @@ -0,0 +1,185 @@ +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) + 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/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 9a7890a5..22452832 100644 --- a/src/phases/phases.jl +++ b/src/phases/phases.jl @@ -1,27 +1,10 @@ - -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) + 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} @@ -34,59 +17,6 @@ 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 - """ fn_ratio(fn::F, rheology::NTuple{N, AbstractMaterialParamsStruct}, ratio) where {N, F} @@ -118,8 +48,27 @@ end end end -# ParallelStencil launch kernel for 2D -@parallel_indices (I...) function phase_ratios_center( +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::JustRelax.PhaseRatio, particles, grid::Geometry, phases +) + return _phase_ratios_center(phase_ratios, particles, grid, phases) +end + +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 + ) + 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 +76,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..553b5af6 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::JustRelax.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..f051d501 100644 --- a/src/rheology/Viscosity.jl +++ b/src/rheology/Viscosity.jl @@ -1,6 +1,23 @@ ## 2D KERNELS +function compute_viscosity!( + stokes::JustRelax.StokesArrays, args, rheology, cutoff; relaxation=1e0 +) + return compute_viscosity!(backend(stokes), stokes, relaxation, args, rheology, cutoff) +end -@parallel_indices (I...) function compute_viscosity!( +function compute_viscosity!(::CPUBackendTrait, stokes, ν, args, rheology, cutoff) + return _compute_viscosity!(stokes, ν, args, rheology, cutoff) +end + +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 + ) + return nothing +end + +@parallel_indices (I...) function compute_viscosity_kernel!( η, ν, εxx, εyy, εxyv, args, rheology, cutoff ) @@ -30,7 +47,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 +72,48 @@ end return nothing end -@parallel_indices (I...) function compute_viscosity!( +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 + ) +end + +function compute_viscosity!( + ::CPUBackendTrait, + stokes::JustRelax.StokesArrays, + ν, + phase_ratios, + args, + rheology, + cutoff, +) + return _compute_viscosity!(stokes, ν, phase_ratios, args, rheology, cutoff) +end + +function _compute_viscosity!( + stokes::JustRelax.StokesArrays, + ν, + phase_ratios::JustRelax.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 +148,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 +181,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 ) @@ -182,23 +231,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} @@ -209,9 +241,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/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/stokes/Stokes2D.jl b/src/stokes/Stokes2D.jl index d9c0c9f4..b792b520 100644 --- a/src/stokes/Stokes2D.jl +++ b/src/stokes/Stokes2D.jl @@ -1,88 +1,27 @@ -## 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::JustRelax.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}, - pt_stokes::PTStokesCoeffs, +# backend trait +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!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{2,T}, flow_bcs::FlowBoundaryConditions, ρg, - η, K, dt, igg::IGG; @@ -90,11 +29,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 @@ -192,6 +133,9 @@ function JustRelax.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, @@ -203,13 +147,12 @@ function JustRelax.solve!( end # visco-elastic solver -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, - η, G, K, dt, @@ -218,11 +161,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 @@ -306,6 +251,9 @@ function JustRelax.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, @@ -318,29 +266,30 @@ end # GeoParams: general (visco-elasto-plastic) solver -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, - η, - η_vep, rheology::MaterialParams, args, dt, igg::IGG; - viscosity_cutoff=(1e16, 1e24), + viscosity_cutoff=(-Inf, Inf), + viscosity_relaxation=1e-2, iterMax=10e3, nout=500, b_width=(4, 4, 0), verbose=true, free_surface=false, -) 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 @@ -378,15 +327,14 @@ function JustRelax.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; relaxation=viscosity_relaxation ) compute_maxloc!(ητ, η; window=(1, 1)) update_halo!(ητ) @@ -408,7 +356,7 @@ function JustRelax.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!( @@ -474,6 +422,9 @@ function JustRelax.solve!( 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) @@ -489,33 +440,33 @@ end ## With phase ratios -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{2,T}, flow_bcs, ρg, - η, - η_vep, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, rheology, args, dt, igg::IGG; - viscosity_cutoff=(1e16, 1e24), + 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), 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 @@ -580,9 +531,7 @@ function JustRelax.solve!( # stokes.P[1, end] = stokes.P[2, end] # stokes.P[end, end] = stokes.P[end - 1, end] - if rem(iter, 5) == 0 - @parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, 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... @@ -592,14 +541,13 @@ function JustRelax.solve!( @copy η0 η end if do_visc - @parallel (@idx ni) compute_viscosity!( - η, - viscosity_relaxation, - phase_ratios.center, - @strain(stokes)..., + compute_viscosity!( + stokes, + phase_ratios, args, rheology, - viscosity_cutoff, + viscosity_cutoff; + relaxation=viscosity_relaxation, ) end @@ -620,9 +568,9 @@ function JustRelax.solve!( dt, θ_dτ, ) - free_surface_bcs!(stokes, flow_bcs) + # free_surface_bcs!(stokes, flow_bcs) - @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!( @@ -696,7 +644,9 @@ function JustRelax.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) @@ -711,16 +661,14 @@ function JustRelax.solve!( ) end -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,2}, - thermal::ThermalArrays, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + thermal::JustRelax.ThermalArrays, + pt_stokes, di::NTuple{2,T}, flow_bcs, ϕ, ρg, - η, - η_vep, phase_v, phase_c, args_η, @@ -731,12 +679,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(η) @@ -856,5 +806,3 @@ function JustRelax.solve!( norm_∇V=norm_∇V, ) end - -end # END OF MODULE diff --git a/src/stokes/Stokes3D.jl b/src/stokes/Stokes3D.jl index 2165f5ec..2e8fd9e2 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,54 +10,42 @@ 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::JustRelax.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) - ]) +## 3D VISCO-ELASTIC STOKES SOLVER +function solve!(stokes::JustRelax.StokesArrays, args...; kwargs) + return solve!(backend(stokes), stokes, args...; kwargs) end -## 3D VISCO-ELASTIC STOKES SOLVER +# entry point for extensions +solve!(::CPUBackendTrait, stokes, args...; kwargs) = _solve!(stokes, args...; kwargs...) -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{3,T}, flow_bcs, ρg, - η, K, G, dt, igg::IGG; + viscosity_relaxation=1e-2, iterMax=10e3, nout=500, b_width=(4, 4, 4), verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # solver related ϵ = pt_stokes.ϵ # geometry _di = @. 1 / di ni = size(stokes.P) + (; η) = stokes.viscosity # ~preconditioner ητ = deepcopy(η) @@ -199,14 +157,12 @@ end ## 3D VISCO-ELASTO-PLASTIC STOKES SOLVER WITH GeoParams.jl -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{3,T}, flow_bcs::FlowBoundaryConditions, ρg, - η, - η_vep, rheology::MaterialParams, args, dt, @@ -214,14 +170,17 @@ function JustRelax.solve!( iterMax=10e3, nout=500, b_width=(4, 4, 4), + viscosity_relaxation=1e-2, verbose=true, -) where {A,B,C,D,T} + kwargs..., +) where {T} # solver related ϵ = pt_stokes.ϵ # geometry _di = @. 1 / di ni = size(stokes.P) + (; η, η_vep) = stokes.viscosity # ~preconditioner ητ = deepcopy(η) @@ -266,18 +225,16 @@ function JustRelax.solve!( stokes.∇V, @strain(stokes)..., @velocity(stokes)..., _di... ) - # # Update buoyancy - # @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) + # Update buoyancy + @parallel (@idx ni) compute_ρg!(ρg[3], rheology, args) - ν = 1e-3 - @parallel (@idx ni) compute_viscosity!( - η, - 1.0, - phase_ratios.center, - @strain(stokes)..., + compute_viscosity!( + stokes, + phase_ratios, args, rheology, - viscosity_cutoff, + viscosity_cutoff; + relaxation=viscosity_relaxation, ) @parallel (@idx ni) compute_τ_nonlinear!( @@ -356,7 +313,12 @@ function JustRelax.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, @@ -372,15 +334,13 @@ function JustRelax.solve!( end # GeoParams and multiple phases -function JustRelax.solve!( - stokes::StokesArrays{ViscoElastic,A,B,C,D,3}, - pt_stokes::PTStokesCoeffs, +function _solve!( + stokes::JustRelax.StokesArrays, + pt_stokes, di::NTuple{3,T}, flow_bc::FlowBoundaryConditions, ρg, - η, - η_vep, - phase_ratios::PhaseRatio, + phase_ratios::JustRelax.PhaseRatio, rheology::NTuple{N,AbstractMaterialParamsStruct}, args, dt, @@ -389,8 +349,10 @@ function JustRelax.solve!( nout=500, b_width=(4, 4, 4), verbose=true, + viscosity_relaxation=1e-2, viscosity_cutoff=(-Inf, Inf), -) where {A,B,C,D,T,N} + kwargs..., +) where {T,N} ## UNPACK @@ -399,6 +361,7 @@ function JustRelax.solve!( # geometry _di = @. 1 / di ni = size(stokes.P) + (; η, η_vep) = stokes.viscosity # errors err = Inf @@ -445,18 +408,16 @@ function JustRelax.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)..., + compute_viscosity!( + stokes, + phase_ratios, args, rheology, - viscosity_cutoff, + viscosity_cutoff; + relaxation=viscosity_relaxation, ) @parallel (@idx ni) compute_τ_nonlinear!( @@ -476,9 +437,9 @@ function JustRelax.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, @@ -545,6 +506,12 @@ function JustRelax.solve!( 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, err_evo1=err_evo1, @@ -557,5 +524,3 @@ function JustRelax.solve!( av_time=av_time, ) end - -end # END OF MODULE diff --git a/src/stokes/StressKernels.jl b/src/stokes/StressKernels.jl index de79d5fc..e48b7ca3 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,42 @@ end return nothing end -@parallel_indices (I...) function tensor_invariant!(II, xx, yy, xyv) +""" + tensor_invariant!(A::JustRelax.SymmetricTensor) + +Compute the tensor invariant of the given symmetric tensor `A`. + +# Arguments +- `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) + return _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 +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/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.jl b/src/thermal_diffusion/DiffusionPT.jl index fdaa0f93..cdb22e20 100644 --- a/src/thermal_diffusion/DiffusionPT.jl +++ b/src/thermal_diffusion/DiffusionPT.jl @@ -1,818 +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") - - thermal_bcs!(thermal.T, thermal_bc) - @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) - @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") - - thermal_bcs!(thermal.T, thermal_bc) - @parallel update_ΔT!(thermal.ΔT, thermal.T, thermal.Told) - @parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T) - @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..c356cb01 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_GeoParams.jl @@ -0,0 +1,103 @@ +## Phases + +@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::JustRelax.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..a566ec17 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_coefficients.jl @@ -0,0 +1,138 @@ +function PTThermalCoeffs( + ::Type{CPUBackend}, + K, + ρCp, + dt, + di::NTuple{nDim,T}, + li::NTuple{nDim,Any}; + ϵ=1e-8, + CFL=0.9 / √3, +) where {nDim,T} + return 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, + 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 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 JustRelax.PTThermalCoeffs(CFL, ϵ, max_lxyz, max_lxyz^2, Vpdτ, θr_dτ, dτ_ρ) +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, +) where {nDim,T} + 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 +) 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 JustRelax.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..bdfbebd5 --- /dev/null +++ b/src/thermal_diffusion/DiffusionPT_solver.jl @@ -0,0 +1,209 @@ +function heatdiffusion_PT!(thermal, args...; kwargs) + return heatdiffusion_PT!(backend(thermal), thermal, args...; kwargs=kwargs) +end + +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) + +Heat diffusion solver using Pseudo-Transient iterations. Both `K` and `ρCp` are n-dimensional arrays. +""" +function _heatdiffusion_PT!( + thermal::JustRelax.ThermalArrays, + pt_thermal::JustRelax.PTThermalCoeffs, + thermal_bc::TemperatureBoundaryConditions, + K::AbstractArray, + ρCp::AbstractArray, + dt, + di; + igg=nothing, + b_width=(4, 4, 1), + iterMax=50e3, + nout=1e3, + verbose=true, + kwargs..., +) + # 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) + temperature2center!(thermal) + + 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::JustRelax.ThermalArrays, + pt_thermal::JustRelax.PTThermalCoeffs, + thermal_bc::TemperatureBoundaryConditions, + rheology, + args::NamedTuple, + dt, + di; + igg=nothing, + phase=nothing, + b_width=(4, 4, 4), + iterMax=50e3, + nout=1e3, + verbose=true, + kwargs..., +) + 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 + !isnothing(phase) && 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) + temperature2center!(thermal) + + 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/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 diff --git a/src/thermal_diffusion/ShearHeating.jl b/src/thermal_diffusion/ShearHeating.jl new file mode 100644 index 00000000..cd1d559b --- /dev/null +++ b/src/thermal_diffusion/ShearHeating.jl @@ -0,0 +1,55 @@ +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) + τij, τij_o, εij = cache_tensors(τ, τ_old, ε, I...) + εij_el = @. 0.5 * ((τij - τij_o) * _Gdt) + shear_heating[I...] = compute_shearheating(rheology, τij, εij, εij_el) + return nothing +end + +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, + @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}, + ε::NTuple{N,T}, + phase_ratios::CellArray, + rheology, + dt, +) where {N,T} + phase = @inbounds phase_ratios[I...] + _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) + τij, τij_o, εij = cache_tensors(τ, τ_old, ε, I...) + εij_el = @. 0.5 * ((τij - τij_o) * _Gdt) + shear_heating[I...] = compute_shearheating(rheology, phase, τij, εij, εij_el) + return nothing +end diff --git a/src/thermal_diffusion/Shearheating.jl b/src/thermal_diffusion/Shearheating.jl deleted file mode 100644 index 8fc1cb9c..00000000 --- a/src/thermal_diffusion/Shearheating.jl +++ /dev/null @@ -1,26 +0,0 @@ -@parallel_indices (I...) function compute_shear_heating!( - shear_heating, τ::NTuple{N,T}, τ_old::NTuple{N,T}, ε::NTuple{N,T}, rheology, dt -) where {N,T} - _Gdt = inv(get_shear_modulus(rheology) * dt) - τij, τij_o, εij = cache_tensors(τ, τ_old, ε, I...) - εij_el = @. 0.5 * ((τij - τij_o) * _Gdt) - shear_heating[I...] = compute_shearheating(rheology, τij, εij, εij_el) - return nothing -end - -@parallel_indices (I...) function compute_shear_heating!( - shear_heating, - τ::NTuple{N,T}, - τ_old::NTuple{N,T}, - ε::NTuple{N,T}, - phase_ratios::CellArray, - rheology, - dt, -) where {N,T} - phase = @inbounds phase_ratios[I...] - _Gdt = inv(fn_ratio(get_shear_modulus, rheology, phase) * dt) - τij, τij_o, εij = cache_tensors(τ, τ_old, ε, I...) - εij_el = @. 0.5 * ((τij - τij_o) * _Gdt) - shear_heating[I...] = compute_shearheating(rheology, phase, τij, εij, εij_el) - return nothing -end diff --git a/src/types/constructors/heat_diffusion.jl b/src/types/constructors/heat_diffusion.jl new file mode 100644 index 00000000..90772dde --- /dev/null +++ b/src/types/constructors/heat_diffusion.jl @@ -0,0 +1,61 @@ +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) + ΔT = @zeros(nx + 3, ny + 1) + ΔTc = @zeros(nx, ny) + 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, + Told, + ΔT, + ΔTc, + 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) + Δ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(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) + return JustRelax.ThermalArrays( + T, Tc, Told, ΔT, ΔTc, 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..432a17cd --- /dev/null +++ b/src/types/constructors/phases.jl @@ -0,0 +1,12 @@ +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,)) + # 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..db12d424 --- /dev/null +++ b/src/types/constructors/stokes.jl @@ -0,0 +1,92 @@ +## 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(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, ny, nz), # yy + @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 + @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(::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...) + ∇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 new file mode 100644 index 00000000..6d855057 --- /dev/null +++ b/src/types/heat_diffusion.jl @@ -0,0 +1,45 @@ +struct ThermalArrays{_T} + T::_T # Temperature @ grid nodes + Tc::_T # Temperature @ cell centers + Told::_T + ΔT::_T + ΔTc::_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 +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")) +end +function ThermalArrays(::Number, ::Number, ::Number) + throw(ArgumentError("ThermalArrays dimensions must be given as integers")) +end + +## 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 diff --git a/src/types/phases.jl b/src/types/phases.jl new file mode 100644 index 00000000..8e7bd9dd --- /dev/null +++ b/src/types/phases.jl @@ -0,0 +1,8 @@ +struct PhaseRatio{T} + vertex::T + center::T + + PhaseRatio(vertex::T, center::T) where {T<:AbstractArray} = new{T}(vertex, center) +end + +@inline PhaseRatio(::Type{CPUBackend}, ni, num_phases) = PhaseRatio(ni, num_phases) diff --git a/src/types/stokes.jl b/src/types/stokes.jl new file mode 100644 index 00000000..e7b7cd48 --- /dev/null +++ b/src/types/stokes.jl @@ -0,0 +1,156 @@ +## 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) + +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...)...) +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 + +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) +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 +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")) +end +function StokesArrays(::Number, ::Number, ::Number) + throw(ArgumentError("StokesArrays dimensions must be given as integers")) +end + +## 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..3eaf37de --- /dev/null +++ b/src/types/traits.jl @@ -0,0 +1,22 @@ +abstract type BackendTrait end +struct CPUBackendTrait <: BackendTrait end +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) = 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")) diff --git a/test/Project.toml b/test/Project.toml index c1444d4f..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] -JustPIC = "0.2" -ParallelStencil = "0.11, 0.12" CellArrays = "0.2" +GeoParams = "0.5.7" +JustPIC = "0.3.2" +ParallelStencil = "0.11, 0.12" StaticArrays = "1" Statistics = "1" julia = "1.9, 1.10" diff --git a/test/runtests.jl b/test/runtests.jl index 6a2032dd..f7bf478e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,6 +12,11 @@ function runtests() nfail = 0 printstyled("Testing package JustRelax.jl\n"; bold=true, color=:white) + f0 = ("test_traits.jl", "test_types.jl", "test_arrays_conversions.jl") + for f in f0 + include(f) + end + for f in testfiles occursin("burstedde", f) && continue @@ -23,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_VanKeken.jl b/test/test_VanKeken.jl index 74d20538..48d1daaa 100644 --- a/test/test_VanKeken.jl +++ b/test/test_VanKeken.jl @@ -5,41 +5,26 @@ 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) - # 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 +41,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 +83,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 +92,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, phase_ratios, args, rheology, (-Inf, Inf)) # Boundary conditions flow_bcs = FlowBoundaryConditions(; @@ -153,10 +133,6 @@ function VanKeken2D(ny=32, nx=32) while it < nt - # Update buoyancy - @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args) - # ------------------------------ - # Stokes solver ---------------- iters = solve!( stokes, @@ -164,23 +140,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 +165,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 @@ -210,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_arrays_conversions.jl b/test/test_arrays_conversions.jl index 88d42fa3..375bbb60 100644 --- a/test/test_arrays_conversions.jl +++ b/test/test_arrays_conversions.jl @@ -1,24 +1,16 @@ -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) - - @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} + ni = 2, 2 + stokes = StokesArrays(CPUBackend, ni) + thermal = ThermalArrays(CPUBackend, ni) - @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") + @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 + diff --git a/test/test_boundary_conditions2D.jl b/test/test_boundary_conditions2D.jl index e0d9dd0b..14e9d93d 100644 --- a/test/test_boundary_conditions2D.jl +++ b/test/test_boundary_conditions2D.jl @@ -1,29 +1,28 @@ -using JustRelax +using JustRelax, JustRelax.JustRelax2D using Test, Suppressor -model = PS_Setup(:cpu, Float64, 2) -environment!(model) +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 724ed5eb..21e3276c 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,9 @@ 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 +115,9 @@ 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 + + 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 end diff --git a/test/test_diffusion2D_multiphase.jl b/test/test_diffusion2D_multiphase.jl index d4cb8d6a..e02ca525 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) - 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,9 @@ 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 + + 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 diff --git a/test/test_diffusion3D.jl b/test/test_diffusion3D.jl index 19661f01..92126ed7 100644 --- a/test/test_diffusion3D.jl +++ b/test/test_diffusion3D.jl @@ -1,14 +1,12 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using GeoParams, CellArrays -using JustRelax, JustRelax.DataIO +using GeoParams +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..6d68a9e2 100644 --- a/test/test_diffusion3D_multiphase.jl +++ b/test/test_diffusion3D_multiphase.jl @@ -1,11 +1,14 @@ 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) +const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend + + using JustPIC using JustPIC._3D # Threads is the default backend, @@ -13,10 +16,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 +48,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(; @@ -121,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...) @@ -149,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 @@ -172,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_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_shearband2D.jl b/test/test_shearband2D.jl index e07b03d1..79aaeac5 100644 --- a/test/test_shearband2D.jl +++ b/test/test_shearband2D.jl @@ -1,14 +1,12 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using GeoParams, 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 / η)) @@ -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 @@ -93,32 +91,29 @@ 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 --------------------------------------------- # 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, 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) @@ -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..a8380598 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 @@ -110,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 - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf) + compute_viscosity!( + stokes, phase_ratios, args, rheology, (-Inf, Inf) ) # Boundary conditions @@ -121,8 +118,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) @@ -143,19 +140,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 +174,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..23516ba2 100644 --- a/test/test_shearheating2D.jl +++ b/test/test_shearheating2D.jl @@ -1,33 +1,28 @@ -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 +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(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2) -environment!(model) +# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend # Load script dependencies -using Printf, LinearAlgebra, GeoParams, CellArrays +using 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) @@ -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,23 +83,22 @@ 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) # ---------------------------------------------------- # TEMPERATURE PROFILE -------------------------------- - thermal = ThermalArrays(ni) + thermal = ThermalArrays(backend_JR, ni) thermal_bc = TemperatureBoundaryConditions(; no_flux = (left = true, right = true, top = false, bot = false), ) @@ -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(η) + args = (; T = thermal.Tc, P = stokes.P, dt = Inf) + compute_viscosity!(stokes, 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(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) @@ -159,13 +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=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[2], phase_ratios.center, rheology, args) - # ------------------------------ # Stokes solver ---------------- iters = solve!( @@ -174,19 +153,18 @@ function Shearheating2D() 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) # ------------------------------ @@ -196,12 +174,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, ) @@ -215,17 +191,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,10 +212,9 @@ 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 +224,6 @@ function Shearheating2D() finalize_global_grid(; finalize_MPI = true) return iters, thermal - - end @testset "Shearheating2D" begin diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl index 4c622f6e..c6ed8764 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 = 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 = 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, 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) @@ -142,13 +136,6 @@ 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 = 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) - # ------------------------------ # Stokes solver ---------------- iters = solve!( @@ -157,19 +144,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 +164,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 +181,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_sinking_block.jl b/test/test_sinking_block.jl index dce3131f..5f44969c 100644 --- a/test/test_sinking_block.jl +++ b/test/test_sinking_block.jl @@ -1,22 +1,16 @@ push!(LOAD_PATH, "..") using Test, Suppressor -using JustRelax -using ParallelStencil +using JustRelax, JustRelax.JustRelax2D +using ParallelStencil, ParallelStencil.FiniteDifferences2D @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 = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend -using Printf, LinearAlgebra, GeoParams, CellArrays +using GeoParams ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT -------------------------------- @@ -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...)) + args = (; T = @ones(ni...), P = stokes.P, dt=Inf) η_cutoff = -Inf, Inf - @parallel (@idx ni) compute_viscosity!( - η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, η_cutoff - ) - η_vep = deepcopy(η) + compute_viscosity!(stokes, phase_ratios, args, rheology, (-Inf, Inf)) # ---------------------------------------------------- # Boundary conditions @@ -160,39 +150,35 @@ 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...)) - local iters, velocity iters = solve!( stokes, pt_stokes, 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) return iters, velocity - - end @testset "Sinking_Block2D" begin 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") diff --git a/test/test_stokes_solvi3D.jl b/test/test_stokes_solvi3D.jl index a06cc519..170eda32 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() @@ -38,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 b/test/test_traits.jl new file mode 100644 index 00000000..bc942702 --- /dev/null +++ b/test/test_traits.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 bk(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