advection_test.py 2.98 KB
Newer Older
mashun1's avatar
veros  
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import numpy as np

from veros.core import advection
from veros.pyom_compat import get_random_state

from test_base import compare_state


TEST_SETTINGS = dict(
    nx=70,
    ny=60,
    nz=50,
    dt_tracer=3600,
    dt_mom=3600,
)


def test_calculate_velocity_on_wgrid(pyom2_lib):
    vs_state, pyom_obj = get_random_state(pyom2_lib, extra_settings=TEST_SETTINGS)
    advection.calculate_velocity_on_wgrid(vs_state)
    pyom_obj.calculate_velocity_on_wgrid()

    compare_state(vs_state, pyom_obj)


def test_adv_flux_2nd(pyom2_lib):
    vs_state, pyom_obj = get_random_state(pyom2_lib, extra_settings=TEST_SETTINGS)

    res = advection.adv_flux_2nd(vs_state, vs_state.variables.Hd[..., 1])

    m = pyom_obj.main_module
    pyom_obj.adv_flux_2nd(
        is_=-1,
        ie_=m.nx + 2,
        js_=-1,
        je_=m.ny + 2,
        nz_=m.nz,
        adv_fe=m.flux_east,
        adv_fn=m.flux_north,
        adv_ft=m.flux_top,
        var=m.hd[..., 1],
    )

    np.testing.assert_allclose(res[0], m.flux_east)
    np.testing.assert_allclose(res[1], m.flux_north)
    np.testing.assert_allclose(res[2], m.flux_top)


def test_adv_flux_superbee(pyom2_lib):
    vs_state, pyom_obj = get_random_state(pyom2_lib, extra_settings=TEST_SETTINGS)

    res = advection.adv_flux_superbee(vs_state, vs_state.variables.Hd[..., 1])

    m = pyom_obj.main_module
    pyom_obj.adv_flux_superbee(
        is_=-1,
        ie_=m.nx + 2,
        js_=-1,
        je_=m.ny + 2,
        nz_=m.nz,
        adv_fe=m.flux_east,
        adv_fn=m.flux_north,
        adv_ft=m.flux_top,
        var=m.hd[..., 1],
    )

    np.testing.assert_allclose(res[0], m.flux_east)
    np.testing.assert_allclose(res[1], m.flux_north)
    np.testing.assert_allclose(res[2], m.flux_top)


def test_adv_flux_upwind_wgrid(pyom2_lib):
    vs_state, pyom_obj = get_random_state(pyom2_lib, extra_settings=TEST_SETTINGS)

    res = advection.adv_flux_upwind_wgrid(vs_state, vs_state.variables.Hd[..., 1])

    m = pyom_obj.main_module
    pyom_obj.adv_flux_upwind_wgrid(
        is_=-1,
        ie_=m.nx + 2,
        js_=-1,
        je_=m.ny + 2,
        nz_=m.nz,
        adv_fe=m.flux_east,
        adv_fn=m.flux_north,
        adv_ft=m.flux_top,
        var=m.hd[..., 1],
    )

    np.testing.assert_allclose(res[0], m.flux_east)
    np.testing.assert_allclose(res[1], m.flux_north)
    np.testing.assert_allclose(res[2], m.flux_top)


def test_adv_flux_superbee_wgrid(pyom2_lib):
    vs_state, pyom_obj = get_random_state(pyom2_lib, extra_settings=TEST_SETTINGS)

    res = advection.adv_flux_superbee_wgrid(vs_state, vs_state.variables.Hd[..., 1])

    m = pyom_obj.main_module
    pyom_obj.adv_flux_superbee_wgrid(
        is_=-1,
        ie_=m.nx + 2,
        js_=-1,
        je_=m.ny + 2,
        nz_=m.nz,
        adv_fe=m.flux_east,
        adv_fn=m.flux_north,
        adv_ft=m.flux_top,
        var=m.hd[..., 1],
    )

    np.testing.assert_allclose(res[0], m.flux_east)
    np.testing.assert_allclose(res[1], m.flux_north)
    np.testing.assert_allclose(res[2], m.flux_top)