diff options
| -rw-r--r-- | doc/release/1.10.0-notes.rst | 6 | ||||
| -rw-r--r-- | doc/source/reference/routines.array-manipulation.rst | 1 | ||||
| -rw-r--r-- | numpy/lib/stride_tricks.py | 159 | ||||
| -rw-r--r-- | numpy/lib/tests/test_stride_tricks.py | 90 |
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() |