Fix a coregistration issue due to unsqueezed dims

For some dimensions the grouped by dimension is sometimes not
squeezed out. Manually squeezing out the dimension works and
produces the desired coregistration successfully.

Closes #684

CHANGES.md

@@ -10,6 +10,9 @@
   Addresses [#701](https://github.com/CCI-Tools/cate/issues/701)
 * Fix a bug where correlation would fail with differing time dimensions
   Addresses [#700](https://github.com/CCI-Tools/cate/issues/700)
+* Fix a bug where coregistration would fail in some cases when the grouped by dimension
+  is not squeezed out automatically.
+  Addresses [#684](https://github.com/CCI-Tools/cate/issues/684)
 
 ## Version 2.0.0.dev15
 

cate/ops/coregistration.py

@@ -206,7 +206,8 @@ def _resample_slice(arr_slice: xr.DataArray, w: int, h: int, ds_method: int, us_
     """
     monitor = parent_monitor.child(1)
     with monitor.observing("resample slice"):
-        result = resampling.resample_2d(np.ma.masked_invalid(arr_slice.values),
+        # TODO (JG, 14.07.18) redundant squeeze should not be neccessary, see <xarray_issue>
+        result = resampling.resample_2d(np.ma.masked_invalid(arr_slice.squeeze().values),
                                         w,
                                         h,
                                         ds_method,
@@ -251,6 +252,7 @@ def _resample_array(array: xr.DataArray, lon: xr.DataArray, lat: xr.DataArray, m
                                 attrs=array.attrs).chunk()
 
     num_steps = 1
+    # print(groupby_list)
     for dim in groupby_list:
         num_steps = num_steps * len(array[dim])
 
@@ -263,6 +265,10 @@ def _resample_array(array: xr.DataArray, lon: xr.DataArray, lat: xr.DataArray, m
             # One spatial slice is one dask chunk, e.g. chunking is
             # (1,1,1..1,len(lat),len(lon))
             chunks[dim] = 1
+            # print(array)
+            # print(array.values.shape)
+            # print(temp_array)
+            # print(temp_array.values.shape)
         return xr.DataArray(temp_array.values,
                             name=array.name,
                             dims=array.dims,
@@ -385,9 +391,9 @@ def _nested_groupby_apply(array: xr.DataArray,
     :return: groupby-split-appy result
     """
     if len(groupby) == 1:
-        return array.groupby(groupby[0]).apply(apply_fn, **kwargs)
+        return array.groupby(groupby[0], squeeze=True).apply(apply_fn, **kwargs)
     else:
-        return array.groupby(groupby[0]).apply(_nested_groupby_apply,
-                                               groupby=groupby[1:],
-                                               apply_fn=apply_fn,
-                                               kwargs=kwargs)
+        return array.groupby(groupby[0], squeeze=True).apply(_nested_groupby_apply,
+                                                             groupby=groupby[1:],
+                                                             apply_fn=apply_fn,
+                                                             kwargs=kwargs)

test/ops/test_coregistration.py

@@ -419,8 +419,16 @@ def test_recursive(self):
         slice_coarse = np.eye(3, 6)
         ndarr_fine = np.zeros([2, 2, 2, 4, 8])
         ndarr_coarse = np.zeros([2, 2, 2, 3, 6])
+        ndarr_fine_l1 = np.zeros([2, 2, 4, 8])
+        ndarr_coarse_l1 = np.zeros([2, 2, 3, 6])
+        ndarr_fine_l2 = np.zeros([2, 2, 4, 8])
+        ndarr_coarse_l2 = np.zeros([2, 2, 3, 6])
         ndarr_fine[:] = slice_fine
         ndarr_coarse[:] = slice_coarse
+        ndarr_fine_l1[:] = slice_fine
+        ndarr_coarse_l1[:] = slice_coarse
+        ndarr_fine_l2[:] = slice_fine
+        ndarr_coarse_l2[:] = slice_coarse
 
         ds_fine = xr.Dataset({
             'first': (['time', 'layer', 'layer2', 'lat', 'lon'], ndarr_fine),
@@ -584,6 +592,41 @@ def test_recursive(self):
 
         assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)
 
+        # Test that coregistering with data arrays with less than all possible
+        # dimensions works
+        ds_fine = xr.Dataset({
+            'first': (['time', 'layer', 'lat', 'lon'], ndarr_fine_l1),
+            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_fine_l2),
+            'lat': np.linspace(-67.5, 67.5, 4),
+            'lon': np.linspace(-157.5, 157.5, 8),
+            'layer': np.array([1, 2]),
+            'layer2': np.array([1, 2]),
+            'time': np.array([1, 2])}).chunk(chunks={'lat': 2, 'lon': 4})
+
+        ds_coarse = xr.Dataset({
+            'first': (['time', 'layer', 'lat', 'lon'], ndarr_coarse_l1),
+            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_coarse_l2),
+            'lat': np.linspace(-60, 60, 3),
+            'lon': np.linspace(-150, 150, 6),
+            'time': np.array([1, 2]),
+            'layer': np.array([1, 2]),
+            'layer2': np.array([1, 2])}).chunk(chunks={'lat': 3, 'lon': 3})
+
+        ds_fine_resampled = coregister(ds_coarse, ds_fine)
+        ndarr_coarse_exp = np.zeros([2, 2, 3, 6])
+        ndarr_coarse_exp[:] = slice_exp
+
+        expected = xr.Dataset({
+            'first': (['time', 'layer', 'lat', 'lon'], ndarr_coarse_exp),
+            'second': (['time', 'layer2', 'lat', 'lon'], ndarr_coarse_exp),
+            'lat': np.linspace(-60, 60, 3),
+            'lon': np.linspace(-150, 150, 6),
+            'layer': np.array([1, 2]),
+            'layer2': np.array([1, 2]),
+            'time': np.array([1, 2])})
+
+        assert_almost_equal(ds_fine_resampled['first'].values, expected['first'].values)
+
     def test_2D(self):
         """
         Test a case where a 2D lat/lon dataset is resampled or used for