Source code for openmdao.util.array_util
""" Some useful array utilities. """
import sys
from six.moves import range, zip
import numpy as np
from numpy import ndarray
from itertools import product
[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
[docs]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 = np.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 = np.zeros(num_divisions, dtype="int")
offsets[1:] = np.cumsum(sizes)[:-1]
return sizes, offsets
[docs]def to_slices(sidxs, didxs):
"""Convert matching src and dest idxs into slices if possible.
Sort the indices together to increase the likelihood of being able
to represent them as slices.
"""
sort_idxs = np.argsort(sidxs)
return _to_slice(sidxs[sort_idxs]), _to_slice(didxs[sort_idxs])
def _to_slice(idxs):
"""Convert an index array to a slice if possible. Otherwise,
return the index array. Indices are assumed to be sorted.
"""
if len(idxs) == 1:
return slice(idxs[0], idxs[0]+1)
elif len(idxs) == 0:
return idxs
stride = idxs[1]-idxs[0]
if stride <= 0:
return idxs
#make sure stride is consistent throughout the array
if any(idxs[1:]-idxs[:-1] != stride):
return idxs
return slice(idxs[0], idxs[-1]+1, stride)