Repository: dask/dask

Issue:
Adding ravel_multi_index and unravel_index
Would be nice to have Dask implementations of [ravel_multi_index]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.ravel_multi_index.html ) and [unravel_index]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.unravel_index.html ). These are particularly nice to have as we are starting to look at Dask index arrays ( e.g. https://github.com/dask/dask/pull/2539 ).


Discussion:
Adding `unravel_index` in PR ( https://github.com/dask/dask/pull/3958 )
👍 to `ravel_multi_index`

Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -56,6 +56,7 @@
         flatnonzero,
         nonzero,
         unravel_index,
+        ravel_multi_index,
         around,
         isin,
         isnull,
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1362,12 +1362,6 @@ def nonzero(a):
         return (ind,)
 
 
-def _int_piecewise(x, *condlist, **kwargs):
-    return np.piecewise(
-        x, list(condlist), kwargs["funclist"], *kwargs["func_args"], **kwargs["func_kw"]
-    )
-
-
 def _unravel_index_kernel(indices, func_kwargs):
     return np.stack(np.unravel_index(indices, **func_kwargs))
 
@@ -1390,6 +1384,27 @@ def unravel_index(indices, shape, order="C"):
     return unraveled_indices
 
 
+def _ravel_multi_index_kernel(multi_index, func_kwargs):
+    return np.ravel_multi_index(multi_index, **func_kwargs)
+
+
+@wraps(np.ravel_multi_index)
+def ravel_multi_index(multi_index, dims, mode="raise", order="C"):
+    return multi_index.map_blocks(
+        _ravel_multi_index_kernel,
+        dtype=np.intp,
+        chunks=(multi_index.shape[-1],),
+        drop_axis=0,
+        func_kwargs=dict(dims=dims, mode=mode, order=order),
+    )
+
+
+def _int_piecewise(x, *condlist, **kwargs):
+    return np.piecewise(
+        x, list(condlist), kwargs["funclist"], *kwargs["func_args"], **kwargs["func_kw"]
+    )
+
+
 @derived_from(np)
 def piecewise(x, condlist, funclist, *args, **kw):
     return map_blocks(


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1503,6 +1503,28 @@ def test_unravel_index():
         assert_eq(darr.vindex[d_indices], arr[indices])
 
 
+@pytest.mark.parametrize(
+    "arr, chunks, kwargs",
+    [
+        # Numpy doctests:
+        ([[3, 6, 6], [4, 5, 1]], (2, 3), dict(dims=(7, 6), order="C")),
+        ([[3, 6, 6], [4, 5, 1]], (2, 1), dict(dims=(7, 6), order="F")),
+        ([[3, 6, 6], [4, 5, 1]], 1, dict(dims=(4, 6), mode="clip")),
+        ([[3, 6, 6], [4, 5, 1]], (2, 3), dict(dims=(4, 4), mode=("clip", "wrap"))),
+        # Shape tests:
+        ([[3, 6, 6]], (1, 1), dict(dims=(7), order="C")),
+        ([[3, 6, 6], [4, 5, 1], [8, 6, 2]], (3, 1), dict(dims=(7, 6, 9), order="C")),
+    ],
+)
+def test_ravel_multi_index(arr, chunks, kwargs):
+    arr = np.asarray(arr)
+    darr = da.from_array(arr, chunks=chunks)
+    assert_eq(
+        da.ravel_multi_index(darr, **kwargs).compute(),
+        np.ravel_multi_index(arr, **kwargs),
+    )
+
+
 def test_coarsen():
     x = np.random.randint(10, size=(24, 24))
     d = da.from_array(x, chunks=(4, 8))