Merge pull request #5371 from shoyer/broadcast_to

ENH: add np.broadcast_to and reimplement np.broadcast_arrays

4 files changed, 199 insertions, 57 deletions

diff --git a/doc/release/1.10.0-notes.rst b/doc/release/1.10.0-notes.rst
index f9d202ec3..a07dca80f 100644
--- a/doc/release/1.10.0-notes.rst
+++ b/doc/release/1.10.0-notes.rst
@@ -94,6 +94,12 @@ number of rows read in a single call. Using this functionality, it is
 possible to read in multiple arrays stored in a single file by making
 repeated calls to the function.
 
+New function *np.broadcast_to* for invoking array broadcasting
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+*np.broadcast_to* manually broadcasts an array to a given shape according to
+numpy's broadcasting rules. The functionality is similar to broadcast_arrays,
+which in fact has been rewritten to use broadcast_to internally, but only a
+single array is necessary.
 
 Improvements
 ============
diff --git a/doc/source/reference/routines.array-manipulation.rst b/doc/source/reference/routines.array-manipulation.rst
index 81af0a315..2b3ba342a 100644
--- a/doc/source/reference/routines.array-manipulation.rst
+++ b/doc/source/reference/routines.array-manipulation.rst
@@ -40,6 +40,7 @@ Changing number of dimensions
    atleast_2d
    atleast_3d
    broadcast
+   broadcast_to
    broadcast_arrays
    expand_dims
    squeeze
diff --git a/numpy/lib/stride_tricks.py b/numpy/lib/stride_tricks.py
index b81307a65..0f46ed335 100644
--- a/numpy/lib/stride_tricks.py
+++ b/numpy/lib/stride_tricks.py
@@ -9,7 +9,8 @@ from __future__ import division, absolute_import, print_function
 
 import numpy as np
 
-__all__ = ['broadcast_arrays']
+__all__ = ['broadcast_to', 'broadcast_arrays']
+
 
 class DummyArray(object):
     """Dummy object that just exists to hang __array_interface__ dictionaries
@@ -20,6 +21,20 @@ class DummyArray(object):
         self.__array_interface__ = interface
         self.base = base
 
+
+def _maybe_view_as_subclass(original_array, new_array):
+    if type(original_array) is not type(new_array):
+        # if input was an ndarray subclass and subclasses were OK,
+        # then view the result as that subclass.
+        new_array = new_array.view(type=type(original_array))
+        # Since we have done something akin to a view from original_array, we
+        # should let the subclass finalize (if it has it implemented, i.e., is
+        # not None).
+        if new_array.__array_finalize__:
+            new_array.__array_finalize__(original_array)
+    return new_array
+
+
 def as_strided(x, shape=None, strides=None, subok=False):
     """ Make an ndarray from the given array with the given shape and strides.
     """
@@ -34,15 +49,80 @@ def as_strided(x, shape=None, strides=None, subok=False):
     # Make sure dtype is correct in case of custom dtype
     if array.dtype.kind == 'V':
         array.dtype = x.dtype
-    if type(x) is not type(array):
-        # if input was an ndarray subclass and subclasses were OK,
-        # then view the result as that subclass.
-        array = array.view(type=type(x))
-        # Since we have done something akin to a view from x, we should let
-        # the subclass finalize (if it has it implemented, i.e., is not None).
-        if array.__array_finalize__:
-            array.__array_finalize__(x)
-    return array
+    return _maybe_view_as_subclass(x, array)
+
+
+def _broadcast_to(array, shape, subok, readonly):
+    shape = tuple(shape) if np.iterable(shape) else (shape,)
+    array = np.array(array, copy=False, subok=subok)
+    if not shape and array.shape:
+        raise ValueError('cannot broadcast a non-scalar to a scalar array')
+    if any(size < 0 for size in shape):
+        raise ValueError('all elements of broadcast shape must be non-'
+                         'negative')
+    broadcast = np.nditer((array,), flags=['multi_index', 'zerosize_ok'],
+                          op_flags=['readonly'], itershape=shape, order='C'
+                          ).itviews[0]
+    result = _maybe_view_as_subclass(array, broadcast)
+    if not readonly and array.flags.writeable:
+        result.flags.writeable = True
+    return result
+
+
+def broadcast_to(array, shape, subok=False):
+    """Broadcast an array to a new shape.
+
+    Parameters
+    ----------
+    array : array_like
+        The array to broadcast.
+    shape : tuple
+        The shape of the desired array.
+    subok : bool, optional
+        If True, then sub-classes will be passed-through, otherwise
+        the returned array will be forced to be a base-class array (default).
+
+    Returns
+    -------
+    broadcast : array
+        A readonly view on the original array with the given shape. It is
+        typically not contiguous. Furthermore, more than one element of a
+        broadcasted array may refer to a single memory location.
+
+    Raises
+    ------
+    ValueError
+        If the array is not compatible with the new shape according to NumPy's
+        broadcasting rules.
+
+    Examples
+    --------
+    >>> x = np.array([1, 2, 3])
+    >>> np.broadcast_to(x, (3, 3))
+    array([[1, 2, 3],
+           [1, 2, 3],
+           [1, 2, 3]])
+    """
+    return _broadcast_to(array, shape, subok=subok, readonly=True)
+
+
+def _broadcast_shape(*args):
+    """Returns the shape of the ararys that would result from broadcasting the
+    supplied arrays against each other.
+    """
+    if not args:
+        raise ValueError('must provide at least one argument')
+    if len(args) == 1:
+        # a single argument does not work with np.broadcast
+        return np.asarray(args[0]).shape
+    # use the old-iterator because np.nditer does not handle size 0 arrays
+    # consistently
+    b = np.broadcast(*args[:32])
+    # unfortunately, it cannot handle 32 or more arguments directly
+    for pos in range(32, len(args), 31):
+        b = np.broadcast(b, *args[pos:(pos + 31)])
+    return b.shape
+
 
 def broadcast_arrays(*args, **kwargs):
     """
@@ -87,55 +167,24 @@ def broadcast_arrays(*args, **kwargs):
            [3, 3, 3]])]
 
     """
+    # nditer is not used here to avoid the limit of 32 arrays.
+    # Otherwise, something like the following one-liner would suffice:
+    # return np.nditer(args, flags=['multi_index', 'zerosize_ok'],
+    #                  order='C').itviews
+
     subok = kwargs.pop('subok', False)
     if kwargs:
         raise TypeError('broadcast_arrays() got an unexpected keyword '
                         'argument {}'.format(kwargs.pop()))
     args = [np.array(_m, copy=False, subok=subok) for _m in args]
-    shapes = [x.shape for x in args]
-    if len(set(shapes)) == 1:
+
+    shape = _broadcast_shape(*args)
+
+    if all(array.shape == shape for array in args):
         # Common case where nothing needs to be broadcasted.
         return args
-    shapes = [list(s) for s in shapes]
-    strides = [list(x.strides) for x in args]
-    nds = [len(s) for s in shapes]
-    biggest = max(nds)
-    # Go through each array and prepend dimensions of length 1 to each of
-    # the shapes in order to make the number of dimensions equal.
-    for i in range(len(args)):
-        diff = biggest - nds[i]
-        if diff > 0:
-            shapes[i] = [1] * diff + shapes[i]
-            strides[i] = [0] * diff + strides[i]
-    # Chech each dimension for compatibility. A dimension length of 1 is
-    # accepted as compatible with any other length.
-    common_shape = []
-    for axis in range(biggest):
-        lengths = [s[axis] for s in shapes]
-        unique = set(lengths + [1])
-        if len(unique) > 2:
-            # There must be at least two non-1 lengths for this axis.
-            raise ValueError("shape mismatch: two or more arrays have "
-                "incompatible dimensions on axis %r." % (axis,))
-        elif len(unique) == 2:
-            # There is exactly one non-1 length. The common shape will take
-            # this value.
-            unique.remove(1)
-            new_length = unique.pop()
-            common_shape.append(new_length)
-            # For each array, if this axis is being broadcasted from a
-            # length of 1, then set its stride to 0 so that it repeats its
-            # data.
-            for i in range(len(args)):
-                if shapes[i][axis] == 1:
-                    shapes[i][axis] = new_length
-                    strides[i][axis] = 0
-        else:
-            # Every array has a length of 1 on this axis. Strides can be
-            # left alone as nothing is broadcasted.
-            common_shape.append(1)
-
-    # Construct the new arrays.
-    broadcasted = [as_strided(x, shape=sh, strides=st, subok=subok)
-                   for (x, sh, st) in zip(args, shapes, strides)]
-    return broadcasted
+
+    # TODO: consider making the results of broadcast_arrays readonly to match
+    # broadcast_to. This will require a deprecation cycle.
+    return [_broadcast_to(array, shape, subok=subok, readonly=False)
+            for array in args]
diff --git a/numpy/lib/tests/test_stride_tricks.py b/numpy/lib/tests/test_stride_tricks.py
index bc7e30ca4..0b73109bc 100644
--- a/numpy/lib/tests/test_stride_tricks.py
+++ b/numpy/lib/tests/test_stride_tricks.py
@@ -5,8 +5,9 @@ from numpy.testing import (
     run_module_suite, assert_equal, assert_array_equal,
     assert_raises, assert_
     )
-from numpy.lib.stride_tricks import as_strided, broadcast_arrays
-
+from numpy.lib.stride_tricks import (
+    as_strided, broadcast_arrays, _broadcast_shape, broadcast_to
+)
 
 def assert_shapes_correct(input_shapes, expected_shape):
     # Broadcast a list of arrays with the given input shapes and check the
@@ -217,6 +218,62 @@ def test_same_as_ufunc():
             assert_same_as_ufunc(input_shapes[0], input_shapes[1], False, True)
             assert_same_as_ufunc(input_shapes[0], input_shapes[1], True, True)
 
+
+def test_broadcast_to_succeeds():
+    data = [
+        [np.array(0), (0,), np.array(0)],
+        [np.array(0), (1,), np.zeros(1)],
+        [np.array(0), (3,), np.zeros(3)],
+        [np.ones(1), (1,), np.ones(1)],
+        [np.ones(1), (2,), np.ones(2)],
+        [np.ones(1), (1, 2, 3), np.ones((1, 2, 3))],
+        [np.arange(3), (3,), np.arange(3)],
+        [np.arange(3), (1, 3), np.arange(3).reshape(1, -1)],
+        [np.arange(3), (2, 3), np.array([[0, 1, 2], [0, 1, 2]])],
+        # test if shape is not a tuple
+        [np.ones(0), 0, np.ones(0)],
+        [np.ones(1), 1, np.ones(1)],
+        [np.ones(1), 2, np.ones(2)],
+        # these cases with size 0 are strange, but they reproduce the behavior
+        # of broadcasting with ufuncs (see test_same_as_ufunc above)
+        [np.ones(1), (0,), np.ones(0)],
+        [np.ones((1, 2)), (0, 2), np.ones((0, 2))],
+        [np.ones((2, 1)), (2, 0), np.ones((2, 0))],
+    ]
+    for input_array, shape, expected in data:
+        actual = broadcast_to(input_array, shape)
+        assert_array_equal(expected, actual)
+
+
+def test_broadcast_to_raises():
+    data = [
+        [(0,), ()],
+        [(1,), ()],
+        [(3,), ()],
+        [(3,), (1,)],
+        [(3,), (2,)],
+        [(3,), (4,)],
+        [(1, 2), (2, 1)],
+        [(1, 1), (1,)],
+        [(1,), -1],
+        [(1,), (-1,)],
+        [(1, 2), (-1, 2)],
+    ]
+    for orig_shape, target_shape in data:
+        arr = np.zeros(orig_shape)
+        assert_raises(ValueError, lambda: broadcast_to(arr, target_shape))
+
+
+def test_broadcast_shape():
+    # broadcast_shape is already exercized indirectly by broadcast_arrays
+    assert_raises(ValueError, _broadcast_shape)
+    assert_equal(_broadcast_shape([1, 2]), (2,))
+    assert_equal(_broadcast_shape(np.ones((1, 1))), (1, 1))
+    assert_equal(_broadcast_shape(np.ones((1, 1)), np.ones((3, 4))), (3, 4))
+    assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 32)), (1, 2))
+    assert_equal(_broadcast_shape(*([np.ones((1, 2))] * 100)), (1, 2))
+
+
 def test_as_strided():
     a = np.array([None])
     a_view = as_strided(a)
@@ -277,6 +334,35 @@ def test_subclasses():
     assert_(type(b_view) is np.ndarray)
     assert_(a_view.shape == b_view.shape)
 
+    # and for broadcast_to
+    shape = (2, 4)
+    a_view = broadcast_to(a, shape)
+    assert_(type(a_view) is np.ndarray)
+    assert_(a_view.shape == shape)
+    a_view = broadcast_to(a, shape, subok=True)
+    assert_(type(a_view) is SimpleSubClass)
+    assert_(a_view.info == 'simple finalized')
+    assert_(a_view.shape == shape)
+
+
+def test_writeable():
+    # broadcast_to should return a readonly array
+    original = np.array([1, 2, 3])
+    result = broadcast_to(original, (2, 3))
+    assert_equal(result.flags.writeable, False)
+    assert_raises(ValueError, result.__setitem__, slice(None), 0)
+
+    # but the result of broadcast_arrays needs to be writeable (for now), to
+    # preserve backwards compatibility
+    for results in [broadcast_arrays(original),
+                    broadcast_arrays(0, original)]:
+        for result in results:
+            assert_equal(result.flags.writeable, True)
+    # keep readonly input readonly
+    original.flags.writeable = False
+    _, result = broadcast_arrays(0, original)
+    assert_equal(result.flags.writeable, False)
+
 
 if __name__ == "__main__":
     run_module_suite()