@@ -25,23 +25,36 @@ Set up a full test problem
25
25
Increasing `more_compute_diagnostics` adds more copies of a compute diagnostic with no output.
26
26
Useful to stress allocations.
27
27
"""
28
- function setup_integrator (output_dir; context, more_compute_diagnostics = 0 )
28
+ function setup_integrator (
29
+ output_dir;
30
+ space = SphericalShellSpace (; context),
31
+ context,
32
+ more_compute_diagnostics = 0 ,
33
+ dict_writer = nothing ,
34
+ )
29
35
t0 = 0.0
30
36
tf = 10.0
31
37
dt = 1.0
32
- space = SphericalShellSpace (; context)
33
38
args, kwargs = create_problem (space; t0, tf, dt)
34
39
35
40
@info " Writing output to $output_dir "
36
41
37
42
dummy_writer = ClimaDiagnostics. Writers. DummyWriter ()
38
43
h5_writer = ClimaDiagnostics. Writers. HDF5Writer (output_dir)
39
- nc_writer = ClimaDiagnostics. Writers. NetCDFWriter (
40
- space,
41
- output_dir;
42
- num_points = (10 , 5 , 3 ),
43
- start_date = Dates. DateTime (2015 , 2 , 2 ),
44
- )
44
+ if space isa ClimaCore. Spaces. PointSpace
45
+ nc_writer = ClimaDiagnostics. Writers. NetCDFWriter (
46
+ space,
47
+ output_dir;
48
+ start_date = Dates. DateTime (2015 , 2 , 2 ),
49
+ )
50
+ else
51
+ nc_writer = ClimaDiagnostics. Writers. NetCDFWriter (
52
+ space,
53
+ output_dir;
54
+ num_points = (10 , 5 , 3 ),
55
+ start_date = Dates. DateTime (2015 , 2 , 2 ),
56
+ )
57
+ end
45
58
46
59
function compute_my_var! (out, u, p, t)
47
60
if isnothing (out)
@@ -81,6 +94,7 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
81
94
variable = simple_var,
82
95
output_writer = h5_writer,
83
96
)
97
+
84
98
inst_every3s_diagnostic_another = ClimaDiagnostics. ScheduledDiagnostic (
85
99
variable = simple_var,
86
100
output_writer = nc_writer,
@@ -96,6 +110,13 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
96
110
inst_every3s_diagnostic,
97
111
inst_every3s_diagnostic_another,
98
112
]
113
+ if ! isnothing (dict_writer)
114
+ inst_diagnostic_dict = ClimaDiagnostics. ScheduledDiagnostic (
115
+ variable = simple_var,
116
+ output_writer = dict_writer,
117
+ )
118
+ scheduled_diagnostics = [scheduled_diagnostics... , inst_diagnostic_dict]
119
+ end
99
120
100
121
@test inst_every3s_diagnostic_another == inst_every3s_diagnostic
101
122
@test ! (inst_every3s_diagnostic_another === inst_every3s_diagnostic)
@@ -116,36 +137,94 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
116
137
)
117
138
end
118
139
119
- @testset " A full problem" begin
120
- mktempdir () do output_dir
121
- output_dir = ClimaComms. bcast (context, output_dir)
122
-
123
- integrator = setup_integrator (output_dir; context)
124
-
125
- SciMLBase. solve! (integrator)
126
-
127
- if ClimaComms. iamroot (context)
128
- NCDatasets. NCDataset (joinpath (output_dir, " YO_1it_inst.nc" )) do nc
129
- @test nc[" YO" ]. attrib[" short_name" ] == " YO"
130
- @test nc[" YO" ]. attrib[" long_name" ] == " YO YO, Instantaneous"
131
- @test size (nc[" YO" ]) == (11 , 10 , 5 , 10 )
132
- @test nc[" YO" ]. attrib[" start_date" ] ==
133
- string (Dates. DateTime (2015 , 2 , 2 ))
134
- end
135
-
136
- NCDatasets. NCDataset (
137
- joinpath (output_dir, " YO_2it_average.nc" ),
138
- ) do nc
139
- @test nc[" YO" ]. attrib[" short_name" ] == " YO"
140
- @test nc[" YO" ]. attrib[" long_name" ] ==
141
- " YO YO, average within every 2 iterations"
142
- @test size (nc[" YO" ]) == (5 , 10 , 5 , 10 )
143
- end
144
-
145
- NCDatasets. NCDataset (joinpath (output_dir, " YO_3s_inst.nc" )) do nc
146
- @test nc[" YO" ]. attrib[" short_name" ] == " YO"
147
- @test nc[" YO" ]. attrib[" long_name" ] == " YO YO, Instantaneous"
148
- @test size (nc[" YO" ]) == (4 , 10 , 5 , 10 )
140
+ sphere_space = SphericalShellSpace (; context)
141
+ purely_horizontal_space = ClimaCore. Spaces. level (sphere_space, 1 )
142
+ # list of tuples of (space, space_name, dimensions of written diagnostics without time)
143
+ spaces_test_list = [
144
+ (sphere_space, " SphericalShellSpace" , (10 , 5 , 10 )),
145
+ (purely_horizontal_space, " purely horizontal space" , (10 , 5 )),
146
+ (
147
+ ClimaCore. Spaces. PointSpace (context, ClimaCore. Geometry. ZPoint (1.0 )),
148
+ " PointSpace" ,
149
+ (),
150
+ ),
151
+ ]
152
+ if context isa ClimaComms. SingletonCommsContext
153
+ spaces_test_list = [
154
+ spaces_test_list... ,
155
+ (
156
+ ColumnCenterFiniteDifferenceSpace (10 , context),
157
+ " purely vertical space" ,
158
+ (10 ,),
159
+ ),
160
+ ]
161
+ end
162
+ for (space, space_name, written_space_dims) in spaces_test_list
163
+ @testset " A full problem using a $space_name " begin
164
+ mktempdir () do output_dir
165
+ output_dir = ClimaComms. bcast (context, output_dir)
166
+ dict_writer = ClimaDiagnostics. Writers. DictWriter ()
167
+ if (space isa ClimaCore. Spaces. PointSpace) &&
168
+ pkgversion (ClimaCore) < v " 0.14.27"
169
+ @test_throws " HDF5Writer only supports Fields with PointSpace for ClimaCore >= 0.14.27" setup_integrator (
170
+ output_dir;
171
+ context,
172
+ space,
173
+ dict_writer,
174
+ )
175
+ else
176
+ integrator =
177
+ setup_integrator (output_dir; context, space, dict_writer)
178
+
179
+ SciMLBase. solve! (integrator)
180
+
181
+ if ClimaComms. iamroot (context)
182
+ NCDatasets. NCDataset (
183
+ joinpath (output_dir, " YO_1it_inst.nc" ),
184
+ ) do nc
185
+ @test nc[" YO" ]. attrib[" short_name" ] == " YO"
186
+ @test nc[" YO" ]. attrib[" long_name" ] ==
187
+ " YO YO, Instantaneous"
188
+ @test size (nc[" YO" ]) == (11 , written_space_dims... )
189
+ @test nc[" YO" ]. attrib[" start_date" ] ==
190
+ string (Dates. DateTime (2015 , 2 , 2 ))
191
+ end
192
+
193
+ NCDatasets. NCDataset (
194
+ joinpath (output_dir, " YO_2it_average.nc" ),
195
+ ) do nc
196
+ @test nc[" YO" ]. attrib[" short_name" ] == " YO"
197
+ @test nc[" YO" ]. attrib[" long_name" ] ==
198
+ " YO YO, average within every 2 iterations"
199
+ @test size (nc[" YO" ]) == (5 , written_space_dims... )
200
+ end
201
+
202
+ NCDatasets. NCDataset (
203
+ joinpath (output_dir, " YO_3s_inst.nc" ),
204
+ ) do nc
205
+ @test nc[" YO" ]. attrib[" short_name" ] == " YO"
206
+ @test nc[" YO" ]. attrib[" long_name" ] ==
207
+ " YO YO, Instantaneous"
208
+ @test size (nc[" YO" ]) == (4 , written_space_dims... )
209
+ end
210
+ end
211
+ @test count (
212
+ occursin .(
213
+ Ref (r" YO_1it_inst_\d *\.\d\. h5" ),
214
+ readdir (output_dir),
215
+ ),
216
+ ) == 11
217
+ reader = ClimaCore. InputOutput. HDF5Reader (
218
+ joinpath (output_dir, " YO_1it_inst_10.0.h5" ),
219
+ context,
220
+ )
221
+ @test parent (
222
+ ClimaCore. InputOutput. read_field (reader, " YO_1it_inst" ),
223
+ ) == parent (integrator. u. my_var)
224
+ close (reader)
225
+ @test length (keys (dict_writer. dict[" YO_1it_inst" ])) == 11
226
+ @test dict_writer. dict[" YO_1it_inst" ][integrator. t] ==
227
+ integrator. u. my_var
149
228
end
150
229
end
151
230
end
0 commit comments