[pypy-commit] pypy missing-ndarray-attributes: shuffle stuff around and implement clip

fijal noreply at buildbot.pypy.org
Mon Oct 29 15:41:14 CET 2012


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: missing-ndarray-attributes
Changeset: r58585:2de8853d1f02
Date: 2012-10-29 15:40 +0100
http://bitbucket.org/pypy/pypy/changeset/2de8853d1f02/

Log:	shuffle stuff around and implement clip

diff --git a/pypy/module/micronumpy/interp_arrayops.py b/pypy/module/micronumpy/interp_arrayops.py
--- a/pypy/module/micronumpy/interp_arrayops.py
+++ b/pypy/module/micronumpy/interp_arrayops.py
@@ -1,8 +1,9 @@
 
 from pypy.module.micronumpy.base import convert_to_array, W_NDimArray
-from pypy.module.micronumpy import loop, interp_ufuncs
+from pypy.module.micronumpy import loop, interp_dtype, interp_ufuncs
 from pypy.module.micronumpy.iter import Chunk, Chunks
-from pypy.module.micronumpy.strides import shape_agreement
+from pypy.module.micronumpy.strides import shape_agreement,\
+     shape_agreement_multiple
 from pypy.module.micronumpy.constants import MODES
 from pypy.interpreter.error import OperationError, operationerrfmt
 from pypy.interpreter.gateway import unwrap_spec
@@ -161,19 +162,9 @@
     if not choices:
         raise OperationError(space.w_ValueError,
                              space.wrap("choices list cannot be empty"))
-    # find the shape agreement
-    shape = arr.get_shape()
-    for choice in choices:
-        shape = shape_agreement(space, shape, choice)
-    if out is not None:
-        shape = shape_agreement(space, shape, out)
-    # find the correct dtype
-    dtype = choices[0].get_dtype()
-    for choice in choices[1:]:
-        dtype = interp_ufuncs.find_binop_result_dtype(space,
-                                                      dtype, choice.get_dtype())
-    if out is None:
-        out = W_NDimArray.from_shape(shape, dtype)
+    shape = shape_agreement_multiple(space, choices + [out])
+    out = interp_dtype.dtype_agreement(space, choices, shape, out)
+    dtype = out.get_dtype()
     if mode not in MODES:
         raise OperationError(space.w_ValueError,
                              space.wrap("mode %s not known" % (mode,)))
diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -5,9 +5,10 @@
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.interpreter.typedef import (TypeDef, GetSetProperty,
     interp_attrproperty, interp_attrproperty_w)
-from pypy.module.micronumpy import types, interp_boxes
+from pypy.module.micronumpy import types, interp_boxes, base
 from pypy.rlib.objectmodel import specialize
 from pypy.rlib.rarithmetic import LONG_BIT, r_longlong, r_ulonglong
+from pypy.rlib import jit
 from pypy.rpython.lltypesystem import rffi
 
 
@@ -27,6 +28,21 @@
     return space.interp_w(W_Dtype,
           space.call_function(space.gettypefor(W_Dtype), w_dtype))
 
+ at jit.unroll_safe
+def dtype_agreement(space, w_arr_list, shape, out=None):
+    """ agree on dtype from a list of arrays. if out is allocated,
+    use it's dtype, otherwise allocate a new one with agreed dtype
+    """
+    from pypy.module.micronumpy.interp_ufuncs import find_binop_result_dtype
+
+    if not space.is_none(out):
+        return out
+    dtype = w_arr_list[0].get_dtype()
+    for w_arr in w_arr_list[1:]:
+        dtype = find_binop_result_dtype(space, dtype, w_arr.get_dtype())
+    out = base.W_NDimArray.from_shape(shape, dtype)
+    return out
+
 class W_Dtype(Wrappable):
     _immutable_fields_ = ["itemtype", "num", "kind"]
 
diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -7,7 +7,8 @@
 from pypy.module.micronumpy import interp_dtype, interp_ufuncs, interp_boxes,\
      interp_arrayops
 from pypy.module.micronumpy.strides import find_shape_and_elems,\
-     get_shape_from_iterable, to_coords, shape_agreement
+     get_shape_from_iterable, to_coords, shape_agreement, \
+     shape_agreement_multiple
 from pypy.module.micronumpy.interp_flatiter import W_FlatIterator
 from pypy.module.micronumpy.interp_support import unwrap_axis_arg
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
@@ -410,8 +411,16 @@
         return interp_arrayops.choose(space, self, w_choices, w_out, mode)
 
     def descr_clip(self, space, w_min, w_max, w_out=None):
-        raise OperationError(space.w_NotImplementedError, space.wrap(
-            "clip not implemented yet"))
+        if w_out is not None and not isinstance(w_out, W_NDimArray):
+            raise OperationError(space.w_TypeError, space.wrap(
+                "return arrays must be of ArrayType"))
+        min = convert_to_array(space, w_min)
+        max = convert_to_array(space, w_max)
+        shape = shape_agreement_multiple(space, [self, min, max, w_out])
+        out = interp_dtype.dtype_agreement(space, [self, min, max], shape,
+                                           w_out)
+        loop.clip(space, self, shape, min, max, out)
+        return out
 
     def descr_conj(self, space):
         raise OperationError(space.w_NotImplementedError, space.wrap(
@@ -801,6 +810,7 @@
     base     = GetSetProperty(W_NDimArray.descr_get_base),
     byteswap = interp2app(W_NDimArray.descr_byteswap),
     choose   = interp2app(W_NDimArray.descr_choose),
+    clip     = interp2app(W_NDimArray.descr_clip),
 
     __array_interface__ = GetSetProperty(W_NDimArray.descr_array_iface),
 )
diff --git a/pypy/module/micronumpy/loop.py b/pypy/module/micronumpy/loop.py
--- a/pypy/module/micronumpy/loop.py
+++ b/pypy/module/micronumpy/loop.py
@@ -522,3 +522,31 @@
             iter.next()
         out_iter.next()
         arr_iter.next()
+
+clip_driver = jit.JitDriver(greens = ['shapelen', 'dtype'],
+                            reds = ['min_iter', 'max_iter', 'arr_iter',
+                                    'out_iter'])
+
+def clip(space, arr, shape, min, max, out):
+    arr_iter = arr.create_iter(shape)
+    dtype = out.get_dtype()
+    shapelen = len(shape)
+    min_iter = min.create_iter(shape)
+    max_iter = max.create_iter(shape)
+    out_iter = out.create_iter(shape)
+    while not arr_iter.done():
+        clip_driver.jit_merge_point(shapelen=shapelen, dtype=dtype,
+                                    min_iter=min_iter, max_iter=max_iter,
+                                    arr_iter=arr_iter, out_iter=out_iter)
+        w_v = arr_iter.getitem().convert_to(dtype)
+        w_min = min_iter.getitem().convert_to(dtype)
+        w_max = max_iter.getitem().convert_to(dtype)
+        if dtype.itemtype.lt(w_v, w_min):
+            w_v = w_min
+        elif dtype.itemtype.gt(w_v, w_max):
+            w_v = w_max
+        out_iter.setitem(w_v)
+        arr_iter.next()
+        max_iter.next()
+        out_iter.next()
+        min_iter.next()
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -132,6 +132,17 @@
         )
     return ret
 
+ at jit.unroll_safe
+def shape_agreement_multiple(space, array_list):
+    """ call shape_agreement recursively, allow elements from array_list to
+    be None (like w_out)
+    """
+    shape = array_list[0].get_shape()
+    for arr in array_list[1:]:
+        if not space.is_none(arr):
+            shape = shape_agreement(space, shape, arr)
+    return shape
+
 def _shape_agreement(shape1, shape2):
     """ Checks agreement about two shapes with respect to broadcasting. Returns
     the resulting shape.
diff --git a/pypy/module/micronumpy/test/test_arrayops.py b/pypy/module/micronumpy/test/test_arrayops.py
--- a/pypy/module/micronumpy/test/test_arrayops.py
+++ b/pypy/module/micronumpy/test/test_arrayops.py
@@ -115,9 +115,15 @@
         r = array([4, 1, 0]).choose([a, b, c], mode='wrap')
         assert (r == [4, 5, 3]).all()
 
-
     def test_choose_dtype(self):
         from _numpypy import array
         a, b, c = array([1.2, 2, 3]), [4, 5, 6], 13
         r = array([2, 1, 0]).choose([a, b, c])
         assert r.dtype == float
+
+    def test_choose_dtype_out(self):
+        from _numpypy import array
+        a, b, c = array([1, 2, 3]), [4, 5, 6], 13
+        x = array([0, 0, 0], dtype='i2')
+        r = array([2, 1, 0]).choose([a, b, c], out=x)
+        assert r.dtype == 'i2'
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1599,6 +1599,15 @@
         assert (a.byteswap(True) == [0x0100, 0x0201, 0x0300]).all()
         assert (a == [0x0100, 0x0201, 0x0300]).all()
 
+    def test_clip(self):
+        from _numpypy import array
+        a = array([1, 2, 17, -3, 12])
+        assert (a.clip(-2, 13) == [1, 2, 13, -2, 12]).all()
+        assert (a.clip(-1, 1) == [1, 1, 1, -1, 1]).all()
+        assert (a.clip(-1, [1, 2, 3, 4, 5]) == [1, 2, 3, -1, 5]).all()
+        assert (a.clip(-2, 13, out=a) == [1, 2, 13, -2, 12]).all()
+        assert (a == [1, 2, 13, -2, 12]).all()
+
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):
         import _numpypy


More information about the pypy-commit mailing list