From f77dedd377c75923ecbcb7d0a34a0570dd7a3682 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 18 Nov 2019 13:42:03 +0000 Subject: [PATCH] slices get ancillary variables --- lib/iris/cube.py | 16 ++++++++++++++ lib/iris/tests/unit/cube/test_Cube.py | 30 +++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/lib/iris/cube.py b/lib/iris/cube.py index 69afbd83bf..2cec1c03cd 100644 --- a/lib/iris/cube.py +++ b/lib/iris/cube.py @@ -2592,6 +2592,13 @@ def new_cell_measure_dims(cm_): if dimension_mapping[d] is not None ] + def new_ancillary_variable_dims(av_): + return [ + dimension_mapping[d] + for d in self.ancillary_variable_dims(av_) + if dimension_mapping[d] is not None + ] + # Fetch the data as a generic array-like object. cube_data = self._data_manager.core_data() @@ -2667,6 +2674,15 @@ def new_cell_measure_dims(cm_): new_cm = cellmeasure[cm_keys] cube.add_cell_measure(new_cm, new_cell_measure_dims(cellmeasure)) + # slice the ancillary variables and add them to the cube + for ancvar in self.ancillary_variables(): + dims = self.ancillary_variable_dims(ancvar) + av_keys = tuple([full_slice[dim] for dim in dims]) + new_av = ancvar[av_keys] + cube.add_ancillary_variable( + new_av, new_ancillary_variable_dims(ancvar) + ) + return cube def subset(self, coord): diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index cd2484370b..5cec50a06a 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -1959,6 +1959,36 @@ def test_cell_measure_1d(self): self.assertEqual(result.shape, result.cell_measures()[0].data.shape) +class Test__getitem_AncillaryVariables(tests.IrisTest): + def setUp(self): + cube = Cube(np.arange(6).reshape(2, 3)) + x_coord = DimCoord(points=np.array([2, 3, 4]), long_name="x") + cube.add_dim_coord(x_coord, 1) + y_coord = DimCoord(points=np.array([5, 6]), long_name="y") + cube.add_dim_coord(y_coord, 0) + z_coord = AuxCoord(points=np.arange(6).reshape(2, 3), long_name="z") + cube.add_aux_coord(z_coord, [0, 1]) + a_ancillary_variable = AncillaryVariable( + data=np.arange(6).reshape(2, 3), long_name="foo" + ) + cube.add_ancillary_variable(a_ancillary_variable, [0, 1]) + self.cube = cube + + def test_ancillary_variables_2d(self): + result = self.cube[0:2, 0:2] + self.assertEqual(len(result.ancillary_variables()), 1) + self.assertEqual( + result.shape, result.ancillary_variables()[0].data.shape + ) + + def test_ancillary_variables_1d(self): + result = self.cube[0, 0:2] + self.assertEqual(len(result.ancillary_variables()), 1) + self.assertEqual( + result.shape, result.ancillary_variables()[0].data.shape + ) + + class TestAncillaryVariables(tests.IrisTest): def setUp(self): cube = Cube(10 * np.arange(6).reshape(2, 3))