Read about our upcoming Code of Conduct on this issue

test_solver.py 11.9 KB
Newer Older
1
2
import unittest

calpe's avatar
calpe committed
3
4
import numpy as np
import matplotlib.pyplot as plt
5

calpe's avatar
calpe committed
6
import fluidsim as fls
7

8
9
10
import fluiddyn.util.mpi as mpi

from fluidsim.solvers.ns2d.strat.solver import Simul
11
from fluidsim.solvers.ns2d.test_solver import TestSimulBase as Base
12
13


Pierre Augier's avatar
Pierre Augier committed
14
class TestSimulBase(Base):
15
16
    Simul = Simul

calpe's avatar
calpe committed
17
18
19
20
21

class TestSolverNS2DTendency(TestSimulBase):
    @classmethod
    def init_params(self):
        params = super().init_params()
22
23
        params.time_stepping.USE_CFL = False
        params.time_stepping.USE_T_END = False
24
25
        # todo: understand why it fails for larger it_end and deltat0. I (pa)
        # guess there could be a bug hidden.
26
        params.time_stepping.it_end = 2
27
        params.time_stepping.deltat0 = 0.02
28
29
        params.output.HAS_TO_SAVE = False

30
    def test_tendency(self):
calpe's avatar
calpe committed
31
        sim = self.sim
32

33
34
        rot_fft = sim.state.get_var("rot_fft")
        b_fft = sim.state.get_var("b_fft")
35

36
37
        assert not np.any(b_fft)

38
        tend = sim.tendencies_nonlin(state_spect=sim.state.state_spect)
calpe's avatar
calpe committed
39
40
        Frot_fft = tend.get_var("rot_fft")
        Fb_fft = tend.get_var("b_fft")
41

calpe's avatar
calpe committed
42
        T_rot = np.real(Frot_fft.conj() * rot_fft)
43

44
        ratio = sim.oper.sum_wavenumbers(T_rot) / sim.oper.sum_wavenumbers(
calpe's avatar
calpe committed
45
46
            abs(T_rot)
        )
47

48
49
50
        self.assertGreater(1e-15, ratio)

        # init_fields = noise gives energy only to rot
51
        sim.time_stepping.start()
52
53
54

        rot_fft = sim.state.get_var("rot_fft")
        b_fft = sim.state.get_var("b_fft")
55

calpe's avatar
calpe committed
56
        # init_fields = noise gives energy only to rot
57
58
59
60
61
62
63
64
65
66
67
68
69
70
        assert np.any(b_fft)

        tend = sim.tendencies_nonlin(state_spect=sim.state.state_spect)
        Frot_fft = tend.get_var("rot_fft")
        Fb_fft = tend.get_var("b_fft")

        assert sim.check_energy_conservation(rot_fft, b_fft, Frot_fft, Fb_fft)


class TestForcingLinearMode(TestSimulBase):
    @classmethod
    def init_params(self):
        params = super().init_params()

Pierre Augier's avatar
Pierre Augier committed
71
        params.forcing.enable = True
72
73
74
75
76
        params.forcing.type = "tcrandom_anisotropic"
        params.forcing.nkmin_forcing = 4
        params.forcing.nkmax_forcing = 6
        params.forcing.key_forced = "ap_fft"

77
    def test_forcing_linear_mode(self):
78
        sim = self.sim
79
        sim.time_stepping.start()
Pierre Augier's avatar
Pierre Augier committed
80
81
        if mpi.nb_proc == 1:
            sim.forcing.forcing_maker.plot_forcing_region()
82
83


84
85
86
87
88
89
90
91
class TestForcingConstantRateEnergy(TestSimulBase):
    @classmethod
    def init_params(self):
        params = super().init_params()
        params.forcing.enable = True
        params.forcing.type = "tcrandom"
        params.forcing.normalized.constant_rate_of = "energy"
        params.forcing.key_forced = "rot_fft"
92
93
94
95
96
97
        params.forcing.tcrandom_anisotropic.kz_negative_enable = True
        params.forcing.forcing_rate = 3.333

        params.output.periods_save.spatial_means = 1e-6

        return params
98
99
100
101

    def test_(self):
        self.sim.time_stepping.start()

102
103
104
105
106
107
        if mpi.rank == 0:
            # Does the energy injection rate have the correct value at all times ?
            means = self.sim.output.spatial_means.load()
            P_tot = means["PK_tot"] + means["PA_tot"]
            assert np.allclose(P_tot, self.sim.params.forcing.forcing_rate)

108

109
class TestForcingConstantRateEnergyAP(TestForcingConstantRateEnergy):
110
111
112
113
    @classmethod
    def init_params(self):
        params = super().init_params()
        params.forcing.key_forced = "ap_fft"
114
        params.forcing.tcrandom_anisotropic.kz_negative_enable = False
115
116


calpe's avatar
calpe committed
117
118
119
class TestForcingOutput(TestSimulBase):
    @classmethod
    def init_params(self):
120

calpe's avatar
calpe committed
121
        params = super().init_params()
122
123
124
125
126
127
128

        # Time stepping parameters
        params.time_stepping.USE_CFL = False
        params.time_stepping.USE_T_END = False
        params.time_stepping.it_end = 2
        params.time_stepping.deltat0 = 0.1

129
        params.forcing.enable = True
calpe's avatar
calpe committed
130
131
        # params.forcing.type = "tcrandom"
        # Forcing also linear mode!!!
132
        params.forcing.type = "tcrandom_anisotropic"
calpe's avatar
calpe committed
133
134
135
        params.forcing.nkmin_forcing = 4
        params.forcing.nkmax_forcing = 6

136
137
138
        # save all outputs!
        periods = params.output.periods_save
        for key in periods._key_attribs:
139
            periods[key] = 0.1
140

calpe's avatar
calpe committed
141
142
143
144
        params.output.ONLINE_PLOT_OK = True
        params.output.periods_print.print_stdout = 0.2
        params.output.periods_plot.phys_fields = 0.2

145
146
147
148
149
150
        Lx, Ly = params.oper.Lx, params.oper.Ly
        nx, ny = params.oper.nx, params.oper.ny
        params.output.temporal_spectra.probes_region = (0.0, Lx, 0.0, Ly)
        params.output.temporal_spectra.probes_deltax = Lx / nx
        params.output.temporal_spectra.probes_deltay = Ly / ny
        params.output.temporal_spectra.SAVE_AS_FLOAT32 = True
151

Pierre Augier's avatar
Pierre Augier committed
152
        params.output.spatiotemporal_spectra.probes_region = (
Pierre Augier's avatar
Pierre Augier committed
153
154
            nx // 2 + 1,
            ny // 2 + 1,
Pierre Augier's avatar
Pierre Augier committed
155
156
        )
        params.output.spatiotemporal_spectra.SAVE_AS_COMPLEX64 = False
calpe's avatar
calpe committed
157
158
159
160
161
162
163
164
165
166
167

        for tag in params.output._tag_children:
            if tag.startswith("periods"):
                continue
            child = params.output[tag]
            if hasattr(child, "HAS_TO_PLOT_SAVED"):
                child["HAS_TO_PLOT_SAVED"] = True

    def test_forcing_output(self):

        sim = self.sim
168

169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
        sim.time_stepping.start()

        sim.state.compute("rot_fft")
        sim.state.compute("rot_fft")

        with self.assertRaises(ValueError):
            sim.state.compute("abc")
        sim.state.compute("abc", RAISE_ERROR=False)

        ux_fft = sim.state.compute("ux_fft")
        uy_fft = sim.state.compute("uy_fft")

        sim.state.init_statespect_from(ux_fft=ux_fft)
        sim.state.init_statespect_from(uy_fft=uy_fft)

        sim.output.compute_enstrophy()

        if mpi.nb_proc == 1:
187

188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
            plt.close("all")

            sim.output.plot_summary()

            sim.output.spectra.plot2d()

            sim.output.spatial_means.plot_energy()
            sim.output.spatial_means.plot_dt_energy()
            sim.output.spatial_means.plot_energy_shear_modes()

            plt.close("all")
            sim.output.spatial_means.compute_time_means()
            sim.output.spatial_means.load_dataset()
            sim.output.spatial_means.time_first_saved()
            sim.output.spatial_means.time_last_saved()

            sim.output.spectra_multidim.plot()
calpe's avatar
calpe committed
205
206

            with self.assertRaises(ValueError):
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
                sim.state.get_var("test")

        # `compute('q')` two times for better coverage...
        sim.state.get_var("q")
        sim.state.get_var("q")
        sim.state.get_var("div")

        path_run = sim.output.path_run
        if mpi.nb_proc > 1:
            path_run = mpi.comm.bcast(path_run)

        sim3 = fls.load_state_phys_file(path_run, modif_save_params=False)
        sim3.params.time_stepping.t_end += 0.2
        sim3.time_stepping.start()

Pierre Augier's avatar
Pierre Augier committed
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        if mpi.nb_proc > 1:
            plt.close("all")
            return

        sim3.output.phys_fields.animate(
            "ux",
            dt_frame_in_sec=1e-6,
            dt_equations=0.3,
            repeat=False,
            clim=(-1, 1),
            save_file=False,
            numfig=1,
        )

        sim2 = fls.load_sim_for_plot(sim.output.path_run)
        sim2.output

        # TODO: add spatiotemporal_spectra calls

        means = sim2.output.spatial_means.load()
Pierre Augier's avatar
Pierre Augier committed
242
        energy_K_mean = means["EK"][:-1].mean()
Pierre Augier's avatar
Pierre Augier committed
243
244

        sim2.output.temporal_spectra.load_time_series()
245
246
        spectra_omega = sim2.output.temporal_spectra.compute_spectra()
        omegas = spectra_omega["omegas"]
Pierre Augier's avatar
Pierre Augier committed
247
        delta_omega = omegas[1]
248

Pierre Augier's avatar
Pierre Augier committed
249
250
        sim2.output.temporal_spectra.plot_spectra()

Pierre Augier's avatar
Pierre Augier committed
251
252
253
        spatiotemporal_spectra = sim2.output.spatiotemporal_spectra
        series_kxky = spatiotemporal_spectra.load_time_series()

254
        spectra_kxkyomega = spatiotemporal_spectra.compute_spectra()
Pierre Augier's avatar
Pierre Augier committed
255
256
257
258
        spectra_omega_from_spatiotemp = (
            spatiotemporal_spectra.compute_temporal_spectra()
        )

259
260
261
262
263
264
265
266
267
268
269
270
271
272
        deltakx = 2 * np.pi / self.params.oper.Lx
        order = spectra_kxkyomega["dims_order"]
        KX = deltakx * spectra_kxkyomega[f"K{order[1]}_adim"]
        kx_max = self.params.oper.nx // 2 * deltakx

        assert kx_max == KX.max()

        from fluidsim.solvers.ns2d.output.spatiotemporal_spectra import (
            _sum_wavenumber2D,
        )

        def sum_wavenumber(field):
            return _sum_wavenumber2D(field, KX, kx_max)

Pierre Augier's avatar
Pierre Augier committed
273
        tspectrum_mean = (
274
            spectra_omega["spectrum_ux"] + spectra_omega["spectrum_uy"]
Pierre Augier's avatar
Pierre Augier committed
275
276
277
278
279
280
281
282
        )

        energy_tspect_mean = 0.5 * delta_omega * tspectrum_mean.sum()

        assert np.allclose(energy_K_mean, energy_tspect_mean), (
            energy_K_mean / energy_tspect_mean
        )

Jason Reneuve's avatar
Jason Reneuve committed
283
284
        _ = spatiotemporal_spectra.save_spectra_kzkhomega()
        spectra_kzkhomega = spatiotemporal_spectra.load_spectra_kzkhomega()
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337

        delta_kz = spectra_kzkhomega["kz_spectra"][1]
        delta_kh = spectra_kzkhomega["kh_spectra"][1]
        delta_omega = spectra_kzkhomega["omegas"][1]
        coef = delta_kz * delta_kh * delta_omega

        for letter in "xy":
            vi_fft = series_kxky[f"u{letter}_Fourier"]
            spectrum_kxkykzomega = spectra_kxkyomega["spectrum_u" + letter]
            spectrum_omega = spectra_omega["spectrum_u" + letter]
            spectrum_omega_from_spatiotemp = spectra_omega_from_spatiotemp[
                "spectrum_u" + letter
            ]
            spectrum_kzkhomega = spectra_kzkhomega["spectrum_u" + letter]

            E_series_kxkykz = 0.5 * sum_wavenumber(
                (abs(vi_fft) ** 2).mean(axis=-1)
            )
            assert E_series_kxkykz > 0, (letter, vi_fft)

            E_kxkykzomega = (
                0.5
                * delta_omega
                * sum_wavenumber(spectrum_kxkykzomega.sum(axis=-1))
            )
            E_omega = 0.5 * delta_omega * spectrum_omega.sum()
            E_omega_from_spatiotemp = (
                0.5 * delta_omega * spectrum_omega_from_spatiotemp.sum()
            )

            E_kzkhomega = 0.5 * coef * spectrum_kzkhomega.sum()

            assert np.allclose(E_omega, E_kxkykzomega), (
                letter,
                E_kxkykzomega / E_omega,
            )

            assert np.allclose(E_series_kxkykz, E_kxkykzomega), (
                letter,
                E_kxkykzomega / E_series_kxkykz,
            )

            assert np.allclose(E_series_kxkykz, E_kzkhomega), (
                letter,
                E_kzkhomega / E_series_kxkykz,
            )

            assert np.allclose(E_omega, E_omega_from_spatiotemp), (
                letter,
                E_omega,
                E_omega_from_spatiotemp,
            )

338
339
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(key_field="ux")
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(
340
            key_field="ux", equation="omega=N/2"
341
        )
342
343
344
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(
            key_field="ux", equation="kh=1"
        )
345
346
347
348
349
350
351
352
353
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(
            key_field="ux", equation="kz=1"
        )
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(
            key_field="ux", equation="ikh=1"
        )
        sim2.output.spatiotemporal_spectra.plot_kzkhomega(
            key_field="ux", equation="ikz=1"
        )
354
        sim2.output.spatiotemporal_spectra.plot_temporal_spectra()
355

Pierre Augier's avatar
Pierre Augier committed
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
        sim2.output.increments.load()
        sim2.output.increments.plot()
        sim2.output.increments.load_pdf_from_file()

        sim2.output.phys_fields.animate(
            "ux",
            dt_frame_in_sec=1e-6,
            dt_equations=0.3,
            repeat=False,
            clim=(-1, 1),
            save_file=False,
            numfig=1,
        )
        sim2.output.phys_fields.plot()

calpe's avatar
calpe committed
371
372
373
374
375
376
377
378
379
        plt.close("all")


class TestSolverNS2DInitLinearMode(TestSimulBase):
    @classmethod
    def init_params(self):
        params = super().init_params()
        params.init_fields.type = "linear_mode"
        params.output.HAS_TO_SAVE = False
380

381
    def test_init_linear_mode(self):
calpe's avatar
calpe committed
382
        pass
383

384

385
if __name__ == "__main__":
386
    unittest.main()