From 5d1180de1591df085c0626153e09df02b0b63d44 Mon Sep 17 00:00:00 2001 From: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> Date: Wed, 15 Jul 2026 13:24:47 +0100 Subject: [PATCH 1/6] `_MeshIndexSet` first pass (#7149) * Use Monotonicity error in Cube indexing, plus test coverage. * Implement MeshIndexSetMetadata. * Changes to cube.mesh type hinting. * Create experimental/mesh_coord_indexing.py . * Implement MeshIndexSet. * Render _MeshIndexSet in docs. * Easy CI fixes. * style: pre-commit fixes * Add TODO for later type hinting. * Missing sphinx-needs item. * Remove linkable UGRID footnotes to avoid duplication warning. * Fix doctest. * TODO comments. * Fix doctest. * TODO comment. * Fix doctest. * Achieve an accurate view of Mesh via Dask arrays and timestamps. * Remove Mesh.is_view_of. * Fixes after rough testing. * Prevent overzealous updates. * Disconnect new meshes from original meshes. * Review suggestions. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/src/conf.py | 2 +- lib/iris/analysis/_interpolation.py | 3 +- lib/iris/common/metadata.py | 108 + lib/iris/coords.py | 25 +- lib/iris/cube.py | 8 +- lib/iris/exceptions.py | 6 + lib/iris/experimental/mesh_coord_indexing.py | 151 ++ lib/iris/experimental/raster.py | 3 +- lib/iris/fileformats/netcdf/saver.py | 10 +- lib/iris/mesh/components.py | 1737 ++++++++++++----- lib/iris/pandas.py | 3 +- .../analysis/maths/test__arith__meshcoords.py | 2 + lib/iris/tests/unit/cube/test_Cube.py | 33 + .../unit/mesh/components/test_MeshCoord.py | 2 + .../mesh/utils/test_recombine_submeshes.py | 1 + 15 files changed, 1566 insertions(+), 528 deletions(-) create mode 100644 lib/iris/experimental/mesh_coord_indexing.py diff --git a/docs/src/conf.py b/docs/src/conf.py index 7b705763a3..ac8bb7fa45 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -252,7 +252,7 @@ def _dotv(version): "members": True, "member-order": "alphabetical", "undoc-members": True, - "private-members": False, + "private-members": "_MeshIndexSet", "special-members": False, # Enums are most valuable when documented as concisely as possible. "inherited-members": "Enum,IntEnum,ReprEnum,StrEnum", diff --git a/lib/iris/analysis/_interpolation.py b/lib/iris/analysis/_interpolation.py index 67b10727ec..1d3a7052cb 100644 --- a/lib/iris/analysis/_interpolation.py +++ b/lib/iris/analysis/_interpolation.py @@ -14,6 +14,7 @@ from iris._lazy_data import map_complete_blocks from iris.coords import AuxCoord, DimCoord +from iris.exceptions import MonotonicityError import iris.util _DEFAULT_DTYPE = np.float16 @@ -478,7 +479,7 @@ def _validate(self): # Check monotonic. if not iris.util.monotonic(coord.points, strict=True): msg = "Cannot interpolate over the non-monotonic coordinate {}." - raise ValueError(msg.format(coord.name())) + raise MonotonicityError(msg.format(coord.name())) def _points(self, sample_points, data, data_dims=None): """Interpolate at the specified points. diff --git a/lib/iris/common/metadata.py b/lib/iris/common/metadata.py index f4af7584b4..e652095547 100644 --- a/lib/iris/common/metadata.py +++ b/lib/iris/common/metadata.py @@ -1594,6 +1594,114 @@ def equal(self, other, lenient=None): return super().equal(other, lenient=lenient) +class MeshIndexSetMetadata(BaseMetadata): + """Metadata container for a :class:`~iris.mesh.components._MeshIndexSet`.""" + + _members = ("mesh", "location", "start_index") + + __slots__ = () + + @wraps(BaseMetadata.__eq__, assigned=("__doc__",), updated=()) + @lenient_service + def __eq__(self, other): + return super().__eq__(other) + + def _combine_lenient(self, other): + """Perform lenient combination of metadata members for _MeshIndexSet. + + Parameters + ---------- + other : MeshIndexSetMetadata + The other metadata participating in the lenient combination. + + Returns + ------- + A list of combined metadata member values. + + """ + + # It is actually "strict" : return None except where members are equal. + def func(field): + left = getattr(self, field) + right = getattr(other, field) + return left if left == right else None + + # Note that, we use "_members" not "_fields". + values = [func(field) for field in self._members] + # Perform lenient combination of the other parent members. + result = super()._combine_lenient(other) + result.extend(values) + + return result + + def _compare_lenient(self, other): + """Perform lenient equality of metadata members for _MeshIndexSet. + + Parameters + ---------- + other : MeshIndexSetMetadata + The other metadata participating in the lenient comparison. + + Returns + ------- + bool + + """ + # Perform "strict" comparison for the _MeshIndexSet specific members + # 'mesh', 'location', 'start_index' : for equality, they must all match. + result = all( + [getattr(self, field) == getattr(other, field) for field in self._members] + ) + if result: + # Perform lenient comparison of the other parent members. + result = super()._compare_lenient(other) + + return result + + def _difference_lenient(self, other): + """Perform lenient difference of metadata members for _MeshIndexSet. + + Parameters + ---------- + other : MeshIndexSetMetadata + The other metadata participating in the lenient difference. + + Returns + ------- + A list of different metadata member values. + + """ + + # Perform "strict" difference for location / axis. + def func(field): + left = getattr(self, field) + right = getattr(other, field) + return None if left == right else (left, right) + + # Note that, we use "_members" not "_fields". + values = [func(field) for field in self._members] + # Perform lenient difference of the other parent members. + result = super()._difference_lenient(other) + result.extend(values) + + return result + + @wraps(BaseMetadata.combine, assigned=("__doc__",), updated=()) + @lenient_service + def combine(self, other, lenient=None): + return super().combine(other, lenient=lenient) + + @wraps(BaseMetadata.difference, assigned=("__doc__",), updated=()) + @lenient_service + def difference(self, other, lenient=None): + return super().difference(other, lenient=lenient) + + @wraps(BaseMetadata.equal, assigned=("__doc__",), updated=()) + @lenient_service + def equal(self, other, lenient=None): + return super().equal(other, lenient=lenient) + + class MeshCoordMetadata(BaseMetadata): """Metadata container for a :class:`~iris.coords.MeshCoord`.""" diff --git a/lib/iris/coords.py b/lib/iris/coords.py index 795c6b802d..012ba2bcd8 100644 --- a/lib/iris/coords.py +++ b/lib/iris/coords.py @@ -277,7 +277,18 @@ def _sanitise_array(self, src, ndmin): @property def _values(self): """The _DimensionalMetadata values as a NumPy array.""" - return self._values_dm.data.view() + + def data_id(): + return id(self._values_dm.core_data()) + + original_id = data_id() + result = self._values_dm.data.view() + if data_id() != original_id: + # Realisation has occurred - potential effect on deferred mesh computations + # (MeshCoord, _MeshIndexSet). + for timestamp in self._mesh_timestamps: + timestamp.update() + return result @_values.setter def _values(self, values): @@ -2425,7 +2436,7 @@ def _guess_bounds(self, bound_position=0.5, monthly=False, yearly=False): # XXX Consider moving into DimCoord # ensure we have monotonic points if not self.is_monotonic(): - raise ValueError( + raise iris.exceptions.MonotonicityError( "Need monotonic points to generate bounds for %s" % self.name() ) @@ -2999,7 +3010,9 @@ def _new_points_requirements(self, points): raise TypeError(emsg.format(self.name(), self.__class__.__name__)) if points.size > 1 and not iris.util.monotonic(points, strict=True): emsg = "The {!r} {} points array must be strictly monotonic." - raise ValueError(emsg.format(self.name(), self.__class__.__name__)) + raise iris.exceptions.MonotonicityError( + emsg.format(self.name(), self.__class__.__name__) + ) @property def _values(self): @@ -3079,7 +3092,7 @@ def _new_bounds_requirements(self, bounds): ) if not monotonic: emsg = "The {!r} {} bounds array must be strictly monotonic." - raise ValueError( + raise iris.exceptions.MonotonicityError( emsg.format(self.name(), self.__class__.__name__) ) directions.add(direction) @@ -3089,7 +3102,9 @@ def _new_bounds_requirements(self, bounds): "The direction of monotonicity for {!r} {} must " "be consistent across all bounds." ) - raise ValueError(emsg.format(self.name(), self.__class__.__name__)) + raise iris.exceptions.MonotonicityError( + emsg.format(self.name(), self.__class__.__name__) + ) if n_bounds == 2: # Make ordering of bounds consistent with coord's direction diff --git a/lib/iris/cube.py b/lib/iris/cube.py index 44be3a63d7..3313e315f3 100644 --- a/lib/iris/cube.py +++ b/lib/iris/cube.py @@ -26,7 +26,7 @@ from functools import partial, reduce import itertools import operator -from typing import TYPE_CHECKING, Any, Optional, TypeAlias, TypeGuard +from typing import TYPE_CHECKING, Any, Literal, Optional, TypeAlias, TypeGuard import warnings from xml.dom.minidom import Document @@ -3289,8 +3289,7 @@ def new_ancillary_variable_dims(av_): coord_keys = tuple([full_slice[dim] for dim in self.coord_dims(coord)]) try: new_coord = coord[coord_keys] - except ValueError: - # TODO make this except more specific to catch monotonic error + except iris.exceptions.MonotonicityError: # Attempt to slice it by converting to AuxCoord first new_coord = iris.coords.AuxCoord.from_coord(coord)[coord_keys] aux_coords.append((new_coord, new_coord_dims(coord))) @@ -3315,8 +3314,7 @@ def new_ancillary_variable_dims(av_): else: dim_coords.append((new_coord, new_dims)) shape += new_coord.core_points().shape - except ValueError: - # TODO make this except more specific to catch monotonic error + except iris.exceptions.MonotonicityError: # Attempt to slice it by converting to AuxCoord first new_coord = iris.coords.AuxCoord.from_coord(coord)[coord_keys] aux_coords.append((new_coord, new_dims)) diff --git a/lib/iris/exceptions.py b/lib/iris/exceptions.py index 5589e03337..c93430227d 100644 --- a/lib/iris/exceptions.py +++ b/lib/iris/exceptions.py @@ -185,3 +185,9 @@ class CFParseError(IrisError): """Raised when a string associated with a CF defined syntax could not be parsed.""" pass + + +class MonotonicityError(ValueError): + """Raised when a coordinate values are not monotonic.""" + + pass diff --git a/lib/iris/experimental/mesh_coord_indexing.py b/lib/iris/experimental/mesh_coord_indexing.py new file mode 100644 index 0000000000..c1e6024263 --- /dev/null +++ b/lib/iris/experimental/mesh_coord_indexing.py @@ -0,0 +1,151 @@ +# Copyright Iris contributors +# +# This file is part of Iris and is released under the BSD license. +# See LICENSE in the root of the repository for full licensing details. +"""Experimental module for alternative modes of indexing a :class:`~iris.mesh.MeshCoord`. + +.. z_reference:: iris.experimental.mesh_coord_indexing + :tags: topic_mesh + + API reference + +:class:`Options` describes the available indexing modes. + +Select the desired option using the run-time setting :data:`SETTING`. + +Examples +-------- +.. testsetup:: + + import iris + from iris.experimental.mesh_coord_indexing import SETTING + my_mesh_cube = iris.load_cube(iris.sample_data_path("mesh_C4_synthetic_float.nc")) + + # Remove non-compliant content. + my_mesh = my_mesh_cube.mesh + wanted_roles = ["edge_node_connectivity", "face_node_connectivity"] + for conn in my_mesh.all_connectivities: + if conn is not None and conn.cf_role not in wanted_roles: + my_mesh.remove_connectivities(conn) + + # Capture original state. + original_setting = SETTING.value + +.. testcleanup:: + + # Restore original state. + SETTING.value = original_setting + +Here is a simple :class:`~iris.cube.Cube` with :class:`~iris.mesh.MeshCoord` s: + +>>> print(my_mesh_cube) +synthetic / (1) (-- : 96) + Mesh coordinates: + latitude x + longitude x + Mesh: + name Topology data of 2D unstructured mesh + location face + Attributes: + NCO 'netCDF Operators version 4.7.5 (Homepage = http://nco.sf.net, Code = h ...' + history 'Mon Apr 12 01:44:41 2021: ncap2 -s synthetic=float(synthetic) mesh_C4_synthetic.nc ...' + nco_openmp_thread_number 1 +>>> print(my_mesh_cube.aux_coords) +(, ) + +Here is the default indexing behaviour: + +>>> from iris.experimental import mesh_coord_indexing +>>> print(mesh_coord_indexing.SETTING.value) +Options.AUX_COORD +>>> indexed_cube = my_mesh_cube[:36] +>>> print(indexed_cube.aux_coords) +(, ) + +Set the indexing mode to return a new mesh: + +>>> mesh_coord_indexing.SETTING.value = mesh_coord_indexing.Options.NEW_MESH +>>> indexed_cube = my_mesh_cube[:36] +>>> print(indexed_cube.aux_coords) +() location(face) [...]+bounds shape(36,)>, ) location(face) [...]+bounds shape(36,)>) + +Or set via a context manager: + +>>> with mesh_coord_indexing.SETTING.context(mesh_coord_indexing.Options.AUX_COORD): +... indexed_cube = my_mesh_cube[:36] +... print(indexed_cube.aux_coords) +(, ) + +""" + +from contextlib import contextmanager +import enum +import threading + +# TODO: update the full documentation + + +class Options(enum.Enum): + """Options for what is returned when a :class:`~iris.mesh.MeshCoord` is indexed. + + See the module docstring for usage instructions: + :mod:`~iris.experimental.mesh_coord_indexing`. + """ + + AUX_COORD = enum.auto() + """The default. Convert the ``MeshCoord`` to a + :class:`~iris.coords.AuxCoord` and index that AuxCoord. + """ + + NEW_MESH = enum.auto() + """Index the :attr:`~iris.mesh.MeshCoord.mesh` of the ``MeshCoord`` to + produce a new :class:`~iris.mesh.MeshXY` instance, then return a new + :class:`~iris.mesh.MeshCoord` instance based on that new mesh. + """ + + MESH_INDEX_SET = enum.auto() + """**EXPERIMENTAL.** Produce a :class:`iris.mesh.components._MeshIndexSet` + instance that references the original :class:`~iris.mesh.MeshXY` instance, + then return a new :class:`~iris.mesh.MeshCoord` instance based on that new + index set. :class:`~iris.mesh.components._MeshIndexSet` is a read-only + indexed 'view' onto its original :class:`~iris.mesh.MeshXY`; behaviour of + this class may change from release to release while the design is + finalised. + """ + + +class _Setting(threading.local): + """Setting for what is returned when a :class:`~iris.mesh.MeshCoord` is indexed. + + See the module docstring for usage instructions: + :mod:`~iris.experimental.mesh_coord_indexing`. + """ + + def __init__(self): + self._value = Options.AUX_COORD + + @property + def value(self): + return self._value + + @value.setter + def value(self, value): + self._value = Options(value) + + @contextmanager + def context(self, value): + new_value = Options(value) + original_value = self._value + try: + self._value = new_value + yield + finally: + self._value = original_value + + +SETTING = _Setting() +""" +Run-time setting for alternative modes of indexing a +:class:`~iris.mesh.MeshCoord`. See the module docstring for usage +instructions: :mod:`~iris.experimental.mesh_coord_indexing`. +""" diff --git a/lib/iris/experimental/raster.py b/lib/iris/experimental/raster.py index 0b5057136c..c491f5bd30 100644 --- a/lib/iris/experimental/raster.py +++ b/lib/iris/experimental/raster.py @@ -24,6 +24,7 @@ import iris from iris._deprecation import warn_deprecated import iris.coord_systems +from iris.exceptions import MonotonicityError wmsg = ( "iris.experimental.raster is deprecated since version 3.2, and will be " @@ -179,7 +180,7 @@ def export_geotiff(cube, fname): raise ValueError(msg) if coord_x.points[0] > coord_x.points[-1]: - raise ValueError( + raise MonotonicityError( "Coordinate {!r} x-points must be monotonically increasing.".format(name) ) diff --git a/lib/iris/fileformats/netcdf/saver.py b/lib/iris/fileformats/netcdf/saver.py index 98c565e990..50704404e8 100644 --- a/lib/iris/fileformats/netcdf/saver.py +++ b/lib/iris/fileformats/netcdf/saver.py @@ -986,7 +986,6 @@ def _add_aux_coords( MeshEdgeCoords, MeshFaceCoords, MeshNodeCoords, - MeshXY, ) # Exclude any mesh coords, which are bundled in with the aux-coords. @@ -995,12 +994,11 @@ def _add_aux_coords( ] # Include any relevant mesh location coordinates. - mesh: MeshXY | None = getattr(cube, "mesh") - mesh_location: str | None = getattr(cube, "location") + mesh = getattr(cube, "mesh") + mesh_location = getattr(cube, "location") + coords_types = typing.Union[MeshNodeCoords, MeshEdgeCoords, MeshFaceCoords] if mesh and mesh_location: - location_coords: MeshNodeCoords | MeshEdgeCoords | MeshFaceCoords = getattr( - mesh, f"{mesh_location}_coords" - ) + location_coords: coords_types = getattr(mesh, f"{mesh_location}_coords") # type: ignore[annotation-unchecked] coords_to_add.extend(list(location_coords)) return self._add_inner_related_vars( diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 3e61d7c0d4..396c05ddcc 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -18,15 +18,24 @@ from collections import namedtuple from collections.abc import Container from contextlib import contextmanager +from copy import deepcopy from datetime import datetime -from typing import Iterable, Literal +import functools +from typing import Any, Iterable, Literal, Optional, TypeAlias import warnings from cf_units import Unit from dask import array as da import numpy as np +from numpy.typing import ArrayLike -from iris.common.metadata import ConnectivityMetadata, MeshCoordMetadata, MeshMetadata +from iris.common.metadata import ( + ConnectivityMetadata, + MeshCoordMetadata, + MeshIndexSetMetadata, + MeshMetadata, +) +import iris.util from .. import _lazy_data as _lazy from ..common import CFVariableMixin, metadata_filter, metadata_manager_factory @@ -617,6 +626,7 @@ class Mesh(CFVariableMixin, ABC): - Move whatever is appropriate from :class:`MeshXY` into this class, leaving behind only those elements specific to the assumption of X and Y node coordinates. + - Set :class:`_MeshIndexSet` to subclass `Mesh` instead of `MeshXY`. - Remove the docstring warning, the NotImplementedError, and the uses of ABC/abstractmethod. - Add a cross-reference in the docstring for :class:`MeshXY`. @@ -635,27 +645,14 @@ def __init__(self): raise NotImplementedError(message) -class MeshXY(Mesh): - """A container representing the UGRID ``cf_role`` ``mesh_topology``. - - A container representing the UGRID [1]_ ``cf_role`` ``mesh_topology``, supporting - 1D network, 2D triangular, and 2D flexible mesh topologies. - - Based on the assumption of 2 :attr:`node_coords` - one associated with the - X-axis (e.g. longitude) and 1 with the Y-axis (e.g. latitude). UGRID - describing alternative node coordinates (e.g. spherical) cannot be - represented. - - Notes - ----- - The 3D layered and fully 3D unstructured mesh topologies are not supported - at this time. - - References - ---------- - .. [1] The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ - - """ +class _MeshXYMixin(Mesh, ABC): + # Subclass __init__ methods must define: + # TODO: Impossible to type hint the return type of metadata_manager_factory(). + _metadata_manager: Any + # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. + _connectivity_manager_attr: Any + # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. + _coord_manager_attr: Any # TBD: for volume and/or z-axis support include axis "z" and/or dimension "3" #: The supported mesh axes. @@ -665,382 +662,96 @@ class MeshXY(Mesh): #: Valid mesh elements. ELEMENTS = ("edge", "node", "face") - def __init__( - self, - topology_dimension, - node_coords_and_axes, - connectivities, - edge_coords_and_axes=None, - face_coords_and_axes=None, - standard_name=None, - long_name=None, - var_name=None, - units=None, - attributes=None, - node_dimension=None, - edge_dimension=None, - face_dimension=None, - ): - """MeshXY initialise. - - .. note:: - - The purpose of the :attr:`node_dimension`, :attr:`edge_dimension` and - :attr:`face_dimension` properties are to preserve the original NetCDF - variable dimension names. Note that, only :attr:`edge_dimension` and - :attr:`face_dimension` are UGRID attributes, and are only present for - :attr:`topology_dimension` ``>=2``. - - """ - # TODO: support volumes. - # TODO: support (coord, "z") - - self._metadata_manager = metadata_manager_factory(MeshMetadata) - - # topology_dimension is read-only, so assign directly to the metadata manager - if topology_dimension not in self.TOPOLOGY_DIMENSIONS: - emsg = f"Expected 'topology_dimension' in range {self.TOPOLOGY_DIMENSIONS!r}, got {topology_dimension!r}." - raise ValueError(emsg) - self._metadata_manager.topology_dimension = topology_dimension - - self.node_dimension = node_dimension - self.edge_dimension = edge_dimension - self.face_dimension = face_dimension - - # assign the metadata to the metadata manager - self.standard_name = standard_name - self.long_name = long_name - self.var_name = var_name - self.units = units - self.attributes = attributes - - # based on the topology_dimension, create the appropriate coordinate manager - def normalise(element, axis): - result = str(axis).lower() - if result not in self.AXES: - emsg = f"Invalid axis specified for {element} coordinate {coord.name()!r}, got {axis!r}." - raise ValueError(emsg) - return f"{element}_{result}" - - if not isinstance(node_coords_and_axes, Iterable): - node_coords_and_axes = [node_coords_and_axes] - - if not isinstance(connectivities, Iterable): - connectivities = [connectivities] + def __eq__(self, other): + result = NotImplemented - kwargs = {} - for coord, axis in node_coords_and_axes: - kwargs[normalise("node", axis)] = coord - if edge_coords_and_axes is not None: - for coord, axis in edge_coords_and_axes: - kwargs[normalise("edge", axis)] = coord - if face_coords_and_axes is not None: - for coord, axis in face_coords_and_axes: - kwargs[normalise("face", axis)] = coord + if isinstance(other, _MeshXYMixin): + result = self.metadata == other.metadata + if result: + result = self.all_coords == other.all_coords + if result: + result = self.all_connectivities == other.all_connectivities - # check the UGRID minimum requirement for coordinates - if "node_x" not in kwargs: - emsg = "Require a node coordinate that is x-axis like to be provided." - raise ValueError(emsg) - if "node_y" not in kwargs: - emsg = "Require a node coordinate that is y-axis like to be provided." - raise ValueError(emsg) + return result - if self.topology_dimension == 1: - self._coord_manager = _Mesh1DCoordinateManager(**kwargs) - self._connectivity_manager = _Mesh1DConnectivityManager(*connectivities) - elif self.topology_dimension == 2: - self._coord_manager = _Mesh2DCoordinateManager(**kwargs) - self._connectivity_manager = _Mesh2DConnectivityManager(*connectivities) - else: - emsg = f"Unsupported 'topology_dimension', got {topology_dimension!r}." - raise NotImplementedError(emsg) + def __hash__(self): + # Allow use in sets and as dictionary keys, as is done for :class:`iris.cube.Cube`. + # See https://github.com/SciTools/iris/pull/1772 + return hash(id(self)) - @classmethod - def from_coords(cls, *coords): - r"""Construct a :class:`MeshXY` by derivation from 1/more :class:`~iris.coords.Coord`. + def __getstate__(self): + return ( + self._metadata_manager, + self._coord_manager, + self._connectivity_manager, + ) - The :attr:`~MeshXY.topology_dimension`, :class:`~iris.coords.Coord` - membership and :class:`Connectivity` membership are all determined - based on the shape of the first :attr:`~iris.coords.Coord.bounds`: + def __ne__(self, other): + result = self.__eq__(other) + if result is not NotImplemented: + result = not result + return result - * ``None`` or ``(n, <2)``: - Not supported - * ``(n, 2)``: - :attr:`~MeshXY.topology_dimension` = ``1``. - :attr:`~MeshXY.node_coords` and :attr:`~MeshXY.edge_node_connectivity` - constructed from :attr:`~iris.coords.Coord.bounds`. - :attr:`~MeshXY.edge_coords` constructed from - :attr:`~iris.coords.Coord.points`. - * ``(n, >=3)``: - :attr:`~MeshXY.topology_dimension` = ``2``. - :attr:`~MeshXY.node_coords` and :attr:`~MeshXY.face_node_connectivity` - constructed from :attr:`~iris.coords.Coord.bounds`. - :attr:`~MeshXY.face_coords` constructed from - :attr:`~iris.coords.Coord.points`. + def summary(self, *args, **kwargs): + """Return a string representation of the MeshXY. Parameters ---------- - *coords : Iterable of :class:`~iris.coords.Coord` - Coordinates to pass into the :class:`MeshXY`. - All :attr:`~iris.coords.Coord.points` must have the same shapes; - all :attr:`~iris.coords.Coord.bounds` must have the same shapes, - and must not be ``None``. + shorten : bool, default=False + If True, produce a oneline string form of the form . + If False, produce a multi-line detailed print output. Returns ------- - :class:`MeshXY` - - Notes - ----- - .. note:: - Any resulting duplicate nodes are not currently removed, due to the - computational intensity. - - .. note:: - :class:`MeshXY` currently requires ``X`` and ``Y`` - :class:`~iris.coords.Coord` specifically. - :meth:`iris.util.guess_coord_axis` is therefore attempted, else the - first two :class:`~iris.coords.Coord` are taken. - - .. testsetup:: - - from iris import load_cube, sample_data_path - from iris.mesh import ( - MeshXY, - MeshCoord, - ) - - file_path = sample_data_path("mesh_C4_synthetic_float.nc") - cube_w_mesh = load_cube(file_path) + str - Examples - -------- - :: + """ + if len(args) > 0: + shorten = args[0] + else: + shorten = kwargs.get("shorten", False) - # Reconstruct a cube-with-mesh after subsetting it. + if shorten: + result = self._summary_oneline() + else: + result = self._summary_multiline() + return result - >>> print(cube_w_mesh.mesh.name()) - Topology data of 2D unstructured mesh - >>> mesh_coord_names = [ - ... coord.name() for coord in cube_w_mesh.coords(mesh_coords=True) - ... ] - >>> print(f"MeshCoords: {mesh_coord_names}") - MeshCoords: ['latitude', 'longitude'] + def __repr__(self): + return self.summary(shorten=True) - # Subsetting converts MeshCoords to AuxCoords. - >>> slices = [slice(None)] * cube_w_mesh.ndim - >>> slices[cube_w_mesh.mesh_dim()] = slice(-1) - >>> cube_sub = cube_w_mesh[tuple(slices)] - >>> print(cube_sub.mesh) - None - >>> orig_coords = [cube_sub.coord(c_name) for c_name in mesh_coord_names] - >>> for coord in orig_coords: - ... print(f"{coord.name()}: {type(coord).__name__}") - latitude: AuxCoord - longitude: AuxCoord + def __str__(self): + return self.summary(shorten=False) - >>> new_mesh = MeshXY.from_coords(*orig_coords) - >>> new_coords = new_mesh.to_MeshCoords(location=cube_w_mesh.location) + def _summary_oneline(self): + # We use the repr output to produce short one-line identity summary, + # similar to the object.__str__ output "". + # This form also used in other str() constructions, like MeshCoord. + # By contrast, __str__ (below) produces a readable multi-line printout. + mesh_name = self.name() + if mesh_name in (None, "", "unknown"): + mesh_name = None + if mesh_name: + # Use a more human-readable form + mesh_string = f"<{self.__class__.__name__}: '{mesh_name}'>" + else: + # Mimic the generic object.__str__ style. + mesh_id = id(self) + mesh_string = f"<{self.__class__.__name__} object at {hex(mesh_id)}>" - # Replace the AuxCoords with MeshCoords. - >>> for ix in range(2): - ... cube_sub.remove_coord(orig_coords[ix]) - ... cube_sub.add_aux_coord(new_coords[ix], cube_w_mesh.mesh_dim()) + return mesh_string - >>> print(cube_sub.mesh.name()) - Topology data of 2D unstructured mesh - >>> for coord_name in mesh_coord_names: - ... coord = cube_sub.coord(coord_name) - ... print(f"{coord_name}: {type(coord).__name__}") - latitude: MeshCoord - longitude: MeshCoord + def _summary_multiline(self): + # Produce a readable multi-line summary of the Mesh content. + lines = [] + n_indent = 4 + indent_str = " " * n_indent - """ + def line(text, i_indent=0): + indent = indent_str * i_indent + lines.append(f"{indent}{text}") - # Validate points and bounds shape match. - def check_shape(array_name): - attr_name = f"core_{array_name}" - arrays = [getattr(coord, attr_name)() for coord in coords] - if any(a is None for a in arrays): - message = f"{array_name} missing from coords[{arrays.index(None)}] ." - raise ValueError(message) - shapes = [array.shape for array in arrays] - if shapes.count(shapes[0]) != len(shapes): - message = f"{array_name} shapes are not identical for all coords." - raise ValueError(message) - - for array in ("points", "bounds"): - check_shape(array) - - # Determine dimensionality, using first coord. - first_coord = coords[0] - - ndim = first_coord.ndim - if ndim != 1: - message = f"Expected coordinate ndim == 1, got: f{ndim} ." - raise ValueError(message) - - bounds_shape = first_coord.core_bounds().shape - bounds_dim1 = bounds_shape[1] - if bounds_dim1 < 2: - message = ( - f"Expected coordinate bounds.shape (n, >=2), got: {bounds_shape} ." - ) - raise ValueError(message) - elif bounds_dim1 == 2: - topology_dimension = 1 - coord_centring = "edge" - conn_cf_role = "edge_node_connectivity" - else: - topology_dimension = 2 - coord_centring = "face" - conn_cf_role = "face_node_connectivity" - - # Create connectivity. - if first_coord.has_lazy_bounds(): - array_lib = da - else: - array_lib = np - indices = array_lib.arange(np.prod(bounds_shape)).reshape(bounds_shape) - masking = array_lib.ma.getmaskarray(first_coord.core_bounds()) - indices = array_lib.ma.masked_array(indices, masking) - connectivity = Connectivity(indices, conn_cf_role) - - # Create coords. - node_coords = [] - centre_coords = [] - for coord in coords: - coord_kwargs = dict( - standard_name=coord.standard_name, - long_name=coord.long_name, - units=coord.units, - attributes=coord.attributes, - ) - node_points = array_lib.ma.filled(coord.core_bounds(), 0.0).flatten() - node_coords.append(AuxCoord(points=node_points, **coord_kwargs)) - - centre_points = coord.core_points() - centre_coords.append(AuxCoord(points=centre_points, **coord_kwargs)) - - ##### - # TODO: remove axis assignment once Mesh supports arbitrary coords. - # TODO: consider filtering coords as the first action in this method. - axes_present = [guess_coord_axis(coord) for coord in coords] - axes_required = ("X", "Y") - if all([req in axes_present for req in axes_required]): - axis_indices = [axes_present.index(req) for req in axes_required] - else: - message = ( - "Unable to find 'X' and 'Y' using guess_coord_axis. Assuming " - "X=coords[0], Y=coords[1] ." - ) - # TODO: reconsider logging level when we have consistent practice. - logger.info(message, extra=dict(cls=None)) - axis_indices = range(len(axes_required)) - - def axes_assign(coord_list): - coords_sorted = [coord_list[ix] for ix in axis_indices] - return zip(coords_sorted, axes_required) - - node_coords_and_axes = axes_assign(node_coords) - centre_coords_and_axes = axes_assign(centre_coords) - ##### - - # Construct the Mesh. - mesh_kwargs = dict( - topology_dimension=topology_dimension, - node_coords_and_axes=node_coords_and_axes, - connectivities=[connectivity], - ) - mesh_kwargs[f"{coord_centring}_coords_and_axes"] = centre_coords_and_axes - return cls(**mesh_kwargs) - - def __eq__(self, other): - result = NotImplemented - - if isinstance(other, MeshXY): - result = self.metadata == other.metadata - if result: - result = self.all_coords == other.all_coords - if result: - result = self.all_connectivities == other.all_connectivities - - return result - - def __hash__(self): - # Allow use in sets and as dictionary keys, as is done for :class:`iris.cube.Cube`. - # See https://github.com/SciTools/iris/pull/1772 - return hash(id(self)) - - def __getstate__(self): - return ( - self._metadata_manager, - self._coord_manager, - self._connectivity_manager, - ) - - def __ne__(self, other): - result = self.__eq__(other) - if result is not NotImplemented: - result = not result - return result - - def summary(self, shorten=False): - """Return a string representation of the MeshXY. - - Parameters - ---------- - shorten : bool, default=False - If True, produce a oneline string form of the form . - If False, produce a multi-line detailed print output. - - Returns - ------- - str - - """ - if shorten: - result = self._summary_oneline() - else: - result = self._summary_multiline() - return result - - def __repr__(self): - return self.summary(shorten=True) - - def __str__(self): - return self.summary(shorten=False) - - def _summary_oneline(self): - # We use the repr output to produce short one-line identity summary, - # similar to the object.__str__ output "". - # This form also used in other str() constructions, like MeshCoord. - # By contrast, __str__ (below) produces a readable multi-line printout. - mesh_name = self.name() - if mesh_name in (None, "", "unknown"): - mesh_name = None - if mesh_name: - # Use a more human-readable form - mesh_string = f"" - else: - # Mimic the generic object.__str__ style. - mesh_id = id(self) - mesh_string = f"" - - return mesh_string - - def _summary_multiline(self): - # Produce a readable multi-line summary of the Mesh content. - lines = [] - n_indent = 4 - indent_str = " " * n_indent - - def line(text, i_indent=0): - indent = indent_str * i_indent - lines.append(f"{indent}{text}") - - line(f"MeshXY : '{self.name()}'") + line(f"{self.__class__.__name__} : '{self.name()}'") line(f"topology_dimension: {self.topology_dimension}", 1) for element in ("node", "edge", "face"): if element == "node": @@ -1124,34 +835,29 @@ def __setstate__(self, state): self._coord_manager = coord_manager self._connectivity_manager = connectivity_manager - def _set_dimension_names(self, node, edge, face, reset=False): - args = (node, edge, face) - currents = ( - self.node_dimension, - self.edge_dimension, - self.face_dimension, - ) - zipped = zip(args, currents) - if reset: - node, edge, face = [None if arg else current for arg, current in zipped] - else: - node, edge, face = [arg or current for arg, current in zipped] + # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. + @property + def _connectivity_manager(self): + # @property enables interruption/customisation in subclasses. + return self._connectivity_manager_attr - self.node_dimension = node - self.edge_dimension = edge - self.face_dimension = face + # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. + @_connectivity_manager.setter + def _connectivity_manager(self, manager) -> None: + # @property enables interruption/customisation in subclasses. + self._connectivity_manager_attr = manager - if self.topology_dimension == 1: - result = Mesh1DNames(self.node_dimension, self.edge_dimension) - elif self.topology_dimension == 2: - result = Mesh2DNames( - self.node_dimension, self.edge_dimension, self.face_dimension - ) - else: - message = f"Unsupported topology_dimension: {self.topology_dimension} ." - raise NotImplementedError(message) + # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. + @property + def _coord_manager(self): + # @property enables interruption/customisation in subclasses. + return self._coord_manager_attr - return result + # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. + @_coord_manager.setter + def _coord_manager(self, manager) -> None: + # @property enables interruption/customisation in subclasses. + self._coord_manager_attr = manager @property def all_connectivities(self): @@ -1187,17 +893,9 @@ def edge_coords(self): return self._coord_manager.edge_coords @property - def edge_dimension(self): - """The *optionally required* UGRID NetCDF variable name for the ``edge`` dimension.""" - return self._metadata_manager.edge_dimension - - @edge_dimension.setter - def edge_dimension(self, name): - if not name or not isinstance(name, str): - edge_dimension = f"Mesh{self.topology_dimension}d_edge" - else: - edge_dimension = name - self._metadata_manager.edge_dimension = edge_dimension + @abstractmethod + def edge_dimension(self) -> str: + raise NotImplementedError() @property def edge_face_connectivity(self): @@ -1229,26 +927,9 @@ def face_coords(self): return self._coord_manager.face_coords @property - def face_dimension(self): - """The *optional* UGRID NetCDF variable name for the ``face`` dimension.""" - return self._metadata_manager.face_dimension - - @face_dimension.setter - def face_dimension(self, name): - if self.topology_dimension < 2: - face_dimension = None - if name: - # Tell the user it is not being set if they expected otherwise. - message = ( - "Not setting face_dimension (inappropriate for " - f"topology_dimension={self.topology_dimension} ." - ) - logger.debug(message, extra=dict(cls=self.__class__.__name__)) - elif not name or not isinstance(name, str): - face_dimension = f"Mesh{self.topology_dimension}d_face" - else: - face_dimension = name - self._metadata_manager.face_dimension = face_dimension + @abstractmethod + def face_dimension(self) -> str: + raise NotImplementedError() @property def face_edge_connectivity(self): @@ -1292,17 +973,9 @@ def node_coords(self): return self._coord_manager.node_coords @property - def node_dimension(self): - """The NetCDF variable name for the ``node`` dimension.""" - return self._metadata_manager.node_dimension - - @node_dimension.setter - def node_dimension(self, name): - if not name or not isinstance(name, str): - node_dimension = f"Mesh{self.topology_dimension}d_node" - else: - node_dimension = name - self._metadata_manager.node_dimension = node_dimension + @abstractmethod + def node_dimension(self) -> str: + raise NotImplementedError() def add_connectivities(self, *connectivities): """Add one or more :class:`~iris.mesh.Connectivity` instances to the :class:`MeshXY`. @@ -1916,26 +1589,429 @@ def to_MeshCoords(self, location): the current :class:`MeshXY`, one for each :attr:`AXES` value, passing through the ``location`` argument. - .. seealso:: + .. seealso:: + + :meth:`to_MeshCoord` for generating a single mesh coord. + + Parameters + ---------- + location : str + The ``location`` argument for :class:`MeshCoord` instantiation. + + Returns + ------- + tuple of :class:`~iris.mesh.mesh.MeshCoord` + Tuple of :class:`~iris.mesh.mesh.MeshCoord` + referencing the current :class:`MeshXY`. One for each value in + :attr:`AXES`, using the value for the ``axis`` argument. + + """ + # factory method + result = [self.to_MeshCoord(location=location, axis=ax) for ax in self.AXES] + return tuple(result) + + @property + @abstractmethod + def cf_role(self) -> str: + raise NotImplementedError() + + @property + @abstractmethod + def topology_dimension(self) -> int: + raise NotImplementedError() + + +class MeshXY(_MeshXYMixin): + """A container representing the UGRID ``cf_role`` ``mesh_topology``. + + A container representing the UGRID ``cf_role`` ``mesh_topology``, supporting + 1D network, 2D triangular, and 2D flexible mesh topologies. + + Based on the assumption of 2 :attr:`node_coords` - one associated with the + X-axis (e.g. longitude) and 1 with the Y-axis (e.g. latitude). UGRID + describing alternative node coordinates (e.g. spherical) cannot be + represented. + + Notes + ----- + The 3D layered and fully 3D unstructured mesh topologies are not supported + at this time. + + References + ---------- + 1. The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ + + """ + + def __init__( + self, + topology_dimension, + node_coords_and_axes, + connectivities, + edge_coords_and_axes=None, + face_coords_and_axes=None, + standard_name=None, + long_name=None, + var_name=None, + units=None, + attributes=None, + node_dimension=None, + edge_dimension=None, + face_dimension=None, + ): + """MeshXY initialise. + + .. note:: + + The purpose of the :attr:`node_dimension`, :attr:`edge_dimension` and + :attr:`face_dimension` properties are to preserve the original NetCDF + variable dimension names. Note that, only :attr:`edge_dimension` and + :attr:`face_dimension` are UGRID attributes, and are only present for + :attr:`topology_dimension` ``>=2``. + + """ + # TODO: support volumes. + # TODO: support (coord, "z") + + self._metadata_manager = metadata_manager_factory(MeshMetadata) + + # topology_dimension is read-only, so assign directly to the metadata manager + if topology_dimension not in self.TOPOLOGY_DIMENSIONS: + emsg = f"Expected 'topology_dimension' in range {self.TOPOLOGY_DIMENSIONS!r}, got {topology_dimension!r}." + raise ValueError(emsg) + self._metadata_manager.topology_dimension = topology_dimension + + self.node_dimension = node_dimension + self.edge_dimension = edge_dimension + self.face_dimension = face_dimension + + # assign the metadata to the metadata manager + self.standard_name = standard_name + self.long_name = long_name + self.var_name = var_name + self.units = units + self.attributes = attributes + + # based on the topology_dimension, create the appropriate coordinate manager + def normalise(element, axis): + result = str(axis).lower() + if result not in self.AXES: + emsg = f"Invalid axis specified for {element} coordinate {coord.name()!r}, got {axis!r}." + raise ValueError(emsg) + return f"{element}_{result}" + + if not isinstance(node_coords_and_axes, Iterable): + node_coords_and_axes = [node_coords_and_axes] + + if not isinstance(connectivities, Iterable): + connectivities = [connectivities] + + kwargs = {} + for coord, axis in node_coords_and_axes: + kwargs[normalise("node", axis)] = coord + if edge_coords_and_axes is not None: + for coord, axis in edge_coords_and_axes: + kwargs[normalise("edge", axis)] = coord + if face_coords_and_axes is not None: + for coord, axis in face_coords_and_axes: + kwargs[normalise("face", axis)] = coord + + # check the UGRID minimum requirement for coordinates + if "node_x" not in kwargs: + emsg = "Require a node coordinate that is x-axis like to be provided." + raise ValueError(emsg) + if "node_y" not in kwargs: + emsg = "Require a node coordinate that is y-axis like to be provided." + raise ValueError(emsg) + + if self.topology_dimension == 1: + self._coord_manager = _Mesh1DCoordinateManager(**kwargs) + self._connectivity_manager = _Mesh1DConnectivityManager(*connectivities) + elif self.topology_dimension == 2: + self._coord_manager = _Mesh2DCoordinateManager(**kwargs) + self._connectivity_manager = _Mesh2DConnectivityManager(*connectivities) + else: + emsg = f"Unsupported 'topology_dimension', got {topology_dimension!r}." + raise NotImplementedError(emsg) + + # TODO: backwards compatibility: make from_coords() perform a no-op if the given + # coords are already MeshCoords. + @classmethod + def from_coords(cls, *coords): + r"""Construct a :class:`MeshXY` by derivation from 1/more :class:`~iris.coords.Coord`. + + The :attr:`~MeshXY.topology_dimension`, :class:`~iris.coords.Coord` + membership and :class:`Connectivity` membership are all determined + based on the shape of the first :attr:`~iris.coords.Coord.bounds`: + + * ``None`` or ``(n, <2)``: + Not supported + * ``(n, 2)``: + :attr:`~MeshXY.topology_dimension` = ``1``. + :attr:`~MeshXY.node_coords` and :attr:`~MeshXY.edge_node_connectivity` + constructed from :attr:`~iris.coords.Coord.bounds`. + :attr:`~MeshXY.edge_coords` constructed from + :attr:`~iris.coords.Coord.points`. + * ``(n, >=3)``: + :attr:`~MeshXY.topology_dimension` = ``2``. + :attr:`~MeshXY.node_coords` and :attr:`~MeshXY.face_node_connectivity` + constructed from :attr:`~iris.coords.Coord.bounds`. + :attr:`~MeshXY.face_coords` constructed from + :attr:`~iris.coords.Coord.points`. + + Parameters + ---------- + *coords : Iterable of :class:`~iris.coords.Coord` + Coordinates to pass into the :class:`MeshXY`. + All :attr:`~iris.coords.Coord.points` must have the same shapes; + all :attr:`~iris.coords.Coord.bounds` must have the same shapes, + and must not be ``None``. + + Returns + ------- + :class:`MeshXY` + + Notes + ----- + .. note:: + Any resulting duplicate nodes are not currently removed, due to the + computational intensity. + + .. note:: + :class:`MeshXY` currently requires ``X`` and ``Y`` + :class:`~iris.coords.Coord` specifically. + :meth:`iris.util.guess_coord_axis` is therefore attempted, else the + first two :class:`~iris.coords.Coord` are taken. + + .. testsetup:: + + from iris import load_cube, sample_data_path + from iris.mesh import ( + MeshXY, + MeshCoord, + ) + + file_path = sample_data_path("mesh_C4_synthetic_float.nc") + cube_w_mesh = load_cube(file_path) + + Examples + -------- + :: + + # Reconstruct a cube-with-mesh after subsetting it. + + >>> print(cube_w_mesh.mesh.name()) + Topology data of 2D unstructured mesh + >>> mesh_coord_names = [ + ... coord.name() for coord in cube_w_mesh.coords(mesh_coords=True) + ... ] + >>> print(f"MeshCoords: {mesh_coord_names}") + MeshCoords: ['latitude', 'longitude'] + + # Subsetting converts MeshCoords to AuxCoords. + >>> slices = [slice(None)] * cube_w_mesh.ndim + >>> slices[cube_w_mesh.mesh_dim()] = slice(-1) + >>> cube_sub = cube_w_mesh[tuple(slices)] + >>> print(cube_sub.mesh) + None + >>> orig_coords = [cube_sub.coord(c_name) for c_name in mesh_coord_names] + >>> for coord in orig_coords: + ... print(f"{coord.name()}: {type(coord).__name__}") + latitude: AuxCoord + longitude: AuxCoord + + >>> new_mesh = MeshXY.from_coords(*orig_coords) + >>> new_coords = new_mesh.to_MeshCoords(location=cube_w_mesh.location) + + # Replace the AuxCoords with MeshCoords. + >>> for ix in range(2): + ... cube_sub.remove_coord(orig_coords[ix]) + ... cube_sub.add_aux_coord(new_coords[ix], cube_w_mesh.mesh_dim()) + + >>> print(cube_sub.mesh.name()) + Topology data of 2D unstructured mesh + >>> for coord_name in mesh_coord_names: + ... coord = cube_sub.coord(coord_name) + ... print(f"{coord_name}: {type(coord).__name__}") + latitude: MeshCoord + longitude: MeshCoord + + """ + + # Validate points and bounds shape match. + def check_shape(array_name): + attr_name = f"core_{array_name}" + arrays = [getattr(coord, attr_name)() for coord in coords] + if any(a is None for a in arrays): + message = f"{array_name} missing from coords[{arrays.index(None)}] ." + raise ValueError(message) + shapes = [array.shape for array in arrays] + if shapes.count(shapes[0]) != len(shapes): + message = f"{array_name} shapes are not identical for all coords." + raise ValueError(message) + + for array in ("points", "bounds"): + check_shape(array) + + # Determine dimensionality, using first coord. + first_coord = coords[0] + + ndim = first_coord.ndim + if ndim != 1: + message = f"Expected coordinate ndim == 1, got: f{ndim} ." + raise ValueError(message) + + bounds_shape = first_coord.core_bounds().shape + bounds_dim1 = bounds_shape[1] + if bounds_dim1 < 2: + message = ( + f"Expected coordinate bounds.shape (n, >=2), got: {bounds_shape} ." + ) + raise ValueError(message) + elif bounds_dim1 == 2: + topology_dimension = 1 + coord_centring = "edge" + conn_cf_role = "edge_node_connectivity" + else: + topology_dimension = 2 + coord_centring = "face" + conn_cf_role = "face_node_connectivity" + + # Create connectivity. + if first_coord.has_lazy_bounds(): + array_lib = da + else: + array_lib = np + indices = array_lib.arange(np.prod(bounds_shape)).reshape(bounds_shape) + masking = array_lib.ma.getmaskarray(first_coord.core_bounds()) + indices = array_lib.ma.masked_array(indices, masking) + connectivity = Connectivity(indices, conn_cf_role) + + # Create coords. + node_coords = [] + centre_coords = [] + for coord in coords: + coord_kwargs = dict( + standard_name=coord.standard_name, + long_name=coord.long_name, + units=coord.units, + attributes=coord.attributes, + ) + node_points = array_lib.ma.filled(coord.core_bounds(), 0.0).flatten() + node_coords.append(AuxCoord(points=node_points, **coord_kwargs)) + + centre_points = coord.core_points() + centre_coords.append(AuxCoord(points=centre_points, **coord_kwargs)) + + ##### + # TODO: remove axis assignment once Mesh supports arbitrary coords. + # TODO: consider filtering coords as the first action in this method. + axes_present = [guess_coord_axis(coord) for coord in coords] + axes_required = ("X", "Y") + if all([req in axes_present for req in axes_required]): + axis_indices = [axes_present.index(req) for req in axes_required] + else: + message = ( + "Unable to find 'X' and 'Y' using guess_coord_axis. Assuming " + "X=coords[0], Y=coords[1] ." + ) + # TODO: reconsider logging level when we have consistent practice. + logger.info(message, extra=dict(cls=None)) + axis_indices = range(len(axes_required)) + + def axes_assign(coord_list): + coords_sorted = [coord_list[ix] for ix in axis_indices] + return zip(coords_sorted, axes_required) + + node_coords_and_axes = axes_assign(node_coords) + centre_coords_and_axes = axes_assign(centre_coords) + ##### + + # Construct the Mesh. + mesh_kwargs = dict( + topology_dimension=topology_dimension, + node_coords_and_axes=node_coords_and_axes, + connectivities=[connectivity], + ) + mesh_kwargs[f"{coord_centring}_coords_and_axes"] = centre_coords_and_axes + return cls(**mesh_kwargs) + + def _set_dimension_names(self, node, edge, face, reset=False): + args = (node, edge, face) + currents = ( + self.node_dimension, + self.edge_dimension, + self.face_dimension, + ) + zipped = zip(args, currents) + if reset: + node, edge, face = [None if arg else current for arg, current in zipped] + else: + node, edge, face = [arg or current for arg, current in zipped] + + self.node_dimension = node + self.edge_dimension = edge + self.face_dimension = face + + if self.topology_dimension == 1: + result = Mesh1DNames(self.node_dimension, self.edge_dimension) + elif self.topology_dimension == 2: + result = Mesh2DNames( + self.node_dimension, self.edge_dimension, self.face_dimension + ) + else: + message = f"Unsupported topology_dimension: {self.topology_dimension} ." + raise NotImplementedError(message) + + return result - :meth:`to_MeshCoord` for generating a single mesh coord. + @property + def edge_dimension(self): + """The *optionally required* UGRID NetCDF variable name for the ``edge`` dimension.""" + return self._metadata_manager.edge_dimension - Parameters - ---------- - location : str - The ``location`` argument for :class:`MeshCoord` instantiation. + @edge_dimension.setter + def edge_dimension(self, name): + if not name or not isinstance(name, str): + edge_dimension = f"Mesh{self.topology_dimension}d_edge" + else: + edge_dimension = name + self._metadata_manager.edge_dimension = edge_dimension - Returns - ------- - tuple of :class:`~iris.mesh.mesh.MeshCoord` - Tuple of :class:`~iris.mesh.mesh.MeshCoord` - referencing the current :class:`MeshXY`. One for each value in - :attr:`AXES`, using the value for the ``axis`` argument. + @property + def face_dimension(self): + """The *optional* UGRID NetCDF variable name for the ``face`` dimension.""" + return self._metadata_manager.face_dimension - """ - # factory method - result = [self.to_MeshCoord(location=location, axis=ax) for ax in self.AXES] - return tuple(result) + @face_dimension.setter + def face_dimension(self, name): + if self.topology_dimension < 2: + face_dimension = None + if name: + # Tell the user it is not being set if they expected otherwise. + message = ( + "Not setting face_dimension (inappropriate for " + f"topology_dimension={self.topology_dimension} ." + ) + logger.debug(message, extra=dict(cls=self.__class__.__name__)) + elif not name or not isinstance(name, str): + face_dimension = f"Mesh{self.topology_dimension}d_face" + else: + face_dimension = name + self._metadata_manager.face_dimension = face_dimension + + @property + def node_dimension(self): + """The NetCDF variable name for the ``node`` dimension.""" + return self._metadata_manager.node_dimension + + @node_dimension.setter + def node_dimension(self, name): + if not name or not isinstance(name, str): + node_dimension = f"Mesh{self.topology_dimension}d_node" + else: + node_dimension = name + self._metadata_manager.node_dimension = node_dimension def dimension_names_reset(self, node=False, edge=False, face=False): """Reset the name used for the NetCDF variable. @@ -2000,6 +2076,33 @@ def topology_dimension(self): return self._metadata_manager.topology_dimension +class _ManagerMembers(dict): + # TODO: docstrings + read_only_message: str + + def _readonly(self, *args, **kwargs): + raise RuntimeError(self.read_only_message) + + def set_mutability(self, mutable: bool, message: Optional[str] = None) -> None: + if not mutable: + self.read_only_message = message or "Members of this manager are read-only." + for op in ( + dict.__setitem__, + dict.__delitem__, + dict.pop, + dict.popitem, + dict.clear, + dict.update, + dict.setdefault, + ): + op_name = op.__name__ + if mutable: + new_op = getattr(super(), op_name) + else: + new_op = self._readonly + setattr(self, op_name, new_op) + + class _Mesh1DCoordinateManager: """TBD: require clarity on coord_systems validation. @@ -2017,11 +2120,17 @@ class _Mesh1DCoordinateManager: "edge_y", ) - def __init__(self, node_x, node_y, edge_x=None, edge_y=None): + def __init__( + self, node_x, node_y, edge_x=None, edge_y=None, view: Optional[str] = None + ): self.timestamp = _Timestamp() + # view = an error message informing that the coordinates of this manager + # are only a 'view' onto the coordinates of another Mesh. Message should carry + # useful user-level info from the calling context. + self._view_message = view # initialise all the coordinates self.ALL = self.REQUIRED + self.OPTIONAL - self._members_dict = {member: None for member in self.ALL} + self._members_dict = _ManagerMembers({member: None for member in self.ALL}) # required coordinates self.node_x = node_x @@ -2040,6 +2149,8 @@ def __init__(self, node_x, node_y, edge_x=None, edge_y=None): if self.edge_y: self.edge_y._mesh_timestamps.append(self.timestamp) + self._set_immutable() + def __eq__(self, other): # TBD: this is a minimalist implementation and requires to be revisited return id(self) == id(other) @@ -2066,7 +2177,7 @@ def __setstate__(self, state): # Create ".timestamp" if missing, as the "._members" setter requires one. # Needing during unpickling, where __setstate__ replaces object __init__. self.timestamp = _Timestamp() - self._members = state + self._members = _ManagerMembers(state) def __str__(self): args = [f"{member}" for member, coord in self if coord is not None] @@ -2091,6 +2202,12 @@ def _remove(self, **kwargs): return result + def _set_immutable(self): + # Factored out to allow subclasses to set immutability after initialisation. + if self.is_view: + view_message = f"Coordinate modifications forbidden: {self._view_message}" + self._members_dict.set_mutability(False, message=view_message) + def _setter(self, element, axis, coord, shape): self.timestamp.update() axis = axis.lower() @@ -2106,6 +2223,15 @@ def _setter(self, element, axis, coord, shape): emsg = f"{member!r} requires to be an 'AuxCoord', got {type(coord)}." raise TypeError(emsg) + if self.is_view: + has_lazy_bounds = coord.has_lazy_bounds() or not coord.has_bounds() + if not (coord.has_lazy_points() and has_lazy_bounds): + message = ( + f"Non-lazy coordinate detected: {member}, which is " + f"inappropriate for a view: {self._view_message}" + ) + raise ValueError(message) + guess_axis = guess_coord_axis(coord) if guess_axis and guess_axis.lower() != axis: @@ -2145,6 +2271,18 @@ def _node_shape(self): @property def _members(self): + if self.is_view: + # This is the appropriate moment to check for continued laziness. + for member, coord in [ + (m, c) for m, c in self._members_dict.items() if c is not None + ]: + has_lazy_bounds = coord.has_lazy_bounds() or not coord.has_bounds() + if not (coord.has_lazy_points() and has_lazy_bounds): + message = ( + f"Non-lazy coordinate detected: {member}, which is " + f"inappropriate for a view: {self._view_message}" + ) + raise ValueError(message) return self._members_dict @_members.setter @@ -2157,7 +2295,7 @@ def all_members(self): return Mesh1DCoords(**self._members) @property - def edge_coords(self): + def edge_coords(self) -> MeshEdgeCoords: return MeshEdgeCoords(edge_x=self.edge_x, edge_y=self.edge_y) @property @@ -2177,7 +2315,11 @@ def edge_y(self, coord): self._setter(element="edge", axis="y", coord=coord, shape=self._edge_shape) @property - def node_coords(self): + def is_view(self): + return self._view_message is not None + + @property + def node_coords(self) -> MeshNodeCoords: return MeshNodeCoords(node_x=self.node_x, node_y=self.node_y) @property @@ -2323,6 +2465,74 @@ def populated_coords(coords_tuple): result_dict = {k: v for k, v in self._members.items() if id(v) in result_ids} return result_dict + def indexed( + self, + node_indices: ArrayLike, + edge_indices: Optional[ArrayLike], + face_indices: Optional[ArrayLike], + mesh_id: int, + frozen: bool = False, + ) -> "_Mesh1DCoordinateManager": + """Return an indexed copy of this coordinate manager. + + Parameters + ---------- + node_indices, edge_indices, face_indices : ArrayLike + The indices to use when indexing member node, edge or face + coordinates respectively. + mesh_id : int + The ID of the mesh that these indices refer to - used to + produce a meaningful read-only error. + frozen : bool, default=False + If True, the returned coordinate manager will comprise indexed copies of the + original coordinates, which do not change if the original coordinates + change. If False, the returned coordinate manager will comprise coordinates + that are indexed 'views' of the original coordinates, which will change if + the original coordinates change; membership of the returned coordinate + manager will also be read-only. + + Returns + ------- + _Mesh1DCoordinateManager + """ + indices_dict = { + "node": node_indices, + "edge": edge_indices, + "face": face_indices, + } + indexed_members = {} + for key, coord in self: + indexing = None + indexed = None + if coord is not None: + indexing = indices_dict[key.split("_")[0]] + if indexing is not None: + if frozen: + indexed = coord.copy()[indexing] + else: + indexed = coord.copy( + # Lazy = deferred calculation. Changes to the original coordinate + # will be reflected in the indexed coordinate. Will be primarily + # used by MeshCoord, which also maintains laziness. + points=coord.lazy_points()[indexing], + bounds=None + if not coord.has_bounds() + else coord.lazy_bounds()[indexing], + ) + indexed_members[key] = indexed + + kwargs = indexed_members + if not frozen: + mesh_index_set, mesh_xy = [c.__name__ for c in (_MeshIndexSet, MeshXY)] + view_message = ( + f"Coordinates on {mesh_index_set} are only 'views' onto the " + f"coordinates of an original {mesh_xy}: id={mesh_id}." + ) + kwargs["view"] = view_message + result = self.__class__(**kwargs) + + return result + def remove( self, item=None, @@ -2360,10 +2570,12 @@ def __init__( edge_y=None, face_x=None, face_y=None, + view: Optional[str] = None, ): - super().__init__(node_x, node_y, edge_x=edge_x, edge_y=edge_y) + super().__init__(node_x, node_y, edge_x=edge_x, edge_y=edge_y, view=view) # optional coordinates + self._members_dict.set_mutability(True) self.face_x = face_x self.face_y = face_y if self.face_x: @@ -2371,6 +2583,8 @@ def __init__( if self.face_y: self.face_y._mesh_timestamps.append(self.timestamp) + self._set_immutable() + @property def _face_shape(self): return self._shape(element="face") @@ -2380,7 +2594,7 @@ def all_members(self): return Mesh2DCoords(**self._members) @property - def face_coords(self): + def face_coords(self) -> MeshFaceCoords: return MeshFaceCoords(face_x=self.face_x, face_y=self.face_y) @property @@ -2432,13 +2646,66 @@ def remove( ) +def _index_conn_array( + array: np.ndarray | da.Array, + indexing: np.ndarray | da.Array, + pre_index: functools.partial | None = None, +) -> np.ndarray | da.Array: + """Index a connectivity array, allowing for laziness and masking. + + Convenience used in constructing MeshCoords and views of _MeshConnectivityManagerBase. + + Parameters + ---------- + array : np.ndarray or dask.array.Array + The array to be indexed. + indexing : np.ndarray or dask.array.Array + The array of indices to use for indexing `array`. + pre_index : functools.partial, optional + A function to apply to the indices before indexing `array`. + + Returns + ------- + np.ndarray or dask.array.Array + """ + # Assumes indices_by_location + n_nodes = array.shape[0] + # Choose real/lazy array library, to suit array types. + lazy = _lazy.is_lazy_data(array) or _lazy.is_lazy_data(indexing) + al = da if lazy else np + # NOTE: Dask cannot index with a multidimensional array, so we + # must flatten it and restore the shape later. + flat_inds = indexing.flatten() + # NOTE: the connectivity array can have masked points, but we can't + # effectively index with those. So use a non-masked index array + # with "safe" index values, and post-mask the results. + flat_inds_nomask = al.ma.filled(flat_inds, -1) + # Note: *also* mask any places where the index is out of range. + missing_inds = (flat_inds_nomask < 0) | (flat_inds_nomask >= n_nodes) + flat_inds_safe = al.where(missing_inds, 0, flat_inds_nomask) + # Execute pre_index operation, if present. + if pre_index is not None: + flat_inds_safe = pre_index(flat_inds_safe) + # Here's the core indexing operation. + # The comma applies all inds-array values to the *first* dimension. + flat_result = array[flat_inds_safe,] + # Fix 'missing' locations, and restore the proper shape. + flat_result_masked = al.ma.masked_array(flat_result, missing_inds) + result = flat_result_masked.reshape(indexing.shape) + return result + + class _MeshConnectivityManagerBase(ABC): # Override these in subclasses. REQUIRED: tuple = NotImplemented OPTIONAL: tuple = NotImplemented - def __init__(self, *connectivities): + def __init__(self, *connectivities, view: Optional[str] = None): self.timestamp = _Timestamp() + # view = an error message informing that the connectivities of this manager + # are only a 'view' onto the connectivities of another Mesh. Message should + # carry useful user-level info from the calling context. + self._view_message = view cf_roles = [c.cf_role for c in connectivities] for requisite in self.REQUIRED: if requisite not in cf_roles: @@ -2446,9 +2713,13 @@ def __init__(self, *connectivities): raise ValueError(message) self.ALL = self.REQUIRED + self.OPTIONAL - self._members_dict = {member: None for member in self.ALL} + self._members_dict = _ManagerMembers({member: None for member in self.ALL}) self.add(*connectivities) + if self.is_view: + view_message = f"Connectivity modifications forbidden: {self._view_message}" + self._members_dict.set_mutability(False, message=view_message) + def __eq__(self, other): # TBD: this is a minimalist implementation and requires to be revisited return id(self) == id(other) @@ -2479,7 +2750,7 @@ def __setstate__(self, state): # Create ".timestamp" if missing, as the "._members" setter requires one. # Needing during unpickling, where __setstate__ replaces object __init__. self.timestamp = _Timestamp() - self._members = state + self._members = _ManagerMembers(state) def __str__(self): args = [ @@ -2492,8 +2763,23 @@ def __str__(self): def all_members(self): return NotImplemented + @property + def is_view(self): + return self._view_message is not None + @property def _members(self): + if self.is_view: + # This is the appropriate moment to check for continued laziness. + for member, connectivity in [ + (m, c) for m, c in self._members_dict.items() if c is not None + ]: + if not connectivity.has_lazy_indices(): + message = ( + f"Non-lazy connectivity detected: {member}, which is " + f"inappropriate for a view: {self._view_message}" + ) + raise ValueError(message) return self._members_dict @_members.setter @@ -2513,6 +2799,14 @@ def add(self, *connectivities): if not isinstance(connectivity, Connectivity): message = f"Expected Connectivity, got: {type(connectivity)} ." raise TypeError(message) + + if self.is_view and not connectivity.has_lazy_indices(): + message = ( + f"Non-lazy connectivity detected: {connectivity!r}, which is " + f"inappropriate for a view: {self._view_message}" + ) + raise ValueError(message) + cf_role = connectivity.cf_role if cf_role not in self.ALL: message = ( @@ -2539,7 +2833,7 @@ def add(self, *connectivities): ) raise ValueError(message) - self._members = proposed_members + self._members = _ManagerMembers(proposed_members) for c in self._members.values(): if c is not None and self.timestamp not in c._mesh_timestamps: c._mesh_timestamps.append(self.timestamp) @@ -2638,6 +2932,95 @@ def element_filter(instances, loc_arg, loc_name): result_dict = {k: v for k, v in self._members.items() if id(v) in result_ids} return result_dict + def indexed( + self, + node_indices: ArrayLike, + edge_indices: Optional[ArrayLike], + face_indices: Optional[ArrayLike], + mesh_id: int, + frozen: bool = False, + ) -> "_MeshConnectivityManagerBase": + """Return an indexed copy of this connectivity manager. + + Parameters + ---------- + node_indices, edge_indices, face_indices : ArrayLike + The indices to use when indexing member connectivities with node, + edge or face :attr:`~Connectivity.location` respectively. + mesh_id : int + The ID of the mesh that these indices refer to - used to + produce a meaningful read-only error. + frozen : bool, default=False + If True, the returned connectivity manager will comprise indexed copies of the + original connectivities, which do not change if the original connectivities + change. If False, the returned connectivity manager will comprise connectivities + that are indexed 'views' of the original connectivities, which will change if + the original connectivities change; membership of the returned connectivity + manager will also be read-only. + + Returns + ------- + _MeshConnectivityManagerBase + """ + indices_dict = { + "node": node_indices, + "edge": edge_indices, + "face": face_indices, + } + # Prep some indexing information for later node indices remapping. + order = node_indices.argsort() + old_sorted = node_indices[order] + al = da if _lazy.is_lazy_data(node_indices) else np + new_ids_sorted = al.arange(len(node_indices))[order] + positions = functools.partial(al.searchsorted, old_sorted) + + indexed_members = {} + for key, connectivity in self: + indexing = None + indexed = None + if connectivity is not None: + indexing = indices_dict[connectivity.location] + if indexing is not None: + if frozen: + connectivity = connectivity.copy() + indices = connectivity.core_indices() + else: + # Lazy = deferred calculation. Changes to the original connectivity + # will be reflected in the indexed connectivity. Will be primarily + # used by MeshCoord, which also maintains laziness. + indices = connectivity.lazy_indices() + reindexed_indices = connectivity.indices_by_location(indices)[indexing] + + # Map node indices in "values" to their new zero-based positions + # in "node_indices". + reindexed_indices = _index_conn_array( + array=new_ids_sorted, + indexing=reindexed_indices, + pre_index=positions, + ) + + if connectivity.location_axis == 1: + reindexed_indices = reindexed_indices.T + if connectivity.start_index == 1: + reindexed_indices += 1 + indexed = connectivity.copy(reindexed_indices) + indexed_members[key] = indexed + + if not frozen: + mesh_index_set, mesh_xy = [c.__name__ for c in (_MeshIndexSet, MeshXY)] + view_message = ( + f"Connectivities on {mesh_index_set} are only 'views' onto the " + f"connectivities of an original {mesh_xy}: id={mesh_id}." + ) + else: + view_message = None + + result = self.__class__( + *[c for c in indexed_members.values() if c is not None], view=view_message + ) + + return result + def remove( self, item=None, @@ -2743,6 +3126,306 @@ def face_node(self): Location = Literal["edge", "node", "face"] +class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): + # TODO: docstring is out of date with the latest code. + # TODO: is a private class more or less appropriate than placing in the experimental + # module? + """A container representing the UGRID ``cf_role``: ``location_index_set``. + + A container representing the UGRID ``cf_role``: + ``location_index_set``. Achieved by referencing an original :class:`MeshXY` + instance (:attr:`super_mesh`), together with a specific :attr:`location` + (``node``/``edge``/``face``) and a set of :attr:`indices`. This is strictly + a view onto the original :class:`MeshXY` - it does not store its own + coordinates or connectivities. + + Warnings + -------- + This class is not yet stable, hence being marked as private. + + References + ---------- + 1. The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ + """ + + # TODO: implement I/O (iris#6123). + # TODO: validation? + # TODO: finish type hinting + # TODO: update the full documentation + # TODO: docstrings + # TODO: informative error when attempting to save (until iris#6123 is implemented). + def __init__( + self, + indices: ArrayLike, + mesh: MeshXY, + location: Literal["node", "edge", "face"], + standard_name: Optional[str] = None, + long_name: Optional[str] = None, + var_name: Optional[str] = None, + units: Optional[str] = None, + attributes: Optional[dict] = None, + start_index: Literal[0, 1] = 0, + ): + self._metadata_manager = metadata_manager_factory(MeshIndexSetMetadata) + # 'structure' is immutable after creation, so assign directly to the + # metadata manager. Desired changes should be made by creating a new + # instance. + self._metadata_manager.mesh = mesh + self._metadata_manager.location = location + self._metadata_manager.start_index = start_index + + _DimensionalMetadata.__init__( + self, + values=indices, + # TODO: for the reviewer: with _MeshIndexSet as its own part of the + # data model, including NetCDF I/O, is it ever appropriate for the + # standard metadata to be inherited from the original MeshXY, or + # calling MeshCoord? + standard_name=standard_name, + long_name=long_name, + var_name=var_name, + units=units or mesh.units, + attributes=attributes, + ) + + # TODO: PyLance reportIncompatibleMethodOverride. Consider make this an abstract + # method then overriding in both subclasses. + def __getstate__(self) -> tuple[ArrayLike, MeshIndexSetMetadata]: + return ( + self.indices, + self._metadata_manager, + ) + + def __setstate__(self, state: tuple[ArrayLike, MeshIndexSetMetadata]): + indices, metadata_manager = state + self._values = indices + self._metadata_manager = metadata_manager + + @property + def cf_role(self) -> str: + return "location_index_set" + + @property + def edge_dimension(self) -> str: + # TODO: standardise string as a constant + return "_MeshIndexSet_NotImplemented" + + @property + def face_dimension(self) -> str: + return "_MeshIndexSet_NotImplemented" + + @property + def node_dimension(self) -> str: + return "_MeshIndexSet_NotImplemented" + + @property + def indices(self) -> ArrayLike: + return self._values + + @property + def mesh(self) -> MeshXY: + return self._metadata_manager.mesh + + @property + def location(self) -> Literal["node", "edge", "face"]: + return self._metadata_manager.location + + @property + def start_index(self) -> Literal[0, 1]: + return self._metadata_manager.start_index + + @property + def topology_dimension(self) -> int: + return self.mesh.topology_dimension + + def _calculate_node_indices(self): + # Use self.location and self.indices to work out the indices to use + # when indexing the nodes of self.mesh. + # TODO: use match-case instead. + if self.location == "node": + result = self.indices + elif self.location in ["edge", "face"]: + (connectivity,) = [ + c + for c in self.mesh.all_connectivities + if ( + c is not None + and c.location == self.location + and c.connected == "node" + ) + ] + # Doesn't matter if connectivity is transposed or not in this case. + # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet + conn_indices = connectivity.core_indices()[self.indices] + node_set = np.unique(conn_indices) + if iris.util.is_masked(node_set): + if _lazy.is_lazy_data(node_set): + # A dask-compatible approach that is compatible with chunking. + # (compressed() is not implemented because of chunking concerns). + node_set_nans = da.where( + da.ma.getmaskarray(node_set), da.ma.getdata(node_set), -1 + ) + node_set_unmasked = node_set_nans[node_set_nans != -1] + else: + node_set_unmasked = node_set.compressed() + else: + node_set_unmasked = node_set + result = node_set_unmasked + else: + result = None + # TODO: should this be validated earlier? + # Maybe even with an Enum? + message = ( + f"Expected location to be one of `node`, `edge` or `face`, " + f"got `{self.location}`" + ) + raise NotImplementedError(message) + + return result + + def _calculate_edge_indices(self): + # Use self.location and self.indices to work out the indices to use + # when indexing the edges of self.mesh. + if self.location == "edge": + result = self.indices + else: + result = None + return result + + def _calculate_face_indices(self): + # Use self.location and self.indices to work out the indices to use + # when indexing the faces of self.mesh. + if self.location == "face": + result = self.indices + else: + result = None + return result + + @property + def _coord_manager(self): + mesh_man: _Mesh1DCoordinateManager = self.mesh._coord_manager + self_man: Optional[_Mesh1DCoordinateManager] = getattr( + self, "_coord_manager_attr", None + ) + update = self_man is None or mesh_man.timestamp._dt != self_man.timestamp._dt + if update: + self._coord_manager_attr = mesh_man.indexed( + self._calculate_node_indices(), + self._calculate_edge_indices(), + self._calculate_face_indices(), + mesh_id=id(self.mesh), + ) + self._coord_manager_attr.timestamp._dt = mesh_man.timestamp._dt + return super()._coord_manager + + @_coord_manager.setter + def _coord_manager(self, manager): + message = ( + f"Modification of {self.__class__.__name__} is forbidden - this is only a " + f"view onto an original MeshXY: id={id(self.mesh)}." + ) + raise NotImplementedError(message) + + @property + def _connectivity_manager(self): + mesh_man: _MeshConnectivityManagerBase = self.mesh._connectivity_manager + self_man: Optional[_MeshConnectivityManagerBase] = getattr( + self, "_connectivity_manager_attr", None + ) + update = self_man is None or mesh_man.timestamp._dt != self_man.timestamp._dt + if update: + self._connectivity_manager_attr = mesh_man.indexed( + self._calculate_node_indices(), + self._calculate_edge_indices(), + self._calculate_face_indices(), + mesh_id=id(self.mesh), + ) + self._connectivity_manager_attr.timestamp._dt = mesh_man.timestamp._dt + return super()._connectivity_manager + + @_connectivity_manager.setter + def _connectivity_manager(self, manager): + message = ( + f"Modification of {self.__class__.__name__} is forbidden - this is only a " + f"view onto an original MeshXY: id={id(self.mesh)}." + ) + raise NotImplementedError(message) + + def _summary_multiline(self) -> str: + # Produce a readable multi-line summary of the MeshIndexSet content. + lines = [] + n_indent = 4 + indent_str = " " * n_indent + + def line(text, i_indent=0): + indent = indent_str * i_indent + lines.append(f"{indent}{text}") + + line(f"{self.__class__.__name__} : '{self.name()}'") + line(f"mesh: {self.mesh.summary(shorten=True)}", 1) + line(f"location: {self.location}", 1) + line(f"start_index: {self.start_index}", 1) + + # Phrasing: because "mesh summary" would imply this is summarising a Mesh object. + line("mesh info summary:", 1) + mesh_summary = super()._summary_multiline() + # TODO: a test for this is important since summary stability is not guaranteed. + for mesh_line in mesh_summary.splitlines()[1:]: + line(mesh_line, 1) + + result = "\n".join(lines) + return result + + def cube_dims(self, cube): + raise NotImplementedError() + + def as_mesh(self) -> MeshXY: + """Return a :class:`MeshXY` representation of this instance. + + The returned instance will no longer be a view onto another + :class:`MeshXY`. All connectivities and coordinates will be deep-copied. + + Returns + ------- + MeshXY + """ + indices = [ + self._calculate_node_indices(), + self._calculate_edge_indices(), + self._calculate_face_indices(), + ] + kwargs = dict(mesh_id=id(self.mesh), frozen=True) + coord_man = self.mesh._coord_manager.indexed(*indices, **kwargs) + conn_man = self.mesh._connectivity_manager.indexed(*indices, **kwargs) + + def _coords_and_axes( + location: Literal["node", "edge", "face"], + ) -> list[tuple[AuxCoord, str]]: + coords = getattr(coord_man, f"{location}_coords") + return [(getattr(coords, f"{location}_{axis}"), axis) for axis in self.AXES] + + return MeshXY( + topology_dimension=self.topology_dimension, + node_coords_and_axes=_coords_and_axes("node"), + connectivities=[conn for conn in conn_man.all_members if conn is not None], + edge_coords_and_axes=_coords_and_axes("edge"), + face_coords_and_axes=_coords_and_axes("face"), + standard_name=self.standard_name, + long_name=self.long_name, + var_name=self.var_name, + units=self.units, + attributes=self.attributes, + ) + + +# TODO: this can only work if the manager classes are defined _before_ the downstream +# use in MeshXY. Have avoided re-ordering so far, to avoid polluting the Git diff. +_CoordinateManagerType: TypeAlias = _Mesh1DCoordinateManager | _Mesh2DCoordinateManager +_ConnectivityManagerType: TypeAlias = ( + _Mesh1DConnectivityManager | _Mesh2DConnectivityManager +) + + class MeshCoord(AuxCoord): """Geographic coordinate values of data on an unstructured mesh. @@ -2793,29 +3476,29 @@ def __init__( self._metadata_manager_temp = metadata_manager_factory(MeshCoordMetadata) # Validate and record the class-specific constructor args. - if not isinstance(mesh, MeshXY): + if not isinstance(mesh, _MeshXYMixin): msg = ( # type: ignore[unreachable] - f"'mesh' must be an {MeshXY.__module__}.{MeshXY.__name__}, got {mesh}." + f"'mesh' must be an {_MeshXYMixin.__module__}.{_MeshXYMixin.__name__}, got {mesh}." ) raise TypeError(msg) # Handled as a readonly ".mesh" property. # NOTE: currently *not* included in metadata. In future it might be. self._mesh = mesh - if location not in MeshXY.ELEMENTS: + if location not in _MeshXYMixin.ELEMENTS: msg = ( f"'location' of {location} is not a valid MeshXY location', " - f"must be one of {MeshXY.ELEMENTS}." + f"must be one of {_MeshXYMixin.ELEMENTS}." ) raise ValueError(msg) # Held in metadata, readable as self.location, but cannot set it. self._metadata_manager_temp.location = location - if axis not in MeshXY.AXES: + if axis not in _MeshXYMixin.AXES: # The valid axes are defined by the MeshXY class. msg = ( f"'axis' of {axis} is not a valid MeshXY axis', " - f"must be one of {MeshXY.AXES}." + f"must be one of {_MeshXYMixin.AXES}." ) raise ValueError(msg) # Held in metadata, readable as self.axis, but cannot set it. @@ -3011,17 +3694,74 @@ def _metadata_manager(self): return self._metadata_manager_temp def __getitem__(self, keys): - # Disallow any sub-indexing, permitting *only* "self[:,]". - # We *don't* intend here to support indexing as such : the exception is - # just sufficient to enable cube slicing, when it does not affect the - # mesh dimension. This works because Cube.__getitem__ passes us keys - # "normalised" with iris.util._build_full_slice_given_keys. - if keys != (slice(None),): - msg = "Cannot index a MeshCoord." - raise ValueError(msg) + # TODO: validate keys as 1-dimensional? + # TODO: handle problems caused by asking for 1 index + # E.g. coord[0:1] works fine, coord[0] causes errors difficult to + # understand. + from ..experimental.mesh_coord_indexing import SETTING, Options + + def get_index_set(): + # Assign to help with typing. + mesh = self.mesh + match mesh: + case MeshXY(): + (length,) = self.shape + kwargs = dict( + mesh=mesh, + location=self.location, + indices=np.arange(length)[keys], + ) + case _MeshIndexSet(): + # Should not base an index set on another index set - base + # on the original mesh instead. + # TODO: do we need to double-check that self.location + # matches mesh_index_set.location? Any other matching to + # check too? + mesh_index_set = mesh + kwargs = dict( + mesh=mesh_index_set.mesh, + location=mesh_index_set.location, + # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet + indices=mesh_index_set.indices[keys], + ) + case _: + message = ( + "Unsupported mesh type for MeshCoord indexing. Expected " + f"{MeshXY.__name__} or {_MeshIndexSet.__name__}, got: " + f"{type(self.mesh).__name__}." + ) + raise NotImplementedError(message) - # Translate "self[:,]" as "self.copy()". - return self.copy() + return _MeshIndexSet(**kwargs) + + match SETTING.value: + case Options.AUX_COORD: + result = AuxCoord.from_coord(self)[keys] + + case Options.NEW_MESH: + new_mesh = get_index_set().as_mesh() + result = self.__class__( + mesh=new_mesh, + location=self.location, + axis=self.axis, + ) + + case Options.MESH_INDEX_SET: + index_set = get_index_set() + result = self.__class__( + mesh=index_set, + location=self.location, + axis=self.axis, + ) + + case _: + message = ( + "Unsupported mesh_coord_indexing setting. Expected one of: " + f"{Options.__members__}, got: {SETTING.value}." + ) + raise NotImplementedError(message) + + return result def update_from_mesh(self): """Fetch and recalculate the points and bounds from the relevant coord on the mesh. @@ -3319,25 +4059,6 @@ def _construct_access_arrays(self): indices = indices - bounds_connectivity.start_index node_points = node_coord.core_points() - n_nodes = node_points.shape[0] - # Choose real/lazy array library, to suit array types. - lazy = _lazy.is_lazy_data(indices) or _lazy.is_lazy_data(node_points) - al = da if lazy else np - # NOTE: Dask cannot index with a multidimensional array, so we - # must flatten it and restore the shape later. - flat_inds = indices.flatten() - # NOTE: the connectivity array can have masked points, but we can't - # effectively index with those. So use a non-masked index array - # with "safe" index values, and post-mask the results. - flat_inds_nomask = al.ma.filled(flat_inds, -1) - # Note: *also* mask any places where the index is out of range. - missing_inds = (flat_inds_nomask < 0) | (flat_inds_nomask >= n_nodes) - flat_inds_safe = al.where(missing_inds, 0, flat_inds_nomask) - # Here's the core indexing operation. - # The comma applies all inds-array values to the *first* dimension. - bounds = node_points[flat_inds_safe,] - # Fix 'missing' locations, and restore the proper shape. - bounds = al.ma.masked_array(bounds, missing_inds) - bounds = bounds.reshape(indices.shape) + bounds = _index_conn_array(node_points, indices) return points, bounds diff --git a/lib/iris/pandas.py b/lib/iris/pandas.py index dcdcbd3bb9..4119db9b9f 100644 --- a/lib/iris/pandas.py +++ b/lib/iris/pandas.py @@ -30,6 +30,7 @@ from iris._deprecation import explicit_copy_checker, warn_deprecated from iris.coords import AncillaryVariable, AuxCoord, CellMeasure, DimCoord from iris.cube import Cube, CubeList +from iris.exceptions import MonotonicityError from iris.util import monotonic, new_axis from iris.warnings import IrisIgnoringWarning @@ -433,7 +434,7 @@ def as_cubes( "Pandas index is not monotonic. Consider using the " "sort_index() method before passing in." ) - raise ValueError(message) + raise MonotonicityError(message) cube_shape = getattr(pandas_index, "levshape", (pandas_index.nunique(),)) n_rows = len(pandas_structure) diff --git a/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py b/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py index 650226eb78..c22b423c58 100644 --- a/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py +++ b/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py @@ -50,6 +50,7 @@ def _base_testcube(self, include_derived=False): return self.cube +# TODO: failing due to _MeshIndexSet work @_shared_utils.skip_data class TestBroadcastingWithMesh( MeshLocationsMixin, @@ -64,6 +65,7 @@ class TestBroadcastingWithMesh( """ +# TODO: failing due to _MeshIndexSet work @_shared_utils.skip_data class TestBroadcastingWithMeshAndDerived( MeshLocationsMixin, diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index d91c7e81c0..b0766d77af 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -3083,6 +3083,39 @@ def test(self): assert (aux / cube).units == "m.s-1" +class Test__getitem_DimCoord: + @pytest.fixture(autouse=True) + 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) + x_coord_aux = DimCoord(points=np.array([2, 3, 4]), long_name="x_aux") + cube.add_aux_coord(x_coord_aux, 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]) + self.cube = cube + + def test_dim_coords_2d(self): + result = self.cube[0:2, 0:2] + assert len(result.dim_coords) == 2 + for ix, size in enumerate(result.shape): + assert result.coord(dim_coords=True, dimensions=ix).shape == (size,) + assert result.coord("x_aux").shape == (result.shape[0],) + + def test_dim_coord_1d(self): + result = self.cube[0, 0:2] + assert len(result.dim_coords) == 1 + assert result.dim_coords[0].shape == result.shape + assert result.coord("x_aux").shape == (result.shape[0],) + + def test_demote(self): + result = self.cube[:, [0, 2, 1]] + assert len(result.dim_coords) == 1 + assert result.coords(dimensions=1, dim_coords=True) == [] + + class Test__getitem_CellMeasure: @pytest.fixture(autouse=True) def _setup(self): diff --git a/lib/iris/tests/unit/mesh/components/test_MeshCoord.py b/lib/iris/tests/unit/mesh/components/test_MeshCoord.py index 8890207a52..4891c3a937 100644 --- a/lib/iris/tests/unit/mesh/components/test_MeshCoord.py +++ b/lib/iris/tests/unit/mesh/components/test_MeshCoord.py @@ -229,6 +229,7 @@ def test_fail_copy_newbounds(self): meshcoord.copy(bounds=meshcoord.bounds) +# TODO: all failing due to _MeshIndexSet work class Test__getitem__: def test_slice_wholeslice_1tuple(self): # The only slicing case that we support, to enable cube slicing. @@ -496,6 +497,7 @@ def test_cube_copy(self): assert meshco2 == meshcoord def test_cube_nonmesh_slice(self): + # TODO: failing due to _MeshIndexSet work # Check that we can slice a cube on a non-mesh dimension, and get a # meshcoord == original. # Note: currently this must have the *same* mesh, as for .copy(). diff --git a/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py b/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py index 519f823a87..46b75f33d3 100644 --- a/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py +++ b/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py @@ -274,6 +274,7 @@ def test_fail_no_regions(self): with pytest.raises(ValueError, match="'submesh_cubes' must be non-empty"): recombine_submeshes(self.mesh_cube, []) + # TODO: failing due to _MeshIndexSet work def test_fail_dims_mismatch_mesh_regions(self): self.mesh_cube = self.mesh_cube[0] with pytest.raises( From 901adaaa946cedff45426e6b325c50306fb982c8 Mon Sep 17 00:00:00 2001 From: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:25:35 +0100 Subject: [PATCH 2/6] Refactor `_MeshIndexSet` for true lazy support (#7209) * Refactor _MeshIndexSet for true lazy support. * Update lib/iris/mesh/components.py * Resolve MyPy failures. * Review comments. * Better naming. * Correct mathematical terminology. --- lib/iris/mesh/components.py | 155 ++++++++++++++++++++++-------------- 1 file changed, 94 insertions(+), 61 deletions(-) diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 7b71c96147..9528e7c508 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -18,9 +18,8 @@ from collections import namedtuple from collections.abc import Container from contextlib import contextmanager -from copy import deepcopy from datetime import datetime -import functools +import math from typing import Any, Iterable, Literal, Mapping, NamedTuple, Optional, TypeAlias import warnings @@ -2087,7 +2086,7 @@ def topology_dimension(self): return self._metadata_manager.topology_dimension -class _ManagerMembers(dict): +class _ManagerMembers[VT](dict[str, VT]): # TODO: docstrings read_only_message: str @@ -2141,7 +2140,9 @@ def __init__( self._view_message = view # initialise all the coordinates self.ALL = self.REQUIRED + self.OPTIONAL - self._members_dict = _ManagerMembers({member: None for member in self.ALL}) + self._members_dict: _ManagerMembers[AuxCoord | None] = _ManagerMembers( + {member: None for member in self.ALL} + ) # required coordinates self.node_x = node_x @@ -2254,10 +2255,18 @@ def _setter(self, element, axis, coord, shape): raise TypeError(emsg) if shape is not None and coord.shape != shape: - emsg = ( - f"{member!r} requires to have shape {shape!r}, got {coord.shape!r}." + # NaN dimensions arise from lazy boolean-mask indexing, + # where the length is not known until computed. + compatible = len(coord.shape) == len(shape) and all( + math.isnan(x) or math.isnan(y) or x == y + for x, y in zip(coord.shape, shape) ) - raise ValueError(emsg) + if not compatible: + emsg = ( + f"{member!r} requires to have shape {shape!r}, " + f"got {coord.shape!r}." + ) + raise ValueError(emsg) # if the coordinate attached to the mesh hasn't been linked if self.timestamp not in coord._mesh_timestamps: coord._mesh_timestamps.append(self.timestamp) @@ -2281,7 +2290,7 @@ def _node_shape(self): return self._shape(element="node") @property - def _members(self) -> dict[str, None] | dict[str, AuxCoord]: + def _members(self) -> _ManagerMembers[AuxCoord | None]: if self.is_view: # This is the appropriate moment to check for continued laziness. for member, coord in [ @@ -2297,7 +2306,7 @@ def _members(self) -> dict[str, None] | dict[str, AuxCoord]: return self._members_dict @_members.setter - def _members(self, value: dict[str, AuxCoord]): + def _members(self, value: _ManagerMembers[AuxCoord | None]): self.timestamp.update() self._members_dict = value @@ -2494,7 +2503,7 @@ def filters( def indexed( self, - node_indices: ArrayLike, + node_bool_index: ArrayLike, edge_indices: Optional[ArrayLike], face_indices: Optional[ArrayLike], mesh_id: int, @@ -2504,8 +2513,11 @@ def indexed( Parameters ---------- - node_indices, edge_indices, face_indices : ArrayLike - The indices to use when indexing member node, edge or face + node_bool_index : ArrayLike + Array of boolean membership, over the full length of original nodes, + to use when indexing member node coordinates. + edge_indices, face_indices : ArrayLike, optional + The indices to use when indexing member edge or face coordinates respectively. mesh_id : int The ID of the mesh that these indices refer to - used to @@ -2523,7 +2535,7 @@ def indexed( _Mesh1DCoordinateManager """ indices_dict = { - "node": node_indices, + "node": node_bool_index, "edge": edge_indices, "face": face_indices, } @@ -2703,7 +2715,6 @@ def get_face(axis: str | None) -> list[AuxCoord]: def _index_conn_array( array: np.ndarray | da.Array, indexing: np.ndarray | da.Array, - pre_index: functools.partial | None = None, ) -> np.ndarray | da.Array: """Index a connectivity array, allowing for laziness and masking. @@ -2715,8 +2726,6 @@ def _index_conn_array( The array to be indexed. indexing : np.ndarray or dask.array.Array The array of indices to use for indexing `array`. - pre_index : functools.partial, optional - A function to apply to the indices before indexing `array`. Returns ------- @@ -2737,9 +2746,6 @@ def _index_conn_array( # Note: *also* mask any places where the index is out of range. missing_inds = (flat_inds_nomask < 0) | (flat_inds_nomask >= n_nodes) flat_inds_safe = al.where(missing_inds, 0, flat_inds_nomask) - # Execute pre_index operation, if present. - if pre_index is not None: - flat_inds_safe = pre_index(flat_inds_safe) # Here's the core indexing operation. # The comma applies all inds-array values to the *first* dimension. flat_result = array[flat_inds_safe,] @@ -2767,7 +2773,9 @@ def __init__(self, *connectivities, view: Optional[str] = None): raise ValueError(message) self.ALL = self.REQUIRED + self.OPTIONAL - self._members_dict = _ManagerMembers({member: None for member in self.ALL}) + self._members_dict: _ManagerMembers[Connectivity | None] = _ManagerMembers( + {member: None for member in self.ALL} + ) self.add(*connectivities) if self.is_view: @@ -2822,7 +2830,7 @@ def is_view(self): return self._view_message is not None @property - def _members(self): + def _members(self) -> _ManagerMembers[Connectivity | None]: if self.is_view: # This is the appropriate moment to check for continued laziness. for member, connectivity in [ @@ -2837,7 +2845,7 @@ def _members(self): return self._members_dict @_members.setter - def _members(self, value): + def _members(self, value: _ManagerMembers[Connectivity | None]): self.timestamp.update() self._members_dict = value @@ -2988,7 +2996,7 @@ def element_filter(instances, loc_arg, loc_name): def indexed( self, - node_indices: ArrayLike, + node_bool_index: ArrayLike, edge_indices: Optional[ArrayLike], face_indices: Optional[ArrayLike], mesh_id: int, @@ -2998,9 +3006,12 @@ def indexed( Parameters ---------- - node_indices, edge_indices, face_indices : ArrayLike - The indices to use when indexing member connectivities with node, - edge or face :attr:`~Connectivity.location` respectively. + node_bool_index : ArrayLike + Array of boolean membership, over the full length of original nodes, + to support the construction of indexed connectivities. + edge_indices, face_indices : ArrayLike, optional + The indices to use when indexing member connectivities with edge or + face :attr:`~Connectivity.location` respectively. mesh_id : int The ID of the mesh that these indices refer to - used to produce a meaningful read-only error. @@ -3017,16 +3028,20 @@ def indexed( _MeshConnectivityManagerBase """ indices_dict = { - "node": node_indices, + "node": node_bool_index, "edge": edge_indices, "face": face_indices, } - # Prep some indexing information for later node indices remapping. - order = node_indices.argsort() - old_sorted = node_indices[order] - al = da if _lazy.is_lazy_data(node_indices) else np - new_ids_sorted = al.arange(len(node_indices))[order] - positions = functools.partial(al.searchsorted, old_sorted) + + # Prep an inverse-lookup table (original node id -> new node id) for later + # node indices remapping. ``node_bool_index`` is a fixed-shape boolean + # membership array over the original nodes (see _calculate_node_bool_index). + node_mask = node_bool_index + al = da if _lazy.is_lazy_data(node_mask) else np + node_lookup_data = al.cumsum(node_mask.astype(np.intp)) - 1 + # Mask unselected nodes - for graceful handling of unexpected scenarios; + # we normally expect downstream indexing to remove these anyway. + node_lookup = al.ma.masked_array(node_lookup_data, ~node_mask) indexed_members = {} for key, connectivity in self: @@ -3043,21 +3058,21 @@ def indexed( # will be reflected in the indexed connectivity. Will be primarily # used by MeshCoord, which also maintains laziness. indices = connectivity.lazy_indices() - reindexed_indices = connectivity.indices_by_location(indices)[indexing] + indices_indexed = connectivity.indices_by_location(indices)[indexing] # Map node indices in "values" to their new zero-based positions - # in "node_indices". - reindexed_indices = _index_conn_array( - array=new_ids_sorted, - indexing=reindexed_indices, - pre_index=positions, + # via the inverse-lookup table. This creates a contiguous + # collection of nodes, and collapses duplicates. + indices_remapped = _index_conn_array( + array=node_lookup, + indexing=indices_indexed, ) if connectivity.location_axis == 1: - reindexed_indices = reindexed_indices.T + indices_remapped = indices_remapped.T if connectivity.start_index == 1: - reindexed_indices += 1 - indexed = connectivity.copy(reindexed_indices) + indices_remapped += 1 + indexed = connectivity.copy(indices_remapped) indexed_members[key] = indexed if not frozen: @@ -3292,12 +3307,36 @@ def start_index(self) -> Literal[0, 1]: def topology_dimension(self) -> int: return self.mesh.topology_dimension - def _calculate_node_indices(self): + def _calculate_node_bool_index(self): # Use self.location and self.indices to work out the indices to use # when indexing the nodes of self.mesh. + # Returns a boolean membership array of fixed shape (n_original_nodes,), + # True where a node is selected. # TODO: use match-case instead. + n_original_nodes = self.mesh.node_coords.node_x.shape[0] if self.location == "node": - result = self.indices + # self.indices is a user-supplied index array over the nodes; convert + # it to a fixed-shape boolean membership mask. + monotonic, direction = iris.util.monotonic( + self.indices, strict=True, return_direction=True + ) + if not (monotonic and direction == 1): + # TODO: boolean 'mask' array precludes non-monotonic indexing, + # but is only needed to support connectivity construction, and + # only causes problems for coordinate construction. Separate + # logic to allow array of integer indices for coordinate + # construction. + message = ( + "Indexing the nodes on a Mesh currently requires strictly " + "increasing indices. Contact the Iris developers if this " + "causes you problems." + ) + raise ValueError(message) + indices = self.indices + al = da if _lazy.is_lazy_data(indices) else np + node_mask = al.zeros(n_original_nodes, dtype=bool) + node_mask[indices] = True + result = node_mask elif self.location in ["edge", "face"]: (connectivity,) = [ c @@ -3311,20 +3350,14 @@ def _calculate_node_indices(self): # Doesn't matter if connectivity is transposed or not in this case. # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet conn_indices = connectivity.core_indices()[self.indices] - node_set = np.unique(conn_indices) - if iris.util.is_masked(node_set): - if _lazy.is_lazy_data(node_set): - # A dask-compatible approach that is compatible with chunking. - # (compressed() is not implemented because of chunking concerns). - node_set_nans = da.where( - da.ma.getmaskarray(node_set), da.ma.getdata(node_set), -1 - ) - node_set_unmasked = node_set_nans[node_set_nans != -1] - else: - node_set_unmasked = node_set.compressed() - else: - node_set_unmasked = node_set - result = node_set_unmasked + al = da if _lazy.is_lazy_data(conn_indices) else np + # Flatten and drop masked padding (ragged connectivities) by scattering + # membership into a fixed-shape boolean mask. + flat = al.ma.filled(conn_indices, -1).flatten() + valid = flat >= 0 + node_mask = al.zeros(n_original_nodes, dtype=bool) + node_mask[flat[valid]] = True + result = node_mask else: result = None # TODO: should this be validated earlier? @@ -3364,7 +3397,7 @@ def _coord_manager(self): update = self_man is None or mesh_man.timestamp._dt != self_man.timestamp._dt if update: self._coord_manager_attr = mesh_man.indexed( - self._calculate_node_indices(), + self._calculate_node_bool_index(), self._calculate_edge_indices(), self._calculate_face_indices(), mesh_id=id(self.mesh), @@ -3389,7 +3422,7 @@ def _connectivity_manager(self): update = self_man is None or mesh_man.timestamp._dt != self_man.timestamp._dt if update: self._connectivity_manager_attr = mesh_man.indexed( - self._calculate_node_indices(), + self._calculate_node_bool_index(), self._calculate_edge_indices(), self._calculate_face_indices(), mesh_id=id(self.mesh), @@ -3444,7 +3477,7 @@ def as_mesh(self) -> MeshXY: MeshXY """ indices = [ - self._calculate_node_indices(), + self._calculate_node_bool_index(), self._calculate_edge_indices(), self._calculate_face_indices(), ] From 8596bbee767f3586120b3b8a6d8a4c509db891a8 Mon Sep 17 00:00:00 2001 From: pt331 <144435193+pt331@users.noreply.github.com> Date: Thu, 23 Jul 2026 10:48:10 +0100 Subject: [PATCH 3/6] (Partially) Finalise meshindexset (#7207) * remove unused import * made _MeshXYMixin get and set state abstract * bonus comments * Add _MeshCoordinateManager as parent ot 1d and 2d * typing with _MeshConnectivityManagerBase * fix typing generator * accounting for dask array in _MeshConnectivityManagerBase indexed (to be discussed) * hacky way to fix test * Sequence check for MeshCoord * _MeshCoordinateManagerBase saves and gets _view_message with state * Raise exception when MeshXY.from_coords is called with MeshCoords * validation to _MeshIndexSet __init__ * Typing _MeshIndexSet. Other areas touched for mypy * Raise exception when saving _MeshIndexSet. Could possibly be done earlier * Added _MeshIndexSet._NOT_IMPLEMENTED * using match-case for _MeshIndexSet._calculate_node_indices and removing the null case because it is handled in the __init__ * Prevent mismatch of location in MeshCoord.__init__ when supplied a _MeshIndexSet * Further typing of _MeshXYMixin * post merge numpy fixes * Added forgotten bonus comment * fixed logic mistake in summary * updating remove_duplicate_nodes in tests to deal with `None`s * Update _MeshIndexSet init validation to use isinstance * remove finished todo * remove completed todos * remove unused import * update _calculate_node_bool_index to use match case * fixed ambiguous truthy and reformatted indexed * updated _MeshIndexSet saving is not yet supported error message * Remove completed todos * fixed from_coords type check and updated error message * improverd _MeshIndexSet.__init__ error messages * removed completed todo * Updated MeshCoord.__getitem__ error message * refactor _MeshCoordinateManagerBase.indexed * Cope with scalar indexing, plus view_message fix. * Updated expected exception messages * Made _MeshIndexSet.__init__ indices more lenient --------- Co-authored-by: Martin Yeo --- lib/iris/common/metadata.py | 3 + lib/iris/coords.py | 4 +- lib/iris/cube.py | 10 +- lib/iris/fileformats/netcdf/saver.py | 5 + lib/iris/mesh/components.py | 423 ++++++++++-------- lib/iris/tests/stock/__init__.py | 3 +- .../analysis/maths/test__arith__meshcoords.py | 2 - .../unit/mesh/components/test_MeshCoord.py | 12 +- .../mesh/utils/test_recombine_submeshes.py | 1 - 9 files changed, 253 insertions(+), 210 deletions(-) diff --git a/lib/iris/common/metadata.py b/lib/iris/common/metadata.py index 3523eddfc7..fe68430db2 100644 --- a/lib/iris/common/metadata.py +++ b/lib/iris/common/metadata.py @@ -119,6 +119,7 @@ class _NamedTupleMeta(ABCMeta): def __new__(mcs, name, bases, namespace): names = [] + # "Inherit" the base classes' `_fields` for base in bases: if hasattr(base, "_fields"): base_names = getattr(base, "_fields") @@ -130,6 +131,8 @@ def __new__(mcs, name, bases, namespace): base_names = (base_names,) names.extend(base_names) + # _members are fields that are specific in this NamedTuple. + # Add these as fields. if "_members" in namespace and not getattr( namespace["_members"], "__isabstractmethod__", False ): diff --git a/lib/iris/coords.py b/lib/iris/coords.py index f5880b654b..68e1ef7aa2 100644 --- a/lib/iris/coords.py +++ b/lib/iris/coords.py @@ -17,6 +17,7 @@ from functools import lru_cache from itertools import zip_longest import operator +from typing import cast import warnings import zlib @@ -330,7 +331,7 @@ def summary( precision=None, convert_dates=True, _section_indices=None, - ): + ) -> str: r"""Make a printable text summary. Parameters @@ -620,6 +621,7 @@ def reindent_data_string(text, n_indent): show = val is not None and val is not False if show: if name == "attributes": + val = cast(dict, val) # Use a multi-line form for this. add_output(newline_indent) add_output("attributes:", section="attributes") diff --git a/lib/iris/cube.py b/lib/iris/cube.py index f23305cf8e..24d0230918 100644 --- a/lib/iris/cube.py +++ b/lib/iris/cube.py @@ -1638,7 +1638,13 @@ def is_mesh_coord(anycoord: iris.coords.Coord) -> TypeGuard[MeshCoord]: ownval=location, ) ) - mesh_dims = (self.mesh_dim(),) + mesh_dim = self.mesh_dim() + mesh_dims: tuple[int] | tuple[()] + if mesh_dim is None: + # Scalar coordinate. + mesh_dims = () + else: + mesh_dims = (mesh_dim,) if data_dims != mesh_dims: raise iris.exceptions.CannotAddError( msg.format( @@ -2605,7 +2611,7 @@ def mesh_dim(self) -> int | None: if coord is None: result = None else: - (result,) = self.coord_dims(coord) # result is a 1-tuple + (result,) = self.coord_dims(coord) or (None,) return result def cell_measures( diff --git a/lib/iris/fileformats/netcdf/saver.py b/lib/iris/fileformats/netcdf/saver.py index fc3f33032c..eb4644e3dc 100644 --- a/lib/iris/fileformats/netcdf/saver.py +++ b/lib/iris/fileformats/netcdf/saver.py @@ -814,12 +814,17 @@ def _add_mesh(self, cube_or_mesh, /, *, compression_kwargs=None): # Do cube- or -mesh-based save from iris.cube import Cube + from iris.mesh.components import _MeshIndexSet if isinstance(cube_or_mesh, Cube): mesh = cube_or_mesh.mesh else: mesh = cube_or_mesh + if isinstance(mesh, _MeshIndexSet): + msg = "_MeshIndexSet saving is not yet supported. See https://github.com/SciTools/iris/issues/6123" + raise ValueError(msg) + if mesh: cf_mesh_name = self._name_coord_map.name(mesh) if cf_mesh_name is None: diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 9528e7c508..69455093f5 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -20,7 +20,18 @@ from contextlib import contextmanager from datetime import datetime import math -from typing import Any, Iterable, Literal, Mapping, NamedTuple, Optional, TypeAlias +from types import NotImplementedType +from typing import ( + Any, + Generator, + Iterable, + Literal, + Mapping, + NamedTuple, + Optional, + Sequence, + TypedDict, +) import warnings from cf_units import Unit @@ -40,7 +51,7 @@ from ..common import CFVariableMixin, metadata_filter, metadata_manager_factory from ..common.metadata import BaseMetadata from ..config import get_logger -from ..coords import AuxCoord, _DimensionalMetadata +from ..coords import AuxCoord, Coord, _DimensionalMetadata from ..exceptions import ConnectivityNotFoundError, CoordinateNotFoundError from ..util import array_equal, clip_string, guess_coord_axis from ..warnings import IrisVagueMetadataWarning @@ -661,10 +672,8 @@ class _MeshXYMixin(Mesh, ABC): # Subclass __init__ methods must define: # TODO: Impossible to type hint the return type of metadata_manager_factory(). _metadata_manager: Any - # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. - _connectivity_manager_attr: Any - # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. - _coord_manager_attr: Any + _connectivity_manager_attr: "_MeshConnectivityManagerBase" + _coord_manager_attr: "_MeshCoordinateManagerBase" # TBD: for volume and/or z-axis support include axis "z" and/or dimension "3" #: The supported mesh axes. @@ -674,8 +683,8 @@ class _MeshXYMixin(Mesh, ABC): #: Valid mesh elements. ELEMENTS = ("edge", "node", "face") - def __eq__(self, other): - result = NotImplemented + def __eq__(self, other) -> bool | NotImplementedType: + result: bool | NotImplementedType = NotImplemented if isinstance(other, _MeshXYMixin): result = self.metadata == other.metadata @@ -686,25 +695,18 @@ def __eq__(self, other): return result - def __hash__(self): + def __hash__(self) -> int: # Allow use in sets and as dictionary keys, as is done for :class:`iris.cube.Cube`. # See https://github.com/SciTools/iris/pull/1772 return hash(id(self)) - def __getstate__(self): - return ( - self._metadata_manager, - self._coord_manager, - self._connectivity_manager, - ) - - def __ne__(self, other): + def __ne__(self, other) -> bool | NotImplementedType: result = self.__eq__(other) if result is not NotImplemented: result = not result return result - def summary(self, *args, **kwargs): + def summary(self, *args, **kwargs) -> str: """Return a string representation of the MeshXY. Parameters @@ -735,12 +737,12 @@ def __repr__(self): def __str__(self): return self.summary(shorten=False) - def _summary_oneline(self): + def _summary_oneline(self) -> str: # We use the repr output to produce short one-line identity summary, # similar to the object.__str__ output "". # This form also used in other str() constructions, like MeshCoord. # By contrast, __str__ (below) produces a readable multi-line printout. - mesh_name = self.name() + mesh_name: str | None = self.name() if mesh_name in (None, "", "unknown"): mesh_name = None if mesh_name: @@ -753,7 +755,7 @@ def _summary_oneline(self): return mesh_string - def _summary_multiline(self): + def _summary_multiline(self) -> str: # Produce a readable multi-line summary of the Mesh content. lines = [] n_indent = 4 @@ -767,29 +769,35 @@ def line(text, i_indent=0): line(f"topology_dimension: {self.topology_dimension}", 1) for element in ("node", "edge", "face"): if element == "node": - element_exists = True + main_conn_string = None else: main_conn_name = f"{element}_node_connectivity" - main_conn = getattr(self, main_conn_name, None) - element_exists = main_conn is not None - if element_exists: - # Include a section for this element - line(element, 1) - # Print element dimension - dim_name = f"{element}_dimension" - dim = getattr(self, dim_name) - line(f"{dim_name}: '{dim}'", 2) - # Print defining connectivity (except node) - if element != "node": - main_conn_string = main_conn.summary(shorten=True, linewidth=0) - line(f"{main_conn_name}: {main_conn_string}", 2) - # Print coords - coords = self.coords(location=element) - if coords: - line(f"{element} coordinates", 2) - for coord in coords: + main_conn: Connectivity | None = getattr(self, main_conn_name, None) + if main_conn is None: + continue + + main_conn_string = main_conn.summary(shorten=True, linewidth=0) + main_conn_string = f"{main_conn_name}: {main_conn_string}" + + # Include a section for this element + line(element, 1) + # Print element dimension + dim_name = f"{element}_dimension" + dim = getattr(self, dim_name) + line(f"{dim_name}: '{dim}'", 2) + # Print defining connectivity (except node) + if main_conn_string: + line(main_conn_string, 2) + # Print coords + coords = self.coords(location=element) + if coords: + line(f"{element} coordinates", 2) + for coord in coords: + if coord: coord_string = coord.summary(shorten=True, linewidth=0) - line(coord_string, 3) + else: + coord_string = "None" + line(coord_string, 3) # Having dealt with essential info, now add any optional connectivities # N.B. includes boundaries: as optional connectivity, not an "element" @@ -799,10 +807,10 @@ def line(text, i_indent=0): "face_edge_connectivity", "edge_face_connectivity", ) - optional_conns = [getattr(self, name, None) for name in optional_conn_names] + optional_conn_keys = [getattr(self, name, None) for name in optional_conn_names] optional_conns = { name: conn - for conn, name in zip(optional_conns, optional_conn_names) + for conn, name in zip(optional_conn_keys, optional_conn_names) if conn is not None } if optional_conns: @@ -841,55 +849,53 @@ def line(text, i_indent=0): result = "\n".join(lines) return result + @abstractmethod + def __getstate__(self): + raise NotImplementedError + + @abstractmethod def __setstate__(self, state): - metadata_manager, coord_manager, connectivity_manager = state - self._metadata_manager = metadata_manager - self._coord_manager = coord_manager - self._connectivity_manager = connectivity_manager + raise NotImplementedError - # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. @property - def _connectivity_manager(self): + def _connectivity_manager(self) -> "_MeshConnectivityManagerBase": # @property enables interruption/customisation in subclasses. return self._connectivity_manager_attr - # TODO: type hint with _ConnectivityManagerType once the file is appropriately re-ordered. @_connectivity_manager.setter - def _connectivity_manager(self, manager) -> None: + def _connectivity_manager(self, manager: "_MeshConnectivityManagerBase") -> None: # @property enables interruption/customisation in subclasses. self._connectivity_manager_attr = manager - # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. @property - def _coord_manager(self): + def _coord_manager(self) -> "_MeshCoordinateManagerBase": # @property enables interruption/customisation in subclasses. return self._coord_manager_attr - # TODO: type hint with _CoordinateManagerType once the file is appropriately re-ordered. @_coord_manager.setter - def _coord_manager(self, manager) -> None: + def _coord_manager(self, manager: "_MeshCoordinateManagerBase") -> None: # @property enables interruption/customisation in subclasses. self._coord_manager_attr = manager @property - def all_connectivities(self): + def all_connectivities(self) -> NamedTuple: """All the :class:`~iris.mesh.Connectivity` instances of the :class:`MeshXY`.""" return self._connectivity_manager.all_members @property - def _last_modified(self): + def _last_modified(self) -> datetime: """The time and date that the mesh coordinates and or connecitivities were last edited.""" return max( self._coord_manager.timestamp._dt, self._connectivity_manager.timestamp._dt ) @property - def all_coords(self): + def all_coords(self) -> NamedTuple: """All the :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" return self._coord_manager.all_members @property - def boundary_node_connectivity(self): + def boundary_node_connectivity(self) -> Connectivity: """The *optional* UGRID ``boundary_node_connectivity`` :class:`~iris.mesh.Connectivity`. The *optional* UGRID ``boundary_node_connectivity`` @@ -897,10 +903,10 @@ def boundary_node_connectivity(self): :class:`MeshXY`. """ - return self._connectivity_manager.boundary_node + return self._connectivity_manager.boundary_node # type:ignore[attr-defined] @property - def edge_coords(self): + def edge_coords(self) -> MeshEdgeCoords: """The *optional* UGRID ``edge`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" return self._coord_manager.edge_coords @@ -910,7 +916,7 @@ def edge_dimension(self) -> str: raise NotImplementedError() @property - def edge_face_connectivity(self): + def edge_face_connectivity(self) -> Connectivity: """The *optional* UGRID ``edge_face_connectivity`` :class:`~iris.mesh.Connectivity`. The *optional* UGRID ``edge_face_connectivity`` @@ -918,10 +924,10 @@ def edge_face_connectivity(self): :class:`MeshXY`. """ - return self._connectivity_manager.edge_face + return self._connectivity_manager.edge_face # type:ignore[attr-defined] @property - def edge_node_connectivity(self): + def edge_node_connectivity(self) -> Connectivity: """The UGRID ``edge_node_connectivity`` :class:`~iris.mesh.Connectivity`. The UGRID ``edge_node_connectivity`` @@ -931,12 +937,12 @@ def edge_node_connectivity(self): :attr:`MeshXY.topology_dimension` ``>=2``. """ - return self._connectivity_manager.edge_node + return self._connectivity_manager.edge_node # type:ignore[attr-defined] @property - def face_coords(self): + def face_coords(self) -> AuxCoord: """The *optional* UGRID ``face`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" - return self._coord_manager.face_coords + return self._coord_manager.face_coords # type:ignore[attr-defined] @property @abstractmethod @@ -944,7 +950,7 @@ def face_dimension(self) -> str: raise NotImplementedError() @property - def face_edge_connectivity(self): + def face_edge_connectivity(self) -> Connectivity: """The *optional* UGRID ``face_edge_connectivity``:class:`~iris.mesh.Connectivity`. The *optional* UGRID ``face_edge_connectivity`` @@ -953,10 +959,10 @@ def face_edge_connectivity(self): """ # optional - return self._connectivity_manager.face_edge + return self._connectivity_manager.face_edge # type:ignore[attr-defined] @property - def face_face_connectivity(self): + def face_face_connectivity(self) -> Connectivity: """The *optional* UGRID ``face_face_connectivity`` :class:`~iris.mesh.Connectivity`. The *optional* UGRID ``face_face_connectivity`` @@ -964,10 +970,10 @@ def face_face_connectivity(self): :class:`MeshXY`. """ - return self._connectivity_manager.face_face + return self._connectivity_manager.face_face # type:ignore[attr-defined] @property - def face_node_connectivity(self): + def face_node_connectivity(self) -> Connectivity: """Return ``face_node_connectivity``:class:`~iris.mesh.Connectivity`. The UGRID ``face_node_connectivity`` @@ -977,10 +983,10 @@ def face_node_connectivity(self): of ``3``. """ - return self._connectivity_manager.face_node + return self._connectivity_manager.face_node # type:ignore[attr-defined] @property - def node_coords(self): + def node_coords(self) -> MeshNodeCoords: """The **required** UGRID ``node`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" return self._coord_manager.node_coords @@ -1744,10 +1750,21 @@ def normalise(element, axis): emsg = f"Unsupported 'topology_dimension', got {topology_dimension!r}." raise NotImplementedError(emsg) - # TODO: backwards compatibility: make from_coords() perform a no-op if the given - # coords are already MeshCoords. + def __getstate__(self) -> tuple[Any, Any, Any]: + return ( + self._metadata_manager, + self._coord_manager, + self._connectivity_manager, + ) + + def __setstate__(self, state: tuple[Any, Any, Any]): + metadata_manager, coord_manager, connectivity_manager = state + self._metadata_manager = metadata_manager + self._coord_manager = coord_manager + self._connectivity_manager = connectivity_manager + @classmethod - def from_coords(cls, *coords): + def from_coords(cls, *coords: Coord): r"""Construct a :class:`MeshXY` by derivation from 1/more :class:`~iris.coords.Coord`. The :attr:`~MeshXY.topology_dimension`, :class:`~iris.coords.Coord` @@ -1847,6 +1864,9 @@ def from_coords(cls, *coords): longitude: MeshCoord """ + if any(isinstance(coord, MeshCoord) for coord in coords): + msg = "Expected coords to be DimCoord or AuxCoord, got: MeshCoord." + raise ValueError(msg) # Validate points and bounds shape match. def check_shape(array_name): @@ -1915,12 +1935,13 @@ def check_shape(array_name): ##### # TODO: remove axis assignment once Mesh supports arbitrary coords. - # TODO: consider filtering coords as the first action in this method. axes_present = [guess_coord_axis(coord) for coord in coords] - axes_required = ("X", "Y") - if all([req in axes_present for req in axes_required]): + # Explicit for mypy + axes_required: tuple[Literal["X"], Literal["Y"]] = ("X", "Y") + axis_indices: list[int] | range + try: axis_indices = [axes_present.index(req) for req in axes_required] - else: + except ValueError: message = ( "Unable to find 'X' and 'Y' using guess_coord_axis. Assuming " "X=coords[0], Y=coords[1] ." @@ -2113,7 +2134,7 @@ def set_mutability(self, mutable: bool, message: Optional[str] = None) -> None: setattr(self, op_name, new_op) -class _Mesh1DCoordinateManager: +class _MeshCoordinateManagerBase(ABC): """TBD: require clarity on coord_systems validation. TBD: require clarity on __eq__ support @@ -2131,7 +2152,12 @@ class _Mesh1DCoordinateManager: ) def __init__( - self, node_x, node_y, edge_x=None, edge_y=None, view: Optional[str] = None + self, + node_x: AuxCoord, + node_y: AuxCoord, + edge_x: AuxCoord | None = None, + edge_y: AuxCoord | None = None, + view: Optional[str] = None, ): self.timestamp = _Timestamp() # view = an error message informing that the coordinates of this manager @@ -2167,10 +2193,10 @@ def __eq__(self, other): # TBD: this is a minimalist implementation and requires to be revisited return id(self) == id(other) - def __getstate__(self): - return self._members + def __getstate__(self) -> tuple[_ManagerMembers, str | None]: + return (self._members, self._view_message) - def __iter__(self): + def __iter__(self) -> Generator[tuple[str, AuxCoord | None], None, None]: for item in self._members.items(): yield item @@ -2184,12 +2210,14 @@ def __repr__(self): args = [f"{member}={coord!r}" for member, coord in self if coord is not None] return f"{self.__class__.__name__}({', '.join(args)})" - def __setstate__(self, state): + def __setstate__(self, state: tuple[_ManagerMembers, str | None]): if not hasattr(self, "timestamp"): # Create ".timestamp" if missing, as the "._members" setter requires one. # Needing during unpickling, where __setstate__ replaces object __init__. self.timestamp = _Timestamp() - self._members = _ManagerMembers(state) + members, view_message = state + self._members = _ManagerMembers(members) + self._view_message = view_message def __str__(self): args = [f"{member}" for member, coord in self if coord is not None] @@ -2508,7 +2536,7 @@ def indexed( face_indices: Optional[ArrayLike], mesh_id: int, frozen: bool = False, - ) -> "_Mesh1DCoordinateManager": + ) -> "_MeshCoordinateManagerBase": """Return an indexed copy of this coordinate manager. Parameters @@ -2539,36 +2567,34 @@ def indexed( "edge": edge_indices, "face": face_indices, } - indexed_members = {} + indexed_members: dict[str, AuxCoord | None] = {} for key, coord in self: - indexing = None indexed = None - if coord is not None: - indexing = indices_dict[key.split("_")[0]] - if indexing is not None: + if ( + coord is not None + and (indexing := indices_dict[key.split("_")[0]]) is not None + ): if frozen: indexed = coord.copy()[indexing] else: - indexed = coord.copy( - # Lazy = deferred calculation. Changes to the original coordinate - # will be reflected in the indexed coordinate. Will be primarily - # used by MeshCoord, which also maintains laziness. - points=coord.lazy_points()[indexing], - bounds=None - if not coord.has_bounds() - else coord.lazy_bounds()[indexing], - ) + lazy_points = coord.lazy_points() + lazy_bounds = coord.lazy_bounds() + points = lazy_points[indexing] if lazy_points is not None else None + bounds = lazy_bounds[indexing] if lazy_bounds is not None else None + # Lazy = deferred calculation. Changes to the original coordinate + # will be reflected in the indexed coordinate. Will be primarily + # used by MeshCoord, which also maintains laziness. + indexed = coord.copy(points, bounds) indexed_members[key] = indexed - kwargs = indexed_members + view_message: str | None = None if not frozen: mesh_index_set, mesh_xy = [c.__name__ for c in (_MeshIndexSet, MeshXY)] view_message = ( f"Coordinates on {mesh_index_set} are only 'views' onto the " f"coordinates of an original {mesh_xy}: id={mesh_id}." ) - kwargs["view"] = view_message - result = self.__class__(**kwargs) + result = self.__class__(**indexed_members, view=view_message) # type: ignore[arg-type] return result @@ -2593,7 +2619,12 @@ def remove( ) -class _Mesh2DCoordinateManager(_Mesh1DCoordinateManager): +class _Mesh1DCoordinateManager(_MeshCoordinateManagerBase): + # The concrete definition of this is in _MeshCoordinateManager + pass + + +class _Mesh2DCoordinateManager(_MeshCoordinateManagerBase): OPTIONAL = ( "edge_x", "edge_y", @@ -2759,6 +2790,7 @@ class _MeshConnectivityManagerBase(ABC): # Override these in subclasses. REQUIRED: tuple = NotImplemented OPTIONAL: tuple = NotImplemented + _view_message: str | None = None def __init__(self, *connectivities, view: Optional[str] = None): self.timestamp = _Timestamp() @@ -2822,8 +2854,8 @@ def __str__(self): @property @abstractmethod - def all_members(self): - return NotImplemented + def all_members(self) -> NamedTuple: + raise NotImplementedError @property def is_view(self): @@ -3197,8 +3229,6 @@ def face_node(self): class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): # TODO: docstring is out of date with the latest code. - # TODO: is a private class more or less appropriate than placing in the experimental - # module? """A container representing the UGRID ``cf_role``: ``location_index_set``. A container representing the UGRID ``cf_role``: @@ -3217,12 +3247,11 @@ class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): 1. The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ """ + _NOT_IMPLEMENTED = "_MeshIndexSet_NotImplemented" + # TODO: implement I/O (iris#6123). - # TODO: validation? - # TODO: finish type hinting # TODO: update the full documentation # TODO: docstrings - # TODO: informative error when attempting to save (until iris#6123 is implemented). def __init__( self, indices: ArrayLike, @@ -3235,6 +3264,16 @@ def __init__( attributes: Optional[dict] = None, start_index: Literal[0, 1] = 0, ): + if not isinstance(mesh, MeshXY): + raise TypeError(f"`mesh` must be `MeshXY`. Got {type(mesh)}") + + if location not in _MeshXYMixin.ELEMENTS: + msg = f"`location` must be in {_MeshXYMixin.ELEMENTS}. Got {location}" + raise ValueError(msg) + + if start_index not in [0, 1]: + raise ValueError(f"`start_index must be 0 or 1. Got {start_index}") + self._metadata_manager = metadata_manager_factory(MeshIndexSetMetadata) # 'structure' is immutable after creation, so assign directly to the # metadata manager. Desired changes should be made by creating a new @@ -3257,8 +3296,6 @@ def __init__( attributes=attributes, ) - # TODO: PyLance reportIncompatibleMethodOverride. Consider make this an abstract - # method then overriding in both subclasses. def __getstate__(self) -> tuple[ArrayLike, MeshIndexSetMetadata]: return ( self.indices, @@ -3276,16 +3313,15 @@ def cf_role(self) -> str: @property def edge_dimension(self) -> str: - # TODO: standardise string as a constant - return "_MeshIndexSet_NotImplemented" + return self._NOT_IMPLEMENTED @property def face_dimension(self) -> str: - return "_MeshIndexSet_NotImplemented" + return self._NOT_IMPLEMENTED @property def node_dimension(self) -> str: - return "_MeshIndexSet_NotImplemented" + return self._NOT_IMPLEMENTED @property def indices(self) -> ArrayLike: @@ -3312,62 +3348,52 @@ def _calculate_node_bool_index(self): # when indexing the nodes of self.mesh. # Returns a boolean membership array of fixed shape (n_original_nodes,), # True where a node is selected. - # TODO: use match-case instead. n_original_nodes = self.mesh.node_coords.node_x.shape[0] - if self.location == "node": - # self.indices is a user-supplied index array over the nodes; convert - # it to a fixed-shape boolean membership mask. - monotonic, direction = iris.util.monotonic( - self.indices, strict=True, return_direction=True - ) - if not (monotonic and direction == 1): - # TODO: boolean 'mask' array precludes non-monotonic indexing, - # but is only needed to support connectivity construction, and - # only causes problems for coordinate construction. Separate - # logic to allow array of integer indices for coordinate - # construction. - message = ( - "Indexing the nodes on a Mesh currently requires strictly " - "increasing indices. Contact the Iris developers if this " - "causes you problems." - ) - raise ValueError(message) - indices = self.indices - al = da if _lazy.is_lazy_data(indices) else np - node_mask = al.zeros(n_original_nodes, dtype=bool) - node_mask[indices] = True - result = node_mask - elif self.location in ["edge", "face"]: - (connectivity,) = [ - c - for c in self.mesh.all_connectivities - if ( - c is not None - and c.location == self.location - and c.connected == "node" + match self.location: + case "node": + # self.indices is a user-supplied index array over the nodes; convert + # it to a fixed-shape boolean membership mask. + monotonic, direction = iris.util.monotonic( + self.indices, strict=True, return_direction=True ) - ] - # Doesn't matter if connectivity is transposed or not in this case. - # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet - conn_indices = connectivity.core_indices()[self.indices] - al = da if _lazy.is_lazy_data(conn_indices) else np - # Flatten and drop masked padding (ragged connectivities) by scattering - # membership into a fixed-shape boolean mask. - flat = al.ma.filled(conn_indices, -1).flatten() - valid = flat >= 0 - node_mask = al.zeros(n_original_nodes, dtype=bool) - node_mask[flat[valid]] = True - result = node_mask - else: - result = None - # TODO: should this be validated earlier? - # Maybe even with an Enum? - message = ( - f"Expected location to be one of `node`, `edge` or `face`, " - f"got `{self.location}`" - ) - raise NotImplementedError(message) - + if not (monotonic and direction == 1): + # TODO: boolean 'mask' array precludes non-monotonic indexing, + # but is only needed to support connectivity construction, and + # only causes problems for coordinate construction. Separate + # logic to allow array of integer indices for coordinate + # construction. + message = ( + "Indexing the nodes on a Mesh currently requires strictly " + "increasing indices. Contact the Iris developers if this " + "causes you problems." + ) + raise ValueError(message) + indices = self.indices + al = da if _lazy.is_lazy_data(indices) else np + node_mask = al.zeros(n_original_nodes, dtype=bool) + node_mask[indices] = True + result = node_mask + case "edge" | "face": + (connectivity,) = [ + c + for c in self.mesh.all_connectivities + if ( + c is not None + and c.location == self.location + and c.connected == "node" + ) + ] + # Doesn't matter if connectivity is transposed or not in this case. + # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet + conn_indices = connectivity.core_indices()[self.indices] + al = da if _lazy.is_lazy_data(conn_indices) else np + # Flatten and drop masked padding (ragged connectivities) by scattering + # membership into a fixed-shape boolean mask. + flat = al.ma.filled(conn_indices, -1).flatten() + valid = flat >= 0 + node_mask = al.zeros(n_original_nodes, dtype=bool) + node_mask[flat[valid]] = True + result = node_mask return result def _calculate_edge_indices(self): @@ -3389,9 +3415,9 @@ def _calculate_face_indices(self): return result @property - def _coord_manager(self): - mesh_man: _Mesh1DCoordinateManager = self.mesh._coord_manager - self_man: Optional[_Mesh1DCoordinateManager] = getattr( + def _coord_manager(self) -> _MeshCoordinateManagerBase: + mesh_man: _MeshCoordinateManagerBase = self.mesh._coord_manager + self_man: Optional[_MeshCoordinateManagerBase] = getattr( self, "_coord_manager_attr", None ) update = self_man is None or mesh_man.timestamp._dt != self_man.timestamp._dt @@ -3414,7 +3440,7 @@ def _coord_manager(self, manager): raise NotImplementedError(message) @property - def _connectivity_manager(self): + def _connectivity_manager(self) -> _MeshConnectivityManagerBase: mesh_man: _MeshConnectivityManagerBase = self.mesh._connectivity_manager self_man: Optional[_MeshConnectivityManagerBase] = getattr( self, "_connectivity_manager_attr", None @@ -3481,7 +3507,8 @@ def as_mesh(self) -> MeshXY: self._calculate_edge_indices(), self._calculate_face_indices(), ] - kwargs = dict(mesh_id=id(self.mesh), frozen=True) + Kwargs = TypedDict("Kwargs", {"mesh_id": int, "frozen": bool}) + kwargs: Kwargs = {"mesh_id": id(self.mesh), "frozen": True} coord_man = self.mesh._coord_manager.indexed(*indices, **kwargs) conn_man = self.mesh._connectivity_manager.indexed(*indices, **kwargs) @@ -3505,14 +3532,6 @@ def _coords_and_axes( ) -# TODO: this can only work if the manager classes are defined _before_ the downstream -# use in MeshXY. Have avoided re-ordering so far, to avoid polluting the Git diff. -_CoordinateManagerType: TypeAlias = _Mesh1DCoordinateManager | _Mesh2DCoordinateManager -_ConnectivityManagerType: TypeAlias = ( - _Mesh1DConnectivityManager | _Mesh2DConnectivityManager -) - - class MeshCoord(AuxCoord): """Geographic coordinate values of data on an unstructured mesh. @@ -3551,9 +3570,9 @@ class MeshCoord(AuxCoord): def __init__( self, - mesh, - location, - axis, + mesh: _MeshXYMixin, + location: Literal["edge", "node", "face"], + axis: Literal["x", "y"], ): self._last_modified = None self._updating = True @@ -3578,6 +3597,9 @@ def __init__( f"must be one of {_MeshXYMixin.ELEMENTS}." ) raise ValueError(msg) + elif isinstance(mesh, _MeshIndexSet) and location != mesh.location: + msg = f"'location' of {location} does not match the location of the mesh location: {mesh.location}" + raise ValueError(msg) # Held in metadata, readable as self.location, but cannot set it. self._metadata_manager_temp.location = location @@ -3781,10 +3803,16 @@ def _metadata_manager(self): return self._metadata_manager_temp def __getitem__(self, keys): - # TODO: validate keys as 1-dimensional? - # TODO: handle problems caused by asking for 1 index - # E.g. coord[0:1] works fine, coord[0] causes errors difficult to - # understand. + if not isinstance(keys, Sequence): + msg = f"MeshCoord indexing expected a Slice or 1-dimensional index array. Got: {type(keys)}" + raise ValueError(msg) + elif len(keys) != 1: + raise ValueError( + f"MeshCoord indexing expected a 1-dimensional index array. Got: {keys}" + ) + elif type(keys[0]) == slice and keys == (slice(None),): + return self.copy() + from ..experimental.mesh_coord_indexing import SETTING, Options def get_index_set(): @@ -3796,20 +3824,17 @@ def get_index_set(): kwargs = dict( mesh=mesh, location=self.location, - indices=np.arange(length)[keys], + indices=np.asanyarray(np.arange(length)[keys]), ) case _MeshIndexSet(): # Should not base an index set on another index set - base # on the original mesh instead. - # TODO: do we need to double-check that self.location - # matches mesh_index_set.location? Any other matching to - # check too? mesh_index_set = mesh kwargs = dict( mesh=mesh_index_set.mesh, location=mesh_index_set.location, # TODO: implement lazy_indices() and core_indices() for _MeshIndexSet - indices=mesh_index_set.indices[keys], + indices=np.asanyarray(mesh_index_set.indices[keys]), ) case _: message = ( diff --git a/lib/iris/tests/stock/__init__.py b/lib/iris/tests/stock/__init__.py index f83d168468..1a71a29849 100644 --- a/lib/iris/tests/stock/__init__.py +++ b/lib/iris/tests/stock/__init__.py @@ -855,7 +855,7 @@ def remove_duplicate_nodes(mesh: ugrid.MeshXY): # Processed nodes: [a, b, c, d] # Processed faces: [[0, 1, 2, 3], [0, 2, 1, 3]] - nodes = np.stack([c.points for c in mesh.node_coords]) + nodes = np.stack([c.points for c in mesh.node_coords if c is not None]) face_node = mesh.face_node_connectivity # first_instances = a full length array but always with the index of @@ -876,6 +876,7 @@ def remove_duplicate_nodes(mesh: ugrid.MeshXY): node_x, node_y = [ AuxCoord(nodes_unique[i], **c.metadata._asdict()) for i, c in enumerate(mesh.node_coords) + if c is not None ] mesh.add_coords(node_x=node_x, node_y=node_y) conn_kwargs = dict(indices=indices_unique, start_index=0) diff --git a/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py b/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py index c22b423c58..650226eb78 100644 --- a/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py +++ b/lib/iris/tests/unit/analysis/maths/test__arith__meshcoords.py @@ -50,7 +50,6 @@ def _base_testcube(self, include_derived=False): return self.cube -# TODO: failing due to _MeshIndexSet work @_shared_utils.skip_data class TestBroadcastingWithMesh( MeshLocationsMixin, @@ -65,7 +64,6 @@ class TestBroadcastingWithMesh( """ -# TODO: failing due to _MeshIndexSet work @_shared_utils.skip_data class TestBroadcastingWithMeshAndDerived( MeshLocationsMixin, diff --git a/lib/iris/tests/unit/mesh/components/test_MeshCoord.py b/lib/iris/tests/unit/mesh/components/test_MeshCoord.py index 4891c3a937..4d1cad6c3f 100644 --- a/lib/iris/tests/unit/mesh/components/test_MeshCoord.py +++ b/lib/iris/tests/unit/mesh/components/test_MeshCoord.py @@ -229,7 +229,6 @@ def test_fail_copy_newbounds(self): meshcoord.copy(bounds=meshcoord.bounds) -# TODO: all failing due to _MeshIndexSet work class Test__getitem__: def test_slice_wholeslice_1tuple(self): # The only slicing case that we support, to enable cube slicing. @@ -243,12 +242,18 @@ def test_slice_wholeslice_1tuple(self): def test_slice_whole_slice_singlekey(self): # A slice(None) also fails, if not presented in a 1-tuple. meshcoord = sample_meshcoord() - with pytest.raises(ValueError, match="Cannot index"): + with pytest.raises( + ValueError, + match="MeshCoord indexing expected a Slice or 1-dimensional index array. Got: ", + ): meshcoord[:] def test_fail_slice_part(self): meshcoord = sample_meshcoord() - with pytest.raises(ValueError, match="Cannot index"): + with pytest.raises( + ValueError, + match="MeshCoord indexing expected a Slice or 1-dimensional index array. Got: ", + ): meshcoord[:1] @@ -497,7 +502,6 @@ def test_cube_copy(self): assert meshco2 == meshcoord def test_cube_nonmesh_slice(self): - # TODO: failing due to _MeshIndexSet work # Check that we can slice a cube on a non-mesh dimension, and get a # meshcoord == original. # Note: currently this must have the *same* mesh, as for .copy(). diff --git a/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py b/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py index 46b75f33d3..519f823a87 100644 --- a/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py +++ b/lib/iris/tests/unit/mesh/utils/test_recombine_submeshes.py @@ -274,7 +274,6 @@ def test_fail_no_regions(self): with pytest.raises(ValueError, match="'submesh_cubes' must be non-empty"): recombine_submeshes(self.mesh_cube, []) - # TODO: failing due to _MeshIndexSet work def test_fail_dims_mismatch_mesh_regions(self): self.mesh_cube = self.mesh_cube[0] with pytest.raises( From 0bcfea6d28044ee0920f41257497afbb0d4355f1 Mon Sep 17 00:00:00 2001 From: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> Date: Thu, 23 Jul 2026 12:35:27 +0100 Subject: [PATCH 4/6] Docstrings for MeshIndexSet work (#7213) * Docstrings for MeshIndexSet work. * More explication. * Indicate default `start_index`. Co-authored-by: pt331 <144435193+pt331@users.noreply.github.com> --------- Co-authored-by: pt331 <144435193+pt331@users.noreply.github.com> --- lib/iris/mesh/components.py | 67 +++++++++++++++++++++++++++++-------- 1 file changed, 53 insertions(+), 14 deletions(-) diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 69455093f5..0532771a5b 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -669,6 +669,15 @@ def __init__(self): class _MeshXYMixin(Mesh, ABC): + """Mixin providing XY-coordinate behaviour for :class:`Mesh` subclasses. + + Implements the shared logic for meshes that carry explicit X and Y + coordinates (nodes, and optionally edges/faces). Concrete subclasses must + populate :attr:`_metadata_manager`, :attr:`_connectivity_manager_attr` and + :attr:`_coord_manager_attr` in their ``__init__``. + + """ + # Subclass __init__ methods must define: # TODO: Impossible to type hint the return type of metadata_manager_factory(). _metadata_manager: Any @@ -2108,7 +2117,13 @@ def topology_dimension(self): class _ManagerMembers[VT](dict[str, VT]): - # TODO: docstrings + """A mutable/immutable dict backing a coordinate or connectivity manager. + + Mutation can be toggled via :meth:`set_mutability`; when immutable every + write operation raises :class:`RuntimeError`. + + """ + read_only_message: str def _readonly(self, *args, **kwargs): @@ -2541,10 +2556,10 @@ def indexed( Parameters ---------- - node_bool_index : ArrayLike + node_bool_index : :obj:`~numpy.typing.ArrayLike` Array of boolean membership, over the full length of original nodes, to use when indexing member node coordinates. - edge_indices, face_indices : ArrayLike, optional + edge_indices, face_indices : :obj:`~numpy.typing.ArrayLike`, optional The indices to use when indexing member edge or face coordinates respectively. mesh_id : int @@ -3038,10 +3053,10 @@ def indexed( Parameters ---------- - node_bool_index : ArrayLike + node_bool_index : :obj:`~numpy.typing.ArrayLike` Array of boolean membership, over the full length of original nodes, to support the construction of indexed connectivities. - edge_indices, face_indices : ArrayLike, optional + edge_indices, face_indices : :obj:`~numpy.typing.ArrayLike`, optional The indices to use when indexing member connectivities with edge or face :attr:`~Connectivity.location` respectively. mesh_id : int @@ -3228,15 +3243,16 @@ def face_node(self): class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): - # TODO: docstring is out of date with the latest code. - """A container representing the UGRID ``cf_role``: ``location_index_set``. + """A sub-mesh defined by an index set into a parent :class:`MeshXY`. - A container representing the UGRID ``cf_role``: - ``location_index_set``. Achieved by referencing an original :class:`MeshXY` - instance (:attr:`super_mesh`), together with a specific :attr:`location` - (``node``/``edge``/``face``) and a set of :attr:`indices`. This is strictly - a view onto the original :class:`MeshXY` - it does not store its own - coordinates or connectivities. + Represents the UGRID ``cf_role``: ``location_index_set``. This is used + to associate a subset of mesh elements (nodes, edges, or faces) with a + data variable, e.g. 3 temperature readings associated with faces 0, 4, and + 15 from the original mesh. Rather than + storing its own coordinates or connectivities, a :class:`_MeshIndexSet` + references a parent :class:`MeshXY` (:attr:`mesh`), a :attr:`location` + (``node`` / ``edge`` / ``face``), and an array of :attr:`indices` that + select the relevant elements from that parent mesh. Warnings -------- @@ -3251,7 +3267,6 @@ class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): # TODO: implement I/O (iris#6123). # TODO: update the full documentation - # TODO: docstrings def __init__( self, indices: ArrayLike, @@ -3264,6 +3279,30 @@ def __init__( attributes: Optional[dict] = None, start_index: Literal[0, 1] = 0, ): + """Create a :class:`_MeshIndexSet`. + + Parameters + ---------- + indices : :obj:`~numpy.typing.ArrayLike` + Indices selecting elements from ``mesh`` at ``location``. + mesh : MeshXY + The parent mesh being indexed. + location : "node", "edge" or "face" + The element type being indexed. + standard_name : str, optional + CF standard name for the index variable. + long_name : str, optional + Descriptive name. + var_name : str, optional + NetCDF variable name. + units : str, optional + Units; defaults to ``mesh.units``. + attributes : dict, optional + Arbitrary metadata attributes. + start_index : 0 or 1, default=0 + Index origin; default is ``0``. + + """ if not isinstance(mesh, MeshXY): raise TypeError(f"`mesh` must be `MeshXY`. Got {type(mesh)}") From a77af27645705e70d5e1364036755a14a91ae4cd Mon Sep 17 00:00:00 2001 From: Martin Yeo <40734014+trexfeathers@users.noreply.github.com> Date: Thu, 23 Jul 2026 14:47:40 +0100 Subject: [PATCH 5/6] Fixes based on early testing. (#7216) --- lib/iris/mesh/components.py | 81 ++++++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 10 deletions(-) diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 0532771a5b..7f513450ed 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -2794,7 +2794,16 @@ def _index_conn_array( flat_inds_safe = al.where(missing_inds, 0, flat_inds_nomask) # Here's the core indexing operation. # The comma applies all inds-array values to the *first* dimension. - flat_result = array[flat_inds_safe,] + try: + flat_result = array[flat_inds_safe,] + except NotImplementedError: + # Protect against Dask NotImplementedError: + # "Slicing an array with unknown chunks with a dask.array of ints is + # not supported" + if not lazy: + raise + array.compute_chunk_sizes() + flat_result = array[flat_inds_safe,] # Fix 'missing' locations, and restore the proper shape. flat_result_masked = al.ma.masked_array(flat_result, missing_inds) result = flat_result_masked.reshape(indexing.shape) @@ -2803,8 +2812,8 @@ def _index_conn_array( class _MeshConnectivityManagerBase(ABC): # Override these in subclasses. - REQUIRED: tuple = NotImplemented - OPTIONAL: tuple = NotImplemented + REQUIRED: tuple[str, ...] | tuple[()] = NotImplemented + OPTIONAL: tuple[str, ...] | tuple[()] = NotImplemented _view_message: str | None = None def __init__(self, *connectivities, view: Optional[str] = None): @@ -3041,6 +3050,18 @@ def element_filter(instances, loc_arg, loc_name): result_dict = {k: v for k, v in self._members.items() if id(v) in result_ids} return result_dict + def _indexed_manager( + self, cf_roles: list[str] + ) -> type["_MeshConnectivityManagerBase"]: + """Identify appropriate Manager for a subsetted list of connectivities. + + Supports :meth:`indexed`. Required because removing connectivities may + violate the minimum requirements for the original manager class, so + we might need to 'downgrade'. Should be overridden as appropriate in + subclasses. + """ + return self.__class__ + def indexed( self, node_bool_index: ArrayLike, @@ -3131,7 +3152,10 @@ def indexed( else: view_message = None - result = self.__class__( + manager_class = self._indexed_manager( + [c.cf_role for c in indexed_members.values() if c is not None] + ) + result = manager_class( *[c for c in indexed_members.values() if c is not None], view=view_message ) @@ -3180,9 +3204,20 @@ def remove( return removal_dict -class _Mesh1DConnectivityManager(_MeshConnectivityManagerBase): - REQUIRED = ("edge_node_connectivity",) - OPTIONAL = () +class _MeshNodeConnectivityManager(_MeshConnectivityManagerBase): + """Manager specifically to support node-based _MeshIndexSet i.e. no connectivities.""" + + REQUIRED: tuple[str, ...] | tuple[()] = () + OPTIONAL: tuple[str, ...] | tuple[()] = () + + @property + def all_members(self): + return () + + +class _Mesh1DConnectivityManager(_MeshNodeConnectivityManager): + REQUIRED: tuple[str, ...] | tuple[()] = ("edge_node_connectivity",) + OPTIONAL: tuple[str, ...] | tuple[()] = () @property def all_members(self): @@ -3192,10 +3227,26 @@ def all_members(self): def edge_node(self): return self._members["edge_node_connectivity"] + def _indexed_manager( + self, cf_roles: list[str] + ) -> type["_MeshConnectivityManagerBase"]: + result: type["_MeshConnectivityManagerBase"] + if "edge_node_connectivity" in cf_roles: + result = _Mesh1DConnectivityManager + elif len(cf_roles) == 0: + result = _MeshNodeConnectivityManager + else: + message = ( + f"Unhandled combination of connectivities: {cf_roles}. " + f"Cannot determine appropriate connectivity manager class." + ) + raise NotImplementedError(message) + return result -class _Mesh2DConnectivityManager(_MeshConnectivityManagerBase): - REQUIRED = ("face_node_connectivity",) - OPTIONAL = ( + +class _Mesh2DConnectivityManager(_Mesh1DConnectivityManager): + REQUIRED: tuple[str, ...] | tuple[()] = ("face_node_connectivity",) + OPTIONAL: tuple[str, ...] | tuple[()] = ( "edge_node_connectivity", "face_edge_connectivity", "face_face_connectivity", @@ -3238,6 +3289,16 @@ def face_face(self): def face_node(self): return self._members["face_node_connectivity"] + def _indexed_manager( + self, cf_roles: list[str] + ) -> type["_MeshConnectivityManagerBase"]: + result: type["_MeshConnectivityManagerBase"] + if "face_node_connectivity" in cf_roles: + result = _Mesh2DConnectivityManager + else: + result = super()._indexed_manager(cf_roles) + return result + Location = Literal["edge", "node", "face"] From 43b82f2588ce360c8b171baa079b06c5e705eaf9 Mon Sep 17 00:00:00 2001 From: pt331 <144435193+pt331@users.noreply.github.com> Date: Fri, 24 Jul 2026 18:40:58 +0100 Subject: [PATCH 6/6] Tests for indexing meshes (#7218) * Barebones tests * assertion fix and allowing for single value index * fixed tests/stock/mesh topology_dimension * Update as_mesh to deal with locations not existing * Copied as_mesh changes from unit tests branch * Applied equality changes offered by @trexfeathers * Remove cube_mesh_node from test_subset_indexing_new_mesh test * updating to full tests - test_subset_indexing_mesh_index_set fails now * Fixed looking at the wrong part of cube and _MeshIndexSet test --- lib/iris/mesh/components.py | 74 +++++-- .../integration/mesh/test_indexing_meshes.py | 189 ++++++++++++++++++ lib/iris/tests/stock/mesh.py | 8 +- 3 files changed, 248 insertions(+), 23 deletions(-) create mode 100644 lib/iris/tests/integration/mesh/test_indexing_meshes.py diff --git a/lib/iris/mesh/components.py b/lib/iris/mesh/components.py index 7f513450ed..dd0ad6d198 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -696,11 +696,16 @@ def __eq__(self, other) -> bool | NotImplementedType: result: bool | NotImplementedType = NotImplemented if isinstance(other, _MeshXYMixin): - result = self.metadata == other.metadata - if result: - result = self.all_coords == other.all_coords - if result: - result = self.all_connectivities == other.all_connectivities + # Using identity speeds up several real-world equality checks + # for _MeshIndexSet and MeshCoord. + result = self is other + + if not result: + result = self.metadata == other.metadata + if result: + result = self.all_coords == other.all_coords + if result: + result = self.all_connectivities == other.all_connectivities return result @@ -3396,6 +3401,18 @@ def __init__( attributes=attributes, ) + def __eq__(self, other) -> bool | NotImplementedType: + result: bool | NotImplementedType = NotImplemented + + if isinstance(other, _MeshIndexSet): + result = self.metadata == other.metadata + # Don't check coords or connectivities as these are + # fully derived using metadata. + if result: + result = self.indices == other.indices + + return result + def __getstate__(self) -> tuple[ArrayLike, MeshIndexSetMetadata]: return ( self.indices, @@ -3453,21 +3470,22 @@ def _calculate_node_bool_index(self): case "node": # self.indices is a user-supplied index array over the nodes; convert # it to a fixed-shape boolean membership mask. - monotonic, direction = iris.util.monotonic( - self.indices, strict=True, return_direction=True - ) - if not (monotonic and direction == 1): - # TODO: boolean 'mask' array precludes non-monotonic indexing, - # but is only needed to support connectivity construction, and - # only causes problems for coordinate construction. Separate - # logic to allow array of integer indices for coordinate - # construction. - message = ( - "Indexing the nodes on a Mesh currently requires strictly " - "increasing indices. Contact the Iris developers if this " - "causes you problems." + if len(self.indices) > 1: # Single value index is fine + monotonic, direction = iris.util.monotonic( + self.indices, strict=True, return_direction=True ) - raise ValueError(message) + if not (monotonic and direction == 1): + # TODO: boolean 'mask' array precludes non-monotonic indexing, + # but is only needed to support connectivity construction, and + # only causes problems for coordinate construction. Separate + # logic to allow array of integer indices for coordinate + # construction. + message = ( + "Indexing the nodes on a Mesh currently requires strictly " + "increasing indices. Contact the Iris developers if this " + "causes you problems." + ) + raise ValueError(message) indices = self.indices al = da if _lazy.is_lazy_data(indices) else np node_mask = al.zeros(n_original_nodes, dtype=bool) @@ -3611,6 +3629,18 @@ def as_mesh(self) -> MeshXY: kwargs: Kwargs = {"mesh_id": id(self.mesh), "frozen": True} coord_man = self.mesh._coord_manager.indexed(*indices, **kwargs) conn_man = self.mesh._connectivity_manager.indexed(*indices, **kwargs) + has_edges = hasattr(conn_man, "edge_node") + has_faces = hasattr(conn_man, "face_node") + if has_faces: + topology_dimension = 2 + elif has_edges: + topology_dimension = 1 + else: + message = ( + f"Cannot create a MeshXY from a MeshIndexSet with no edge or face " + f"connectivities. This is a *{self.location}* MeshIndexSet." + ) + raise NotImplementedError(message) def _coords_and_axes( location: Literal["node", "edge", "face"], @@ -3619,11 +3649,11 @@ def _coords_and_axes( return [(getattr(coords, f"{location}_{axis}"), axis) for axis in self.AXES] return MeshXY( - topology_dimension=self.topology_dimension, + topology_dimension=topology_dimension, node_coords_and_axes=_coords_and_axes("node"), connectivities=[conn for conn in conn_man.all_members if conn is not None], - edge_coords_and_axes=_coords_and_axes("edge"), - face_coords_and_axes=_coords_and_axes("face"), + edge_coords_and_axes=_coords_and_axes("edge") if has_edges else None, + face_coords_and_axes=_coords_and_axes("face") if has_faces else None, standard_name=self.standard_name, long_name=self.long_name, var_name=self.var_name, diff --git a/lib/iris/tests/integration/mesh/test_indexing_meshes.py b/lib/iris/tests/integration/mesh/test_indexing_meshes.py new file mode 100644 index 0000000000..38d1effbbe --- /dev/null +++ b/lib/iris/tests/integration/mesh/test_indexing_meshes.py @@ -0,0 +1,189 @@ +# Copyright Iris contributors +# +# This file is part of Iris and is released under the BSD license. +# See LICENSE in the root of the repository for full licensing details. + +import pytest + +from iris.coords import AuxCoord, Coord +from iris.cube import Cube +from iris.experimental import mesh_coord_indexing +from iris.loading import load_cube +from iris.mesh.components import MeshCoord, MeshXY, _MeshIndexSet +from iris.tests import _shared_utils +from iris.tests.stock.mesh import sample_mesh, sample_mesh_cube + + +@pytest.fixture +def cube_mesh_from_file(): + file_path = _shared_utils.get_data_path( + [ + "NetCDF", + "unstructured_grid", + "lfric_ngvat_2D_72t_face_half_levels_main_conv_rain.nc", + ] + ) + return (load_cube(file_path, "conv_rain"), "face") + + +@pytest.fixture +def cube_mesh_node(): + location = "node" + mesh = sample_mesh(n_nodes=15, n_edges=3, n_faces=3) # Cannot have a node only mesh + return (sample_mesh_cube(location=location, mesh=mesh), location) + + +@pytest.fixture +def cube_mesh_edge(): + location = "edge" + mesh = sample_mesh(n_nodes=15, n_edges=3, n_faces=0) + return (sample_mesh_cube(location=location, mesh=mesh), location) + + +@pytest.fixture +def cube_mesh_face(): + location = "face" + mesh = sample_mesh(n_nodes=15, n_edges=0, n_faces=3) + return (sample_mesh_cube(location=location, mesh=mesh), location) + + +def change_and_assert_no_change_in_right(left: Coord, right: Coord, value: int): + left.points[0] = value + assert right.points[0] != value + + +@pytest.mark.parametrize( + "fixture", + ["cube_mesh_from_file", "cube_mesh_node", "cube_mesh_edge", "cube_mesh_face"], +) +def test_subset_indexing_auxcoord(fixture, request): + cube: Cube + location: str + (cube, location) = request.getfixturevalue(fixture) + with mesh_coord_indexing.SETTING.context(mesh_coord_indexing.Options.AUX_COORD): + indexed_cube = cube[0, 0:1] + + indexed_lat = indexed_cube.coord(standard_name="latitude") + indexed_lon = indexed_cube.coord(standard_name="longitude") + # The mesh's lat/lon should be represented as AuxCoord + assert isinstance(indexed_lat, AuxCoord) + assert isinstance(indexed_lon, AuxCoord) + # And no mesh in cube + assert indexed_cube.mesh is None + + assert cube.mesh is not None + original_lat = cube.mesh.coord(standard_name="latitude", location=location) + original_lon = cube.mesh.coord(standard_name="longitude", location=location) + assert original_lat is not None + assert original_lon is not None + # Any changes to the indexed cube's mesh should not be reflected in the original + value = 9999 + change_and_assert_no_change_in_right(indexed_lat, original_lat, value) + change_and_assert_no_change_in_right(indexed_lon, original_lon, value) + # and vice versa + value = -9999 + change_and_assert_no_change_in_right(original_lat, indexed_lat, value) + change_and_assert_no_change_in_right(original_lon, indexed_lon, value) + + +# Excluding "cube_mesh_node" from this test because it is an invalid thing to do +# Test for the raised exception exists in unit tests +@pytest.mark.parametrize( + "fixture", + ["cube_mesh_from_file", "cube_mesh_edge", "cube_mesh_face"], +) +def test_subset_indexing_new_mesh(fixture, request): + cube: Cube + location: str + (cube, location) = request.getfixturevalue(fixture) + + with mesh_coord_indexing.SETTING.context(mesh_coord_indexing.Options.NEW_MESH): + indexed_cube = cube[0, 0:1] + # The mesh's lat/lon should be represented as MeshCoord with a Mesh as the mesh + indexed_lat = indexed_cube.coord(standard_name="latitude") + indexed_lon = indexed_cube.coord(standard_name="longitude") + assert isinstance(indexed_lat, MeshCoord) + assert isinstance(indexed_lon, MeshCoord) + assert isinstance(indexed_lat.mesh, MeshXY) + assert isinstance(indexed_lon.mesh, MeshXY) + + assert cube.mesh is not None + original_lat = cube.mesh.coord(standard_name="latitude", location=location) + original_lon = cube.mesh.coord(standard_name="longitude", location=location) + assert original_lat is not None + assert original_lon is not None + # Any changes to the indexed cube's mesh should not be reflected in the original + value = 9999 + change_and_assert_no_change_in_right(indexed_lat, original_lat, value) + change_and_assert_no_change_in_right(indexed_lon, original_lon, value) + # and vice versa + value = -9999 + change_and_assert_no_change_in_right(original_lat, indexed_lat, value) + change_and_assert_no_change_in_right(original_lon, indexed_lon, value) + + +@pytest.mark.parametrize( + "fixture", + ["cube_mesh_from_file", "cube_mesh_node", "cube_mesh_edge", "cube_mesh_face"], +) +def test_subset_indexing_mesh_index_set(fixture, request): + cube: Cube + location: str + (cube, location) = request.getfixturevalue(fixture) + with mesh_coord_indexing.SETTING.context( + mesh_coord_indexing.Options.MESH_INDEX_SET + ): + indexed_cube = cube[0, 0:1] + # The mesh's lat/lon should be represented as MeshCoord, + # with a _MeshIndexSet as the mesh + indexed_cube_lat = indexed_cube.coord(standard_name="latitude") + indexed_cube_lon = indexed_cube.coord(standard_name="longitude") + assert isinstance(indexed_cube_lat, MeshCoord) + assert isinstance(indexed_cube_lon, MeshCoord) + assert isinstance(indexed_cube_lat.mesh, _MeshIndexSet) + assert isinstance(indexed_cube_lon.mesh, _MeshIndexSet) + # The indexed_cube's mesh is also a _MeshIndexSet + assert isinstance(indexed_cube.mesh, _MeshIndexSet) + + assert isinstance(cube.mesh, MeshXY) + original_lat = cube.mesh.coord(standard_name="latitude", location=location) + original_lon = cube.mesh.coord(standard_name="longitude", location=location) + assert original_lat is not None + assert original_lon is not None + + indexed_mesh_lat = indexed_cube.mesh.coord( + standard_name="latitude", location=location + ) + indexed_mesh_lon = indexed_cube.mesh.coord( + standard_name="longitude", location=location + ) + assert indexed_mesh_lat is not None + assert indexed_mesh_lon is not None + # Changing the values of a _MeshIndexSet's coords does nothing to the original + value = 9999 + change_and_assert_no_change_in_right(indexed_mesh_lat, original_lat, value) + change_and_assert_no_change_in_right(indexed_mesh_lon, original_lon, value) + + # Changing the original mesh is reflected in the _MeshIndexSet + # Importantly, the _MeshIndexSet must be declared after the changes to the original + value = -9999 + + def change_original_mesh_reflected_in_index_set( + original: Cube, indexed: Cube, name: str, loc: str, value: int + ): + assert original.mesh is not None + original_coord = original.mesh.coord(standard_name=name, location=loc) + assert original_coord is not None + original_coord.points[0] = value + + assert indexed.mesh is not None + indexed_coord = indexed.mesh.coord(standard_name=name, location=loc) + assert indexed_coord is not None + assert indexed_coord.points[0] == value + + change_original_mesh_reflected_in_index_set( + cube, indexed_cube, "latitude", location, value + ) + change_original_mesh_reflected_in_index_set( + cube, indexed_cube, "longitude", location, value + ) diff --git a/lib/iris/tests/stock/mesh.py b/lib/iris/tests/stock/mesh.py index 3824ba84fc..3d920d3630 100644 --- a/lib/iris/tests/stock/mesh.py +++ b/lib/iris/tests/stock/mesh.py @@ -81,6 +81,8 @@ def sample_mesh( ) node_y = AuxCoord(1200 + arr.arange(n_nodes), standard_name="latitude") + topology_dimension = 0 + connectivities = [] if n_edges == 0: edge_coords_and_axes = None @@ -100,6 +102,8 @@ def sample_mesh( edge_y = AuxCoord(2200 + arr.arange(n_edges), standard_name="latitude") edge_coords_and_axes = [(edge_x, "x"), (edge_y, "y")] + topology_dimension = 1 + if n_faces == 0: face_coords_and_axes = None else: @@ -118,8 +122,10 @@ def sample_mesh( face_y = AuxCoord(3200 + arr.arange(n_faces), standard_name="latitude") face_coords_and_axes = [(face_x, "x"), (face_y, "y")] + topology_dimension = 2 + mesh = MeshXY( - topology_dimension=2, + topology_dimension=topology_dimension, node_coords_and_axes=[(node_x, "x"), (node_y, "y")], connectivities=connectivities, edge_coords_and_axes=edge_coords_and_axes,