diff --git a/docs/src/conf.py b/docs/src/conf.py index 0baf04c5dd..0e491f3ddc 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -255,7 +255,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 5058cfbb69..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 ): @@ -1596,6 +1599,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 06c796fbab..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 @@ -277,7 +278,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): @@ -319,7 +331,7 @@ def summary( precision=None, convert_dates=True, _section_indices=None, - ): + ) -> str: r"""Make a printable text summary. Parameters @@ -609,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") @@ -2473,7 +2486,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() ) @@ -3047,7 +3060,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): @@ -3127,7 +3142,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) @@ -3137,7 +3152,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 6653018055..24d0230918 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 @@ -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( @@ -3286,8 +3292,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))) @@ -3312,8 +3317,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 9655089d81..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: @@ -994,7 +999,6 @@ def _add_aux_coords( MeshEdgeCoords, MeshFaceCoords, MeshNodeCoords, - MeshXY, ) # Exclude any mesh coords, which are bundled in with the aux-coords. @@ -1003,12 +1007,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 c8e097b801..dd0ad6d198 100644 --- a/lib/iris/mesh/components.py +++ b/lib/iris/mesh/components.py @@ -19,20 +19,39 @@ from collections.abc import Container from contextlib import contextmanager from datetime import datetime -from typing import Iterable, Literal, Mapping, NamedTuple +import math +from types import NotImplementedType +from typing import ( + Any, + Generator, + Iterable, + Literal, + Mapping, + NamedTuple, + Optional, + Sequence, + TypedDict, +) 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 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 @@ -630,6 +649,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`. @@ -648,28 +668,22 @@ 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. +class _MeshXYMixin(Mesh, ABC): + """Mixin providing XY-coordinate behaviour for :class:`Mesh` subclasses. - References - ---------- - .. [1] The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ + 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 + _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. AXES = ("x", "y") @@ -678,657 +692,335 @@ 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. + def __eq__(self, other) -> bool | NotImplementedType: + result: bool | NotImplementedType = NotImplemented - .. note:: + if isinstance(other, _MeshXYMixin): + # Using identity speeds up several real-world equality checks + # for _MeshIndexSet and MeshCoord. + result = self is other - 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``. + 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 - """ - # TODO: support volumes. - # TODO: support (coord, "z") + return result - self._metadata_manager = metadata_manager_factory(MeshMetadata) + 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)) - # 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 + def __ne__(self, other) -> bool | NotImplementedType: + result = self.__eq__(other) + if result is not NotImplemented: + result = not result + return result - self.node_dimension = node_dimension - self.edge_dimension = edge_dimension - self.face_dimension = face_dimension + def summary(self, *args, **kwargs) -> str: + """Return a string representation of the MeshXY. - # 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 + Parameters + ---------- + shorten : bool, default=False + If True, produce a oneline string form of the form . + If False, produce a multi-line detailed print output. - # 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}" + Returns + ------- + str - if not isinstance(node_coords_and_axes, Iterable): - node_coords_and_axes = [node_coords_and_axes] + """ + if len(args) > 0: + shorten = args[0] + else: + shorten = kwargs.get("shorten", False) - if not isinstance(connectivities, Iterable): - connectivities = [connectivities] + if shorten: + result = self._summary_oneline() + else: + result = self._summary_multiline() + return result - 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 + def __repr__(self): + return self.summary(shorten=True) - # 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) + def __str__(self): + return self.summary(shorten=False) - 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) + 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: str | None = 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: - emsg = f"Unsupported 'topology_dimension', got {topology_dimension!r}." - raise NotImplementedError(emsg) + # Mimic the generic object.__str__ style. + mesh_id = id(self) + mesh_string = f"<{self.__class__.__name__} object at {hex(mesh_id)}>" - @classmethod - def from_coords(cls, *coords): - r"""Construct a :class:`MeshXY` by derivation from 1/more :class:`~iris.coords.Coord`. + return mesh_string - 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 _summary_multiline(self) -> str: + # Produce a readable multi-line summary of the Mesh content. + lines = [] + n_indent = 4 + indent_str = " " * n_indent - * ``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 line(text, i_indent=0): + indent = indent_str * i_indent + lines.append(f"{indent}{text}") - 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``. + line(f"{self.__class__.__name__} : '{self.name()}'") + line(f"topology_dimension: {self.topology_dimension}", 1) + for element in ("node", "edge", "face"): + if element == "node": + main_conn_string = None + else: + main_conn_name = f"{element}_node_connectivity" + main_conn: Connectivity | None = getattr(self, main_conn_name, None) + if main_conn is None: + continue - Returns - ------- - :class:`MeshXY` + 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) + else: + coord_string = "None" + line(coord_string, 3) - Notes - ----- - .. note:: - Any resulting duplicate nodes are not currently removed, due to the - computational intensity. + # Having dealt with essential info, now add any optional connectivities + # N.B. includes boundaries: as optional connectivity, not an "element" + optional_conn_names = ( + "boundary_connectivity", + "face_face_connectivity", + "face_edge_connectivity", + "edge_face_connectivity", + ) + optional_conn_keys = [getattr(self, name, None) for name in optional_conn_names] + optional_conns = { + name: conn + for conn, name in zip(optional_conn_keys, optional_conn_names) + if conn is not None + } + if optional_conns: + line("optional connectivities", 1) + for name, conn in optional_conns.items(): + conn_string = conn.summary(shorten=True, linewidth=0) + line(f"{name}: {conn_string}", 2) - .. 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. + # Output the detail properties, basically those from CFVariableMixin + for name in BaseMetadata._members: + val = getattr(self, name, None) + if val is not None: + if name == "units": + show = val.origin != Unit(None) + elif isinstance(val, Container): + show = bool(val) + else: + show = val is not None + if show: + if name == "attributes": + # Use a multi-line form for this. + line("attributes:", 1) + max_attname_len = max(len(attr) for attr in val.keys()) + for attrname, attrval in val.items(): + attrname = attrname.ljust(max_attname_len) + if isinstance(attrval, str): + # quote strings + attrval = repr(attrval) + # and abbreviate really long ones + attrval = clip_string(attrval) + attr_string = f"{attrname} {attrval}" + line(attr_string, 2) + else: + line(f"{name}: {val!r}", 1) - .. testsetup:: + result = "\n".join(lines) + return result - from iris import load_cube, sample_data_path - from iris.mesh import ( - MeshXY, - MeshCoord, - ) + @abstractmethod + def __getstate__(self): + raise NotImplementedError - file_path = sample_data_path("mesh_C4_synthetic_float.nc") - cube_w_mesh = load_cube(file_path) + @abstractmethod + def __setstate__(self, state): + raise NotImplementedError - Examples - -------- - :: + @property + def _connectivity_manager(self) -> "_MeshConnectivityManagerBase": + # @property enables interruption/customisation in subclasses. + return self._connectivity_manager_attr - # Reconstruct a cube-with-mesh after subsetting it. + @_connectivity_manager.setter + def _connectivity_manager(self, manager: "_MeshConnectivityManagerBase") -> None: + # @property enables interruption/customisation in subclasses. + self._connectivity_manager_attr = manager - >>> 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 + @property + def _coord_manager(self) -> "_MeshCoordinateManagerBase": + # @property enables interruption/customisation in subclasses. + return self._coord_manager_attr - >>> new_mesh = MeshXY.from_coords(*orig_coords) - >>> new_coords = new_mesh.to_MeshCoords(location=cube_w_mesh.location) + @_coord_manager.setter + def _coord_manager(self, manager: "_MeshCoordinateManagerBase") -> None: + # @property enables interruption/customisation in subclasses. + self._coord_manager_attr = manager - # 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()) + @property + def all_connectivities(self) -> NamedTuple: + """All the :class:`~iris.mesh.Connectivity` instances of the :class:`MeshXY`.""" + return self._connectivity_manager.all_members - >>> 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 + @property + 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) -> NamedTuple: + """All the :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" + return self._coord_manager.all_members - # 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) + @property + def boundary_node_connectivity(self) -> Connectivity: + """The *optional* UGRID ``boundary_node_connectivity`` :class:`~iris.mesh.Connectivity`. - for array in ("points", "bounds"): - check_shape(array) + The *optional* UGRID ``boundary_node_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`. - # Determine dimensionality, using first coord. - first_coord = coords[0] + """ + return self._connectivity_manager.boundary_node # type:ignore[attr-defined] - ndim = first_coord.ndim - if ndim != 1: - message = f"Expected coordinate ndim == 1, got: f{ndim} ." - raise ValueError(message) + @property + def edge_coords(self) -> MeshEdgeCoords: + """The *optional* UGRID ``edge`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" + return self._coord_manager.edge_coords - 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" + @property + @abstractmethod + def edge_dimension(self) -> str: + raise NotImplementedError() - # 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) + @property + def edge_face_connectivity(self) -> Connectivity: + """The *optional* UGRID ``edge_face_connectivity`` :class:`~iris.mesh.Connectivity`. - # 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)) + The *optional* UGRID ``edge_face_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`. - centre_points = coord.core_points() - centre_coords.append(AuxCoord(points=centre_points, **coord_kwargs)) + """ + return self._connectivity_manager.edge_face # type:ignore[attr-defined] - ##### - # 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)) + @property + def edge_node_connectivity(self) -> Connectivity: + """The UGRID ``edge_node_connectivity`` :class:`~iris.mesh.Connectivity`. - def axes_assign(coord_list): - coords_sorted = [coord_list[ix] for ix in axis_indices] - return zip(coords_sorted, axes_required) + The UGRID ``edge_node_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`, which is **required** for :attr:`MeshXY.topology_dimension` + of ``1``, and *optionally required* for + :attr:`MeshXY.topology_dimension` ``>=2``. - node_coords_and_axes = axes_assign(node_coords) - centre_coords_and_axes = axes_assign(centre_coords) - ##### + """ + return self._connectivity_manager.edge_node # type:ignore[attr-defined] - # 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) + @property + def face_coords(self) -> AuxCoord: + """The *optional* UGRID ``face`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" + return self._coord_manager.face_coords # type:ignore[attr-defined] - def __eq__(self, other): - result = NotImplemented + @property + @abstractmethod + def face_dimension(self) -> str: + raise NotImplementedError() - 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 + @property + def face_edge_connectivity(self) -> Connectivity: + """The *optional* UGRID ``face_edge_connectivity``:class:`~iris.mesh.Connectivity`. - return result + The *optional* UGRID ``face_edge_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`. - 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)) + """ + # optional + return self._connectivity_manager.face_edge # type:ignore[attr-defined] - def __getstate__(self): - return ( - self._metadata_manager, - self._coord_manager, - self._connectivity_manager, - ) + @property + def face_face_connectivity(self) -> Connectivity: + """The *optional* UGRID ``face_face_connectivity`` :class:`~iris.mesh.Connectivity`. - def __ne__(self, other): - result = self.__eq__(other) - if result is not NotImplemented: - result = not result - return result + The *optional* UGRID ``face_face_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`. - def summary(self, shorten=False): - """Return a string representation of the MeshXY. + """ + return self._connectivity_manager.face_face # type:ignore[attr-defined] - Parameters - ---------- - shorten : bool, default=False - If True, produce a oneline string form of the form . - If False, produce a multi-line detailed print output. + @property + def face_node_connectivity(self) -> Connectivity: + """Return ``face_node_connectivity``:class:`~iris.mesh.Connectivity`. - Returns - ------- - str + The UGRID ``face_node_connectivity`` + :class:`~iris.mesh.Connectivity` of the + :class:`MeshXY`, which is **required** for :attr:`MeshXY.topology_dimension` + of ``2``, and *optionally required* for :attr:`MeshXY.topology_dimension` + of ``3``. """ - if shorten: - result = self._summary_oneline() - else: - result = self._summary_multiline() - return result - - def __repr__(self): - return self.summary(shorten=True) + return self._connectivity_manager.face_node # type:ignore[attr-defined] - 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"" + @property + def node_coords(self) -> MeshNodeCoords: + """The **required** UGRID ``node`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" + return self._coord_manager.node_coords - return mesh_string + @property + @abstractmethod + def node_dimension(self) -> str: + raise NotImplementedError() - def _summary_multiline(self): - # Produce a readable multi-line summary of the Mesh content. - lines = [] - n_indent = 4 - indent_str = " " * n_indent + def add_connectivities(self, *connectivities): + """Add one or more :class:`~iris.mesh.Connectivity` instances to the :class:`MeshXY`. - def line(text, i_indent=0): - indent = indent_str * i_indent - lines.append(f"{indent}{text}") + Parameters + ---------- + *connectivities : iterable of object + A collection of one or more + :class:`~iris.mesh.Connectivity` instances to + add to the :class:`MeshXY`. - line(f"MeshXY : '{self.name()}'") - line(f"topology_dimension: {self.topology_dimension}", 1) - for element in ("node", "edge", "face"): - if element == "node": - element_exists = True - 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: - coord_string = coord.summary(shorten=True, linewidth=0) - line(coord_string, 3) - - # Having dealt with essential info, now add any optional connectivities - # N.B. includes boundaries: as optional connectivity, not an "element" - optional_conn_names = ( - "boundary_connectivity", - "face_face_connectivity", - "face_edge_connectivity", - "edge_face_connectivity", - ) - optional_conns = [getattr(self, name, None) for name in optional_conn_names] - optional_conns = { - name: conn - for conn, name in zip(optional_conns, optional_conn_names) - if conn is not None - } - if optional_conns: - line("optional connectivities", 1) - for name, conn in optional_conns.items(): - conn_string = conn.summary(shorten=True, linewidth=0) - line(f"{name}: {conn_string}", 2) - - # Output the detail properties, basically those from CFVariableMixin - for name in BaseMetadata._members: - val = getattr(self, name, None) - if val is not None: - if name == "units": - show = val.origin != Unit(None) - elif isinstance(val, Container): - show = bool(val) - else: - show = val is not None - if show: - if name == "attributes": - # Use a multi-line form for this. - line("attributes:", 1) - max_attname_len = max(len(attr) for attr in val.keys()) - for attrname, attrval in val.items(): - attrname = attrname.ljust(max_attname_len) - if isinstance(attrval, str): - # quote strings - attrval = repr(attrval) - # and abbreviate really long ones - attrval = clip_string(attrval) - attr_string = f"{attrname} {attrval}" - line(attr_string, 2) - else: - line(f"{name}: {val!r}", 1) - - result = "\n".join(lines) - return result - - 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 - - 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 - - @property - def all_connectivities(self): - """All the :class:`~iris.mesh.Connectivity` instances of the :class:`MeshXY`.""" - return self._connectivity_manager.all_members - - @property - def _last_modified(self): - """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): - """All the :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" - return self._coord_manager.all_members - - @property - def boundary_node_connectivity(self): - """The *optional* UGRID ``boundary_node_connectivity`` :class:`~iris.mesh.Connectivity`. - - The *optional* UGRID ``boundary_node_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`. - - """ - return self._connectivity_manager.boundary_node - - @property - def edge_coords(self): - """The *optional* UGRID ``edge`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" - 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 - - @property - def edge_face_connectivity(self): - """The *optional* UGRID ``edge_face_connectivity`` :class:`~iris.mesh.Connectivity`. - - The *optional* UGRID ``edge_face_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`. - - """ - return self._connectivity_manager.edge_face - - @property - def edge_node_connectivity(self): - """The UGRID ``edge_node_connectivity`` :class:`~iris.mesh.Connectivity`. - - The UGRID ``edge_node_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`, which is **required** for :attr:`MeshXY.topology_dimension` - of ``1``, and *optionally required* for - :attr:`MeshXY.topology_dimension` ``>=2``. - - """ - return self._connectivity_manager.edge_node - - @property - def face_coords(self): - """The *optional* UGRID ``face`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" - 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 - - @property - def face_edge_connectivity(self): - """The *optional* UGRID ``face_edge_connectivity``:class:`~iris.mesh.Connectivity`. - - The *optional* UGRID ``face_edge_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`. - - """ - # optional - return self._connectivity_manager.face_edge - - @property - def face_face_connectivity(self): - """The *optional* UGRID ``face_face_connectivity`` :class:`~iris.mesh.Connectivity`. - - The *optional* UGRID ``face_face_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`. - - """ - return self._connectivity_manager.face_face - - @property - def face_node_connectivity(self): - """Return ``face_node_connectivity``:class:`~iris.mesh.Connectivity`. - - The UGRID ``face_node_connectivity`` - :class:`~iris.mesh.Connectivity` of the - :class:`MeshXY`, which is **required** for :attr:`MeshXY.topology_dimension` - of ``2``, and *optionally required* for :attr:`MeshXY.topology_dimension` - of ``3``. - - """ - return self._connectivity_manager.face_node - - @property - def node_coords(self): - """The **required** UGRID ``node`` :class:`~iris.coords.AuxCoord` coordinates of the :class:`MeshXY`.""" - 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 - - def add_connectivities(self, *connectivities): - """Add one or more :class:`~iris.mesh.Connectivity` instances to the :class:`MeshXY`. - - Parameters - ---------- - *connectivities : iterable of object - A collection of one or more - :class:`~iris.mesh.Connectivity` instances to - add to the :class:`MeshXY`. - - """ - self._connectivity_manager.add(*connectivities) + """ + self._connectivity_manager.add(*connectivities) def add_coords( self, @@ -1948,6 +1640,424 @@ def to_MeshCoords(self, location): 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) + + 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: Coord): + 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 + + """ + 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): + 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. + axes_present = [guess_coord_axis(coord) for coord in coords] + # 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] + except ValueError: + 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 + + @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 + + @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 + + @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. @@ -2011,7 +2121,40 @@ def topology_dimension(self): return self._metadata_manager.topology_dimension -class _Mesh1DCoordinateManager: +class _ManagerMembers[VT](dict[str, VT]): + """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): + 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 _MeshCoordinateManagerBase(ABC): """TBD: require clarity on coord_systems validation. TBD: require clarity on __eq__ support @@ -2028,11 +2171,24 @@ class _Mesh1DCoordinateManager: "edge_y", ) - def __init__(self, node_x, node_y, edge_x=None, edge_y=None): + def __init__( + 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 + # 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[AuxCoord | None] = _ManagerMembers( + {member: None for member in self.ALL} + ) # required coordinates self.node_x = node_x @@ -2051,14 +2207,16 @@ 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) - 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 @@ -2072,12 +2230,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 = 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] @@ -2102,6 +2262,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() @@ -2117,6 +2283,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: @@ -2128,10 +2303,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) @@ -2155,11 +2338,23 @@ 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 [ + (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 - def _members(self, value: dict[str, AuxCoord]): + def _members(self, value: _ManagerMembers[AuxCoord | None]): self.timestamp.update() self._members_dict = value @@ -2168,7 +2363,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 @@ -2188,7 +2383,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 @@ -2350,6 +2549,75 @@ def filters( result_dict = {k: v for k, v in self._members.items() if id(v) in result_ids} return result_dict + def indexed( + self, + node_bool_index: ArrayLike, + edge_indices: Optional[ArrayLike], + face_indices: Optional[ArrayLike], + mesh_id: int, + frozen: bool = False, + ) -> "_MeshCoordinateManagerBase": + """Return an indexed copy of this coordinate manager. + + Parameters + ---------- + 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 : :obj:`~numpy.typing.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 + 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_bool_index, + "edge": edge_indices, + "face": face_indices, + } + indexed_members: dict[str, AuxCoord | None] = {} + for key, coord in self: + indexed = None + if ( + coord is not None + and (indexing := indices_dict[key.split("_")[0]]) is not None + ): + if frozen: + indexed = coord.copy()[indexing] + else: + 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 + + 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}." + ) + result = self.__class__(**indexed_members, view=view_message) # type: ignore[arg-type] + + return result + def remove( self, item=None, @@ -2371,7 +2639,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", @@ -2387,10 +2660,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: @@ -2398,6 +2673,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") @@ -2407,7 +2684,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 @@ -2486,13 +2763,70 @@ def get_face(axis: str | None) -> list[AuxCoord]: ) +def _index_conn_array( + array: np.ndarray | da.Array, + indexing: np.ndarray | da.Array, +) -> 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`. + + 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) + # Here's the core indexing operation. + # The comma applies all inds-array values to the *first* dimension. + 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) + return result + + 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): + 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: @@ -2500,9 +2834,15 @@ 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[Connectivity | None] = _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) @@ -2533,7 +2873,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 = [ @@ -2543,15 +2883,30 @@ def __str__(self): @property @abstractmethod - def all_members(self): - return NotImplemented + def all_members(self) -> NamedTuple: + raise NotImplementedError + + @property + 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 [ + (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 - def _members(self, value): + def _members(self, value: _ManagerMembers[Connectivity | None]): self.timestamp.update() self._members_dict = value @@ -2567,6 +2922,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 = ( @@ -2593,7 +2956,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) @@ -2692,6 +3055,117 @@ 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, + 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_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 : :obj:`~numpy.typing.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. + 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_bool_index, + "edge": edge_indices, + "face": face_indices, + } + + # 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: + 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() + indices_indexed = connectivity.indices_by_location(indices)[indexing] + + # Map node indices in "values" to their new zero-based 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: + indices_remapped = indices_remapped.T + if connectivity.start_index == 1: + indices_remapped += 1 + indexed = connectivity.copy(indices_remapped) + 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 + + 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 + ) + + return result + def remove( self, item=None, @@ -2735,9 +3209,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): @@ -2747,10 +3232,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", @@ -2793,10 +3294,374 @@ 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"] +class _MeshIndexSet(_MeshXYMixin, _DimensionalMetadata): + """A sub-mesh defined by an index set into a parent :class:`MeshXY`. + + 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 + -------- + This class is not yet stable, hence being marked as private. + + References + ---------- + 1. The UGRID Conventions, https://ugrid-conventions.github.io/ugrid-conventions/ + """ + + _NOT_IMPLEMENTED = "_MeshIndexSet_NotImplemented" + + # TODO: implement I/O (iris#6123). + # TODO: update the full documentation + 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, + ): + """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)}") + + 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 + # 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, + ) + + 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, + 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: + return self._NOT_IMPLEMENTED + + @property + def face_dimension(self) -> str: + return self._NOT_IMPLEMENTED + + @property + def node_dimension(self) -> str: + return self._NOT_IMPLEMENTED + + @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_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. + n_original_nodes = self.mesh.node_coords.node_x.shape[0] + 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. + if len(self.indices) > 1: # Single value index is fine + 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 + 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): + # 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) -> _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 + if update: + self._coord_manager_attr = mesh_man.indexed( + self._calculate_node_bool_index(), + 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) -> _MeshConnectivityManagerBase: + 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_bool_index(), + 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_bool_index(), + self._calculate_edge_indices(), + self._calculate_face_indices(), + ] + 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) + 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"], + ) -> 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=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") 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, + units=self.units, + attributes=self.attributes, + ) + + class MeshCoord(AuxCoord): """Geographic coordinate values of data on an unstructured mesh. @@ -2835,9 +3700,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 @@ -2847,29 +3712,32 @@ 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) + 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 - 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. @@ -3065,17 +3933,77 @@ 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." + 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(): + # Assign to help with typing. + mesh = self.mesh + match mesh: + case MeshXY(): + (length,) = self.shape + kwargs = dict( + mesh=mesh, + location=self.location, + 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. + 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=np.asanyarray(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. @@ -3373,25 +4301,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/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/__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/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, 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..4d1cad6c3f 100644 --- a/lib/iris/tests/unit/mesh/components/test_MeshCoord.py +++ b/lib/iris/tests/unit/mesh/components/test_MeshCoord.py @@ -242,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]