import numpy as np __all__ = ['bin_array'] def bin_array(arr: np.ndarray, new_shape: any, pad_zeros:bool=True) -> np.ndarray: """ Reduce the size of an array by binning :param arr: original :param new_shape: tuple which must be an integer divisor of the original shape, or integer to bin by that factor :param pad_zeros: pad array with zeros to enable binning by the given factor :return: new array """ # make tuple with new shape if type(new_shape) == int: # binning factor is given _shape = tuple([i // new_shape for i in arr.shape]) binfactor = tuple([new_shape for i in _shape]) else: _shape = new_shape binfactor = tuple([s // _shape[i] for i, s in enumerate(arr.shape)]) # determine if padding is needed padding = tuple([(0, (binfactor[i] - s % binfactor[i]) % binfactor[i]) for i, s in enumerate(arr.shape)]) if pad_zeros and np.any(np.array(padding) != 0): _arr = np.pad(arr, padding, mode='constant', constant_values=0) # pad array _shape = tuple([s // binfactor[i] for i, s in enumerate(_arr.shape)]) # update binned size due to padding else: _arr = arr # expected to fail if padding has non-zeros # send to 2d or 3d padding functions try: if len(arr.shape) == 2: out = bin_2d_array(_arr, _shape) elif len(arr.shape) == 3: out = bin_3d_array(_arr, _shape) else: raise NotImplementedError('Cannot only bin 3d or 2d arrays') return out except ValueError: raise ValueError("Cannot bin data with this shape. Try setting pad_zeros=True, or change the binning.") def bin_2d_array(arr: np.ndarray, new_shape: tuple) -> np.ndarray: """ bins a 2D numpy array Args: arr: input array to be binned new_shape: shape after binning, must be an integer divisor of the original shape Returns: binned np array """ shape = (new_shape[0], arr.shape[0] // new_shape[0], new_shape[1], arr.shape[1] // new_shape[1]) if np.any(np.isnan(arr)): binfactor = 1 for i, s in enumerate(arr.shape): binfactor *= new_shape[i] / s return np.nanmean(arr.reshape(shape), axis=(3, 1)) * binfactor else: return arr.reshape(shape).sum(-1).sum(1) def bin_3d_array(arr: np.ndarray, new_shape: tuple) -> np.ndarray: """" bins a 3D numpy array Args: arr: input array to be binned new_shape: shape after binning, must be an integer divisor of the original shape Returns: binned np array """ shape = (new_shape[0], arr.shape[0] // new_shape[0], new_shape[1], arr.shape[1] // new_shape[1], new_shape[2], arr.shape[2] // new_shape[2]) if np.any(np.isnan(arr)): binfactor = 1 for i, s in enumerate(arr.shape): binfactor *= new_shape[i] / s return np.nanmean(arr.reshape(shape), axis=(5, 3, 1)) * binfactor else: return np.sum(arr.reshape(shape), axis=(5, 3, 1))