PyMPDATA.options

  1"""
  2MPDATA variants, iterations, data-type and jit-flags settings
  3"""
  4
  5import numpy as np
  6from pystrict import strict
  7
  8
  9class HashableDict(dict):
 10    """serialization enabler"""
 11
 12    def __hash__(self):
 13        return hash(tuple(sorted(self.items())))
 14
 15
 16@strict
 17class Options:
 18    """representation of MPDATA algorithm variant choice, for an overview of
 19    MPDATA options implemented in PyMPDATA, see
 20    [Olesik et al. 2020](https://doi.org/10.5194/gmd-15-3879-2022);
 21    equipped with meaningful `__str__` `__hash__`, `__eq__`.
 22    """
 23
 24    def __init__(
 25        self,
 26        *,
 27        n_iters: int = 2,
 28        infinite_gauge: bool = False,
 29        divergent_flow: bool = False,
 30        nonoscillatory: bool = False,
 31        third_order_terms: bool = False,
 32        DPDC: bool = False,  # pylint: disable=invalid-name
 33        epsilon: float = 1e-15,
 34        non_zero_mu_coeff: bool = False,
 35        dynamic_advector: bool = False,
 36        dimensionally_split: bool = False,
 37        dtype: [np.float32, np.float64] = np.float64
 38    ):
 39        self._values = HashableDict(
 40            {
 41                "n_iters": n_iters,
 42                "infinite_gauge": infinite_gauge,
 43                "epsilon": epsilon,
 44                "divergent_flow": divergent_flow,
 45                "nonoscillatory": nonoscillatory,
 46                "third_order_terms": third_order_terms,
 47                "non_zero_mu_coeff": non_zero_mu_coeff,
 48                "dynamic_advector": dynamic_advector,
 49                "dimensionally_split": dimensionally_split,
 50                "dtype": dtype,
 51                "DPDC": DPDC,
 52            }
 53        )
 54
 55        if (
 56            any(
 57                (
 58                    infinite_gauge,
 59                    divergent_flow,
 60                    nonoscillatory,
 61                    third_order_terms,
 62                    DPDC,
 63                )
 64            )
 65            and n_iters < 2
 66        ):
 67            raise ValueError()
 68        if n_iters < 1:
 69            raise ValueError()
 70
 71    @property
 72    def dtype(self):
 73        """data type (e.g., np.float64)"""
 74        return self._values["dtype"]
 75
 76    @property
 77    def n_iters(self) -> int:
 78        """Number of corrective iterations in the MPDATA algorithm + 1
 79        e.g. (1: upwind, 2: upwind + one corrective iteration, ...).
 80        Bigger values mean smaller error, but more computational cost.
 81        It does not change the order of the method.
 82        The order of the method depends on the variant of antidiffusive
 83        velocity used, see for example `third_order_terms` option.
 84        Note: not to confuse with n_steps in the Stepper."""
 85        return self._values["n_iters"]
 86
 87    @property
 88    def infinite_gauge(self) -> bool:
 89        """flag enabling the infinite-gauge option, see e.g.:
 90        [Margolin & Shashkov, 2006](https://doi.org/10.1002/fld.1070),
 91        [Smolarkiewicz & Clark, 1986](https://doi.org/10.1016/0021-9991(86)90270-6)
 92        """
 93        return self._values["infinite_gauge"]
 94
 95    @property
 96    def epsilon(self) -> float:
 97        """value of constant used to prevent from divisions by zero
 98        in statements such as (a - b)/(a + b + eps)"""
 99        return self._values["epsilon"]
100
101    @property
102    def divergent_flow(self) -> bool:
103        """flag enabling the divergent-flow option, see e.g.:
104        [Smolarkiewicz, 1984](https://doi.org/10.1016/0021-9991(84)90121-9),
105        [Margolin & Smolarkiewicz, 1998](https://doi.org/10.1137/S106482759324700X)
106        """
107        return self._values["divergent_flow"]
108
109    @property
110    def nonoscillatory(self) -> bool:
111        """flag enabling the nonoscillatory option, see
112        [Smolarkiewicz & Grabowski 1990](https://doi.org/10.1016/0021-9991(90)90105-A)
113        """
114        return self._values["nonoscillatory"]
115
116    @property
117    def third_order_terms(self) -> bool:
118        """flag enabling the third-order-terms option, see
119        [Margolin & Smolarkiewicz 1998](https://doi.org/10.1137/S106482759324700X)"""
120        return self._values["third_order_terms"]
121
122    @property
123    def DPDC(self) -> bool:  # pylint: disable=invalid-name
124        """flag enabling the double-pass donor-cell option, see:
125        [Beason & Margolin, 1988](https://osti.gov/biblio/7049237),
126        [Margolin & Shashkov, 2006](https://doi.org/10.1002/fld.1070),
127        [Margolin & Smolarkiewicz, 1998](https://doi.org/10.1137/S106482759324700X)
128        """
129        return self._values["DPDC"]
130
131    @property
132    def non_zero_mu_coeff(self) -> bool:
133        """flag enabling handling of Fickian diffusion term"""
134        return self._values["non_zero_mu_coeff"]
135
136    @property
137    def dynamic_advector(self) -> bool:
138        """flag enabling (todo desc)"""
139        return self._values["dynamic_advector"]
140
141    @property
142    def dimensionally_split(self) -> bool:
143        """flag disabling cross-dimensional terms in antidiffusive velocities"""
144        return self._values["dimensionally_split"]
145
146    def __str__(self):
147        return str(self._values)
148
149    def __hash__(self):
150        value = hash(self._values) + hash(self.jit_flags)
151        return value
152
153    def __eq__(self, other):
154        return other.__hash__() == self.__hash__()
155
156    @property
157    def n_halo(self) -> int:
158        """Halo extent for a given options set.
159        The halo extent is the number of 'ghost layers' that need to be added
160        to the outside of the domain to ensure that the MPDATA stencil operations can be
161        applied to the edges of the domain.
162        It is similar to
163        [array padding](https://numpy.org/doc/stable/reference/generated/numpy.pad.html).
164        The halo extent is determined by the options set."""
165        if self.divergent_flow or self.nonoscillatory or self.third_order_terms:
166            return 2
167        return 1
168
169    @property
170    def jit_flags(self) -> HashableDict:
171        """options passed [to numba.njit()](
172        https://numba.pydata.org/numba-doc/dev/user/jit.html#compilation-options)"""
173        return HashableDict(
174            {
175                "fastmath": True,
176                "error_model": "numpy",
177                "boundscheck": False,
178            }
179        )