[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