PyMPDATA_examples.Shipway_and_Hill_2012.droplet_activation

 1from functools import lru_cache
 2
 3import numba
 4import numpy as np
 5
 6from PyMPDATA.impl.enumerations import (
 7    ARG_DATA,
 8    ARG_FOCUS,
 9    META_AND_DATA_DATA,
10    META_AND_DATA_META,
11    OUTER,
12    SIGN_RIGHT,
13)
14from PyMPDATA.impl.traversals_common import make_fill_halos_loop
15
16
17@lru_cache()
18def _make_scalar(value, set_value, halo, dtype, jit_flags, _):
19    @numba.njit(**jit_flags)
20    def impl(psi, __, sign):
21        if sign == SIGN_RIGHT:
22            return 0
23        z = psi[ARG_FOCUS][OUTER]
24        activated = np.sum(psi[ARG_DATA][z : z + 1, halo:-halo])
25        # assert activated < value
26        result = max(0, value - activated)
27        return result
28
29    if dtype == complex:
30
31        @numba.njit(**jit_flags)
32        def fill_halos_scalar(psi, n, sign):
33            return complex(
34                impl((psi[META_AND_DATA_META], psi[META_AND_DATA_DATA].real), n, sign),
35                impl((psi[META_AND_DATA_META], psi[META_AND_DATA_DATA].imag), n, sign),
36            )
37
38    else:
39
40        @numba.njit(**jit_flags)
41        def fill_halos_scalar(psi, n, sign):
42            return impl(psi, n, sign)
43
44    return make_fill_halos_loop(jit_flags, set_value, fill_halos_scalar)
45
46
47# pylint: disable=too-few-public-methods
48class DropletActivation:
49    def __init__(self, value, dr, dz):
50        self._value = value / dz / dr
51
52    def make_scalar(self, indexers, halo, dtype, jit_flags, dimension_index):
53        return _make_scalar(
54            self._value, indexers.set, halo, dtype, jit_flags, dimension_index
55        )