PyMPDATA_examples.Jarecka_et_al_2015.simulation

 1import numpy as np
 2from PyMPDATA_examples.Jarecka_et_al_2015 import formulae
 3
 4from PyMPDATA import ScalarField, Solver, Stepper, VectorField
 5from PyMPDATA.boundary_conditions import Constant
 6
 7
 8class Simulation:
 9    def __init__(self, settings):
10        self.settings = settings
11        s = settings
12
13        halo = settings.options.n_halo
14        grid = (s.nx, s.ny)
15        bcs = [Constant(value=0)] * len(grid)
16
17        self.advector = VectorField(
18            (np.zeros((s.nx + 1, s.ny)), np.zeros((s.nx, s.ny + 1))), halo, bcs
19        )
20
21        xi, yi = np.indices(grid, dtype=float)
22        xi -= (s.nx - 1) / 2
23        yi -= (s.ny - 1) / 2
24        x = xi * s.dx
25        y = yi * s.dy
26        h0 = formulae.amplitude(x, y, s.lx0, s.ly0)
27
28        advectees = {
29            "h": ScalarField(h0, halo, bcs),
30            "uh": ScalarField(np.zeros(grid), halo, bcs),
31            "vh": ScalarField(np.zeros(grid), halo, bcs),
32        }
33
34        stepper = Stepper(options=s.options, grid=grid)
35        self.solvers = {
36            k: Solver(stepper, v, self.advector) for k, v in advectees.items()
37        }
38
39    @staticmethod
40    def interpolate(psi, axis):
41        idx = (
42            (slice(None, -1), slice(None, None)),
43            (slice(None, None), slice(None, -1)),
44        )
45        return np.diff(psi, axis=axis) / 2 + psi[idx[axis]]
46
47    def run(self):
48        s = self.settings
49        grid_step = (s.dx, s.dy)
50        idx = ((slice(1, -1), slice(None, None)), (slice(None, None), slice(1, -1)))
51        output = []
52        for it in range(s.nt + 1):
53            if it != 0:
54                h = self.solvers["h"].advectee.get()
55                for xy, k in enumerate(("uh", "vh")):
56                    mask = h > s.eps
57                    vel = np.where(mask, np.nan, 0)
58                    np.divide(self.solvers[k].advectee.get(), h, where=mask, out=vel)
59                    self.advector.get_component(xy)[idx[xy]] = (
60                        self.interpolate(vel, axis=xy) * s.dt / grid_step[xy]
61                    )
62                self.solvers["h"].advance(1)
63                assert h.ctypes.data == self.solvers["h"].advectee.get().ctypes.data
64                for xy, k in enumerate(("uh", "vh")):
65                    psi = self.solvers[k].advectee.get()
66                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
67                    self.solvers[k].advance(1)
68                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
69            if it % s.outfreq == 0:
70                output.append(
71                    {
72                        k: self.solvers[k].advectee.get().copy()
73                        for k in self.solvers.keys()
74                    }
75                )
76        return output
class Simulation:
 9class Simulation:
10    def __init__(self, settings):
11        self.settings = settings
12        s = settings
13
14        halo = settings.options.n_halo
15        grid = (s.nx, s.ny)
16        bcs = [Constant(value=0)] * len(grid)
17
18        self.advector = VectorField(
19            (np.zeros((s.nx + 1, s.ny)), np.zeros((s.nx, s.ny + 1))), halo, bcs
20        )
21
22        xi, yi = np.indices(grid, dtype=float)
23        xi -= (s.nx - 1) / 2
24        yi -= (s.ny - 1) / 2
25        x = xi * s.dx
26        y = yi * s.dy
27        h0 = formulae.amplitude(x, y, s.lx0, s.ly0)
28
29        advectees = {
30            "h": ScalarField(h0, halo, bcs),
31            "uh": ScalarField(np.zeros(grid), halo, bcs),
32            "vh": ScalarField(np.zeros(grid), halo, bcs),
33        }
34
35        stepper = Stepper(options=s.options, grid=grid)
36        self.solvers = {
37            k: Solver(stepper, v, self.advector) for k, v in advectees.items()
38        }
39
40    @staticmethod
41    def interpolate(psi, axis):
42        idx = (
43            (slice(None, -1), slice(None, None)),
44            (slice(None, None), slice(None, -1)),
45        )
46        return np.diff(psi, axis=axis) / 2 + psi[idx[axis]]
47
48    def run(self):
49        s = self.settings
50        grid_step = (s.dx, s.dy)
51        idx = ((slice(1, -1), slice(None, None)), (slice(None, None), slice(1, -1)))
52        output = []
53        for it in range(s.nt + 1):
54            if it != 0:
55                h = self.solvers["h"].advectee.get()
56                for xy, k in enumerate(("uh", "vh")):
57                    mask = h > s.eps
58                    vel = np.where(mask, np.nan, 0)
59                    np.divide(self.solvers[k].advectee.get(), h, where=mask, out=vel)
60                    self.advector.get_component(xy)[idx[xy]] = (
61                        self.interpolate(vel, axis=xy) * s.dt / grid_step[xy]
62                    )
63                self.solvers["h"].advance(1)
64                assert h.ctypes.data == self.solvers["h"].advectee.get().ctypes.data
65                for xy, k in enumerate(("uh", "vh")):
66                    psi = self.solvers[k].advectee.get()
67                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
68                    self.solvers[k].advance(1)
69                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
70            if it % s.outfreq == 0:
71                output.append(
72                    {
73                        k: self.solvers[k].advectee.get().copy()
74                        for k in self.solvers.keys()
75                    }
76                )
77        return output
Simulation(settings)
10    def __init__(self, settings):
11        self.settings = settings
12        s = settings
13
14        halo = settings.options.n_halo
15        grid = (s.nx, s.ny)
16        bcs = [Constant(value=0)] * len(grid)
17
18        self.advector = VectorField(
19            (np.zeros((s.nx + 1, s.ny)), np.zeros((s.nx, s.ny + 1))), halo, bcs
20        )
21
22        xi, yi = np.indices(grid, dtype=float)
23        xi -= (s.nx - 1) / 2
24        yi -= (s.ny - 1) / 2
25        x = xi * s.dx
26        y = yi * s.dy
27        h0 = formulae.amplitude(x, y, s.lx0, s.ly0)
28
29        advectees = {
30            "h": ScalarField(h0, halo, bcs),
31            "uh": ScalarField(np.zeros(grid), halo, bcs),
32            "vh": ScalarField(np.zeros(grid), halo, bcs),
33        }
34
35        stepper = Stepper(options=s.options, grid=grid)
36        self.solvers = {
37            k: Solver(stepper, v, self.advector) for k, v in advectees.items()
38        }
settings
advector
solvers
@staticmethod
def interpolate(psi, axis):
40    @staticmethod
41    def interpolate(psi, axis):
42        idx = (
43            (slice(None, -1), slice(None, None)),
44            (slice(None, None), slice(None, -1)),
45        )
46        return np.diff(psi, axis=axis) / 2 + psi[idx[axis]]
def run(self):
48    def run(self):
49        s = self.settings
50        grid_step = (s.dx, s.dy)
51        idx = ((slice(1, -1), slice(None, None)), (slice(None, None), slice(1, -1)))
52        output = []
53        for it in range(s.nt + 1):
54            if it != 0:
55                h = self.solvers["h"].advectee.get()
56                for xy, k in enumerate(("uh", "vh")):
57                    mask = h > s.eps
58                    vel = np.where(mask, np.nan, 0)
59                    np.divide(self.solvers[k].advectee.get(), h, where=mask, out=vel)
60                    self.advector.get_component(xy)[idx[xy]] = (
61                        self.interpolate(vel, axis=xy) * s.dt / grid_step[xy]
62                    )
63                self.solvers["h"].advance(1)
64                assert h.ctypes.data == self.solvers["h"].advectee.get().ctypes.data
65                for xy, k in enumerate(("uh", "vh")):
66                    psi = self.solvers[k].advectee.get()
67                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
68                    self.solvers[k].advance(1)
69                    psi[:] -= s.dt / 2 * h * np.gradient(h, grid_step[xy], axis=xy)
70            if it % s.outfreq == 0:
71                output.append(
72                    {
73                        k: self.solvers[k].advectee.get().copy()
74                        for k in self.solvers.keys()
75                    }
76                )
77        return output