Skip to content

Commit

Permalink
Avoid realising auxiliary coordinates in concatenate (#5142)
Browse files Browse the repository at this point in the history 
* Use core_points instead of points

* Remove unused import

* Use core_bounds

* Correct test

* Remove unused import

* Correct test

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
@sloosvel @pre-commit-ci
sloosvel and pre-commit-ci[bot] committed Jun 15, 2023
1 parent eaa7ee4 commit 438f292
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 23 deletions.
75 changes: 53 additions & 22 deletions lib/iris/_concatenate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,19 +101,23 @@ def __new__(mcs, coord, dims):
"""
defn = coord.metadata
points_dtype = coord.points.dtype
bounds_dtype = coord.bounds.dtype if coord.bounds is not None else None
points_dtype = coord.core_points().dtype
bounds_dtype = (
coord.core_bounds().dtype
if coord.core_bounds() is not None
else None
)
kwargs = {}
# Add scalar flag metadata.
kwargs["scalar"] = coord.points.size == 1
kwargs["scalar"] = coord.core_points().size == 1
# Add circular flag metadata for dimensional coordinates.
if hasattr(coord, "circular"):
kwargs["circular"] = coord.circular
if isinstance(coord, iris.coords.DimCoord):
# Mix the monotonic ordering into the metadata.
if coord.points[0] == coord.points[-1]:
if coord.core_points()[0] == coord.core_points()[-1]:
order = _CONSTANT
elif coord.points[-1] > coord.points[0]:
elif coord.core_points()[-1] > coord.core_points()[0]:
order = _INCREASING
else:
order = _DECREASING
Expand Down Expand Up @@ -700,18 +704,27 @@ def _cmp(coord, other):
"""
# A candidate axis must have non-identical coordinate points.
candidate_axis = not array_equal(coord.points, other.points)
candidate_axis = not array_equal(
coord.core_points(), other.core_points()
)

if candidate_axis:
# Ensure both have equal availability of bounds.
result = (coord.bounds is None) == (other.bounds is None)
result = (coord.core_bounds() is None) == (
other.core_bounds() is None
)
else:
if coord.bounds is not None and other.bounds is not None:
if (
coord.core_bounds() is not None
and other.core_bounds() is not None
):
# Ensure equality of bounds.
result = array_equal(coord.bounds, other.bounds)
result = array_equal(coord.core_bounds(), other.core_bounds())
else:
# Ensure both have equal availability of bounds.
result = coord.bounds is None and other.bounds is None
result = (
coord.core_bounds() is None and other.core_bounds() is None
)

return result, candidate_axis

Expand Down Expand Up @@ -762,21 +775,37 @@ def _calculate_extents(self):
self.dim_extents = []
for coord, order in zip(self.dim_coords, self.dim_order):
if order == _CONSTANT or order == _INCREASING:
points = _Extent(coord.points[0], coord.points[-1])
if coord.bounds is not None:
points = _Extent(
coord.core_points()[0], coord.core_points()[-1]
)
if coord.core_bounds() is not None:
bounds = (
_Extent(coord.bounds[0, 0], coord.bounds[-1, 0]),
_Extent(coord.bounds[0, 1], coord.bounds[-1, 1]),
_Extent(
coord.core_bounds()[0, 0],
coord.core_bounds()[-1, 0],
),
_Extent(
coord.core_bounds()[0, 1],
coord.core_bounds()[-1, 1],
),
)
else:
bounds = None
else:
# The order must be decreasing ...
points = _Extent(coord.points[-1], coord.points[0])
if coord.bounds is not None:
points = _Extent(
coord.core_points()[-1], coord.core_points()[0]
)
if coord.core_bounds() is not None:
bounds = (
_Extent(coord.bounds[-1, 0], coord.bounds[0, 0]),
_Extent(coord.bounds[-1, 1], coord.bounds[0, 1]),
_Extent(
coord.core_bounds()[-1, 0],
coord.core_bounds()[0, 0],
),
_Extent(
coord.core_bounds()[-1, 1],
coord.core_bounds()[0, 1],
),
)
else:
bounds = None
Expand Down Expand Up @@ -1095,7 +1124,7 @@ def _build_aux_coordinates(self):
# Concatenate the points together.
dim = dims.index(self.axis)
points = [
skton.signature.aux_coords_and_dims[i].coord.points
skton.signature.aux_coords_and_dims[i].coord.core_points()
for skton in skeletons
]
points = np.concatenate(tuple(points), axis=dim)
Expand All @@ -1104,7 +1133,9 @@ def _build_aux_coordinates(self):
bnds = None
if coord.has_bounds():
bnds = [
skton.signature.aux_coords_and_dims[i].coord.bounds
skton.signature.aux_coords_and_dims[
i
].coord.core_bounds()
for skton in skeletons
]
bnds = np.concatenate(tuple(bnds), axis=dim)
Expand Down Expand Up @@ -1307,7 +1338,7 @@ def _build_dim_coordinates(self):

# Concatenate the points together for the nominated dimension.
points = [
skeleton.signature.dim_coords[dim_ind].points
skeleton.signature.dim_coords[dim_ind].core_points()
for skeleton in skeletons
]
points = np.concatenate(tuple(points))
Expand All @@ -1316,7 +1347,7 @@ def _build_dim_coordinates(self):
bounds = None
if self._cube_signature.dim_coords[dim_ind].has_bounds():
bounds = [
skeleton.signature.dim_coords[dim_ind].bounds
skeleton.signature.dim_coords[dim_ind].core_bounds()
for skeleton in skeletons
]
bounds = np.concatenate(tuple(bounds))
Expand Down
25 changes: 25 additions & 0 deletions lib/iris/tests/test_concatenate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
# before importing anything else.
import iris.tests as tests # isort:skip

import dask.array as da
import numpy as np
import numpy.ma as ma

Expand Down Expand Up @@ -800,6 +801,30 @@ def test_concat_2y2d_derived_x_y_xy(self):
self.assertEqual(result[0].shape, (6, 2))
self.assertEqual(result[0], com)

def test_concat_lazy_aux_coords(self):
cubes = []
y = (0, 2)
cube = _make_cube((2, 4), y, 2, aux="xy")
cubes.append(cube)
cubes.append(_make_cube((0, 2), y, 1, aux="xy"))
for cube in cubes:
cube.data = cube.lazy_data()
cube.coord("xy-aux").points = cube.coord("xy-aux").lazy_points()
bounds = da.arange(
4 * cube.coord("xy-aux").core_points().size
).reshape(cube.shape + (4,))
cube.coord("xy-aux").bounds = bounds
result = concatenate(cubes)

self.assertTrue(cubes[0].coord("xy-aux").has_lazy_points())
self.assertTrue(cubes[0].coord("xy-aux").has_lazy_bounds())

self.assertTrue(cubes[1].coord("xy-aux").has_lazy_points())
self.assertTrue(cubes[1].coord("xy-aux").has_lazy_bounds())

self.assertTrue(result[0].coord("xy-aux").has_lazy_points())
self.assertTrue(result[0].coord("xy-aux").has_lazy_bounds())


class TestMulti2D(tests.IrisTest):
def test_concat_4x2d_aux_xy(self):
Expand Down
25 changes: 24 additions & 1 deletion lib/iris/tests/unit/concatenate/test_concatenate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -355,14 +355,23 @@ def test_desc_bounds_all_singleton(self):


class TestConcatenate__dask(tests.IrisTest):
def build_lazy_cube(self, points, bounds=None, nx=4):
def build_lazy_cube(self, points, bounds=None, nx=4, aux_coords=False):
data = np.arange(len(points) * nx).reshape(len(points), nx)
data = as_lazy_data(data)
cube = iris.cube.Cube(data, standard_name="air_temperature", units="K")
lat = iris.coords.DimCoord(points, "latitude", bounds=bounds)
lon = iris.coords.DimCoord(np.arange(nx), "longitude")
cube.add_dim_coord(lat, 0)
cube.add_dim_coord(lon, 1)
if aux_coords:
bounds = np.arange(len(points) * nx * 4).reshape(
len(points), nx, 4
)
bounds = as_lazy_data(bounds)
aux_coord = iris.coords.AuxCoord(
data, var_name="aux_coord", bounds=bounds
)
cube.add_aux_coord(aux_coord, (0, 1))
return cube

def test_lazy_concatenate(self):
Expand All @@ -372,6 +381,20 @@ def test_lazy_concatenate(self):
self.assertTrue(cube.has_lazy_data())
self.assertFalse(ma.isMaskedArray(cube.data))

def test_lazy_concatenate_aux_coords(self):
c1 = self.build_lazy_cube([1, 2], aux_coords=True)
c2 = self.build_lazy_cube([3, 4, 5], aux_coords=True)
(result,) = concatenate([c1, c2])

self.assertTrue(c1.coord("aux_coord").has_lazy_points())
self.assertTrue(c1.coord("aux_coord").has_lazy_bounds())

self.assertTrue(c2.coord("aux_coord").has_lazy_points())
self.assertTrue(c2.coord("aux_coord").has_lazy_bounds())

self.assertTrue(result.coord("aux_coord").has_lazy_points())
self.assertTrue(result.coord("aux_coord").has_lazy_bounds())

def test_lazy_concatenate_masked_array_mixed_deferred(self):
c1 = self.build_lazy_cube([1, 2])
c2 = self.build_lazy_cube([3, 4, 5])
Expand Down

0 comments on commit 438f292

Please sign in to comment.