Scipy二维逆快速傅里叶变换(ifft2)在傅里叶光学仿真中失效问题求助

我最近跟着一个傅里叶光学的Python教程做单缝衍射的仿真，目标是实现get_U(z, k)函数并绘制对应的衍射图样，但始终得不到正确的结果——绘图窗口里全是白色，检查后发现计算得到的U变量全是(nan+nanj)值。

我对照了教程里的公式，看起来是完全一致的。考虑到np.sqrt()可能会遇到负数输入，我试过加np.abs()包裹根号里的内容，也用过np.where()把负数转换成0，但都没能得到预期的衍射图样。

以下是我的代码：

import numpy as np
import scipy as sp
from scipy.fft import fft2, ifft2, fftfreq, fftshift
import matplotlib.pyplot as plt
import pint

plt.style.use(['grayscale'])
u = pint.UnitRegistry()

D = 0.1 * u.mm
lam = 660 * u.mm

x = np.linspace(-2, 2, 1600) * u.mm
xv, yv = np.meshgrid(x, x)

U0 = (np.abs(xv) < D/2) * (np.abs(yv) < 0.5 * u.mm)
U0 = U0.astype(float)

A = fft2(U0)
kx = fftfreq(len(x), np.diff(x)[0]) * 2 * np.pi
kxv, kyv = np.meshgrid(kx, kx)

def get_U(z, k):
  return ifft2(A*np.exp(1j*z.magnitude*np.sqrt(k.magnitude**2 - kxv.magnitude**2 - kyv.magnitude**2)))
k = 2*np.pi/(lam)
d = 3 * u.cm
U = get_U(d, k)

plt.figure(figsize=(5, 5))
plt.pcolormesh(xv, yv, np.abs(U), cmap='inferno')
plt.xlabel('$x$ [mm]')
plt.ylabel('$y$ [mm]')
plt.title('Single slit diffraction')
plt.show()

麻烦各位帮我看看问题出在哪里？为什么会出现全是nan的复数数组，该怎么调整才能得到正确的衍射图样？

备注：内容来源于stack exchange，提问作者Hrishav Saha