1from typing import Any, List, Tuple 

2 

3from pystencils import Assignment 

4from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo 

5from pystencils.data_types import create_type 

6 

7 

8class Boundary: 

9 """Base class all boundaries should derive from""" 

10 

11 inner_or_boundary = True 

12 single_link = False 

13 

14 def __init__(self, name=None): 

15 self._name = name 

16 

17 def __call__(self, field, direction_symbol, index_field) -> List[Assignment]: 

18 """Defines the boundary behavior and must therefore be implemented by all boundaries. 

19 

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

21 

22 Args: 

23 field: pystencils field where boundary condition should be applied. 

24 The current cell is cell next to the boundary, which is influenced by the boundary 

25 cell i.e. has a link from the boundary cell to itself. 

26 direction_symbol: a sympy symbol that can be used as index to the pdf_field. It describes 

27 the direction pointing from the fluid to the boundary cell 

28 index_field: the boundary index field that can be used to retrieve and update boundary data 

29 """ 

30 raise NotImplementedError("Boundary class has to overwrite __call__") 

31 

32 @property 

33 def additional_data(self) -> Tuple[str, Any]: 

34 """Return a list of (name, type) tuples for additional data items required in this boundary 

35 These data items can either be initialized in separate kernel see additional_data_kernel_init or by 

36 Python callbacks - see additional_data_callback """ 

37 return () 

38 

39 @property 

40 def additional_data_init_callback(self): 

41 """Return a callback function called with a boundary data setter object and returning a dict of 

42 data-name to data for each element that should be initialized""" 

43 return None 

44 

45 @property 

46 def name(self): 

47 if self._name: 

48 return self._name 

49 else: 

50 return type(self).__name__ 

51 

52 @name.setter 

53 def name(self, new_value): 

54 self._name = new_value 

55 

56 

57class Neumann(Boundary): 

58 

59 inner_or_boundary = False 

60 single_link = True 

61 

62 def __call__(self, field, direction_symbol, **kwargs): 

63 

64 neighbor = BoundaryOffsetInfo.offset_from_dir(direction_symbol, field.spatial_dimensions) 

65 if field.index_dimensions == 0: 

66 return [Assignment(field.center, field[neighbor])] 

67 else: 

68 from itertools import product 

69 if not field.has_fixed_index_shape: 

70 raise NotImplementedError("Neumann boundary works only for fields with fixed index shape") 

71 index_iter = product(*(range(i) for i in field.index_shape)) 

72 return [Assignment(field(*idx), field[neighbor](*idx)) for idx in index_iter] 

73 

74 def __hash__(self): 

75 # All boundaries of these class behave equal -> should also be equal 

76 return hash("Neumann") 

77 

78 def __eq__(self, other): 

79 return type(other) == Neumann 

80 

81 

82class Dirichlet(Boundary): 

83 

84 inner_or_boundary = False 

85 single_link = True 

86 

87 def __init__(self, value, name=None): 

88 super().__init__(name) 

89 self._value = value 

90 

91 @property 

92 def additional_data(self): 

93 if callable(self._value): 

94 return [('value', create_type("double"))] 

95 else: 

96 return [] 

97 

98 @property 

99 def additional_data_init_callback(self): 

100 if callable(self._value): 

101 return self._value 

102 

103 def __call__(self, field, direction_symbol, index_field, **kwargs): 

104 

105 if field.index_dimensions == 0: 

106 return [Assignment(field.center, index_field("value") if self.additional_data else self._value)] 

107 elif field.index_dimensions == 1: 

108 assert not self.additional_data 

109 if not field.has_fixed_index_shape: 

110 raise NotImplementedError("Field needs fixed index shape") 

111 assert len(self._value) == field.index_shape[0], "Dirichlet value does not match index shape of field" 

112 return [Assignment(field(i), self._value[i]) for i in range(field.index_shape[0])] 

113 raise NotImplementedError("Dirichlet boundary not implemented for fields with more than one index dimension")