Mesh Loading (AVD-1813)

lib/iris/experimental/ugrid/__init__.py

@@ -15,7 +15,9 @@
 from collections.abc import Iterable
 from contextlib import contextmanager
 from functools import wraps
+from itertools import groupby
 import logging
+from pathlib import Path
 import re
 import threading
 
@@ -36,15 +38,22 @@
 from ...common.mixin import CFVariableMixin
 from ...config import get_logger
 from ...coords import AuxCoord, _DimensionalMetadata
-from ...exceptions import ConnectivityNotFoundError, CoordinateNotFoundError
+from ...exceptions import (
+    ConnectivityNotFoundError,
+    ConstraintMismatchError,
+    CoordinateNotFoundError,
+)
 from ...fileformats import cf, netcdf
 from ...fileformats._nc_load_rules.helpers import get_attr_units, get_names
+from ...io import decode_uri, expand_filespecs
 from ...util import guess_coord_axis
 
 __all__ = [
     "CFUGridReader",
     "Connectivity",
     "ConnectivityMetadata",
+    "load_mesh",
+    "load_meshes",
     "Mesh",
     "Mesh1DConnectivities",
     "Mesh1DCoords",
@@ -3274,6 +3283,134 @@ def context(self):
 PARSE_UGRID_ON_LOAD = ParseUGridOnLoad()
 
 
+def _meshes_from_cf(cf_reader):
+    """
+    Common behaviour for extracting meshes from a CFReader.
+
+    Simple now, but expected to increase in complexity as Mesh sharing develops.
+
+    """
+    # Mesh instances are shared between file phenomena.
+    # TODO: more sophisticated Mesh sharing between files.
+    # TODO: access external Mesh cache?
+    mesh_vars = cf_reader.cf_group.meshes
+    meshes = {
+        name: _build_mesh(cf_reader, var, cf_reader.filename)
+        for name, var in mesh_vars.items()
+    }
+    return meshes
+
+
+def load_mesh(uris, var_name=None):
+    """
+    Create a single :class:`Mesh` object from one or more NetCDF files.
+
+    Raises an error if more/less than one :class:`Mesh` is found.
+
+    Parameters
+    ----------
+    uris : str or iterable of str
+        One or more filenames/URI's. Any URI's must support OpenDAP.
+    var_name : str, optional
+        Only return a :class:`Mesh` if its var_name matches this value.
+
+    Returns
+    -------
+    :class:`Mesh`
+
+    """
+    meshes_result = load_meshes(uris, var_name)
+    result = [mesh for file in meshes_result.values() for mesh in file]
+    mesh_count = len(result)
+    if mesh_count != 1:
+        message = (
+            f"Expecting 1 mesh, but input file(s) produced: {mesh_count} ."
+        )
+        raise ConstraintMismatchError(message)
+    return result[0]
+
+
+def load_meshes(uris, var_name=None):
+    """
+    Create :class:`Mesh` objects from one or more NetCDF files.
+
+    Parameters
+    ----------
+    uris : str or iterable of str
+        One or more filenames/URI's. Any URI's must support OpenDAP.
+    var_name : str, optional
+        Only return a :class:`Mesh` if its var_name matches this value.
+
+    Returns
+    -------
+    dict
+        A dictionary of file paths/URL's and lists of the :class:`Mesh`es
+        returned from each.
+
+    """
+    # No constraints or callbacks supported - these assume they are operating
+    #  on a Cube.
+
+    from iris.fileformats import FORMAT_AGENT
+
+    # TODO: rationalise UGRID/mesh handling once experimental.ugrid is folded
+    #  into standard behaviour.
+
+    if not PARSE_UGRID_ON_LOAD:
+        # Explicit behaviour, consistent with netcdf.load_cubes(), rather than
+        #  an invisible assumption.
+        message = (
+            f"PARSE_UGRID_ON_LOAD is {bool(PARSE_UGRID_ON_LOAD)}. Must be "
+            f"True to enable mesh loading."
+        )
+        raise ValueError(message)
+
+    if isinstance(uris, str):
+        uris = [uris]
+
+    # Group collections of uris by their iris handler
+    # Create list of tuples relating schemes to part names.
+    uri_tuples = sorted(decode_uri(uri) for uri in uris)
+
+    valid_sources = []
+    for scheme, groups in groupby(uri_tuples, key=lambda x: x[0]):
+        # Call each scheme handler with the appropriate URIs
+        if scheme == "file":
+            filenames = [x[1] for x in groups]
+            sources = expand_filespecs(filenames)
+        elif scheme in ["http", "https"]:
+            sources = [":".join(x) for x in groups]
+        else:
+            message = f"Iris cannot handle the URI scheme: {scheme}"
+            raise ValueError(message)
+
+        for source in sources:
+            if scheme == "file":
+                with open(source, "rb") as fh:
+                    handling_format_spec = FORMAT_AGENT.get_spec(
+                        Path(source).name, fh
+                    )
+            else:
+                handling_format_spec = FORMAT_AGENT.get_spec(source, None)
+
+            if handling_format_spec.handler == netcdf.load_cubes:
+                valid_sources.append(source)
+            else:
+                message = f"Ignoring non-NetCDF file: {source}"
+                logger.info(msg=message, extra=dict(cls=None))
+
+    result = {}
+    for source in valid_sources:
+        meshes_dict = _meshes_from_cf(CFUGridReader(source))
+        meshes = list(meshes_dict.values())
+        if var_name is not None:
+            meshes = list(filter(lambda m: m.var_name == var_name, meshes))
+        if meshes:
+            result[source] = meshes
+
+    return result
+
+
 ############
 # CF Overrides.
 # These are not included in __all__ since they are not [currently] needed
@@ -3469,7 +3606,17 @@ def identify(cls, variables, ignore=None, target=None, warn=True):
         log_level = logging.WARNING if warn else logging.DEBUG
 
         # Identify all CF-UGRID mesh variables.
+        all_vars = target == variables
         for nc_var_name, nc_var in target.items():
+            if all_vars:
+                # SPECIAL BEHAVIOUR FOR MESH VARIABLES.
+                # We are looking for all mesh variables. Check if THIS variable
+                #  is a mesh using its own attributes.
+                if getattr(nc_var, "cf_role", "") == "mesh_topology":
+                    result[nc_var_name] = CFUGridMeshVariable(
+                        nc_var_name, nc_var
+                    )
+
             # Check for mesh variable references.
             nc_var_att = getattr(nc_var, cls.cf_identity, None)
 

lib/iris/fileformats/cf.py

@@ -1067,6 +1067,11 @@ def __init__(self, filename, warn=False, monotonic=False):
         self._build_cf_groups()
         self._reset()
 
+    @property
+    def filename(self):
+        """The file that the CFReader is reading."""
+        return self._filename
+
     def __repr__(self):
         return "%s(%r)" % (self.__class__.__name__, self._filename)
 

lib/iris/fileformats/netcdf.py

@@ -792,8 +792,8 @@ def load_cubes(filenames, callback=None):
     from iris.experimental.ugrid import (
         PARSE_UGRID_ON_LOAD,
         CFUGridReader,
-        _build_mesh,
         _build_mesh_coords,
+        _meshes_from_cf,
     )
     from iris.io import run_callback
 
@@ -808,15 +808,7 @@ def load_cubes(filenames, callback=None):
         meshes = {}
         if PARSE_UGRID_ON_LOAD:
             cf = CFUGridReader(filename)
-
-            # Mesh instances are shared between file phenomena.
-            # TODO: more sophisticated Mesh sharing between files.
-            # TODO: access external Mesh cache?
-            mesh_vars = cf.cf_group.meshes
-            meshes = {
-                name: _build_mesh(cf, var, filename)
-                for name, var in mesh_vars.items()
-            }
+            meshes = _meshes_from_cf(cf)
         else:
             cf = iris.fileformats.cf.CFReader(filename)
 

lib/iris/tests/unit/experimental/ugrid/test_CFUGridMeshVariable.py

@@ -31,6 +31,22 @@ class TestIdentify(tests.IrisTest):
     def setUp(self):
         self.cf_identity = "mesh"
 
+    def test_cf_role(self):
+        match_name = "match"
+        match = named_variable(match_name)
+        setattr(match, "cf_role", "mesh_topology")
+
+        not_match_name = f"not_{match_name}"
+        not_match = named_variable(not_match_name)
+        setattr(not_match, "cf_role", "foo")
+
+        vars_all = {match_name: match, not_match_name: not_match}
+
+        # ONLY expecting match, excluding not_match.
+        expected = {match_name: CFUGridMeshVariable(match_name, match)}
+        result = CFUGridMeshVariable.identify(vars_all)
+        self.assertDictEqual(expected, result)
+
     def test_cf_identity(self):
         subject_name = "ref_subject"
         ref_subject = named_variable(subject_name)
@@ -49,6 +65,31 @@ def test_cf_identity(self):
         result = CFUGridMeshVariable.identify(vars_all)
         self.assertDictEqual(expected, result)
 
+    def test_cf_role_and_identity(self):
+        role_match_name = "match"
+        role_match = named_variable(role_match_name)
+        setattr(role_match, "cf_role", "mesh_topology")
+
+        subject_name = "ref_subject"
+        ref_subject = named_variable(subject_name)
+        ref_source = named_variable("ref_source")
+        setattr(ref_source, self.cf_identity, subject_name)
+
+        vars_all = {
+            role_match_name: role_match,
+            subject_name: ref_subject,
+            "ref_not_subject": named_variable("ref_not_subject"),
+            "ref_source": ref_source,
+        }
+
+        # Expecting role_match and ref_subject but excluding other variables.
+        expected = {
+            role_match_name: CFUGridMeshVariable(role_match_name, role_match),
+            subject_name: CFUGridMeshVariable(subject_name, ref_subject),
+        }
+        result = CFUGridMeshVariable.identify(vars_all)
+        self.assertDictEqual(expected, result)
+
     def test_duplicate_refs(self):
         subject_name = "ref_subject"
         ref_subject = named_variable(subject_name)

lib/iris/tests/unit/experimental/ugrid/test_load_mesh.py

@@ -0,0 +1,57 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the LGPL license.
+# See COPYING and COPYING.LESSER in the root of the repository for full
+# licensing details.
+"""
+Unit tests for the :func:`iris.experimental.ugrid.load_mesh` function.
+
+"""
+
+# Import iris.tests first so that some things can be initialised before
+# importing anything else.
+import iris.tests as tests  # isort:skip
+
+from unittest import mock
+
+from iris.exceptions import ConstraintMismatchError
+from iris.experimental.ugrid import PARSE_UGRID_ON_LOAD, load_mesh
+
+
+class Tests(tests.IrisTest):
+    # All 'real' tests have been done for load_meshes(). Here we just check
+    #  that load_mesh() works with load_meshes() correctly, using mocking.
+    def setUp(self):
+        patcher = mock.patch("iris.experimental.ugrid.load_meshes")
+        self.addCleanup(patcher.stop)
+
+        self.load_meshes_mock = patcher.start()
+        # The expected return from load_meshes - a dict of files, each with
+        #  a list of meshes.
+        self.load_meshes_mock.return_value = {"file": ["mesh"]}
+
+    def test_calls_load_meshes(self):
+        args = [("file_1", "file_2"), "my_var_name"]
+        with PARSE_UGRID_ON_LOAD.context():
+            _ = load_mesh(args)
+        self.assertTrue(self.load_meshes_mock.called_with(args))
+
+    def test_returns_mesh(self):
+        with PARSE_UGRID_ON_LOAD.context():
+            mesh = load_mesh([])
+        self.assertEqual(mesh, "mesh")
+
+    def test_single_mesh(self):
+        # Override the load_meshes_mock return values to provoke errors.
+        def common(ret_val):
+            self.load_meshes_mock.return_value = ret_val
+            with self.assertRaisesRegex(
+                ConstraintMismatchError, "Expecting 1 mesh.*"
+            ):
+                with PARSE_UGRID_ON_LOAD.context():
+                    _ = load_mesh([])
+
+        # Too many.
+        common({"file": ["mesh1", "mesh2"]})
+        # Too few.
+        common({"file": []})