Rework: Extensions

.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."

.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 "[email protected]"
+            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

.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()'

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"

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,44 +266,49 @@ 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)
     lines!(ax4, ttot, sol, color = :red)
     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
 

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"

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",
     ],
 )