Migrate Data.transpose method from LAMA to Dask

Author: sadielbartholomew

cf/data/data.py

@@ -12079,6 +12079,7 @@ def tolist(self):
         """
         return self.array.tolist()
 
+    @daskified(1)
     @_deprecated_kwarg_check("i")
     @_inplace_enabled(default=False)
     def transpose(self, axes=None, inplace=False, i=False):
@@ -12119,9 +12120,8 @@ def transpose(self, axes=None, inplace=False, i=False):
         """
         d = _inplace_enabled_define_and_cleanup(self)
 
-        ndim = d._ndim
-
         # Parse the axes. By default, reverse the order of the axes.
+        ndim = d.ndim
         if axes is None:
             if ndim <= 1:
                 return d
@@ -12142,21 +12142,14 @@ def transpose(self, axes=None, inplace=False, i=False):
                 )
         # --- End: if
 
-        # Permute the axes.
-        data_axes = d._axes
-        d._axes = [data_axes[i] for i in iaxes]
-
-        # Permute the shape
-        shape = d._shape
-        d._shape = tuple([shape[i] for i in iaxes])
+        # TODODASK Q) is d._axes still relevant? Or can the following go?
+        # Test passes with it commented out, but it could still be important?
+        # data_axes = d._axes
+        # d._axes = [data_axes[i] for i in iaxes]
 
-        # Permute the locations map
-        for partition in d.partitions.matrix.flat:
-            location = partition.location
-            shape = partition.shape
-
-            partition.location = [location[i] for i in iaxes]
-            partition.shape = [shape[i] for i in iaxes]
+        dx = d._get_dask()
+        dx = da.transpose(dx, axes=axes)
+        d._set_dask(dx)
 
         return d
 

cf/test/test_Data.py

@@ -859,6 +859,7 @@ def test_Data_squeeze_insert_dimension(self):
     def test_Data___getitem__(self):
         if self.test_only and inspect.stack()[0][3] not in self.test_only:
             return
+        # TODODASK - why is there no testing here?
 
     def test_Data___setitem__(self):
         if self.test_only and inspect.stack()[0][3] not in self.test_only:
@@ -1522,29 +1523,24 @@ def test_Data_swapaxes(self):
                         self.assertEqual(b.shape, e.shape, message)
                         self.assertTrue((b == e.array).all(), message)
 
-    @unittest.skipIf(TEST_DASKIFIED_ONLY, "no attribute 'chunk_sizes'")
     def test_Data_transpose(self):
         if self.test_only and inspect.stack()[0][3] not in self.test_only:
             return
 
         a = np.arange(10 * 15 * 19).reshape(10, 1, 15, 19)
 
-        for chunksize in self.chunk_sizes:
-            with cf.chunksize(chunksize):
-                d = cf.Data(a.copy())
+        d = cf.Data(a.copy())
 
-                for indices in (range(a.ndim), range(-a.ndim, 0)):
-                    for axes in itertools.permutations(indices):
-                        a = np.transpose(a, axes)
-                        d.transpose(axes, inplace=True)
-                        message = (
-                            "cf.Data.transpose({}) failed: "
-                            "d.shape={}, a.shape={}".format(
-                                axes, d.shape, a.shape
-                            )
-                        )
-                        self.assertEqual(d.shape, a.shape, message)
-                        self.assertTrue((d.array == a).all(), message)
+        for indices in (range(a.ndim), range(-a.ndim, 0)):
+            for axes in itertools.permutations(indices):
+                a = np.transpose(a, axes)
+                d.transpose(axes, inplace=True)
+                message = (
+                    "cf.Data.transpose({}) failed: "
+                    "d.shape={}, a.shape={}".format(axes, d.shape, a.shape)
+                )
+                self.assertEqual(d.shape, a.shape, message)
+                self.assertTrue((d.array == a).all(), message)
 
     @unittest.skipIf(TEST_DASKIFIED_ONLY, "no attribute 'chunk_sizes'")
     def test_Data_unique(self):