Repository: numba/numba

Issue:
Request np.diag() implementation in nopython mode
It would be great to have np.diag() implemented in nopython mode. Thank you for the very good work!



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
@@ -2433,6 +2433,54 @@ def eye(n, m, k, dtype):
     res = context.compile_internal(builder, eye, sig, args)
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
+@lower_builtin(numpy.diag, types.Array)
+def numpy_diag(context, builder, sig, args):
+    def diag_impl(val):
+        return numpy.diag(val, k=0)
+    return context.compile_internal(builder, diag_impl, sig, args)
+
+@lower_builtin(numpy.diag, types.Array, types.Integer)
+def numpy_diag_kwarg(context, builder, sig, args):
+    arg = sig.args[0]
+    if arg.ndim == 1:
+        # vector context
+        def diag_impl(arr, k=0):
+            s = arr.shape
+            n = s[0] + abs(k)
+            ret = numpy.zeros((n, n), arr.dtype)
+            if k >= 0:
+                for i in range(n - k):
+                    ret[i, k + i] = arr[i]
+            else:
+                for i in range(n + k):
+                    ret[i - k, i] = arr[i]
+            return ret
+    elif arg.ndim == 2:
+        # matrix context
+        def diag_impl(arr, k=0):
+            #Will return arr.diagonal(v, k) when axis args are supported
+            rows, cols = arr.shape
+            r = rows
+            c = cols
+            if k < 0:
+                rows = rows + k
+            if k > 0:
+                cols = cols - k
+            n = max(min(rows, cols), 0)
+            ret = numpy.empty(n, arr.dtype)
+            if k >= 0:
+                for i in range(n):
+                    ret[i] = arr[i, k + i]
+            else:
+                for i in range(n):
+                    ret[i] = arr[i - k, i]
+            return ret
+    else:
+        #invalid input
+        raise ValueError("Input must be 1- or 2-d.")
+
+    res = context.compile_internal(builder, diag_impl, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 @lower_builtin(numpy.arange, types.Number)
 def numpy_arange_1(context, builder, sig, args):
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -904,3 +904,20 @@ def typer(ref, deg = False):
             else:
                 return types.float64
         return typer
+
+@infer_global(numpy.diag)
+class DiagCtor(CallableTemplate):
+    """
+    Typing template for np.diag()
+    """
+    def generic(self):
+        def typer(ref, k = 0):
+            if(isinstance(ref, types.Array)):
+                if ref.ndim == 1:
+                    rdim = 2
+                elif ref.ndim == 2:
+                    rdim = 1
+                else:
+                    return None
+                return types.Array(ndim=rdim, dtype=ref.dtype, layout='C')
+        return typer


Test patch:
diff --git a/numba/tests/test_dyn_array.py b/numba/tests/test_dyn_array.py
--- a/numba/tests/test_dyn_array.py
+++ b/numba/tests/test_dyn_array.py
@@ -852,6 +852,73 @@ def func(n, m, k):
         self.check_eye_n_m_k(func)
 
 
+class TestNdDiag(TestCase):
+
+    def setUp(self):
+        v = np.array([1, 2, 3])
+        hv = np.array([[1, 2, 3]])
+        vv = np.transpose(hv)
+        self.vectors = [v, hv, vv]
+        a3x4 = np.arange(12).reshape(3, 4)
+        a4x3 = np.arange(12).reshape(4, 3)
+        self.matricies = [a3x4, a4x3]
+        def func(q):
+            return np.diag(q)
+        self.py = func
+        self.jit = nrtjit(func)
+
+        def func_kwarg(q, k=0):
+            return np.diag(q, k=k)
+        self.py_kw = func_kwarg
+        self.jit_kw = nrtjit(func_kwarg)
+
+    def check_diag(self, pyfunc, nrtfunc, *args, **kwargs):
+        expected = pyfunc(*args, **kwargs)
+        computed = nrtfunc(*args, **kwargs)
+        self.assertEqual(computed.size, expected.size)
+        self.assertEqual(computed.dtype, expected.dtype)
+        # NOTE: stride not tested as np returns a RO view, nb returns new data
+        np.testing.assert_equal(expected, computed)
+
+    # create a diag matrix from a vector
+    def test_diag_vect_create(self):
+        for d in self.vectors:
+            self.check_diag(self.py, self.jit, d)
+
+    # create a diag matrix from a vector at a given offset
+    def test_diag_vect_create_kwarg(self):
+        for k in range(-10, 10):
+            for d in self.vectors:
+                self.check_diag(self.py_kw, self.jit_kw, d, k=k)
+
+    # extract the diagonal
+    def test_diag_extract(self):
+        for d in self.matricies:
+            self.check_diag(self.py, self.jit, d)
+
+    # extract a diagonal at a given offset
+    def test_diag_extract_kwarg(self):
+        for k in range(-4, 4):
+            for d in self.matricies:
+                self.check_diag(self.py_kw, self.jit_kw, d, k=k)
+
+    # check error handling
+    def test_error_handling(self):
+        from numba.errors import TypingError
+        d = np.array([[[1.]]])
+        cfunc = nrtjit(self.py)
+
+        # missing arg
+        with self.assertRaises(TypeError):
+            cfunc()
+
+        # > 2d
+        with self.assertRaises(TypingError):
+            cfunc(d)
+        with self.assertRaises(TypingError):
+            dfunc = nrtjit(self.py_kw)
+            dfunc(d, k=3)
+
 class TestNdArange(BaseTest):
 
     def test_linspace_2(self):