Author: Matti Picus <[email protected]>
Branch: missing-ndarray-attributes
Changeset: r60852:33c8a9999dba
Date: 2013-02-04 01:08 +0200
http://bitbucket.org/pypy/pypy/changeset/33c8a9999dba/
Log: simplify complex, extend argsort to work for complex
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
@@ -16,8 +16,9 @@
INT_SIZE = rffi.sizeof(lltype.Signed)
-def make_sort_function(space, itemtype):
+def make_sort_function(space, itemtype, count=1):
TP = itemtype.T
+ step = rffi.sizeof(TP)
class Repr(object):
def __init__(self, index_stride_size, stride_size, size, values,
@@ -31,16 +32,29 @@
self.indexes = indexes
def getitem(self, item):
- v = raw_storage_getitem(TP, self.values, item * self.stride_size
+ if count < 2:
+ v = raw_storage_getitem(TP, self.values, item *
self.stride_size
+ self.start)
- v = itemtype.for_computation(v)
+ v = itemtype.for_computation(v)
+ else:
+ v = []
+ for i in range(count):
+ _v = raw_storage_getitem(TP, self.values, item *
self.stride_size
+ + self.start + step * i)
+ v.append(_v)
+ v = itemtype.for_computation(v)
return (v, raw_storage_getitem(lltype.Signed, self.indexes,
item * self.index_stride_size +
self.index_start))
def setitem(self, idx, item):
- raw_storage_setitem(self.values, idx * self.stride_size +
+ if count < 2:
+ raw_storage_setitem(self.values, idx * self.stride_size +
self.start, rffi.cast(TP, item[0]))
+ else:
+ for i in range(count):
+ raw_storage_setitem(self.values, idx * self.stride_size +
+ self.start + i*step, rffi.cast(TP, item[0][i]))
raw_storage_setitem(self.indexes, idx * self.index_stride_size +
self.index_start, item[1])
@@ -49,7 +63,8 @@
start = 0
dtype = interp_dtype.get_dtype_cache(space).w_longdtype
self.indexes = dtype.itemtype.malloc(size*dtype.get_size())
- self.values = alloc_raw_storage(size*rffi.sizeof(TP),
track_allocation=False)
+ self.values = alloc_raw_storage(size * stride_size,
+ track_allocation=False)
Repr.__init__(self, index_stride_size, stride_size,
size, self.values, self.indexes, start, start)
@@ -152,6 +167,9 @@
self.built = True
cache = {}
for cls in all_types._items:
- cache[cls] = make_sort_function(space, cls)
+ if cls in types.all_complex_types:
+ cache[cls] = make_sort_function(space, cls, 2)
+ else:
+ cache[cls] = make_sort_function(space, cls)
self.cache = cache
self._lookup = specialize.memo()(lambda tp : cache[tp])
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -1080,7 +1080,7 @@
return bool(v[0]) or bool(v[1])
def get_element_size(self):
- return 2 * rffi.sizeof(self._COMPONENTS_T)
+ return 2 * rffi.sizeof(self.T)
def byteswap(self, w_v):
real, imag = self.unbox(w_v)
@@ -1089,19 +1089,19 @@
@specialize.argtype(1)
def box(self, value):
return self.BoxType(
- rffi.cast(self._COMPONENTS_T, value),
- rffi.cast(self._COMPONENTS_T, 0.0))
+ rffi.cast(self.T, value),
+ rffi.cast(self.T, 0.0))
@specialize.argtype(1)
def box_component(self, value):
return self.ComponentBoxType(
- rffi.cast(self._COMPONENTS_T, value))
+ rffi.cast(self.T, value))
@specialize.argtype(1, 2)
def box_complex(self, real, imag):
return self.BoxType(
- rffi.cast(self._COMPONENTS_T, real),
- rffi.cast(self._COMPONENTS_T, imag))
+ rffi.cast(self.T, real),
+ rffi.cast(self.T, imag))
def unbox(self, box):
assert isinstance(box, self.BoxType)
@@ -1113,12 +1113,12 @@
real, imag = self.unbox(box)
raw_storage_setitem(arr.storage, i+offset, real)
raw_storage_setitem(arr.storage,
- i+offset+rffi.sizeof(self._COMPONENTS_T), imag)
+ i+offset+rffi.sizeof(self.T), imag)
def _read(self, storage, i, offset):
- real = raw_storage_getitem(self._COMPONENTS_T, storage, i + offset)
- imag = raw_storage_getitem(self._COMPONENTS_T, storage,
- i + offset + rffi.sizeof(self._COMPONENTS_T))
+ real = raw_storage_getitem(self.T, storage, i + offset)
+ imag = raw_storage_getitem(self.T, storage,
+ i + offset + rffi.sizeof(self.T))
return real, imag
def read(self, arr, i, offset, dtype=None):
@@ -1536,8 +1536,7 @@
class Complex64(ComplexFloating, BaseType):
_attrs_ = ()
- T = rffi.CHAR
- _COMPONENTS_T = rffi.FLOAT
+ T = rffi.FLOAT
BoxType = interp_boxes.W_Complex64Box
ComponentBoxType = interp_boxes.W_Float32Box
@@ -1547,8 +1546,7 @@
class Complex128(ComplexFloating, BaseType):
_attrs_ = ()
- T = rffi.CHAR
- _COMPONENTS_T = rffi.DOUBLE
+ T = rffi.DOUBLE
BoxType = interp_boxes.W_Complex128Box
ComponentBoxType = interp_boxes.W_Float64Box
@@ -1580,8 +1578,7 @@
class Complex192(ComplexFloating, BaseType):
_attrs_ = ()
- T = rffi.CHAR
- _COMPONENTS_T = rffi.LONGDOUBLE
+ T = rffi.LONGDOUBLE
BoxType = interp_boxes.W_Complex192Box
ComponentBoxType = interp_boxes.W_Float96Box
@@ -1612,8 +1609,7 @@
class Complex256(ComplexFloating, BaseType):
_attrs_ = ()
- T = rffi.CHAR
- _COMPONENTS_T = rffi.LONGDOUBLE
+ T = rffi.LONGDOUBLE
BoxType = interp_boxes.W_Complex256Box
ComponentBoxType = interp_boxes.W_Float128Box
_______________________________________________
pypy-commit mailing list
[email protected]
http://mail.python.org/mailman/listinfo/pypy-commit
