Ensure collapsed/aggregated_by/rolling_window treat ancillary variabl…

…es properly (#3549)

docs/iris/src/whatsnew/contributions_3.0.0/newfeature_2019-Nov-27_cell_measure_statistics.txt

@@ -0,0 +1,5 @@
+* Statistical operations :meth:`iris.cube.Cube.collapsed`,
+ :meth:`iris.cube.Cube.aggregated_by` and :meth:`iris.cube.Cube.rolling_window`
+ previously removed every :class:`iris.coord.CellMeasure` attached to the cube.
+ Now, a :class:`iris.coord.CellMeasure` will only be removed if it is associated
+ with an axis over which the statistic is being run.

lib/iris/cube.py

@@ -3786,11 +3786,15 @@ def collapsed(self, coords, aggregator, **kwargs):
 
  untouched_dims = set(range(self.ndim)) - set(dims_to_collapse)
 
+ collapsed_cube = iris.util._strip_metadata_from_dims(
+ self, dims_to_collapse
+ )
+
  # Remove the collapsed dimension(s) from the metadata
  indices = [slice(None, None)] * self.ndim
  for dim in dims_to_collapse:
  indices[dim] = 0
- collapsed_cube = self[tuple(indices)]
+ collapsed_cube = collapsed_cube[tuple(indices)]
 
  # Collapse any coords that span the dimension(s) being collapsed
  for coord in self.dim_coords + self.aux_coords:
@@ -3999,11 +4003,14 @@ def aggregated_by(self, coords, aggregator, **kwargs):
 
  # Create the resulting aggregate-by cube and remove the original
  # coordinates that are going to be groupedby.
+ aggregateby_cube = iris.util._strip_metadata_from_dims(
+ self, [dimension_to_groupby]
+ )
  key = [slice(None, None)] * self.ndim
  # Generate unique index tuple key to maintain monotonicity.
  key[dimension_to_groupby] = tuple(range(len(groupby)))
  key = tuple(key)
- aggregateby_cube = self[key]
+ aggregateby_cube = aggregateby_cube[key]
  for coord in groupby_coords + shared_coords:
  aggregateby_cube.remove_coord(coord)
 
@@ -4202,9 +4209,10 @@ def rolling_window(self, coord, aggregator, window, **kwargs):
  # some sort of `cube.prepare()` method would be handy to allow
  # re-shaping with given data, and returning a mapping of
  # old-to-new-coords (to avoid having to use metadata identity)?
+ new_cube = iris.util._strip_metadata_from_dims(self, [dimension])
  key = [slice(None, None)] * self.ndim
  key[dimension] = slice(None, self.shape[dimension] - window + 1)
- new_cube = self[tuple(key)]
+ new_cube = new_cube[tuple(key)]
 
  # take a view of the original data using the rolling_window function
  # this will add an extra dimension to the data at dimension + 1 which

lib/iris/tests/unit/cube/test_Cube.py

@@ -313,6 +313,32 @@ def test_non_lazy_aggregator(self):
  self.assertArrayEqual(result.data, np.mean(self.data, axis=1))
 
 
+class Test_collapsed__cellmeasure_ancils(tests.IrisTest):
+ def setUp(self):
+ cube = Cube(np.arange(6.0).reshape((2, 3)))
+ for i_dim, name in enumerate(("y", "x")):
+ npts = cube.shape[i_dim]
+ coord = DimCoord(np.arange(npts), long_name=name)
+ cube.add_dim_coord(coord, i_dim)
+ self.ancillary_variable = AncillaryVariable([0, 1], long_name="foo")
+ cube.add_ancillary_variable(self.ancillary_variable, 0)
+ self.cell_measure = CellMeasure([0, 1], long_name="bar")
+ cube.add_cell_measure(self.cell_measure, 0)
+ self.cube = cube
+
+ def test_ancillary_variables_and_cell_measures_kept(self):
+ cube_collapsed = self.cube.collapsed("x", MEAN)
+ self.assertEqual(
+ cube_collapsed.ancillary_variables(), [self.ancillary_variable]
+ )
+ self.assertEqual(cube_collapsed.cell_measures(), [self.cell_measure])
+
+ def test_ancillary_variables_and_cell_measures_removed(self):
+ cube_collapsed = self.cube.collapsed("y", MEAN)
+ self.assertEqual(cube_collapsed.ancillary_variables(), [])
+ self.assertEqual(cube_collapsed.cell_measures(), [])
+
+
 class Test_collapsed__warning(tests.IrisTest):
  def setUp(self):
  self.cube = Cube([[1, 2], [1, 2]])
@@ -508,6 +534,13 @@ def setUp(self):
  self.mock_agg.lazy_func = None
  self.mock_agg.post_process = mock.Mock(side_effect=lambda x, y, z: x)
 
+ self.ancillary_variable = AncillaryVariable(
+ [0, 1, 2, 3], long_name="foo"
+ )
+ self.cube.add_ancillary_variable(self.ancillary_variable, 0)
+ self.cell_measure = CellMeasure([0, 1, 2, 3], long_name="bar")
+ self.cube.add_cell_measure(self.cell_measure, 0)
+
  def test_2d_coord_simple_agg(self):
  # For 2d coords, slices of aggregated coord should be the same as
  # aggregated slices.
@@ -597,6 +630,18 @@ def test_single_string_aggregation(self):
  result.coord("bar"), AuxCoord(["a|a", "a"], long_name="bar")
  )
 
+ def test_ancillary_variables_and_cell_measures_kept(self):
+ cube_agg = self.cube.aggregated_by("val", self.mock_agg)
+ self.assertEqual(
+ cube_agg.ancillary_variables(), [self.ancillary_variable]
+ )
+ self.assertEqual(cube_agg.cell_measures(), [self.cell_measure])
+
+ def test_ancillary_variables_and_cell_measures_removed(self):
+ cube_agg = self.cube.aggregated_by("simple_agg", self.mock_agg)
+ self.assertEqual(cube_agg.ancillary_variables(), [])
+ self.assertEqual(cube_agg.cell_measures(), [])
+
 
 class Test_aggregated_by__lazy(tests.IrisTest):
  def setUp(self):
@@ -709,12 +754,23 @@ def test_single_string_aggregation__lazy(self):
 class Test_rolling_window(tests.IrisTest):
  def setUp(self):
  self.cube = Cube(np.arange(6))
+ self.multi_dim_cube = Cube(np.arange(36).reshape(6, 6))
  val_coord = DimCoord([0, 1, 2, 3, 4, 5], long_name="val")
  month_coord = AuxCoord(
  ["jan", "feb", "mar", "apr", "may", "jun"], long_name="month"
  )
+ extra_coord = AuxCoord([0, 1, 2, 3, 4, 5], long_name="extra")
  self.cube.add_dim_coord(val_coord, 0)
  self.cube.add_aux_coord(month_coord, 0)
+ self.multi_dim_cube.add_dim_coord(val_coord, 0)
+ self.multi_dim_cube.add_aux_coord(extra_coord, 1)
+ self.ancillary_variable = AncillaryVariable(
+ [0, 1, 2, 0, 1, 2], long_name="foo"
+ )
+ self.multi_dim_cube.add_ancillary_variable(self.ancillary_variable, 1)
+ self.cell_measure = CellMeasure([0, 1, 2, 0, 1, 2], long_name="bar")
+ self.multi_dim_cube.add_cell_measure(self.cell_measure, 1)
+
  self.mock_agg = mock.Mock(spec=Aggregator)
  self.mock_agg.aggregate = mock.Mock(return_value=np.empty([4]))
  self.mock_agg.post_process = mock.Mock(side_effect=lambda x, y, z: x)
@@ -760,6 +816,20 @@ def test_kwargs(self):
  )
  self.assertMaskedArrayEqual(expected_result, res_cube.data)
 
+ def test_ancillary_variables_and_cell_measures_kept(self):
+ res_cube = self.multi_dim_cube.rolling_window("val", self.mock_agg, 3)
+ self.assertEqual(
+ res_cube.ancillary_variables(), [self.ancillary_variable]
+ )
+ self.assertEqual(res_cube.cell_measures(), [self.cell_measure])
+
+ def test_ancillary_variables_and_cell_measures_removed(self):
+ res_cube = self.multi_dim_cube.rolling_window(
+ "extra", self.mock_agg, 3
+ )
+ self.assertEqual(res_cube.ancillary_variables(), [])
+ self.assertEqual(res_cube.cell_measures(), [])
+
 
 class Test_slices_dim_order(tests.IrisTest):
  """

lib/iris/util.py

@@ -1828,3 +1828,30 @@ def equalise_attributes(cubes):
  for key in list(cube.attributes.keys()):
  if key not in common_keys:
  del cube.attributes[key]
+
+
+def _strip_metadata_from_dims(cube, dims):
+ """
+ Remove ancillary variables and cell measures that map to specific dimensions.
+
+ Returns a cube copy with (possibly) some cell-measures and ancillary variables removed.
+
+ To be used by operations that modify or remove dimensions.
+ Note: does nothing to (aux)-coordinates. Those would be handled explicitly by the calling operation.
+
+ """
+ reduced_cube = cube.copy()
+
+ # Remove any ancillary variables that span the dimension(s) being collapsed
+ for ancil in reduced_cube.ancillary_variables():
+ ancil_dims = reduced_cube.ancillary_variable_dims(ancil)
+ if set(dims).intersection(ancil_dims):
+ reduced_cube.remove_ancillary_variable(ancil)
+
+ # Remove any cell measures that span the dimension(s) being collapsed
+ for cm in reduced_cube.cell_measures():
+ cm_dims = reduced_cube.cell_measure_dims(cm)
+ if set(dims).intersection(cm_dims):
+ reduced_cube.remove_cell_measure(cm)
+
+ return reduced_cube