Author: mattip <matti.pi...@gmail.com>
Branch: 
Changeset: r78690:c1b066339482
Date: 2015-07-27 23:32 +0300
http://bitbucket.org/pypy/pypy/changeset/c1b066339482/

Log:    merge nditer-revisited which provides 'buffered' flag and other
        small improvements

diff --git a/pypy/module/micronumpy/nditer.py b/pypy/module/micronumpy/nditer.py
--- a/pypy/module/micronumpy/nditer.py
+++ b/pypy/module/micronumpy/nditer.py
@@ -9,7 +9,8 @@
 from pypy.module.micronumpy.iterators import ArrayIter
 from pypy.module.micronumpy.strides import (calculate_broadcast_strides,
                                             shape_agreement, 
shape_agreement_multiple)
-from pypy.module.micronumpy.casting import find_binop_result_dtype
+from pypy.module.micronumpy.casting import (find_binop_result_dtype, 
+                    can_cast_array, can_cast_type)
 
 
 def parse_op_arg(space, name, w_op_flags, n, parse_one_arg):
@@ -108,9 +109,7 @@
         if item == 'external_loop':
             nditer.external_loop = True
         elif item == 'buffered':
-            raise oefmt(space.w_NotImplementedError,
-                'nditer buffered not implemented yet')
-            # For numpy compatability
+            # Each iterator should be 1d
             nditer.buffered = True
         elif item == 'c_index':
             nditer.tracked_index = 'C'
@@ -213,30 +212,6 @@
         return arr
 
 
-def get_iter(space, order, arr, shape, dtype, op_flags, base):
-    imp = arr.implementation
-    backward = is_backward(imp, order)
-    if arr.is_scalar():
-        return ConcreteIter(imp, 1, [], [], [], op_flags, base)
-    if (abs(imp.strides[0]) < abs(imp.strides[-1]) and not backward) or \
-       (abs(imp.strides[0]) > abs(imp.strides[-1]) and backward):
-        # flip the strides. Is this always true for multidimension?
-        strides = imp.strides[:]
-        backstrides = imp.backstrides[:]
-        shape = imp.shape[:]
-        strides.reverse()
-        backstrides.reverse()
-        shape.reverse()
-    else:
-        strides = imp.strides
-        backstrides = imp.backstrides
-    r = calculate_broadcast_strides(strides, backstrides, imp.shape,
-                                    shape, backward)
-    if len(shape) != len(r[0]):
-        # shape can be shorter when using an external loop, just return a view
-        return ConcreteIter(imp, imp.get_size(), imp.shape, r[0], r[1], 
op_flags, base)
-    return ConcreteIter(imp, imp.get_size(), shape, r[0], r[1], op_flags, base)
-
 def calculate_ndim(op_in, oa_ndim):
     if oa_ndim >=0:
         return oa_ndim
@@ -308,6 +283,8 @@
     of shape (4,3) by setting the offset to the beginning of the data at each 
iteration
     '''
     shape = [s+1 for s in old_iter.shape_m1]
+    if len(shape) < 1:
+        return old_iter
     strides = old_iter.strides
     backstrides = old_iter.backstrides
     if order == 'F':
@@ -328,6 +305,8 @@
         _shape = [shape[-1]]  + old_iter.slice_shape
         _backstride = [(shape[-1] - 1) * strides[-1]] + 
old_iter.slice_backstride
         fastest = shape[-1]
+    if fastest == 0:
+        return old_iter
     if flat:
         _shape = [support.product(_shape)]
         if len(_stride) > 1:
@@ -383,6 +362,10 @@
         self.done = False
         self.first_next = True
         self.op_axes = []
+        if not space.is_w(w_casting, space.w_None):
+            self.casting = space.str_w(w_casting)
+        else:
+            self.casting = 'safe'
         # convert w_seq operands to a list of W_NDimArray
         if space.isinstance_w(w_seq, space.w_tuple) or \
            space.isinstance_w(w_seq, space.w_list):
@@ -445,8 +428,15 @@
                     self.seq[i] = W_NDimArray.from_shape(space, self.shape, 
out_dtype)
                 else:
                     if not self.op_flags[i].broadcast:
-                        # Raises if ooutput cannot be broadcast
-                        shape_agreement(space, self.shape, self.seq[i], False)
+                        # Raises if output cannot be broadcast
+                        try:
+                            shape_agreement(space, self.shape, self.seq[i], 
False)
+                        except OperationError as e:
+                            raise oefmt(space.w_ValueError, "non-broadcastable"
+                                " output operand with shape %s doesn't match "
+                                "the broadcast shape %s", 
+                                str(self.seq[i].get_shape()),
+                                str(self.shape)) 
 
         if self.tracked_index != "":
             if self.order == "K":
@@ -465,14 +455,41 @@
                 if not self_d:
                     self.dtypes[i] = seq_d
                 elif self_d != seq_d:
-                    if not 'r' in self.op_flags[i].tmp_copy:
-                        raise oefmt(space.w_TypeError,
-                                    "Iterator operand required copying or "
-                                    "buffering for operand %d", i)
-                    impl = self.seq[i].implementation
-                    order = support.get_order_as_CF(impl.order, self.order)
-                    new_impl = impl.astype(space, self_d, order)
-                    self.seq[i] = W_NDimArray(new_impl)
+                        impl = self.seq[i].implementation
+                        order = support.get_order_as_CF(impl.order, self.order)
+                        if self.buffered or 'r' in self.op_flags[i].tmp_copy:
+                            if not can_cast_array(
+                                    space, self.seq[i], self_d, self.casting):
+                                raise oefmt(space.w_TypeError, "Iterator 
operand %d"
+                                    " dtype could not be cast from %s to %s"
+                                    " according to the rule '%s'", i, 
+                                    space.str_w(seq_d.descr_repr(space)),
+                                    space.str_w(self_d.descr_repr(space)),
+                                    self.casting)
+ 
+                            new_impl = impl.astype(space, self_d, 
order).copy(space)
+                            self.seq[i] = W_NDimArray(new_impl)
+                        else:
+                            raise oefmt(space.w_TypeError, "Iterator "
+                                "operand required copying or buffering, "
+                                "but neither copying nor buffering was "
+                                "enabled")
+                        if 'w' in self.op_flags[i].rw:
+                            if not can_cast_type(
+                                    space, self_d, seq_d, self.casting):
+                                raise oefmt(space.w_TypeError, "Iterator"
+                                    " requested dtype could not be cast from "
+                                    " %s to %s, the operand %d dtype, accord"
+                                    "ing to the rule '%s'", 
+                                    space.str_w(self_d.descr_repr(space)),
+                                    space.str_w(seq_d.descr_repr(space)),
+                                    i, self.casting)
+        elif self.buffered:
+            for i in range(len(self.seq)):
+                if i not in outargs:
+                    self.seq[i] = self.seq[i].descr_copy(space,
+                                     w_order=space.wrap(self.order))
+            self.dtypes = [s.get_dtype() for s in self.seq]
         else:
             #copy them from seq
             self.dtypes = [s.get_dtype() for s in self.seq]
@@ -480,14 +497,43 @@
         # create an iterator for each operand
         self.iters = []
         for i in range(len(self.seq)):
-            it = get_iter(space, self.order, self.seq[i], self.shape,
-                          self.dtypes[i], self.op_flags[i], self)
+            it = self.get_iter(space, i)
             it.contiguous = False
             self.iters.append((it, it.reset()))
 
         if self.external_loop:
             coalesce_axes(self, space)
 
+    def get_iter(self, space, i):
+        arr = self.seq[i]
+        dtype = self.dtypes[i]
+        shape = self.shape
+        imp = arr.implementation
+        backward = is_backward(imp, self.order)
+        if arr.is_scalar():
+            return ConcreteIter(imp, 1, [], [], [], self.op_flags[i], self)
+        if (abs(imp.strides[0]) < abs(imp.strides[-1]) and not backward) or \
+           (abs(imp.strides[0]) > abs(imp.strides[-1]) and backward):
+            # flip the strides. Is this always true for multidimension?
+            strides = imp.strides[:]
+            backstrides = imp.backstrides[:]
+            shape = imp.shape[:]
+            strides.reverse()
+            backstrides.reverse()
+            shape.reverse()
+        else:
+            strides = imp.strides
+            backstrides = imp.backstrides
+        r = calculate_broadcast_strides(strides, 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 set_op_axes(self, space, w_op_axes):
         if space.len_w(w_op_axes) != len(self.seq):
             raise oefmt(space.w_ValueError,
@@ -520,8 +566,8 @@
         return space.wrap(self)
 
     def getitem(self, it, st):
-        res = it.getoperand(st)
-        return W_NDimArray(res)
+        w_res = W_NDimArray(it.getoperand(st))
+        return w_res
 
     def descr_getitem(self, space, w_idx):
         idx = space.int_w(w_idx)
diff --git a/pypy/module/micronumpy/test/dummy_module.py 
b/pypy/module/micronumpy/test/dummy_module.py
--- a/pypy/module/micronumpy/test/dummy_module.py
+++ b/pypy/module/micronumpy/test/dummy_module.py
@@ -38,3 +38,6 @@
     a = zeros(*args, **kwargs)
     a.fill(1)
     return a
+
+def isscalar(a):
+    return type(a) in [typeinfo[t] for t in types]
diff --git a/pypy/module/micronumpy/test/test_nditer.py 
b/pypy/module/micronumpy/test/test_nditer.py
--- a/pypy/module/micronumpy/test/test_nditer.py
+++ b/pypy/module/micronumpy/test/test_nditer.py
@@ -64,18 +64,15 @@
         a = arange(24).reshape(2, 3, 4)
         import sys
         r = []
-        n = 0
         for x in nditer(a, flags=['external_loop']):
             r.append(x)
-            n += 1
-        assert n == 1
+        assert len(r) == 1
+        assert r[0].shape == (24,)
         assert (array(r) == range(24)).all()
         r = []
-        n = 0
         for x in nditer(a, flags=['external_loop'], order='F'):
             r.append(x)
-            n += 1
-        assert n == 12
+        assert len(r) == 12
         assert (array(r) == [[ 0, 12], [ 4, 16], [ 8, 20], [ 1, 13], [ 5, 17], 
[ 9, 21],
                              [ 2, 14], [ 6, 18], [10, 22], [ 3, 15], [ 7, 19], 
[11, 23],
                             ]).all()
@@ -149,19 +146,60 @@
         # assert str(exc.value).startswith("Iterator flag EXTERNAL_LOOP 
cannot")
 
     def test_buffered(self):
-        from numpy import arange, nditer, array
-        a = arange(6).reshape(2,3)
-        import sys
-        if '__pypy__' in sys.builtin_module_names:
-            raises(NotImplementedError, nditer, a, flags=['buffered'])
-            skip('nditer buffered not implmented')
+        from numpy import arange, nditer, array, isscalar
+        a = arange(24).reshape(2, 3, 4)
+        r = []
+        for x in nditer(a, flags=['external_loop'], order='F'):
+            r.append(x)
+        array_r = array(r)
+        assert len(array_r.shape) == 2
+        assert array_r.shape == (12, 2)
+        assert (array_r == [[0, 12], [4, 16], [8, 20], [1, 13], [5, 17], [9, 
21],
+                      [2, 14], [6, 18], [10, 22], [3, 15], [7, 19], [11, 
23]]).all
+        assert (a == arange(24).reshape(2, 3, 4)).all()
+        a[0,0,0] = 100
+        assert r[0][0] == 100
+
+        r = []
+        try:
+            it = nditer(a, flags=['buffered'], order='F')
+        except NotImplementedError as e:
+            assert 'unsupported value for order' in str(e)
+            skip('buffered with order="F" requires fortran tmp array creation')
+        for x in it:
+            r.append(x)
+        array_r = array(r)
+        assert len(array_r.shape) == 1
+        assert array_r.shape == (24,)
+        assert r[0].shape == ()
+        assert not isscalar(r[0])
+        assert (array_r == [0, 12, 4, 16, 8, 20, 1, 13, 5, 17, 9, 21,
+                      2, 14, 6, 18, 10, 22, 3, 15, 7, 19, 11, 23]).all
+        assert a.shape == (2, 3, 4)
+        a[0,0,0] = 0
+        # buffered copies the data into a tmp array
+        assert r[0] == 100
+        assert (a == arange(24).reshape(2, 3, 4)).all()
+
         r = []
         for x in nditer(a, flags=['external_loop', 'buffered'], order='F'):
             r.append(x)
-        array_r = array(r)
-        assert len(array_r.shape) == 2
-        assert array_r.shape == (1, 6)
-        assert (array_r == [0, 3, 1, 4, 2, 5]).all()
+        assert r[0].shape == (24,)
+        assert (array_r == [0, 12, 4, 16, 8, 20, 1, 13, 5, 17, 9, 21,
+                      2, 14, 6, 18, 10, 22, 3, 15, 7, 19, 11, 23]).all
+        assert a.shape == (2, 3, 4)
+        assert (a == arange(24).reshape(2, 3, 4)).all()
+
+    def test_zerosize(self):
+        from numpy import nditer, array
+        for a in [ array([]), array([1]), array([1, 2]) ]:
+            buffersize = max(16 * 1024 ** 2 // a.itemsize, 1)
+            r = []
+            for chunk in nditer(a, 
+                    flags=['external_loop', 'buffered', 'zerosize_ok'],
+                    buffersize=buffersize, order='C'):
+                r.append(chunk)
+            assert (r == a).all()
 
     def test_op_dtype(self):
         from numpy import arange, nditer, sqrt, array
@@ -188,11 +226,10 @@
         from numpy import arange, nditer
         import sys
         a = arange(6.)
-        if '__pypy__' in sys.builtin_module_names:
-            raises(NotImplementedError, nditer, a, flags=['buffered'], 
op_dtypes=['float32'])
-            skip('nditer casting not implemented yet')
         exc = raises(TypeError, nditer, a, flags=['buffered'], 
op_dtypes=['float32'])
-        assert str(exc.value).startswith("Iterator operand 0 dtype could not 
be cast")
+        assert str(exc.value) == "Iterator operand 0 dtype could not be " + \
+            "cast from dtype('float64') to dtype('float32') according to the" 
+\
+            " rule 'safe'"
         r = []
         for x in nditer(a, flags=['buffered'], op_dtypes=['float32'],
                                 casting='same_kind'):
@@ -232,9 +269,6 @@
             return it.operands[1]
         assert (square1([1, 2, 3]) == [1, 4, 9]).all()
 
-        if '__pypy__' in sys.builtin_module_names:
-            raises(NotImplementedError, nditer, [1, 2], flags=['buffered'])
-            skip('nditer buffered not implmented')
         def square2(a, out=None):
             it = nditer([a, out], flags=['external_loop', 'buffered'],
                         op_flags=[['readonly'],
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to