I am new to Python and have seen some articles and examples of vectorizing nested for loops, but I am unclear on how to go about that with some FDTD code that I wrote. The goal is to make the code cleaner and run more efficiently. Can someone help me vectorize this code and explain to me the thought process behind the changes? I know this code isn't the cleanest by a long shot so any help is very appreciated.
def fdtd(prl, pim, nsteps, T, xpml, ypml, zpml, Ptime_rl, Ptime_im,
Phi_rl, Phi_im, omegaT, win_Phi, PF_TIME, xpos, ypos, MDM_time):
print('n_steps =', nsteps)
print('B_0 =', B0)
for _ in range(nsteps):
T = T + 1
for n in range(1, NN - 1):
for m in range(1, MM - 1):
for k in range(1, KK - 1):
prl[n, m, k] = (xpml[n] * ypml[m] * zpml[k] * prl[n, m, k]
+ rd * V[n, m, k] * pim[n, m, k]
- ra * (-6 * pim[n, m, k]
+ pim[n + 1, m, k] + pim[n - 1, m, k]
+ pim[n, m + 1, k] + pim[n, m - 1, k]
+ pim[n, m, k - 1] + pim[n, m, k + 1]
)
+ rd * ((B0 + Btime[T]) ** 2) * VoB[n, m] * pim[n, m, k]
+ K_dfB * (B0 + Btime[T]) * 0.5 * (-(m - MC) * (prl[n + 1, m, k] -
prl[n - 1, m, k])
+ (n - NC) * (prl[n, m + 1, k] -
prl[n, m - 1, k]))
)
for n in range(1, NN - 1):
for m in range(1, MM - 1):
for k in range(1, KK - 1):
pim[n, m, k] = (xpml[n] * ypml[m] * zpml[k] * pim[n, m, k]
- rd * V[n, m, k] * prl[n, m, k]
+ ra * (-6 * prl[n, m, k]
+ prl[n + 1, m, k] + prl[n - 1, m, k]
+ prl[n, m + 1, k] + prl[n, m - 1, k]
+ prl[n, m, k - 1] + prl[n, m, k + 1]
)
- rd * ((B0 + Btime[T]) ** 2) * VoB[n, m] * prl[n, m, k]
+ K_dfB * (B0 + Btime[T]) * 0.5 * (-(m - MC) * (pim[n + 1, m, k] -
pim[n - 1, m, k])
+ (n - NC) * (pim[n, m + 1, k] - pim[n, m - 1,k])))
Ptime_rl[T] = prl[NC + 35, MC, KC]
Ptime_im[T] = pim[NC + 35, MC, KC]
for n in range(NN):
for m in range(MM):
for k in range(KK):
xpos[T] = xpos[T] + (n - NC) * (prl[n, m, k] ** 2 + pim[n, m, k] ** 2)
ypos[T] = ypos[T] + (m - MC) * (prl[n, m, k] ** 2 + pim[n, m, k] ** 2)
# print(T,'xpos,ypos:',round(xpos[T],3),round(ypos[T],3))
# This section calculates the magnetic dipole moment
mdm = 0 + 1j * 0
for n in range(1, NN - 1):
for m in range(1, MM - 1):
for k in range(KK):
pmag = prl[n, m, k] ** 2 + pim[n, m, k] ** 2
der_y = 0.5 * (n - NC) * (prl[n, m + 1, k] - prl[n, m - 1, k]
+ 1j * (pim[n, m + 1, k] - pim[n, m - 1, k]))
der_x = 0.5 * (m - MC) * (prl[n + 1, m, k] - prl[n - 1, m, k]
+ 1j * (pim[n + 1, m, k] - pim[n - 1, m, k]))
# The 0.5 are because the spatial derivative is over two cells.
mdm = mdm + (Kmdm * (1j * prl[n, m, k] + pim[n, m, k]) * (der_y - der_x)
+ (B0 + Btime[T]) * KdelB * (del_x ** 2) * ((n - NC) ** 2 + (m - MC)
** 2) * pmag)
MDM_time[T] = 1e23 * mdm.real
# This section constructs an eigenfunction
if T < PF_TIME:
cos_term = win_Phi[T] * cos(omegaT * T)
sin_term = win_Phi[T] * sin(omegaT * T)
for n in range(NN):
for m in range(MM):
for k in range(KK):
Phi_rl[n, m, k] = (Phi_rl[n, m, k]
+ cos_term * prl[n, m, k] - sin_term * pim[n, m, k])
Phi_im[n, m, k] = (Phi_im[n, m, k]
+ sin_term * prl[n, m, k] + cos_term * pim[n, m, k])
return prl, pim, T, Ptime_rl, Ptime_im, Phi_rl, Phi_im, xpos, ypos, MDM_time
