diff --git a/pysteps/io/exporters.py b/pysteps/io/exporters.py index 0d26b0fd3..38b2159a7 100644 --- a/pysteps/io/exporters.py +++ b/pysteps/io/exporters.py @@ -392,9 +392,10 @@ def initialize_forecast_exporter_netcdf( Start date of the forecast. timestep: int Time step of the forecast (minutes). - n_timesteps: int - Number of time steps in the forecast this argument is ignored if - incremental is set to 'timestep'. + n_timesteps: int or list of integers + Number of time steps to forecast or a list of time steps for which the + forecasts are computed (relative to the input time step). The elements of + the list are required to be in ascending order. shape: tuple of int Two-element tuple defining the shape (height,width) of the forecast grids. @@ -460,8 +461,14 @@ def initialize_forecast_exporter_netcdf( + "'timestep' or 'member'" ) + n_timesteps_is_list = isinstance(n_timesteps, list) + if n_timesteps_is_list: + num_timesteps = len(n_timesteps) + else: + num_timesteps = n_timesteps + if incremental == "timestep": - n_timesteps = None + num_timesteps = None elif incremental == "member": n_ens_members = None elif incremental is not None: @@ -498,7 +505,7 @@ def initialize_forecast_exporter_netcdf( h, w = shape ncf.createDimension("ens_number", size=n_ens_members) - ncf.createDimension("time", size=n_timesteps) + ncf.createDimension("time", size=num_timesteps) ncf.createDimension("y", size=h) ncf.createDimension("x", size=w) @@ -585,7 +592,10 @@ def initialize_forecast_exporter_netcdf( var_time = ncf.createVariable("time", int, dimensions=("time",)) if incremental != "timestep": - var_time[:] = [i * timestep * 60 for i in range(1, n_timesteps + 1)] + if n_timesteps_is_list: + var_time[:] = np.array(n_timesteps) * timestep * 60 + else: + var_time[:] = [i * timestep * 60 for i in range(1, n_timesteps + 1)] var_time.long_name = "forecast time" startdate_str = datetime.strftime(startdate, "%Y-%m-%d %H:%M:%S") var_time.units = "seconds since %s" % startdate_str @@ -635,7 +645,8 @@ def initialize_forecast_exporter_netcdf( exporter["timestep"] = timestep exporter["metadata"] = metadata exporter["incremental"] = incremental - exporter["num_timesteps"] = n_timesteps + exporter["num_timesteps"] = num_timesteps + exporter["timesteps"] = n_timesteps exporter["num_ens_members"] = n_ens_members exporter["shape"] = shape @@ -853,7 +864,12 @@ def _export_netcdf(field, exporter): else: var_f[var_f.shape[0], :, :] = field var_time = exporter["var_time"] - var_time[len(var_time) - 1] = len(var_time) * exporter["timestep"] * 60 + if isinstance(exporter["timesteps"], list): + var_time[len(var_time) - 1] = ( + exporter["timesteps"][len(var_time) - 1] * exporter["timestep"] * 60 + ) + else: + var_time[len(var_time) - 1] = len(var_time) * exporter["timestep"] * 60 else: var_f[var_f.shape[0], :, :, :] = field var_ens_num = exporter["var_ens_num"] diff --git a/pysteps/tests/test_exporters.py b/pysteps/tests/test_exporters.py index 010d7b75a..f72c248bb 100644 --- a/pysteps/tests/test_exporters.py +++ b/pysteps/tests/test_exporters.py @@ -23,14 +23,17 @@ "fill_value", "scale_factor", "offset", + "n_timesteps", ) exporter_arg_values = [ - (1, None, np.float32, None, None, None), - (1, "timestep", np.float32, 65535, None, None), - (2, None, np.float32, 65535, None, None), - (2, "timestep", np.float32, None, None, None), - (2, "member", np.float64, None, 0.01, 1.0), + (1, None, np.float32, None, None, None, 3), + (1, "timestep", np.float32, 65535, None, None, 3), + (2, None, np.float32, 65535, None, None, 3), + (2, None, np.float32, 65535, None, None, [1, 2, 4]), + (2, "timestep", np.float32, None, None, None, 3), + (2, "timestep", np.float32, None, None, None, [1, 2, 4]), + (2, "member", np.float64, None, 0.01, 1.0, 3), ] @@ -54,7 +57,7 @@ def test_get_geotiff_filename(): @pytest.mark.parametrize(exporter_arg_names, exporter_arg_values) def test_io_export_netcdf_one_member_one_time_step( - n_ens_members, incremental, datatype, fill_value, scale_factor, offset + n_ens_members, incremental, datatype, fill_value, scale_factor, offset, n_timesteps ): """ Test the export netcdf. @@ -75,7 +78,6 @@ def test_io_export_netcdf_one_member_one_time_step( file_path = os.path.join(outpath, outfnprefix + ".nc") startdate = metadata["timestamps"][0] timestep = metadata["accutime"] - n_timesteps = 3 shape = precip.shape[1:] exporter = initialize_forecast_exporter_netcdf( @@ -100,7 +102,11 @@ def test_io_export_netcdf_one_member_one_time_step( if incremental == None: export_forecast_dataset(precip, exporter) if incremental == "timestep": - for t in range(n_timesteps): + if isinstance(n_timesteps, list): + timesteps = len(n_timesteps) + else: + timesteps = n_timesteps + for t in range(timesteps): if n_ens_members > 1: export_forecast_dataset(precip[:, t, :, :], exporter) else: