Source code for iris.experimental.ugrid.utils

# 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.

"""Utility operations specific to unstructured data."""

from typing import AnyStr, Iterable, Union

import dask.array as da
import numpy as np

from iris.cube import Cube


[docs] def recombine_submeshes( mesh_cube: Cube, submesh_cubes: Union[Iterable[Cube], Cube], index_coord_name: AnyStr = "i_mesh_index", ) -> Cube: """Put data from sub-meshes back onto the original full mesh. The result is a cube like ``mesh_cube``, but with its data replaced by a combination of the data in the ``submesh_cubes``. Parameters ---------- mesh_cube : Cube Describes the mesh and mesh-location onto which the all the ``submesh-cubes``' data are mapped, and acts as a template for the result. Must have a :class:`~iris.experimental.ugrid.mesh.Mesh`. submesh_cubes : iterable of Cube, or Cube Cubes, each with data on a _subset_ of the ``mesh_cube`` datapoints (within the mesh dimension). The submesh cubes do not need to have a mesh. There must be at least 1 of them, to determine the result phenomenon. Their metadata (names, units and attributes) must all be the same, _except_ that 'var_name' is ignored. Their dtypes must all be the same. Their shapes and dimension-coords must all match those of ``mesh_cube``, except in the mesh dimension, which can have different sizes between the submeshes, and from the ``mesh_cube``. The mesh dimension of each must have a 1-D coord named by ``index_coord_name``. These "index coords" can vary in length, but they must all have matching metadata (units, attributes and names except 'var_name'), and must also match the coord of that name in ``mesh_cube``, if there is one. The ".points" values of the index coords specify, for each datapoint, its location in the original mesh -- i.e. they are indices into the relevant mesh-location dimension. index_coord_name : Cube Coord name of an index coord containing the mesh location indices, in every submesh cube. Returns ------- result_cube A cube with the same mesh, location, and shape as ``mesh_cube``, but with its data replaced by that from the``submesh_cubes``. The result phenomeon identity is also that of the``submesh_cubes``, i.e. units, attributes and names (except 'var_name', which is None). Notes ----- Where regions overlap, the result data comes from the submesh cube containing that location which appears _last_ in ``submesh_cubes``. Where no region contains a datapoint, it will be masked in the result. HINT: alternatively, values covered by no region can be set to the original 'full_mesh_cube' data value, if 'full_mesh_cube' is *also* passed as the first of the 'region_cubes'. The ``result_cube`` dtype is that of the ``submesh_cubes``. """ if not submesh_cubes: raise ValueError("'submesh_cubes' must be non-empty.") mesh_dim = mesh_cube.mesh_dim() if mesh_dim is None: raise ValueError("'mesh_cube' has no \".mesh\".") # # Perform consistency checks on all the region-cubes. # if not isinstance(submesh_cubes, Iterable): # Treat a single submesh cube input as a list-of-1. submesh_cubes = [submesh_cubes] result_metadata = None result_dtype = None indexcoord_metadata = None for i_sub, cube in enumerate(submesh_cubes): sub_str = f"Submesh cube #{i_sub + 1}/{len(submesh_cubes)}, " f'"{cube.name()}"' # Check dimensionality. if cube.ndim != mesh_cube.ndim: err = ( f"{sub_str} has {cube.ndim} dimensions, but " f"'mesh_cube' has {mesh_cube.ndim}." ) raise ValueError(err) # Get cube metadata + dtype : must match, and will apply to the result dtype = cube.dtype metadata = cube.metadata._replace(var_name=None) if i_sub == 0: # Store the first-cube metadata + dtype as reference result_metadata = metadata result_dtype = dtype else: # Check subsequent region-cubes metadata + dtype against the first if metadata != result_metadata: err = ( f"{sub_str} has metadata {metadata}, " "which does not match that of the other region_cubes, " f"which is {result_metadata}." ) raise ValueError(err) elif dtype != result_dtype: err = ( f"{sub_str} has a dtype of {dtype}, " "which does not match that of the other region_cubes, " f"which is {result_dtype}." ) raise ValueError(err) # For each dim, check that coords match other regions, and full-cube for i_dim in range(mesh_cube.ndim): if i_dim == mesh_dim: # mesh dim : look for index coords (by name) full_coord = mesh_cube.coords( name_or_coord=index_coord_name, dimensions=(i_dim,) ) sub_coord = cube.coords( name_or_coord=index_coord_name, dimensions=(i_dim,) ) else: # non-mesh dims : look for dim-coords (only) full_coord = mesh_cube.coords(dim_coords=True, dimensions=(i_dim,)) sub_coord = cube.coords(dim_coords=True, dimensions=(i_dim,)) if full_coord: (full_coord,) = full_coord full_dimname = full_coord.name() full_metadata = full_coord.metadata._replace(var_name=None) if sub_coord: (sub_coord,) = sub_coord sub_dimname = sub_coord.name() sub_metadata = sub_coord.metadata._replace(var_name=None) err = None # N.B. checks for mesh- and non-mesh-dims are different if i_dim != mesh_dim: # i_dim == mesh_dim : checks for non-mesh dims if full_coord and not sub_coord: err = ( f"{sub_str} has no dim-coord for dimension " f"{i_dim}, to match the 'mesh_cube' dimension " f'"{full_dimname}".' ) elif sub_coord and not full_coord: err = ( f'{sub_str} has a dim-coord "{sub_dimname}" for ' f"dimension {i_dim}, but 'mesh_cube' has none." ) elif sub_coord != full_coord: err = ( f'{sub_str} has a dim-coord "{sub_dimname}" for ' f"dimension {i_dim}, which does not match that " f"of 'mesh_cube', \"{full_dimname}\"." ) else: # i_dim == mesh_dim : different rules for this one if not sub_coord: # Must have an index coord on the mesh dimension err = ( f'{sub_str} has no "{index_coord_name}" coord on ' f"the mesh dimension (dimension {mesh_dim})." ) elif full_coord and sub_metadata != full_metadata: # May *not* have full-cube index, but if so it must match err = ( f"{sub_str} has an index coord " f'"{index_coord_name}" whose ".metadata" does not ' f"match that of the same name in 'mesh_cube' : " f"{sub_metadata} != {full_metadata}." ) else: # At this point, we know we *have* an index coord, and it does # not conflict with the one on 'mesh_cube' (if any). # Now check for matches between the region cubes. if indexcoord_metadata is None: # Store first occurrence (from first region-cube) indexcoord_metadata = sub_metadata elif sub_metadata != indexcoord_metadata: # Compare subsequent occurrences (from other region-cubes) err = ( f"{sub_str} has an index coord " f'"{index_coord_name}" whose ".metadata" does not ' f"match that of the other submesh-cubes : " f"{sub_metadata} != {indexcoord_metadata}." ) if err: raise ValueError(err) # Use the mesh_dim to transpose inputs + outputs, if required, as it is # simpler for all the array operations to always have the mesh dim *last*. if mesh_dim == mesh_cube.ndim - 1: # Mesh dim is already the last one : no transpose required untranspose_dims = None else: dim_range = np.arange(mesh_cube.ndim, dtype=int) # Transpose all inputs to mesh-last order tranpose_dims = [i_dim for i_dim in dim_range if i_dim != mesh_dim] + [ mesh_dim ] # chop out mesh_dim + put it at the end def transposed_copy(cube, dim_order): cube = cube.copy() cube.transpose(dim_order) return cube mesh_cube = transposed_copy(mesh_cube, tranpose_dims) submesh_cubes = [ transposed_copy(region_cube, tranpose_dims) for region_cube in submesh_cubes ] # Also prepare for transforming the output back to the original order untranspose_dims = dim_range.copy() # Neat trick to produce the reverse operation untranspose_dims[tranpose_dims] = dim_range # # Here's the core operation.. # def fill_region(target, regiondata, regioninds): if not target.flags.writeable: # The initial input can be a section of a da.zeros(), which has no # real array "behind" it. This means that real arrays created in # memory are only chunk-sized, but it also means that 'target' may # not be writeable. So take a copy to fix that, where needed. target = target.copy() # N.B. Indices are basically 1D, but may have leading *1 dims for # alignment, to satisfy da.map_blocks assert all(size == 1 for size in regioninds.shape[:-1]) inds = regioninds.flatten() # Assign blocks with indexing on the last dim only target[..., inds] = regiondata return target # Create an initially 'empty' (all-masked) dask array matching the input. # N.B. this does not use the mesh_cube.lazy_data() array, but only its # shape and dtype, since the data itself is not used in the calculation. # N.B. chunking matches the input cube, allowing performance control. input_data = mesh_cube.lazy_data() result_array = da.ma.masked_array( da.zeros( input_data.shape, dtype=result_dtype, chunks=input_data.chunksize, ), True, ) # Wrap this repeatedly with a lazy operation to assign each region. # It is done this way because we couldn't get map_blocks to correctly wrap # a function which does all regions in a single operation. # TODO: replace with a single-stage solution: Probably better, if possible. # Notes on resultant calculation properties: # 1. map_blocks is chunk-mapped, so it is parallelisable and space-saving # 2. However, fetching less than a whole chunk is not efficient for cube in submesh_cubes: # Lazy data array from the region cube sub_data = cube.lazy_data() # Lazy indices from the mesh-dim coord mesh_dimcoord = cube.coord( name_or_coord=index_coord_name, dimensions=cube.ndim - 1 ) indarr = mesh_dimcoord.lazy_points() # Extend indarr dimensions to align it with the 'target' array dims assert indarr.ndim == 1 shape = (1,) * (cube.ndim - 1) + indarr.shape indarr = indarr.reshape(shape) # Apply the operation to paste from one region into the target # N.B. replacing 'result_array' each time around the loop result_array = da.map_blocks( fill_region, result_array, sub_data, indarr, dtype=result_array.dtype, meta=np.ndarray, ) # Construct the result cube result_cube = mesh_cube.copy() result_cube.data = result_array # Copy names, units + attributes from region data (N.B. but not var_name) result_cube.metadata = result_metadata if untranspose_dims is not None: # Re-order dims as in the original input result_cube.transpose(untranspose_dims) return result_cube