"""Provides ESMF representations of grids/UGRID meshes and a modified regridder."""
import numpy as np
from numpy import ma
import scipy.sparse
import esmf_regrid
from esmf_regrid import check_method, check_norm, Constants
from . import esmpy
from ._esmf_sdo import GridInfo, RefinedGridInfo
__all__ = [
"GridInfo",
"RefinedGridInfo",
"Regridder",
]
def _get_regrid_weights_dict(src_field, tgt_field, regrid_method):
# The value, in array form, that ESMF should treat as an affirmative mask.
expected_mask = np.array([True])
regridder = esmpy.Regrid(
src_field,
tgt_field,
ignore_degenerate=True,
regrid_method=regrid_method,
unmapped_action=esmpy.UnmappedAction.IGNORE,
# Choosing the norm_type DSTAREA allows for mdtol type operations
# to be performed using the weights information later on.
norm_type=esmpy.NormType.DSTAREA,
src_mask_values=expected_mask,
dst_mask_values=expected_mask,
factors=True,
)
# Without specifying deep_copy=true, the information in weights_dict
# would be corrupted when the ESMF regridder is destoyed.
weights_dict = regridder.get_weights_dict(deep_copy=True)
# The weights_dict contains all the information needed for regridding,
# the ESMF objects can be safely removed.
regridder.destroy()
return weights_dict
def _weights_dict_to_sparse_array(weights, shape, index_offsets):
matrix = scipy.sparse.csr_matrix(
(
weights["weights"],
(
weights["row_dst"] - index_offsets[0],
weights["col_src"] - index_offsets[1],
),
),
shape=shape,
)
return matrix
[docs]
class Regridder:
"""Regridder for directly interfacing with :mod:`esmpy`."""
def __init__(
self, src, tgt, method=Constants.Method.CONSERVATIVE, precomputed_weights=None
):
"""
Create a regridder from descriptions of horizontal grids/meshes.
Weights will be calculated using :mod:`esmpy` and stored as a
:class:`scipy.sparse.csr_matrix`
for use in regridding. If precomputed weights are provided,
these will be used instead of calculating via :mod:`esmpy`.
Parameters
----------
src : :class:`~esmf_regrid.experimental.unstructured_regrid.MeshInfo` or :class:`GridInfo`
Describes the source mesh/grid.
Data supplied to this regridder should be in a :class:`numpy.ndarray`
whose shape is compatible with ``src``.
tgt : :class:`~esmf_regrid.experimental.unstructured_regrid.MeshInfo` or :class:`GridInfo`
Describes the target mesh/grid.
Data output by this regridder will be a :class:`numpy.ndarray` whose
shape is compatible with ``tgt``.
method : :class:`Constants.Method`
The method to be used to calculate weights.
precomputed_weights : :class:`scipy.sparse.spmatrix`, optional
If ``None``, :mod:`esmpy` will be used to
calculate regridding weights. Otherwise, :mod:`esmpy` will be bypassed
and ``precomputed_weights`` will be used as the regridding weights.
"""
self.src = src
self.tgt = tgt
# type checks method
self.method = check_method(method)
self.esmf_regrid_version = esmf_regrid.__version__
if precomputed_weights is None:
self.esmf_version = esmpy.__version__
weights_dict = _get_regrid_weights_dict(
src.make_esmf_field(),
tgt.make_esmf_field(),
regrid_method=method.value,
)
self.weight_matrix = _weights_dict_to_sparse_array(
weights_dict,
(self.tgt._refined_size, self.src._refined_size),
(self.tgt.index_offset, self.src.index_offset),
)
if isinstance(tgt, RefinedGridInfo):
# At this point, the weight matrix represents more target points than
# tgt respresents. In order to collapse these points, we collapse the
# weights matrix by the appropriate matrix multiplication.
self.weight_matrix = (
tgt._collapse_weights(is_tgt=True) @ self.weight_matrix
)
if isinstance(src, RefinedGridInfo):
# At this point, the weight matrix represents more source points than
# src respresents. In order to collapse these points, we collapse the
# weights matrix by the appropriate matrix multiplication.
self.weight_matrix = self.weight_matrix @ src._collapse_weights(
is_tgt=False
)
else:
if not scipy.sparse.isspmatrix(precomputed_weights):
raise ValueError(
"Precomputed weights must be given as a sparse matrix."
)
if precomputed_weights.shape != (self.tgt.size, self.src.size):
msg = "Expected precomputed weights to have shape {}, got shape {} instead."
raise ValueError(
msg.format(
(self.tgt.size, self.src.size),
precomputed_weights.shape,
)
)
self.esmf_version = None
self.weight_matrix = precomputed_weights
[docs]
def regrid(self, src_array, norm_type=Constants.NormType.FRACAREA, mdtol=1):
"""
Perform regridding on an array of data.
Parameters
----------
src_array : :obj:`~numpy.typing.ArrayLike`
Array whose shape is compatible with ``self.src``
norm_type : :class:`Constants.NormType`
Either ``Constants.NormType.FRACAREA`` or ``Constants.NormType.DSTAREA``.
Determines the type of normalisation applied to the weights.
mdtol : float, default=1
A number between 0 and 1 describing the missing data tolerance.
Depending on the value of ``mdtol``, if a cell in the target grid is not
sufficiently covered by unmasked cells of the source grid, then it will
be masked. ``mdtol=1`` means that only target cells which are not
covered at all will be masked, ``mdtol=0`` means that all target
cells that are not entirely covered will be masked, and ``mdtol=0.5``
means that all target cells that are less than half covered will
be masked.
Returns
-------
:obj:`~numpy.typing.ArrayLike`
An array whose shape is compatible with ``self.tgt``.
"""
# Sets default value, as this can't be done with class attributes within method call
norm_type = check_norm(norm_type)
array_shape = src_array.shape
main_shape = array_shape[-self.src.dims :]
if main_shape != self.src.shape:
raise ValueError(
f"Expected an array whose shape ends in {self.src.shape}, "
f"got an array with shape ending in {main_shape}."
)
extra_shape = array_shape[: -self.src.dims]
extra_size = max(1, np.prod(extra_shape))
src_inverted_mask = self.src._array_to_matrix(~ma.getmaskarray(src_array))
weight_sums = self.weight_matrix @ src_inverted_mask
# Set the minimum mdtol to be slightly higher than 0 to account for rounding
# errors.
mdtol = max(mdtol, 1e-8)
tgt_mask = weight_sums > 1 - mdtol
masked_weight_sums = weight_sums * tgt_mask
normalisations = np.ones([self.tgt.size, extra_size])
if norm_type == Constants.NormType.FRACAREA:
normalisations[tgt_mask] /= masked_weight_sums[tgt_mask]
elif norm_type == Constants.NormType.DSTAREA:
pass
normalisations = ma.array(normalisations, mask=np.logical_not(tgt_mask))
flat_src = self.src._array_to_matrix(ma.filled(src_array, 0.0))
flat_tgt = self.weight_matrix @ flat_src
flat_tgt = flat_tgt * normalisations
tgt_array = self.tgt._matrix_to_array(flat_tgt, extra_shape)
return tgt_array