[pypy-commit] pypy numpy-dtype-refactor-complex: initial work on complex values

alex_gaynor noreply at buildbot.pypy.org
Sun Nov 20 20:26:33 CET 2011


Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: numpy-dtype-refactor-complex
Changeset: r49587:7a565e62ba58
Date: 2011-11-19 17:47 -0500
http://bitbucket.org/pypy/pypy/changeset/7a565e62ba58/

Log:	initial work on complex values

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -27,6 +27,8 @@
         'inexact': 'interp_boxes.W_InexactBox',
         'floating': 'interp_boxes.W_FloatingBox',
         'float64': 'interp_boxes.W_Float64Box',
+        'complexfloating': 'interp_boxes.W_ComplexFloatingBox',
+        'complex128': 'interp_boxes.W_Complex128Box',
     }
 
     # ufuncs
diff --git a/pypy/module/micronumpy/interp_boxes.py b/pypy/module/micronumpy/interp_boxes.py
--- a/pypy/module/micronumpy/interp_boxes.py
+++ b/pypy/module/micronumpy/interp_boxes.py
@@ -2,6 +2,7 @@
 from pypy.interpreter.error import operationerrfmt
 from pypy.interpreter.gateway import interp2app
 from pypy.interpreter.typedef import TypeDef
+from pypy.objspace.std.complextype import complex_typedef
 from pypy.objspace.std.floattype import float_typedef
 from pypy.objspace.std.inttype import int_typedef
 from pypy.objspace.std.typeobject import W_TypeObject
@@ -29,6 +30,12 @@
     def convert_to(self, dtype):
         return dtype.box(self.value)
 
+class CompositeBox(object):
+    _mixin_ = True
+
+    def __init__(self, subboxes):
+        self.subboxes = subboxes
+
 class W_GenericBox(Wrappable):
     _attrs_ = ()
 
@@ -150,7 +157,11 @@
 class W_Float64Box(W_FloatingBox, PrimitiveBox):
     descr__new__, get_dtype = new_dtype_getter("float64")
 
+class W_ComplexFloatingBox(W_InexactBox):
+    pass
 
+class W_Complex128Box(W_ComplexFloatingBox, CompositeBox):
+    descr__new__, get_dtype = new_dtype_getter("complex128")
 
 W_GenericBox.typedef = TypeDef("generic",
     __module__ = "numpy",
@@ -259,4 +270,12 @@
     __module__ = "numpy",
 
     __new__ = interp2app(W_Float64Box.descr__new__.im_func),
+)
+
+W_ComplexFloatingBox.typedef = TypeDef("complexfloating", W_InexactBox.typedef,
+    __module__ = "numpy",
+)
+
+W_Complex128Box.typedef = TypeDef("complex128", (W_ComplexFloatingBox.typedef, complex_typedef),
+    __module__ = "numpy",
 )
\ No newline at end of file
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
@@ -13,6 +13,7 @@
 SIGNEDLTR = "i"
 BOOLLTR = "b"
 FLOATINGLTR = "f"
+COMPLEXLTR = "c"
 
 class W_Dtype(Wrappable):
     def __init__(self, itemtype, num, kind, name, char, w_box_type, alternate_constructors=[]):
@@ -90,6 +91,8 @@
     num = interp_attrproperty("num", cls=W_Dtype),
     kind = interp_attrproperty("kind", cls=W_Dtype),
     type = interp_attrproperty_w("w_box_type", cls=W_Dtype),
+    name = interp_attrproperty("name", cls=W_Dtype),
+    char = interp_attrproperty("char", cls=W_Dtype),
     itemsize = GetSetProperty(W_Dtype.descr_get_itemsize),
     shape = GetSetProperty(W_Dtype.descr_get_shape),
 )
@@ -103,7 +106,7 @@
             kind=BOOLLTR,
             name="bool",
             char="?",
-            w_box_type = space.gettypefor(interp_boxes.W_BoolBox),
+            w_box_type=space.gettypefor(interp_boxes.W_BoolBox),
             alternate_constructors=[space.w_bool],
         )
         self.w_int8dtype = W_Dtype(
@@ -112,7 +115,7 @@
             kind=SIGNEDLTR,
             name="int8",
             char="b",
-            w_box_type = space.gettypefor(interp_boxes.W_Int8Box)
+            w_box_type=space.gettypefor(interp_boxes.W_Int8Box)
         )
         self.w_uint8dtype = W_Dtype(
             types.UInt8(),
@@ -120,7 +123,7 @@
             kind=UNSIGNEDLTR,
             name="uint8",
             char="B",
-            w_box_type = space.gettypefor(interp_boxes.W_UInt8Box),
+            w_box_type=space.gettypefor(interp_boxes.W_UInt8Box),
         )
         self.w_int16dtype = W_Dtype(
             types.Int16(),
@@ -128,7 +131,7 @@
             kind=SIGNEDLTR,
             name="int16",
             char="h",
-            w_box_type = space.gettypefor(interp_boxes.W_Int16Box),
+            w_box_type=space.gettypefor(interp_boxes.W_Int16Box),
         )
         self.w_uint16dtype = W_Dtype(
             types.UInt16(),
@@ -136,7 +139,7 @@
             kind=UNSIGNEDLTR,
             name="uint16",
             char="H",
-            w_box_type = space.gettypefor(interp_boxes.W_UInt16Box),
+            w_box_type=space.gettypefor(interp_boxes.W_UInt16Box),
         )
         self.w_int32dtype = W_Dtype(
             types.Int32(),
@@ -144,7 +147,7 @@
             kind=SIGNEDLTR,
             name="int32",
             char="i",
-             w_box_type = space.gettypefor(interp_boxes.W_Int32Box),
+            w_box_type=space.gettypefor(interp_boxes.W_Int32Box),
        )
         self.w_uint32dtype = W_Dtype(
             types.UInt32(),
@@ -152,7 +155,7 @@
             kind=UNSIGNEDLTR,
             name="uint32",
             char="I",
-            w_box_type = space.gettypefor(interp_boxes.W_UInt32Box),
+            w_box_type=space.gettypefor(interp_boxes.W_UInt32Box),
         )
         if LONG_BIT == 32:
             name = "int32"
@@ -164,7 +167,7 @@
             kind=SIGNEDLTR,
             name=name,
             char="l",
-            w_box_type = space.gettypefor(interp_boxes.W_LongBox),
+            w_box_type=space.gettypefor(interp_boxes.W_LongBox),
             alternate_constructors=[space.w_int],
         )
         self.w_ulongdtype = W_Dtype(
@@ -173,7 +176,7 @@
             kind=UNSIGNEDLTR,
             name="u" + name,
             char="L",
-            w_box_type = space.gettypefor(interp_boxes.W_ULongBox),
+            w_box_type=space.gettypefor(interp_boxes.W_ULongBox),
         )
         self.w_int64dtype = W_Dtype(
             types.Int64(),
@@ -181,7 +184,7 @@
             kind=SIGNEDLTR,
             name="int64",
             char="q",
-            w_box_type = space.gettypefor(interp_boxes.W_Int64Box),
+            w_box_type=space.gettypefor(interp_boxes.W_Int64Box),
             alternate_constructors=[space.w_long],
         )
         self.w_uint64dtype = W_Dtype(
@@ -190,7 +193,7 @@
             kind=UNSIGNEDLTR,
             name="uint64",
             char="Q",
-            w_box_type = space.gettypefor(interp_boxes.W_UInt64Box),
+            w_box_type=space.gettypefor(interp_boxes.W_UInt64Box),
         )
         self.w_float32dtype = W_Dtype(
             types.Float32(),
@@ -198,7 +201,7 @@
             kind=FLOATINGLTR,
             name="float32",
             char="f",
-            w_box_type = space.gettypefor(interp_boxes.W_Float32Box),
+            w_box_type=space.gettypefor(interp_boxes.W_Float32Box),
         )
         self.w_float64dtype = W_Dtype(
             types.Float64(),
@@ -206,16 +209,26 @@
             kind=FLOATINGLTR,
             name="float64",
             char="d",
-            w_box_type = space.gettypefor(interp_boxes.W_Float64Box),
+            w_box_type=space.gettypefor(interp_boxes.W_Float64Box),
             alternate_constructors=[space.w_float],
         )
 
+        self.w_complex128dtype = W_Dtype(
+            types.Complex([types.Float64(), types.Float64()]),
+            num=15,
+            kind=COMPLEXLTR,
+            name="complex128",
+            char="D",
+            w_box_type=space.gettypefor(interp_boxes.W_Complex128Box),
+            alternate_constructors=[space.w_complex],
+        )
+
         self.builtin_dtypes = [
             self.w_booldtype, self.w_int8dtype, self.w_uint8dtype,
             self.w_int16dtype, self.w_uint16dtype, self.w_int32dtype,
             self.w_uint32dtype, self.w_longdtype, self.w_ulongdtype,
             self.w_int64dtype, self.w_uint64dtype, self.w_float32dtype,
-            self.w_float64dtype
+            self.w_float64dtype, self.w_complex128dtype
         ]
         self.dtypes_by_num_bytes = sorted(
             (dtype.itemtype.get_element_size(), dtype)
diff --git a/pypy/module/micronumpy/test/test_dtypes.py b/pypy/module/micronumpy/test/test_dtypes.py
--- a/pypy/module/micronumpy/test/test_dtypes.py
+++ b/pypy/module/micronumpy/test/test_dtypes.py
@@ -166,6 +166,17 @@
         # You can't subclass dtype
         raises(TypeError, type, "Foo", (dtype,), {})
 
+    def test_complex128(self):
+        import numpypy as numpy
+
+        dtype = numpy.dtype(complex)
+        assert dtype.name == "complex128"
+        assert dtype.char == "D"
+        assert dtype.itemsize == 16
+        assert dtype.kind == "c"
+        assert dtype.num == 15
+        assert dtype.type is numpy.complex128
+
 class AppTestTypes(BaseNumpyAppTest):
     def test_abstract_types(self):
         import numpypy as numpy
@@ -224,6 +235,11 @@
 
         assert numpy.float64(2.0) == 2.0
 
+    def test_complex128(self):
+        import numpypy as numpy
+
+        assert numpy.complex128.mro() == [numpy.complex128, numpy.complexfloating, numpy.inexact, numpy.number, numpy.generic, complex, object]
+
     def test_subclass_type(self):
         import numpypy as numpy
 
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
@@ -591,6 +591,8 @@
         assert array([int8(3)]).dtype is dtype("int8")
         assert array([bool_(True)]).dtype is dtype(bool)
         assert array([bool_(True), 3.0]).dtype is dtype(float)
+        assert array([1 + 2j]).dtype is dtype(complex)
+        assert array([1, 1 + 2j]).dtype is dtype(complex)
 
     def test_comparison(self):
         import operator
@@ -612,6 +614,15 @@
             for i in xrange(5):
                 assert c[i] == func(b[i], 3)
 
+    def test_complex_basic(self):
+        from numpypy import array
+
+        x = array([1, 2, 3], complex)
+        assert x[0] == complex(1, 0)
+        x[0] = 1 + 3j
+        assert x[0] == 1 + 3j
+        assert x[2].real == 3
+        assert x[2].imag == 0
 
 class AppTestSupport(BaseNumpyAppTest):
     def setup_class(cls):
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
@@ -364,4 +364,50 @@
 
 class Float64(BaseType, Float):
     T = rffi.DOUBLE
-    BoxType = interp_boxes.W_Float64Box
\ No newline at end of file
+    BoxType = interp_boxes.W_Float64Box
+
+
+class BaseCompositeType(BaseType):
+    def __init__(self, itemtypes):
+        self.itemtypes = itemtypes
+
+    def get_element_size(self):
+        s = 0
+        for itemtype in self.itemtypes:
+            s += itemtype.get_element_size()
+        return s
+
+    def box(self, value):
+        return self.BoxType(value)
+
+    def unbox(self, box):
+        assert isinstance(box, self.BoxType)
+        return box.subboxes
+
+    def store(self, storage, width, i, offset, box):
+        subboxes = self.unbox(box)
+        i = 0
+        for box in subboxes:
+            self.itemtypes[i].store(storage, width, i, offset, box)
+            offset += self.itemtypes[i].get_element_size()
+            i += 1
+
+    def read(self, storage, width, i, offset):
+        boxes = []
+        for itemtype in self.itemtypes:
+            boxes.append(itemtype.read(storage, width, i, offset))
+            offset += itemtype.get_element_size()
+        return self.box(boxes)
+
+class Complex(BaseCompositeType):
+    BoxType = interp_boxes.W_Complex128Box
+
+    def __init__(self, itemtypes):
+        BaseCompositeType.__init__(self, itemtypes)
+        [self.real, self.imag] = self.itemtypes
+
+    def coerce(self, space, w_item):
+        if isinstance(w_item, self.BoxType):
+            return w_item
+        real, imag = space.unpackcomplex(w_item)
+        return self.box([self.real.box(real), self.imag.box(imag)])
\ No newline at end of file


More information about the pypy-commit mailing list