PyMPDATA.boundary_conditions.extrapolated

  1"""boundary condition extrapolating values from the edge to the halo for scalars
  2and returning edge-of-the-domain value for vectors (with all negative scalar values
  3set to zero)"""
  4
  5# pylint: disable=too-many-arguments
  6from functools import lru_cache
  7
  8import numba
  9
 10from PyMPDATA.impl.enumerations import (
 11    ARG_FOCUS,
 12    INNER,
 13    META_AND_DATA_DATA,
 14    META_AND_DATA_META,
 15    SIGN_LEFT,
 16)
 17from PyMPDATA.impl.traversals_common import (
 18    make_fill_halos_loop,
 19    make_fill_halos_loop_vector,
 20)
 21
 22
 23class Extrapolated:
 24    """class which instances are to be passed in boundary_conditions tuple to the
 25    `PyMPDATA.scalar_field.ScalarField` and
 26    `PyMPDATA.vector_field.VectorField` __init__ methods"""
 27
 28    def __init__(self, dim=INNER, eps=1e-10):
 29        self.eps = eps
 30        self.dim = dim
 31
 32    def make_scalar(self, indexers, halo, dtype, jit_flags, dimension_index):
 33        """returns (lru-cached) Numba-compiled scalar halo-filling callable"""
 34        return _make_scalar_extrapolated(
 35            self.dim,
 36            self.eps,
 37            indexers.ats[dimension_index],
 38            indexers.set,
 39            halo,
 40            dtype,
 41            jit_flags,
 42        )
 43
 44    @staticmethod
 45    def make_vector(indexers, _, __, jit_flags, dimension_index):
 46        """returns (lru-cached) Numba-compiled vector halo-filling callable"""
 47        return _make_vector_extrapolated(
 48            indexers.atv[dimension_index], indexers.set, jit_flags, dimension_index
 49        )
 50
 51
 52@lru_cache()
 53# pylint: disable=too-many-arguments
 54def _make_scalar_extrapolated(dim, eps, ats, set_value, halo, dtype, jit_flags):
 55    @numba.njit(**jit_flags)
 56    def impl(focus_psi, span, sign):
 57        if sign == SIGN_LEFT:
 58            edg = halo - focus_psi[ARG_FOCUS][dim]
 59            nom = ats(*focus_psi, edg + 1) - ats(*focus_psi, edg)
 60            den = ats(*focus_psi, edg + 2) - ats(*focus_psi, edg + 1)
 61            cnst = nom / den if abs(den) > eps else 0
 62            return max(
 63                ats(*focus_psi, 1) - (ats(*focus_psi, 2) - ats(*focus_psi, 1)) * cnst, 0
 64            )
 65        edg = span + halo - 1 - focus_psi[ARG_FOCUS][dim]
 66        den = ats(*focus_psi, edg - 1) - ats(*focus_psi, edg - 2)
 67        nom = ats(*focus_psi, edg) - ats(*focus_psi, edg - 1)
 68        cnst = nom / den if abs(den) > eps else 0
 69        return max(
 70            ats(*focus_psi, -1) + (ats(*focus_psi, -1) - ats(*focus_psi, -2)) * cnst, 0
 71        )
 72
 73    if dtype == complex:
 74
 75        @numba.njit(**jit_flags)
 76        def fill_halos_scalar(psi, span, sign):
 77            return complex(
 78                impl(
 79                    (psi[META_AND_DATA_META], psi[META_AND_DATA_DATA].real), span, sign
 80                ),
 81                impl(
 82                    (psi[META_AND_DATA_META], psi[META_AND_DATA_DATA].imag), span, sign
 83                ),
 84            )
 85
 86    else:
 87
 88        @numba.njit(**jit_flags)
 89        def fill_halos_scalar(psi, span, sign):
 90            return impl(psi, span, sign)
 91
 92    return make_fill_halos_loop(jit_flags, set_value, fill_halos_scalar)
 93
 94
 95@lru_cache()
 96def _make_vector_extrapolated(atv, set_value, jit_flags, dimension_index):
 97    @numba.njit(**jit_flags)
 98    def fill_halos_parallel(focus_psi, _, sign):
 99        return atv(*focus_psi, sign + 0.5)
100
101    @numba.njit(**jit_flags)
102    def fill_halos_normal(focus_psi, _, sign, __):
103        return atv(*focus_psi, sign, 0.5)
104
105    return make_fill_halos_loop_vector(
106        jit_flags, set_value, fill_halos_parallel, fill_halos_normal, dimension_index
107    )