array_tools.py 1.33 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import numpy as np
from scipy.ndimage import measurements
import scipy.fftpack as fft


def shift_array(arr, dy, dx):
    temp = np.roll(arr, (dy, dx), (0, 1))
    return temp


def roll_to_pos(arr, y=0, x=0, move_maximum=False, by_abs_val=True):
    if move_maximum:
        if by_abs_val or arr.dtype in [np.complex64, np.complex128]:
            old = np.floor(measurements.maximum_position(abs(arr)))
        else:
            old = np.floor(measurements.maximum_position(arr))
    else:
        if by_abs_val or arr.dtype in [np.complex64, np.complex128]:
            old = np.floor(measurements.center_of_mass(abs(arr)))
        else:
            old = np.floor(measurements.center_of_mass(arr))
    temp = shift_array(arr, int(y - old[0]), int(x - old[1]))
    return temp


def shifted_fft(arr):
    return fft.ifftshift(fft.fft2(fft.fftshift(arr)))


def shifted_ifft(arr):
    return fft.fftshift(fft.ifft2(fft.ifftshift(arr)))


def fourier_interpolate(arr):
    """
    Interpolate 2d complex array by a factor 2

    :param arr: numpy array
    :return: interpolated array
    """
    ny, nx = arr.shape
    assert ny == nx
    pd = int(ny / 2)
    _arr = roll_to_pos(arr, pd, pd, move_maximum=True)
    _arr = roll_to_pos(_arr, pd, pd)
    fd = shifted_fft(arr)
    tmp = np.pad(fd, ((pd, pd), (pd, pd)), mode='constant')
    return shifted_ifft(tmp)