PTsolvers · albert-de-montserrat · Mar 7, 2024 · Mar 4, 2024 · Mar 4, 2024 · Mar 5, 2024
diff --git a/Blackenbach.jl b/Blackenbach.jl
diff --git a/Blackenbach_Rheology.jl b/Blackenbach_Rheology.jl
@@ -0,0 +1,36 @@
+# 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=4000.0,α = 2.5e-5, β = 0.0),
+            HeatCapacity      = ConstantHeatCapacity(; Cp=1250.0),
+            Conductivity      = ConstantConductivity(;k=5.0),
+            # CompositeRheology = CompositeRheology((LinearViscous(; η=1e23),)),
+            CompositeRheology = CompositeRheology((LinearViscous(; η=1e21),)),
+            RadioactiveHeat   = ConstantRadioactiveHeat(0.0),
+            Gravity           = ConstantGravity(; g=10.0),
+        ),        
+    )
+end
+
+function init_phases!(phases, particles, Lx, d, r, thick_air)
+    ni = size(phases)
+
+    @parallel_indices (i, j) function init_phases!(phases, px, py, index, r, Lx)
+        @inbounds for ip in JustRelax.cellaxes(phases)
+            # quick escape
+            JustRelax.@cell(index[ip, i, j]) == 0 && continue
+
+            @cell phases[ip, i, j] = 1.0
+
+        end
+        return nothing
+    end
+
+    @parallel (JustRelax.@idx ni) init_phases!(phases, particles.coords..., particles.index, r, Lx)
+end
diff --git a/miniapps/convection/Particles2D/Layered_convection2D.jl b/miniapps/convection/Particles2D/Layered_convection2D.jl
@@ -96,7 +96,7 @@ end
 function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     # Physical domain ------------------------------------
-    thick_air    = 10e3              # thickness of sticky air layer
+    thick_air    = 0                 # thickness of sticky air layer
     ly           = 700e3 + thick_air # domain length in y
     lx           = ly * ar           # domain length in x
     ni           = nx, ny            # number of cells
@@ -114,14 +114,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     # ----------------------------------------------------
 
     # Initialize particles -------------------------------
-    nxcell, max_xcell, min_xcell = 20, 40, 1
-    particles = init_particles(
+    nxcell, max_xcell, min_xcell = 25, 30, 8
+    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))
+    pT, pPhases      = init_cell_arrays(particles, Val(2))
     particle_args    = (pT, pPhases)
 
     # Elliptical temperature anomaly
@@ -136,7 +137,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     # STOKES ---------------------------------------------
     # Allocate arrays needed for every Stokes problem
     stokes           = StokesArrays(ni, ViscoElastic)
-    pt_stokes        = PTStokesCoeffs(li, di; ϵ=1e-4,  CFL = 0.25 / √2.1)
+    pt_stokes        = PTStokesCoeffs(li, di; ϵ=1e-4,  CFL = 0.75 / √2.1)
     # ----------------------------------------------------
 
     # TEMPERATURE PROFILE --------------------------------
@@ -148,7 +149,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     # initialize thermal profile - Half space cooling
     @parallel (@idx ni .+ 1) init_T!(thermal.T, xvi[2], thick_air)
     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)
     # ----------------------------------------------------
@@ -169,7 +170,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     # PT coefficients for thermal diffusion
     pt_thermal       = PTThermalCoeffs(
-        rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1
+        rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-2 / √2.1
     )
 
     # Boundary conditions
@@ -206,6 +207,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 
     T_buffer    = @zeros(ni.+1)
     Told_buffer = similar(T_buffer)
+    dt₀         = similar(stokes.P)
     for (dst, src) in zip((T_buffer, Told_buffer), (thermal.T, thermal.Told))
         copyinn_x!(dst, src)
     end
@@ -218,9 +220,15 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
     end
     # Time loop
     t, it = 0.0, 0
-    while it < 50 # run only for 5 Myrs
-        # while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs
+    while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs
 
+        # interpolate fields from particle to grid vertices
+        particle2grid!(T_buffer, pT, xvi, particles)
+        @views T_buffer[:, end]      .= 273.0
+        @views T_buffer[:, 1]        .= Tbot
+        @views thermal.T[2:end-1, :] .= T_buffer
+        temperature2center!(thermal)
+
         # Update buoyancy and viscosity -
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
         @parallel (@idx ni) compute_viscosity!(
@@ -251,12 +259,6 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
         dt   = compute_dt(stokes, di, dt_diff)
         # ------------------------------
 
-        # interpolate fields from particle to grid vertices
-        particle2grid!(T_buffer, pT, xvi, particles)
-        @views T_buffer[:, end]      .= 273.0
-        @views thermal.T[2:end-1, :] .= T_buffer
-        temperature2center!(thermal)
-
         # Thermal solver ---------------
         heatdiffusion_PT!(
             thermal,
@@ -272,18 +274,23 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
             nout    = 1e2,
             verbose = true,
         )
+        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
+        )
         # ------------------------------
 
         # 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)
@@ -294,7 +301,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
         t        += dt
 
         # Data I/O and plotting ---------------------
-        if it == 1 || rem(it, 10) == 0
+        if it == 1 || rem(it, 1) == 0
             checkpointing(figdir, stokes, thermal.T, η, t)
 
             if save_vtk
@@ -329,6 +336,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
             pxv      = ppx.data[:]./1e3
             pyv      = ppy.data[:]./1e3
             clr      = pPhases.data[:]
+            clr      = pT.data[:]
             idxv     = particles.index.data[:];
 
             # Make Makie figure
@@ -364,12 +372,11 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", save_vtk =false)
 end
 ## END OF MAIN SCRIPT ----------------------------------------------------------------
 
-
 # (Path)/folder where output data and figures are stored
 figdir   = "Plume2D"
 save_vtk = false # set to true to generate VTK files for ParaView
 ar       = 1 # aspect ratio
-n        = 256
+n        = 128
 nx       = n*ar - 2
 ny       = n - 2
 igg      = if !(JustRelax.MPI.Initialized()) # initialize (or not) MPI grid
@@ -380,3 +387,4 @@ end
 
 # run main script
 main2D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, save_vtk = save_vtk);
+
diff --git a/miniapps/convection/Particles3D/Layered_convection3D.jl b/miniapps/convection/Particles3D/Layered_convection3D.jl
@@ -82,12 +82,12 @@ end
 function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
 
     # Physical domain ------------------------------------
-    lz            = 700e3                # domain length in z
-    lx = ly       = lz * ar              # domain length in x and y
-    ni            = nx, ny, nz           # number of cells
-    li            = lx, ly, lz           # domain length
-    di            = @. li / ni           # grid steps
-    origin        = 0.0, 0.0, -lz        # origin coordinates (15km of sticky air layer)
+    lz           = 700e3                # domain length in z
+    lx = ly      = lz * ar              # domain length in x and y
+    ni           = nx, ny, nz           # number of cells
+    li           = lx, ly, lz           # domain length
+    di           = @. li / ni           # grid steps
+    origin       = 0.0, 0.0, -lz        # origin coordinates (15km of sticky air layer)
     grid         = Geometry(ni, li; origin = origin)
     (; xci, xvi) = grid # nodes at the center and vertices of the cells
     # ----------------------------------------------------
@@ -99,10 +99,11 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
     # ----------------------------------------------------
 
     # Initialize particles -------------------------------
-    nxcell, max_xcell, min_xcell = 20, 20, 1
-    particles = init_particles(
+    nxcell, max_xcell, min_xcell = 25, 35, 8
+    particles                   = init_particles(
         backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni...
     )
+    subgrid_arrays              = SubgridDiffusionCellArrays(particles)
     # velocity grids
     grid_vx, grid_vy, grid_vz   = velocity_grids(xci, xvi, di)
     # temperature
@@ -113,7 +114,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
     xc_anomaly       = lx/2   # origin of thermal anomaly
     yc_anomaly       = ly/2   # origin of thermal anomaly
     zc_anomaly       = -610e3 # origin of thermal anomaly
-    r_anomaly        = 25e3   # radius of perturbation
+    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))
     @parallel (@idx ni) phase_ratios_center(phase_ratios.center, particles.coords, xci, di, pPhases)
@@ -133,9 +134,8 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
     )
     # initialize thermal profile - Half space cooling
     @parallel init_T!(thermal.T, xvi[3])
-    thermal_bcs!(thermal.T, thermal_bc)
-
     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)
     # ----------------------------------------------------
 
@@ -193,6 +193,7 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
     end
 
     grid2particle!(pT, xvi, thermal.T, particles)
+    dt₀         = similar(stokes.P)
 
     local Vx_v, Vy_v, Vz_v
     if do_vtk
@@ -203,6 +204,11 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
     # Time loop
     t, it = 0.0, 0
     while (t/(1e6 * 3600 * 24 *365.25)) < 5 # run only for 5 Myrs
+
+        # interpolate fields from particle to grid vertices
+        particle2grid!(thermal.T, pT, xvi, particles)
+        temperature2center!(thermal)
+
         # Update buoyancy and viscosity -
         args = (; T = thermal.Tc, P = stokes.P,  dt=Inf)
         @parallel (@idx ni) compute_viscosity!(
@@ -228,14 +234,10 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
             nout             = 1e3,
             viscosity_cutoff = (1e18, 1e24)
         );
-        @parallel (JustRelax.@idx ni) tensor_invariant!(stokes.ε.II, @strain(stokes)...)
+        @parallel (JustRelax.@idx ni) JustRelax.Stokes3D.tensor_invariant!(stokes.ε.II, @strain(stokes)...)
         dt   = compute_dt(stokes, di, dt_diff) / 2
         # ------------------------------
 
-        # interpolate fields from particle to grid vertices
-        particle2grid!(thermal.T, pT, xvi, particles)
-        temperature2center!(thermal)
-
         # Thermal solver ---------------
         heatdiffusion_PT!(
             thermal,
@@ -251,15 +253,20 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
             nout    = 1e2,
             verbose = true,
         )
+        subgrid_characteristic_time!(
+            subgrid_arrays, particles, dt₀, phase_ratios, rheology, thermal, stokes, xci, di
+        )
+        centroid2particle!(subgrid_arrays.dt₀, xci, dt₀, particles)
+        subgrid_diffusion!(
+            pT, thermal.T, thermal.ΔT, subgrid_arrays, particles, xvi,  di, dt
+        )
         # ------------------------------
 
         # Advection --------------------
         # advect particles in space
         advection_RK!(particles, @velocity(stokes), grid_vx, grid_vy, grid_vz, dt, 2 / 3)
         # 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)
@@ -308,9 +315,9 @@ function main3D(igg; ar=1, nx=16, ny=16, nz=16, figdir="figs3D", do_vtk =false)
             ax3 = Axis(fig[1,3], aspect = ar, title = "log10(εII)")
             ax4 = Axis(fig[2,3], aspect = ar, title = "log10(η)")
             # Plot temperature
-            h1  = heatmap!(ax1, xvi[1].*1e-3, xvi[3].*1e-3, Array(thermal.T[:, slice_j, :]) , colormap=:batlow)
+            h1  = heatmap!(ax1, xvi[1].*1e-3, xvi[3].*1e-3, Array(thermal.T[:, slice_j, :]) , colormap=:lajolla)
             # Plot particles phase
-            h2  = heatmap!(ax2, xci[1].*1e-3, xci[3].*1e-3, Array(stokes.τ.II[:, slice_j, :].*1-e6) , colormap=:batlow)
+            h2  = heatmap!(ax2, xci[1].*1e-3, xci[3].*1e-3, Array(stokes.τ.II[:, slice_j, :].*1e-6) , colormap=:batlow)
             # Plot 2nd invariant of strain rate
             h3  = heatmap!(ax3, xci[1].*1e-3, xci[3].*1e-3, Array(log10.(stokes.ε.II[:, slice_j, :])) , colormap=:batlow)
             # Plot effective viscosity
@@ -335,7 +342,7 @@ end
 
 do_vtk   = true # set to true to generate VTK files for ParaView
 ar       = 1 # aspect ratio
-n        = 128
+n        = 32
 nx       = n
 ny       = n
 nz       = n
@@ -347,4 +354,4 @@ end
 
 # (Path)/folder where output data and figures are stored
 figdir   = "Plume3D_$n"
-main3D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, do_vtk = do_vtk);
+main3D(igg; figdir = figdir, ar = ar, nx = nx, ny = ny, nz = nz, do_vtk = do_vtk);
diff --git a/src/MetaJustRelax.jl b/src/MetaJustRelax.jl
@@ -124,21 +124,24 @@ function environment!(model::PS_Setup{T,N}) where {T,N}
         include(joinpath(@__DIR__, "rheology/Viscosity.jl"))
         export compute_viscosity!
 
-        include(joinpath(@__DIR__, "stokes/StressKernels.jl"))
-        export tensor_invariant!
+        # include(joinpath(@__DIR__, "stokes/StressKernels.jl"))
+        # export tensor_invariant!
 
         include(joinpath(@__DIR__, "stokes/Stokes2D.jl"))
-        export solve!
+        export solve!, tensor_invariant!
 
         include(joinpath(@__DIR__, "stokes/Stokes3D.jl"))
-        export solve!
+        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!