Repository: dask/dask

Issue:
Add da.diag 
Today I needed a dask.array equivalent for `np.diag`, in particular for the 1d to 2d case

``` python
In [1]: import numpy as np

In [2]: np.diag([1, 2, 3])
Out[2]: 
array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]])
```

Creating `da.diag` might be a good starter problem for @jcrist .  This would take a 1d dask Array object (with metadata like `.chunks` and `.dtype`), create a new dask graph with diagonal blocks on the diagonal and zero blocks elsewhere, and then create a new Array object with new `.chunks` (same `.dtype`) and a new graph caused by merging together the old graph and the new bit for the diag call.



Discussion:
This doc might be of use: http://dask.readthedocs.org/en/latest/array-design.html

Would there be a use case ever for passing in a non `dask.Array` for the 1d sequence?

I'd skip it for now


Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -25,4 +25,4 @@
 from .rechunk import rechunk
 from ..context import set_options
 from .optimization import optimize
-from .creation import arange, linspace
+from .creation import arange, linspace, diag
diff --git a/dask/array/creation.py b/dask/array/creation.py
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -8,6 +8,7 @@
 
 linspace_names = ('linspace-%d' % i for i in count(1))
 arange_names = ('arange-%d' % i for i in count(1))
+diag_names = ('diag-%d' % i for i in count(1))
 
 
 def _get_blocksizes(num, blocksize):
@@ -143,3 +144,46 @@ def arange(*args, **kwargs):
         elem_count += bs
 
     return Array(dsk, name, chunks, dtype=dtype)
+
+
+def diag(v):
+    """Construct a diagonal array, with ``v`` on the diagonal.
+
+    Currently only implements diagonal array creation on the zeroth diagonal.
+    Support for the ``k``th diagonal or diagonal extraction, as per the numpy
+    interface, is not yet implemented.
+
+    Parameters
+    ----------
+    v : dask array
+
+    Returns
+    -------
+    out_array : dask array
+
+    Examples
+    --------
+
+    >>> diag(arange(3, chunks=3)).compute()
+    array([[0, 0, 0],
+           [0, 1, 0],
+           [0, 0, 2]])
+    """
+    if not isinstance(v, Array):
+        raise TypeError("v must be a dask array")
+    if v.ndim != 1:
+        raise NotImplementedError("Extracting diagonals with `diag` is not "
+                                  "implemented.")
+    chunks_1d = v.chunks[0]
+    name = next(diag_names)
+
+    blocks = v._keys()
+    dsk = v.dask.copy()
+    for i, m in enumerate(chunks_1d):
+        for j, n in enumerate(chunks_1d):
+            key = (name, i, j)
+            if i == j:
+                dsk[key] = (np.diag, blocks[i])
+            else:
+                dsk[key] = (np.zeros, (m, n))
+    return Array(dsk, name, (chunks_1d, chunks_1d), dtype=v._dtype)


Test patch:
diff --git a/dask/array/tests/test_creation.py b/dask/array/tests/test_creation.py
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -94,3 +94,20 @@ def test_arange_float_step():
     darr = da.arange(7.7, 1.5, -.8, chunks=3)
     nparr = np.arange(7.7, 1.5, -.8)
     eq(darr, nparr)
+
+
+def test_diag():
+    v = da.arange(11, chunks=3)
+    darr = da.diag(v)
+    nparr = np.diag(v)
+    eq(darr, nparr)
+
+    v = v + v + 3
+    darr = da.diag(v)
+    nparr = np.diag(v)
+    eq(darr, nparr)
+
+    v = da.arange(11, chunks=11)
+    darr = da.diag(v)
+    nparr = np.diag(v)
+    eq(darr, nparr)