PyMPDATA.boundary_conditions.periodic

periodic/cyclic boundary condition logic

 1"""periodic/cyclic boundary condition logic"""
 2
 3from functools import lru_cache
 4
 5import numba
 6
 7from PyMPDATA.impl.enumerations import SIGN_LEFT, SIGN_RIGHT
 8from PyMPDATA.impl.traversals_common import (
 9    make_fill_halos_loop,
10    make_fill_halos_loop_vector,
11)
12
13
14class Periodic:
15    """class which instances are to be passed in boundary_conditions tuple to the
16    `PyMPDATA.scalar_field.ScalarField` and
17    `PyMPDATA.vector_field.VectorField` __init__ methods"""
18
19    def __init__(self):
20        assert SIGN_RIGHT == -1
21        assert SIGN_LEFT == +1
22
23    @staticmethod
24    def make_scalar(indexers, __, ___, jit_flags, dimension_index):
25        """returns (lru-cached) Numba-compiled scalar halo-filling callable"""
26        return _make_scalar_periodic(
27            indexers.ats[dimension_index], indexers.set, jit_flags
28        )
29
30    @staticmethod
31    def make_vector(indexers, __, ___, jit_flags, dimension_index):
32        """returns (lru-cached) Numba-compiled vector halo-filling callable"""
33        return _make_vector_periodic(
34            indexers.atv, indexers.set, jit_flags, dimension_index, indexers.n_dims
35        )
36
37
38@lru_cache()
39def _make_scalar_periodic(ats, set_value, jit_flags):
40    @numba.njit(**jit_flags)
41    def fill_halos_scalar(focus_psi, span, sign):
42        return ats(*focus_psi, sign * span)
43
44    return make_fill_halos_loop(jit_flags, set_value, fill_halos_scalar)
45
46
47@lru_cache()
48def _make_vector_periodic(atv, set_value, jit_flags, dimension_index, n_dims):
49    @numba.njit(**jit_flags)
50    def fill_halos_parallel(focus_psi, span, sign):
51        offset = 0.5 if sign == SIGN_LEFT else 1.5
52        return atv[dimension_index](*focus_psi, sign * (span - offset))
53
54    @numba.njit(**jit_flags)
55    def fill_halos_normal(focus_psi, span, sign, dim):
56        if n_dims == 3 and dimension_index == dim + 1:
57            return atv[dim](*focus_psi, 0.5, sign * span)
58        return atv[dimension_index](*focus_psi, sign * span, 0.5)
59
60    return make_fill_halos_loop_vector(
61        jit_flags, set_value, fill_halos_parallel, fill_halos_normal, dimension_index
62    )
class Periodic: 
15class Periodic:
16    """class which instances are to be passed in boundary_conditions tuple to the
17    `PyMPDATA.scalar_field.ScalarField` and
18    `PyMPDATA.vector_field.VectorField` __init__ methods"""
19
20    def __init__(self):
21        assert SIGN_RIGHT == -1
22        assert SIGN_LEFT == +1
23
24    @staticmethod
25    def make_scalar(indexers, __, ___, jit_flags, dimension_index):
26        """returns (lru-cached) Numba-compiled scalar halo-filling callable"""
27        return _make_scalar_periodic(
28            indexers.ats[dimension_index], indexers.set, jit_flags
29        )
30
31    @staticmethod
32    def make_vector(indexers, __, ___, jit_flags, dimension_index):
33        """returns (lru-cached) Numba-compiled vector halo-filling callable"""
34        return _make_vector_periodic(
35            indexers.atv, indexers.set, jit_flags, dimension_index, indexers.n_dims
36        )

class which instances are to be passed in boundary_conditions tuple to the PyMPDATA.scalar_field.ScalarField and PyMPDATA.vector_field.VectorField __init__ methods

@staticmethod
def make_scalar(indexers, __, ___, jit_flags, dimension_index): 
24    @staticmethod
25    def make_scalar(indexers, __, ___, jit_flags, dimension_index):
26        """returns (lru-cached) Numba-compiled scalar halo-filling callable"""
27        return _make_scalar_periodic(
28            indexers.ats[dimension_index], indexers.set, jit_flags
29        )

returns (lru-cached) Numba-compiled scalar halo-filling callable

@staticmethod
def make_vector(indexers, __, ___, jit_flags, dimension_index): 
31    @staticmethod
32    def make_vector(indexers, __, ___, jit_flags, dimension_index):
33        """returns (lru-cached) Numba-compiled vector halo-filling callable"""
34        return _make_vector_periodic(
35            indexers.atv, indexers.set, jit_flags, dimension_index, indexers.n_dims
36        )

returns (lru-cached) Numba-compiled vector halo-filling callable