array_tools.py 2.32 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
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


Matthijs's avatar
Matthijs committed
11
12
13
14
15
16
17
18
19
20
21
22
23
def roll_to_pos(arr: np.ndarray, y: int = 0, x: int = 0, pos: tuple = None, move_maximum: bool = False,
                by_abs_val: bool = True) -> np.ndarray:
    """
    Shift the center of mass of an array to the given position by cyclic permutation

    :param arr: 2d array, works best for well-centered feature with limited support
    :param y: position parameter
    :param x: position parameter for second dimension
    :param pos: tuple with the new position, overriding y,x values. should be used for higher-dimensional arrays
    :param move_maximum: if true, look only at max-value
    :param by_abs_val: take abs value for the determination of max-val or center-of-mass
    :return: array like original
    """
24
25
26
27
28
29
30
31
32
33
    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))
Matthijs's avatar
Matthijs committed
34
35
36
37
38
39
    if pos is not None:  # dimension-independent method
        shifts = tuple([int(np.round(pos[i]-old[i])) for i in range(len(pos))])
        dims = tuple([i for i in range(len(pos))])
        temp = np.roll(arr, shift=shifts, axis=dims)
    else:  # old method
        temp = shift_array(arr, int(y - old[0]), int(x - old[1]))
40
41
42
    return temp


43
44
45
46
47
48
49
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
50

51
52
53
    Returns:
        transformed array
    """
54

55
    return fft.ifftshift(fft.fftn(fft.fftshift(arr, axes=axes), axes=axes), axes=axes)
56
57


58
def shifted_ifft(arr, axes=None):
59
    """
60
61
62
63
64
    Combined fftshift and fft routine, based on scipy.fftpack

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

66
67
    Returns:
        transformed array
68
    """
69
    return fft.fftshift(fft.ifftn(fft.ifftshift(arr, axes=axes), axes=axes), axes=axes)