crikit.utils.general — CRIKit2 0.4.3 documentation

# Source code for crikit.utils.general

    """
    General utilities

        expand_1d_to_ndim_data : Match 1D data array dimensionality to that of \
            another array

        expand_1d_to_ndim : Expand 1D data array dimensionality to ndim

        find_nearest : Given a vector and a value, find the index and value
            of the closest match

        pad : Wrapper around numpy.pad that also returns a window defining the
            original signal

    Notes
    -----
    """
    import numpy as _np

    def pad(y, pad_width, mode):
        """
        Pad array with either constants or edge values.

        Note: For N-D arrays, pads the -1 axis

        Parameters
        ----------
        y : ndarray
            Input array

        pad_width : int
            Size of padding on each side of y

        mode : str
            'constant' (0), 'edge' currently accepted

        Returns
        -------
        y_pad, window
            Padded array and window. Window defines the region of the original signal
        """
        if pad_width <= 0:
            return y, _np.ones(y.shape[-1])
        else:
            shaper = list(y.shape)
            shaper_out = list(y.shape)
            shaper_out[-1] += 2*pad_width
            y_pad = _np.zeros(shaper_out, dtype=y.dtype)
            window = _np.zeros(shaper_out[-1], dtype=_np.int32)

            y_pad[...,pad_width:shaper[-1]+pad_width] = 1*y
            window[pad_width:shaper[-1]+pad_width] = 1

            if (mode == 'zeros') | (mode == 'constant') | (mode == 'zero'):
                pass
            elif mode == 'edge':
                y_pad[...,:pad_width] = _np.dot(y[...,0:1], _np.ones((1, pad_width)))
                y_pad[..., -pad_width:] = _np.dot(y[...,-1:-2:-1], _np.ones((1, pad_width)))

            return y_pad, window

    def pad_dual(y, edge_pad_width, constant_pad_width):
        """
        Pad array with edge values followed by constant 0's.

        Note: For N-D arrays, pads the -1 axis

        Parameters
        ----------
        y : ndarray
            Input array

        edge_pad_width : int
            Size of edge-value padding on each side of y

        constant_pad_width : int
            Size of 0-padding on each side of y after edge-value padding

        Returns
        -------
        y_pad, window
            Padded array and window. Window defines the region of the original signal
        """
        y_pad_edge, win_edge = pad(y, edge_pad_width, 'edge')
        y_pad, win_constant = pad(y_pad_edge, constant_pad_width, 'constant')

        window = 0*win_constant
        window[_np.where(win_constant == 1)[0][win_edge == 1]] = 1

        return y_pad, window

    def pad_edge_mean(y, pad_width, n_edge=1, axis=-1):
        """
        Pad data y with edge-values or near-edge mean values along axis

        Parameters
        ----------

        y : ndarray
            Input array

        pad_width : int
            Size of padding on each side of y

        n_edge : int
            Number of edge points to average for the pad value

        axis : int
            Axis to pad

        Returns
        -------
        (y_pad, window)

        y_pad : ndarray
            Padded y

        window : ndarray (1D)
            Mask with 0's for pad regions, 1's for original size

        """
        if pad_width == 0:  # No padding
            window = _np.ones((y.shape[axis]), dtype=_np.int32)
            y_pad = y
        elif pad_width > 0:
            orig_shape = y.shape
            pad_shape = list(orig_shape)
            pad_shape[axis] += pad_width*2

            window = _np.zeros((pad_shape[axis]), dtype=_np.int32)
            window[pad_width:-pad_width] = 1

            y_pad = _np.zeros(pad_shape, dtype=y.dtype)
            slice_vec = y.ndim*[slice(None)]
            slice_vec[axis] = slice(pad_width,-pad_width)
            y_pad[tuple(slice_vec)] = y

            y_slice_vec_low = y.ndim*[slice(None)]
            y_slice_vec_low[axis] = slice(0,n_edge)
            y_slice_vec_high = y.ndim*[slice(None)]
            y_slice_vec_high[axis] = slice(-n_edge,None)

            y_pad_slice_vec_low = y.ndim*[slice(None)]
            y_pad_slice_vec_low[axis] = slice(0,pad_width)
            y_pad_slice_vec_high = y.ndim*[slice(None)]
            y_pad_slice_vec_high[axis] = slice(-pad_width,None)

            y_pad[tuple(y_pad_slice_vec_low)] += y[tuple(y_slice_vec_low)].mean(axis=axis, keepdims=True)
            y_pad[tuple(y_pad_slice_vec_high)] += y[tuple(y_slice_vec_high)].mean(axis=axis, keepdims=True)
        else:
            raise ValueError('pad_width must be >= 0')

        return y_pad, window

    def np_fcn_nd_to_1d(fcn, data, axis=-1):
        """
        Take in an n-dimensional array and return a 1D version operated on by fcn.\
        Works with many numpy functions that can take an "axis" parameter
        """
        if data.ndim > 1:
            dims = list(range(data.ndim))
            dims.pop(axis)
            vec = fcn(data, axis=tuple(dims))
        else:
            vec = data

        return vec

    def mean_nd_to_1d(data, axis=-1):
        """
        Take the mean of an nd array, except axis, returning a 1D array
        """
        vec = np_fcn_nd_to_1d(_np.mean, data, axis=axis)

        return vec

    def std_nd_to_1d(data, axis=-1):
        """
        Take the mean of an nd array, except axis, returning a 1D array
        """
        vec = np_fcn_nd_to_1d(_np.std, data, axis=axis)

        return vec

    def arange_nonzero(start, stop, dtype=float):
        """
        Similar to numpy arange but only returns non-zero elements
        """
        vec = _np.arange(start, stop+1)
        vec = vec[vec != 0]
        return vec

    def expand_1d_to_ndim_data(data, data_to_match):
        """
        Make 1D data array equal in dimensions to data_to_match
        """
        if data.ndim > 1:
            print('data must be 1D')
        else:
            nd = data_to_match.ndim
            return expand_1d_to_ndim(data, nd)

    def expand_1d_to_ndim(data, ndim):
        """
        Make 1D array into ndim dimensions
        """
        if data.ndim > 1:
            print('data must be 1D')
        else:
            sh = _np.ones((ndim-1),dtype=int)
            sh = _np.append(sh,-1)
            return data.reshape(sh)

    def find_nearest(np_vec,to_find = 0):
        """
        Given a vector and a value (or list/vector of values), find the index and
        value of the closest match

        Parameters
        ----------
        np_vec : numpy.ndarray
            Numpy array (list) of values

        to_find : int, float, numpy.ndarray, or list

        Returns
        -------
        out : tuple (nearest_value(s), index(es))
            Closest value (nearest_value) and index (index)

        """

        # Number of values (to_find) to find
        len_to_find = 0

        if isinstance(to_find, int) or isinstance(to_find, float):
            len_to_find = 1
        elif isinstance(to_find, list) or isinstance(to_find, tuple):
            len_to_find = len(to_find)
        elif isinstance(to_find, _np.ndarray):
            len_to_find = to_find.size
        else:
            pass

        if len_to_find == 0:
            return (None, None)
        elif len_to_find == 1:  # Single value
            test = _np.abs(_np.array(np_vec)-to_find)
            nearest_loc = test.argmin()
            nearest_val = np_vec[nearest_loc]
        else:  # Series of values
            nearest_val = []
            nearest_loc = []

            for val in to_find:
                loc = _np.argmin(_np.abs(_np.array(np_vec)-val))
                nearest_loc.append(loc)
                nearest_val.append(np_vec[loc])

        return (nearest_val, nearest_loc)

    def row_col_from_lin(ct, sh):
        """
        Convert a 1D counter into a col and row counter
        """

        assert len(sh) == 2, 'Shape must be 2D'

        tot_rows = sh[0]
        tot_cols = sh[1]

        if isinstance(ct, _np.ndarray):
            if (ct > tot_rows*tot_cols).any():
                print('Count is out-of-range. Returning None.')
                return None
        else:
            if ct > tot_rows*tot_cols:
                print('Count is out-of-range. Returning None.')
                return None

        row = _np.mod(ct, tot_rows)
        col = ct//tot_rows

        return [row, col]

    def lin_from_row_col(row, col, sh):
        """
        Convert a col and row counter to 1D linear count
        """

        assert len(sh) == 2, 'Shape must be 2D'

        tot_rows = sh[0]
        # tot_cols = sh[1]

        ct = col*tot_rows + row

        return ct

    if __name__ == '__main__':
        import timeit as _timeit

        print('Test 1.....')
        x = _np.random.rand(10,11)

        for ct in range(x.size):
            row, col = row_col_from_lin(ct, x.shape)
            print('R: {} C: {}'.format(row,col))
        print('Total number iterated through: {}'.format(ct+1))

        print('Test 2...')
        x = _np.random.rand(100,100,878)
        y = _np.zeros(x.shape, dtype=complex)

        tmr = _timeit.default_timer()
        for rc, blk in enumerate(x):
            for cc, sp in enumerate(blk):
                y[rc,cc,:] = _np.fft.fft(sp)
        tmr -= _timeit.default_timer()
        print('Time with 2 for-loops: {:.3g} sec'.format(-tmr))

        tmr = _timeit.default_timer()
        shp = x.shape
        x = x.reshape((-1, shp[-1]))
        y = _np.zeros(x.shape, dtype=complex)
        for num, sp in enumerate(x):
            y[num,:] = _np.fft.fft(sp)
        y = y.reshape(shp)
        tmr -= _timeit.default_timer()
        print('Time with reshaping and 1 for-loops: {:.3g} sec'.format(-tmr))
        x = x.reshape(shp)

        tmr = _timeit.default_timer()
        space_shp = _np.array(x.shape)[0:-1]
        num_sp = space_shp.prod()
        for num in range(num_sp):
            rc, cc = row_col_from_lin(num, space_shp)
            y[rc, cc, :] = _np.fft.fft(x[rc, cc, :])
        tmr -= _timeit.default_timer()
        print('Time with 1 for-loops: {:.3g} sec'.format(-tmr))