Speedup area weighted regridding

lib/iris/experimental/regrid.py

@@ -18,7 +18,6 @@
 Regridding functions.
 
 """
-
 from __future__ import (absolute_import, division, print_function)
 from six.moves import (filter, input, map, range, zip)  # noqa
 import six
@@ -350,7 +349,7 @@ def _get_bounds_in_units(coord, units, dtype):
     return coord.units.convert(coord.bounds.astype(dtype), units).astype(dtype)
 
 
-def _weighted_mean_with_mdtol(data, weights, axis=None, mdtol=0):
+def _weighted_mean_with_mdtol_old(data, weights, axis=None, mdtol=0):
     """
     Return the weighted mean of an array over the specified axis
     using the provided weights (if any) and a permitted fraction of
@@ -402,6 +401,62 @@ def _weighted_mean_with_mdtol(data, weights, axis=None, mdtol=0):
     return res
 
 
+def _weighted_mean_with_mdtol_new(data, weights, axis=None, mdtol=0):
+    """
+    Return the weighted mean of an array over the specified axis
+    using the provided weights (if any) and a permitted fraction of
+    masked data.
+
+    Args:
+
+    * data (array-like):
+        Data to be averaged.
+
+    * weights (array-like):
+        An array of the same shape as the data that specifies the contribution
+        of each corresponding data element to the calculated mean.
+
+    Kwargs:
+
+    * axis (int or tuple of ints):
+        Axis along which the mean is computed. The default is to compute
+        the mean of the flattened array.
+
+    * mdtol (float):
+        Tolerance of missing data. The value returned in each element of the
+        returned array will be masked if the fraction of masked data exceeds
+        mdtol. This fraction is weighted by the `weights` array if one is
+        provided. mdtol=0 means no missing data is tolerated
+        while mdtol=1 will mean the resulting element will be masked if and
+        only if all the contributing elements of data are masked.
+        Defaults to 0.
+
+    Returns:
+        Numpy array (possibly masked) or scalar.
+
+    """
+    if ma.is_masked(data):
+        res, unmasked_weights_sum = ma.average(data, weights=weights,
+                                               axis=axis, returned=True)
+        if mdtol < 1:
+            weights_sum = weights.sum(axis=axis)
+            frac_masked = 1 - np.true_divide(unmasked_weights_sum, weights_sum)
+            mask_pt = frac_masked > mdtol
+            if np.any(mask_pt):
+                if np.isscalar(res):
+                    res = ma.masked
+                elif ma.isMaskedArray(res):
+                    res.mask |= mask_pt
+                else:
+                    res = ma.masked_array(res, mask=mask_pt)
+    else:
+        res = np.average(data, weights=weights, axis=axis)
+    return res
+
+
+_weighted_mean_with_mdtol = _weighted_mean_with_mdtol_new
+
+
 def _regrid_area_weighted_array(src_data, x_dim, y_dim,
                                 src_x_bounds, src_y_bounds,
                                 grid_x_bounds, grid_y_bounds,