Commit 0795b636 authored by Matthijs's avatar Matthijs
Browse files

odd and uncommon error due to center_of_mass fixed

parent 2ae1ddb6
from numpy import roll, ndarray, floor, iscomplexobj, round
from numpy import roll, ndarray, floor, iscomplexobj, round, any, isnan, nan_to_num
from scipy.ndimage.measurements import maximum_position, center_of_mass
from scipy.fftpack import fftn, fftshift, ifftn, ifftshift
from warnings import warn
__all__ = ["shift_array", 'roll_to_pos', 'shifted_ifft', 'shifted_fft']
......@@ -40,6 +41,9 @@ def roll_to_pos(arr: ndarray, y: int = 0, x: int = 0, pos: tuple = None, move_ma
old = floor(center_of_mass(abs(arr)))
old = floor(center_of_mass(arr))
if any(isnan(old)):
old = nan_to_num(old)
warn(Warning("Unexpected error in the calculation of the center of mass, casting NaNs to num"))
if pos is not None: # dimension-independent method
shifts = tuple([int(round(pos[i] - old[i])) for i in range(len(pos))])
dims = tuple([i for i in range(len(pos))])
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment