From d150fe625409f1e0ca0879a1fb978db3821299af Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Thu, 9 Mar 2017 11:44:38 +0000
Subject: [PATCH 1/7] Reuse multidim_daskstack in merge in fast um loading.

---
 lib/iris/_lazy_data.py                        | 28 +++++++++
 lib/iris/_merge.py                            | 26 +--------
 .../um/_fast_load_structured_fields.py        | 19 +++---
 .../lazy_data/test_array_masked_to_nans.py    |  2 +-
 .../tests/unit/lazy_data/test_is_lazy_data.py |  2 +-
 .../unit/lazy_data/test_multidim_daskarray.py | 58 +++++++++++++++++++
 6 files changed, 98 insertions(+), 37 deletions(-)
 create mode 100644 lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py

diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py
index e96f24da86..fde7437fba 100644
--- a/lib/iris/_lazy_data.py
+++ b/lib/iris/_lazy_data.py
@@ -87,3 +87,31 @@ def array_masked_to_nans(array, mask=None):
         array = array.astype(np.dtype('f8'))
     array[mask] = np.nan
     return array
+
+
+def multidim_daskstack(stack):
+    """
+    Recursively build a multidimensional stacked dask array.
+
+    This is needed because dask.array.stack only accepts a 1-dimensional list.
+
+    Args:
+
+    * stack:
++        An ndarray of dask arrays.
+
+    Returns:
+        The input array converted to a lazy dask array.
+
+    """
+    if stack.ndim == 0:
+        # A 0-d array cannot be stacked.
+        result = stack.item()
+    elif stack.ndim == 1:
+        # Another base case : simple 1-d goes direct in dask.
+        result = da.stack(list(stack))
+    else:
+        # Recurse because dask.stack does not do multi-dimensional.
+        result = da.stack([multidim_daskstack(subarray)
+                           for subarray in stack])
+    return result
diff --git a/lib/iris/_merge.py b/lib/iris/_merge.py
index bb943b9415..c963c82730 100644
--- a/lib/iris/_merge.py
+++ b/lib/iris/_merge.py
@@ -33,7 +33,8 @@
 import numpy as np
 import numpy.ma as ma
 
-from iris._lazy_data import array_masked_to_nans, as_lazy_data, is_lazy_data
+from iris._lazy_data import (array_masked_to_nans, as_lazy_data, is_lazy_data,
+                             multidim_daskstack)
 import iris.cube
 import iris.coords
 import iris.exceptions
@@ -1069,27 +1070,6 @@ def derive_space(groups, relation_matrix, positions, function_matrix=None):
     return space
 
 
-def _multidim_daskstack(stack):
-    """
-    Recursively build a multidensional stacked dask array.
-
-    The argument is an ndarray of dask arrays.
-    This is needed because dask.array.stack only accepts a 1-dimensional list.
-
-    """
-    if stack.ndim == 0:
-        # A 0-d array cannot be merged.
-        result = stack.item()
-    elif stack.ndim == 1:
-        # 'Another' base case : simple 1-d goes direct in dask.
-        result = da.stack(list(stack))
-    else:
-        # Recurse because dask.stack does not do multi-dimensional.
-        result = da.stack([_multidim_daskstack(subarray)
-                           for subarray in stack])
-    return result
-
-
 class ProtoCube(object):
     """
     Framework for merging source-cubes into one or more higher
@@ -1234,7 +1214,7 @@ def merge(self, unique=True):
                     data = as_lazy_data(data, chunks=data.shape)
                 stack[nd_index] = data
 
-            merged_data = _multidim_daskstack(stack)
+            merged_data = multidim_daskstack(stack)
             if all_have_data:
                 # All inputs were concrete, so turn the result back into a
                 # normal array.
diff --git a/lib/iris/fileformats/um/_fast_load_structured_fields.py b/lib/iris/fileformats/um/_fast_load_structured_fields.py
index 25c7f21059..f17df36659 100644
--- a/lib/iris/fileformats/um/_fast_load_structured_fields.py
+++ b/lib/iris/fileformats/um/_fast_load_structured_fields.py
@@ -32,12 +32,11 @@
 from netCDF4 import netcdftime
 import numpy as np
 
+from iris._lazy_data import (as_lazy_data, array_masked_to_nans, is_lazy_data,
+                             multidim_daskstack)
 from iris.fileformats.um._optimal_array_structuring import \
     optimal_array_structure
 
-from iris.fileformats.pp import PPField3
-from iris._lazy_data import as_lazy_data
-
 
 class FieldCollation(object):
     """
@@ -89,15 +88,11 @@ def data(self):
         if not self._structure_calculated:
             self._calculate_structure()
         if self._data_cache is None:
-            data_arrays = [as_lazy_data(f._data, chunks=f._data.shape)
-                           for f in self.fields]
-            vector_dims_list = list(self.vector_dims_shape)
-            vector_dims_list.reverse()
-            self._data_cache = data_arrays
-            for size in vector_dims_list:
-                self._data_cache = [da.stack(self._data_cache[i:i+size]) for i
-                                    in range(0, len(self._data_cache), size)]
-            self._data_cache, = self._data_cache
+            data_arrays = np.array([as_lazy_data(f._data, chunks=f._data.shape)
+                for f in self.fields], dtype=object)
+            data_arrays = data_arrays.reshape(
+                self.vector_dims_shape + data_arrays.shape[1:])
+            self._data_cache = multidim_daskstack(data_arrays)            
         return self._data_cache
 
     @property
diff --git a/lib/iris/tests/unit/lazy_data/test_array_masked_to_nans.py b/lib/iris/tests/unit/lazy_data/test_array_masked_to_nans.py
index de55026e55..55ead790f0 100644
--- a/lib/iris/tests/unit/lazy_data/test_array_masked_to_nans.py
+++ b/lib/iris/tests/unit/lazy_data/test_array_masked_to_nans.py
@@ -14,7 +14,7 @@
 #
 # You should have received a copy of the GNU Lesser General Public License
 # along with Iris.  If not, see <http://www.gnu.org/licenses/>.
-"""Test :meth:`iris._lazy data.array_masked_to_nans` method."""
+"""Test function :func:`iris._lazy data.array_masked_to_nans`."""
 
 from __future__ import (absolute_import, division, print_function)
 from six.moves import (filter, input, map, range, zip)  # noqa
diff --git a/lib/iris/tests/unit/lazy_data/test_is_lazy_data.py b/lib/iris/tests/unit/lazy_data/test_is_lazy_data.py
index 1c40d9a475..5a9595fa1f 100644
--- a/lib/iris/tests/unit/lazy_data/test_is_lazy_data.py
+++ b/lib/iris/tests/unit/lazy_data/test_is_lazy_data.py
@@ -14,7 +14,7 @@
 #
 # You should have received a copy of the GNU Lesser General Public License
 # along with Iris.  If not, see <http://www.gnu.org/licenses/>.
-"""Test :meth:`iris._lazy data.is_lazy_data` method."""
+"""Test function :func:`iris._lazy data.is_lazy_data`."""
 
 from __future__ import (absolute_import, division, print_function)
 from six.moves import (filter, input, map, range, zip)  # noqa
diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py b/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
new file mode 100644
index 0000000000..6ad2ebf010
--- /dev/null
+++ b/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
@@ -0,0 +1,58 @@
+# (C) British Crown Copyright 2017, Met Office
+#
+# This file is part of Iris.
+#
+# Iris is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Iris is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
+"""Test function :func:`iris._lazy data.multidim_daskstack`."""
+
+from __future__ import (absolute_import, division, print_function)
+from six.moves import (filter, input, map, range, zip)  # noqa
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests
+
+import numpy as np
+import dask.array as da
+
+from iris._lazy_data import multidim_daskstack
+
+
+class Test_multidim_daskstack(tests.IrisTest):
+    def test_0d(self):
+        value = 4
+        data = np.array(da.from_array(np.array(value), chunks=1), dtype=object)
+        result = multidim_daskstack(data)
+        self.assertEqual(result, value)
+
+    def test_1d(self):
+        vals = [4, 11]
+        data = np.array([da.from_array(val*np.ones((3, 3)), chunks=3) for val
+                         in vals], dtype=object)
+        result = multidim_daskstack(data)
+        self.assertEqual(result.shape, (2, 3, 3))
+        self.assertArrayEqual(result[:, 0, 0], np.array(vals))
+
+    def test_2d(self):
+        vals = [4, 8, 11]
+        data = np.array([da.from_array(val*np.ones((2, 2)), chunks=2) for val
+                         in vals], dtype=object)
+        data = data.reshape(3, 1, 2, 2)
+        result = multidim_daskstack(data)
+        self.assertEqual(result.shape, (3, 1, 2, 2))
+        self.assertArrayEqual(result[:, :, 0, 0], np.array(vals).reshape(3, 1))
+
+
+if __name__ == '__main__':
+    tests.main()

From 975105b8b284e7cbae2892119ea2f2eaf280d4d7 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Thu, 9 Mar 2017 18:09:44 +0000
Subject: [PATCH 2/7] Use as lazy data.

---
 .../unit/lazy_data/test_multidim_daskarray.py  | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py b/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
index 6ad2ebf010..20fe9e4736 100644
--- a/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
+++ b/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
@@ -26,28 +26,28 @@
 import numpy as np
 import dask.array as da
 
-from iris._lazy_data import multidim_daskstack
+from iris._lazy_data import as_lazy_data, multidim_daskstack
 
 
 class Test_multidim_daskstack(tests.IrisTest):
-    def test_0d(self):
+    def test_0d_dask_stack(self):
         value = 4
-        data = np.array(da.from_array(np.array(value), chunks=1), dtype=object)
+        data = np.array(as_lazy_data(np.array(value)), dtype=object)
         result = multidim_daskstack(data)
         self.assertEqual(result, value)
 
-    def test_1d(self):
+    def test_1d_dask_stack(self):
         vals = [4, 11]
-        data = np.array([da.from_array(val*np.ones((3, 3)), chunks=3) for val
-                         in vals], dtype=object)
+        data = np.array([as_lazy_data(val*np.ones((3, 3))) for val in vals],
+                        dtype=object)
         result = multidim_daskstack(data)
         self.assertEqual(result.shape, (2, 3, 3))
         self.assertArrayEqual(result[:, 0, 0], np.array(vals))
 
-    def test_2d(self):
+    def test_2d_dask_stack(self):
         vals = [4, 8, 11]
-        data = np.array([da.from_array(val*np.ones((2, 2)), chunks=2) for val
-                         in vals], dtype=object)
+        data = np.array([as_lazy_data(val*np.ones((2, 2))) for val in vals],
+                        dtype=object)
         data = data.reshape(3, 1, 2, 2)
         result = multidim_daskstack(data)
         self.assertEqual(result.shape, (3, 1, 2, 2))

From 717e7cfd0771fa822f0bbf51f923ff7cd0618bc9 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Fri, 10 Mar 2017 13:56:31 +0000
Subject: [PATCH 3/7] Create empty arrays then fill with dasks.

---
 lib/iris/_lazy_data.py                        |  4 +-
 lib/iris/_merge.py                            |  5 +-
 .../um/_fast_load_structured_fields.py        | 13 ++---
 .../unit/lazy_data/test_multidim_daskarray.py | 58 -------------------
 .../lazy_data/test_multidim_lazy_stack.py     | 57 ++++++++++++++++++
 5 files changed, 67 insertions(+), 70 deletions(-)
 delete mode 100644 lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
 create mode 100644 lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py

diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py
index fde7437fba..2d643eb808 100644
--- a/lib/iris/_lazy_data.py
+++ b/lib/iris/_lazy_data.py
@@ -89,7 +89,7 @@ def array_masked_to_nans(array, mask=None):
     return array
 
 
-def multidim_daskstack(stack):
+def multidim_lazy_stack(stack):
     """
     Recursively build a multidimensional stacked dask array.
 
@@ -112,6 +112,6 @@ def multidim_daskstack(stack):
         result = da.stack(list(stack))
     else:
         # Recurse because dask.stack does not do multi-dimensional.
-        result = da.stack([multidim_daskstack(subarray)
+        result = da.stack([multidim_lazy_stack(subarray)
                            for subarray in stack])
     return result
diff --git a/lib/iris/_merge.py b/lib/iris/_merge.py
index c963c82730..972e962d81 100644
--- a/lib/iris/_merge.py
+++ b/lib/iris/_merge.py
@@ -33,8 +33,7 @@
 import numpy as np
 import numpy.ma as ma
 
-from iris._lazy_data import (array_masked_to_nans, as_lazy_data, is_lazy_data,
-                             multidim_daskstack)
+from iris._lazy_data import as_lazy_data, is_lazy_data, multidim_lazy_stack
 import iris.cube
 import iris.coords
 import iris.exceptions
@@ -1214,7 +1213,7 @@ def merge(self, unique=True):
                     data = as_lazy_data(data, chunks=data.shape)
                 stack[nd_index] = data
 
-            merged_data = multidim_daskstack(stack)
+            merged_data = multidim_lazy_stack(stack)
             if all_have_data:
                 # All inputs were concrete, so turn the result back into a
                 # normal array.
diff --git a/lib/iris/fileformats/um/_fast_load_structured_fields.py b/lib/iris/fileformats/um/_fast_load_structured_fields.py
index f17df36659..5f0760bce4 100644
--- a/lib/iris/fileformats/um/_fast_load_structured_fields.py
+++ b/lib/iris/fileformats/um/_fast_load_structured_fields.py
@@ -32,8 +32,7 @@
 from netCDF4 import netcdftime
 import numpy as np
 
-from iris._lazy_data import (as_lazy_data, array_masked_to_nans, is_lazy_data,
-                             multidim_daskstack)
+from iris._lazy_data import as_lazy_data, multidim_lazy_stack
 from iris.fileformats.um._optimal_array_structuring import \
     optimal_array_structure
 
@@ -88,11 +87,11 @@ def data(self):
         if not self._structure_calculated:
             self._calculate_structure()
         if self._data_cache is None:
-            data_arrays = np.array([as_lazy_data(f._data, chunks=f._data.shape)
-                for f in self.fields], dtype=object)
-            data_arrays = data_arrays.reshape(
-                self.vector_dims_shape + data_arrays.shape[1:])
-            self._data_cache = multidim_daskstack(data_arrays)            
+            stack = np.empty(np.prod(self.vector_dims_shape), 'object')
+            for index, f in enumerate(self.fields):
+                stack[index] = as_lazy_data(f._data, chunks=f._data.shape)
+            stack = stack.reshape(self.vector_dims_shape)
+            self._data_cache = multidim_lazy_stack(stack)
         return self._data_cache
 
     @property
diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py b/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
deleted file mode 100644
index 20fe9e4736..0000000000
--- a/lib/iris/tests/unit/lazy_data/test_multidim_daskarray.py
+++ /dev/null
@@ -1,58 +0,0 @@
-# (C) British Crown Copyright 2017, Met Office
-#
-# This file is part of Iris.
-#
-# Iris is free software: you can redistribute it and/or modify it under
-# the terms of the GNU Lesser General Public License as published by the
-# Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# Iris is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU Lesser General Public License for more details.
-#
-# You should have received a copy of the GNU Lesser General Public License
-# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
-"""Test function :func:`iris._lazy data.multidim_daskstack`."""
-
-from __future__ import (absolute_import, division, print_function)
-from six.moves import (filter, input, map, range, zip)  # noqa
-
-# Import iris.tests first so that some things can be initialised before
-# importing anything else.
-import iris.tests as tests
-
-import numpy as np
-import dask.array as da
-
-from iris._lazy_data import as_lazy_data, multidim_daskstack
-
-
-class Test_multidim_daskstack(tests.IrisTest):
-    def test_0d_dask_stack(self):
-        value = 4
-        data = np.array(as_lazy_data(np.array(value)), dtype=object)
-        result = multidim_daskstack(data)
-        self.assertEqual(result, value)
-
-    def test_1d_dask_stack(self):
-        vals = [4, 11]
-        data = np.array([as_lazy_data(val*np.ones((3, 3))) for val in vals],
-                        dtype=object)
-        result = multidim_daskstack(data)
-        self.assertEqual(result.shape, (2, 3, 3))
-        self.assertArrayEqual(result[:, 0, 0], np.array(vals))
-
-    def test_2d_dask_stack(self):
-        vals = [4, 8, 11]
-        data = np.array([as_lazy_data(val*np.ones((2, 2))) for val in vals],
-                        dtype=object)
-        data = data.reshape(3, 1, 2, 2)
-        result = multidim_daskstack(data)
-        self.assertEqual(result.shape, (3, 1, 2, 2))
-        self.assertArrayEqual(result[:, :, 0, 0], np.array(vals).reshape(3, 1))
-
-
-if __name__ == '__main__':
-    tests.main()
diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
new file mode 100644
index 0000000000..24a743c464
--- /dev/null
+++ b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
@@ -0,0 +1,57 @@
+# (C) British Crown Copyright 2017, Met Office
+#
+# This file is part of Iris.
+#
+# Iris is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Lesser General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# Iris is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with Iris.  If not, see <http://www.gnu.org/licenses/>.
+"""Test function :func:`iris._lazy data.multidim_lazy_stack`."""
+
+from __future__ import (absolute_import, division, print_function)
+from six.moves import (filter, input, map, range, zip)  # noqa
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests
+
+import numpy as np
+import dask.array as da
+
+from iris._lazy_data import as_lazy_data, multidim_lazy_stack
+
+
+class Test_multidim_lazy_stack(tests.IrisTest):
+    def _check(self, stack_shape):
+        vals = np.arange(np.prod(stack_shape)).reshape(stack_shape)
+        stack = np.empty(stack_shape, 'object')
+        stack_element_shape = (4, 5)
+        for index, val in np.ndenumerate(vals):
+            stack[index] = as_lazy_data(val*np.ones(stack_element_shape))
+        result = multidim_lazy_stack(stack)
+        self.assertEqual(result.shape, stack_shape + stack_element_shape)
+        self.assertIsInstance(result, da.core.Array)
+
+    def test_0d_lazy_stack(self):
+        shape = ()
+        result = self._check(shape)
+
+    def test_1d_lazy_stack(self):
+        shape = (2,)
+        result = self._check(shape)
+
+    def test_2d_lazy_stack(self):
+        shape = (3, 2)
+        result = self._check(shape)
+
+
+if __name__ == '__main__':
+    tests.main()

From 62e25e5aa7d007ee86b4e6a9d2fecf66b1365a16 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Fri, 10 Mar 2017 14:59:44 +0000
Subject: [PATCH 4/7] Rename iteration variable.

---
 lib/iris/fileformats/um/_fast_load_structured_fields.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/iris/fileformats/um/_fast_load_structured_fields.py b/lib/iris/fileformats/um/_fast_load_structured_fields.py
index 5f0760bce4..1cf556ef29 100644
--- a/lib/iris/fileformats/um/_fast_load_structured_fields.py
+++ b/lib/iris/fileformats/um/_fast_load_structured_fields.py
@@ -88,7 +88,7 @@ def data(self):
             self._calculate_structure()
         if self._data_cache is None:
             stack = np.empty(np.prod(self.vector_dims_shape), 'object')
-            for index, f in enumerate(self.fields):
+            for index, field in enumerate(self.fields):
                 stack[index] = as_lazy_data(f._data, chunks=f._data.shape)
             stack = stack.reshape(self.vector_dims_shape)
             self._data_cache = multidim_lazy_stack(stack)

From e69c46f320349c3af73a4808ba4fc9ae4dbd4929 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Fri, 10 Mar 2017 15:08:49 +0000
Subject: [PATCH 5/7] Add comment explaining numbers.

---
 lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
index 24a743c464..00589b3b0e 100644
--- a/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
+++ b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
@@ -33,6 +33,7 @@ class Test_multidim_lazy_stack(tests.IrisTest):
     def _check(self, stack_shape):
         vals = np.arange(np.prod(stack_shape)).reshape(stack_shape)
         stack = np.empty(stack_shape, 'object')
+        # Define the shape of each element in the stack.
         stack_element_shape = (4, 5)
         for index, val in np.ndenumerate(vals):
             stack[index] = as_lazy_data(val*np.ones(stack_element_shape))

From 2b9fe09c566393ce6230b0b24f505e8d59459bd6 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Fri, 10 Mar 2017 15:11:59 +0000
Subject: [PATCH 6/7] Use iteration variable

---
 lib/iris/fileformats/um/_fast_load_structured_fields.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/iris/fileformats/um/_fast_load_structured_fields.py b/lib/iris/fileformats/um/_fast_load_structured_fields.py
index 1cf556ef29..4b2820d205 100644
--- a/lib/iris/fileformats/um/_fast_load_structured_fields.py
+++ b/lib/iris/fileformats/um/_fast_load_structured_fields.py
@@ -89,7 +89,8 @@ def data(self):
         if self._data_cache is None:
             stack = np.empty(np.prod(self.vector_dims_shape), 'object')
             for index, field in enumerate(self.fields):
-                stack[index] = as_lazy_data(f._data, chunks=f._data.shape)
+                stack[index] = as_lazy_data(field._data,
+                                            chunks=field._data.shape)
             stack = stack.reshape(self.vector_dims_shape)
             self._data_cache = multidim_lazy_stack(stack)
         return self._data_cache

From 34f810dc79f6bb50e5cd653df601f6abf270bef0 Mon Sep 17 00:00:00 2001
From: lbdreyer <laura.dreyer@metoffice.gov.uk>
Date: Mon, 13 Mar 2017 14:10:37 +0000
Subject: [PATCH 7/7] add test of values; use np.ndindex

---
 .../fileformats/um/_fast_load_structured_fields.py     | 10 +++++-----
 .../tests/unit/lazy_data/test_multidim_lazy_stack.py   |  7 ++++++-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/lib/iris/fileformats/um/_fast_load_structured_fields.py b/lib/iris/fileformats/um/_fast_load_structured_fields.py
index 4b2820d205..402c93803f 100644
--- a/lib/iris/fileformats/um/_fast_load_structured_fields.py
+++ b/lib/iris/fileformats/um/_fast_load_structured_fields.py
@@ -87,11 +87,11 @@ def data(self):
         if not self._structure_calculated:
             self._calculate_structure()
         if self._data_cache is None:
-            stack = np.empty(np.prod(self.vector_dims_shape), 'object')
-            for index, field in enumerate(self.fields):
-                stack[index] = as_lazy_data(field._data,
-                                            chunks=field._data.shape)
-            stack = stack.reshape(self.vector_dims_shape)
+            stack = np.empty(self.vector_dims_shape, 'object')
+            for nd_index, field in zip(np.ndindex(self.vector_dims_shape),
+                                       self.fields):
+                stack[nd_index] = as_lazy_data(field._data,
+                                               chunks=field._data.shape)
             self._data_cache = multidim_lazy_stack(stack)
         return self._data_cache
 
diff --git a/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
index 00589b3b0e..80839b02b5 100644
--- a/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
+++ b/lib/iris/tests/unit/lazy_data/test_multidim_lazy_stack.py
@@ -35,11 +35,16 @@ def _check(self, stack_shape):
         stack = np.empty(stack_shape, 'object')
         # Define the shape of each element in the stack.
         stack_element_shape = (4, 5)
+        expected = np.empty(stack_shape + stack_element_shape,
+                            dtype=int)
         for index, val in np.ndenumerate(vals):
-            stack[index] = as_lazy_data(val*np.ones(stack_element_shape))
+            stack[index] = as_lazy_data(val * np.ones(stack_element_shape))
+
+            expected[index] = val
         result = multidim_lazy_stack(stack)
         self.assertEqual(result.shape, stack_shape + stack_element_shape)
         self.assertIsInstance(result, da.core.Array)
+        self.assertArrayAllClose(result.compute(), expected)
 
     def test_0d_lazy_stack(self):
         shape = ()