Skip to content

Commit

Permalink
up MPI test & MPI 3D miniapp
Browse files Browse the repository at this point in the history
  • Loading branch information
aelligp committed Nov 12, 2024
1 parent 8e5a7a2 commit c725f41
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 87 deletions.
121 changes: 59 additions & 62 deletions miniapps/benchmarks/stokes3D/shear_band/ShearBand3D_MPI.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,25 +15,29 @@ const backend = JustPIC.CPUBackend
solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η))

# Initialize phases on the particles
function init_phases!(phase_ratios, xci, xvi, radius)
ni = size(phase_ratios.center)
function init_phases!(phases, particles, radius)
ni = size(phases)
origin = 0.5, 0.5, 0.5

@parallel_indices (i, j, k) function init_phases!(phases, xc, yc, zc, o_x, o_y, o_z)
x, y, z = xc[i], yc[j], zc[k]
if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius^2
@index phases[1, i, j, k] = 1.0
@index phases[2, i, j, k] = 0.0
@parallel_indices (I...) function init_phases!(phases, index, xc, yc, zc, o_x, o_y, o_z)

else
@index phases[1, i, j, k] = 0.0
@index phases[2, i, j, k] = 1.0
for ip in cellaxes(xc)
(@index index[ip, I...]) || continue

x = @index xc[ip, I...]
y = @index yc[ip, I...]
z = @index zc[ip, I...]

if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius^2
@index phases[ip, I...] = 1.0
else
@index phases[ip, I...] = 2.0
end
end
return nothing
end

@parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin...)
@parallel (@idx ni .+ 1) init_phases!(phase_ratios.vertex, xvi..., origin...)
@parallel (@idx ni) init_phases!(phases, particles.index, particles.coords...,origin...)
return nothing
end

Expand All @@ -56,13 +60,15 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)
η0 = 1.0 # viscosity
G0 = 1.0 # elastic shear modulus
Gi = G0/(6.0-4.0) # elastic shear modulus perturbation
# Gi = G0 # 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
dt /= 2
el_bg = ConstantElasticity(; G=G0, ν=0.5)
el_inc = ConstantElasticity(; G=Gi, ν=0.5)
visc = LinearViscous(; η=η0)
visc_inc = LinearViscous(; η=η0/10)
pl = DruckerPrager_regularised(; # non-regularized plasticity
C = C,
ϕ = ϕ,
Expand All @@ -75,7 +81,6 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)
Phase = 1,
Density = ConstantDensity(; ρ = 0.0),
Gravity = ConstantGravity(; g = 0.0),
# CompositeRheology = CompositeRheology((visc, el_bg, )),
CompositeRheology = CompositeRheology((visc, el_bg, pl)),
Elasticity = el_bg,

Expand All @@ -84,21 +89,26 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)
SetMaterialParams(;
Density = ConstantDensity(; ρ = 0.0),
Gravity = ConstantGravity(; g = 0.0),
# CompositeRheology = CompositeRheology((visc, el_inc, )),
CompositeRheology = CompositeRheology((visc, el_inc, pl)),
Elasticity = el_inc,
),
)

# Initialize phase ratios -------------------------------
radius = 0.1
nxcell, max_xcell, min_xcell = 125, 150, 75
particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi, di, ni)
radius = 0.1
phase_ratios = PhaseRatios(backend, length(rheology), ni)
pPhases, = init_cell_arrays(particles, Val(1))
# Assign particles phases anomaly
init_phases!(pPhases, particles, radius)
phase_ratios = PhaseRatios(backend, length(rheology), ni)
init_phases!(phase_ratios, xci, xvi, radius)
update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases)

# STOKES ---------------------------------------------
# STOKES ---------------------------------------------
# Allocate arrays needed for every Stokes problem
stokes = StokesArrays(backend_JR, ni)
pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-6, CFL = 0.5 / 3.1)
pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-5, CFL = 0.5 / 3.1)
# Buoyancy forces
ρg = @zeros(ni...), @zeros(ni...), @zeros(ni...)
args = (; T = @zeros(ni...), P = stokes.P, dt = Inf)
Expand All @@ -112,16 +122,8 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)
no_slip = (left = false, right = false, top = false, bot = false, back = false, front = false),
)

Vx = PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2])
Vz = PTArray(backend_JR)([-z*εbg for _ in 1:nx+2, _ in 1:ny+2, z in xvi[3]])
stokes.V.Vx .= Vx
stokes.V.Vz .= Vz

# println("Rank $(igg.me):
# extrema Vx = $(extrema(Vx))
# extrema Vz = $(extrema(Vz))
# ")

stokes.V.Vx .= PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2])
stokes.V.Vz .= PTArray(backend_JR)([-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!(@velocity(stokes)...)

Expand Down Expand Up @@ -206,18 +208,16 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)

igg.me == 0 && println("it = $it; t = $t \n")

if do_vtk
velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)
vertex2center!(Vx, Vx_v)
vertex2center!(Vy, Vy_v)
vertex2center!(Vz, Vz_v)
@views Vx_nohalo .= Array(Vx[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vy_nohalo .= Array(Vy[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vz_nohalo .= Array(Vz[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(Vx_nohalo, Vxv_v)
gather!(Vy_nohalo, Vyv_v)
gather!(Vz_nohalo, Vzv_v)
end
velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)
vertex2center!(Vx, Vx_v)
vertex2center!(Vy, Vy_v)
vertex2center!(Vz, Vz_v)
@views Vx_nohalo .= Array(Vx[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vy_nohalo .= Array(Vy[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vz_nohalo .= Array(Vz[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(Vx_nohalo, Vxv_v)
gather!(Vy_nohalo, Vyv_v)
gather!(Vz_nohalo, Vzv_v)

# MPI
phase_center = [argmax(p) for p in Array(phase_ratios.center)]
Expand All @@ -234,27 +234,24 @@ function main(igg; nx=64, ny=64, nz=64, figdir="model_figs", do_vtk=false)

if igg.me == 0

if do_vtk == true

data_c = (;
τII = τII_v,
εII = εII_v,
η = η_vep_v,
phases = phases_c_v
)
velocity = (
Array(Vxv_v),
Array(Vyv_v),
Array(Vzv_v),
)
save_vtk(
joinpath(vtk_dir, "vtk_" * lpad("$(it)_$(igg.me)", 6, "0")),
xci_v,
data_c,
velocity,
t
)
end
data_c = (;
τII = τII_v,
εII = εII_v,
η = η_vep_v,
phases = phases_c_v
)
velocity = (
Array(Vxv_v),
Array(Vyv_v),
Array(Vzv_v),
)
save_vtk(
joinpath(vtk_dir, "vtk_" * lpad("$(it)_$(igg.me)", 6, "0")),
xci_v,
data_c,
velocity,
t
)

# visualisation
th = 0:pi/50:3*pi;
Expand Down
58 changes: 33 additions & 25 deletions test/test_shearband3D_MPI.jl
Original file line number Diff line number Diff line change
Expand Up @@ -35,25 +35,29 @@ end
solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η))

# Initialize phases on the particles
function init_phases!(phase_ratios, xci, xvi, radius)
ni = size(phase_ratios.center)
function init_phases!(phases, particles, radius)
ni = size(phases)
origin = 0.5, 0.5, 0.5

@parallel_indices (i, j, k) function init_phases!(phases, xc, yc, zc, o_x, o_y, o_z)
x, y, z = xc[i], yc[j], zc[k]
if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius^2
@index phases[1, i, j, k] = 1.0
@index phases[2, i, j, k] = 0.0
@parallel_indices (I...) function init_phases!(phases, index, xc, yc, zc, o_x, o_y, o_z)

else
@index phases[1, i, j, k] = 0.0
@index phases[2, i, j, k] = 1.0
for ip in cellaxes(xc)
(@index index[ip, I...]) || continue

x = @index xc[ip, I...]
y = @index yc[ip, I...]
z = @index zc[ip, I...]

if ((x-o_x)^2 + (y-o_y)^2 + (z-o_z)^2) > radius^2
@index phases[ip, I...] = 1.0
else
@index phases[ip, I...] = 2.0
end
end
return nothing
end

@parallel (@idx ni) init_phases!(phase_ratios.center, xci..., origin...)
@parallel (@idx ni .+ 1) init_phases!(phase_ratios.vertex, xvi..., origin...)
@parallel (@idx ni) init_phases!(phases, particles.index, particles.coords...,origin...)
return nothing
end

Expand All @@ -76,13 +80,15 @@ function main(igg; nx=64, ny=64, nz=64)
η0 = 1.0 # viscosity
G0 = 1.0 # elastic shear modulus
Gi = G0/(6.0-4.0) # elastic shear modulus perturbation
# Gi = G0 # 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
dt /= 2
el_bg = ConstantElasticity(; G=G0, ν=0.5)
el_inc = ConstantElasticity(; G=Gi, ν=0.5)
visc = LinearViscous(; η=η0)
visc_inc = LinearViscous(; η=η0/10)
pl = DruckerPrager_regularised(; # non-regularized plasticity
C = C,
ϕ = ϕ,
Expand All @@ -95,7 +101,6 @@ function main(igg; nx=64, ny=64, nz=64)
Phase = 1,
Density = ConstantDensity(; ρ = 0.0),
Gravity = ConstantGravity(; g = 0.0),
# CompositeRheology = CompositeRheology((visc, el_bg, )),
CompositeRheology = CompositeRheology((visc, el_bg, pl)),
Elasticity = el_bg,

Expand All @@ -104,21 +109,26 @@ function main(igg; nx=64, ny=64, nz=64)
SetMaterialParams(;
Density = ConstantDensity(; ρ = 0.0),
Gravity = ConstantGravity(; g = 0.0),
# CompositeRheology = CompositeRheology((visc, el_inc, )),
CompositeRheology = CompositeRheology((visc, el_inc, pl)),
Elasticity = el_inc,
),
)

# Initialize phase ratios -------------------------------
radius = 0.1
nxcell, max_xcell, min_xcell = 125, 150, 75
particles = init_particles(backend, nxcell, max_xcell, min_xcell, xvi, di, ni)
radius = 0.1
phase_ratios = PhaseRatios(backend, length(rheology), ni)
pPhases, = init_cell_arrays(particles, Val(1))
# Assign particles phases anomaly
init_phases!(pPhases, particles, radius)
phase_ratios = PhaseRatios(backend, length(rheology), ni)
init_phases!(phase_ratios, xci, xvi, radius)
update_phase_ratios!(phase_ratios, particles, xci, xvi, pPhases)

# STOKES ---------------------------------------------
# STOKES ---------------------------------------------
# Allocate arrays needed for every Stokes problem
stokes = StokesArrays(backend_JR, ni)
pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-6, CFL = 0.5 / 3.1)
pt_stokes = PTStokesCoeffs(li, di; ϵ = 1e-5, CFL = 0.5 / 3.1)
# Buoyancy forces
ρg = @zeros(ni...), @zeros(ni...), @zeros(ni...)
args = (; T = @zeros(ni...), P = stokes.P, dt = Inf)
Expand All @@ -132,11 +142,8 @@ function main(igg; nx=64, ny=64, nz=64)
no_slip = (left = false, right = false, top = false, bot = false, back = false, front = false),
)

Vx = PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2])
Vz = PTArray(backend_JR)([-z*εbg for _ in 1:nx+2, _ in 1:ny+2, z in xvi[3]])
stokes.V.Vx .= Vx
stokes.V.Vz .= Vz

stokes.V.Vx .= PTArray(backend_JR)([ x*εbg for x in xvi[1], _ in 1:ny+2, _ in 1:nz+2])
stokes.V.Vz .= PTArray(backend_JR)([-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!(@velocity(stokes)...)

Expand All @@ -154,7 +161,7 @@ function main(igg; nx=64, ny=64, nz=64)

# Time loop
t, it = 0.0, 0
tmax = 4
tmax = 3.5
τII = Float64[]
sol = Float64[]
ttot = Float64[]
Expand All @@ -180,6 +187,7 @@ function main(igg; nx=64, ny=64, nz=64)
)

tensor_invariant!(stokes.ε)
tensor_invariant!(stokes.ε_pl)
push!(τII, maximum(stokes.τ.xx))

it += 1
Expand Down

0 comments on commit c725f41

Please sign in to comment.