Skip to content

Commit a41fc1f

Browse files
elenafuengarelenafuengar
committed
feature: first working example with MPI (needs more iterations to be perfect)
1 parent bf87d9c commit a41fc1f

File tree

1 file changed

+122
-133
lines changed

1 file changed

+122
-133
lines changed

‎examples/003_MPI_wakefield_simulation.py‎

Lines changed: 122 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@
77
import numpy as np
88
import pyvista as pv
99
import matplotlib.pyplot as plt
10+
from mpl_toolkits.axes_grid1 import make_axes_locatable
11+
12+
import sys
13+
sys.path.append('../')
1014

1115
from wakis import SolverFIT3D
1216
from wakis import GridFIT3D
@@ -47,42 +51,28 @@
4751
Nz = 141
4852

4953
# Adjust for MPI & ompute local Z-slice range
50-
Nz += Nz%(size-1)
54+
Nz += Nz%(size) #if rank 0 is common, then size-1
5155
dz = (zmax - zmin) / Nz
5256
z = np.linspace(zmin, zmax, Nz+1)
5357

54-
Nz_mpi = Nz // (size-1)
55-
zmin_mpi = (rank-1) * Nz_mpi * dz
56-
zmax_mpi= rank * Nz_mpi * dz
58+
# Allocate mpi node cells
59+
Nz_mpi = Nz // (size)
60+
zmin_mpi = rank * Nz_mpi * dz + zmin
61+
zmax_mpi= (rank+1) * Nz_mpi * dz + zmin
5762

5863
print(f"Process {rank}: Handling Z range {zmin_mpi} to {zmax_mpi}")
5964

60-
# set grid and geometry
61-
if rank == 0:
62-
grid = GridFIT3D(xmin, xmax, ymin, ymax, zmin, zmax,
63-
Nx, Ny, Nz,
64-
stl_solids=stl_solids,
65-
stl_materials=stl_materials,
66-
stl_scale=1.0,
67-
stl_rotate=[0,0,0],
68-
stl_translate=[0,0,0],
69-
verbose=1)
70-
else:
71-
grid = GridFIT3D(xmin, xmax, ymin, ymax,
72-
zmin_mpi, zmax_mpi,
73-
Nx, Ny, Nz_mpi,
74-
stl_solids=stl_solids,
75-
stl_materials=stl_materials,
76-
stl_scale=1.0,
77-
stl_rotate=[0,0,0],
78-
stl_translate=[0,0,0],
79-
verbose=1)
80-
# BONUS: Visualize grid - Uncomment for plotting!
81-
# grid.inspect(add_stl=[solid_1, solid_2], stl_opacity=1.0)
82-
83-
# BONUS: Visualize imported solids - Uncomment for plotting!
84-
#grid.plot_solids()
65+
grid = GridFIT3D(xmin, xmax, ymin, ymax,
66+
zmin_mpi, zmax_mpi,
67+
Nx, Ny, Nz_mpi,
68+
stl_solids=stl_solids,
69+
stl_materials=stl_materials,
70+
stl_scale=1.0,
71+
stl_rotate=[0,0,0],
72+
stl_translate=[0,0,0],
73+
verbose=1)
8574

75+
# ------------ Beam source & Wake ----------------
8676
# ------------ Beam source & Wake ----------------
8777
# Beam parameters
8878
sigmaz = 10e-2 #[m] -> 2 GHz
@@ -98,133 +88,132 @@
9888
wakelength = 10. # [m]
9989
add_space = 10 # no. cells to skip from boundaries - removes BC artifacts
10090

91+
from wakis.sources import Beam
92+
from scipy.constants import c
93+
beam = Beam(q=q, sigmaz=sigmaz, beta=beta,
94+
xsource=xs, ysource=ys, ti=3*sigmaz/c)
95+
96+
results_folder = f'001_results_n{rank}/'
97+
98+
'''
10199
wake = WakeSolver(q=q,
102100
sigmaz=sigmaz,
103101
beta=beta,
104102
xsource=xs, ysource=ys,
105103
xtest=xt, ytest=yt,
106104
add_space=add_space,
107-
results_folder='001_results/',
108-
Ez_file='001_results/001_Ez.h5')
105+
results_folder=results_folder,
106+
Ez_file=results_folder+'001_Ez.h5')
107+
'''
109108

110109
# ----------- Solver & Simulation ----------
111-
# boundary conditions``
110+
# boundary conditions
112111
bc_low=['pec', 'pec', 'pec']
113112
bc_high=['pec', 'pec', 'pec']
114113

115-
# on-the-fly plotting parameters
116-
if not os.path.exists('001_img/'):
117-
os.mkdir('001_img/')
114+
if rank > 0:
115+
bc_low=['pec', 'pec', 'periodic']
118116

119-
plotkw2D = {'title':'001_img/Ez',
120-
'add_patch':'cavity', 'patch_alpha':1.0,
121-
'patch_reverse' : True, # patch logical_not('cavity')
122-
'vmin':-1e3, 'vmax':1e3, # colormap limits
123-
'cmap': 'rainbow',
124-
'plane': [int(Nx/2), # x
125-
slice(0, Ny), # y
126-
slice(add_space, -add_space)]} # z
117+
if rank < size - 1:
118+
bc_high=['pec', 'pec', 'periodic']
127119

128120
# Solver setup
129-
solver = SolverFIT3D(grid, wake,
121+
solver = SolverFIT3D(grid,
130122
bc_low=bc_low,
131123
bc_high=bc_high,
132124
use_stl=True,
133125
bg='pec' # Background material
134126
)
135-
# [TODO]: domain should be split after the tensors are built
136-
# to avoid issues with boundary conditions / geometry
137-
# Maybe need to add >1 ghost cells
127+
# Communication between ghost cells
128+
def communicate_ghost_cells():
129+
if rank > 0:
130+
for d in ['x','y','z']:
131+
comm.Sendrecv(solver.E[:, :, 1,d],
132+
recvbuf=solver.E[:, :, 0,d],
133+
dest=rank-1, sendtag=0,
134+
source=rank-1, recvtag=1)
135+
136+
comm.Sendrecv(solver.H[:, :, 1,d],
137+
recvbuf=solver.H[:, :, 0,d],
138+
dest=rank-1, sendtag=0,
139+
source=rank-1, recvtag=1)
140+
141+
comm.Sendrecv(solver.J[:, :, 1,d],
142+
recvbuf=solver.J[:, :, 0,d],
143+
dest=rank-1, sendtag=0,
144+
source=rank-1, recvtag=1)
145+
146+
if rank < size - 1:
147+
for d in ['x','y','z']:
148+
comm.Sendrecv(solver.E[:, :, -2,d],
149+
recvbuf=solver.E[:, :, -1, d],
150+
dest=rank+1, sendtag=1,
151+
source=rank+1, recvtag=0)
152+
153+
comm.Sendrecv(solver.H[:, :, -2,d],
154+
recvbuf=solver.H[:, :, -1, d],
155+
dest=rank+1, sendtag=1,
156+
source=rank+1, recvtag=0)
157+
158+
comm.Sendrecv(solver.J[:, :, -2,d],
159+
recvbuf=solver.J[:, :, -1, d],
160+
dest=rank+1, sendtag=1,
161+
source=rank+1, recvtag=0)
162+
163+
def compose_field(field, d='z'):
164+
165+
if field == 'E':
166+
local = solver.E[Nx//2, :, :,d].ravel()
167+
elif field == 'H':
168+
local = solver.H[Nx//2, :, :,d].ravel()
169+
elif field == 'J':
170+
local = solver.J[Nx//2, :, :,d].ravel()
171+
172+
buffer = comm.gather(local, root=0)
173+
field = None
174+
175+
if rank == 0:
176+
field = np.zeros((Ny, Nz)) # Reinitialize global array
177+
for r in range(size):
178+
field[:, r*Nz_mpi:(r+1)*Nz_mpi] = np.reshape(buffer[r], (Ny, Nz_mpi))
179+
180+
return field
181+
182+
def plot_field(field, name='E', n=None, results_folder='img/'):
183+
extent = (zmin, zmax, ymin, ymax)
184+
fig, ax = plt.subplots()
185+
im = ax.imshow(field, cmap='bwr', extent=extent, vmin=-500, vmax=500)
186+
fig.colorbar(im, cax=make_axes_locatable(ax).append_axes('right', size='5%', pad=0.05))
187+
ax.set_title(f'{name} at timestep={n}')
188+
ax.set_xlabel('z [m]')
189+
ax.set_ylabel('y [m]')
190+
191+
fig.tight_layout(h_pad=0.3)
192+
fig.savefig(results_folder+name+str(n).zfill(4)+'.png')
193+
194+
plt.clf()
195+
plt.close(fig)
138196

139-
# Solver run
140-
'''
141-
solver.wakesolve(wakelength=wakelength,
142-
add_space=add_space,
143-
plot=True, # turn False for speedup
144-
plot_every=30, plot_until=3000, **plotkw2D
145-
)
146-
'''
147-
from wakis.sources import Beam
148-
beam = Beam(q=q, sigmaz=sigmaz, beta=beta,
149-
xsource=xs, ysource=ys)
150-
151-
Nt = 1/solver.dt * (wake.wakelength + wake.ti*wake.v \
152-
+ (solver.z.max()-solver.z.min()))/wake.v
197+
if rank == 0:
198+
img_folder = '003_img/'
199+
if not os.path.exists(img_folder):
200+
os.mkdir(img_folder)
153201

202+
Nt = 3000
154203
for n in tqdm(range(Nt)):
155204

156-
beam.update_mpi(solver, n*solver.dt, zmin, z)
205+
beam.update_mpi(solver, n*solver.dt, zmin)
157206

158207
solver.one_step()
159208

160-
#Communicate slices [TODO]
161-
162-
163-
# ----------- 1d plot results --------------------
164-
# Plot longitudinal wake potential and impedance
165-
fig1, ax = plt.subplots(1,2, figsize=[12,4], dpi=150)
166-
ax[0].plot(wake.s*1e2, wake.WP, c='r', lw=1.5, label='Wakis')
167-
ax[0].set_xlabel('s [cm]')
168-
ax[0].set_ylabel('Longitudinal wake potential [V/pC]', color='r')
169-
ax[0].legend()
170-
ax[0].set_xlim(xmax=wakelength*1e2)
171-
172-
ax[1].plot(wake.f*1e-9, np.abs(wake.Z), c='b', lw=1.5, label='Wakis')
173-
ax[1].set_xlabel('f [GHz]')
174-
ax[1].set_ylabel('Longitudinal impedance [Abs][$\Omega$]', color='b')
175-
ax[1].legend()
176-
177-
fig1.tight_layout()
178-
fig1.savefig('001_results/001_longitudinal.png')
179-
#plt.show()
180-
181-
# Plot transverse x wake potential and impedance
182-
fig2, ax = plt.subplots(1,2, figsize=[12,4], dpi=150)
183-
ax[0].plot(wake.s*1e2, wake.WPx, c='r', lw=1.5, label='Wakis')
184-
ax[0].set_xlabel('s [cm]')
185-
ax[0].set_ylabel('Transverse wake potential X [V/pC]', color='r')
186-
ax[0].legend()
187-
ax[0].set_xlim(xmax=wakelength*1e2)
188-
189-
ax[1].plot(wake.f*1e-9, np.abs(wake.Zx), c='b', lw=1.5, label='Wakis')
190-
ax[1].set_xlabel('f [GHz]')
191-
ax[1].set_ylabel('Transverse impedance X [Abs][$\Omega$]', color='b')
192-
ax[1].legend()
193-
194-
fig2.tight_layout()
195-
fig2.savefig('001_results/001_transverse_x.png')
196-
#plt.show()
197-
198-
# Plot transverse y wake potential and impedance
199-
fig3, ax = plt.subplots(1,2, figsize=[12,4], dpi=150)
200-
ax[0].plot(wake.s*1e2, wake.WPy, c='r', lw=1.5, label='Wakis')
201-
ax[0].set_xlabel('s [cm]')
202-
ax[0].set_ylabel('Transverse wake potential Y [V/pC]', color='r')
203-
ax[0].legend()
204-
ax[0].set_xlim(xmax=wakelength*1e2)
205-
206-
ax[1].plot(wake.f*1e-9, np.abs(wake.Zy), c='b', lw=1.5, label='Wakis')
207-
ax[1].set_xlabel('f [GHz]')
208-
ax[1].set_ylabel('Transverse impedance Y [Abs][$\Omega$]', color='b')
209-
ax[1].legend()
210-
211-
fig3.tight_layout()
212-
fig3.savefig('001_results/001_transverse_y.png')
213-
#plt.show()
214-
215-
# Plot Electric field component in 2D using imshow
216-
solver.plot1D(field='E', component='z',
217-
line='z', pos=0.5, xscale='linear', yscale='linear',
218-
off_screen=True, title='001_img/Ez1d')
219-
#plt.show()
220-
221-
# ----------- 2d plots results --------------------
222-
from matplotlib.colors import LinearSegmentedColormap
223-
cmap = LinearSegmentedColormap.from_list('name', plt.cm.jet(np.linspace(0.1, 0.9))) # CST's colormap
224-
225-
# Plot Electric field component in 2D using imshow
226-
solver.plot2D(field='E', component='z',
227-
plane='XY', pos=0.5,
228-
cmap=cmap, vmin=-500, vmax=500., interpolation='hanning',
229-
add_patch='cavity', patch_reverse=True, patch_alpha=0.8,
230-
off_screen=True, title='001_img/Ez2d')
209+
#Communicate slices
210+
communicate_ghost_cells()
211+
212+
if n%20 == 0:
213+
Ez = compose_field('E')
214+
if rank == 0:
215+
plot_field(Ez, 'Ez', n=n, results_folder=img_folder)
216+
217+
218+
# Run with:
219+
# mpiexec -n 4 python 003_MPI_wakefield_simulation.py

0 commit comments

Comments
 (0)