Обратное быстрое преобразование Фурье (ifft2) scipy не работает для оптики Фурье ⇐ Python

[Anonymous]

Я слежу за уроком на YouTube по оптике Фурье на Python, чтобы имитировать дифракцию света через щель.

Это видео

Исходный код видео

Теперь я пытаюсь реализовать функцию get_U(z, k), а затем отобразить соответствующий график под ней, как показано на видео (у меня есть скудные знания по этой теме), однако я просто не могу заставить сюжет работать (белый сюжет виден все время). При проверке я обнаружил, что переменная U состоит всего лишь из набора значений (nan+nanj), чего, по моему мнению, быть не должно. Я перепроверил формулу, и она выглядит идеально. Я также понимаю, что иногда np.sqrt() имеет дело с отрицательными значениями, но не добавляет ни np.abs(), ни np.where()(для преобразования отрицательные значения равны нулю) дает мне желаемый результат.

Мой код:
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()


Подробнее здесь: https://stackoverflow.com/questions/793 ... ier-optics