Author: mattip <matti.pi...@gmail.com>
Branch: 
Changeset: r83715:bbef742722ea
Date: 2016-04-17 16:37 +0300
http://bitbucket.org/pypy/pypy/changeset/bbef742722ea/

Log:    merge branch which provides numpy.broadcast

diff --git a/pypy/module/micronumpy/__init__.py 
b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -32,6 +32,7 @@
         'set_string_function': 'appbridge.set_string_function',
         'typeinfo': 'descriptor.get_dtype_cache(space).w_typeinfo',
         'nditer': 'nditer.W_NDIter',
+        'broadcast': 'broadcast.W_Broadcast',
 
         'set_docstring': 'support.descr_set_docstring',
         'VisibleDeprecationWarning': 'support.W_VisibleDeprecationWarning',
diff --git a/pypy/module/micronumpy/broadcast.py 
b/pypy/module/micronumpy/broadcast.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/micronumpy/broadcast.py
@@ -0,0 +1,110 @@
+import pypy.module.micronumpy.constants as NPY
+from nditer import ConcreteIter, parse_op_flag, parse_op_arg
+from pypy.interpreter.error import OperationError, oefmt
+from pypy.interpreter.gateway import interp2app
+from pypy.interpreter.typedef import TypeDef, GetSetProperty
+from pypy.module.micronumpy import support
+from pypy.module.micronumpy.base import W_NDimArray, convert_to_array, 
W_NumpyObject
+from rpython.rlib import jit
+from strides import calculate_broadcast_strides, shape_agreement_multiple
+
+def descr_new_broadcast(space, w_subtype, __args__):
+    return W_Broadcast(space, __args__.arguments_w)
+
+class W_Broadcast(W_NumpyObject):
+    """
+    Implementation of numpy.broadcast.
+    This class is a simplified version of nditer.W_NDIter with fixed iteration 
for broadcasted arrays.
+    """
+
+    def __init__(self, space, args):
+        num_args = len(args)
+        if not (2 <= num_args <= NPY.MAXARGS):
+            raise oefmt(space.w_ValueError,
+                                 "Need at least two and fewer than (%d) array 
objects.", NPY.MAXARGS)
+
+        self.seq = [convert_to_array(space, w_elem)
+                    for w_elem in args]
+
+        self.op_flags = parse_op_arg(space, 'op_flags', space.w_None,
+                                     len(self.seq), parse_op_flag)
+
+        self.shape = shape_agreement_multiple(space, self.seq, shape=None)
+        self.order = NPY.CORDER
+
+        self.iters = []
+        self.index = 0
+
+        try:
+            self.size = support.product_check(self.shape)
+        except OverflowError as e:
+            raise oefmt(space.w_ValueError, "broadcast dimensions too large.")
+        for i in range(len(self.seq)):
+            it = self.get_iter(space, i)
+            it.contiguous = False
+            self.iters.append((it, it.reset()))
+
+        self.done = False
+        pass
+
+    def get_iter(self, space, i):
+        arr = self.seq[i]
+        imp = arr.implementation
+        if arr.is_scalar():
+            return ConcreteIter(imp, 1, [], [], [], self.op_flags[i], self)
+        shape = self.shape
+
+        backward = imp.order != self.order
+
+        r = calculate_broadcast_strides(imp.strides, imp.backstrides, 
imp.shape,
+                                        shape, backward)
+
+        iter_shape = shape
+        if len(shape) != len(r[0]):
+            # shape can be shorter when using an external loop, just return a 
view
+            iter_shape = imp.shape
+        return ConcreteIter(imp, imp.get_size(), iter_shape, r[0], r[1],
+                            self.op_flags[i], self)
+
+    def descr_iter(self, space):
+        return space.wrap(self)
+
+    def descr_get_shape(self, space):
+        return space.newtuple([space.wrap(i) for i in self.shape])
+
+    def descr_get_size(self, space):
+        return space.wrap(self.size)
+
+    def descr_get_index(self, space):
+        return space.wrap(self.index)
+
+    def descr_get_numiter(self, space):
+        return space.wrap(len(self.iters))
+
+    @jit.unroll_safe
+    def descr_next(self, space):
+        if self.index >= self.size:
+            self.done = True
+            raise OperationError(space.w_StopIteration, space.w_None)
+        self.index += 1
+        res = []
+        for i, (it, st) in enumerate(self.iters):
+            res.append(self._get_item(it, st))
+            self.iters[i] = (it, it.next(st))
+        if len(res) < 2:
+            return res[0]
+        return space.newtuple(res)
+
+    def _get_item(self, it, st):
+        return W_NDimArray(it.getoperand(st))
+
+
+W_Broadcast.typedef = TypeDef("numpy.broadcast",
+                              __new__=interp2app(descr_new_broadcast),
+                              __iter__=interp2app(W_Broadcast.descr_iter),
+                              next=interp2app(W_Broadcast.descr_next),
+                              
shape=GetSetProperty(W_Broadcast.descr_get_shape),
+                              size=GetSetProperty(W_Broadcast.descr_get_size),
+                              
index=GetSetProperty(W_Broadcast.descr_get_index),
+                              
numiter=GetSetProperty(W_Broadcast.descr_get_numiter),
+                              )
diff --git a/pypy/module/micronumpy/constants.py 
b/pypy/module/micronumpy/constants.py
--- a/pypy/module/micronumpy/constants.py
+++ b/pypy/module/micronumpy/constants.py
@@ -77,6 +77,8 @@
 WRAP = 1
 RAISE = 2
 
+MAXARGS = 32
+
 # These can be requested in constructor functions and tested for
 ARRAY_C_CONTIGUOUS = 0x0001
 ARRAY_F_CONTIGUOUS = 0x0002
diff --git a/pypy/module/micronumpy/test/test_broadcast.py 
b/pypy/module/micronumpy/test/test_broadcast.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/micronumpy/test/test_broadcast.py
@@ -0,0 +1,97 @@
+# -*- encoding: utf-8 -*-
+
+from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
+
+
+class AppTestArrayBroadcast(BaseNumpyAppTest):
+    def test_broadcast_for_row_and_column(self):
+        import numpy as np
+        x = np.array([[1], [2], [3]])
+        y = np.array([4, 5])
+        b = list(np.broadcast(x, y))
+        assert b == [(1, 4), (1, 5), (2, 4), (2, 5), (3, 4), (3, 5)]
+
+    def test_broadcast_properties(self):
+        import numpy as np
+        x = np.array([[1], [2], [3]])
+        y = np.array([4, 5])
+        b = np.broadcast(x, y)
+
+        assert b.shape == (3, 2)
+        assert b.size == 6
+        assert b.index == 0
+
+        b.next()
+        b.next()
+
+        assert b.shape == (3, 2)
+        assert b.size == 6
+        assert b.index == 2
+
+    def test_broadcast_from_doctest(self):
+        """
+        Test from numpy.broadcast doctest.
+        """
+        import numpy as np
+        x = np.array([[1], [2], [3]])
+        y = np.array([4, 5, 6])
+        reference = np.array([[5., 6., 7.],
+                              [6., 7., 8.],
+                              [7., 8., 9.]])
+
+        b = np.broadcast(x, y)
+        out = np.empty(b.shape)
+        out.flat = [u + v for (u, v) in b]
+
+        assert (reference == out).all()
+        assert out.dtype == reference.dtype
+        assert b.shape == reference.shape
+
+    def test_broadcast_linear(self):
+        import numpy as np
+        x = np.array([1, 2, 3])
+        y = np.array([4, 5, 6])
+        b = list(np.broadcast(x, y))
+        assert b == [(1, 4), (2, 5), (3, 6)]
+        assert b[0][0].dtype == x.dtype
+
+    def test_broadcast_failures(self):
+        import numpy as np
+        import sys
+        x = np.array([1, 2, 3])
+        y = np.array([4, 5])
+        raises(ValueError, np.broadcast, x, y)
+        a = np.empty(2**16,dtype='int8')
+        a = a.reshape(-1, 1, 1, 1)
+        b = a.reshape(1, -1, 1, 1)
+        c = a.reshape(1, 1, -1, 1)
+        d = a.reshape(1, 1, 1, -1)
+        exc = raises(ValueError, np.broadcast, a, b, c, d)
+        assert exc.value[0] == ('broadcast dimensions too large.')
+
+    def test_broadcast_3_args(self):
+        import numpy as np
+        x = np.array([[[1]], [[2]], [[3]]])
+        y = np.array([[[40], [50]]])
+        z = np.array([[[700, 800]]])
+
+        b = list(np.broadcast(x, y, z))
+
+        assert b == [(1, 40, 700), (1, 40, 800), (1, 50, 700), (1, 50, 800),
+                     (2, 40, 700), (2, 40, 800), (2, 50, 700), (2, 50, 800),
+                     (3, 40, 700), (3, 40, 800), (3, 50, 700), (3, 50, 800)]
+
+    def test_number_of_arguments(self):
+        """
+        Test from numpy unit tests.
+        """
+        import numpy as np
+        arr = np.empty((5,))
+        for j in range(35):
+            arrs = [arr] * j
+            if j < 2 or j > 32:
+                exc = raises(ValueError, np.broadcast, *arrs)
+                assert exc.value[0] == ('Need at least two and fewer than (32) 
array objects.')
+            else:
+                mit = np.broadcast(*arrs)
+                assert mit.numiter == j
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to