Source code for openmdao.util.arrayutil
""" Some useful array utilities. """
from six.moves import range
#from numpy import ndarray, ones, zeros, cumsum
from itertools import product
#def to_slice(idxs):
#"""Convert an index array or list to a slice if possible. Otherwise,
#return the index array or list.
#"""
#if isinstance(idxs, slice):
#return idxs
#elif isinstance(idxs, ndarray) or isinstance(idxs, list):
#if len(idxs) == 1:
#return slice(idxs[0], idxs[0]+1)
#elif len(idxs) == 0:
#return slice(0,0)
#if isinstance(idxs, ndarray):
#imin = idxs.min()
#imax = idxs.max()
#else:
#imin = min(idxs)
#imax = max(idxs)
#stride = idxs[1]-idxs[0]
#if stride == 0:
#return idxs
#for i in range(len(idxs)):
#if i and idxs[i] - idxs[i-1] != stride:
#return idxs
#if stride < 0:
### negative strides cause some failures, so just do positive for now
##return slice(imax+1, imin, stride)
#return idxs
#else:
#return slice(imin, imax+1, stride)
#elif isinstance(idxs, int_types):
#return slice(idxs, idxs+1)
#else:
#raise RuntimeError("can't convert indices of type '%s' to a slice" %
#str(type(idxs)))
[docs]def array_idx_iter(shape):
"""
Return an iterator over the indices into a n-dimensional array.
Args
----
shape : tuple
shape of the array.
"""
for p in product(*[range(s) for s in shape]):
yield p
#def evenly_distrib_idxs(num_divisions, arr_size):
#"""
#Given a number of divisions and the size of an array, chop the array up
#into pieces according to number of divisions, keeping the distribution
#of entries as even as possible.
#Args
#----
#num_divisions : int
#Number of parts to divide the array into.
#arr_size : int
#Number of entries in the array.
#Returns
#-------
#tuple
#a tuple of (sizes, offsets), where sizes and offsets contain values for all
#divisions.
#"""
#base = arr_size / num_divisions
#leftover = arr_size % num_divisions
#sizes = ones(num_divisions, dtype="int") * base
## evenly distribute the remainder across size-leftover procs,
## instead of giving the whole remainder to one proc
#sizes[:leftover] += 1
#offsets = zeros(num_divisions, dtype="int")
#offsets[1:] = cumsum(sizes)[:-1]
#return sizes, offsets