PyMPDATA.solver

  1"""
  2class grouping user-supplied stepper, fields and post-step/post-iter hooks,
  3as well as self-initialised temporary storage
  4"""
  5
  6from typing import Hashable, Iterable, Mapping, Union
  7
  8import numba
  9
 10from .impl.meta import META_IS_NULL
 11from .scalar_field import ScalarField
 12from .stepper import Stepper
 13from .vector_field import VectorField
 14
 15
 16@numba.experimental.jitclass([])
 17class AnteStepNull:  # pylint: disable=too-few-public-methods
 18    """do-nothing version of the ante-step hook"""
 19
 20    def __init__(self):
 21        pass
 22
 23    def call(
 24        self,
 25        traversals_data,
 26        advectee,
 27        advector,
 28        step,
 29        index,
 30        dynamic_advector_stash_outer,
 31        dynamic_advector_stash_mid3d,
 32        dynamic_advector_stash_inner,
 33    ):  # pylint: disable-next=unused-argument,disable=too-many-arguments
 34        """think of it as a `__call__` method (which Numba does not allow)"""
 35
 36
 37@numba.experimental.jitclass([])
 38class PostStepNull:  # pylint: disable=too-few-public-methods
 39    """do-nothing version of the post-step hook"""
 40
 41    def __init__(self):
 42        pass
 43
 44    def call(
 45        self, traversals_data, psi, step, index
 46    ):  # pylint: disable-next=unused-argument
 47        """think of it as a `__call__` method (which Numba does not allow)"""
 48
 49
 50@numba.experimental.jitclass([])
 51class PostIterNull:  # pylint: disable=too-few-public-methods
 52    """do-nothing version of the post-iter hook"""
 53
 54    def __init__(self):
 55        pass
 56
 57    def call(
 58        self, traversals_data, flux, g_factor, step, iteration
 59    ):  # pylint: disable=unused-argument,disable=too-many-arguments
 60        """think of it as a `__call__` method (which Numba does not allow)"""
 61
 62
 63class Solver:
 64    """Solution orchestrator requiring prior instantiation of: a `Stepper`,
 65    a scalar advectee field (that which is advected),
 66    a vector advector field (that which advects),
 67    and optionally a scalar g_factor field (used in some cases of the advection equation).
 68    Note: in some cases of advection, i.e. momentum advection,
 69    the advectee can act upon the advector.
 70    See `PyMPDATA_examples.Jarecka_et_al_2015` for an example of this.
 71    """
 72
 73    def __init__(
 74        self,
 75        stepper: Stepper,
 76        advectee: [ScalarField, Iterable[ScalarField], Mapping[Hashable, ScalarField]],
 77        advector: VectorField,
 78        g_factor: [ScalarField, None] = None,
 79    ):
 80        self.advectee = advectee
 81        n_dims = advector.n_dims
 82        if isinstance(advectee, ScalarField):
 83            self.__fields = {"advectee": (advectee,)}
 84
 85        elif isinstance(advectee, Mapping):
 86            self.__fields = {"advectee": tuple(advectee.values())}
 87
 88        elif isinstance(advectee, Iterable):
 89            self.__fields = {"advectee": tuple(advectee)}
 90
 91        def scalar_field(dtype=None):
 92            return ScalarField.clone(self.__fields["advectee"][0], dtype=dtype)
 93
 94        def null_scalar_field():
 95            return ScalarField.make_null(n_dims, stepper.traversals)
 96
 97        def vector_field():
 98            return VectorField.clone(advector)
 99
100        def null_vector_field():
101            return VectorField.make_null(n_dims, stepper.traversals)
102
103        for field in [advector, *self.__fields["advectee"]] + (
104            [g_factor] if g_factor is not None else []
105        ):
106            assert field.halo == stepper.options.n_halo
107            assert field.dtype == stepper.options.dtype
108            assert field.grid == advector.grid
109
110        self.__fields.update(
111            {
112                "advector": advector,
113                "g_factor": g_factor or null_scalar_field(),
114                "vectmp_a": vector_field(),
115                "vectmp_b": vector_field(),
116                "vectmp_c": (
117                    vector_field()
118                    if stepper.options.non_zero_mu_coeff
119                    else null_vector_field()
120                ),
121                "dynamic_advector_stash_outer": (
122                    scalar_field()
123                    if stepper.options.dynamic_advector and n_dims > 1
124                    else null_scalar_field()
125                ),
126                "dynamic_advector_stash_mid3d": (
127                    scalar_field()
128                    if stepper.options.dynamic_advector and n_dims > 2
129                    else null_scalar_field()
130                ),
131                "dynamic_advector_stash_inner": (
132                    scalar_field()
133                    if stepper.options.dynamic_advector
134                    else null_scalar_field()
135                ),
136                "nonosc_xtrm": (
137                    scalar_field(dtype=complex)
138                    if stepper.options.nonoscillatory
139                    else null_scalar_field()
140                ),
141                "nonosc_beta": (
142                    scalar_field(dtype=complex)
143                    if stepper.options.nonoscillatory
144                    else null_scalar_field()
145                ),
146            }
147        )
148        for key, value in self.__fields.items():
149            for field in (value,) if key != "advectee" else value:
150                field.assemble(stepper.traversals)
151
152        self.__stepper = stepper
153
154    @property
155    def advector(self) -> VectorField:
156        """advector vector field , dynamic_advector_stash_outer,
157        dynamic_advector_stash_mid3d, dynamic_advector_stash_inner(with halo),
158        unmodified by advance(), may be modified from user code"""
159        return self.__fields["advector"]
160
161    @property
162    def g_factor(self) -> ScalarField:
163        """G_factor field (with halo), unmodified by advance(), assumed to be constant-in-time.
164        Can be used as a Jacobian for coordinate transformations,
165        e.g. into spherical coordinates.
166        For this type of usage, see
167        `PyMPDATA_examples.Williamson_and_Rasch_1989_as_in_Jaruga_et_al_2015_Fig_14`.
168        Can also be used to account for spatial variability of fluid density, see
169        `PyMPDATA_examples.Shipway_and_Hill_2012`.
170        e.g. the changing density of a fluid."""
171        return self.__fields["g_factor"]
172
173    def advance(
174        self,
175        n_steps: int,
176        *,
177        mu_coeff: Union[tuple, None] = None,
178        ante_step=None,
179        post_step=None,
180        post_iter=None,
181    ):
182        """advances solution by `n_steps` steps, optionally accepts: a tuple of diffusion
183        coefficients (one value per dimension) as well as `post_iter` and `post_step`
184        callbacks expected to be `numba.jitclass`es with a `call` method, for
185        signature see `PostStepNull` and `PostIterNull`;
186        returns CPU time per timestep (units as returned by `clock.clock()`)"""
187        if mu_coeff is not None:
188            assert self.__stepper.options.non_zero_mu_coeff
189        else:
190            mu_coeff = (0.0, 0.0, 0.0)
191        if (
192            self.__stepper.options.non_zero_mu_coeff
193            and not self.__fields["g_factor"].meta[META_IS_NULL]
194        ):
195            raise NotImplementedError()
196
197        ante_step = ante_step or AnteStepNull()
198        post_step = post_step or PostStepNull()
199        post_iter = post_iter or PostIterNull()
200
201        return self.__stepper(
202            n_steps=n_steps,
203            mu_coeff=mu_coeff,
204            ante_step=ante_step,
205            post_step=post_step,
206            post_iter=post_iter,
207            fields=self.__fields,
208        )