geometry.py 11.6 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
import numpy as np
2

Martin Bauer's avatar
Martin Bauer committed
3
4
5
from lbmpy.boundaries import UBB, NoSlip
from pystencils.slicing import (
    normalize_slice, shift_slice, slice_from_direction, slice_intersection)
6
7


8
9
def add_box_boundary(boundary_handling, boundary=NoSlip(), replace=True):
    """ Adds wall boundary conditions at all domain boundaries.
Martin Bauer's avatar
Martin Bauer committed
10

11
12
13
14
    Args:
        boundary_handling: boundary handling object
        boundary: the boundary to set at domain boundaries
        replace: see BoundaryHandling.set_boundary , True overwrites flag field, False only adds the boundary flag
Martin Bauer's avatar
Martin Bauer committed
15

16
17
18
    Returns:
        flag used for the wall boundary
    """
19
    borders = ['N', 'S', 'E', 'W']
Martin Bauer's avatar
Martin Bauer committed
20
    if boundary_handling.dim == 3:
21
        borders += ['T', 'B']
22
    flag = None
23
    for d in borders:
Martin Bauer's avatar
Martin Bauer committed
24
        flag = boundary_handling.set_boundary(boundary, slice_from_direction(d, boundary_handling.dim), replace=replace)
25
    assert flag is not None
Martin Bauer's avatar
Martin Bauer committed
26
    return flag
27
28


29
30
31
32
33
34
35
def add_sphere(boundary_handling, midpoint, radius, boundary=NoSlip(), replace=True):
    """Sets boundary in spherical region."""
    def set_sphere(x, y):
        return (x - midpoint[0]) ** 2 + (y - midpoint[1]) ** 2 < radius ** 2
    return boundary_handling.set_boundary(boundary, mask_callback=set_sphere, replace=replace)


36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def add_pipe_inflow_boundary(boundary_handling, u_max, slice_obj, flow_direction=0, diameter=None):
    """Adds velocity inflow UBB boundary for pipe flow.

    Args:
        boundary_handling: boundary handling object
        u_max: maximum velocity at center of the pipe
        slice_obj: slice describing where the boundary should be set
        flow_direction: dimension index, e.g. 0 for a channel that flows along x,
        diameter: pipe/diameter, if None is passed the maximum diameter is assumed

    Returns:
        flag used for the inflow boundary

    Examples:
        >>> from lbmpy.lbstep import LatticeBoltzmannStep
        >>> from pystencils import make_slice
        >>> pipe = LatticeBoltzmannStep(domain_size=(20, 10, 10), relaxation_rate=1.8, periodicity=(True, False, False))
        >>> flag = add_pipe_inflow_boundary(pipe.boundary_handling, u_max=0.05,
        ...                                 slice_obj=make_slice[0, :, :], flow_direction=0)
    """
Martin Bauer's avatar
Martin Bauer committed
56
    dim = boundary_handling.dim
57

Martin Bauer's avatar
Martin Bauer committed
58
    def velocity_info_callback(boundary_data):
59
        for i, name in enumerate(['vel_0', 'vel_1', 'vel_2'][:dim]):
60
            if i != flow_direction:
Martin Bauer's avatar
Martin Bauer committed
61
                boundary_data[name] = 0.0
62
        if diameter is None:
63
            radius = int(round(min(sh for i, sh in enumerate(boundary_handling.shape) if i != flow_direction) / 2))
64
        else:
Martin Bauer's avatar
Martin Bauer committed
65
            radius = int(round(diameter / 2))
66

67
        if dim == 3:
68
69
            normal_coord1 = (flow_direction + 1) % 3
            normal_coord2 = (flow_direction + 2) % 3
Martin Bauer's avatar
Martin Bauer committed
70
71
72
73
            y, z = boundary_data.link_positions(normal_coord1), boundary_data.link_positions(normal_coord2)
            centered_normal1 = y - int(round(boundary_handling.shape[normal_coord1] / 2))
            centered_normal2 = z - int(round(boundary_handling.shape[normal_coord2] / 2))
            dist_to_center = np.sqrt(centered_normal1 ** 2 + centered_normal2 ** 2)
74
        elif dim == 2:
75
            normal_coord = (flow_direction + 1) % 2
Martin Bauer's avatar
Martin Bauer committed
76
77
            centered_normal = boundary_data.link_positions(normal_coord) - radius
            dist_to_center = np.sqrt(centered_normal ** 2)
78
79
80
        else:
            raise ValueError("Invalid dimension")

Martin Bauer's avatar
Martin Bauer committed
81
        vel_profile = u_max * (1 - (dist_to_center / radius)**2)
82
        boundary_data['vel_%d' % (flow_direction,)] = vel_profile
83

Martin Bauer's avatar
Martin Bauer committed
84
    inflow = UBB(velocity_info_callback, dim=boundary_handling.dim)
85
    return boundary_handling.set_boundary(inflow, slice_obj=slice_obj, ghost_layers=True)
86
87


88
89
def add_pipe_walls(boundary_handling, diameter, boundary=NoSlip()):
    """Sets boundary for the wall of a pipe with flow in x direction.
90

91
92
93
94
95
96
    Args:
        boundary_handling: boundary handling object, works for serial and parallel versions
        diameter: pipe diameter, can be either a constant value or a callback function.
                  the callback function has the signature (x_coord_array, domain_shape_in_cells) and has to return
                  a array of same shape as the received x_coord_array, with the diameter for each x position
        boundary: boundary object that is set at the wall, defaults to NoSlip (bounce back)
Martin Bauer's avatar
Martin Bauer committed
97

98
99
    Returns:
        flag used for the wall boundary
Martin Bauer's avatar
Martin Bauer committed
100
    """
Martin Bauer's avatar
Martin Bauer committed
101
102
    domain_shape = boundary_handling.shape
    dim = len(domain_shape)
Martin Bauer's avatar
Martin Bauer committed
103
104
105
    assert dim in (2, 3)

    def callback(*coordinates):
Martin Bauer's avatar
Martin Bauer committed
106
        flow_coord = coordinates[0]
Martin Bauer's avatar
Martin Bauer committed
107
108

        if callable(diameter):
Martin Bauer's avatar
Martin Bauer committed
109
110
111
            flow_coord_line = flow_coord[:, 0, 0] if dim == 3 else flow_coord[:, 0]
            diameter_value = diameter(flow_coord_line, domain_shape)
            diameter_value = diameter_value[:, np.newaxis, np.newaxis] if dim == 3 else diameter_value[:, np.newaxis]
Martin Bauer's avatar
Martin Bauer committed
112
        else:
Martin Bauer's avatar
Martin Bauer committed
113
            diameter_value = diameter
Martin Bauer's avatar
Martin Bauer committed
114

Martin Bauer's avatar
Martin Bauer committed
115
        radius_sq = (diameter_value / 2) ** 2
Martin Bauer's avatar
Martin Bauer committed
116

Martin Bauer's avatar
Martin Bauer committed
117
        mid = [domain_shape[i] // 2 for i in range(1, dim)]
Martin Bauer's avatar
Martin Bauer committed
118
        distance = sum((c_i - mid_i) ** 2 for c_i, mid_i in zip(coordinates[1:], mid))
Martin Bauer's avatar
Martin Bauer committed
119
        return distance > radius_sq
Martin Bauer's avatar
Martin Bauer committed
120

121
    return boundary_handling.set_boundary(boundary, mask_callback=callback)
122
123


124
125
def get_pipe_velocity_field(domain_size, u_max, flow_direction=0, diameter=None):
    """Analytic velocity field for 2D or 3D pipe flow.
126

127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
    Args:
        domain_size: 2-tuple or 3-tuple: shape of the domain
        u_max: velocity in the middle of the pipe
        flow_direction: dimension index, e.g. 0 for a channel that flows along x,
        diameter: pipe/diameter, if None is passed the maximum diameter is assumed

    Returns:
        numpy array with velocity field, velocity values outside the pipe are invalid

    Examples:
        Set up channel flow with analytic solution
        >>> from lbmpy.scenarios import create_channel
        >>> domain_size = (10, 10, 5)
        >>> u_max = 0.05
        >>> initial_velocity = get_pipe_velocity_field(domain_size, u_max)
        >>> scenario = create_channel(domain_size=domain_size, u_max=u_max, relaxation_rate=1.8,
        ...                           initial_velocity=initial_velocity)
    """
    if diameter is None:
        radius = int(round(min(sh for i, sh in enumerate(domain_size) if i != flow_direction) / 2))
    else:
        radius = int(round(diameter / 2))

    params = [np.arange(s) + 0.5 for s in domain_size]
151
    grid = np.meshgrid(*params, indexing='ij')
152
153
154
155
156
157
158
159
160
161
162
163

    dist = 0
    for i in range(len(domain_size)):
        if i == flow_direction:
            continue
        center = int(round(domain_size[i] / 2))
        dist += (grid[i] - center) ** 2
    dist = np.sqrt(dist)

    u = np.zeros(tuple(domain_size) + (len(domain_size),))
    u[..., flow_direction] = u_max * (1 - (dist / radius) ** 2)
    return u
164
165


Martin Bauer's avatar
Martin Bauer committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def get_shear_flow_velocity_field(domain_size, u_max, random_y_factor=0.01):
    """Returns a velocity field with two shear layers.

    -----------------------------
    | ->  ->  ->  ->  ->  ->  ->| Upper layer moves to the right
    | ->  ->  ->  ->  ->  ->  ->|
    | <-  <-  <-  <-  <-  <-  <-| Middle layer to the left
    | <-  <-  <-  <-  <-  <-  <-|
    | ->  ->  ->  ->  ->  ->  ->| Lower layer to the right again
    | ->  ->  ->  ->  ->  ->  ->|
    -----------------------------

    Args:
        domain_size: size of the domain, tuple with (x,y) or (x,y,z)
        u_max: magnitude of x component of the velocity
        random_y_factor: to break the symmetry the y velocity component is initialized randomly
                         its maximum value is u_max * random_y_factor
    """
    height = domain_size[2] if len(domain_size) == 3 else domain_size[1]
    velocity = np.empty(tuple(domain_size) + (len(domain_size),))
186
187
188
    velocity[..., 0] = u_max
    velocity[..., height // 3: height // 3 * 2, 0] = -u_max
    for j in range(1, len(domain_size)):
189
        velocity[..., j] = random_y_factor * u_max * (np.random.rand(*domain_size) - 0.5)
Martin Bauer's avatar
Martin Bauer committed
190
191
192
    return velocity


Martin Bauer's avatar
Martin Bauer committed
193
194
def add_black_and_white_image(boundary_handling, image_file, target_slice=None, plane=(0, 1), boundary=NoSlip(),
                              keep_aspect_ratio=False):
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
    """Sets boundary from a black and white image, setting boundary where image is black.

    For 3D domains the image is extruded along a coordinate. Requires scipy for image resizing.

    Args:
        boundary_handling: boundary handling object
        image_file: path to image file
        target_slice: rectangular sub-region where the image should be set, or None for everywhere
        plane: 2-tuple with coordinate indices, (0, 1) means that the image is mapped into the x-y plane and
               extruded in z direction
        boundary: boundary object to set, where the image has black pixels
        keep_aspect_ratio: by default the image is reshaped to match the target_slice. If this parameter is True the
                           aspect ratio is kept, effectively making the target_slice smaller

    Returns:
        numpy array with velocity field, velocity values outside the pipe are invalid
    """
212
    from scipy.ndimage import zoom
213

Martin Bauer's avatar
Martin Bauer committed
214
215
216
    domain_size = boundary_handling.shape
    if target_slice is None:
        target_slice = [slice(None, None, None)] * len(domain_size)
217

Martin Bauer's avatar
Martin Bauer committed
218
    dim = boundary_handling.dim
219

Martin Bauer's avatar
Martin Bauer committed
220
221
    image_slice = normalize_slice(target_slice, domain_size)
    target_size = [image_slice[i].stop - image_slice[i].start for i in plane]
222

223
    img_arr = _read_image(image_file, flatten=True).astype(int)
Martin Bauer's avatar
Martin Bauer committed
224
    img_arr = np.rot90(img_arr, 3)
225

Martin Bauer's avatar
Martin Bauer committed
226
227
228
229
    zoom_factor = [target_size[i] / img_arr.shape[i] for i in range(2)]
    if keep_aspect_ratio:
        zoom_factor = min(zoom_factor)
    zoomed_image = zoom(img_arr, zoom_factor, order=0)
230
231

    # binarize
Martin Bauer's avatar
Martin Bauer committed
232
233
    zoomed_image[zoomed_image <= 254] = 0
    zoomed_image[zoomed_image > 254] = 1
Markus Holzer's avatar
Markus Holzer committed
234
    zoomed_image = np.logical_not(zoomed_image.astype(bool))
235
236

    # resize necessary if aspect ratio should be constant
Martin Bauer's avatar
Martin Bauer committed
237
    if zoomed_image.shape != target_size:
Markus Holzer's avatar
Markus Holzer committed
238
        resized_image = np.zeros(target_size, dtype=bool)
Martin Bauer's avatar
Martin Bauer committed
239
240
        mid = [(ts - s) // 2 for ts, s in zip(target_size, zoomed_image.shape)]
        resized_image[mid[0]:zoomed_image.shape[0] + mid[0], mid[1]:zoomed_image.shape[1] + mid[1]] = zoomed_image
Martin Bauer's avatar
Martin Bauer committed
241
        zoomed_image = resized_image
242
243

    def callback(*coordinates):
Markus Holzer's avatar
Markus Holzer committed
244
        result = np.zeros_like(coordinates[0], dtype=bool)
Martin Bauer's avatar
Martin Bauer committed
245
246
        mask_start = [int(coordinates[i][(0,) * dim] - 0.5) for i in range(dim)]
        mask_end = [int(coordinates[i][(-1,) * dim] + 1 - 0.5) for i in range(dim)]
247

Martin Bauer's avatar
Martin Bauer committed
248
249
250
        mask_slice = [slice(start, stop) for start, stop in zip(mask_start, mask_end)]
        intersection_slice = slice_intersection(mask_slice, image_slice)
        if intersection_slice is None:
251
252
            return result
        else:
Martin Bauer's avatar
Martin Bauer committed
253
254
255
256
            mask_target_slice = shift_slice(intersection_slice, [-e for e in mask_start])
            image_target_slice = shift_slice(intersection_slice, [-s.start for s in image_slice])
            mask_target_slice = [mask_target_slice[i] if i in plane else slice(None, None) for i in range(dim)]
            image_target_slice = [image_target_slice[i] if i in plane else np.newaxis for i in range(dim)]
257
            result[tuple(mask_target_slice)] = zoomed_image[tuple(image_target_slice)]
258
259
            return result

260
261
262
263
264
265
266
267
268
269
270
271
272
273
    return boundary_handling.set_boundary(boundary, slice_obj=image_slice, mask_callback=callback,
                                          ghost_layers=False, inner_ghost_layers=True)


def _read_image(path, flatten=False):
    try:
        from PIL import Image
    except ImportError:
        raise ImportError("Image loading failed. Required package 'pillow' is missing")

    im = Image.open(path)
    if flatten:
        im = im.convert('F')
    return np.array(im)