15
15
from pysteps .io .exporters import initialize_forecast_exporter_netcdf
16
16
from pysteps .tests .helpers import get_precipitation_fields , get_invalid_mask
17
17
18
+ # Test arguments
19
+ exporter_arg_names = ("n_ens_members" , "incremental" )
20
+
21
+ exporter_arg_values = [
22
+ (1 , None ),
23
+ (1 , "timestep" ),
24
+ (2 , None ),
25
+ (2 , "timestep" ),
26
+ (2 , "member" ),
27
+ ]
28
+
18
29
19
30
def test_get_geotiff_filename ():
20
31
"""Test the geotif name generator."""
@@ -34,7 +45,8 @@ def test_get_geotiff_filename():
34
45
assert expected == file_name
35
46
36
47
37
- def test_io_export_netcdf_one_member_one_time_step ():
48
+ @pytest .mark .parametrize (exporter_arg_names , exporter_arg_values )
49
+ def test_io_export_netcdf_one_member_one_time_step (n_ens_members , incremental ):
38
50
"""
39
51
Test the export netcdf.
40
52
Also, test that the exported file can be read by the importer.
@@ -43,20 +55,20 @@ def test_io_export_netcdf_one_member_one_time_step():
43
55
pytest .importorskip ("pyproj" )
44
56
45
57
precip , metadata = get_precipitation_fields (
46
- return_raw = True , metadata = True , source = "fmi"
58
+ num_prev_files = 2 , return_raw = True , metadata = True , source = "fmi"
47
59
)
48
- precip = precip .squeeze ()
49
60
50
61
invalid_mask = get_invalid_mask (precip )
51
62
52
- # save it back to disk
53
63
with tempfile .TemporaryDirectory () as outpath :
64
+ # save it back to disk
54
65
outfnprefix = "test_netcdf_out"
55
66
file_path = os .path .join (outpath , outfnprefix + ".nc" )
56
67
startdate = metadata ["timestamps" ][0 ]
57
68
timestep = metadata ["accutime" ]
58
- n_timesteps = 1
59
- shape = precip .shape
69
+ n_timesteps = 3
70
+ shape = precip .shape [1 :]
71
+
60
72
exporter = initialize_forecast_exporter_netcdf (
61
73
outpath ,
62
74
outfnprefix ,
@@ -65,9 +77,25 @@ def test_io_export_netcdf_one_member_one_time_step():
65
77
n_timesteps ,
66
78
shape ,
67
79
metadata ,
68
- n_ens_members = 1 ,
80
+ n_ens_members = n_ens_members ,
81
+ incremental = incremental ,
69
82
)
70
- export_forecast_dataset (precip [np .newaxis , :], exporter )
83
+
84
+ if n_ens_members > 1 :
85
+ precip = np .repeat (precip [np .newaxis , :, :, :], n_ens_members , axis = 0 )
86
+
87
+ if incremental == None :
88
+ export_forecast_dataset (precip , exporter )
89
+ if incremental == "timestep" :
90
+ for t in range (n_timesteps ):
91
+ if n_ens_members > 1 :
92
+ export_forecast_dataset (precip [:, t , :, :], exporter )
93
+ else :
94
+ export_forecast_dataset (precip [t , :, :], exporter )
95
+ if incremental == "member" :
96
+ for ens_mem in range (n_ens_members ):
97
+ export_forecast_dataset (precip [ens_mem , :, :, :], exporter )
98
+
71
99
close_forecast_files (exporter )
72
100
73
101
# assert if netcdf file was saved and file size is not zero
@@ -78,11 +106,11 @@ def test_io_export_netcdf_one_member_one_time_step():
78
106
79
107
precip_new , _ = import_netcdf_pysteps (output_file_path )
80
108
81
- assert_array_almost_equal (precip , precip_new .data )
109
+ assert_array_almost_equal (precip . squeeze () , precip_new .data )
82
110
assert precip_new .dtype == "single"
83
111
84
112
precip_new , _ = import_netcdf_pysteps (output_file_path , dtype = "double" )
85
- assert_array_almost_equal (precip , precip_new .data )
113
+ assert_array_almost_equal (precip . squeeze () , precip_new .data )
86
114
assert precip_new .dtype == "double"
87
115
88
116
precip_new , _ = import_netcdf_pysteps (output_file_path , fillna = - 1000 )
0 commit comments