[pypy-commit] pypy missing-ndarray-attributes: cumsum and cumprod
fijal
noreply at buildbot.pypy.org
Wed Oct 31 16:04:52 CET 2012
Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: missing-ndarray-attributes
Changeset: r58642:9010050912c8
Date: 2012-10-31 17:04 +0200
http://bitbucket.org/pypy/pypy/changeset/9010050912c8/
Log: cumsum and cumprod
diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -220,8 +220,8 @@
backstrides)
return loop.setslice(self.get_shape(), impl, self)
- def create_axis_iter(self, shape, dim):
- return iter.AxisIterator(self, shape, dim)
+ def create_axis_iter(self, shape, dim, cum):
+ return iter.AxisIterator(self, shape, dim, cum)
def create_dot_iter(self, shape, skip):
r = calculate_dot_strides(self.get_strides(), self.get_backstrides(),
diff --git a/pypy/module/micronumpy/arrayimpl/scalar.py b/pypy/module/micronumpy/arrayimpl/scalar.py
--- a/pypy/module/micronumpy/arrayimpl/scalar.py
+++ b/pypy/module/micronumpy/arrayimpl/scalar.py
@@ -83,7 +83,7 @@
def reshape(self, space, orig_array, new_shape):
return self.set_shape(space, orig_array, new_shape)
- def create_axis_iter(self, shape, dim):
+ def create_axis_iter(self, shape, dim, cum):
raise Exception("axis iter should not happen on scalar")
def swapaxes(self, axis1, axis2):
diff --git a/pypy/module/micronumpy/arrayimpl/sort.py b/pypy/module/micronumpy/arrayimpl/sort.py
--- a/pypy/module/micronumpy/arrayimpl/sort.py
+++ b/pypy/module/micronumpy/arrayimpl/sort.py
@@ -100,9 +100,9 @@
if axis < 0 or axis > len(shape):
raise OperationError(space.w_IndexError("Wrong axis %d" % axis))
iterable_shape = shape[:axis] + [0] + shape[axis + 1:]
- iter = AxisIterator(arr, iterable_shape, axis)
+ iter = AxisIterator(arr, iterable_shape, axis, False)
index_impl = index_arr.implementation
- index_iter = AxisIterator(index_impl, iterable_shape, axis)
+ index_iter = AxisIterator(index_impl, iterable_shape, axis, False)
stride_size = arr.strides[axis]
index_stride_size = index_impl.strides[axis]
axis_size = arr.shape[axis]
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
@@ -210,8 +210,8 @@
def create_iter(self, shape=None):
return self.implementation.create_iter(shape)
- def create_axis_iter(self, shape, dim):
- return self.implementation.create_axis_iter(shape, dim)
+ def create_axis_iter(self, shape, dim, cum):
+ return self.implementation.create_axis_iter(shape, dim, cum)
def create_dot_iter(self, shape, skip):
return self.implementation.create_dot_iter(shape, skip)
@@ -430,10 +430,6 @@
raise OperationError(space.w_NotImplementedError, space.wrap(
"cumprod not implemented yet"))
- def descr_cumsum(self, space, w_axis=None, w_dtype=None, w_out=None):
- raise OperationError(space.w_NotImplementedError, space.wrap(
- "cumsum not implemented yet"))
-
def descr_data(self, space):
raise OperationError(space.w_NotImplementedError, space.wrap(
"data not implemented yet"))
@@ -644,7 +640,8 @@
# ----------------------- reduce -------------------------------
- def _reduce_ufunc_impl(ufunc_name, promote_to_largest=False):
+ def _reduce_ufunc_impl(ufunc_name, promote_to_largest=False,
+ cumultative=False):
def impl(self, space, w_axis=None, w_out=None, w_dtype=None):
if space.is_none(w_out):
out = None
@@ -653,9 +650,9 @@
'output must be an array'))
else:
out = w_out
- return getattr(interp_ufuncs.get(space), ufunc_name).reduce(space,
- self, True, promote_to_largest, w_axis,
- False, out, w_dtype)
+ return getattr(interp_ufuncs.get(space), ufunc_name).reduce(
+ space, self, True, promote_to_largest, w_axis,
+ False, out, w_dtype, cumultative=cumultative)
return func_with_new_name(impl, "reduce_%s_impl" % ufunc_name)
descr_sum = _reduce_ufunc_impl("add")
@@ -666,6 +663,9 @@
descr_all = _reduce_ufunc_impl('logical_and')
descr_any = _reduce_ufunc_impl('logical_or')
+ descr_cumsum = _reduce_ufunc_impl('add', cumultative=True)
+ descr_cumprod = _reduce_ufunc_impl('multiply', cumultative=True)
+
def descr_mean(self, space, w_axis=None, w_out=None):
if space.is_none(w_axis):
w_denom = space.wrap(self.get_size())
@@ -779,6 +779,9 @@
var = interp2app(W_NDimArray.descr_var),
std = interp2app(W_NDimArray.descr_std),
+ cumsum = interp2app(W_NDimArray.descr_cumsum),
+ cumprod = interp2app(W_NDimArray.descr_cumprod),
+
copy = interp2app(W_NDimArray.descr_copy),
reshape = interp2app(W_NDimArray.descr_reshape),
T = GetSetProperty(W_NDimArray.descr_get_transpose),
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -144,7 +144,7 @@
w_dtype)
def reduce(self, space, w_obj, multidim, promote_to_largest, w_axis,
- keepdims=False, out=None, dtype=None):
+ keepdims=False, out=None, dtype=None, cumultative=False):
if self.argcount != 2:
raise OperationError(space.w_ValueError, space.wrap("reduce only "
"supported for binary functions"))
@@ -158,7 +158,7 @@
return obj.get_scalar_value()
shapelen = len(obj_shape)
axis = unwrap_axis_arg(space, shapelen, w_axis)
- assert axis>=0
+ assert axis >= 0
size = obj.get_size()
dtype = interp_dtype.decode_w_dtype(space, dtype)
if dtype is None:
@@ -175,7 +175,14 @@
raise operationerrfmt(space.w_ValueError, "zero-size array to "
"%s.reduce without identity", self.name)
if shapelen > 1 and axis < shapelen:
- if keepdims:
+ temp = None
+ if cumultative:
+ shape = obj_shape[:]
+ temp_shape = obj_shape[:axis] + obj_shape[axis + 1:]
+ if out:
+ dtype = out.get_dtype()
+ temp = W_NDimArray.from_shape(temp_shape, dtype)
+ elif keepdims:
shape = obj_shape[:axis] + [1] + obj_shape[axis + 1:]
else:
shape = obj_shape[:axis] + obj_shape[axis + 1:]
@@ -202,7 +209,17 @@
else:
out = W_NDimArray.from_shape(shape, dtype)
return loop.do_axis_reduce(shape, self.func, obj, dtype, axis, out,
- self.identity)
+ self.identity, cumultative, temp)
+ if cumultative:
+ if out:
+ if out.get_shape() != [obj.get_size()]:
+ raise OperationError(space.w_ValueError, space.wrap(
+ "out of incompatible size"))
+ else:
+ out = W_NDimArray.from_shape([obj.get_size()], dtype)
+ loop.compute_reduce_cumultative(obj, out, dtype, self.func,
+ self.identity)
+ return out
if out:
if len(out.get_shape())>0:
raise operationerrfmt(space.w_ValueError, "output parameter "
diff --git a/pypy/module/micronumpy/iter.py b/pypy/module/micronumpy/iter.py
--- a/pypy/module/micronumpy/iter.py
+++ b/pypy/module/micronumpy/iter.py
@@ -259,11 +259,14 @@
self.offset %= self.size
class AxisIterator(base.BaseArrayIterator):
- def __init__(self, array, shape, dim):
+ def __init__(self, array, shape, dim, cumultative):
self.shape = shape
strides = array.strides
backstrides = array.backstrides
- if len(shape) == len(strides):
+ if cumultative:
+ self.strides = strides
+ self.backstrides = backstrides
+ elif len(shape) == len(strides):
# keepdims = True
self.strides = strides[:dim] + [0] + strides[dim + 1:]
self.backstrides = backstrides[:dim] + [0] + backstrides[dim + 1:]
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
@@ -112,6 +112,24 @@
obj_iter.next()
return cur_value
+reduce_cum_driver = jit.JitDriver(greens = ['shapelen', 'func', 'dtype'],
+ reds = ['obj_iter', 'out_iter'])
+
+def compute_reduce_cumultative(obj, out, calc_dtype, func, identity):
+ obj_iter = obj.create_iter()
+ out_iter = out.create_iter()
+ cur_value = identity.convert_to(calc_dtype)
+ shapelen = len(obj.get_shape())
+ while not obj_iter.done():
+ reduce_cum_driver.jit_merge_point(shapelen=shapelen, func=func,
+ dtype=calc_dtype, obj_iter=obj_iter,
+ out_iter=out_iter)
+ rval = obj_iter.getitem().convert_to(calc_dtype)
+ cur_value = func(calc_dtype, cur_value, rval)
+ out_iter.setitem(cur_value)
+ out_iter.next()
+ obj_iter.next()
+
def fill(arr, box):
arr_iter = arr.create_iter()
while not arr_iter.done():
@@ -156,13 +174,20 @@
return out
axis_reduce__driver = jit.JitDriver(name='numpy_axis_reduce',
- greens=['shapelen', 'func', 'dtype',
+ greens=['shapelen', 'cumultative',
+ 'func', 'dtype',
'identity'],
reds=['axis', 'arr', 'out', 'shape',
- 'out_iter', 'arr_iter'])
+ 'out_iter', 'arr_iter',
+ 'temp_iter'])
-def do_axis_reduce(shape, func, arr, dtype, axis, out, identity):
- out_iter = out.create_axis_iter(arr.get_shape(), axis)
+def do_axis_reduce(shape, func, arr, dtype, axis, out, identity, cumultative,
+ temp):
+ out_iter = out.create_axis_iter(arr.get_shape(), axis, cumultative)
+ if cumultative:
+ temp_iter = temp.create_axis_iter(arr.get_shape(), axis, False)
+ else:
+ temp_iter = out_iter # hack
arr_iter = arr.create_iter()
if identity is not None:
identity = identity.convert_to(dtype)
@@ -172,15 +197,20 @@
dtype=dtype, identity=identity,
axis=axis, arr=arr, out=out,
shape=shape, out_iter=out_iter,
- arr_iter=arr_iter)
+ arr_iter=arr_iter,
+ cumultative=cumultative,
+ temp_iter=temp_iter)
w_val = arr_iter.getitem().convert_to(dtype)
if out_iter.first_line:
if identity is not None:
w_val = func(dtype, identity, w_val)
else:
- cur = out_iter.getitem()
+ cur = temp_iter.getitem()
w_val = func(dtype, cur, w_val)
out_iter.setitem(w_val)
+ if cumultative:
+ temp_iter.setitem(w_val)
+ temp_iter.next()
arr_iter.next()
out_iter.next()
return out
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
@@ -2077,6 +2077,29 @@
a[array([0, 2]), slice(0, 2)] = [[10, 11], [12, 13]]
assert (a == [[10, 11], [3, 4], [12, 13]]).all()
+ def test_cumsum(self):
+ from _numpypy import arange
+ a = arange(6).reshape(3, 2)
+ b = arange(6)
+ assert (a.cumsum() == [0, 1, 3, 6, 10, 15]).all()
+ a.cumsum(out=b)
+ assert (b == [0, 1, 3, 6, 10, 15]).all()
+ raises(ValueError, "a.cumsum(out=arange(6).reshape(3, 2))")
+
+ def test_cumprod(self):
+ from _numpypy import array
+ a = array([[1, 2], [3, 4], [5, 6]])
+ assert (a.cumprod() == [1, 2, 6, 24, 120, 720]).all()
+
+ def test_cumsum_axis(self):
+ from _numpypy import arange, array
+ a = arange(6).reshape(3, 2)
+ assert (a.cumsum(0) == [[0, 1], [2, 4], [6, 9]]).all()
+ assert (a.cumsum(1) == [[0, 1], [2, 5], [4, 9]]).all()
+ a = array([[1, 1], [2, 2], [3, 4]])
+ assert (a.cumsum(1) == [[1, 2], [2, 4], [3, 7]]).all()
+ assert (a.cumsum(0) == [[1, 1], [3, 3], [6, 7]]).all()
+
class AppTestSupport(BaseNumpyAppTest):
def setup_class(cls):
import struct
More information about the pypy-commit
mailing list