boundaryconditions.py 4.21 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
2
from typing import Any, List, Tuple

3
from pystencils import Assignment
4
from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
5
from pystencils.data_types import create_type
6
7


8
class Boundary:
9
10
    """Base class all boundaries should derive from"""

Martin Bauer's avatar
Martin Bauer committed
11
12
13
    inner_or_boundary = True
    single_link = False

14
15
16
    def __init__(self, name=None):
        self._name = name

Martin Bauer's avatar
Martin Bauer committed
17
18
19
20
21
22
23
24
25
    def __call__(self, field, direction_symbol, index_field) -> List[Assignment]:
        """Defines the boundary behavior and must therefore be implemented by all boundaries.

        Here the boundary is defined as a list of sympy assignments, from which a boundary kernel is generated.

        Args:
            field: pystencils field where boundary condition should be applied.
                   The current cell is cell next to the boundary, which is influenced by the boundary
                   cell i.e. has a link from the boundary cell to itself.
Martin Bauer's avatar
Martin Bauer committed
26
            direction_symbol: a sympy symbol that can be used as index to the pdf_field. It describes
Martin Bauer's avatar
Martin Bauer committed
27
28
                              the direction pointing from the fluid to the boundary cell
            index_field: the boundary index field that can be used to retrieve and update boundary data
29
30
31
32
        """
        raise NotImplementedError("Boundary class has to overwrite __call__")

    @property
Martin Bauer's avatar
Martin Bauer committed
33
    def additional_data(self) -> Tuple[str, Any]:
34
        """Return a list of (name, type) tuples for additional data items required in this boundary
Martin Bauer's avatar
Martin Bauer committed
35
36
        These data items can either be initialized in separate kernel see additional_data_kernel_init or by
        Python callbacks - see additional_data_callback """
37
        return ()
38
39

    @property
Martin Bauer's avatar
Martin Bauer committed
40
    def additional_data_init_callback(self):
41
42
43
44
45
46
47
48
49
50
51
52
        """Return a callback function called with a boundary data setter object and returning a dict of
        data-name to data for each element that should be initialized"""
        return None

    @property
    def name(self):
        if self._name:
            return self._name
        else:
            return type(self).__name__

    @name.setter
Martin Bauer's avatar
Martin Bauer committed
53
54
    def name(self, new_value):
        self._name = new_value
55
56
57


class Neumann(Boundary):
Martin Bauer's avatar
Martin Bauer committed
58
59
60
61

    inner_or_boundary = False
    single_link = True

Martin Bauer's avatar
Martin Bauer committed
62
    def __call__(self, field, direction_symbol, **kwargs):
63

Martin Bauer's avatar
Martin Bauer committed
64
65
        neighbor = BoundaryOffsetInfo.offset_from_dir(direction_symbol, field.spatial_dimensions)
        if field.index_dimensions == 0:
Martin Bauer's avatar
Martin Bauer committed
66
            return [Assignment(field.center, field[neighbor])]
67
68
        else:
            from itertools import product
Martin Bauer's avatar
Martin Bauer committed
69
            if not field.has_fixed_index_shape:
70
                raise NotImplementedError("Neumann boundary works only for fields with fixed index shape")
Martin Bauer's avatar
Martin Bauer committed
71
            index_iter = product(*(range(i) for i in field.index_shape))
Martin Bauer's avatar
Martin Bauer committed
72
            return [Assignment(field(*idx), field[neighbor](*idx)) for idx in index_iter]
73
74
75
76
77
78
79

    def __hash__(self):
        # All boundaries of these class behave equal -> should also be equal
        return hash("Neumann")

    def __eq__(self, other):
        return type(other) == Neumann
80
81
82


class Dirichlet(Boundary):
Martin Bauer's avatar
Martin Bauer committed
83
84
85
86

    inner_or_boundary = False
    single_link = True

87
    def __init__(self, value, name=None):
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        super().__init__(name)
        self._value = value

    @property
    def additional_data(self):
        if callable(self._value):
            return [('value', create_type("double"))]
        else:
            return []

    @property
    def additional_data_init_callback(self):
        if callable(self._value):
            return self._value

    def __call__(self, field, direction_symbol, index_field, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
104

105
        if field.index_dimensions == 0:
106
            return [Assignment(field.center, index_field("value") if self.additional_data else self._value)]
Martin Bauer's avatar
Martin Bauer committed
107
108
109
110
111
112
113
        elif field.index_dimensions == 1:
            assert not self.additional_data
            if not field.has_fixed_index_shape:
                raise NotImplementedError("Field needs fixed index shape")
            assert len(self._value) == field.index_shape[0], "Dirichlet value does not match index shape of field"
            return [Assignment(field(i), self._value[i]) for i in range(field.index_shape[0])]
        raise NotImplementedError("Dirichlet boundary not implemented for fields with more than one index dimension")