Repository: numba/numba

Issue:
Support .real and .imag attribute for complex array
to get the view of the real and imag part

(reference: https://github.com/numba/numba/issues/2131#issuecomment-248909125)

Support .real and .imag attribute for complex array
to get the view of the real and imag part

(reference: https://github.com/numba/numba/issues/2131#issuecomment-248909125)



Discussion:

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -1885,6 +1885,87 @@ def array_ctypes_data(context, builder, typ, value):
     return impl_ret_untracked(context, builder, typ, res)
 
 
+#-------------------------------------------------------------------------------
+# .real / .imag
+
+@lower_getattr(types.Array, "real")
+def array_real_part(context, builder, typ, value):
+    if typ.dtype in types.complex_domain:
+        return array_complex_attr(context, builder, typ, value, attr='real')
+    elif typ.dtype in types.number_domain:
+        # as an identity function
+        return impl_ret_borrowed(context, builder, typ, value)
+    else:
+        raise NotImplementedError('unsupported .real for {}'.format(type.dtype))
+
+
+@lower_getattr(types.Array, "imag")
+def array_imag_part(context, builder, typ, value):
+    if typ.dtype in types.complex_domain:
+        return array_complex_attr(context, builder, typ, value, attr='imag')
+    elif typ.dtype in types.number_domain:
+        # return a readonly zero array
+        sig = signature(typ.copy(readonly=True), typ)
+        return numpy_zeros_like_nd(context, builder, sig, [value])
+    else:
+        raise NotImplementedError('unsupported .imag for {}'.format(type.dtype))
+
+
+def array_complex_attr(context, builder, typ, value, attr):
+    """
+    Given a complex array, it's memory layout is:
+
+        R C R C R C
+        ^   ^   ^
+
+    (`R` indicates a float for the real part;
+     `C` indicates a float for the imaginery part;
+     the `^` indicates the start of each element)
+
+    To get the real part, we can simply change the dtype and itemsize to that
+    of the underlying float type.  The new layout is:
+
+        R x R x R x
+        ^   ^   ^
+
+    (`x` indicates unused)
+
+    A load operation will use the dtype to determine the number of bytes to
+    load.
+
+    To get the imaginary part, we shift the pointer by 1 float offset and
+    change the dtype and itemsize.  The new layout is:
+
+        x C x C x C
+          ^   ^   ^
+    """
+    if attr not in ['real', 'imag'] or typ.dtype not in types.complex_domain:
+        raise NotImplementedError("cannot get attribute `{}`".format(attr))
+
+    arrayty = make_array(typ)
+    array = arrayty(context, builder, value)
+
+    # sizeof underlying float type
+    flty = typ.dtype.underlying_float
+    sizeof_flty = context.get_abi_sizeof(context.get_data_type(flty))
+    itemsize = array.itemsize.type(sizeof_flty)
+
+    # cast data pointer to float type
+    llfltptrty = context.get_value_type(flty).as_pointer()
+    dataptr = builder.bitcast(array.data, llfltptrty)
+
+    # add offset
+    if attr == 'imag':
+        dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
+
+    # make result
+    resultty = typ.copy(dtype=flty, layout='A')
+    result = make_array(resultty)(context, builder)
+    repl = dict(data=dataptr, itemsize=itemsize)
+    cgutils.copy_struct(result, array, repl)
+    return impl_ret_borrowed(context, builder, resultty, result._getvalue())
+
+
 #-------------------------------------------------------------------------------
 # Structured / record lookup
 
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -9,6 +9,7 @@
 # import time side effect: array operations requires typing support of sequence
 # defined in collections: e.g. array.shape[i]
 from numba.typing import collections
+from numba.errors import TypingError
 
 Indexing = namedtuple("Indexing", ("index", "result", "advanced"))
 
@@ -261,6 +262,24 @@ def resolve_T(self, ary):
             retty = ary.copy(layout=layout)
         return retty
 
+    def resolve_real(self, ary):
+        return self._resolve_real_imag(ary, attr='real')
+
+    def resolve_imag(self, ary):
+        return self._resolve_real_imag(ary, attr='imag')
+
+    def _resolve_real_imag(self, ary, attr):
+        if ary.dtype in types.complex_domain:
+            return ary.copy(dtype=ary.dtype.underlying_float, layout='A')
+        elif ary.dtype in types.number_domain:
+            res = ary.copy(dtype=ary.dtype)
+            if attr == 'imag':
+                res = res.copy(readonly=True)
+            return res
+        else:
+            msg = "cannot access .{} of array of {}"
+            raise TypingError(msg.format(attr, ary.dtype))
+
     @bound_function("array.transpose")
     def resolve_transpose(self, ary, args, kws):
         assert not args


Test patch:
diff --git a/numba/tests/test_array_attr.py b/numba/tests/test_array_attr.py
--- a/numba/tests/test_array_attr.py
+++ b/numba/tests/test_array_attr.py
@@ -1,7 +1,5 @@
 from __future__ import print_function
 
-import sys
-
 import numpy as np
 
 import numba.unittest_support as unittest
@@ -9,6 +7,7 @@
 from numba.numpy_support import from_dtype
 from numba import types, njit, typeof, numpy_support
 from .support import TestCase, CompilationCache, MemoryLeakMixin, tag
+from numba.errors import TypingError
 
 
 def array_dtype(a):
@@ -79,6 +78,14 @@ def array_ctypes_data(arr):
     return arr.ctypes.data
 
 
+def array_real(arr):
+    return arr.real
+
+
+def array_imag(arr):
+    return arr.imag
+
+
 class TestArrayAttr(MemoryLeakMixin, TestCase):
 
     def setUp(self):
@@ -225,5 +232,102 @@ def test_array_ctypes_data(self):
         self.assertEqual(pyfunc(arr), cfunc(arr))
 
 
+class TestRealImagAttr(MemoryLeakMixin, TestCase):
+    def check_complex(self, pyfunc):
+        cfunc = njit(pyfunc)
+        # test 1D
+        size = 10
+        arr = np.arange(size) + np.arange(size) * 10j
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+        # test 2D
+        arr = arr.reshape(2, 5)
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+
+    def test_complex_real(self):
+        self.check_complex(array_real)
+
+    def test_complex_imag(self):
+        self.check_complex(array_imag)
+
+    def check_number_real(self, dtype):
+        pyfunc = array_real
+        cfunc = njit(pyfunc)
+        # test 1D
+        size = 10
+        arr = np.arange(size, dtype=dtype)
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+        # test 2D
+        arr = arr.reshape(2, 5)
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+        # test identity
+        self.assertEqual(arr.data, pyfunc(arr).data)
+        self.assertEqual(arr.data, cfunc(arr).data)
+        # test writable
+        real = cfunc(arr)
+        self.assertNotEqual(arr[0, 0], 5)
+        real[0, 0] = 5
+        self.assertEqual(arr[0, 0], 5)
+
+    def test_number_real(self):
+        """
+        Testing .real of non-complex dtypes
+        """
+        for dtype in [np.uint8, np.int32, np.float32, np.float64]:
+            self.check_number_real(dtype)
+
+    def check_number_imag(self, dtype):
+        pyfunc = array_imag
+        cfunc = njit(pyfunc)
+        # test 1D
+        size = 10
+        arr = np.arange(size, dtype=dtype)
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+        # test 2D
+        arr = arr.reshape(2, 5)
+        self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+        # test are zeros
+        self.assertEqual(cfunc(arr).tolist(), np.zeros_like(arr).tolist())
+        # test readonly
+        imag = cfunc(arr)
+        with self.assertRaises(ValueError) as raises:
+            imag[0] = 1
+        self.assertEqual('assignment destination is read-only',
+                         str(raises.exception))
+
+    def test_number_imag(self):
+        """
+        Testing .imag of non-complex dtypes
+        """
+        for dtype in [np.uint8, np.int32, np.float32, np.float64]:
+            self.check_number_imag(dtype)
+
+    def test_record_real(self):
+        rectyp = np.dtype([('real', np.float32), ('imag', np.complex64)])
+        arr = np.zeros(3, dtype=rectyp)
+        arr['real'] = np.random.random(arr.size)
+        arr['imag'] = np.random.random(arr.size) * 1.3j
+
+        # check numpy behavior
+        # .real is identity
+        self.assertIs(array_real(arr), arr)
+        # .imag is zero_like
+        self.assertEqual(array_imag(arr).tolist(), np.zeros_like(arr).tolist())
+
+        # check numba behavior
+        # it's most likely a user error, anyway
+        jit_array_real = njit(array_real)
+        jit_array_imag = njit(array_imag)
+
+        with self.assertRaises(TypingError) as raises:
+            jit_array_real(arr)
+        self.assertIn("cannot access .real of array of Record",
+                      str(raises.exception))
+
+        with self.assertRaises(TypingError) as raises:
+            jit_array_imag(arr)
+        self.assertIn("cannot access .imag of array of Record",
+                      str(raises.exception))
+
+
 if __name__ == '__main__':
     unittest.main()