"""
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
[docs]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
[docs]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
[docs]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
[docs]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
[docs]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
[docs]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
[docs]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
[docs]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)
[docs]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)
[docs]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)
[docs]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]
[docs]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))