PyMPDATA_examples.Shipway_and_Hill_2012.mpdata

 1import numpy as np
 2
 3from PyMPDATA import ScalarField, Solver, Stepper, VectorField
 4from PyMPDATA.boundary_conditions import Constant, Extrapolated
 5from PyMPDATA.impl.enumerations import INNER, OUTER
 6
 7from .arakawa_c import arakawa_c
 8
 9
10class MPDATA:
11    # pylint: disable=too-few-public-methods
12    def __init__(
13        self, nz, dt, qv_of_zZ_at_t0, g_factor_of_zZ, nr, options, activation_bc
14    ):
15        self.t = 0
16        self.dt = dt
17        self.fields = ("qv", "ql")
18
19        self.options = options
20
21        self._solvers = {}
22        for k in self.fields:
23            grid = (nz, nr) if nr > 1 and k == "ql" else (nz,)
24
25            bcs_extrapol = tuple(
26                Extrapolated(dim=d)
27                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
28            )
29            bcs_zero = tuple(
30                Extrapolated(dim=d)
31                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
32            )
33
34            stepper = Stepper(
35                options=self.options, n_dims=len(grid), non_unit_g_factor=True
36            )
37
38            data = g_factor_of_zZ(arakawa_c.z_scalar_coord(grid))
39            if nr > 1 and k == "ql":
40                data = np.repeat(data.reshape(-1, 1), nr, axis=1).squeeze()
41            g_factor = ScalarField(
42                data=data, halo=self.options.n_halo, boundary_conditions=bcs_extrapol
43            )
44
45            if nr == 1 or k == "qv":
46                data = (np.zeros(nz + 1),)
47            else:
48                data = (np.zeros((nz + 1, nr)), np.zeros((nz, nr + 1)))
49            advector = VectorField(
50                data=data, halo=self.options.n_halo, boundary_conditions=bcs_zero
51            )
52            if k == "qv":
53                data = qv_of_zZ_at_t0(arakawa_c.z_scalar_coord(grid))
54                bcs = (Constant(value=data[0]),)
55            else:
56                data = np.zeros(grid)
57                if nr == 1:
58                    bcs = (Constant(value=0),)
59                else:
60                    bcs = (Constant(value=0), activation_bc)
61            advectee = ScalarField(
62                data=data, halo=self.options.n_halo, boundary_conditions=bcs
63            )
64            self._solvers[k] = Solver(
65                stepper=stepper, advectee=advectee, advector=advector, g_factor=g_factor
66            )
67
68    def __getitem__(self, k):
69        return self._solvers[k]
class MPDATA:
11class MPDATA:
12    # pylint: disable=too-few-public-methods
13    def __init__(
14        self, nz, dt, qv_of_zZ_at_t0, g_factor_of_zZ, nr, options, activation_bc
15    ):
16        self.t = 0
17        self.dt = dt
18        self.fields = ("qv", "ql")
19
20        self.options = options
21
22        self._solvers = {}
23        for k in self.fields:
24            grid = (nz, nr) if nr > 1 and k == "ql" else (nz,)
25
26            bcs_extrapol = tuple(
27                Extrapolated(dim=d)
28                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
29            )
30            bcs_zero = tuple(
31                Extrapolated(dim=d)
32                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
33            )
34
35            stepper = Stepper(
36                options=self.options, n_dims=len(grid), non_unit_g_factor=True
37            )
38
39            data = g_factor_of_zZ(arakawa_c.z_scalar_coord(grid))
40            if nr > 1 and k == "ql":
41                data = np.repeat(data.reshape(-1, 1), nr, axis=1).squeeze()
42            g_factor = ScalarField(
43                data=data, halo=self.options.n_halo, boundary_conditions=bcs_extrapol
44            )
45
46            if nr == 1 or k == "qv":
47                data = (np.zeros(nz + 1),)
48            else:
49                data = (np.zeros((nz + 1, nr)), np.zeros((nz, nr + 1)))
50            advector = VectorField(
51                data=data, halo=self.options.n_halo, boundary_conditions=bcs_zero
52            )
53            if k == "qv":
54                data = qv_of_zZ_at_t0(arakawa_c.z_scalar_coord(grid))
55                bcs = (Constant(value=data[0]),)
56            else:
57                data = np.zeros(grid)
58                if nr == 1:
59                    bcs = (Constant(value=0),)
60                else:
61                    bcs = (Constant(value=0), activation_bc)
62            advectee = ScalarField(
63                data=data, halo=self.options.n_halo, boundary_conditions=bcs
64            )
65            self._solvers[k] = Solver(
66                stepper=stepper, advectee=advectee, advector=advector, g_factor=g_factor
67            )
68
69    def __getitem__(self, k):
70        return self._solvers[k]
MPDATA(nz, dt, qv_of_zZ_at_t0, g_factor_of_zZ, nr, options, activation_bc)
13    def __init__(
14        self, nz, dt, qv_of_zZ_at_t0, g_factor_of_zZ, nr, options, activation_bc
15    ):
16        self.t = 0
17        self.dt = dt
18        self.fields = ("qv", "ql")
19
20        self.options = options
21
22        self._solvers = {}
23        for k in self.fields:
24            grid = (nz, nr) if nr > 1 and k == "ql" else (nz,)
25
26            bcs_extrapol = tuple(
27                Extrapolated(dim=d)
28                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
29            )
30            bcs_zero = tuple(
31                Extrapolated(dim=d)
32                for d in ((OUTER, INNER) if k == "ql" and nr > 1 else (INNER,))
33            )
34
35            stepper = Stepper(
36                options=self.options, n_dims=len(grid), non_unit_g_factor=True
37            )
38
39            data = g_factor_of_zZ(arakawa_c.z_scalar_coord(grid))
40            if nr > 1 and k == "ql":
41                data = np.repeat(data.reshape(-1, 1), nr, axis=1).squeeze()
42            g_factor = ScalarField(
43                data=data, halo=self.options.n_halo, boundary_conditions=bcs_extrapol
44            )
45
46            if nr == 1 or k == "qv":
47                data = (np.zeros(nz + 1),)
48            else:
49                data = (np.zeros((nz + 1, nr)), np.zeros((nz, nr + 1)))
50            advector = VectorField(
51                data=data, halo=self.options.n_halo, boundary_conditions=bcs_zero
52            )
53            if k == "qv":
54                data = qv_of_zZ_at_t0(arakawa_c.z_scalar_coord(grid))
55                bcs = (Constant(value=data[0]),)
56            else:
57                data = np.zeros(grid)
58                if nr == 1:
59                    bcs = (Constant(value=0),)
60                else:
61                    bcs = (Constant(value=0), activation_bc)
62            advectee = ScalarField(
63                data=data, halo=self.options.n_halo, boundary_conditions=bcs
64            )
65            self._solvers[k] = Solver(
66                stepper=stepper, advectee=advectee, advector=advector, g_factor=g_factor
67            )
t
dt
fields
options