Skip to content

Commit 6a50076

Browse files
duncanwpduncanwp
committed
Adding support for partial collapse of multi-dimensional coordinates
1 parent 0b58150 commit 6a50076

File tree

4 files changed

+85
-57
lines changed

4 files changed

+85
-57
lines changed

‎lib/iris/coords.py‎

Lines changed: 11 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1144,20 +1144,12 @@ def collapsed(self, dims_to_collapse=None):
11441144
the specified dimensions.
11451145
11461146
Replaces the points & bounds with a simple bounded region.
1147-
1148-
.. note::
1149-
You cannot partially collapse a multi-dimensional coordinate. See
1150-
:ref:`cube.collapsed <partially_collapse_multi-dim_coord>` for more
1151-
information.
1152-
11531147
"""
1148+
# Ensure dims_to_collapse is a tuple to be able to pass through to numpy
11541149
if isinstance(dims_to_collapse, (int, np.integer)):
1155-
dims_to_collapse = [dims_to_collapse]
1156-
1157-
if dims_to_collapse is not None and \
1158-
set(range(self.ndim)) != set(dims_to_collapse):
1159-
raise ValueError('Cannot partially collapse a coordinate (%s).'
1160-
% self.name())
1150+
dims_to_collapse = (dims_to_collapse, )
1151+
if isinstance(dims_to_collapse, list):
1152+
dims_to_collapse = tuple(dims_to_collapse)
11611153

11621154
if np.issubdtype(self.dtype, np.str):
11631155
# Collapse the coordinate by serializing the points and
@@ -1191,27 +1183,16 @@ def serialize(x):
11911183
warnings.warn(msg.format(self.name()))
11921184

11931185
# Create bounds for the new collapsed coordinate.
1194-
item = self.core_bounds() if self.has_bounds() \
1186+
item = np.concatenate(self.core_bounds()) if self.has_bounds() \
11951187
else self.core_points()
1196-
lower, upper = item.min(), item.max()
1197-
bounds_dtype = item.dtype
1198-
# Ensure 2D shape of new bounds.
1199-
bounds = np.empty((1, 2), 'object')
1200-
bounds[0, 0] = lower
1201-
bounds[0, 1] = upper
1202-
# Create points for the new collapsed coordinate.
1203-
points_dtype = self.dtype
1204-
points = (float(lower) + float(upper)) * 0.5
1188+
1189+
# Calculate the bounds and points along the right dims
1190+
bounds = np.stack([item.min(axis=dims_to_collapse),
1191+
item.max(axis=dims_to_collapse)]).T
1192+
points = item.mean(axis=dims_to_collapse, dtype=self.dtype)
12051193

12061194
# Create the new collapsed coordinate.
1207-
if is_lazy_data(item):
1208-
bounds = multidim_lazy_stack(bounds)
1209-
coord = self.copy(points=points, bounds=bounds)
1210-
else:
1211-
bounds = np.concatenate(bounds)
1212-
bounds = np.array(bounds, dtype=bounds_dtype)
1213-
coord = self.copy(points=np.array(points, dtype=points_dtype),
1214-
bounds=bounds)
1195+
coord = self.copy(points=points, bounds=bounds)
12151196
return coord
12161197

12171198
def _guess_bounds(self, bound_position=0.5):

‎lib/iris/cube.py‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3111,22 +3111,6 @@ def collapsed(self, coords, aggregator, **kwargs):
31113111
31123112
cube.collapsed(['latitude', 'longitude'],
31133113
iris.analysis.VARIANCE)
3114-
3115-
.. _partially_collapse_multi-dim_coord:
3116-
3117-
.. note::
3118-
You cannot partially collapse a multi-dimensional coordinate. Doing
3119-
so would result in a partial collapse of the multi-dimensional
3120-
coordinate. Instead you must either:
3121-
* collapse in a single operation all cube axes that the
3122-
multi-dimensional coordinate spans,
3123-
* remove the multi-dimensional coordinate from the cube before
3124-
performing the collapse operation, or
3125-
* not collapse the coordinate at all.
3126-
3127-
Multi-dimensional derived coordinates will not prevent a successful
3128-
collapse operation.
3129-
31303114
"""
31313115
# Convert any coordinate names to coordinates
31323116
coords = self._as_list_of_coords(coords)

‎lib/iris/tests/test_analysis.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -310,6 +310,60 @@ def test_sum(self):
310310
np.testing.assert_array_equal(cube.data, np.array([6, 18, 17]))
311311

312312

313+
class TestAuxCoordCollapse(tests.IrisTest):
314+
315+
def setUp(self):
316+
from iris.analysis.cartography import area_weights
317+
self.cube_with_aux_coord = tests.stock.simple_4d_with_hybrid_height()
318+
319+
# Guess bounds to get the weights
320+
self.cube_with_aux_coord.coord('grid_latitude').guess_bounds()
321+
self.cube_with_aux_coord.coord('grid_longitude').guess_bounds()
322+
323+
self.weights = area_weights(self.cube_with_aux_coord, normalize=False)
324+
self.normalized_weights = area_weights(self.cube_with_aux_coord, normalize=True)
325+
326+
self.original_alt = self.cube_with_aux_coord.coord('altitude')
327+
# [[100, 101, 102, 103, 104, 105],
328+
# [106, 107, 108, 109, 110, 111],
329+
# [112, 113, 114, 115, 116, 117],
330+
# [118, 119, 120, 121, 122, 123],
331+
# [124, 125, 126, 127, 128, 129]]
332+
333+
def test_max(self):
334+
cube = self.cube_with_aux_coord.collapsed('grid_latitude', iris.analysis.MAX)
335+
np.testing.assert_array_equal(cube.coord('surface_altitude').points,
336+
np.array([112, 113, 114, 115, 116, 117]))
337+
338+
np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
339+
np.array([[100, 124],
340+
[101, 125],
341+
[102, 126],
342+
[103, 127],
343+
[104, 128],
344+
[105, 129]]))
345+
346+
# Check collapsing over the whole coord still works
347+
cube = self.cube_with_aux_coord.collapsed('altitude', iris.analysis.MAX)
348+
349+
np.testing.assert_array_equal(cube.coord('surface_altitude').points,
350+
np.array([114]))
351+
352+
np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
353+
np.array([[100, 129]]))
354+
355+
cube = self.cube_with_aux_coord.collapsed('grid_longitude', iris.analysis.MAX)
356+
357+
np.testing.assert_array_equal(cube.coord('surface_altitude').points,
358+
np.array([102, 108, 114, 120, 126]))
359+
360+
np.testing.assert_array_equal(cube.coord('surface_altitude').bounds,
361+
np.array([[100, 105],
362+
[106, 111],
363+
[112, 117],
364+
[118, 123],
365+
[124, 129]]))
366+
313367
class TestAggregator_mdtol_keyword(tests.IrisTest):
314368
def setUp(self):
315369
data = ma.array([[1, 2], [4, 5]], dtype=np.float32,

‎lib/iris/tests/unit/coords/test_Coord.py‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -277,19 +277,28 @@ def test_dim_1d(self):
277277
[[coord.bounds.min(), coord.bounds.max()]])
278278

279279
def test_numeric_nd(self):
280-
# Contiguous only defined for 2d bounds.
281-
coord = AuxCoord(points=np.array([3, 6, 9]),
282-
bounds=np.array([[1, 2, 4, 5],
280+
coord = AuxCoord(points=np.array([[1, 2, 4, 5],
283281
[4, 5, 7, 8],
284282
[7, 8, 10, 11]]))
285-
with self.assertRaises(ValueError):
286-
coord.collapsed()
287-
288-
def test_collapsed_overflow(self):
289-
coord = DimCoord(points=np.array([1493892000, 1493895600, 1493899200],
290-
dtype=np.int32))
291-
result = coord.collapsed()
292-
self.assertEqual(result.points, 1493895600)
283+
284+
collapsed_coord = coord.collapsed()
285+
self.assertArrayEqual(collapsed_coord.points, np.array([6]))
286+
self.assertArrayEqual(collapsed_coord.bounds, np.array([[1, 11]]))
287+
288+
# Test partially collapsing one dimension...
289+
collapsed_coord = coord.collapsed(1)
290+
self.assertArrayEqual(collapsed_coord.points, np.array([3., 6., 9.]))
291+
self.assertArrayEqual(collapsed_coord.bounds, np.array([[1, 5],
292+
[4, 8],
293+
[7, 11]]))
294+
295+
# ... and the other
296+
collapsed_coord = coord.collapsed(0)
297+
self.assertArrayEqual(collapsed_coord.points, np.array([4, 5, 7, 8]))
298+
self.assertArrayEqual(collapsed_coord.bounds, np.array([[1, 7],
299+
[2, 8],
300+
[4, 10],
301+
[5, 11]]))
293302

294303

295304
class Test_is_compatible(tests.IrisTest):

0 commit comments

Comments
 (0)