Skip to content

Commit

Permalink
massive rework
Browse files Browse the repository at this point in the history
  • Loading branch information
@albert-de-montserrat
albert-de-montserrat committed Apr 2, 2024
1 parent cb858ad commit da5a3e8
Show file tree
Hide file tree
Showing 27 changed files with 3,149 additions and 1,257 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ HDF5 = "0.17.1"
ImplicitGlobalGrid = "0.14.0, 0.15.0"
MPI = "0.20"
MuladdMacro = "0.2"
ParallelStencil = "0.9, 0.10, 0.11"
ParallelStencil = "0.12"
Reexport = "1.2.2"
StaticArrays = "1"
Statistics = "1"
Expand Down
86 changes: 40 additions & 46 deletions miniapps/benchmarks/stokes2D/shear_heating/Shearheating2D.jl
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
# Benchmark of Duretz et al. 2014
# http://dx.doi.org/10.1002/2014GL060438
using JustRelax, JustRelax.DataIO
import JustRelax.@cell
using ParallelStencil
using JustRelax, JustRelax.JustRelax2D
using JustRelax.DataIO
# import JustRelax.@cell
using ParallelStencil, ParallelStencil.FiniteDifferences2D
@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2)

using JustPIC
Expand All @@ -13,8 +14,8 @@ using JustPIC._2D
const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend

# setup ParallelStencil.jl environment
model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2)
environment!(model)
# model = PS_Setup(:cpu, Float64, 2) #or (:CUDA, Float64, 2) or (:AMDGPU, Float64, 2)
# environment!(model)

# Load script dependencies
using Printf, LinearAlgebra, GeoParams, CellArrays
Expand All @@ -23,8 +24,7 @@ using GLMakie
# Load file with all the rheology configurations
include("Shearheating_rheology.jl")


# ## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT --------------------------------
## SET OF HELPER FUNCTIONS PARTICULAR FOR THIS SCRIPT --------------------------------

function copyinn_x!(A, B)

Expand Down Expand Up @@ -72,7 +72,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
# ----------------------------------------------------

# Initialize particles -------------------------------
nxcell, max_xcell, min_xcell = 20, 40, 1
nxcell, max_xcell, min_xcell = 20, 40, 10
particles = init_particles(
backend, nxcell, max_xcell, min_xcell, xvi..., di..., ni...
)
Expand All @@ -89,64 +89,58 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
r_anomaly = 3e3 # radius of perturbation
init_phases!(pPhases, particles, lx/2, yc_anomaly, r_anomaly)
phase_ratios = PhaseRatio(ni, length(rheology))
@parallel (@idx ni) phase_ratios_center(phase_ratios.center, pPhases)
phase_ratios_center(phase_ratios, particles, grid, pPhases)
# ----------------------------------------------------

# STOKES ---------------------------------------------
# Allocate arrays needed for every Stokes problem
stokes = StokesArrays(ni, ViscoElastic)
stokes = StokesArrays(ni)
pt_stokes = PTStokesCoeffs(li, di; ϵ=1e-4, CFL = 0.9 / √2.1)
# ----------------------------------------------------

# TEMPERATURE PROFILE --------------------------------
thermal = ThermalArrays(ni)
thermal_bc = TemperatureBoundaryConditions(;
no_flux = (left = true, right = true, top = false, bot = false),
periodicity = (left = false, right = false, top = false, bot = false),
)

# Initialize constant temperature
@views thermal.T .= 273.0 + 400
thermal_bcs!(thermal.T, thermal_bc)

@parallel (JustRelax.@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T)
temperature2center!(thermal)
# ----------------------------------------------------

# Buoyancy forces
ρg = @zeros(ni...), @zeros(ni...)

@parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, (T=thermal.Tc, P=stokes.P))
ρg = @zeros(ni...), @zeros(ni...)
compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P))
@parallel init_P!(stokes.P, ρg[2], xci[2])

# Rheology
η = @ones(ni...)
args = (; T = thermal.Tc, P = stokes.P, dt = Inf)
@parallel (@idx ni) compute_viscosity!(
η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf)
args = (; T = thermal.Tc, P = stokes.P, dt = Inf)
compute_viscosity!(
stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)
)
η_vep = copy(η)

# PT coefficients for thermal diffusion
pt_thermal = PTThermalCoeffs(
pt_thermal = PTThermalCoeffs(
rheology, phase_ratios, args, dt, ni, di, li; ϵ=1e-5, CFL= 1e-3 / √2.1
)

# Boundary conditions
flow_bcs = FlowBoundaryConditions(;
free_slip = (left = true, right=true, top=true, bot=true),
periodicity = (left = false, right = false, top = false, bot = false),
flow_bcs = FlowBoundaryConditions(;
free_slip = (left = true, right=true, top=true, bot=true),
)
## Compression and not extension - fix this
εbg = 5e-14
stokes.V.Vx .= PTArray([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2])
stokes.V.Vy .= PTArray([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]])
εbg = 5e-14
stokes.V.Vx .= PTArray()([ -(x - lx/2) * εbg for x in xvi[1], _ in 1:ny+2])
stokes.V.Vy .= PTArray()([ (ly - abs(y)) * εbg for _ in 1:nx+2, y in xvi[2]])
flow_bcs!(stokes, flow_bcs) # apply boundary conditions
update_halo!(stokes.V.Vx, stokes.V.Vy)

# IO ----- -------------------------------------------
# if it does not exist, make folder where figures are stored
if do_vtk
vtk_dir = figdir*"\\vtk"
vtk_dir = joinpath(figdir, "vtk")
take(vtk_dir)
end
take(figdir)
Expand All @@ -160,7 +154,7 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
ax1 = Axis(fig[1,1], aspect = 2/3, title = "T")
ax2 = Axis(fig[1,2], aspect = 2/3, title = "log10(η)")
scatter!(ax1, Array(thermal.T[2:end-1,:][:]), Yv./1e3)
scatter!(ax2, Array(log10.(η[:])), Y./1e3)
scatter!(ax2, Array(log10.(stokes.viscosity.η[:])), Y./1e3)
ylims!(ax1, minimum(xvi[2])./1e3, 0)
ylims!(ax2, minimum(xvi[2])./1e3, 0)
hideydecorations!(ax2)
Expand All @@ -185,10 +179,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
while it < 100
# Update buoyancy and viscosity -
args = (; T = thermal.Tc, P = stokes.P, dt=Inf)
@parallel (@idx ni) compute_viscosity!(
η, 1.0, phase_ratios.center, @strain(stokes)..., args, rheology, (-Inf, Inf)
compute_ρg!(ρg[2], phase_ratios, rheology, (T=thermal.Tc, P=stokes.P))
compute_viscosity!(
stokes, 1.0, phase_ratios, args, rheology, (-Inf, Inf)
)
@parallel (JustRelax.@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args)
# ------------------------------

# Stokes solver ----------------
Expand All @@ -198,18 +192,18 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
di,
flow_bcs,
ρg,
η,
η_vep,
phase_ratios,
rheology,
args,
dt,
igg;
iterMax = 75e3,
nout=1e3,
viscosity_cutoff=(-Inf, Inf)
kwargs = (;
iterMax = 75e3,
nout=1e3,
viscosity_cutoff=(-Inf, Inf)
)
)
@parallel (JustRelax.@idx ni) JustRelax.Stokes2D.tensor_invariant!(stokes.ε.II, @strain(stokes)...)
tensor_invariant!(stokes.ε)
dt = compute_dt(stokes, di, dt_diff)
# ------------------------------

Expand All @@ -219,12 +213,10 @@ function main2D(igg; ar=8, ny=16, nx=ny*8, figdir="figs2D", do_vtk =false)
@views thermal.T[2:end-1, :] .= T_buffer
temperature2center!(thermal)

@parallel (@idx ni) compute_shear_heating!(
thermal.shear_heating,
@tensor_center(stokes.τ),
@tensor_center(stokes.τ_o),
@strain(stokes),
phase_ratios.center,
compute_shear_heating!(
thermal,
stokes,
phase_ratios,
rheology, # needs to be a tuple
dt,
)
Expand Down Expand Up @@ -347,4 +339,6 @@ else
igg
end

main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk)
# main2D(igg; ar=ar, ny=ny, nx=nx, figdir=figdir, do_vtk=do_vtk)

foo(args...; x=1) = x
4 changes: 2 additions & 2 deletions miniapps/benchmarks/stokes2D/shear_heating/Shearheating_rheology.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
29 changes: 23 additions & 6 deletions src/Interpolations.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@ end
# From cell vertices to cell center

function temperature2center!(thermal::ThermalArrays)
@parallel (@idx size(thermal.Tc)...) temperature2center!(thermal.Tc, thermal.T)
@parallel (@idx size(thermal.Tc)...) temperature2center_kernel!(thermal.Tc, thermal.T)
return nothing
end

@parallel_indices (i, j) function temperature2center!(
@parallel_indices (i, j) function temperature2center_kernel!(
T_center::T, T_vertex::T
) where {T<:AbstractArray{_T,2} where {_T<:Real}}
T_center[i, j] =
Expand All @@ -34,7 +34,7 @@ end
return nothing
end

@parallel_indices (i, j, k) function temperature2center!(
@parallel_indices (i, j, k) function temperature2center_kernel!(
T_center::T, T_vertex::T
) where {T<:AbstractArray{_T,3} where {_T<:Real}}
@inline av_T() = _av(T_vertex, i, j, k)
Expand All @@ -44,17 +44,34 @@ end
return nothing
end

@parallel function vertex2center!(center, vertex)
function vertex2center!(center, vertex)
@parallel vertex2center_kernel!(center, vertex)
return nothing
end

@parallel function vertex2center_kernel!(center, vertex)
@all(center) = @av(vertex)
return nothing
end

@parallel function center2vertex!(vertex, center)
function center2vertex!(vertex, center)
@parallel center2vertex_kernel!(vertex, center)
return nothing
end

@parallel function center2vertex_kernel!(vertex, center)
@inn(vertex) = @av(center)
return nothing
end

@parallel_indices (i, j, k) function center2vertex!(
function center2vertex!(vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy)
@parallel center2vertex_kernel!(
vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy
)
return nothing
end

@parallel_indices (i, j, k) function center2vertex_kernel!(
vertex_yz, vertex_xz, vertex_xy, center_yz, center_xz, center_xy
)
i1, j1, k1 = (i, j, k) .+ 1
Expand Down
55 changes: 11 additions & 44 deletions src/JustRelax.jl
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
module JustRelax

using Reexport
@reexport using ParallelStencil
@reexport using ImplicitGlobalGrid
using LinearAlgebra
using Printf
Expand All @@ -13,60 +12,28 @@ using StaticArrays

function solve!() end

include("types/traits.jl")

include("topology/Topology.jl")
export IGG, lazy_grid, Geometry, velocity_grids, x_g, y_g, z_g

include("MiniKernels.jl")
export _d_xa,
_d_ya,
_d_za,
_d_za,
_d_xi,
_d_yi,
_d_zi,
_d_zi,
_av,
_av_a,
_av_xa,
_av_ya,
_av_x,
_av_y,
_av_z,
_av_yz,
_av_xz,
_av_xy,
_av_yzi,
_av_xzi,
_av_xyi,
_gather,
_gather_yz,
_gather_xz,
_gather_xy,
_harm_x,
_harm_y,
_harm_z,
_harm_yz,
_harm_xz,
_harm_xy,
_harm_xyi,
_harm_xzi,
_harm_yzi,
_current
# include("MiniKernels.jl")

include("phases/CellArrays.jl")
export @cell, element, setelement!, cellnum, cellaxes, new_empty_cell, setindex!

include("rheology/StressUpdate.jl")
export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear!
# include("rheology/StressUpdate.jl")
# export plastic_params, plastic_params_phase, compute_dτ_r, _compute_τ_nonlinear!

include("MetaJustRelax.jl")
include("JustRelax_CPU.jl")
# include("MetaJustRelax.jl")

include("stokes/MetaStokes.jl")
export PS_Setup, environment!, ps_reset!
# include("stokes/MetaStokes.jl")
# export PS_Setup, environment!, ps_reset!

include("thermal_diffusion/MetaDiffusion.jl")
# include("thermal_diffusion/MetaDiffusion.jl")

include("thermal_diffusion/Rheology.jl")
# include("thermal_diffusion/Rheology.jl")

include("IO/DataIO.jl")

Expand Down
Loading

0 comments on commit da5a3e8

Please sign in to comment.