dask: `Data.convolution_filter` by davidhassell · Pull Request #294 · NCAS-CMS/cf-python

Expand Up @@ -8,11 +8,6 @@ from numbers import Integral from operator import mul
try: from scipy.ndimage.filters import convolve1d as scipy_convolve1d except ImportError: pass
import cfdm import cftime import dask.array as da Expand Down Expand Up @@ -2855,6 +2850,7 @@ def ceil(self, inplace=False, i=False): """ return self.func(np.ceil, out=True, inplace=inplace)
@daskified(1) @_inplace_enabled(default=False) def convolution_filter( self, Expand Down Expand Up @@ -2935,6 +2931,9 @@ def convolution_filter( ``'wrap'`` The input is extended by ``(a b c | a b c | a b c)`` wrapping around to the opposite edge.
``'periodic'`` This is a synonym for ``'wrap'``. ============== ========================== ============================
The position of the window relative to each value can be Expand Down Expand Up @@ -2973,75 +2972,80 @@ def convolution_filter( in-place.
""" # TODODSAK - map_overlap
try: scipy_convolve1d except NameError: raise ImportError( "Must install scipy to use the convolution_filter method." ) from .dask_utils import cf_convolve1d
d = _inplace_enabled_define_and_cleanup(self)
iaxis = d._parse_axes(axis) if len(iaxis) != 1: raise ValueError( "Must specify a unique domain axis with the 'axis' " "parameter. {!r} specifies axes {!r}".format(axis, iaxis) f"parameter. {axis!r} specifies axes {iaxis!r}" )
iaxis = iaxis[0]
# Default mode to 'wrap' if the axis is cyclic if mode is None: # Default mode is 'wrap' if the axis is cyclic, or else # 'constant'. if iaxis in d.cyclic(): mode = "wrap" boundary = "periodic" else: mode = "constant" # --- End: if boundary = cval elif mode == "wrap": boundary = "periodic" elif mode == "constant": boundary = cval elif mode == "mirror": raise ValueError( "'mirror' mode is no longer available. Please raise an " "issue at https://github.com/NCAS-CMS/cf-python/issues " "if you would like it to be re-implemented." ) # This re-implementation would involve getting a 'mirror' # function added to dask.array.overlap, along similar # lines to the existing 'reflect' function in that module. else: boundary = mode
# Set cval to NaN if it is currently None, so that the edges # will be filled with missing data if the mode is 'constant' if cval is None: cval = np.nan # Set the overlap depth large enough to accommodate the # filter. # # For instance, for a 5-point window, the calculated value at # each point requires 2 points either side if the filter is # centred (i.e. origin is 0) and (up to) 3 points either side # if origin is 1 or -1. # # It is a restriction of dask.array.map_overlap that we can't # use asymmetric halos for general 'boundary' types. size = len(window) depth = int(size / 2) if not origin and not size % 2: depth += 1
# Section the data into sections up to a chunk in size sections = self.section([iaxis], chunks=True) depth += abs(origin)
# Filter each section replacing masked points with numpy # NaNs and then remasking after filtering. for key, data in sections.items(): data.dtype = float input_array = data.array masked = np.ma.is_masked(input_array) if masked: input_array = input_array.filled(np.nan)
output_array = scipy_convolve1d( input_array, window, axis=iaxis, mode=mode, cval=cval, origin=origin, ) if masked or (mode == "constant" and np.isnan(cval)): with np.errstate(invalid="ignore"): output_array = np.ma.masked_invalid(output_array) # --- End: if dx = d._get_dask()
sections[key] = type(self)( output_array, units=self.Units, fill_value=self.fill_value ) # Cast to float to ensure that NaNs can be stored (as required # by cf_convolve1d) if dx.dtype != float: dx = dx.astype(float, copy=False)
# Glue the sections back together again out = self.reconstruct_sectioned_data(sections, cyclic=self.cyclic()) # Convolve each chunk convolve1d = partial( cf_convolve1d, window=window, axis=iaxis, origin=origin )
if inplace: d.__dict__ = out.__dict__ else: d = out dx = dx.map_overlap( convolve1d, depth={iaxis: depth}, boundary=boundary, trim=True, meta=np.array((), dtype=float), )
d._set_dask(dx, reset_mask_hardness=True)
return d
Expand Down