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 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 PostStepNull:  # pylint: disable=too-few-public-methods
 18    """do-nothing version of the post-step hook"""
 19
 20    def __init__(self):
 21        pass
 22
 23    def call(self, psi, step):  # pylint: disable-next=unused-argument
 24        """think of it as a `__call__` method (which Numba does not allow)"""
 25
 26
 27@numba.experimental.jitclass([])
 28class PostIterNull:  # pylint: disable=too-few-public-methods
 29    """do-nothing version of the post-iter hook"""
 30
 31    def __init__(self):
 32        pass
 33
 34    def call(self, flux, g_factor, step, iteration):  # pylint: disable=unused-argument
 35        """think of it as a `__call__` method (which Numba does not allow)"""
 36
 37
 38class Solver:
 39    """Solution orchestrator requiring prior instantiation of: a `Stepper`,
 40    a scalar advectee field (that which is advected),
 41    a vector advector field (that which advects),
 42    and optionally a scalar g_factor field (used in some cases of the advection equation).
 43    Note: in some cases of advection, i.e. momentum advection,
 44    the advectee can act upon the advector.
 45    See `PyMPDATA_examples.Jarecka_et_al_2015` for an example of this.
 46    """
 47
 48    def __init__(
 49        self,
 50        stepper: Stepper,
 51        advectee: ScalarField,
 52        advector: VectorField,
 53        g_factor: [ScalarField, None] = None,
 54    ):
 55        def scalar_field(dtype=None):
 56            return ScalarField.clone(advectee, dtype=dtype)
 57
 58        def null_scalar_field():
 59            return ScalarField.make_null(advectee.n_dims, stepper.traversals)
 60
 61        def vector_field():
 62            return VectorField.clone(advector)
 63
 64        def null_vector_field():
 65            return VectorField.make_null(advector.n_dims, stepper.traversals)
 66
 67        for field in [advector, advectee] + (
 68            [g_factor] if g_factor is not None else []
 69        ):
 70            assert field.halo == stepper.options.n_halo
 71            assert field.dtype == stepper.options.dtype
 72            assert field.grid == advector.grid
 73
 74        self.__fields = {
 75            "advectee": advectee,
 76            "advector": advector,
 77            "g_factor": g_factor or null_scalar_field(),
 78            "vectmp_a": vector_field(),
 79            "vectmp_b": vector_field(),
 80            "vectmp_c": (
 81                vector_field()
 82                if stepper.options.non_zero_mu_coeff
 83                else null_vector_field()
 84            ),
 85            "nonosc_xtrm": (
 86                scalar_field(dtype=complex)
 87                if stepper.options.nonoscillatory
 88                else null_scalar_field()
 89            ),
 90            "nonosc_beta": (
 91                scalar_field(dtype=complex)
 92                if stepper.options.nonoscillatory
 93                else null_scalar_field()
 94            ),
 95        }
 96        for field in self.__fields.values():
 97            field.assemble(stepper.traversals)
 98
 99        self.__stepper = stepper
100
101    @property
102    def advectee(self) -> ScalarField:
103        """advectee scalar field (with halo), modified by advance(),
104        may be modified from user code (e.g., source-term handling)"""
105        return self.__fields["advectee"]
106
107    @property
108    def advector(self) -> VectorField:
109        """advector vector field (with halo), unmodified by advance(),
110        may be modified from user code"""
111        return self.__fields["advector"]
112
113    @property
114    def g_factor(self) -> ScalarField:
115        """G_factor field (with halo), unmodified by advance(), assumed to be constant-in-time.
116        Can be used as a Jacobian for coordinate transformations,
117        e.g. into spherical coordinates.
118        For this type of usage, see
119        `PyMPDATA_examples.Williamson_and_Rasch_1989_as_in_Jaruga_et_al_2015_Fig_14`.
120        Can also be used to account for spatial variability of fluid density, see
121        `PyMPDATA_examples.Shipway_and_Hill_2012`.
122        e.g. the changing density of a fluid."""
123        return self.__fields["g_factor"]
124
125    def advance(
126        self,
127        n_steps: int,
128        mu_coeff: Union[tuple, None] = None,
129        post_step=None,
130        post_iter=None,
131    ):
132        """advances solution by `n_steps` steps, optionally accepts: a tuple of diffusion
133        coefficients (one value per dimension) as well as `post_iter` and `post_step`
134        callbacks expected to be `numba.jitclass`es with a `call` method, for
135        signature see `PostStepNull` and `PostIterNull`;
136        returns CPU time per timestep (units as returned by `clock.clock()`)"""
137        if mu_coeff is not None:
138            assert self.__stepper.options.non_zero_mu_coeff
139        else:
140            mu_coeff = (0.0, 0.0, 0.0)
141        if (
142            self.__stepper.options.non_zero_mu_coeff
143            and not self.__fields["g_factor"].meta[META_IS_NULL]
144        ):
145            raise NotImplementedError()
146
147        post_step = post_step or PostStepNull()
148        post_iter = post_iter or PostIterNull()
149
150        return self.__stepper(
151            n_steps=n_steps,
152            mu_coeff=mu_coeff,
153            post_step=post_step,
154            post_iter=post_iter,
155            fields=self.__fields,
156        )