1import itertools 

2import warnings 

3 

4import numpy as np 

5 

6try: 

7 # Try to import right away - assume compiled code is available 

8 # compile with: python setup.py build_ext --inplace --use-cython 

9 from pystencils.boundaries.createindexlistcython import create_boundary_neighbor_index_list_2d, \ 

10 create_boundary_neighbor_index_list_3d, create_boundary_cell_index_list_2d, create_boundary_cell_index_list_3d 

11 

12 cython_funcs_available = True 

13except ImportError: 

14 try: 

15 # If not, try development mode and import via pyximport 

16 import pyximport 

17 

18 pyximport.install(language_level=3) 

19 cython_funcs_available = True 

20 except ImportError: 

21 cython_funcs_available = False 

22 if cython_funcs_available: 22 ↛ 27line 22 didn't jump to line 27, because the condition on line 22 was never false

23 from pystencils.boundaries.createindexlistcython import create_boundary_neighbor_index_list_2d, \ 

24 create_boundary_neighbor_index_list_3d, create_boundary_cell_index_list_2d, \ 

25 create_boundary_cell_index_list_3d 

26 

27boundary_index_array_coordinate_names = ["x", "y", "z"] 

28direction_member_name = "dir" 

29 

30 

31def numpy_data_type_for_boundary_object(boundary_object, dim): 

32 coordinate_names = boundary_index_array_coordinate_names[:dim] 

33 return np.dtype([(name, np.int32) for name in coordinate_names] 

34 + [(direction_member_name, np.int32)] 

35 + [(i[0], i[1].numpy_dtype) for i in boundary_object.additional_data], align=True) 

36 

37 

38def _create_boundary_neighbor_index_list_python(flag_field_arr, nr_of_ghost_layers, boundary_mask, 

39 fluid_mask, stencil, single_link): 

40 coordinate_names = boundary_index_array_coordinate_names[:len(flag_field_arr.shape)] 

41 index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)]) 

42 

43 result = [] 

44 gl = nr_of_ghost_layers 

45 for cell in itertools.product(*reversed([range(gl, i - gl) for i in flag_field_arr.shape])): 

46 cell = cell[::-1] 

47 if not flag_field_arr[cell] & fluid_mask: 

48 continue 

49 for dir_idx, direction in enumerate(stencil): 

50 neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)]) 

51 if flag_field_arr[neighbor_cell] & boundary_mask: 

52 result.append(cell + (dir_idx,)) 

53 if single_link: 

54 break 

55 

56 return np.array(result, dtype=index_arr_dtype) 

57 

58 

59def _create_boundary_cell_index_list_python(flag_field_arr, boundary_mask, 

60 fluid_mask, stencil, single_link): 

61 coordinate_names = boundary_index_array_coordinate_names[:len(flag_field_arr.shape)] 

62 index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)]) 

63 

64 result = [] 

65 for cell in itertools.product(*reversed([range(0, i) for i in flag_field_arr.shape])): 

66 cell = cell[::-1] 

67 if not flag_field_arr[cell] & boundary_mask: 

68 continue 

69 for dir_idx, direction in enumerate(stencil): 

70 neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)]) 

71 if any(not 0 <= e < upper for e, upper in zip(neighbor_cell, flag_field_arr.shape)): 

72 continue 

73 if flag_field_arr[neighbor_cell] & fluid_mask: 

74 result.append(cell + (dir_idx,)) 

75 if single_link: 

76 break 

77 

78 return np.array(result, dtype=index_arr_dtype) 

79 

80 

81def create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask, 

82 nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False): 

83 """Creates a numpy array storing links (connections) between domain cells and boundary cells. 

84 

85 Args: 

86 flag_field: flag integer array where boundary and domain cells are marked (interpreted as bit vector) 

87 stencil: list of directions, for possible links. When single_link is set to true the order matters, because 

88 then only the first link is added to the list 

89 boundary_mask: cells where (cell & mask) is true are considered boundary cells 

90 fluid_mask: cells where (cell & mask) is true are considered fluid/inner cells cells 

91 nr_of_ghost_layers: only relevant if neighbors is True 

92 inner_or_boundary: if true, the result contains the cell coordinates of the domain cells - 

93 if false the boundary cells are listed 

94 single_link: if true only the first link is reported from this cell 

95 

96 """ 

97 dim = len(flag_field.shape) 

98 coordinate_names = boundary_index_array_coordinate_names[:dim] 

99 index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)]) 

100 

101 stencil = np.array(stencil, dtype=np.int32) 

102 args = (flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil, single_link) 

103 args_no_gl = (flag_field, boundary_mask, fluid_mask, stencil, single_link) 

104 

105 if cython_funcs_available: 

106 if dim == 2: 

107 if inner_or_boundary: 

108 idx_list = create_boundary_neighbor_index_list_2d(*args) 

109 else: 

110 idx_list = create_boundary_cell_index_list_2d(*args_no_gl) 

111 elif dim == 3: 

112 if inner_or_boundary: 

113 idx_list = create_boundary_neighbor_index_list_3d(*args) 

114 else: 

115 idx_list = create_boundary_cell_index_list_3d(*args_no_gl) 

116 else: 

117 raise ValueError("Flag field has to be a 2 or 3 dimensional numpy array") 

118 return np.array(idx_list, dtype=index_arr_dtype) 

119 else: 

120 if flag_field.size > 1e6: 

121 warnings.warn("Boundary setup may take very long! Consider installing cython to speed it up") 

122 if inner_or_boundary: 

123 return _create_boundary_neighbor_index_list_python(*args) 

124 else: 

125 return _create_boundary_cell_index_list_python(*args_no_gl) 

126 

127 

128def create_boundary_index_array(flag_field, stencil, boundary_mask, fluid_mask, boundary_object, 

129 nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False): 

130 idx_array = create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask, 

131 nr_of_ghost_layers, inner_or_boundary, single_link) 

132 dim = len(flag_field.shape) 

133 

134 if boundary_object.additional_data: 

135 coordinate_names = boundary_index_array_coordinate_names[:dim] 

136 index_arr_dtype = numpy_data_type_for_boundary_object(boundary_object, dim) 

137 extended_idx_field = np.empty(len(idx_array), dtype=index_arr_dtype) 

138 for prop in coordinate_names + ['dir']: 

139 extended_idx_field[prop] = idx_array[prop] 

140 

141 idx_array = extended_idx_field 

142 

143 return idx_array