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