Commit 73a6277f by Matthijs

### rearrange imports, documentation

parent 387f4ab0
 import numpy as np from numpy import roll, ndarray, floor, iscomplexobj, round from scipy.ndimage import measurements from scipy.ndimage.measurements import maximum_position, center_of_mass import scipy.fftpack as fft from scipy.fftpack import fftn, fftshift, ifftn, ifftshift def shift_array(arr, dy, dx): def shift_array(arr: ndarray, dy: int, dx: int): temp = np.roll(arr, (dy, dx), (0, 1)) """ Use numpy.roll to shift an array in the first and second dimensions :param arr: numpy array :param dy: shift in first dimension :param dx: shift in second dimension :return: array like arr """ temp = roll(arr, (dy, dx), (0, 1)) return temp return temp def roll_to_pos(arr: np.ndarray, y: int = 0, x: int = 0, pos: tuple = None, move_maximum: bool = False, def roll_to_pos(arr: ndarray, y: int = 0, x: int = 0, pos: tuple = None, move_maximum: bool = False, by_abs_val: bool = True) -> np.ndarray: by_abs_val: bool = True) -> ndarray: """ """ Shift the center of mass of an array to the given position by cyclic permutation Shift the center of mass of an array to the given position by cyclic permutation ... @@ -22,19 +29,19 @@ def roll_to_pos(arr: np.ndarray, y: int = 0, x: int = 0, pos: tuple = None, move ... @@ -22,19 +29,19 @@ def roll_to_pos(arr: np.ndarray, y: int = 0, x: int = 0, pos: tuple = None, move :return: array like original :return: array like original """ """ if move_maximum: if move_maximum: if by_abs_val or arr.dtype in [np.complex64, np.complex128]: if by_abs_val or iscomplexobj(arr): old = np.floor(measurements.maximum_position(abs(arr))) old = floor(maximum_position(abs(arr))) else: else: old = np.floor(measurements.maximum_position(arr)) old = floor(maximum_position(arr)) else: else: if by_abs_val or arr.dtype in [np.complex64, np.complex128]: if by_abs_val or iscomplexobj(arr): old = np.floor(measurements.center_of_mass(abs(arr))) old = floor(center_of_mass(abs(arr))) else: else: old = np.floor(measurements.center_of_mass(arr)) old = floor(center_of_mass(arr)) if pos is not None: # dimension-independent method if pos is not None: # dimension-independent method shifts = tuple([int(np.round(pos[i]-old[i])) for i in range(len(pos))]) shifts = tuple([int(round(pos[i] - old[i])) for i in range(len(pos))]) dims = tuple([i for i in range(len(pos))]) dims = tuple([i for i in range(len(pos))]) temp = np.roll(arr, shift=shifts, axis=dims) temp = roll(arr, shift=shifts, axis=dims) else: # old method else: # old method temp = shift_array(arr, int(y - old[0]), int(x - old[1])) temp = shift_array(arr, int(y - old[0]), int(x - old[1])) if temp.shape != arr.shape: if temp.shape != arr.shape: ... @@ -54,7 +61,7 @@ def shifted_fft(arr, axes=None): ... @@ -54,7 +61,7 @@ def shifted_fft(arr, axes=None): transformed array transformed array """ """ return fft.ifftshift(fft.fftn(fft.fftshift(arr, axes=axes), axes=axes), axes=axes) return ifftshift(fftn(fftshift(arr, axes=axes), axes=axes), axes=axes) def shifted_ifft(arr, axes=None): def shifted_ifft(arr, axes=None): ... @@ -68,4 +75,4 @@ def shifted_ifft(arr, axes=None): ... @@ -68,4 +75,4 @@ def shifted_ifft(arr, axes=None): Returns: Returns: transformed array transformed array """ """ return fft.fftshift(fft.ifftn(fft.ifftshift(arr, axes=axes), axes=axes), axes=axes) return fftshift(ifftn(ifftshift(arr, axes=axes), axes=axes), axes=axes)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!