def test_boxsum_with_automatic_cumsum(array_size_5):
"""Test that boxsum correctly calculates neighbourhood sums using raw array."""
result = boxsum(array_size_5, 3)
expected = np.array(
[[np.sum(array_size_5[i - 1:i + 2, j - 1:j + 2]) for j in [2, 3]]
for i in [2, 3]])
np.testing.assert_array_equal(result, expected)
def test_boxsum_with_padding(array_size_5):
"""Test that boxsum correctly calculates neighbourhood sums when adding padding to array."""
result = boxsum(array_size_5, 3, mode="constant", constant_values=0)
expected = np.array([[
np.sum(array_size_5[max(0, i - 1):i + 2,
max(0, j - 1):j + 2]) for j in range(5)
] for i in range(5)])
np.testing.assert_array_equal(result, expected)
def test_boxsum_non_square(array_size_5):
"""Test that boxsum correctly calculates neighbourhood sums using
non-square box."""
result = boxsum(array_size_5, (1, 3))
expected = np.array(
[[np.sum(array_size_5[i, j - 1:j + 2]) for j in [2, 3]]
for i in range(1, 5)])
np.testing.assert_array_equal(result, expected)
def test_boxsum_with_precalculated_cumsum(array_size_5):
"""Test that boxsum correctly calculates neighbourhood sums using
pre-calculated cumsum."""
cumsum_arr = np.array(
[[np.sum(array_size_5[:i + 1, :j + 1]) for j in range(5)]
for i in range(5)])
result = boxsum(cumsum_arr, 3, cumsum=False)
expected = np.array(
[[np.sum(array_size_5[i - 1:i + 2, j - 1:j + 2]) for j in [2, 3]]
for i in [2, 3]])
np.testing.assert_array_equal(result, expected)
def _calculate_neighbourhood(data, mask, nb_size, sum_only, re_mask):
"""
Apply neighbourhood processing.
Args:
data (numpy.ndarray):
Input data array.
mask (numpy.ndarray):
Mask of valid input data elements.
nb_size (int):
Size of the square neighbourhood as the number of grid cells.
sum_only (bool):
If true, return neighbourhood sum instead of mean.
re_mask (bool):
If true, reapply the original mask and return
`numpy.ma.MaskedArray`.
Returns:
numpy.ndarray:
Array containing the smoothed field after the square
neighbourhood method has been applied.
"""
if not sum_only:
min_val = np.nanmin(data)
max_val = np.nanmax(data)
# Use 64-bit types for enough precision in accumulations.
area_mask_dtype = np.int64
if mask is None:
area_mask = np.ones(data.shape, dtype=area_mask_dtype)
else:
area_mask = np.array(mask, dtype=area_mask_dtype, copy=False)
# Data mask to be eventually used for re-masking.
# (This is OK even if mask is None, it gives a scalar False mask then.)
data_mask = mask == 0
if isinstance(data, np.ma.MaskedArray):
# Include data mask if masked array.
data_mask = data_mask | data.mask
data = data.data
# Working type.
if issubclass(data.dtype.type, np.complexfloating):
data_dtype = np.complex128
else:
data_dtype = np.float64
data = np.array(data, dtype=data_dtype)
# Replace invalid elements with zeros.
nan_mask = np.isnan(data)
zero_mask = nan_mask | data_mask
np.copyto(area_mask, 0, where=zero_mask)
np.copyto(data, 0, where=zero_mask)
# Calculate neighbourhood totals for input data.
data = boxsum(data, nb_size, mode="constant")
if not sum_only:
# Calculate neighbourhood totals for mask.
area_sum = boxsum(area_mask, nb_size, mode="constant")
with np.errstate(divide="ignore", invalid="ignore"):
# Calculate neighbourhood mean.
data = data / area_sum
mask_invalid = (area_sum == 0) | nan_mask
np.copyto(data, np.nan, where=mask_invalid)
data = data.clip(min_val, max_val)
# Output type.
if issubclass(data.dtype.type, np.complexfloating):
data_dtype = np.complex64
else:
data_dtype = np.float32
data = data.astype(data_dtype)
if re_mask:
data = np.ma.masked_array(data, data_mask, copy=False)
return data
def test_boxsum_exception_not_odd(array_size_5):
"""Test that an exception is raised if `boxsize` contains a number that is not odd."""
msg = "The size of the neighbourhood must be an odd number."
with pytest.raises(ValueError) as exc_info:
boxsum(array_size_5, (1, 2))
assert msg in str(exc_info.value)
def test_boxsum_exception_non_integer(array_size_5):
"""Test that an exception is raised if `boxsize` is not an integer."""
msg = "The size of the neighbourhood must be of an integer type."
with pytest.raises(ValueError) as exc_info:
boxsum(array_size_5, 1.5)
assert msg in str(exc_info.value)
def _calculate_neighbourhood(
self,
data: ndarray,
mask: ndarray = None) -> Union[ndarray, np.ma.MaskedArray]:
"""
Apply neighbourhood processing. Ensures that masked data does not
contribute to the neighbourhood result. Masked data is either data that
is masked in the input data array or that corresponds to zeros in the
input mask.
Args:
data:
Input data array.
mask:
Mask of valid input data elements.
Returns:
Array containing the smoothed field after the
neighbourhood method has been applied.
"""
if not self.sum_only:
min_val = np.nanmin(data)
max_val = np.nanmax(data)
# Data mask to be eventually used for re-masking.
# (This is OK even if mask is None, it gives a scalar False mask then.)
# Invalid data where the mask provided == 0.
data_mask = mask == 0
if isinstance(data, np.ma.MaskedArray):
# Include data mask if masked array.
data_mask = data_mask | data.mask
data = data.data
# Working type.
if issubclass(data.dtype.type, np.complexfloating):
data_dtype = np.complex128
else:
# Use 64-bit types for enough precision in accumulations.
data_dtype = np.float64
data = np.array(data, dtype=data_dtype)
# Replace invalid elements with zeros so they don't count towards
# neighbourhood sum
valid_data_mask = np.ones(data.shape, dtype=np.int64)
valid_data_mask[data_mask] = 0
data[data_mask] = 0
# Calculate neighbourhood totals for input data.
if self.neighbourhood_method == "square":
data = boxsum(data, self.nb_size, mode="constant")
elif self.neighbourhood_method == "circular":
data = correlate(data, self.kernel, mode="nearest")
if not self.sum_only:
# Calculate neighbourhood totals for valid mask.
if self.neighbourhood_method == "square":
area_sum = boxsum(valid_data_mask,
self.nb_size,
mode="constant")
elif self.neighbourhood_method == "circular":
area_sum = correlate(valid_data_mask.astype(np.float32),
self.kernel,
mode="nearest")
with np.errstate(divide="ignore", invalid="ignore"):
# Calculate neighbourhood mean.
data = data / area_sum
# For points where all data in the neighbourhood is masked,
# set result to nan
data[area_sum == 0] = np.nan
data = data.clip(min_val, max_val)
# Output type.
if issubclass(data.dtype.type, np.complexfloating):
data_dtype = np.complex64
else:
data_dtype = np.float32
data = data.astype(data_dtype)
if self.re_mask:
data = np.ma.masked_array(data, data_mask, copy=False)
return data
