array_tools.py 1.42 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
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


26
27
28
29
30
31
32
def shifted_fft(arr, axes=None):
    """
    Combined fftshift and fft routine, based on scipy.fftpack

    Args:
        arr: numpy array
        axes: identical to  argument for scipy.fftpack.fft
33

34
35
36
    Returns:
        transformed array
    """
37

38
    return fft.ifftshift(fft.fftn(fft.fftshift(arr, axes=axes), axes=axes), axes=axes)
39
40


41
def shifted_ifft(arr, axes=None):
42
    """
43
44
45
46
47
    Combined fftshift and fft routine, based on scipy.fftpack

    Args:
        arr: numpy array
        axes: identical to  argument for scipy.fftpack.fft
48

49
50
    Returns:
        transformed array
51
    """
52
    return fft.fftshift(fft.ifftn(fft.ifftshift(arr, axes=axes), axes=axes), axes=axes)