Author: Brian Kearns <[email protected]>
Branch:
Changeset: r69940:bb85257cb544
Date: 2014-03-13 15:17 -0400
http://bitbucket.org/pypy/pypy/changeset/bb85257cb544/
Log: avoid loops in numpy array creation
diff --git a/pypy/module/micronumpy/ctors.py b/pypy/module/micronumpy/ctors.py
--- a/pypy/module/micronumpy/ctors.py
+++ b/pypy/module/micronumpy/ctors.py
@@ -5,7 +5,6 @@
from pypy.module.micronumpy import descriptor, loop, ufuncs
from pypy.module.micronumpy.base import W_NDimArray, convert_to_array
from pypy.module.micronumpy.converters import shape_converter
-from pypy.module.micronumpy.strides import find_shape_and_elems
def build_scalar(space, w_dtype, w_state):
@@ -27,6 +26,8 @@
@unwrap_spec(ndmin=int, copy=bool, subok=bool)
def array(space, w_object, w_dtype=None, copy=True, w_order=None, subok=False,
ndmin=0):
+ from pypy.module.micronumpy import strides
+
# for anything that isn't already an array, try __array__ method first
if not isinstance(w_object, W_NDimArray):
w___array__ = space.lookup(w_object, "__array__")
@@ -68,12 +69,9 @@
return w_ret
# not an array or incorrect dtype
- shape, elems_w = find_shape_and_elems(space, w_object, dtype)
+ shape, elems_w = strides.find_shape_and_elems(space, w_object, dtype)
if dtype is None or (dtype.is_str_or_unicode() and dtype.elsize < 1):
- for w_elem in elems_w:
- if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
- w_elem = w_elem.get_scalar_value()
- dtype = ufuncs.find_dtype_for_scalar(space, w_elem, dtype)
+ dtype = strides.find_dtype_for_seq(space, elems_w, dtype)
if dtype is None:
dtype = descriptor.get_dtype_cache(space).w_float64dtype
elif dtype.is_str_or_unicode() and dtype.elsize < 1:
@@ -83,10 +81,10 @@
if ndmin > len(shape):
shape = [1] * (ndmin - len(shape)) + shape
w_arr = W_NDimArray.from_shape(space, shape, dtype, order=order)
- arr_iter = w_arr.create_iter()
- for w_elem in elems_w:
- arr_iter.setitem(dtype.coerce(space, w_elem))
- arr_iter.next()
+ if len(elems_w) == 1:
+ w_arr.set_scalar_value(dtype.coerce(space, elems_w[0]))
+ else:
+ loop.assign(space, w_arr, elems_w)
return w_arr
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
@@ -164,6 +164,13 @@
arr_iter.setitem(box)
arr_iter.next()
+def assign(space, arr, seq):
+ arr_iter = arr.create_iter()
+ arr_dtype = arr.get_dtype()
+ for item in seq:
+ arr_iter.setitem(arr_dtype.coerce(space, item))
+ arr_iter.next()
+
where_driver = jit.JitDriver(name='numpy_where',
greens = ['shapelen', 'dtype', 'arr_dtype'],
reds = 'auto')
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
@@ -181,6 +181,10 @@
return [], [w_iterable]
if isinstance(w_iterable, W_NDimArray) and w_iterable.is_scalar():
return [], [w_iterable]
+ return _find_shape_and_elems(space, w_iterable, is_rec_type)
+
+
+def _find_shape_and_elems(space, w_iterable, is_rec_type):
shape = [space.len_w(w_iterable)]
batch = space.listview(w_iterable)
while True:
@@ -210,6 +214,25 @@
batch = new_batch
+def find_dtype_for_seq(space, elems_w, dtype):
+ from pypy.module.micronumpy.ufuncs import find_dtype_for_scalar
+ if len(elems_w) == 1:
+ w_elem = elems_w[0]
+ if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
+ w_elem = w_elem.get_scalar_value()
+ return find_dtype_for_scalar(space, w_elem, dtype)
+ return _find_dtype_for_seq(space, elems_w, dtype)
+
+
+def _find_dtype_for_seq(space, elems_w, dtype):
+ from pypy.module.micronumpy.ufuncs import find_dtype_for_scalar
+ for w_elem in elems_w:
+ if isinstance(w_elem, W_NDimArray) and w_elem.is_scalar():
+ w_elem = w_elem.get_scalar_value()
+ dtype = find_dtype_for_scalar(space, w_elem, dtype)
+ return dtype
+
+
def to_coords(space, shape, size, order, w_item_or_slice):
'''Returns a start coord, step, and length.
'''
_______________________________________________
pypy-commit mailing list
[email protected]
https://mail.python.org/mailman/listinfo/pypy-commit
