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
