fix: make sure netcdf exporters can handle list of timesteps (#369)

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"]

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: