PySDM_examples.Szumowski_et_al_1998.mpdata_2d

  1import inspect
  2from functools import cached_property
  3from threading import Thread
  4
  5import numpy as np
  6from PyMPDATA import Options, ScalarField, Solver, Stepper, VectorField
  7from PyMPDATA.boundary_conditions import Periodic
  8from PySDM_examples.Szumowski_et_al_1998.fields import (
  9    nondivergent_vector_field_2d,
 10    x_vec_coord,
 11    z_vec_coord,
 12)
 13
 14from PySDM.backends.impl_numba import conf
 15from PySDM.impl.arakawa_c import make_rhod
 16
 17
 18class MPDATA_2D:
 19    def __init__(
 20        self,
 21        *,
 22        advectees,
 23        stream_function,
 24        rhod_of_zZ,
 25        dt,
 26        grid,
 27        size,
 28        n_iters=2,
 29        infinite_gauge=True,
 30        nonoscillatory=True,
 31        third_order_terms=False
 32    ):
 33        self._grid = grid
 34        self.size = size
 35        self.dt = dt
 36        self.stream_function = stream_function
 37        self.stream_function_time_dependent = (
 38            "t" in inspect.signature(stream_function).parameters
 39        )
 40        self.asynchronous = False
 41        self.thread: (Thread, None) = None
 42        self.t = 0
 43        self.advectees = advectees
 44
 45        self._options = Options(
 46            n_iters=n_iters,
 47            infinite_gauge=infinite_gauge,
 48            nonoscillatory=nonoscillatory,
 49            third_order_terms=third_order_terms,
 50        )
 51
 52        self.g_factor = make_rhod(grid, rhod_of_zZ)
 53        self.g_factor_vec = (
 54            rhod_of_zZ(zZ=x_vec_coord(grid)[-1]),
 55            rhod_of_zZ(zZ=z_vec_coord(grid)[-1]),
 56        )
 57
 58    @cached_property
 59    def mpdatas(self):
 60        disable_threads_if_needed = {}
 61        if not conf.JIT_FLAGS["parallel"]:
 62            disable_threads_if_needed["n_threads"] = 1
 63
 64        stepper = Stepper(
 65            options=self._options,
 66            grid=self._grid,
 67            non_unit_g_factor=True,
 68            **disable_threads_if_needed
 69        )
 70
 71        advector_impl = VectorField(
 72            (
 73                np.full((self._grid[0] + 1, self._grid[1]), np.nan),
 74                np.full((self._grid[0], self._grid[1] + 1), np.nan),
 75            ),
 76            halo=self._options.n_halo,
 77            boundary_conditions=(Periodic(), Periodic()),
 78        )
 79
 80        g_factor_impl = ScalarField(
 81            self.g_factor.astype(dtype=self._options.dtype),
 82            halo=self._options.n_halo,
 83            boundary_conditions=(Periodic(), Periodic()),
 84        )
 85
 86        mpdatas = {}
 87        for k, v in self.advectees.items():
 88            advectee_impl = ScalarField(
 89                np.asarray(v, dtype=self._options.dtype),
 90                halo=self._options.n_halo,
 91                boundary_conditions=(Periodic(), Periodic()),
 92            )
 93            mpdatas[k] = Solver(
 94                stepper=stepper,
 95                advectee=advectee_impl,
 96                advector=advector_impl,
 97                g_factor=g_factor_impl,
 98            )
 99        return mpdatas
100
101    def __getitem__(self, key: str):
102        if "mpdatas" in self.__dict__:
103            return self.mpdatas[key].advectee.get()
104        return self.advectees[key]
105
106    def __call__(self, displacement):
107        if self.asynchronous:
108            self.thread = Thread(target=self.step, args=())
109            self.thread.start()
110        else:
111            self.step(displacement)
112
113    def wait(self):
114        if self.asynchronous:
115            if self.thread is not None:
116                self.thread.join()
117
118    def refresh_advector(self, displacement):
119        for mpdata in self.mpdatas.values():
120            advector = nondivergent_vector_field_2d(
121                self._grid, self.size, self.dt, self.stream_function, t=self.t
122            )
123            for d in range(len(self._grid)):
124                np.testing.assert_array_less(np.abs(advector[d]), 1)
125                mpdata.advector.get_component(d)[:] = advector[d]
126            if displacement is not None:
127                for d in range(len(self._grid)):
128                    advector[d] /= self.g_factor_vec[d]
129                displacement.upload_courant_field(advector)
130            break  # the advector field is shared
131
132    def step(self, displacement):
133        if not self.stream_function_time_dependent and self.t == 0:
134            self.refresh_advector(displacement)
135
136        self.t += 0.5 * self.dt
137        if self.stream_function_time_dependent:
138            self.refresh_advector(displacement)
139        for mpdata in self.mpdatas.values():
140            mpdata.advance(1)
141        self.t += 0.5 * self.dt