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_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_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 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(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)