def check_input_coords(cube, require_time=False): """ Checks an input cube has precisely two non-scalar dimension coordinates (spatial x/y), or raises an error. If "require_time" is set to True, raises an error if no scalar time coordinate is present. Args: cube (iris.cube.Cube): Cube to be checked require_time (bool): Flag to check for a scalar time coordinate Raises: InvalidCubeError if coordinate requirements are not met """ # check that cube has both x and y axes try: check_for_x_and_y_axes(cube) except ValueError as msg: raise InvalidCubeError(msg) # check that cube data has only two non-scalar dimensions data_shape = np.array(cube.shape) non_scalar_coords = np.sum(np.where(data_shape > 1, 1, 0)) if non_scalar_coords > 2: raise InvalidCubeError('Cube has {:d} (more than 2) non-scalar ' 'coordinates'.format(non_scalar_coords)) if require_time: try: _ = cube.coord("time") except CoordinateNotFoundError: raise InvalidCubeError('Input cube has no time coordinate')
def test_fail_dimension_requirement(self): """Test that the expected exception is raised, if there the x and y coordinates are not dimensional coordinates.""" msg = "The cube does not contain the expected" cube = self.cube[0, 0, :, 0] with self.assertRaisesRegex(ValueError, msg): check_for_x_and_y_axes(cube, require_dim_coords=True)
def _check_template_cube(cube): """ The template cube is expected to contain a leading threshold dimension followed by spatial (y/x) dimensions. This check raises an error if this is not the case. If the cube contains the expected dimensions, a threshold leading order is enforced. Args: cube (iris.cube.Cube): A cube whose dimensions are checked to ensure they match what is expected. Raises: ValueError: If cube is not of the expected dimensions. """ check_for_x_and_y_axes(cube, require_dim_coords=True) dim_coords = [coord.name() for coord in cube.coords(dim_coords=True)] if 'threshold' in dim_coords and len(dim_coords) < 4: enforce_coordinate_ordering(cube, 'threshold') return msg = ('GenerateProbabilitiesFromMeanAndVariance expects a cube with ' 'only a leading threshold dimension, followed by spatial (y/x) ' 'dimensions. Got dimensions: {}'.format(dim_coords)) raise ValueError(msg)
def remove_halo_from_cube(cube, width_x, width_y): """ Method to remove rows/columns from the edge of an iris cube. Used to 'unpad' cubes which have been previously padded by pad_cube_with_halo. Args: cube (iris.cube.Cube): The original cube to be trimmed of edge data. The cube should contain only x and y dimensions, so will generally be a slice of a cube. width_x, width_y (int): The width in x and y directions of the neighbourhood radius in grid cells. This will be the width removed from the numpy array. Returns: trimmed_cube (iris.cube.Cube): Cube containing the new trimmed cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) end_y = -width_y if width_y != 0 else None end_x = -width_x if width_x != 0 else None trimmed_data = cube.data[width_y:end_y, width_x:end_x] coord_x = cube.coord(axis='x') trimmed_x_coord = pad_coord(coord_x, width_x, 'remove') coord_y = cube.coord(axis='y') trimmed_y_coord = pad_coord(coord_y, width_y, 'remove') trimmed_cube = _create_cube_with_padded_data( cube, trimmed_data, trimmed_x_coord, trimmed_y_coord) return trimmed_cube
def test_no_y_coordinate(self): """Test that the expected exception is raised, if there is no y coordinate.""" sliced_cube = next(self.cube.slices(["projection_x_coordinate"])) sliced_cube.remove_coord("projection_y_coordinate") msg = "The cube does not contain the expected" with self.assertRaisesRegex(ValueError, msg): check_for_x_and_y_axes(sliced_cube)
def pad_cube_with_halo(self, cube, width_x, width_y, masked_halo=False): """ Method to pad a halo around the data in an iris cube. Normally the masked_halo should be zero as it is considered masked data however if masked_halo is False then the padding calculates the mean within the neighbourhood radius in grid cells i.e. the neighbourhood width at the edge of the data and uses this mean value as the padding value. Args: cube (iris.cube.Cube): The original cube prior to applying padding. The cube should contain only x and y dimensions, so will generally be a slice of a cube. width_x, width_y (int): The width in x and y directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. masked_halo (bool): masked_halo = True means that the halo will be treated as masked points (i.e. set to 0.0) otherwise the halo will be filled with mean values. Default is set to False for backwards compatability as this function is used outside of SquareNeighbourhooding. Returns: padded_cube (iris.cube.Cube): Cube containing the new padded cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) # Pad a halo around the original data with the extent of the halo # given by width_y and width_x. Assumption to pad using the mean # value within the neighbourhood width for backwards compatability # as this function is used outside of SquareNeighbourhood. if masked_halo: padded_data = np.pad(cube.data, ((2 * width_y, 2 * width_y), (2 * width_x, 2 * width_x)), "constant", constant_values=(0.0, 0.0)) else: padded_data = np.pad(cube.data, ((2 * width_y, 2 * width_y), (2 * width_x, 2 * width_x)), "mean", stat_length=((width_y, width_y), (width_x, width_x))) coord_x = cube.coord(axis='x') padded_x_coord = (SquareNeighbourhood.pad_coord( coord_x, width_x, 'add')) coord_y = cube.coord(axis='y') padded_y_coord = (SquareNeighbourhood.pad_coord( coord_y, width_y, 'add')) padded_cube = self._create_cube_with_new_data(cube, padded_data, padded_x_coord, padded_y_coord) return padded_cube
def pad_cube_with_halo(cube: Cube, width_x: int, width_y: int, pad_method: str = "constant") -> Cube: """ Method to pad a halo around the data in an iris cube. If halo_with_data is False, the halo is filled with zeros. Otherwise the padding calculates a mean within half the padding width with which to fill the halo region. Args: cube: The original cube prior to applying padding. The cube should contain only x and y dimensions, so will generally be a slice of a cube. width_x: The width in x directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. width_y: The width in y directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. pad_method: The numpy.pad method with which to populate the halo. The default is 'constant' which will populate the region with zeros. All other np.pad methods are accepted, though they are not fully configurable. Returns: Cube containing the new padded cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) # Pad a halo around the original data with the extent of the halo # given by width_y and width_x. kwargs = { "stat_length": ((width_y // 2, width_y // 2), (width_x // 2, width_x // 2)) } if pad_method == "constant": kwargs = {"constant_values": (0.0, 0.0)} if pad_method == "symmetric": kwargs = {} padded_data = np.pad(cube.data, ((width_y, width_y), (width_x, width_x)), mode=pad_method, **kwargs) coord_x = cube.coord(axis="x") padded_x_coord = pad_coord(coord_x, width_x, "add") coord_y = cube.coord(axis="y") padded_y_coord = pad_coord(coord_y, width_y, "add") padded_cube = _create_cube_with_padded_data(cube, padded_data, padded_x_coord, padded_y_coord) return padded_cube
def _create_cube_with_padded_data(source_cube: Cube, data: ndarray, coord_x: DimCoord, coord_y: DimCoord) -> Cube: """ Create a cube with newly created data where the metadata is copied from the input cube and the supplied x and y coordinates are added to the cube. Args: source_cube: Template cube used for copying metadata and non x and y axes coordinates. data: Data to be put into the new cube. coord_x: Coordinate to be added to the new cube to represent the x axis. coord_y: Coordinate to be added to the new cube to represent the y axis. Returns: Cube built from the template cube using the requested data and the supplied x and y axis coordinates. """ check_for_x_and_y_axes(source_cube) yname = source_cube.coord(axis="y").name() xname = source_cube.coord(axis="x").name() ycoord_dim = source_cube.coord_dims(yname) xcoord_dim = source_cube.coord_dims(xname) # inherit metadata (cube name, units, attributes etc) metadata_dict = deepcopy(source_cube.metadata._asdict()) new_cube = iris.cube.Cube(data, **metadata_dict) # inherit non-spatial coordinates for coord in source_cube.coords(): if coord.name() not in [yname, xname]: if source_cube.coords(coord, dim_coords=True): coord_dim = source_cube.coord_dims(coord) new_cube.add_dim_coord(coord, coord_dim) else: new_cube.add_aux_coord(coord) # update spatial coordinates if len(xcoord_dim) > 0: new_cube.add_dim_coord(coord_x, xcoord_dim) else: new_cube.add_aux_coord(coord_x) if len(ycoord_dim) > 0: new_cube.add_dim_coord(coord_y, ycoord_dim) else: new_cube.add_aux_coord(coord_y) return new_cube
def pad_cube_with_halo(cube, width_x, width_y, halo_mean_data=True): """ Method to pad a halo around the data in an iris cube. If halo_with_data is False, the halo is filled with zeros. Otherwise the padding calculates a mean within half the padding width with which to fill the halo region. Args: cube (iris.cube.Cube): The original cube prior to applying padding. The cube should contain only x and y dimensions, so will generally be a slice of a cube. width_x (int): The width in x directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. width_y (int): The width in y directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. halo_mean_data (bool): Flag whether to populate the halo region with 0.0 or to fill with mean values derived from the existing data matrix. By default the mean data is used. Returns: padded_cube (iris.cube.Cube): Cube containing the new padded cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) # Pad a halo around the original data with the extent of the halo # given by width_y and width_x. if halo_mean_data: padded_data = np.pad( cube.data, ((width_y, width_y), (width_x, width_x)), "mean", stat_length=((0.5*width_y, 0.5*width_y), (0.5*width_x, 0.5*width_x))) else: padded_data = np.pad( cube.data, ((width_y, width_y), (width_x, width_x)), "constant", constant_values=(0.0, 0.0)) coord_x = cube.coord(axis='x') padded_x_coord = pad_coord(coord_x, width_x, 'add') coord_y = cube.coord(axis='y') padded_y_coord = pad_coord(coord_y, width_y, 'add') padded_cube = _create_cube_with_padded_data( cube, padded_data, padded_x_coord, padded_y_coord) return padded_cube
def _create_cube_with_new_data(cube, data, coord_x, coord_y): """ Create a cube with newly created data where the metadata is copied from the input cube and the supplied x and y coordinates are added to the cube. Parameters ---------- cube : Iris.cube.Cube Template cube used for copying metadata and non x and y axes coordinates. data : Numpy array Data to be put into the new cube. coord_x : Iris.coords.DimCoord Coordinate to be added to the new cube to represent the x axis. coord_y : Iris.coords.DimCoord Coordinate to be added to the new cube to represent the y axis. Returns ------- new_cube : Iris.cube.Cube Cube built from the template cube using the requested data and the supplied x and y axis coordinates. """ check_for_x_and_y_axes(cube) yname = cube.coord(axis='y').name() xname = cube.coord(axis='x').name() ycoord_dim = cube.coord_dims(yname) xcoord_dim = cube.coord_dims(xname) metadata_dict = copy.deepcopy(cube.metadata._asdict()) new_cube = iris.cube.Cube(data, **metadata_dict) for coord in cube.coords(): if coord.name() not in [yname, xname]: if cube.coords(coord, dim_coords=True): coord_dim = cube.coord_dims(coord) new_cube.add_dim_coord(coord, coord_dim) else: new_cube.add_aux_coord(coord) if len(xcoord_dim) > 0: new_cube.add_dim_coord(coord_x, xcoord_dim) else: new_cube.add_aux_coord(coord_x) if len(ycoord_dim) > 0: new_cube.add_dim_coord(coord_y, ycoord_dim) else: new_cube.add_aux_coord(coord_y) return new_cube
def pad_cube_with_halo(self, cube, width_x, width_y): """ Method to pad a halo around the data in an iris cube. The padding calculates the mean within the neighbourhood radius in grid cells i.e. the neighbourhood width at the edge of the data and uses this mean value as the padding value. Parameters ---------- cube : iris.cube.Cube The original cube prior to applying padding. width_x, width_y : integer The width in x and y directions of the neighbourhood radius in grid cells. This will be the width of padding to be added to the numpy array. Returns ------- iris.cube.Cube Cube containing the new padded cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) yname = cube.coord(axis='y').name() xname = cube.coord(axis='x').name() cubelist = iris.cube.CubeList([]) for slice_2d in cube.slices([yname, xname]): # Pad a halo around the original data with the extent of the halo # given by width_y and width_x. Assumption to pad using the mean # value within the neighbourhood width. padded_data = np.pad(slice_2d.data, ((2 * width_y, 2 * width_y), (2 * width_x, 2 * width_x)), "mean", stat_length=((width_y, width_y), (width_x, width_x))) coord_x = cube.coord(axis='x') padded_x_coord = (SquareNeighbourhood.pad_coord( coord_x, width_x, 'add')) coord_y = cube.coord(axis='y') padded_y_coord = (SquareNeighbourhood.pad_coord( coord_y, width_y, 'add')) cubelist.append( self._create_cube_with_new_data(slice_2d, padded_data, padded_x_coord, padded_y_coord)) return cubelist.merge_cube()
def remove_halo_from_cube(self, cube, width_x, width_y): """ Method to remove rows/columns from the edge of an iris cube. Used to 'unpad' cubes which have been previously padded by pad_cube_with_halo. Parameters ---------- cube : iris.cube.Cube The original cube to be trimmed of edge data. width_x, width_y : integer The width in x and y directions of the neighbourhood radius in grid cells. This will be the width removed from the numpy array. Returns ------- iris.cube.Cube Cube containing the new trimmed cube, with appropriate changes to the cube's dimension coordinates. """ check_for_x_and_y_axes(cube) yname = cube.coord(axis='y') xname = cube.coord(axis='x') cubelist = iris.cube.CubeList([]) for slice_2d in cube.slices([yname, xname]): end_y = -2 * width_y if width_y != 0 else None end_x = -2 * width_x if width_x != 0 else None trimmed_data = slice_2d.data[2 * width_y:end_y, 2 * width_x:end_x] coord_x = slice_2d.coord(axis='x') trimmed_x_coord = (SquareNeighbourhood.pad_coord( coord_x, width_x, 'remove')) coord_y = slice_2d.coord(axis='y') trimmed_y_coord = (SquareNeighbourhood.pad_coord( coord_y, width_y, 'remove')) cubelist.append( self._create_cube_with_new_data(slice_2d, trimmed_data, trimmed_x_coord, trimmed_y_coord)) return cubelist.merge_cube()
def test_pass_dimension_requirement(self): """Pass in compatible cubes that should not raise an exception. No assert statement required as any other input will raise an exception.""" check_for_x_and_y_axes(self.cube, require_dim_coords=True)
def mean_over_neighbourhood(self, summed_cube, summed_mask, cells_x, cells_y, iscomplex=False): """ Method to calculate the average value in a square neighbourhood using the 4-point algorithm to find the total sum over the neighbourhood. The output from the cumulate_array method can be used to calculate the sum over a neighbourhood of size (2*cells_x+1)*(2*cells_y+1). This sum is then divided by the area of the neighbourhood to calculate the mean value in the neighbourhood. For all points, a fast vectorised approach is taken: 1. The displacements between the four points used to calculate the neighbourhood total sum and the central grid point are calculated. 2. Within the function calculate_neighbourhood... Four copies of the cumulate array output are flattened and rolled by these displacements to align the four terms used in the neighbourhood total sum calculation. 3. The neighbourhood total at all points can then be calculated simultaneously in a single vector sum. Neighbourhood mean = Neighbourhood sum / Neighbourhood area Neighbourhood area = (2 * nb_width +1)^2 if there are no missing points, nb_width is the neighbourhood width, which is equal to 1 for a 3x3 neighbourhood. Args: summed_cube (iris.cube.Cube): Summed Cube to which neighbourhood processing is being applied. Must be passed through cumulate_array method first. The cube should contain only x and y dimensions, so will generally be a slice of a cube. summed_mask (iris.cube.Cube): Summed Mask used to calculate neighbourhood size. Must be passed through cumulate_array method first. The cube should contain only x and y dimensions, so will generally be a slice of a cube. cells_x, cells_y (int): The radius of the neighbourhood in grid points, in the x and y directions (excluding the central grid point). Kwargs: iscomplex (bool): Flag indicating whether cube.data contains complex values. Returns: cube (iris.cube.Cube): Cube to which square neighbourhood has been applied. """ cube = summed_cube check_for_x_and_y_axes(summed_cube) # Calculate displacement factors to find 4-points after flattening the # array. n_rows = len(cube.coord(axis="y").points) n_columns = len(cube.coord(axis="x").points) # Displacements from the point at the centre of the neighbourhood. # Equivalent to point B in the docstring example. ymax_xmax_disp = (cells_y * n_columns) + cells_x # Equivalent to point A in the docstring example. ymax_xmin_disp = (cells_y * n_columns) - cells_x - 1 # Equivalent to point D in the docstring example. ymin_xmax_disp = (-1 * (cells_y + 1) * n_columns) + cells_x # Equivalent to point C in the docstring example. ymin_xmin_disp = (-1 * (cells_y + 1) * n_columns) - cells_x - 1 # Flatten the cube data and create 4 copies of the flattened # array which are rolled to align the 4-points which are needed # for the calculation. neighbourhood_total = self.calculate_neighbourhood( summed_cube, ymax_xmax_disp, ymin_xmax_disp, ymin_xmin_disp, ymax_xmin_disp, n_rows, n_columns) if self.sum_or_fraction == "fraction": # Initialise and calculate the neighbourhood area. neighbourhood_area = self.calculate_neighbourhood( summed_mask, ymax_xmax_disp, ymin_xmax_disp, ymin_xmin_disp, ymax_xmin_disp, n_rows, n_columns) with np.errstate(invalid='ignore', divide='ignore'): if iscomplex: cube.data = (neighbourhood_total.astype(complex) / neighbourhood_area.astype(complex)) else: cube.data = (neighbourhood_total.astype(float) / neighbourhood_area.astype(float)) cube.data[~np.isfinite(cube.data)] = np.nan elif self.sum_or_fraction == "sum": if iscomplex: cube.data = neighbourhood_total.astype(complex) else: cube.data = neighbourhood_total.astype(float) return cube
def process(self, temperature, humidity, pressure, uwind, vwind, topography): """ Calculate precipitation enhancement over orography on standard and high resolution grids. Input variables are expected to be on the same grid (either standard or high resolution). Args: temperature (iris.cube.Cube): Temperature at top of boundary layer humidity (iris.cube.Cube): Relative humidity at top of boundary layer pressure (iris.cube.Cube): Pressure at top of boundary layer uwind (iris.cube.Cube): Positive eastward wind vector component at top of boundary layer vwind (iris.cube.Cube): Positive northward wind vector component at top of boundary layer topography (iris.cube.Cube): Height of topography above sea level on high resolution (1 km) UKPP domain grid Returns: (tuple): tuple containing: **orogenh** (iris.cube.Cube): Precipitation enhancement due to orography in mm/h on the 1 km Transverse Mercator UKPP grid domain **orogenh_standard_grid** (iris.cube.Cube): Precipitation enhancement due to orography in mm/h on the UK standard grid, padded with masked np.nans where outside the UKPP domain """ # check input variable cube coordinates match unmatched_coords = compare_coords( [temperature, pressure, humidity, uwind, vwind]) if any(item.keys() for item in unmatched_coords): msg = 'Input cube coordinates {} are unmatched' raise ValueError(msg.format(unmatched_coords)) # check one of the input variable cubes is a 2D spatial field (this is # equivalent to checking all cubes whose coords are matched above) msg = 'Require 2D fields as input; found {} dimensions' if temperature.ndim > 2: raise ValueError(msg.format(temperature.ndim)) check_for_x_and_y_axes(temperature) # check the topography cube is a 2D spatial field if topography.ndim > 2: raise ValueError(msg.format(topography.ndim)) check_for_x_and_y_axes(topography) # regrid variables to match topography and populate class instance self._regrid_and_populate(temperature, humidity, pressure, uwind, vwind, topography) # calculate saturation vapour pressure wbt = WetBulbTemperature() self.svp = wbt.pressure_correct_svp(wbt.lookup_svp(self.temperature), self.temperature, self.pressure) # calculate site-specific orographic enhancement point_orogenh_data = self._point_orogenh() # integrate upstream component grid_coord_km = self.topography.coord(axis='x').copy() grid_coord_km.convert_units('km') self.grid_spacing_km = (grid_coord_km.points[1] - grid_coord_km.points[0]) orogenh_data = self._add_upstream_component(point_orogenh_data) # create data cubes on the two required output grids orogenh, orogenh_standard_grid = self._create_output_cubes( orogenh_data, temperature) return orogenh, orogenh_standard_grid
def process(self, temperature, humidity, pressure, uwind, vwind, topography): """ Calculate precipitation enhancement over orography on high resolution grid. Input diagnostics are all expected to be on the same grid, and are regridded to match the orography. Args: temperature (iris.cube.Cube): Temperature at top of boundary layer humidity (iris.cube.Cube): Relative humidity at top of boundary layer pressure (iris.cube.Cube): Pressure at top of boundary layer uwind (iris.cube.Cube): Positive eastward wind vector component at top of boundary layer vwind (iris.cube.Cube): Positive northward wind vector component at top of boundary layer topography (iris.cube.Cube): Height of topography above sea level on high resolution (1 km) UKPP domain grid Returns: iris.cube.Cube: Precipitation enhancement due to orography in m/s. """ # check input variable cube coordinates match unmatched_coords = compare_coords( [temperature, pressure, humidity, uwind, vwind]) if any(item.keys() for item in unmatched_coords): msg = 'Input cube coordinates {} are unmatched' raise ValueError(msg.format(unmatched_coords)) # check one of the input variable cubes is a 2D spatial field (this is # equivalent to checking all cubes whose coords are matched above) msg = 'Require 2D fields as input; found {} dimensions' if temperature.ndim > 2: raise ValueError(msg.format(temperature.ndim)) check_for_x_and_y_axes(temperature) # check the topography cube is a 2D spatial field if topography.ndim > 2: raise ValueError(msg.format(topography.ndim)) check_for_x_and_y_axes(topography) # regrid variables to match topography and populate class instance self._regrid_and_populate(temperature, humidity, pressure, uwind, vwind, topography) # calculate saturation vapour pressure self.svp = calculate_svp_in_air( self.temperature.data, self.pressure.data) # calculate site-specific orographic enhancement point_orogenh_data = self._point_orogenh() # integrate upstream component grid_coord_km = self.topography.coord(axis='x').copy() grid_coord_km.convert_units('km') self.grid_spacing_km = ( grid_coord_km.points[1] - grid_coord_km.points[0]) orogenh_data = self._add_upstream_component(point_orogenh_data) # create data cubes on the two required output grids orogenh = self._create_output_cube(orogenh_data, temperature) return orogenh
def mean_over_neighbourhood(self, cube, cells_x, cells_y, nan_masks): """ Method to calculate the average value in a square neighbourhood using the 4-point algorithm to find the total sum over the neighbourhood. The output from the cumulate_array method can be used to calculate the sum over a neighbourhood of size (2*cells_x+1)*(2*cells_y+1). This sum is then divided by the area of the neighbourhood to calculate the mean value in the neighbourhood. For all points, a fast vectorised approach is taken: 1. The displacements between the four points used to calculate the neighbourhood total sum and the central grid point are calculated. 2. Four copies of the cumulate array output are flattened and rolled by these displacements to align the four terms used in the neighbourhood total sum calculation. 3. The neighbourhood total at all points can then be calculated simultaneously in a single vector sum. Displacements are calculated as follows for the following input array, where the accumulation has occurred from left to right and top to bottom:: | 2 | 4 | 6 | 7 | | 2 | 4 | 5 | 6 | | 1 | 3 | 4 | 4 | | 1 | 2 | 2 | 2 | For a 3x3 neighbourhood centred around the point with a value of 5:: | 2 (A) | 4 | 6 | 7 (B) | | 2 | 4 | 5 (Central point) | 6 | | 1 | 3 | 4 | 4 | | 1 (C) | 2 | 2 | 2 (D) | To calculate the value for the neighbourhood sum at the "Central point" with a value of 5, calculate:: Neighbourhood sum = B - A - D + C At the central point, this will yield:: Neighbourhood sum = 7 - 2 - 2 +1 => 4 Neighbourhood mean = Neighbourhood sum ----------------- (2 * nb_width +1) where nb_width is the neighbourhood width, which is equal to 1 for a 3x3 neighbourhood. This example gives:: Neighbourhood mean = 4. / 9. Args: cube (iris.cube.Cube): Cube to which neighbourhood processing is being applied. Must be passed through cumulate_array method first. cells_x, cells_y (int): The radius of the neighbourhood in grid points, in the x and y directions (excluding the central grid point). nan_masks (list): List of numpy arrays to be used to set the values within the data of the output cube to be NaN. Returns: cube (iris.cube.Cube): Cube to which square neighbourhood has been applied. """ check_for_x_and_y_axes(cube) yname = cube.coord(axis="y").name() xname = cube.coord(axis="x").name() # Calculate displacement factors to find 4-points after flattening the # array. n_rows = len(cube.coord(axis="y").points) n_columns = len(cube.coord(axis="x").points) # Displacements from the point at the centre of the neighbourhood. # Equivalent to point B in the docstring example. ymax_xmax_disp = (cells_y * n_columns) + cells_x # Equivalent to point A in the docstring example. ymax_xmin_disp = (cells_y * n_columns) - cells_x - 1 # Equivalent to point D in the docstring example. ymin_xmax_disp = (-1 * (cells_y + 1) * n_columns) + cells_x # Equivalent to point C in the docstring example. ymin_xmin_disp = (-1 * (cells_y + 1) * n_columns) - cells_x - 1 cubelist = iris.cube.CubeList([]) for slice_2d, nan_mask in zip(cube.slices([yname, xname]), nan_masks): # Flatten the 2d slice and create 4 copies of the flattened # array which are rolled to align the 4-points which are needed # for the calculation. flattened = slice_2d.data.flatten() ymax_xmax_array = np.roll(flattened, -ymax_xmax_disp) ymin_xmax_array = np.roll(flattened, -ymin_xmax_disp) ymin_xmin_array = np.roll(flattened, -ymin_xmin_disp) ymax_xmin_array = np.roll(flattened, -ymax_xmin_disp) neighbourhood_total = (ymax_xmax_array - ymin_xmax_array + ymin_xmin_array - ymax_xmin_array) neighbourhood_total.resize(n_rows, n_columns) if self.sum_or_fraction == "fraction": # Initialise and calculate the neighbourhood area. neighbourhood_area = np.zeros(neighbourhood_total.shape) neighbourhood_area.fill((2 * cells_x + 1) * (2 * cells_y + 1)) with np.errstate(invalid='ignore', divide='ignore'): slice_2d.data = (neighbourhood_total.astype(float) / neighbourhood_area.astype(float)) elif self.sum_or_fraction == "sum": slice_2d.data = neighbourhood_total.astype(float) slice_2d.data[nan_mask.astype(bool)] = np.NaN cubelist.append(slice_2d) return cubelist.merge_cube()