Repository: dask/dask

Issue:
Add rollaxis and moveaxis
Would be nice to have Dask Array implementations of [`rollaxis`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.rollaxis.html ) and [`moveaxis`]( https://docs.scipy.org/doc/numpy/reference/generated/numpy.moveaxis.html ). This can probably be done relatively easily by building off of [`transpose`]( http://dask.pydata.org/en/latest/array-api.html#dask.array.transpose ).


Discussion:
These functions already work on dask arrays, since they are defined in terms of the transpose method:
```
In [1]: import numpy as np

In [2]: import dask.array as da

In [3]: x = da.ones((10, 11, 12), chunks=7)

In [4]: np.rollaxis(x, 0, -1)
Out[4]: dask.array<transpose, shape=(11, 10, 12), dtype=float64, chunksize=(7, 7, 7)>

In [5]: np.moveaxis(x, 0, -1)
Out[5]: dask.array<transpose, shape=(11, 12, 10), dtype=float64, chunksize=(7, 7, 7)>
```

So we could probably just add appropriate aliases here.
I see. Should we be relying on an implementation detail of these functions though? Not that I'd expect them to change much (if at all). Just being cautious. FWIW it should be equally trivial to write them using Dask alone.
I'm not 100% sure about rollaxis (you should check the numpy test suite),
but certainly this is absolutely intended for moveaxis -- I wrote that test
case myself. If you prefer, you can also copy there implementations
verbatim -- this is why we use compatible licenses.
On Wed, Jul 26, 2017 at 8:17 PM jakirkham <notifications@github.com> wrote:

> I see. Should we be relying on an implementation detail of these functions
> though? Not that I'd expect them to change much (if at all). Just being
> cautious. FWIW it should be equally trivial to write them using Dask alone.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/2559#issuecomment-318247246>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/ABKS1q9yBlWoKPp5CpAUWFOh1w4GN0_Yks5sSAFNgaJpZM4OjdPH>
> .
>

Certainly having test cases for it is good. I'm just wondering do the docs spell this out somewhere? Ideally the API description would say this, but I'm not seeing it. Am thinking about this from the perspective of a C/C++ programmer where the docs really are the gold standard about what the library does.

No strong preference as long as these behaviors are documented and tested this way. Sounds like we know `moveaxis` is tested this way. The doc change should be easy if it is agreeable. `rollaxis` will take some research and at least a doc addition (if not a test addition as well).
Also this [comment]( https://github.com/numpy/numpy/issues/9473#issuecomment-318371050 ) is interesting.
> I'm just wondering do the docs spell this out somewhere? Ideally the API description would say this, but I'm not seeing it. Am thinking about this from the perspective of a C/C++ programmer where the docs really are the gold standard about what the library does.

Agreed, see https://github.com/numpy/numpy/issues/9476.
I'm labeling this as good first issue.  It sounds like providing an alias like

```
rollaxis = np.rollaxis
```

providing some tests and adding this to the api docs in `docs/source/array-api.rst` would be enough
I could spend some time on this to make a PR in the style @mrocklin proposes, but am somewhat confused by @jakirkham 's work on this. Can I just start a new PR trying the aliassing?
Also, I saw this: https://github.com/numpy/numpy/issues/11658, should I take this into account in any way?

Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -23,6 +23,7 @@
         to_zarr,
     )
     from .tiledb_io import from_tiledb, to_tiledb
+    from .numpy_compat import rollaxis, moveaxis
     from .routines import (
         take,
         choose,
diff --git a/dask/array/numpy_compat.py b/dask/array/numpy_compat.py
--- a/dask/array/numpy_compat.py
+++ b/dask/array/numpy_compat.py
@@ -5,6 +5,8 @@
 import numpy as np
 import warnings
 
+from ..utils import derived_from
+
 _numpy_116 = LooseVersion(np.__version__) >= "1.16.0"
 _numpy_117 = LooseVersion(np.__version__) >= "1.17.0"
 
@@ -289,3 +291,47 @@ def walk(self, x, index=()):
     _unravel_index_keyword = "shape"
 else:
     _unravel_index_keyword = "dims"
+
+
+# Implementation taken directly from numpy:
+# https://github.com/numpy/numpy/blob/d9b1e32cb8ef90d6b4a47853241db2a28146a57d/numpy/core/numeric.py#L1336-L1405
+@derived_from(np)
+def moveaxis(a, source, destination):
+    source = np.core.numeric.normalize_axis_tuple(source, a.ndim, "source")
+    destination = np.core.numeric.normalize_axis_tuple(
+        destination, a.ndim, "destination"
+    )
+    if len(source) != len(destination):
+        raise ValueError(
+            "`source` and `destination` arguments must have "
+            "the same number of elements"
+        )
+
+    order = [n for n in range(a.ndim) if n not in source]
+
+    for dest, src in sorted(zip(destination, source)):
+        order.insert(dest, src)
+
+    result = a.transpose(order)
+    return result
+
+
+# Implementation adapted directly from numpy:
+# https://github.com/numpy/numpy/blob/v1.17.0/numpy/core/numeric.py#L1107-L1204
+def rollaxis(a, axis, start=0):
+    n = a.ndim
+    axis = np.core.numeric.normalize_axis_index(axis, n)
+    if start < 0:
+        start += n
+    msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
+    if not (0 <= start < n + 1):
+        raise ValueError(msg % ("start", -n, "start", n + 1, start))
+    if axis < start:
+        # it's been removed
+        start -= 1
+    if axis == start:
+        return a[...]
+    axes = list(range(0, n))
+    axes.remove(axis)
+    axes.insert(start, axis)
+    return a.transpose(axes)


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
@@ -139,6 +139,41 @@ def test_swapaxes():
     assert d.swapaxes(0, 1).name != d.swapaxes(1, 0).name
 
 
+@pytest.mark.parametrize("funcname", ["moveaxis", "rollaxis"])
+@pytest.mark.parametrize("shape", [(), (5,), (3, 5, 7, 3)])
+def test_moveaxis_rollaxis(funcname, shape):
+    x = np.random.random(shape)
+    d = da.from_array(x, chunks=(len(shape) * (2,)))
+    np_func = getattr(np, funcname)
+    da_func = getattr(da, funcname)
+    for axis1 in range(-x.ndim, x.ndim):
+        assert isinstance(da_func(d, 0, axis1), da.Array)
+        for axis2 in range(-x.ndim, x.ndim):
+            assert_eq(np_func(x, axis1, axis2), da_func(d, axis1, axis2))
+
+
+def test_moveaxis_rollaxis_keyword():
+    x = np.random.random((10, 12, 7))
+    d = da.from_array(x, chunks=(4, 5, 2))
+    assert_eq(
+        np.moveaxis(x, destination=1, source=0), da.moveaxis(d, destination=1, source=0)
+    )
+    assert_eq(np.rollaxis(x, 2), da.rollaxis(d, 2))
+    assert isinstance(da.rollaxis(d, 1), da.Array)
+    assert_eq(np.rollaxis(x, start=1, axis=2), da.rollaxis(d, start=1, axis=2))
+
+
+def test_moveaxis_rollaxis_numpy_api():
+    a = da.random.random((4, 4, 4), chunks=2)
+    result = np.moveaxis(a, 2, 0)
+    assert isinstance(result, da.Array)
+    assert_eq(result, np.moveaxis(a.compute(), 2, 0))
+
+    result = np.rollaxis(a, 2, 0)
+    assert isinstance(result, da.Array)
+    assert_eq(result, np.rollaxis(a.compute(), 2, 0))
+
+
 @pytest.mark.parametrize(
     "funcname, kwargs",
     [