Skip to content

Commit

Permalink
MPI progress
Browse files Browse the repository at this point in the history
  • Loading branch information
@aelligp
aelligp committed Nov 6, 2024
1 parent 1bc962c commit e5c74d7
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 87 deletions.
58 changes: 28 additions & 30 deletions miniapps/subduction/2D/Subduction2D_MPI.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,11 +63,11 @@ end
## END OF HELPER FUNCTION ------------------------------------------------------------

## BEGIN OF MAIN SCRIPT --------------------------------------------------------------
function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false)
function main(x_global, z_global,li, origin, phases_GMG, T_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk =false)

# Physical domain ------------------------------------
ni = nx, ny # number of cells
di = @. li / ni # grid steps
di = @. li / (nx_g(), ny_g()) # grid steps
grid = Geometry(ni, li; origin = origin)
(; xci, xvi) = grid # nodes at the center and vertices of the cells
# ----------------------------------------------------
Expand Down Expand Up @@ -166,9 +166,11 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
if do_vtk
Vx_v = @zeros(ni.+1...)
Vy_v = @zeros(ni.+1...)
Vx = @zeros(ni...)
Vy = @zeros(ni...)
end

#MPI
#MPI
# global array
nx_v = (nx - 2) * igg.dims[1]
ny_v = (ny - 2) * igg.dims[2]
Expand All @@ -185,20 +187,16 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
εII_nohalo = zeros(nx-2, ny-2)
phases_c_nohalo = zeros(nx-2, ny-2)
#vertex
Vxv_v = zeros(nx_v+1, ny_v+1)
Vyv_v = zeros(nx_v+1, ny_v+1)
Vzv_v = zeros(nx_v+1, ny_v+1)
T_v = zeros(nx_v+1, ny_v+1)
phases_v_v = zeros(nx_v+1, ny_v+1)
Vxv_v = zeros(nx_v, ny_v)
Vyv_v = zeros(nx_v, ny_v)
T_v = zeros(nx_v, ny_v)
#vertex nohalo
Vxv_nohalo = zeros(nx-1, ny-1)
Vyv_nohalo = zeros(nx-1, ny-1)
Vzv_nohalo = zeros(nx-1, ny-1)
T_nohalo = zeros(nx-1, ny-1)
phase_v_nohalo = zeros(nx-1, ny-1)
Vxv_nohalo = zeros(nx-2, ny-2)
Vyv_nohalo = zeros(nx-2, ny-2)
Vzv_nohalo = zeros(nx-2, ny-2)
T_nohalo = zeros(nx-2, ny-2)

xci_v = LinRange(minimum(x_global).*1e3, maximum(x_global).*1e3, nx_v), LinRange(minimum(y_global).*1e3, maximum(y_global).*1e3, ny_v)
xvi_v = LinRange(minimum(x_global).*1e3, maximum(x_global).*1e3, nx_v+1), LinRange(minimum(y_global).*1e3, maximum(y_global).*1e3, ny_v+1)
xci_v = LinRange(minimum(x_global).*1e3, maximum(x_global).*1e3, nx_v), LinRange(minimum(z_global).*1e3, maximum(z_global).*1e3, ny_v)


T_buffer = @zeros(ni.+1)
Expand Down Expand Up @@ -351,16 +349,16 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
gather!(phases_c_nohalo, phases_c_v)
#vertices
if do_vtk
velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)
@views Vxv_nohalo .= Array(Vx_v[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vyv_nohalo .= Array(Vy_v[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
# gather!(Vxv_nohalo, Vxv_v)
# gather!(Vyv_nohalo, Vyv_v)
velocity2vertex!(Vx_v, Vy_v, @velocity(stokes)...)
vertex2center!(Vx, Vx_v)
vertex2center!(Vy, Vy_v)
@views Vxv_nohalo .= Array(Vx[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vyv_nohalo .= Array(Vy[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(Vxv_nohalo, Vxv_v)
gather!(Vyv_nohalo, Vyv_v)
end
# @views T_nohalo .= Array(thermal.T[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
# @views phase_v_nohalo .= Array(phase_vertex[2:end-1, 2:end-1, 2:end-1])
# gather!(T_nohalo, T_v)
# gather!(phase_v_nohalo, phases_v_v)
@views T_nohalo .= Array(thermal.Tc[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(T_nohalo, T_v)

# Data I/O and plotting ---------------------
if it == 1 || rem(it, 10) == 0
Expand All @@ -372,6 +370,7 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
# phases_v = phases_v_v,
)
data_c = (;
T = T_v,
P = P_v,
τII = τII_v,
εII = εII_v,
Expand All @@ -385,10 +384,8 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
Array(Vyv_v),
)
save_vtk(
joinpath(vtk_dir, "vtk_" * lpad("$(it)_$(igg.me)", 6, "0")),
xvi_v./1e3,
joinpath(vtk_dir, "vtk_" * lpad("$(it)", 6, "0")),
xci_v./1e3,
data_v,
data_c,
velocity_v
)
Expand All @@ -411,9 +408,10 @@ function main(li, origin, phases_GMG, igg; nx=16, ny=16, figdir="figs2D", do_vtk
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_v[1].*1e-3, xvi_v[2].*1e-3, Array(T_v) , colormap=:batlow)
h1 = heatmap!(ax1, xci_v[1].*1e-3, xci_v[2].*1e-3, Array(T_v) , colormap=:batlow)
# Plot particles phase
h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1)
# h2 = scatter!(ax2, Array(pxv[idxv]), Array(pyv[idxv]), color=Array(clr[idxv]), markersize = 1)
h2 = heatmap!(ax2, xci_v[1].*1e-3, xci_v[2].*1e-3, Array(phases_c_v) , colormap=:batlow)
# Plot 2nd invariant of strain rate
# h3 = heatmap!(ax3, xci_v[1].*1e-3, xci_v[2].*1e-3, Array(log10.(εII)) , colormap=:batlow)
h3 = heatmap!(ax3, xci_v[1].*1e-3, xci_v[2].*1e-3, Array(τII_v) , colormap=:batlow)
Expand Down Expand Up @@ -463,4 +461,4 @@ figdir = "Subduction3D_$(nx_g())x$(ny_g())"

li_GMG, origin_GMG, phases_GMG, T_GMG = GMG_subduction_2D(model_depth, grid_global.xvi,nx+1, ny+1)

# main(x_global, z_global,li_GMG, origin_GMG, phases_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk);
main(x_global, z_global,li_GMG, origin_GMG, phases_GMG, T_GMG, igg; figdir = figdir, nx = nx, ny = ny, do_vtk = do_vtk);
58 changes: 29 additions & 29 deletions miniapps/subduction/2D/Subduction2D_setup.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ function GMG_subduction_2D(nx, ny)
model_depth = 660
# Our starting basis is the example above with ridge and overriding slab
nx, nz = nx, ny
Tbot = 1474.0
Tbot = 1474.0
x = range(0, 3000, nx);
air_thickness = 15.0
z = range(-model_depth, air_thickness, nz);
Expand All @@ -13,71 +13,71 @@ function GMG_subduction_2D(nx, ny)
Temp = fill(Tbot, nx, 1, nz);
Tlab = 1300
# lith = LithosphericPhases(Layers=[80], Phases=[1 0], Tlab=Tlab)

# phases
# 0: asthenosphere
# 1: lithosphere
# 2: subduction lithosphere
# 1: lithosphere
# 2: subduction lithosphere
# 3: oceanic crust
# 4: air
add_box!(
Phases,
Temp,
Grid2D;
Phases,
Temp,
Grid2D;
xlim =(0, 3000),
zlim =(-model_depth, 0.0),
zlim =(-model_depth, 0.0),
Origin = nothing, StrikeAngle=0, DipAngle=0,
phase = LithosphericPhases(Layers=[], Phases=[0], Tlab=Tlab),
phase = LithosphericPhases(Layers=[], Phases=[0], Tlab=Tlab),
T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0)
)

# Add left oceanic plate
add_box!(
Phases,
Temp,
Grid2D;
Phases,
Temp,
Grid2D;
xlim =(100, 3000-100),
zlim =(-model_depth, 0.0),
zlim =(-model_depth, 0.0),
Origin = nothing, StrikeAngle=0, DipAngle=0,
phase = LithosphericPhases(Layers=[80], Phases=[1 0], Tlab=Tlab),
phase = LithosphericPhases(Layers=[80], Phases=[1 0], Tlab=Tlab),
T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0)
)

# Add right oceanic plate crust
add_box!(
Phases,
Temp,
Grid2D;
xlim =(3000-1430, 3000-200),
zlim =(-model_depth, 0.0),
Phases,
Temp,
Grid2D;
xlim =(3000-1430, 3000-200),
zlim =(-model_depth, 0.0),
Origin = nothing, StrikeAngle=0, DipAngle=0,
phase = LithosphericPhases(Layers=[8 72], Phases=[2 1 0], Tlab=Tlab),
phase = LithosphericPhases(Layers=[8 72], Phases=[2 1 0], Tlab=Tlab),
T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0)
)

# Add slab
add_box!(
Phases,
Temp,
Grid2D;
xlim = (3000-1430, 3000-1430-250),
zlim =(-80, 0.0),
Phases,
Temp,
Grid2D;
xlim = (3000-1430, 3000-1430-250),
zlim =(-80, 0.0),
Origin = nothing, StrikeAngle=0, DipAngle=-30,
phase = LithosphericPhases(Layers=[8 80], Phases=[2 1 0], Tlab=Tlab),
phase = LithosphericPhases(Layers=[8 80], Phases=[2 1 0], Tlab=Tlab),
T = HalfspaceCoolingTemp(Tsurface=20, Tmantle=Tbot, Age=50, Adiabat=0)
)

surf = Grid2D.z.val .> 0.0
surf = Grid2D.z.val .> 0.0
Temp[surf] .= 20.0
Phases[surf] .= 3

Grid2D = addfield(Grid2D,(;Phases, Temp))

write_paraview(Grid2D,"Initial_Setup_Subduction_rank");
li = (abs(last(x)-first(x)), abs(last(z)-first(z))).* 1e3
origin = (x[1], z[1]) .* 1e3

ph = Phases[:,1,:] .+ 1
T = Temp[:,1,:]

return li, origin, ph, T
end
end
51 changes: 23 additions & 28 deletions miniapps/subduction/3D/Subduction3D_MPI.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,9 @@ function main3D(x_global, y_global, z_global, li, origin, phases_GMG, igg; nx=16
Vx_v = @zeros(ni.+1...)
Vy_v = @zeros(ni.+1...)
Vz_v = @zeros(ni.+1...)
Vx = @zeros(ni...)
Vy = @zeros(ni...)
Vz = @zeros(ni...)
end


Expand All @@ -128,20 +131,17 @@ function main3D(x_global, y_global, z_global, li, origin, phases_GMG, igg; nx=16
εII_nohalo = zeros(nx-2, ny-2, nz-2)
phases_c_nohalo = zeros(nx-2, ny-2, nz-2)
#vertex
Vxv_v = zeros(nx_v+1, ny_v+1, nz_v+1)
Vyv_v = zeros(nx_v+1, ny_v+1, nz_v+1)
Vzv_v = zeros(nx_v+1, ny_v+1, nz_v+1)
T_v = zeros(nx_v+1, ny_v+1, nz_v+1)
phases_v_v = zeros(nx_v+1, ny_v+1, nz_v+1)
Vxv_v = zeros(nx_v, ny_v, nz_v)
Vyv_v = zeros(nx_v, ny_v, nz_v)
Vzv_v = zeros(nx_v, ny_v, nz_v)
T_v = zeros(nx_v, ny_v, nz_v)
#vertex nohalo
Vxv_nohalo = zeros(nx-1, ny-1, nz-1)
Vyv_nohalo = zeros(nx-1, ny-1, nz-1)
Vzv_nohalo = zeros(nx-1, ny-1, nz-1)
T_nohalo = zeros(nx-1, ny-1, nz-1)
phase_v_nohalo = zeros(nx-1, ny-1, nz-1)
Vxv_nohalo = zeros(nx-2, ny-2, nz-2)
Vyv_nohalo = zeros(nx-2, ny-2, nz-2)
Vzv_nohalo = zeros(nx-2, ny-2, nz-2)
T_nohalo = zeros(nx-2, ny-2, nz-2)

xci_v = LinRange(minimum(x_global).*1e3, maximum(x_global).*1e3, nx_v), LinRange(minimum(y_global).*1e3, maximum(y_global).*1e3, ny_v), LinRange(minimum(z_global).*1e3, maximum(z_global).*1e3, nz_v)
xvi_v = LinRange(minimum(x_global).*1e3, maximum(x_global).*1e3, nx_v+1), LinRange(minimum(y_global).*1e3, maximum(y_global).*1e3, ny_v+1), LinRange(minimum(z_global).*1e3, maximum(z_global).*1e3, nz_v+1)

# Time loop
t, it = 0.0, 0
Expand Down Expand Up @@ -224,7 +224,6 @@ function main3D(x_global, y_global, z_global, li, origin, phases_GMG, igg; nx=16
end

#MPI gathering
phase_vertex = [argmax(p) for p in Array(phase_ratios.vertex)]
phase_center = [argmax(p) for p in Array(phase_ratios.center)]
#centers
@views P_nohalo .= Array(stokes.P[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
Expand All @@ -240,28 +239,26 @@ function main3D(x_global, y_global, z_global, li, origin, phases_GMG, igg; nx=16
#vertices
if do_vtk
velocity2vertex!(Vx_v, Vy_v, Vz_v, @velocity(stokes)...)
@views Vxv_nohalo .= Array(Vx_v[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vyv_nohalo .= Array(Vy_v[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vzv_nohalo .= Array(Vz_v[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
# gather!(Vxv_nohalo, Vxv_v)
# gather!(Vyv_nohalo, Vyv_v)
# gather!(Vzv_nohalo, Vzv_v)
vertex2center!(Vx, Vx_v)
vertex2center!(Vy, Vy_v)
vertex2center!(Vz, Vz_v)
@views Vxv_nohalo .= Array(Vx[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vyv_nohalo .= Array(Vy[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
@views Vzv_nohalo .= Array(Vz[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(Vxv_nohalo, Vxv_v)
gather!(Vyv_nohalo, Vyv_v)
gather!(Vzv_nohalo, Vzv_v)
end
# @views T_nohalo .= Array(thermal.T[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
# @views phase_v_nohalo .= Array(phase_vertex[2:end-1, 2:end-1, 2:end-1])
# gather!(T_nohalo, T_v)
# gather!(phase_v_nohalo, phases_v_v)
@views T_nohalo .= Array(thermal.Tc[2:end-1, 2:end-1, 2:end-1]) # Copy data to CPU removing the halo
gather!(T_nohalo, T_v)

# Data I/O and plotting ---------------------
if igg.me == 0 && (it == 1 || rem(it, 1) == 0)
# checkpointing(figdir, stokes, thermal.T, η, t)
if do_vtk

data_v = (;
# T = T_v,
# phases_v = phases_v_v,
)
data_c = (;
T = T_v,
P = P_v,
τII = τII_v,
εII = εII_v,
Expand All @@ -277,9 +274,7 @@ function main3D(x_global, y_global, z_global, li, origin, phases_GMG, igg; nx=16
)
save_vtk(
joinpath(vtk_dir, "vtk_" * lpad("$(it)_$(igg.me)", 6, "0")),
xvi_v./1e3,
xci_v./1e3,
data_v,
data_c,
velocity_v
)
Expand Down

0 comments on commit e5c74d7

Please sign in to comment.