array_tools.py 2.32 KB
 jansen31 committed Jun 12, 2019 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 committed Mar 29, 2020 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 """ `````` jansen31 committed Jun 12, 2019 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 committed Mar 29, 2020 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])) `````` jansen31 committed Jun 12, 2019 40 41 42 `````` return temp `````` Matthijs committed Mar 27, 2020 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 `````` jansen31 committed Jun 12, 2019 50 `````` `````` Matthijs committed Mar 27, 2020 51 52 53 `````` Returns: transformed array """ `````` jansen31 committed Jun 12, 2019 54 `````` `````` Matthijs committed Mar 27, 2020 55 `````` return fft.ifftshift(fft.fftn(fft.fftshift(arr, axes=axes), axes=axes), axes=axes) `````` jansen31 committed Jun 12, 2019 56 57 `````` `````` Matthijs committed Mar 27, 2020 58 ``````def shifted_ifft(arr, axes=None): `````` jansen31 committed Jun 12, 2019 59 `````` """ `````` Matthijs committed Mar 27, 2020 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 `````` jansen31 committed Jun 12, 2019 65 `````` `````` Matthijs committed Mar 27, 2020 66 67 `````` Returns: transformed array `````` jansen31 committed Jun 12, 2019 68 `````` """ `````` Matthijs committed Mar 27, 2020 69 `` return fft.fftshift(fft.ifftn(fft.ifftshift(arr, axes=axes), axes=axes), axes=axes)``