Commit 51b0a44d authored by Leon Merten Lohse's avatar Leon Merten Lohse
Browse files

use scipy.fft instead of pyfftw

parent 6f163c72
Pipeline #175251 passed with stages
in 1 minute and 23 seconds
import numpy as np
from functools import reduce
from pyfftw.interfaces.numpy_fft import fftn, ifftn, ifftshift, fftshift
from scipy.fft import fftn, ifftn, ifftshift, fftshift
from . import finite_differences as fd
from .misc import fftfreqn
......@@ -62,16 +62,18 @@ class MultislicePropagator:
self.u = ifftshift(u0)
def step(self, n):
def step(self, n, workers=-1):
nshift = ifftshift(n)
F = _compute_potential(nshift, self._k) * self._ones
self._realKernel = np.exp(-1j * self._dz / (2 * self._k) * F)
realKernel = np.exp(-1j * self._dz / (2 * self._k) * F)
up = self.u
self.u = ifftn(self._fourierKernel * fftn(self._realKernel * up))
up *= realKernel
ft_u = fftn(up, workers=workers)
ft_u *= self._fourierKernel
self.u = ifftn(ft_u, workers=workers)
return fftshift(self.u)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment