Chris@87: """
Chris@87: Utilities that manipulate strides to achieve desirable effects.
Chris@87: 
Chris@87: An explanation of strides can be found in the "ndarray.rst" file in the
Chris@87: NumPy reference guide.
Chris@87: 
Chris@87: """
Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: import numpy as np
Chris@87: 
Chris@87: __all__ = ['broadcast_arrays']
Chris@87: 
Chris@87: class DummyArray(object):
Chris@87:     """Dummy object that just exists to hang __array_interface__ dictionaries
Chris@87:     and possibly keep alive a reference to a base array.
Chris@87:     """
Chris@87: 
Chris@87:     def __init__(self, interface, base=None):
Chris@87:         self.__array_interface__ = interface
Chris@87:         self.base = base
Chris@87: 
Chris@87: def as_strided(x, shape=None, strides=None):
Chris@87:     """ Make an ndarray from the given array with the given shape and strides.
Chris@87:     """
Chris@87:     interface = dict(x.__array_interface__)
Chris@87:     if shape is not None:
Chris@87:         interface['shape'] = tuple(shape)
Chris@87:     if strides is not None:
Chris@87:         interface['strides'] = tuple(strides)
Chris@87:     array = np.asarray(DummyArray(interface, base=x))
Chris@87:     # Make sure dtype is correct in case of custom dtype
Chris@87:     if array.dtype.kind == 'V':
Chris@87:         array.dtype = x.dtype
Chris@87:     return array
Chris@87: 
Chris@87: def broadcast_arrays(*args):
Chris@87:     """
Chris@87:     Broadcast any number of arrays against each other.
Chris@87: 
Chris@87:     Parameters
Chris@87:     ----------
Chris@87:     `*args` : array_likes
Chris@87:         The arrays to broadcast.
Chris@87: 
Chris@87:     Returns
Chris@87:     -------
Chris@87:     broadcasted : list of arrays
Chris@87:         These arrays are views on the original arrays.  They are typically
Chris@87:         not contiguous.  Furthermore, more than one element of a
Chris@87:         broadcasted array may refer to a single memory location.  If you
Chris@87:         need to write to the arrays, make copies first.
Chris@87: 
Chris@87:     Examples
Chris@87:     --------
Chris@87:     >>> x = np.array([[1,2,3]])
Chris@87:     >>> y = np.array([[1],[2],[3]])
Chris@87:     >>> np.broadcast_arrays(x, y)
Chris@87:     [array([[1, 2, 3],
Chris@87:            [1, 2, 3],
Chris@87:            [1, 2, 3]]), array([[1, 1, 1],
Chris@87:            [2, 2, 2],
Chris@87:            [3, 3, 3]])]
Chris@87: 
Chris@87:     Here is a useful idiom for getting contiguous copies instead of
Chris@87:     non-contiguous views.
Chris@87: 
Chris@87:     >>> [np.array(a) for a in np.broadcast_arrays(x, y)]
Chris@87:     [array([[1, 2, 3],
Chris@87:            [1, 2, 3],
Chris@87:            [1, 2, 3]]), array([[1, 1, 1],
Chris@87:            [2, 2, 2],
Chris@87:            [3, 3, 3]])]
Chris@87: 
Chris@87:     """
Chris@87:     args = [np.asarray(_m) for _m in args]
Chris@87:     shapes = [x.shape for x in args]
Chris@87:     if len(set(shapes)) == 1:
Chris@87:         # Common case where nothing needs to be broadcasted.
Chris@87:         return args
Chris@87:     shapes = [list(s) for s in shapes]
Chris@87:     strides = [list(x.strides) for x in args]
Chris@87:     nds = [len(s) for s in shapes]
Chris@87:     biggest = max(nds)
Chris@87:     # Go through each array and prepend dimensions of length 1 to each of
Chris@87:     # the shapes in order to make the number of dimensions equal.
Chris@87:     for i in range(len(args)):
Chris@87:         diff = biggest - nds[i]
Chris@87:         if diff > 0:
Chris@87:             shapes[i] = [1] * diff + shapes[i]
Chris@87:             strides[i] = [0] * diff + strides[i]
Chris@87:     # Chech each dimension for compatibility. A dimension length of 1 is
Chris@87:     # accepted as compatible with any other length.
Chris@87:     common_shape = []
Chris@87:     for axis in range(biggest):
Chris@87:         lengths = [s[axis] for s in shapes]
Chris@87:         unique = set(lengths + [1])
Chris@87:         if len(unique) > 2:
Chris@87:             # There must be at least two non-1 lengths for this axis.
Chris@87:             raise ValueError("shape mismatch: two or more arrays have "
Chris@87:                 "incompatible dimensions on axis %r." % (axis,))
Chris@87:         elif len(unique) == 2:
Chris@87:             # There is exactly one non-1 length. The common shape will take
Chris@87:             # this value.
Chris@87:             unique.remove(1)
Chris@87:             new_length = unique.pop()
Chris@87:             common_shape.append(new_length)
Chris@87:             # For each array, if this axis is being broadcasted from a
Chris@87:             # length of 1, then set its stride to 0 so that it repeats its
Chris@87:             # data.
Chris@87:             for i in range(len(args)):
Chris@87:                 if shapes[i][axis] == 1:
Chris@87:                     shapes[i][axis] = new_length
Chris@87:                     strides[i][axis] = 0
Chris@87:         else:
Chris@87:             # Every array has a length of 1 on this axis. Strides can be
Chris@87:             # left alone as nothing is broadcasted.
Chris@87:             common_shape.append(1)
Chris@87: 
Chris@87:     # Construct the new arrays.
Chris@87:     broadcasted = [as_strided(x, shape=sh, strides=st) for (x, sh, st) in
Chris@87:                    zip(args, shapes, strides)]
Chris@87:     return broadcasted