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