import numpy
[docs]def fft_wavefunction(psi, nx, ny, ns, sin):
return numpy.fft.fft2(psi.reshape(nx, ny, ns), axes=(0, 1)).reshape(sin)
[docs]def ifft_wavefunction(psi, nx, ny, ns, sin):
return numpy.fft.ifft2(psi.reshape(nx, ny, ns), axes=(0, 1)).reshape(sin)