Coverage for src/CSET/operators/improver/nbhood.py: 10%
162 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-10 13:27 +0000
1# (C) Crown Copyright, Met Office. All rights reserved.
2#
3# This file is part of 'IMPROVER' and is released under the BSD 3-Clause license.
4# See LICENSE in the root of the repository for full licensing details.
5"""Module containing plugin base class."""
7# This code is borrowed from https://github.com/metoppv/improver/ with some modifications.
9from abc import ABC, abstractmethod
10from collections.abc import Iterable
11from typing import Any, Optional, Tuple, Union
13import numpy as np
14from iris.cube import Cube, CubeList
15from numpy import ndarray
16from scipy.ndimage import correlate
19def pad_boxsum(
20 data: ndarray, boxsize: Union[int, Tuple[int, int]], **pad_options: Any
21) -> ndarray:
22 """Pad an array to shape suitable for `boxsum`.
24 Note that padding is not symmetric: there is an extra row/column at
25 the top/left (as required for calculating the boxsum).
27 Args:
28 data:
29 The input data array.
30 boxsize:
31 The size of the neighbourhood.
32 pad_options:
33 Additional keyword arguments passed to `numpy.pad` function.
35 Returns
36 -------
37 Array padded to shape suitable for `boxsum`.
38 """
39 boxsize = np.atleast_1d(boxsize)
40 ih, jh = boxsize[0] // 2, boxsize[-1] // 2
41 padding = [(0, 0)] * (data.ndim - 2) + [(ih + 1, ih), (jh + 1, jh)]
42 padded = np.pad(data, padding, **pad_options)
43 return padded
46def boxsum(
47 data: ndarray,
48 boxsize: Union[int, Tuple[int, int]],
49 cumsum: bool = True,
50 **pad_options: Any,
51) -> ndarray:
52 """Fast vectorised approach to calculating neighbourhood totals.
54 This function makes use of the summed-area table method. An input
55 array is accumulated top to bottom and left to right. This accumulated
56 array can then be used to efficiently calculate the total within a
57 neighbourhood about any point. An example input data array::
59 | 1 | 1 | 1 | 1 | 1 |
60 | 1 | 1 | 1 | 1 | 1 |
61 | 1 | 1 | 1 | 1 | 1 |
62 | 1 | 1 | 1 | 1 | 1 |
64 is accumulated to become::
66 | 1 | 2 | 3 | 4 | 5 |
67 | 2 | 4 | 6 | 8 | 10 |
68 | 3 | 6 | 9 | 12 | 15 |
69 | 4 | 8 | 12 | 16 | 20 |
70 | 5 | 10 | 15 | 20 | 25 |
72 If we wish to calculate the total in a 3x3 neighbourhood about
73 some point (*) of our array we use the following points::
75 | 1 (C) | 2 | 3 | 4 (D) | 5 |
76 | 2 | 4 | 6 | 8 | 10 |
77 | 3 | 6 | 9 (*) | 12 | 15 |
78 | 4 (A) | 8 | 12 | 16 (B) | 20 |
79 | 5 | 10 | 15 | 20 | 25 |
81 And the calculation is::
83 Neighbourhood sum = C - A - D + B
84 = 1 - 4 - 4 + 16
85 = 9
87 This is the value we would expect for a 3x3 neighbourhood
88 in an array filled with ones.
90 Args:
91 data:
92 The input data array.
93 boxsize:
94 The size of the neighbourhood. Must be an odd number.
95 cumsum:
96 If False, assume the input data is already cumulative. If True
97 (default), calculate cumsum along the last two dimensions of
98 the input array.
99 pad_options:
100 Additional keyword arguments passed to `numpy.pad` function.
101 If given, the returned result will have the same shape as the input
102 array.
104 Returns
105 -------
106 Array containing the calculated neighbourhood total.
108 Raises
109 ------
110 ValueError: If `boxsize` has non-integer type.
111 ValueError: If any member of `boxsize` is not an odd number.
112 """
113 boxsize = np.atleast_1d(boxsize)
114 if not issubclass(boxsize.dtype.type, np.integer):
115 raise ValueError("The size of the neighbourhood must be of an integer type.")
116 if not np.all(boxsize % 2):
117 raise ValueError("The size of the neighbourhood must be an odd number.")
118 if pad_options:
119 data = pad_boxsum(data, boxsize, **pad_options)
120 if cumsum:
121 data = data.cumsum(-2).cumsum(-1)
122 i, j = boxsize[0], boxsize[-1]
123 m, n = data.shape[-2] - i, data.shape[-1] - j
124 result = (
125 data[..., i : i + m, j : j + n]
126 - data[..., :m, j : j + n]
127 + data[..., :m, :n]
128 - data[..., i : i + m, :n]
129 )
130 return result
133class BasePlugin(ABC):
134 """An abstract class for IMPROVER plugins.
135 Subclasses must be callable. We preserve the process
136 method by redirecting to __call__.
137 """
139 def __call__(self, *args, **kwargs) -> Any:
140 """Makes subclasses callable to use process
141 Args:
142 *args:
143 Positional arguments.
144 **kwargs:
145 Keyword arguments.
147 Returns
148 -------
149 Output of self.process()
150 """
151 return self.process(*args, **kwargs)
153 @abstractmethod
154 def process(self, *args, **kwargs) -> Any:
155 """Abstract class for rest to implement."""
156 pass
159class PostProcessingPlugin(BasePlugin):
160 """An abstract class for IMPROVER post-processing plugins.
161 Makes generalised changes to metadata relating to post-processing.
162 """
164 def __call__(self, *args, **kwargs) -> Any:
165 """Makes subclasses callable to use process
166 Args:
167 *args:
168 Positional arguments.
169 **kwargs:
170 Keyword arguments.
172 Returns
173 -------
174 Output of self.process() with updated title attribute
175 """
176 result = super().__call__(*args, **kwargs)
177 if isinstance(result, Cube):
178 self.post_processed_title(result)
179 elif isinstance(result, Iterable) and not isinstance(result, str):
180 for item in result:
181 if isinstance(item, Cube):
182 self.post_processed_title(item)
183 return result
185 @staticmethod
186 def post_processed_title(cube):
187 """Updates title attribute on output cube to include
188 "Post-Processed".
190 """
191 MANDATORY_ATTRIBUTE_DEFAULTS = {
192 "title": "unknown",
193 "source": "IMPROVER",
194 "institution": "unknown",
195 }
197 default_title = MANDATORY_ATTRIBUTE_DEFAULTS["title"]
198 if (
199 "title" in cube.attributes.keys()
200 and cube.attributes["title"] != default_title
201 and "Post-Processed" not in cube.attributes["title"]
202 ):
203 title = cube.attributes["title"]
204 cube.attributes["title"] = f"Post-Processed {title}"
207def circular_kernel(ranges: int, weighted_mode: bool) -> ndarray:
208 """
209 Method to create a circular kernel.
211 Args:
212 ranges:
213 Number of grid cells in the x and y direction used to create
214 the kernel.
215 weighted_mode:
216 If True, use a circle for neighbourhood kernel with
217 weighting decreasing with radius.
218 If False, use a circle with constant weighting.
220 Returns
221 -------
222 Array containing the circular smoothing kernel.
223 This will have the same number of dimensions as fullranges.
224 """
225 # The range is square
227 area = ranges * ranges
228 # Define the size of the kernel based on the number of grid cells
229 # contained within the desired radius.
230 kernel = np.ones((int(1 + ranges * 2), (int(1 + ranges * 2))))
231 # Create an open multi-dimensional meshgrid.
232 open_grid = np.array(
233 np.ogrid[[slice(-x, x + 1) for x in (ranges, ranges)]], dtype=object
234 )
235 if weighted_mode:
236 # Create a kernel, such that the central grid point has the
237 # highest weighting, with the weighting decreasing with distance
238 # away from the central grid point.
239 open_grid_summed_squared = np.sum(open_grid**2.0).astype(float)
240 kernel[:] = (area - open_grid_summed_squared) / area
241 mask = kernel < 0.0
242 else:
243 mask = np.reshape(np.sum(open_grid**2) > area, np.shape(kernel))
244 kernel[mask] = 0.0
245 return kernel
248class NeighbourhoodProcessing(PostProcessingPlugin):
249 """Class for applying neighbourhood processing to produce a smoothed field
250 within the chosen neighbourhood.
251 """
253 def __init__(
254 self,
255 neighbourhood_method: str,
256 radii: float,
257 weighted_mode: bool = False,
258 sum_only: bool = False,
259 re_mask: bool = True,
260 ) -> None:
261 """
262 Initialise class.
264 Args:
265 neighbourhood_method:
266 Name of the neighbourhood method to use. Options: 'circular',
267 'square'.
268 radii:
269 The radii in grid points of the neighbourhood to apply.
270 Rounded up to convert into integer number of grid
271 points east and north, based on the characteristic spacing
272 at the zero indices of the cube projection-x and y coords.
273 weighted_mode:
274 If True, use a circle for neighbourhood kernel with
275 weighting decreasing with radius.
276 If False, use a circle with constant weighting.
277 sum_only:
278 If true, return neighbourhood sum instead of mean.
279 re_mask:
280 If re_mask is True, the original un-neighbourhood processed
281 mask is applied to mask out the neighbourhood processed cube.
282 If re_mask is False, the original un-neighbourhood processed
283 mask is not applied. Therefore, the neighbourhood processing
284 may result in values being present in areas that were
285 originally masked.
287 Raises
288 ------
289 ValueError: If the neighbourhood_method is not either
290 "square" or "circular".
291 ValueError: If the weighted_mode is used with a
292 neighbourhood_method that is not "circular".
293 """
294 self.radius = float(radii)
296 if neighbourhood_method in ["square", "circular"]:
297 self.neighbourhood_method = neighbourhood_method
298 else:
299 msg = "{} is not a valid neighbourhood_method.".format(neighbourhood_method)
300 raise ValueError(msg)
301 if weighted_mode and neighbourhood_method != "circular":
302 msg = (
303 "weighted_mode can only be used if neighbourhood_method is circular."
304 f" weighted_mode provided: {weighted_mode}, "
305 f"neighbourhood_method provided: {neighbourhood_method}."
306 )
307 raise ValueError(msg)
308 self.weighted_mode = weighted_mode
309 self.sum_only = sum_only
310 self.re_mask = re_mask
312 def _calculate_neighbourhood(
313 self, data: ndarray, mask: ndarray = None
314 ) -> Union[ndarray, np.ma.MaskedArray]:
315 """
316 Apply neighbourhood processing.
318 Ensures that masked data does not
319 contribute to the neighbourhood result. Masked data is either data that
320 is masked in the input data array or that corresponds to zeros in the
321 input mask.
323 Args:
324 data:
325 Input data array.
326 mask:
327 Mask of valid input data elements.
329 Returns
330 -------
331 Array containing the smoothed field after the
332 neighbourhood method has been applied.
333 """
334 if not self.sum_only:
335 min_val = np.nanmin(data)
336 max_val = np.nanmax(data)
338 # Data mask to be eventually used for re-masking.
339 # (This is OK even if mask is None, it gives a scalar False mask then.)
340 # Invalid data where the mask provided == 0.
341 data_mask = mask == 0
342 if isinstance(data, np.ma.MaskedArray):
343 # Include data mask if masked array.
344 data_mask = data_mask | data.mask
345 data = data.data
347 # Define working type and output type.
348 if issubclass(data.dtype.type, np.complexfloating):
349 loc_data_dtype = np.complex128
350 out_data_dtype = np.complex64
351 else:
352 # Use 64-bit types for enough precision in accumulations.
353 loc_data_dtype = np.float64
354 out_data_dtype = np.float32
355 data = np.array(data, dtype=loc_data_dtype)
357 # Replace invalid elements with zeros so they don't count towards
358 # neighbourhood sum
359 if self.neighbourhood_method == "circular":
360 mask_type = np.float32
361 else:
362 mask_type = np.int64
363 valid_data_mask = np.ones(data.shape, dtype=mask_type)
364 valid_data_mask[data_mask] = 0
365 data[data_mask] = 0
367 if self.sum_only:
368 max_extreme_data = None
369 else:
370 area_sum = self._do_nbhood_sum(valid_data_mask)
371 max_extreme_data = area_sum.astype(loc_data_dtype)
372 # Where data are all ones in nbhood, result will be same as area_sum
373 data = self._do_nbhood_sum(data, max_extreme=max_extreme_data)
375 if not self.sum_only:
376 with np.errstate(divide="ignore", invalid="ignore"):
377 # Calculate neighbourhood mean.
378 data = data / area_sum
379 # For points where all data in the neighbourhood is masked,
380 # set result to nan
381 data[area_sum == 0] = np.nan
382 data = data.clip(min_val, max_val)
384 if self.re_mask:
385 data = np.ma.masked_array(data, data_mask, copy=False)
387 return data.astype(out_data_dtype)
389 def _do_nbhood_sum(
390 self, data: np.ndarray, max_extreme: Optional[np.ndarray] = None
391 ) -> np.ndarray:
392 """Calculate the sum-in-area from an array.
394 As this can be expensive, the method first checks for the extreme cases where the data are:
395 All zeros (result will be all zeros too)
396 All ones (result will be max_extreme, if supplied)
397 Contains outer rows / columns that are completely zero or completely one, these
398 rows and columns are trimmed before calculating the area sum and their contents
399 will be as for the appropriate all case above.
401 Args:
402 data:
403 Input data array where any masking has already been replaced with zeroes.
404 max_extreme:
405 Used as the result for any large areas of data that are all ones, allowing an
406 optimisation to be used. If not supplied, the optimisation will only be used for
407 large areas of zeroes, where a return of zero can be safely predicted.
409 Returns
410 -------
411 Array containing the sum of data within the usable neighbourhood of each point.
412 """
413 # Determine the smallest box containing all non-zero or all non-one values with a
414 # neighbourhood-sized buffer and quit if there are none.
415 data_shape = data.shape
416 ystart = xstart = 0
417 ystop, xstop = data.shape
418 size = data.size
419 extreme = 0
420 fill_value = 0
421 half_nb_size = self.nb_size // 2
422 # For the two extreme values, 0 and 1, find the size and position of the smallest array
423 # that includes all other values with a buffer of the neighbourhood radius.
424 # The smallest box from either extreme will be passed to the neighbourhood method.
425 for _extreme, _fill_value in {0: 0, 1: max_extreme}.items():
426 if _fill_value is None or issubclass(data.dtype.type, np.complexfloating):
427 # We can't take this shortcut if we don't have either a default value/array,
428 # or the data values are complex, as comparisons with non-complex values are
429 # tricky.
430 continue
431 nonextreme_indices = np.argwhere(data != _extreme)
432 if nonextreme_indices.size == 0:
433 # No non-extreme values, so result will be _fill_value if set
434 _ystart = _ystop = _xstart = _xstop = 0
435 else:
436 (_ystart, _xstart), (_ystop, _xstop) = (
437 nonextreme_indices.min(0),
438 nonextreme_indices.max(0) + 1,
439 )
440 _ystart = max(0, _ystart - half_nb_size)
441 _ystop = min(data_shape[0], _ystop + half_nb_size)
442 _xstart = max(0, _xstart - half_nb_size)
443 _xstop = min(data_shape[1], _xstop + half_nb_size)
444 _size = (_ystop - _ystart) * (_xstop - _xstart)
445 if _size < size:
446 size, extreme, fill_value, ystart, ystop, xstart, xstop = (
447 _size,
448 _extreme,
449 _fill_value,
450 _ystart,
451 _ystop,
452 _xstart,
453 _xstop,
454 )
455 if size != data.size:
456 # If our chosen extreme allows us to process a subset of data, define the default array
457 # of neighbourhood sums that we know we will get for regions of extreme data values.
458 if isinstance(fill_value, np.ndarray):
459 untrimmed = fill_value.astype(data.dtype)
460 else:
461 untrimmed = np.full(data_shape, fill_value, dtype=data.dtype)
462 if size:
463 # The subset of data is non-zero in size, so calculate the neighbourhood sums in the
464 # subset.
465 data = data[ystart:ystop, xstart:xstop]
467 # Calculate neighbourhood totals for input data.
468 if self.neighbourhood_method == "square":
469 data = boxsum(
470 data, self.nb_size, mode="constant", constant_values=extreme
471 )
472 elif self.neighbourhood_method == "circular":
473 data = correlate(data, self.kernel, mode="nearest")
474 else:
475 data = untrimmed
477 # Expand data to the full size again
478 if data.shape != data_shape:
479 untrimmed[ystart:ystop, xstart:xstop] = data
480 data = untrimmed
481 return data
483 def process(self, cube: Cube, mask_cube: Optional[Cube] = None) -> Cube:
484 """
485 Call the methods required to apply a neighbourhood processing to a cube.
487 Applies neighbourhood processing to each 2D x-y-slice of the input cube.
489 If the input cube is masked the neighbourhood sum is calculated from
490 the total of the unmasked data in the neighbourhood around each grid
491 point. The neighbourhood mean is then calculated by dividing the
492 neighbourhood sum at each grid point by the total number of valid grid
493 points that contributed to that sum. If a mask_cube is provided then
494 this is used to mask each x-y-slice prior to the neighbourhood sum
495 or mean being calculated.
498 Args:
499 cube:
500 Cube containing the array to which the neighbourhood processing
501 will be applied. Usually thresholded data.
502 mask_cube:
503 Cube containing the array to be used as a mask. Zero values in
504 this array are taken as points to be masked.
506 Returns
507 -------
508 Cube containing the smoothed field after the
509 neighbourhood method has been applied.
510 """
511 if np.isnan(cube.data).any():
512 raise ValueError("Error: NaN detected in input cube data")
514 # check_if_grid_is_equal_area(cube)
516 # If the data is masked, the mask will be processed as well as the
517 # original_data * mask array.
518 # check_radius_against_distance(cube, self.radius)
520 # grid_cells = distance_to_number_of_grid_cells(cube, self.radius)
521 grid_cells = int(self.radius)
523 if self.neighbourhood_method == "circular":
524 self.kernel = circular_kernel(grid_cells, self.weighted_mode)
525 self.nb_size = max(self.kernel.shape)
526 else:
527 self.nb_size = grid_cells
529 try:
530 mask_cube_data = mask_cube.data
531 except AttributeError:
532 mask_cube_data = None
534 result_slices = CubeList()
535 for cube_slice in cube.slices([cube.coord(axis="y"), cube.coord(axis="x")]):
536 cube_slice.data = self._calculate_neighbourhood(
537 cube_slice.data, mask_cube_data
538 )
539 result_slices.append(cube_slice)
540 neighbourhood_averaged_cube = result_slices.merge_cube()
542 return neighbourhood_averaged_cube