def test_without_orographic_enhancement(self):
"""Test plugin returns the correct advected forecast cube.
In this case we have 600m grid spacing in our cubes, and 1m/s
advection velocities in the x and y direction, so after 10 minutes,
our precipitation will have moved exactly one grid square along
each axis."""
input_cube = self.precip_cube.copy()
input_cube.rename("air_temperature")
input_cube.units = "K"
plugin = CreateExtrapolationForecast(
input_cube, self.vel_x, self.vel_y)
result = plugin.extrapolate(10)
expected_result = np.array([[np.nan, np.nan, np.nan],
[np.nan, 1, 2],
[np.nan, 1, 1],
[np.nan, 0, 2]], dtype=np.float32)
expected_result = np.ma.masked_invalid(expected_result)
expected_forecast_period = np.array([600], dtype=np.int64)
# Check we get the expected result, and the correct time coordinates.
self.assertArrayEqual(expected_result.mask, result.data.mask)
self.assertArrayAlmostEqual(expected_result.data, result.data.data)
self.assertArrayAlmostEqual(
result.coord("forecast_period").points, expected_forecast_period)
self.assertEqual(result.coord("forecast_period").units, "seconds")
self.assertEqual(result.coord("forecast_reference_time").points,
input_cube.coord("forecast_reference_time").points)
self.assertEqual(result.coord("time").points,
input_cube.coord("time").points+600)
def test_with_orographic_enhancement(self):
"""Test plugin returns the correct advected forecast cube, with
orographic enhancement.
In this case we have 600m grid spacing in our cubes, and 1m/s
advection velocities in the x and y direction, so after 10 minutes,
our precipitation will have moved exactly one grid square along
each axis.
The orographic enhancement has been removed before advecting, then
added back on afterwards, leading to a different end result."""
plugin = CreateExtrapolationForecast(
self.precip_cube, self.vel_x, self.vel_y,
orographic_enhancement_cube=self.oe_cube)
result = plugin.extrapolate(10)
expected_result = np.array([[np.nan, np.nan, np.nan],
[np.nan, 1.03125, 1.0],
[np.nan, 1.0, 0.03125],
[np.nan, 0, 2.0]], dtype=np.float32)
expected_result = np.ma.masked_invalid(expected_result)
expected_forecast_period = np.array([600], dtype=np.int64)
# Check we get the expected result, and the correct time coordinates.
self.assertArrayEqual(expected_result.mask, result.data.mask)
self.assertArrayAlmostEqual(expected_result.data, result.data.data)
self.assertArrayAlmostEqual(
result.coord("forecast_period").points, expected_forecast_period)
self.assertEqual(result.coord("forecast_period").units, "seconds")
self.assertEqual(
result.coord("forecast_reference_time").points,
self.precip_cube.coord("forecast_reference_time").points)
self.assertEqual(result.coord("time").points,
self.precip_cube.coord("time").points+600)
def test_raises_error(self):
"""Test an error is raised if no leadtime is provided"""
plugin = CreateExtrapolationForecast(
self.precip_cube, self.vel_x, self.vel_y,
orographic_enhancement_cube=self.oe_cube)
message = ("leadtime_minutes must be provided in order to "
"produce an extrapolated forecast")
with self.assertRaisesRegex(ValueError, message):
plugin.extrapolate()
def test_no_orographic_enhancement(self):
"""Test what happens if no orographic enhancement cube is provided"""
message = ("For precipitation fields, orographic enhancement cube "
"must be supplied.")
with self.assertRaisesRegex(ValueError, message):
CreateExtrapolationForecast(
self.precip_cube, self.vel_x, self.vel_y)
def test_basic(self):
"""Test for simple case where __init__ does not change the input."""
# Change the input cube so no orographic enhancement is expected.
input_cube = self.precip_cube.copy()
input_cube.rename("air_temperature")
input_cube.units = "K"
plugin = CreateExtrapolationForecast(input_cube.copy(), self.vel_x, self.vel_y)
self.assertEqual(input_cube, plugin.input_cube)
self.assertIsNone(plugin.orographic_enhancement_cube)
self.assertIsInstance(plugin.advection_plugin, AdvectField)
def process(cube: cli.inputcube,
advection_velocity: inputadvection,
orographic_enhancement: cli.inputcube = None,
*,
attributes_config: cli.inputjson = None,
max_lead_time: int = 360, lead_time_interval: int = 15):
"""Module to extrapolate input cubes given advection velocity fields.
Args:
cube (iris.cube.Cube):
The data to be advected.
advection_velocity (iris.cube.CubeList):
Advection cubes of U and V.
These must have the names of.
precipitation_advection_x_velocity
precipitation_advection_y_velocity
orographic_enhancement (iris.cube.Cube):
Cube containing orographic enhancement forecasts for the lead times
at which an extrapolation nowcast is required.
attributes_config (dict):
Dictionary containing the required changes to the attributes.
max_lead_time (int):
Maximum lead time required (mins).
lead_time_interval (int):
Interval between required lead times (mins).
Returns:
iris.cube.CubeList:
New cubes with updated time and extrapolated data.
"""
from improver.nowcasting.forecasting import CreateExtrapolationForecast
from improver.utilities.cube_manipulation import merge_cubes
u_cube, v_cube = advection_velocity
# extrapolate input data to required lead times
forecast_plugin = CreateExtrapolationForecast(
cube, u_cube, v_cube, orographic_enhancement,
attributes_dict=attributes_config)
forecast_cubes = forecast_plugin.process(lead_time_interval, max_lead_time)
return merge_cubes(forecast_cubes)
def test_orographic_enhancement(self):
"""Test what happens if an orographic enhancement cube is provided"""
plugin = CreateExtrapolationForecast(
self.precip_cube, self.vel_x, self.vel_y,
orographic_enhancement_cube=self.oe_cube.copy())
expected_data = np.array([[0.03125, 1.0, 0.03125],
[1.0, 0.03125, 1.0],
[0.0, 2.0, 0.0],
[1.0, 2.0, 1.0]], dtype=np.float32)
self.assertEqual(self.precip_cube.metadata, plugin.input_cube.metadata)
self.assertArrayAlmostEqual(plugin.input_cube.data, expected_data)
self.assertEqual(plugin.orographic_enhancement_cube, self.oe_cube)
self.assertIsInstance(plugin.advection_plugin, AdvectField)
def test_basic_with_metadata_dict(self):
"""Test for simple case where __init__ does not change the input and
we amend the attributes."""
# Change the input cube so no orographic enhancement is expected.
input_cube = self.precip_cube.copy()
input_cube.rename("air_temperature")
input_cube.units = "K"
attributes = {"source": "IMPROVER"}
plugin = CreateExtrapolationForecast(
input_cube.copy(), self.vel_x, self.vel_y,
attributes_dict=attributes)
self.assertEqual(input_cube, plugin.input_cube)
self.assertEqual(plugin.orographic_enhancement_cube, None)
self.assertIsInstance(plugin.advection_plugin, AdvectField)
self.assertEqual(plugin.advection_plugin.attributes_dict, attributes)
def test_basic(self):
"""Test string representation"""
plugin = CreateExtrapolationForecast(
self.precip_cube, self.vel_x, self.vel_y,
orographic_enhancement_cube=self.oe_cube)
result = str(plugin)
expected_result = (
"<CreateExtrapolationForecast: input_cube = <iris 'Cube' of "
"lwe_precipitation_rate / (mm/hr) (projection_y_coordinate: 4; "
"projection_x_coordinate: 3)>, orographic_enhancement_cube = "
"<iris 'Cube' of orographic_enhancement / (mm/hr) "
"(projection_y_coordinate: 4; projection_x_coordinate: 3)>, "
"advection_plugin = <AdvectField: vel_x=<iris "
"'Cube' of advection_velocity_x / (m s-1) "
"(projection_y_coordinate: 4; projection_x_coordinate: 3)>, "
"vel_y=<iris 'Cube' of advection_velocity_y / (m s-1) "
"(projection_y_coordinate: 4; projection_x_coordinate: 3)>, "
"attributes_dict={}>>"
)
self.assertEqual(result, expected_result)
def test_velocity_mismatch(self):
"""Tests that x and y velocities are both used."""
with self.assertRaisesRegex(TypeError, ".*x velocity or y.*"):
CreateExtrapolationForecast(
self.precip_cube, self.vel_x, None)
def process(original_cube_list,
orographic_enhancement_cube=None,
metadata_dict=None,
ofc_box_size=30,
smart_smoothing_iterations=100,
extrapolate=False,
max_lead_time=360,
lead_time_interval=15):
"""Calculates optical flow and can (optionally) extrapolate data.
Calculates optical flow components from input fields and (optionally)
extrapolate to required lead times.
Args:
original_cube_list (iris.cube.CubeList):
Cubelist from which to calculate optical flow velocities.
The cubes require a 'time' coordinate on which they are sorted,
so the order of cubes does not matter.
orographic_enhancement_cube (iris.cube.Cube):
Cube containing the orographic enhancement fields.
Default is None.
metadata_dict (dict):
Dictionary containing required changes to the metadata.
Information describing the intended contents of the dictionary is
available in improver.utilities.cube_metadata.amend_metadata.
Every output cube will have the metadata_dict applied.
Default is None.
ofc_box_size (int):
Size of square 'box' (in grid spaces) within which to solve
the optical flow equations.
Default is 30.
smart_smoothing_iterations (int):
Number of iterations to perform in enforcing smoothness constraint
for optical flow velocities.
Default is 100.
extrapolate (bool):
If True, advects current data forward to specified lead times.
Default is False.
max_lead_time (int):
Maximum lead time required (mins). Ignored unless extrapolate is
True.
Default is 360.
lead_time_interval (int):
Interval between required lead times (mins). Ignored unless
extrapolate is True.
Default is 15.
Returns:
(tuple): tuple containing:
**forecast_cubes** (list<Cube>):
List of Cubes if extrapolate is True, else None.
**u_and_v_mean** (list<Cube>):
List of the umean and vmean cubes.
Raises:
ValueError:
If there is no oe_cube but a cube is called 'precipitation_rate'.
"""
if orographic_enhancement_cube:
cube_list = ApplyOrographicEnhancement("subtract").process(
original_cube_list, orographic_enhancement_cube)
else:
cube_list = original_cube_list
if any("precipitation_rate" in cube.name() for cube in cube_list):
cube_names = [cube.name() for cube in cube_list]
msg = ("For precipitation fields, orographic enhancement "
"filepaths must be supplied. The names of the cubes "
"supplied were: {}".format(cube_names))
raise ValueError(msg)
# order input files by validity time
cube_list.sort(key=lambda x: x.coord("time").points[0])
time_coord = cube_list[-1].coord("time")
# calculate optical flow velocities from T-1 to T and T-2 to T-1
ofc_plugin = OpticalFlow(iterations=smart_smoothing_iterations,
metadata_dict=metadata_dict)
u_cubes = iris.cube.CubeList([])
v_cubes = iris.cube.CubeList([])
for older_cube, newer_cube in zip(cube_list[:-1], cube_list[1:]):
ucube, vcube = ofc_plugin.process(older_cube,
newer_cube,
boxsize=ofc_box_size)
u_cubes.append(ucube)
v_cubes.append(vcube)
# average optical flow velocity components
u_cube = u_cubes.merge_cube()
u_mean = u_cube.collapsed("time", iris.analysis.MEAN)
u_mean.coord("time").points = time_coord.points
u_mean.coord("time").units = time_coord.units
v_cube = v_cubes.merge_cube()
v_mean = v_cube.collapsed("time", iris.analysis.MEAN)
v_mean.coord("time").points = time_coord.points
v_mean.coord("time").units = time_coord.units
u_and_v_mean = [u_mean, v_mean]
forecast_cubes = []
if extrapolate:
# generate list of lead times in minutes
lead_times = np.arange(0, max_lead_time + 1, lead_time_interval)
forecast_plugin = CreateExtrapolationForecast(
original_cube_list[-1],
u_mean,
v_mean,
orographic_enhancement_cube=orographic_enhancement_cube,
metadata_dict=metadata_dict)
# extrapolate input data to required lead times
for lead_time in lead_times:
forecast_cubes.append(
forecast_plugin.extrapolate(leadtime_minutes=lead_time))
return forecast_cubes, u_and_v_mean
def main(argv=None):
"""Extrapolate data forward in time."""
parser = ArgParser(
description="Extrapolate input data to required lead times.")
parser.add_argument("input_filepath", metavar="INPUT_FILEPATH",
type=str, help="Path to input NetCDF file.")
group = parser.add_mutually_exclusive_group()
group.add_argument("--output_dir", metavar="OUTPUT_DIR", type=str,
default="", help="Directory to write output files.")
group.add_argument("--output_filepaths", nargs="+", type=str,
help="List of full paths to output nowcast files, in "
"order of increasing lead time.")
optflw = parser.add_argument_group('Advect using files containing the x '
' and y components of the velocity')
optflw.add_argument("--eastward_advection_filepath", type=str, help="Path"
" to input file containing Eastward advection "
"velocities.")
optflw.add_argument("--northward_advection_filepath", type=str, help="Path"
" to input file containing Northward advection "
"velocities.")
speed = parser.add_argument_group('Advect using files containing speed and'
' direction')
speed.add_argument("--advection_speed_filepath", type=str, help="Path"
" to input file containing advection speeds,"
" usually wind speeds, on multiple pressure levels.")
speed.add_argument("--advection_direction_filepath", type=str,
help="Path to input file containing the directions from"
" which advection speeds are coming (180 degrees from"
" the direction in which the speed is directed). The"
" directions should be on the same grid as the input"
" speeds, including the same vertical levels.")
speed.add_argument("--pressure_level", type=int, default=75000, help="The"
" pressure level in Pa to extract from the multi-level"
" advection_speed and advection_direction files. The"
" velocities at this level are used for advection.")
parser.add_argument("--orographic_enhancement_filepaths", nargs="+",
type=str, default=None, help="List or wildcarded "
"file specification to the input orographic "
"enhancement files. Orographic enhancement files are "
"compulsory for precipitation fields.")
parser.add_argument("--json_file", metavar="JSON_FILE", default=None,
help="Filename for the json file containing "
"required changes to the metadata. Information "
"describing the intended contents of the json file "
"is available in "
"improver.utilities.cube_metadata.amend_metadata."
"Every output cube will have the metadata_dict "
"applied. Defaults to None.", type=str)
parser.add_argument("--max_lead_time", type=int, default=360,
help="Maximum lead time required (mins).")
parser.add_argument("--lead_time_interval", type=int, default=15,
help="Interval between required lead times (mins).")
accumulation_args = parser.add_argument_group(
'Calculate accumulations from advected fields')
accumulation_args.add_argument(
"--accumulation_fidelity", type=int, default=0,
help="If set, this CLI will additionally return accumulations"
" calculated from the advected fields. This fidelity specifies the"
" time interval in minutes between advected fields that is used to"
" calculate these accumulations. This interval must be a factor of"
" the lead_time_interval.")
accumulation_args.add_argument(
"--accumulation_units", type=str, default='m',
help="Desired units in which the accumulations should be expressed,"
"e.g. mm")
args = parser.parse_args(args=argv)
upath, vpath = (args.eastward_advection_filepath,
args.northward_advection_filepath)
spath, dpath = (args.advection_speed_filepath,
args.advection_direction_filepath)
# load files and initialise advection plugin
input_cube = load_cube(args.input_filepath)
if (upath and vpath) and not (spath or dpath):
ucube = load_cube(upath)
vcube = load_cube(vpath)
elif (spath and dpath) and not (upath or vpath):
level_constraint = Constraint(pressure=args.pressure_level)
try:
scube = load_cube(spath, constraints=level_constraint)
dcube = load_cube(dpath, constraints=level_constraint)
except ValueError as err:
raise ValueError(
'{} Unable to extract specified pressure level from given '
'speed and direction files.'.format(err))
ucube, vcube = ResolveWindComponents().process(scube, dcube)
else:
raise ValueError('Cannot mix advection component velocities with speed'
' and direction')
oe_cube = None
if args.orographic_enhancement_filepaths:
oe_cube = load_cube(args.orographic_enhancement_filepaths)
metadata_dict = None
if args.json_file:
# Load JSON file for metadata amendments.
with open(args.json_file, 'r') as input_file:
metadata_dict = json.load(input_file)
# generate list of lead times in minutes
lead_times = np.arange(0, args.max_lead_time+1,
args.lead_time_interval)
if args.output_filepaths:
if len(args.output_filepaths) != len(lead_times):
raise ValueError("Require exactly one output file name for each "
"forecast lead time")
# determine whether accumulations are also to be returned.
time_interval = args.lead_time_interval
if args.accumulation_fidelity > 0:
fraction, _ = np.modf(args.lead_time_interval /
args.accumulation_fidelity)
if fraction != 0:
msg = ("The specified lead_time_interval ({}) is not cleanly "
"divisible by the specified accumulation_fidelity ({}). As "
"a result the lead_time_interval cannot be constructed from"
" accumulation cubes at this fidelity.".format(
args.lead_time_interval, args.accumulation_fidelity))
raise ValueError(msg)
time_interval = args.accumulation_fidelity
lead_times = np.arange(0, args.max_lead_time+1, time_interval)
lead_time_filter = args.lead_time_interval // time_interval
forecast_plugin = CreateExtrapolationForecast(
input_cube, ucube, vcube, orographic_enhancement_cube=oe_cube,
metadata_dict=metadata_dict)
# extrapolate input data to required lead times
forecast_cubes = iris.cube.CubeList()
for i, lead_time in enumerate(lead_times):
forecast_cubes.append(
forecast_plugin.extrapolate(leadtime_minutes=lead_time))
# return rate cubes
for i, cube in enumerate(forecast_cubes[::lead_time_filter]):
# save to a suitably-named output file
if args.output_filepaths:
file_name = args.output_filepaths[i]
else:
file_name = os.path.join(
args.output_dir, generate_file_name(cube))
save_netcdf(cube, file_name)
# calculate accumulations if required
if args.accumulation_fidelity > 0:
plugin = Accumulation(accumulation_units=args.accumulation_units,
accumulation_period=args.lead_time_interval * 60)
accumulation_cubes = plugin.process(forecast_cubes)
# return accumulation cubes
for i, cube in enumerate(accumulation_cubes):
file_name = os.path.join(args.output_dir, generate_file_name(cube))
save_netcdf(cube, file_name)
def process(cube: cli.inputcube,
advection_velocity: inputadvection,
orographic_enhancement: cli.inputcube,
*,
attributes_config: cli.inputjson = None,
max_lead_time=360,
lead_time_interval=15,
accumulation_period=15,
accumulation_units='m'):
"""Module to extrapolate and accumulate the weather with 1 min fidelity.
Args:
cube (iris.cube.Cube):
The input Cube to be processed.
advection_velocity (iris.cube.CubeList):
Advection cubes of U and V.
These must have the names of either:
precipitation_advection_x_velocity or grid_eastward_wind
precipitation_advection_y_velocity or grid_northward_wind
orographic_enhancement (iris.cube.Cube):
Cube containing the orographic enhancement fields. May have data
for multiple times in the cube.
attributes_config (dict):
Dictionary containing the required changes to the attributes.
max_lead_time (int):
Maximum lead time required (mins).
lead_time_interval (int):
Interval between required lead times (mins).
accumulation_period (int):
The period over which the accumulation is calculated (mins).
Only full accumulation periods will be computed. At lead times
that are shorter than the accumulation period, no accumulation
output will be produced.
accumulation_units (str):
Desired units in which the accumulations should be expressed.
e.g. 'mm'
Returns:
iris.cube.CubeList:
New cubes with accumulated data.
Raises:
ValueError:
If advection_velocity doesn't contain x and y velocity.
"""
import numpy as np
from improver.nowcasting.accumulation import Accumulation
from improver.nowcasting.forecasting import CreateExtrapolationForecast
from improver.utilities.cube_manipulation import MergeCubes
u_cube, v_cube = advection_velocity
if not (u_cube and v_cube):
raise ValueError("Neither u_cube or v_cube can be None")
# extrapolate input data to the maximum required lead time
plugin = CreateExtrapolationForecast(cube,
u_cube,
v_cube,
orographic_enhancement,
attributes_dict=attributes_config)
forecast_cubes = plugin(ACCUMULATION_FIDELITY, max_lead_time)
lead_times = (np.arange(lead_time_interval, max_lead_time + 1,
lead_time_interval))
# Accumulate high frequency rate into desired accumulation intervals.
plugin = Accumulation(accumulation_units=accumulation_units,
accumulation_period=accumulation_period * 60,
forecast_periods=lead_times * 60)
result = plugin(forecast_cubes)
return MergeCubes()(result)
def process(input_cube,
u_cube,
v_cube,
speed_cube,
direction_cube,
orographic_enhancement_cube=None,
metadata_dict=None,
max_lead_time=360,
lead_time_interval=15,
accumulation_fidelity=0,
accumulation_period=15,
accumulation_units='m'):
"""Module to extrapolate input cubes given advection velocity fields.
Args:
input_cube (iris.cube.Cube):
The input Cube to be processed.
u_cube (iris.cube.Cube):
Cube with the velocities in the x direction.
Must be used with v_cube.
s_cube and d_cube must be None.
v_cube (iris.cube.Cube):
Cube with the velocities in the y direction.
Must be used with u_cube.
s_cube and d_cube must be None.
speed_cube (iris.cube.Cube):
Cube containing advection speeds, usually wind speed.
Must be used with d_cube.
u_cube and v_cube must be None.
direction_cube (iris.cube.Cube):
Cube from which advection speeds are coming. The directions
should be on the same grid as the input speeds, including the same
vertical levels.
Must be used with d_cube.
u_cube and v_cube must be None.
orographic_enhancement_cube (iris.cube.Cube):
Cube containing the orographic enhancement fields. May have data
for multiple times in the cube.
Default is None.
metadata_dict (dict):
Dictionary containing the required changes to the metadata.
Information describing the intended contents of the dictionary
is available in improver.utilities.cube_metadata.amend_metadata.
Every output cube will have the metadata_dict applied.
Default is None.
max_lead_time (int):
Maximum lead time required (mins).
Default is 360.
lead_time_interval (int):
Interval between required lead times (mins).
Default is 15.
accumulation_fidelity (int):
If set, this will additionally return accumulations calculated
from the advected fields. This fidelity specifies the time
interval in minutes between advected fields that is used to
calculate these accumulations. This interval must be a factor of
the lead_time_interval.
Default is 0.
accumulation_period (int):
The period over which the accumulation is calculated (mins).
Only full accumulation periods will be computed. At lead times
that are shorter than the accumulation period, no accumulation
output will be produced.
accumulation_units (str):
Desired units in which the accumulations should be expressed.
e.g. 'mm'
Default is 'm'.
Returns:
(tuple) tuple containing:
**accumulation_cubes** (iris.cube.Cubelist):
A cubelist containing precipitation accumulation cubes where
the accumulation periods are determined by the
lead_time_interval.
**forecast_to_return** (iris.cube.Cubelist):
New cubes with updated time and extrapolated data.
Raises:
ValueError:
can either use s_cube and d_cube or u_cube and v_cube.
Therefore: (s and d)⊕(u and v)
ValueError:
If accumulation_fidelity is greater than 0 and max_lead_time is not
cleanly divisible by accumulation_fidelity.
"""
if (speed_cube and direction_cube) and not (u_cube or v_cube):
u_cube, v_cube = ResolveWindComponents().process(
speed_cube, direction_cube)
elif (u_cube or v_cube) and (speed_cube or direction_cube):
raise ValueError('Cannot mix advection component velocities with speed'
' and direction')
# generate list of lead times in minutes
lead_times = np.arange(0, max_lead_time + 1, lead_time_interval)
# determine whether accumulations are also to be returned.
time_interval = lead_time_interval
if accumulation_fidelity > 0:
fraction, _ = np.modf(max_lead_time / accumulation_fidelity)
if fraction != 0:
msg = ("The specified lead_time_interval ({}) is not cleanly "
"divisible by the specified accumulation_fidelity ({}). As "
"a result the lead_time_interval cannot be constructed from"
" accumulation cubes at this fidelity.")
raise ValueError(
msg.format(lead_time_interval, accumulation_fidelity))
time_interval = accumulation_fidelity
lead_times = np.arange(0, max_lead_time + 1, time_interval)
lead_time_filter = lead_time_interval // time_interval
forecast_plugin = CreateExtrapolationForecast(
input_cube,
u_cube,
v_cube,
orographic_enhancement_cube=orographic_enhancement_cube,
metadata_dict=metadata_dict)
# extrapolate input data to required lead times
forecast_cubes = iris.cube.CubeList()
for lead_time in lead_times:
forecast_cubes.append(
forecast_plugin.extrapolate(leadtime_minutes=lead_time))
forecast_to_return = forecast_cubes[::lead_time_filter].copy()
# return rate cubes
# calculate accumulations if required
accumulation_cubes = None
if accumulation_fidelity > 0:
lead_times = (np.arange(lead_time_interval, max_lead_time + 1,
lead_time_interval))
plugin = Accumulation(accumulation_units=accumulation_units,
accumulation_period=accumulation_period * 60,
forecast_periods=lead_times * 60)
accumulation_cubes = plugin.process(forecast_cubes)
return accumulation_cubes, forecast_to_return
def main(argv=None):
"""Calculate optical flow advection velocities and (optionally)
extrapolate data."""
parser = ArgParser(
description="Calculate optical flow components from input fields "
"and (optionally) extrapolate to required lead times.")
parser.add_argument("input_filepaths", metavar="INPUT_FILEPATHS",
nargs=3, type=str, help="Paths to the input radar "
"files. There should be 3 input files at T, T-1 and "
"T-2 from which to calculate optical flow velocities. "
"The files require a 'time' coordinate on which they "
"are sorted, so the order of inputs does not matter.")
parser.add_argument("--output_dir", metavar="OUTPUT_DIR", type=str,
default='', help="Directory to write all output files,"
" or only advection velocity components if "
"NOWCAST_FILEPATHS is specified.")
parser.add_argument("--nowcast_filepaths", nargs="+", type=str,
default=None, help="Optional list of full paths to "
"output nowcast files. Overrides OUTPUT_DIR. Ignored "
"unless '--extrapolate' is set.")
parser.add_argument("--orographic_enhancement_filepaths", nargs="+",
type=str, default=None, help="List or wildcarded "
"file specification to the input orographic "
"enhancement files. Orographic enhancement files are "
"compulsory for precipitation fields.")
parser.add_argument("--json_file", metavar="JSON_FILE", default=None,
help="Filename for the json file containing "
"required changes to the metadata. Information "
"describing the intended contents of the json file "
"is available in "
"improver.utilities.cube_metadata.amend_metadata."
"Every output cube will have the metadata_dict "
"applied. Defaults to None.", type=str)
# OpticalFlow plugin configurable parameters
parser.add_argument("--ofc_box_size", type=int, default=30, help="Size of "
"square 'box' (in grid squares) within which to solve "
"the optical flow equations.")
parser.add_argument("--smart_smoothing_iterations", type=int, default=100,
help="Number of iterations to perform in enforcing "
"smoothness constraint for optical flow velocities.")
# AdvectField options
parser.add_argument("--extrapolate", action="store_true", default=False,
help="Optional flag to advect current data forward to "
"specified lead times.")
parser.add_argument("--max_lead_time", type=int, default=360,
help="Maximum lead time required (mins). Ignored "
"unless '--extrapolate' is set.")
parser.add_argument("--lead_time_interval", type=int, default=15,
help="Interval between required lead times (mins). "
"Ignored unless '--extrapolate' is set.")
args = parser.parse_args(args=argv)
# read input data
original_cube_list = load_cubelist(args.input_filepaths)
if args.orographic_enhancement_filepaths:
# Subtract orographic enhancement
oe_cube = load_cube(args.orographic_enhancement_filepaths)
cube_list = ApplyOrographicEnhancement("subtract").process(
original_cube_list, oe_cube)
else:
cube_list = original_cube_list
if any("precipitation_rate" in cube.name() for cube in cube_list):
cube_names = [cube.name() for cube in cube_list]
msg = ("For precipitation fields, orographic enhancement "
"filepaths must be supplied. The names of the cubes "
"supplied were: {}".format(cube_names))
raise ValueError(msg)
# order input files by validity time
cube_list.sort(key=lambda x: x.coord("time").points[0])
time_coord = cube_list[-1].coord("time")
metadata_dict = None
if args.json_file:
# Load JSON file for metadata amendments.
with open(args.json_file, 'r') as input_file:
metadata_dict = json.load(input_file)
# calculate optical flow velocities from T-1 to T and T-2 to T-1
ofc_plugin = OpticalFlow(iterations=args.smart_smoothing_iterations,
metadata_dict=metadata_dict)
ucubes = iris.cube.CubeList([])
vcubes = iris.cube.CubeList([])
for older_cube, newer_cube in zip(cube_list[:-1], cube_list[1:]):
ucube, vcube = ofc_plugin.process(older_cube, newer_cube,
boxsize=args.ofc_box_size)
ucubes.append(ucube)
vcubes.append(vcube)
# average optical flow velocity components
ucube = ucubes.merge_cube()
umean = ucube.collapsed("time", iris.analysis.MEAN)
umean.coord("time").points = time_coord.points
umean.coord("time").units = time_coord.units
vcube = vcubes.merge_cube()
vmean = vcube.collapsed("time", iris.analysis.MEAN)
vmean.coord("time").points = time_coord.points
vmean.coord("time").units = time_coord.units
# save mean optical flow components as netcdf files
for wind_cube in [umean, vmean]:
file_name = generate_file_name(wind_cube)
save_netcdf(wind_cube, os.path.join(args.output_dir, file_name))
# advect latest input data to the required lead times
if args.extrapolate:
# generate list of lead times in minutes
lead_times = np.arange(0, args.max_lead_time+1,
args.lead_time_interval)
if args.nowcast_filepaths:
if len(args.nowcast_filepaths) != len(lead_times):
raise ValueError("Require exactly one output file name for "
"each forecast lead time")
forecast_plugin = CreateExtrapolationForecast(
original_cube_list[-1], umean, vmean,
orographic_enhancement_cube=oe_cube, metadata_dict=metadata_dict)
# extrapolate input data to required lead times
for i, lead_time in enumerate(lead_times):
forecast_cube = forecast_plugin.extrapolate(
leadtime_minutes=lead_time)
# save to a suitably-named output file
if args.nowcast_filepaths:
file_name = args.nowcast_filepaths[i]
else:
file_name = os.path.join(
args.output_dir, generate_file_name(forecast_cube))
save_netcdf(forecast_cube, file_name)
