def _get_advection_time(cube1: Cube, cube2: Cube) -> None: """Get time over which the advection has occurred, in seconds, using the difference in time or forecast reference time between input cubes""" time_diff_seconds = ( cube2.coord("time").cell(0).point - cube1.coord("time").cell(0).point).total_seconds() time_diff_seconds = int(time_diff_seconds) if time_diff_seconds == 0: # second cube should be an observation; first cube should have a # non-zero forecast period which describes the advection time if (cube2.coords("forecast_period") and cube2.coord("forecast_period").points[0] != 0): raise InvalidCubeError( "Second input cube must be a current observation") # get the time difference from the first cube's forecast period fp_coord = cube1.coord("forecast_period").copy() fp_coord.convert_units("seconds") (time_diff_seconds, ) = fp_coord.points if time_diff_seconds <= 0: error_msg = "Expected positive time difference cube2 - cube1: got {} s" raise InvalidCubeError(error_msg.format(time_diff_seconds)) return time_diff_seconds
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 __init__(self, vel_x, vel_y): """ Initialises the plugin. Velocities are expected to be on a regular grid (such that grid spacing in metres is the same at all points in the domain). Args: vel_x (iris.cube.Cube): Cube containing a 2D array of velocities along the x coordinate axis vel_y (iris.cube.Cube): Cube containing a 2D array of velocities along the y coordinate axis """ # check each input velocity cube has precisely two non-scalar # dimension coordinates (spatial x/y) check_input_coords(vel_x) check_input_coords(vel_y) # check input velocity cubes have the same spatial coordinates if (vel_x.coord(axis="x") != vel_y.coord(axis="x") or vel_x.coord(axis="y") != vel_y.coord(axis="y")): raise InvalidCubeError("Velocity cubes on unmatched grids") vel_x.convert_units('m s-1') vel_y.convert_units('m s-1') self.vel_x = vel_x self.vel_y = vel_y self.x_coord = vel_x.coord(axis="x") self.y_coord = vel_x.coord(axis="y")
def _check_input_cubes(cube1: Cube, cube2: Cube) -> None: """Check that input cubes have appropriate and matching dimensions""" # check the nature of the input cubes, and raise a warning if they are # not both precipitation if cube1.name() != cube2.name(): msg = "Input cubes contain different data types {} and {}" raise ValueError(msg.format(cube1.name(), cube2.name())) data_name = cube1.name().lower() if "rain" not in data_name and "precipitation" not in data_name: msg = ("Input data are of non-precipitation type {}. Plugin " "parameters have not been tested and may not be appropriate" " for this variable.") warnings.warn(msg.format(cube1.name())) # check cubes have exactly two spatial dimension coordinates and a # scalar time coordinate check_input_coords(cube1, require_time=True) check_input_coords(cube2, require_time=True) # check cube dimensions match if cube1.coord(axis="x") != cube2.coord(axis="x") or cube1.coord( axis="y") != cube2.coord(axis="y"): raise InvalidCubeError("Input cubes on unmatched grids") # check grids are equal area check_if_grid_is_equal_area(cube1) check_if_grid_is_equal_area(cube2)
def process(self, cube, timestep, fill_value=0.0): """ Extrapolates input cube data and updates validity time. The input cube should have precisely two non-scalar dimension coordinates (spatial x/y), and is expected to be in a projection such that grid spacing is the same (or very close) at all points within the spatial domain. The input cube should also have a "time" coordinate. Args: cube (iris.cube.Cube): The 2D cube containing data to be advected timestep (datetime.timedelta): Advection time step fill_value (float): Default output value for spatial points where data cannot be extrapolated (source is out of bounds) Returns: advected_cube (iris.cube.Cube): New cube with updated time and extrapolated data """ # check that the input cube has precisely two non-scalar dimension # coordinates (spatial x/y) and a scalar time coordinate check_input_coords(cube, require_time=True) # check spatial coordinates match those of plugin velocities if (cube.coord(axis="x") != self.x_coord or cube.coord(axis="y") != self.y_coord): raise InvalidCubeError("Input data grid does not match advection " "velocities") # derive velocities in "grid squares per second" def grid_spacing(coord): """Calculate grid spacing along a given spatial axis""" new_coord = coord.copy() new_coord.convert_units('m') return float(np.diff((new_coord).points)[0]) grid_vel_x = self.vel_x.data / grid_spacing(cube.coord(axis="x")) grid_vel_y = self.vel_y.data / grid_spacing(cube.coord(axis="y")) # perform advection and create output cube advected_data = self._advect_field(cube.data, grid_vel_x, grid_vel_y, timestep.total_seconds(), fill_value) advected_cube = cube.copy(data=advected_data) # increment output cube time original_datetime, = \ (cube.coord("time").units).num2date(cube.coord("time").points) new_datetime = original_datetime + timestep new_time = (cube.coord("time").units).date2num(new_datetime) advected_cube.coord("time").points = new_time return advected_cube
def process(self, cube, timestep): """ Extrapolates input cube data and updates validity time. The input cube should have precisely two non-scalar dimension coordinates (spatial x/y), and is expected to be in a projection such that grid spacing is the same (or very close) at all points within the spatial domain. The input cube should also have a "time" coordinate. Args: cube (iris.cube.Cube): The 2D cube containing data to be advected timestep (datetime.timedelta): Advection time step Returns: iris.cube.Cube: New cube with updated time and extrapolated data. New data are filled with np.nan and masked where source data were out of bounds (ie where data could not be advected from outside the cube domain). """ # check that the input cube has precisely two non-scalar dimension # coordinates (spatial x/y) and a scalar time coordinate check_input_coords(cube, require_time=True) # check spatial coordinates match those of plugin velocities if cube.coord(axis="x") != self.x_coord or cube.coord(axis="y") != self.y_coord: raise InvalidCubeError( "Input data grid does not match advection " "velocities" ) # derive velocities in "grid squares per second" def grid_spacing(coord): """Calculate grid spacing along a given spatial axis""" new_coord = coord.copy() new_coord.convert_units("m") return np.float32(np.diff((new_coord).points)[0]) grid_vel_x = self.vel_x.data / grid_spacing(cube.coord(axis="x")) grid_vel_y = self.vel_y.data / grid_spacing(cube.coord(axis="y")) # raise a warning if data contains unmasked NaNs nan_count = np.count_nonzero(~np.isfinite(cube.data)) if nan_count > 0: warnings.warn("input data contains unmasked NaNs") # perform advection and create output cube advected_data = self._advect_field( cube.data, grid_vel_x, grid_vel_y, round(timestep.total_seconds()) ) advected_cube = self._create_output_cube(cube, advected_data, timestep) return advected_cube
def process(cube, ensemble_realization_numbers=None): """ Rebadge percentiles as ensemble realizations. The ensemble realization numbering will depend upon the number of percentiles in the input cube i.e. 0, 1, 2, 3, ..., n-1, if there are n percentiles. Args: cube (iris.cube.Cube): Cube containing a percentile coordinate, which will be rebadged as ensemble realization. Keyword Args: ensemble_realization_numbers (numpy.ndarray): An array containing the ensemble numbers required in the output realization coordinate. Default is None, meaning the realization coordinate will be numbered 0, 1, 2 ... n-1 for n percentiles on the input cube. Raises: InvalidCubeError: If the realization coordinate already exists on the cube. """ percentile_coord_name = ( find_percentile_coordinate(cube).name()) if ensemble_realization_numbers is None: ensemble_realization_numbers = ( np.arange( len(cube.coord(percentile_coord_name).points))) cube.coord(percentile_coord_name).points = ( ensemble_realization_numbers) # we can't rebadge if the realization coordinate already exists: try: realization_coord = cube.coord('realization') except CoordinateNotFoundError: realization_coord = None if realization_coord: raise InvalidCubeError( "Cannot rebadge percentile coordinate to realization " "coordinate because a realization coordinate already exists.") cube.coord(percentile_coord_name).rename("realization") cube.coord("realization").units = "1" return cube
def __init__(self, vel_x, vel_y, metadata_dict=None): """ Initialises the plugin. Velocities are expected to be on a regular grid (such that grid spacing in metres is the same at all points in the domain). Args: vel_x (iris.cube.Cube): Cube containing a 2D array of velocities along the x coordinate axis vel_y (iris.cube.Cube): Cube containing a 2D array of velocities along the y coordinate axis Keyword Args: metadata_dict (dict): Dictionary containing information for amending the metadata of the output cube. Please see the :func:`improver.utilities.cube_metadata.amend_metadata` for information regarding the allowed contents of the metadata dictionary. """ # check each input velocity cube has precisely two non-scalar # dimension coordinates (spatial x/y) check_input_coords(vel_x) check_input_coords(vel_y) # check input velocity cubes have the same spatial coordinates if (vel_x.coord(axis="x") != vel_y.coord(axis="x") or vel_x.coord(axis="y") != vel_y.coord(axis="y")): raise InvalidCubeError("Velocity cubes on unmatched grids") vel_x.convert_units('m s-1') vel_y.convert_units('m s-1') self.vel_x = vel_x self.vel_y = vel_y self.x_coord = vel_x.coord(axis="x") self.y_coord = vel_x.coord(axis="y") # Initialise metadata dictionary. if metadata_dict is None: metadata_dict = {} self.metadata_dict = metadata_dict
def __init__(self, vel_x: Cube, vel_y: Cube, attributes_dict: Optional[Dict] = None) -> None: """ Initialises the plugin. Velocities are expected to be on a regular grid (such that grid spacing in metres is the same at all points in the domain). Args: vel_x: Cube containing a 2D array of velocities along the x coordinate axis vel_y: Cube containing a 2D array of velocities along the y coordinate axis attributes_dict: Dictionary containing information for amending the attributes of the output cube. """ # check each input velocity cube has precisely two non-scalar # dimension coordinates (spatial x/y) check_input_coords(vel_x) check_input_coords(vel_y) # check input velocity cubes have the same spatial coordinates if vel_x.coord(axis="x") != vel_y.coord(axis="x") or vel_x.coord( axis="y") != vel_y.coord(axis="y"): raise InvalidCubeError("Velocity cubes on unmatched grids") vel_x.convert_units("m s-1") vel_y.convert_units("m s-1") self.vel_x = vel_x self.vel_y = vel_y self.x_coord = vel_x.coord(axis="x") self.y_coord = vel_x.coord(axis="y") # Initialise metadata dictionary. if attributes_dict is None: attributes_dict = {} self.attributes_dict = attributes_dict
def process(self, orography, thresholds_dict, landmask=None): """Calculate the weights depending upon where the orography point is within the topographic zones. Args: orography (iris.cube.Cube): Orography on standard grid. thresholds_dict (dict): Definition of orography bands required. The expected format of the dictionary is e.g. `{'bounds': [[0, 50], [50, 200]], 'units': 'm'}` landmask (iris.cube.Cube): Land mask on standard grid, with land points set to one and sea points set to zero. If provided sea points are masked out in the output array. Returns: iris.cube.Cube: Cube containing the weights depending upon where the orography point is within the topographic zones. """ # Check that orography is a 2d cube. if len(orography.shape) != 2: msg = ("The input orography cube should be two-dimensional." "The input orography cube has {} dimensions".format( len(orography.shape))) raise InvalidCubeError(msg) # Find bands and midpoints from bounds. bands = np.array(thresholds_dict["bounds"], dtype=np.float32) threshold_units = thresholds_dict["units"] # Create topographic_zone_cube first, so that a cube is created for # each band. This will allow the data for neighbouring bands to be # put into the cube. mask_data = np.zeros(orography.shape, dtype=np.float32) topographic_zone_cubes = iris.cube.CubeList([]) for band in bands: sea_points_included = not landmask topographic_zone_cube = _make_mask_cube( mask_data, orography.coords(), band, threshold_units, sea_points_included=sea_points_included, ) topographic_zone_cubes.append(topographic_zone_cube) topographic_zone_weights = topographic_zone_cubes.concatenate_cube() topographic_zone_weights.data = topographic_zone_weights.data.astype( np.float32) # Ensure topographic_zone coordinate units is equal to orography units. topographic_zone_weights.coord("topographic_zone").convert_units( orography.units) # Read bands from cube, now that they can be guaranteed to be in the # same units as the orography. The bands are converted to a list, so # that they can be iterated through. bands = list(topographic_zone_weights.coord("topographic_zone").bounds) midpoints = topographic_zone_weights.coord("topographic_zone").points # Raise a warning, if orography extremes are outside the extremes of # the bands. if np.max(orography.data) > np.max(bands): msg = ("The maximum orography is greater than the uppermost band. " "This will potentially cause the topographic zone weights " "to not sum to 1 for a given grid point.") warnings.warn(msg) if np.min(orography.data) < np.min(bands): msg = ("The minimum orography is lower than the lowest band. " "This will potentially cause the topographic zone weights " "to not sum to 1 for a given grid point.") warnings.warn(msg) # Insert the appropriate weights into the topographic zone cube. This # includes the weights from the band that a point is in, as well as # the contribution from an adjacent band. for band_number, band in enumerate(bands): # Determine the points that are within the specified band. mask_y, mask_x = np.where((orography.data > band[0]) & (orography.data <= band[1])) orography_band = np.full(orography.shape, np.nan, dtype=np.float32) orography_band[mask_y, mask_x] = orography.data[mask_y, mask_x] # Calculate the weights. This involves calculating the # weights for all the orography but only inserting weights # that are within the band into the topographic_zone_weights cube. weights = self.calculate_weights(orography_band, band) topographic_zone_weights.data[band_number, mask_y, mask_x] = weights[mask_y, mask_x] # Calculate the contribution to the weights from the adjacent # lower band. topographic_zone_weights.data = self.add_weight_to_lower_adjacent_band( topographic_zone_weights.data, orography_band, midpoints[band_number], band_number, ) # Calculate the contribution to the weights from the adjacent # upper band. topographic_zone_weights.data = self.add_weight_to_upper_adjacent_band( topographic_zone_weights.data, orography_band, midpoints[band_number], band_number, len(bands) - 1, ) # Metadata updates topographic_zone_weights.rename("topographic_zone_weights") topographic_zone_weights.units = Unit("1") # Mask output weights using a land-sea mask. topographic_zone_masked_weights = iris.cube.CubeList([]) for topographic_zone_slice in topographic_zone_weights.slices_over( "topographic_zone"): if landmask: topographic_zone_slice.data = GenerateOrographyBandAncils( ).sea_mask(landmask.data, topographic_zone_slice.data) topographic_zone_masked_weights.append(topographic_zone_slice) topographic_zone_weights = topographic_zone_masked_weights.merge_cube() return topographic_zone_weights
def process(self, cube1, cube2): """ Extracts data from input cubes, performs dimensionless advection displacement calculation, and creates new cubes with advection velocities in metres per second. Each input cube should have precisely two non-scalar dimension coordinates (spatial x/y), and are expected to be in a projection such that grid spacing is the same (or very close) at all points within the spatial domain. Each input cube must also have a scalar "time" coordinate. Args: cube1 (iris.cube.Cube): 2D cube from (earlier) time 1 cube2 (iris.cube.Cube): 2D cube from (later) time 2 Returns: ucube (iris.cube.Cube): 2D cube of advection velocities in the x-direction vcube (iris.cube.Cube): 2D cube of advection velocities in the y-direction """ # check cubes have exactly two spatial dimension coordinates and a # scalar time coordinate check_input_coords(cube1, require_time=True) check_input_coords(cube2, require_time=True) # check cube dimensions match if (cube1.coord(axis="x") != cube2.coord(axis="x") or cube1.coord(axis="y") != cube2.coord(axis="y")): raise InvalidCubeError("Input cubes on unmatched grids") # check time difference is positive time1 = (cube1.coord("time").units).num2date( cube1.coord("time").points[0]) time2 = (cube2.coord("time").units).num2date( cube2.coord("time").points[0]) cube_time_diff = time2 - time1 if cube_time_diff.total_seconds() <= 0: msg = "Expected positive time difference cube2 - cube1: got {} s" raise InvalidCubeError(msg.format(cube_time_diff.seconds)) # extract spatial grid length new_coord = cube1.coord(axis='x').copy() new_coord.convert_units('km') grid_length_km = float(np.diff((new_coord).points)[0]) # check x and y have the same grid length - fail if not new_coord = cube1.coord(axis='y').copy() new_coord.convert_units('km') grid_length_y_km = float(np.diff((new_coord).points)[0]) if not np.isclose(grid_length_y_km, grid_length_km): raise InvalidCubeError("Input cube has different grid spacing in " "x and y") # calculate plugin parameters in grid square units self.data_smoothing_radius = \ int(self.data_smoothing_radius_km / grid_length_km) self.boxsize = int(self.boxsize_km / grid_length_km) # Fail verbosely if self.data_smoothing_radius is too small and will # trigger silent failures downstream. (Note that self.boxsize < 3 will # also trigger silent failures, but this is caught indirectly by # enoforcing boxsize > data_smoothing_radius at initialisation.) if self.data_smoothing_radius < 3: msg = ("Input data smoothing radius {} too small (minimum 3 " "grid squares)") raise ValueError(msg.format(self.data_smoothing_radius)) # TODO raise further warnings if data_smoothing_radius or boxsize # have values (to be determined) that are scientifically questionable. # Here if dimensionless; at initialisation if dimensioned. # calculate dimensionless displacement between the two input fields data1 = next(cube1.slices([cube1.coord(axis='y'), cube1.coord(axis='x')])).data data2 = next(cube2.slices([cube2.coord(axis='y'), cube2.coord(axis='x')])).data ucomp, vcomp = self.process_dimensionless(data1, data2, 1, 0) # convert displacements to velocities in metres per second for vel in [ucomp, vcomp]: vel *= (1000.*grid_length_km) vel /= cube_time_diff.total_seconds() # create velocity output cubes based on metadata from later input cube x_coord = cube2.coord(axis="x") y_coord = cube2.coord(axis="y") t_coord = cube2.coord("time") ucube = iris.cube.Cube( ucomp, long_name="advection_velocity_x", units="m s-1", dim_coords_and_dims=[(y_coord, 0), (x_coord, 1)]) ucube.add_aux_coord(t_coord) vcube = iris.cube.Cube( vcomp, long_name="advection_velocity_y", units="m s-1", dim_coords_and_dims=[(y_coord, 0), (x_coord, 1)]) vcube.add_aux_coord(t_coord) return ucube, vcube
def process(self, cube1, cube2, boxsize=30): """ Extracts data from input cubes, performs dimensionless advection displacement calculation, and creates new cubes with advection velocities in metres per second. Each input cube should have precisely two non-scalar dimension coordinates (spatial x/y), and are expected to be in a projection such that grid spacing is the same (or very close) at all points within the spatial domain. Each input cube must also have a scalar "time" coordinate. Args: cube1 (iris.cube.Cube): 2D cube from (earlier) time 1 cube2 (iris.cube.Cube): 2D cube from (later) time 2 Kwargs: boxsize (int): The side length of the square box over which to solve the optical flow constraint. This should be greater than the data smoothing radius. Returns: (tuple) : tuple containing: **ucube** (iris.cube.Cube): 2D cube of advection velocities in the x-direction **vcube** (iris.cube.Cube): 2D cube of advection velocities in the y-direction """ # clear existing parameters self.data_smoothing_radius = None self.boxsize = None # check the nature of the input cubes, and raise a warning if they are # not both precipitation if cube1.name() != cube2.name(): msg = 'Input cubes contain different data types {} and {}' raise ValueError(msg.format(cube1.name(), cube2.name())) data_name = cube1.name().lower() if "rain" not in data_name and "precipitation" not in data_name: msg = ('Input data are of non-precipitation type {}. Plugin ' 'parameters have not been tested and may not be appropriate' ' for this variable.') warnings.warn(msg.format(cube1.name())) # check cubes have exactly two spatial dimension coordinates and a # scalar time coordinate check_input_coords(cube1, require_time=True) check_input_coords(cube2, require_time=True) # check cube dimensions match if (cube1.coord(axis="x") != cube2.coord(axis="x") or cube1.coord(axis="y") != cube2.coord(axis="y")): raise InvalidCubeError("Input cubes on unmatched grids") # check grids are equal area check_if_grid_is_equal_area(cube1) check_if_grid_is_equal_area(cube2) # convert units to mm/hr as these avoid the need to manipulate tiny # decimals try: cube1 = cube1.copy() cube2 = cube2.copy() cube1.convert_units('mm/hr') cube2.convert_units('mm/hr') except ValueError as err: msg = ('Input data are in units that cannot be converted to mm/hr ' 'which are the required units for use with optical flow.') raise ValueError(msg) from err # check time difference is positive time1 = (cube1.coord("time").units).num2date( cube1.coord("time").points[0]) time2 = (cube2.coord("time").units).num2date( cube2.coord("time").points[0]) cube_time_diff = time2 - time1 if cube_time_diff.total_seconds() <= 0: msg = "Expected positive time difference cube2 - cube1: got {} s" raise InvalidCubeError(msg.format(cube_time_diff.total_seconds())) # if time difference is greater 15 minutes, increase data smoothing # radius so that larger advection displacements can be resolved if cube_time_diff.total_seconds() > 900: data_smoothing_radius_km = self.data_smoothing_radius_km * ( cube_time_diff.total_seconds()/900.) else: data_smoothing_radius_km = self.data_smoothing_radius_km # calculate smoothing radius in grid square units new_coord = cube1.coord(axis='x').copy() new_coord.convert_units('km') grid_length_km = np.float32(np.diff((new_coord).points)[0]) data_smoothing_radius = \ int(data_smoothing_radius_km / grid_length_km) # Fail verbosely if data smoothing radius is too small and will # trigger silent failures downstream if data_smoothing_radius < 3: msg = ("Input data smoothing radius {} too small (minimum 3 " "grid squares)") raise ValueError(msg.format(data_smoothing_radius)) # Fail if self.boxsize is less than data smoothing radius self.boxsize = boxsize if self.boxsize < data_smoothing_radius: msg = ("Box size {} too small (should not be less than data " "smoothing radius {})") raise ValueError( msg.format(self.boxsize, data_smoothing_radius)) # extract 2-dimensional data arrays data1 = next(cube1.slices([cube1.coord(axis='y'), cube1.coord(axis='x')])).data data2 = next(cube2.slices([cube2.coord(axis='y'), cube2.coord(axis='x')])).data # fill any mask with 0 values so fill_values are not spread into the # domain when smoothing the fields. if np.ma.is_masked(data1): data1 = data1.filled(0) if np.ma.is_masked(data2): data2 = data2.filled(0) # if input arrays have no non-zero values, set velocities to zero here # and raise a warning if (np.allclose(data1, np.zeros(data1.shape)) or np.allclose(data2, np.zeros(data2.shape))): msg = ("No non-zero data in input fields: setting optical flow " "velocities to zero") warnings.warn(msg) ucomp = np.zeros(data1.shape, dtype=np.float32) vcomp = np.zeros(data2.shape, dtype=np.float32) else: # calculate dimensionless displacement between the two input fields ucomp, vcomp = self.process_dimensionless(data1, data2, 1, 0, data_smoothing_radius) # convert displacements to velocities in metres per second for vel in [ucomp, vcomp]: vel *= np.float32(1000.*grid_length_km) vel /= cube_time_diff.total_seconds() # create velocity output cubes based on metadata from later input cube x_coord = cube2.coord(axis="x") y_coord = cube2.coord(axis="y") t_coord = cube2.coord("time") ucube = iris.cube.Cube( ucomp, long_name="precipitation_advection_x_velocity", units="m s-1", dim_coords_and_dims=[(y_coord, 0), (x_coord, 1)]) ucube.add_aux_coord(t_coord) ucube = amend_metadata(ucube, **self.metadata_dict) vcube = iris.cube.Cube( vcomp, long_name="precipitation_advection_y_velocity", units="m s-1", dim_coords_and_dims=[(y_coord, 0), (x_coord, 1)]) vcube.add_aux_coord(t_coord) vcube = amend_metadata(vcube, **self.metadata_dict) return ucube, vcube
def process(self, cube, timestep): """ Extrapolates input cube data and updates validity time. The input cube should have precisely two non-scalar dimension coordinates (spatial x/y), and is expected to be in a projection such that grid spacing is the same (or very close) at all points within the spatial domain. The input cube should also have a "time" coordinate. Args: cube (iris.cube.Cube): The 2D cube containing data to be advected timestep (datetime.timedelta): Advection time step Returns: advected_cube (iris.cube.Cube): New cube with updated time and extrapolated data. New data are filled with np.nan and masked where source data were out of bounds (ie where data could not be advected from outside the cube domain). """ # check that the input cube has precisely two non-scalar dimension # coordinates (spatial x/y) and a scalar time coordinate check_input_coords(cube, require_time=True) # check spatial coordinates match those of plugin velocities if (cube.coord(axis="x") != self.x_coord or cube.coord(axis="y") != self.y_coord): raise InvalidCubeError("Input data grid does not match advection " "velocities") # derive velocities in "grid squares per second" def grid_spacing(coord): """Calculate grid spacing along a given spatial axis""" new_coord = coord.copy() new_coord.convert_units('m') return float(np.diff((new_coord).points)[0]) grid_vel_x = self.vel_x.data / grid_spacing(cube.coord(axis="x")) grid_vel_y = self.vel_y.data / grid_spacing(cube.coord(axis="y")) # raise a warning if data contains unmasked NaNs nan_count = np.count_nonzero(~np.isfinite(cube.data)) if nan_count > 0: warnings.warn("input data contains unmasked NaNs") # perform advection and create output cube advected_data = self._advect_field(cube.data, grid_vel_x, grid_vel_y, timestep.total_seconds()) advected_cube = cube.copy(data=advected_data) # increment output cube time and add a "forecast_period" coordinate original_datetime, = \ (cube.coord("time").units).num2date(cube.coord("time").points) new_datetime = original_datetime + timestep new_time = (cube.coord("time").units).date2num(new_datetime) advected_cube.coord("time").points = new_time forecast_period_seconds = timestep.total_seconds() forecast_period_coord = AuxCoord(forecast_period_seconds, standard_name="forecast_period", units="s") try: advected_cube.remove_coord("forecast_period") except CoordinateNotFoundError: pass advected_cube.add_aux_coord(forecast_period_coord) return advected_cube
def check_cube_coordinates(cube, new_cube, exception_coordinates=None): """Find and promote to dimension coordinates any scalar coordinates in new_cube that were originally dimension coordinates in the progenitor cube. If coordinate is in new_cube that is not in the old cube, keep coordinate in its current position. Args: cube (iris.cube.Cube): The input cube that will be checked to identify the preferred coordinate order for the output cube. new_cube (iris.cube.Cube): The cube that must be checked and adjusted using the coordinate order from the original cube. exception_coordinates (List of strings or None): The names of the coordinates that are permitted to be within the new_cube but are not available within the original cube. Returns: new_cube (iris.cube.Cube): Modified cube with relevant scalar coordinates promoted to dimension coordinates with the dimension coordinates re-ordered, as best as can be done based on the original cube. Raises: CoordinateNotFoundError : Raised if the final dimension coordinates of the returned cube do not match the input cube. InvalidCubeError : If the coordinate is not within the original cube and it is not within the list of permitted exceptions. """ if exception_coordinates is None: exception_coordinates = [] # Promote available and relevant scalar coordinates for coord in new_cube.aux_coords[::-1]: if coord in cube.dim_coords: new_cube = iris.util.new_axis(new_cube, coord) # Ensure dimension order matches; if lengths are unequal a coordinate # is missing, so raise an appropriate error. cube_dimension_order = {coord.name(): cube.coord_dims(coord.name())[0] for coord in cube.dim_coords} correct_order = [] new_cube_only_dims = [] new_cube_dim_names = [coord.name() for coord in new_cube.dim_coords] cube_dim_names = [coord.name() for coord in cube.dim_coords] for coord_name in new_cube_dim_names: if coord_name in cube_dim_names: correct_order.append(cube_dimension_order[coord_name]) else: if coord_name in exception_coordinates: new_coord_dim = new_cube.coord_dims(coord_name)[0] new_cube_only_dims.append(new_coord_dim) else: msg = ("The coordinate: {0} is within new_cube, " "however, this is not within the original " "cube. As {0} is not within the permitted " "exceptions: {1}, this is not allowed. " "\nnew_cube: {2}" "\ncube: {3}").format( coord_name, exception_coordinates, new_cube, cube) raise InvalidCubeError(msg) correct_order = np.array(correct_order) for dim in new_cube_only_dims: correct_order[correct_order >= dim] += 1 correct_order = np.insert(correct_order, dim, dim) if (len(cube_dimension_order.keys())+len(exception_coordinates) == len(correct_order)): new_cube.transpose(correct_order) else: msg = ('The number of dimension coordinates within the new cube ' 'do not match the number of dimension coordinates within the ' 'original cube plus the number of exception coordinates. ' '\n input cube shape {} returned cube shape {}'.format( cube.shape, new_cube.shape)) raise CoordinateNotFoundError(msg) return new_cube