diff --git a/lib/iris/analysis/_regrid.py b/lib/iris/analysis/_regrid.py index 1802d70f86..195aadd7ef 100644 --- a/lib/iris/analysis/_regrid.py +++ b/lib/iris/analysis/_regrid.py @@ -288,13 +288,15 @@ def _regrid(src_data, x_dim, y_dim, interp_coords = np.dstack(interp_coords) + weights = interpolator.compute_interp_weights(interp_coords) + def interpolate(data): # Update the interpolator for this data slice. data = data.astype(interpolator.values.dtype) if y_dim < x_dim: data = data.T interpolator.values = data - data = interpolator(interp_coords) + data = interpolator.interp_using_pre_computed_weights(weights) if y_dim > x_dim: data = data.T return data diff --git a/lib/iris/analysis/_scipy_interpolate.py b/lib/iris/analysis/_scipy_interpolate.py index 5bd663a7aa..26bc7aeef1 100644 --- a/lib/iris/analysis/_scipy_interpolate.py +++ b/lib/iris/analysis/_scipy_interpolate.py @@ -1,9 +1,17 @@ - from __future__ import (absolute_import, division, print_function) +import itertools +import time + +from scipy.sparse import csc_matrix +from scipy.sparse import csr_matrix import numpy as np +SPARSE = 1 +TIMINGS = 0 + + # ============================================================================ # | Copyright SciPy | # | Code from this point unto the termination banner is copyright SciPy. | @@ -137,17 +145,41 @@ def __call__(self, xi, method=None): Parameters ---------- xi : ndarray of shape (..., ndim) - The coordinates to sample the gridded data at + The coordinates to sample the gridded data at. method : str The method of interpolation to perform. Supported are "linear" and "nearest". """ - method = self.method if method is None else method - if method not in ["linear", "nearest"]: - raise ValueError("Method '%s' is not defined" % method) + # Note: No functionality should live in this method. It should all be + # decomposed into the two interfaces (compute weights + use weights). + weights = self.compute_interp_weights(xi, method) + return self.interp_using_pre_computed_weights(weights) + + def compute_interp_weights(self, xi, method=None): + """ + Prepare the interpolator for interpolation to the given sample points. + .. note:: + For normal interpolation, simply call the interpolator with the + sample points, rather using this decomposed interface. + + Parameters + ---------- + xi : ndarray of shape (..., ndim) + The coordinates to sample the gridded data at. + + Returns + ------- + An the items necessary for passing to + :meth:`interp_using_pre_computed_weights`. The contents of this return + value are not guaranteed to be consistent across Iris versions, and + should only be used for passing to + :meth:`interp_using_pre_computed_weights`. + + """ + t0 = time.time() ndim = len(self.grid) xi = _ndim_coords_from_arrays(xi, ndim=ndim) if xi.shape[-1] != len(self.grid): @@ -165,10 +197,76 @@ def __call__(self, xi, method=None): raise ValueError("One of the requested xi is out of " "bounds in dimension %d" % i) - indices, norm_distances, out_of_bounds = self._find_indices(xi.T) + prepared = (xi_shape, method) + self._find_indices(xi.T) + if TIMINGS: + print('Prepare:', time.time() - t0) + + if SPARSE: + t0 = time.time() + + xi_shape, method, indices, norm_distances, out_of_bounds = prepared + + # Allocate arrays for describing the sparse matrix. + n_src_values_per_result_value = 2 ** len(self.grid) + n_result_values = len(indices[0]) + n_non_zero = n_result_values * n_src_values_per_result_value + weights = np.ones(n_non_zero, dtype=norm_distances[0].dtype) + col_indices = np.empty(n_non_zero) + row_ptrs = np.arange(0, n_non_zero + n_src_values_per_result_value, + n_src_values_per_result_value) + + corners = itertools.product(*[[(i, 1 - n), (i + 1, n)] + for i, n in zip(indices, + norm_distances)]) + for i, corner in enumerate(corners): + corner_indices = [ci for ci, cw in corner] + n_indices = np.ravel_multi_index(corner_indices, + self.values.shape) + col_indices[i::n_src_values_per_result_value] = n_indices + for ci, cw in corner: + weights[i::n_src_values_per_result_value] *= cw + + n_src_values = np.prod(map(len, self.grid)) + sparse_matrix = csr_matrix((weights, col_indices, row_ptrs), + shape=(n_result_values, n_src_values)) + if TIMINGS: + print('Convert:', time.time() - t0) + + # XXX Hacky-Mc-Hack-Hack + prepared = (xi_shape, method, sparse_matrix, None, out_of_bounds) + + return prepared + + def interp_using_pre_computed_weights(self, computed_weights): + """ + Do the interpolation, given pre-computed interpolation weights. + + .. note:: + For normal interpolation, simply call the interpolator with the + sample points, rather using this decomposed interface. + + Parameters + ---------- + computed_weights : *intentionally undefined interface* + The pre-computed interpolation weights which come from calling + :meth:`compute_interp_weights`. + + """ + [xi_shape, method, indices, + norm_distances, out_of_bounds] = computed_weights + + method = self.method if method is None else method + if method not in ["linear", "nearest"]: + raise ValueError("Method '%s' is not defined" % method) + + ndim = len(self.grid) + if method == "linear": - result = self._evaluate_linear( - indices, norm_distances, out_of_bounds) + if SPARSE: + result = self._evaluate_linear_sparse(indices) + else: + result = self._evaluate_linear( + indices, norm_distances, out_of_bounds) elif method == "nearest": result = self._evaluate_nearest( indices, norm_distances, out_of_bounds) @@ -177,7 +275,17 @@ def __call__(self, xi, method=None): return result.reshape(xi_shape[:-1] + self.values.shape[ndim:]) + def _evaluate_linear_sparse(self, sparse_matrix): + t0 = time.time() + # TODO: Consider handling of non-grid dimensions. + result = sparse_matrix * self.values.reshape(-1) + if TIMINGS: + print('Apply:', time.time() - t0) + + return result + def _evaluate_linear(self, indices, norm_distances, out_of_bounds): + t0 = time.time() # slice for broadcasting over trailing dimensions in self.values vslice = (slice(None),) + (np.newaxis,) * \ (self.values.ndim - len(indices)) @@ -190,8 +298,14 @@ def _evaluate_linear(self, indices, norm_distances, out_of_bounds): for edge_indices in edges: weight = 1. for ei, i, yi in zip(edge_indices, indices, norm_distances): - weight *= np.where(ei == i, 1 - yi, yi) + #weight *= np.where(ei == i, 1 - yi, yi) + if ei[0] == i[0]: + weight *= 1 - yi + else: + weight *= yi values += np.asarray(self.values[edge_indices]) * weight[vslice] + if TIMINGS: + print('Apply:', time.time() - t0) return values def _evaluate_nearest(self, indices, norm_distances, out_of_bounds): diff --git a/lib/iris/coord_systems.py b/lib/iris/coord_systems.py index 9dc8613855..f6a9cbddad 100644 --- a/lib/iris/coord_systems.py +++ b/lib/iris/coord_systems.py @@ -798,3 +798,73 @@ def as_cartopy_crs(self): def as_cartopy_projection(self): return self.as_cartopy_crs() + + +class LambertAzimuthalEqualArea(CoordSystem): + """ + A coordinate system in the Lambert Azimuthal Equal-Area projection. + + """ + + grid_mapping_name = "lambert_azimuthal_equal_area" + + def __init__(self, central_lat=0.0, central_lon=0.0, + false_easting=0.0, false_northing=0.0, + ellipsoid=None): + """ + Constructs a Lambert Azimuthal Equal Area coord system. + + Args: + + * central_lat + The central latitude. + + * central_lon + The central longitude. + + * false_easting + X offset from planar origin in metres. + + * false_northing + Y offset from planar origin in metres. + + Kwargs: + + * ellipsoid + :class:`GeogCS` defining the ellipsoid. + + """ + + #: True latitude of planar origin in degrees. + self.central_lat = central_lat + #: True longitude of planar origin in degrees. + self.central_lon = central_lon + #: X offset from planar origin in metres. + self.false_easting = false_easting + #: Y offset from planar origin in metres. + self.false_northing = false_northing + #: Ellipsoid definition. + self.ellipsoid = ellipsoid + + def __repr__(self): + return "LambertAzimuthalEqualArea(central_lat={!r}, central_lon={!r}, "\ + "false_easting={!r}, false_northing={!r}, "\ + "ellipsoid={!r})".format(self.central_lat, self.central_lon, + self.false_easting, + self.false_northing, + self.ellipsoid) + + def as_cartopy_crs(self): + if self.ellipsoid is not None: + globe = self.ellipsoid.as_cartopy_globe() + else: + globe = ccrs.Globe() + return ccrs.LambertAzimuthalEqualArea( + central_longitude=self.central_lon, + central_latitude=self.central_lat, + false_easting=self.false_easting, + false_northing=self.false_northing, + globe=globe) + + def as_cartopy_projection(self): + return self.as_cartopy_crs()