createindexlist.py 7.82 KB
Newer Older
1
2
import warnings

Martin Bauer's avatar
Martin Bauer committed
3
4
import numpy as np

5
try:
Martin Bauer's avatar
Martin Bauer committed
6
7
    # Try to import right away - assume compiled code is available
    # compile with: python setup.py build_ext --inplace --use-cython
Martin Bauer's avatar
Martin Bauer committed
8
9
    from pystencils.boundaries.createindexlistcython import create_boundary_neighbor_index_list_2d, \
        create_boundary_neighbor_index_list_3d, create_boundary_cell_index_list_2d, create_boundary_cell_index_list_3d
10

Martin Bauer's avatar
Martin Bauer committed
11
    cython_funcs_available = True
Martin Bauer's avatar
Martin Bauer committed
12
13
14
15
16
17
18
19
20
21
22
23
24
except ImportError:
    try:
        # If not, try development mode and import via pyximport
        import pyximport

        pyximport.install(language_level=3)
        cython_funcs_available = True
    except ImportError:
        cython_funcs_available = False
    if cython_funcs_available:
        from pystencils.boundaries.createindexlistcython import create_boundary_neighbor_index_list_2d, \
            create_boundary_neighbor_index_list_3d, create_boundary_cell_index_list_2d, \
            create_boundary_cell_index_list_3d
25

Martin Bauer's avatar
Martin Bauer committed
26
27
boundary_index_array_coordinate_names = ["x", "y", "z"]
direction_member_name = "dir"
28
29


Martin Bauer's avatar
Martin Bauer committed
30
31
def numpy_data_type_for_boundary_object(boundary_object, dim):
    coordinate_names = boundary_index_array_coordinate_names[:dim]
Martin Bauer's avatar
Martin Bauer committed
32
33
34
    return np.dtype([(name, np.int32) for name in coordinate_names]
                    + [(direction_member_name, np.int32)]
                    + [(i[0], i[1].numpy_dtype) for i in boundary_object.additional_data], align=True)
35
36


37
38
39
40
41
42
43
44
45
46
def _create_index_list_python(flag_field_arr, boundary_mask,
                              fluid_mask, stencil, single_link, inner_or_boundary=False, nr_of_ghost_layers=None):

    if inner_or_boundary and nr_of_ghost_layers is None:
        raise ValueError("If inner_or_boundary is set True the number of ghost layers "
                         "around the inner domain has to be specified")

    if nr_of_ghost_layers is None:
        nr_of_ghost_layers = 0

Martin Bauer's avatar
Martin Bauer committed
47
48
    coordinate_names = boundary_index_array_coordinate_names[:len(flag_field_arr.shape)]
    index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)])
49

50
51
52
53
54
55
56
57
58
59
    # boundary cells are extracted via np.where. To ensure continous memory access in the compute kernel these cells
    # have to be sorted.
    boundary_cells = np.transpose(np.nonzero(flag_field_arr == boundary_mask))
    for i in range(len(flag_field_arr.shape)):
        boundary_cells = boundary_cells[boundary_cells[:, i].argsort(kind='mergesort')]

    # First a set is created to save all fluid cells which are near boundary
    fluid_cells = set()
    for cell in boundary_cells:
        cell = tuple(cell)
Martin Bauer's avatar
Martin Bauer committed
60
61
        for dir_idx, direction in enumerate(stencil):
            neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)])
62
63
64
65
66
67
            # prevent out ouf bounds access. If boundary cell is at the border, some stencil directions would be out.
            if any(not 0 + nr_of_ghost_layers <= e < upper - nr_of_ghost_layers
                   for e, upper in zip(neighbor_cell, flag_field_arr.shape)):
                continue
            if flag_field_arr[neighbor_cell] & fluid_mask:
                fluid_cells.add(neighbor_cell)
68

69
70
71
72
73
74
    # then this is set is transformed to a list to make it sortable. This ensures continoous memory access later.
    fluid_cells = list(fluid_cells)
    if len(flag_field_arr.shape) == 3:
        fluid_cells.sort(key=lambda tup: (tup[-1], tup[-2], tup[0]))
    else:
        fluid_cells.sort(key=lambda tup: (tup[-1], tup[0]))
75

76
77
    cells_to_iterate = fluid_cells if inner_or_boundary else boundary_cells
    checkmask = boundary_mask if inner_or_boundary else fluid_mask
Martin Bauer's avatar
Martin Bauer committed
78
79

    result = []
80
81
82
    for cell in cells_to_iterate:
        cell = tuple(cell)
        sum_cells = np.zeros(len(cell))
Martin Bauer's avatar
Martin Bauer committed
83
84
        for dir_idx, direction in enumerate(stencil):
            neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)])
85
            # prevent out ouf bounds access. If boundary cell is at the border, some stencil directions would be out.
86
87
            if any(not 0 <= e < upper for e, upper in zip(neighbor_cell, flag_field_arr.shape)):
                continue
88
            if flag_field_arr[neighbor_cell] & checkmask:
Martin Bauer's avatar
Martin Bauer committed
89
                if single_link:
90
91
92
93
94
95
96
97
                    sum_cells += np.array(direction)
                else:
                    result.append(tuple(cell) + (dir_idx,))

        # the discrete normal direction is the one which gives the maximum inner product to the stencil direction
        if single_link and any(sum_cells != 0):
            idx = np.argmax(np.inner(sum_cells, stencil))
            result.append(tuple(cell) + (idx,))
Martin Bauer's avatar
Martin Bauer committed
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

    return np.array(result, dtype=index_arr_dtype)


def create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask,
                               nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False):
    """Creates a numpy array storing links (connections) between domain cells and boundary cells.

    Args:
        flag_field: flag integer array where boundary and domain cells are marked (interpreted as bit vector)
        stencil: list of directions, for possible links. When single_link is set to true the order matters, because
                 then only the first link is added to the list
        boundary_mask: cells where (cell & mask) is true are considered boundary cells
        fluid_mask: cells where (cell & mask) is true are considered fluid/inner cells cells
        nr_of_ghost_layers: only relevant if neighbors is True
        inner_or_boundary: if true, the result contains the cell coordinates of the domain cells -
                    if false the boundary cells are listed
115
        single_link: if true only the link in normal direction to this cell is reported
Martin Bauer's avatar
Martin Bauer committed
116
117

    """
Martin Bauer's avatar
Martin Bauer committed
118
119
120
    dim = len(flag_field.shape)
    coordinate_names = boundary_index_array_coordinate_names[:dim]
    index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)])
121

Martin Bauer's avatar
Martin Bauer committed
122
123
    stencil = np.array(stencil, dtype=np.int32)
    args = (flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil, single_link)
124
    args_no_gl = (flag_field, boundary_mask, fluid_mask, stencil, single_link)
Martin Bauer's avatar
Martin Bauer committed
125

Martin Bauer's avatar
Martin Bauer committed
126
    if cython_funcs_available:
127
        if dim == 2:
Martin Bauer's avatar
Martin Bauer committed
128
129
130
            if inner_or_boundary:
                idx_list = create_boundary_neighbor_index_list_2d(*args)
            else:
131
                idx_list = create_boundary_cell_index_list_2d(*args_no_gl)
132
        elif dim == 3:
Martin Bauer's avatar
Martin Bauer committed
133
134
135
            if inner_or_boundary:
                idx_list = create_boundary_neighbor_index_list_3d(*args)
            else:
136
                idx_list = create_boundary_cell_index_list_3d(*args_no_gl)
137
138
        else:
            raise ValueError("Flag field has to be a 2 or 3 dimensional numpy array")
Martin Bauer's avatar
Martin Bauer committed
139
        return np.array(idx_list, dtype=index_arr_dtype)
140
    else:
Martin Bauer's avatar
Martin Bauer committed
141
        if flag_field.size > 1e6:
142
            warnings.warn("Boundary setup may take very long! Consider installing cython to speed it up")
143
144
        return _create_index_list_python(*args_no_gl, inner_or_boundary=inner_or_boundary,
                                         nr_of_ghost_layers=nr_of_ghost_layers)
145
146


Martin Bauer's avatar
Martin Bauer committed
147
148
149
150
def create_boundary_index_array(flag_field, stencil, boundary_mask, fluid_mask, boundary_object,
                                nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False):
    idx_array = create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask,
                                           nr_of_ghost_layers, inner_or_boundary, single_link)
Martin Bauer's avatar
Martin Bauer committed
151
    dim = len(flag_field.shape)
152

Martin Bauer's avatar
Martin Bauer committed
153
154
155
156
157
158
    if boundary_object.additional_data:
        coordinate_names = boundary_index_array_coordinate_names[:dim]
        index_arr_dtype = numpy_data_type_for_boundary_object(boundary_object, dim)
        extended_idx_field = np.empty(len(idx_array), dtype=index_arr_dtype)
        for prop in coordinate_names + ['dir']:
            extended_idx_field[prop] = idx_array[prop]
159

Martin Bauer's avatar
Martin Bauer committed
160
        idx_array = extended_idx_field
161

Martin Bauer's avatar
Martin Bauer committed
162
    return idx_array