Proper fft velocity

.gitignore

@@ -0,0 +1,10 @@
+PhaseUtils.egg-info/dependency_links.txt
+PhaseUtils.egg-info/PKG-INFO
+PhaseUtils.egg-info/requires.txt
+PhaseUtils.egg-info/SOURCES.txt
+PhaseUtils.egg-info/top_level.txt
+build/lib/PhaseUtils/contrast.py
+build/lib/PhaseUtils/fast_interp.py
+build/lib/PhaseUtils/SLM.py
+build/lib/PhaseUtils/velocity.py
+tests/fft.wisdom

PhaseUtils/velocity.py

@@ -135,30 +135,28 @@ def velocity_cp(phase: cp.ndarray, dx: float = 1) -> cp.ndarray:
         velo[1, :, :] = cp.gradient(phase_unwrap[1, :, :], dx, axis=0)
         return velo
 
-    def velocity_fft_cp(phase: cp.ndarray, dx: float = 1) -> cp.ndarray:
-        """Returns the velocity from the phase using an fft to compute
-        the gradient
+    def velocity_fft_cp(field: cp.ndarray, dx: float = 1) -> cp.ndarray:
+        """Compute velocity from the field.
 
         Args:
-            phase (cp.ndarray): The field phase
-            dx (float, optional): the pixel size in m. Defaults to 1 (adimensional).
+            field (cp.ndarray): The field to compute the velocity
+            dx (float, optional): pixel size in m. Defaults to 1.
 
         Returns:
-            cp.ndarray: The velocity field [vx, vy]
+            cp.ndarray: the velocity field [vx, vy]
         """
-        # 1D unwrap
-        phase_unwrap = cp.empty((2, phase.shape[-2], phase.shape[-1]), dtype=np.float32)
-        phase_unwrap[0, :, :] = cp.unwrap(phase, axis=-1)
-        phase_unwrap[1, :, :] = cp.unwrap(phase, axis=-2)
+        rho = field.real * field.real + field.imag * field.imag
         # prepare K matrix
-        kx = 2 * np.pi * cp.fft.fftfreq(phase.shape[-1], dx)
-        ky = 2 * np.pi * cp.fft.fftfreq(phase.shape[-2], dx)
+        kx = 2 * np.pi * cp.fft.fftfreq(field.shape[-1], dx)
+        ky = 2 * np.pi * cp.fft.fftfreq(field.shape[-2], dx)
+        K = cp.array(cp.meshgrid(kx, ky))
         K = cp.array(cp.meshgrid(kx, ky))
         # gradient reconstruction
-        velo = cp.empty((2, phase.shape[-2], phase.shape[-1]), dtype=np.float32)
-        velo = cp.fft.irfft2(
-            1j * K[:, :, 0 : K.shape[-1] // 2 + 1] * cp.fft.rfft2(phase_unwrap)
-        )
+        velo = cp.fft.ifft2(1j * K * cp.fft.fft2(field))
+        velo[0, :, :] = cp.imag(cp.conj(field) * velo[0, :, :]) / rho
+        velo[1, :, :] = cp.imag(cp.conj(field) * velo[1, :, :]) / rho
+        velo[cp.isnan(velo)] = 0
+        velo = velo.astype(np.float32)
         return velo
 
     def helmholtz_decomp_cp(

tests/test_contrast.py

@@ -13,9 +13,33 @@
 
 
 def main():
+    from PhaseUtils import velocity
+    from cupyx.scipy import ndimage
+
     im = np.array(Image.open("../examples/dn_ref.tiff"))
     field = contrast.im_osc_fast(im)
     field_t = contrast.im_osc_fast_t(im)
+    field_t = ndimage.gaussian_filter(cp.asarray(field_t), 2)
+    velo = velocity.velocity_cp(cp.asarray(np.angle(field_t)))
+    velo_field = velocity.velocity_fft_cp(cp.asarray(field_t))
+    fig, ax = plt.subplots(2, 2)
+    ax[0, 0].imshow(velo[0, :, :].get(), vmin=-0.5, vmax=0.5)
+    ax[0, 1].imshow(velo[1, :, :].get(), vmin=-0.5, vmax=0.5)
+    ax[0, 0].set_title("Vx")
+    ax[0, 1].set_title("Vy")
+    ax[1, 0].imshow(velo_field[0, :, :].get(), vmin=-0.5, vmax=0.5)
+    ax[1, 1].imshow(velo_field[1, :, :].get(), vmin=-0.5, vmax=0.5)
+    ax[1, 0].set_title("Vx field")
+    ax[1, 1].set_title("Vy field")
+    plt.show()
+    fig, ax = plt.subplots(1, 2)
+    ax[0].imshow((velo[0] - velo_field[0]).get(), cmap="coolwarm", vmin=-0.5, vmax=0.5)
+    ax[1].imshow((velo[1] - velo_field[1]).get(), cmap="coolwarm", vmin=-0.5, vmax=0.5)
+    ax[0].set_title("Vx diff")
+    ax[1].set_title("Vy diff")
+    plt.show()
+    velocity.helmholtz_decomp_cp(cp.asarray(field), plot=True)
+    field_t = contrast.im_osc_fast_t(im)
     contrast.im_osc(im, plot=True)
     fig, ax = plt.subplots(2, 2)
     ax[0, 0].imshow(np.abs(field) ** 2)