Coupling 3D solvers with JustPIC (#62)

* Pin to ParallelStencik 0.9.0 * update density kernel syntax * generalize Delta function * use triple equality in Delta function * explicitly add fma's * update indexing syntax * move Delta to Utils.jl * format & add @inbounds * format changes * 3d multiphase sxript * more updates to 3d script * generalize compute_maxloc! * generalize phase_ratios_center * fix 3D viscosity kernel * some cosmetic changes * include StressKernels.jl; export tensor_invariant! * export only tensor_invariant! * generalize multi_copyto! * update PS syntax of pressure kernels * syntax updates in stress kernels * syntax updates in pressure kernels * do not consider P in isyielding * move @idx to the top * include StressKernels.jl in the appropiate file. * some fixes in 2D convection miniapp * updates in 3D scripts * minor fixes in Stokes2D * re-structuring and cleaning of Stokes3D.jl * formatting * minimal changes to some Utils funs * dimension agnostic stress update * fix NaNs on eff visc when strain is zero * some fixes * update 3d scripts * make plasticity great again * Plasticity is great again * fixed shear stress at the vertices * 3D script * call correct viscosity kernel in the 3D case * 3D script * update 3D stress kernels * update Stokes3D * format * update scripts * fix 3D vertex stress kernel * simplify _compute_τ_nonlinear! * format and remove unused arguments * update P kernels syntax * updates * improve multi_copy! * fix 3D center2vertex tensor interpolation * fix/improve few thins * update PS syntax * fix 3D stokes * format * add volumetric plasticity component * fix voluemtric plasticity * copy P into P0 * fix small bug * few tweaks * force fma * format and remove commented lines * format * add 1.10 to CI * fix conflicts * replace delta by allzero * update convection miniapp * update scripts * fix merge errors * rework @idx * fix merge errors * take function * fix stress update * update viscosity calls * update 3D scripts * remove dead code * fix 3D stuff * move 3D scripts to own folder * tweak compats * format * fix IGG dimensions * remove JustPIC from Project.toml * remove duplicated lines * format
PTsolvers · Oct 20, 2023 · c02d41a · c02d41a
1 parent 0c467fc
commit c02d41a
Show file tree

Hide file tree

Showing 20 changed files with 823 additions and 278 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -15,6 +15,7 @@ jobs:
       matrix:
         version:
           - '1.9'
+          - '1.10'
           # - 'nightly'
         os:
           - ubuntu-latest

diff --git a/Project.toml b/Project.toml
@@ -20,7 +20,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [compat]
-julia = "1.9"
 AMDGPU = "0.5.4"
 CUDA = "4.4.1"
 CellArrays = "0.1.3"
@@ -31,7 +30,9 @@ MPI = "0.20.14"
 ParallelStencil = "0.9.0"
 Reexport = "1.2.2"
 StaticArrays = "1.6.5"
+Statistics = "1.9"
 WriteVTK = "1.18.0"
+julia = "1.9"
 
 [extras]
 CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
diff --git a/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl b/miniapps/benchmarks/stokes2D/shear_band/ShearBand2D.jl
@@ -100,7 +100,7 @@ function main(igg; nx=64, ny=64, figdir="model_figs")
     # Rheology
     η         = @ones(ni...)
     η_vep     = similar(η) # effective visco-elasto-plastic viscosity
-    compute_viscosity!(
+    @parallel (@idx ni) compute_viscosity!(
         η, 1.0, phase_ratios.center, stokes.ε.xx, stokes.ε.yy, stokes.ε.xy, args, rheology, (-Inf, Inf)
     )
 

diff --git a/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl b/miniapps/benchmarks/stokes3D/shear_band/ShearBand3D.jl
@@ -213,7 +213,7 @@ n            = 64 + 2
 nx = ny = nz = n - 2
 figdir       = "ShearBand3D"
 igg          = if !(JustRelax.MPI.Initialized())
-    IGG(init_global_grid(nx, ny, 0; init_MPI = true)...)
+    IGG(init_global_grid(nx, ny, nz; init_MPI = true)...)
 else
     igg
 end

diff --git a/miniapps/convection/GlobalConvection3D.jl b/miniapps/convection/GlobalConvection3D.jl
@@ -182,7 +182,9 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th
     η               = @ones(ni...)
     depth           = PTArray([abs(z) for x in xci[1], y in xci[2], z in xci[3]])
     args            = (; T = thermal.Tc, P = stokes.P, depth = depth, dt = Inf)
-    @parallel (@idx ni) compute_viscosity!(η, 1, 1e-15, args, rheology)
+    @parallel (@idx ni) compute_viscosity!(
+        η, 1.0,  @strain(stokes)..., args, rheology, (1e18, 1e24)
+    )
     η_vep           = deepcopy(η)
     # Boundary conditions
     flow_bcs = FlowBoundaryConditions(; 
@@ -225,7 +227,7 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th
         args = (; T=thermal.Tc, P=stokes.P, depth=depth, dt=Inf)
 
         # Stokes solver ----------------
-        iters = solve!(
+        solve!(
             stokes,
             pt_stokes,
             di,
@@ -299,15 +301,15 @@ function thermal_convection3D(; ar=8, nz=16, nx=ny*8, ny=nx, figdir="figs3D", th
     return (ni=ni, xci=xci, li=li, di=di), thermal
 end
 
-function run()
+# function run()
     figdir = "figs3D_test"
     ar     = 3 # aspect ratio
     n      = 32
     nx     = n*ar - 2
     ny     = nx
     nz     = n - 2
 
-    thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz);
-end
+#     thermal_convection3D(; figdir=figdir, ar=ar,nx=nx, ny=ny, nz=nz);
+# end
 
-run()
+# run()
diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl
@@ -1,5 +1,5 @@
-# using CUDA
-# CUDA.allowscalar(false) # for safety
+using CUDA
+CUDA.allowscalar(false) # for safety
 
 using JustRelax, JustRelax.DataIO, JustPIC
 import JustRelax.@cell
@@ -9,12 +9,11 @@ import JustRelax.@cell
 # set_backend("CUDA_Float64_2D")
 
 # setup ParallelStencil.jl environment
-# model = PS_Setup(:CUDA, Float64, 2)
-model = PS_Setup(:Threads, Float64, 2)
+model = PS_Setup(:CUDA, Float64, 2)
 environment!(model)
 
 # Load script dependencies
-using Printf, LinearAlgebra, GeoParams, GLMakie, CellArrays
+using Printf, LinearAlgebra, GeoParams, GLMakie, SpecialFunctions, CellArrays
 
 # Load file with all the rheology configurations
 include("Layered_rheology.jl")
@@ -90,7 +89,7 @@ end
 @parallel_indices (i, j) function init_T!(T, z)
     depth = -z[j]
 
-    # (depth) because we have 15km of sticky air
+    # (depth - 15e3) because we have 15km of sticky air
     if depth < 0e0
         T[i + 1, j] = 273e0
 
@@ -201,14 +200,14 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     # Rheology
     η                = @zeros(ni...)
     args             = (; T = thermal.Tc, P = stokes.P, dt = Inf)
-    compute_viscosity!(
+    @parallel (@idx ni) compute_viscosity!(
         η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e16, 1e24)
     )
     η_vep            = copy(η)
 
     # PT coefficients for thermal diffusion
     pt_thermal       = PTThermalCoeffs(
-        rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2
+        rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL=1e-3 / √2.1
     )
 
     # Boundary conditions
@@ -221,9 +220,9 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     # if it does not exist, make folder where figures are stored
     if save_vtk
         vtk_dir      = figdir*"\\vtk"
-        !isdir(vtk_dir) && mkpath(vtk_dir)
+        take(vtk_dir)
     end
-    !isdir(figdir) && mkpath(figdir)
+    take(figdir)
     # ----------------------------------------------------
 
     # Plot initial T and η profiles
@@ -261,7 +260,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
         # Update buoyancy and viscosity -
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
-        compute_viscosity!(
+        @parallel (@idx ni) compute_viscosity!(
             η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (1e18, 1e24)
         )
         @parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args)
@@ -407,7 +406,7 @@ end
 figdir   = "Plume2D"
 save_vtk = false # set to true to generate VTK files for ParaView
 ar       = 1 # aspect ratio
-n        = 128
+n        = 64
 nx       = n*ar - 2
 ny       = n - 2
 igg      = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid
@@ -417,4 +416,4 @@ else
 end
 
 # run main script
-main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, save_vtk = save_vtk);
+# main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, save_vtk = save_vtk);
diff --git a/miniapps/convection/Particles2D/Layered_rheology.jl b/miniapps/convection/Particles2D/Layered_rheology.jl
@@ -90,18 +90,18 @@ function init_rheologies(; is_plastic = true)
             Elasticity        = el_lithospheric_mantle,
         ),
         # Name              = "SubLithosphericMantle",
-        SetMaterialParams(;
-            Phase             = 4,
-            Density           = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5),
-            HeatCapacity      = ConstantHeatCapacity(; cp=1.25e3),
-            Conductivity      = K_mantle,
-            RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
-            CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)),
-            Elasticity        = el_sublithospheric_mantle,
-        ),
+        # SetMaterialParams(;
+        #     Phase             = 4,
+        #     Density           = PT_Density(; ρ0=3.3e3, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5),
+        #     HeatCapacity      = ConstantHeatCapacity(; cp=1.25e3),
+        #     Conductivity      = K_mantle,
+        #     RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
+        #     CompositeRheology = CompositeRheology((disl_sublithospheric_mantle, diff_sublithospheric_mantle, el_sublithospheric_mantle)),
+        #     Elasticity        = el_sublithospheric_mantle,
+        # ),
         # Name              = "Plume",
         SetMaterialParams(;
-            Phase             = 5,
+            Phase             = 4,
             Density           = PT_Density(; ρ0=3.3e3-50, β=β_sublithospheric_mantle, T0=0.0, α = 3e-5),
             HeatCapacity      = ConstantHeatCapacity(; cp=1.25e3),
             Conductivity      = K_mantle,
@@ -111,7 +111,7 @@ function init_rheologies(; is_plastic = true)
         ),
         # Name              = "StickyAir",
         SetMaterialParams(;
-            Phase             = 6,
+            Phase             = 5,
             Density           = ConstantDensity(; ρ=1e3),
             HeatCapacity      = ConstantHeatCapacity(; cp=1.25e3),
             RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
@@ -125,7 +125,7 @@ function init_phases!(phases, particles, Lx; d=650e3, r=50e3)
     ni = size(phases)
 
     @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx)
-        @inbounds for ip in JustRelax.JustRelax.cellaxes(phases)
+        @inbounds for ip in JustRelax.cellaxes(phases)
             # quick escape
             JustRelax.@cell(index[ip, i, j]) == 0 && continue
 
@@ -143,14 +143,14 @@ function init_phases!(phases, particles, Lx; d=650e3, r=50e3)
             elseif depth > 90e3
                 @cell phases[ip, i, j] = 3.0
 
-            elseif 0e0 > depth 
-                @cell phases[ip, i, j] = 6.0
+            elseif depth < 0e0 
+                @cell phases[ip, i, j] = 5.0
 
             end
 
             # plume - rectangular
             if ((x - Lx * 0.5)^2 ≤ r^2) && ((depth - d)^2 ≤ r^2)
-                JustRelax.@cell phases[ip, i, j] = 5.0
+                JustRelax.@cell phases[ip, i, j] = 4.0
             end
         end
         return nothing
-Original file line number
+Diff line change
@@ Expand Up / @@ -15,6 +15,7 @@ jobs: @@
           matrix:
             version:
               - '1.9'
+              - '1.10'
               # - 'nightly'
             os:
               - ubuntu-latest
@@ Expand Down @@