test_boundary.py 4.55 KB
Newer Older
1
import os
Martin Bauer's avatar
Martin Bauer committed
2
3
from tempfile import TemporaryDirectory

4
import numpy as np
Michael Kuron's avatar
Michael Kuron committed
5
6
import pytest

Martin Bauer's avatar
Martin Bauer committed
7
from pystencils import Assignment, create_kernel
8
from pystencils.boundaries import BoundaryHandling, Dirichlet, Neumann, add_neumann_boundary
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from pystencils.datahandling import SerialDataHandling
from pystencils.slicing import slice_from_direction


def test_kernel_vs_copy_boundary():
    dh = SerialDataHandling(domain_size=(7, 7))
    src = dh.add_array('src')
    dst_builtin = dh.add_array_like('dst_builtin', 'src')
    dst_python_copy = dh.add_array_like('dst_python_copy', 'src')
    dst_handling = dh.add_array_like('dst_handling', 'src')

    src_arr = np.arange(dh.shape[0] * dh.shape[1]).reshape(dh.shape)

    def reset_src():
        for block in dh.iterate(ghost_layers=True, inner_ghost_layers=True):
            np.copyto(block['src'], np.random.rand(*block.shape))

        for block in dh.iterate(ghost_layers=False, inner_ghost_layers=True):
            np.copyto(block['src'], src_arr)

    for b in dh.iterate(ghost_layers=False, inner_ghost_layers=True):
        np.copyto(b['dst_builtin'], 42)
        np.copyto(b['dst_python_copy'], 43)
        np.copyto(b['dst_handling'], 44)

    flags = dh.add_array('flags', dtype=np.uint8)
    dh.fill(flags.name, 0)
    borders = ['N', 'S', 'E', 'W']
    for d in borders:
        dh.fill(flags.name, 1, slice_obj=slice_from_direction(d, dim=2), ghost_layers=True, inner_ghost_layers=True)

    rhs = sum(src.neighbors([(1, 0), (-1, 0), (0, 1), (0, -1)]))

    simple_kernel = create_kernel([Assignment(dst_python_copy.center, rhs)]).compile()
    kernel_handling = create_kernel([Assignment(dst_handling.center, rhs)]).compile()

    assignments_with_boundary = add_neumann_boundary([Assignment(dst_builtin.center, rhs)],
                                                     fields=[src], flag_field=flags, boundary_flag=1)
    kernel_with_boundary = create_kernel(assignments_with_boundary).compile()

    # ------ Method 1: Built-in boundary
    reset_src()
    dh.run_kernel(kernel_with_boundary)

    # ------ Method 2: Using python to copy out the values (reference)
    reset_src()
    for b in dh.iterate():
        arr = b['src']
        arr[:, 0] = arr[:, 1]
        arr[:, -1] = arr[:, -2]
        arr[0, :] = arr[1, :]
        arr[-1, :] = arr[-2, :]
    dh.run_kernel(simple_kernel)

    # ------ Method 3: Using boundary handling to copy out the values
    reset_src()
    boundary_stencil = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    boundary_handling = BoundaryHandling(dh, src.name, boundary_stencil)
    neumann = Neumann()
    assert neumann.name == 'Neumann'
    neumann.name = "wall"
    assert neumann.name == 'wall'
    assert neumann.additional_data_init_callback is None
    assert len(neumann.additional_data) == 0

    for d in ('N', 'S', 'W', 'E'):
        boundary_handling.set_boundary(neumann, slice_from_direction(d, dim=2))
    boundary_handling()
    dh.run_kernel(kernel_handling)

    python_copy_result = dh.gather_array('dst_python_copy')
    builtin_result = dh.gather_array('dst_builtin')
    handling_result = dh.gather_array('dst_handling')

    np.testing.assert_almost_equal(python_copy_result, builtin_result)
    np.testing.assert_almost_equal(python_copy_result, handling_result)

    with TemporaryDirectory() as tmp_dir:
Michael Kuron's avatar
Michael Kuron committed
87
        pytest.importorskip('pyevtk')
88
89
        boundary_handling.geometry_to_vtk(file_name=os.path.join(tmp_dir, 'test_output1'), ghost_layers=False)
        boundary_handling.geometry_to_vtk(file_name=os.path.join(tmp_dir, 'test_output2'), ghost_layers=True)
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113


@pytest.mark.parametrize('with_indices', ('with_indices', False))
def test_dirichlet(with_indices):
    value = (1, 20, 3) if with_indices else 1

    dh = SerialDataHandling(domain_size=(7, 7))
    src = dh.add_array('src', values_per_cell=3 if with_indices else 1)
    dh.cpu_arrays.src[...] = np.random.rand(*src.shape)
    boundary_stencil = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    boundary_handling = BoundaryHandling(dh, src.name, boundary_stencil)
    dirichlet = Dirichlet(value)
    assert dirichlet.name == 'Dirichlet'
    dirichlet.name = "wall"
    assert dirichlet.name == 'wall'

    for d in ('N', 'S', 'W', 'E'):
        boundary_handling.set_boundary(dirichlet, slice_from_direction(d, dim=2))
    boundary_handling()

    assert all([np.allclose(a, np.array(value)) for a in dh.cpu_arrays.src[1:-2, 0]])
    assert all([np.allclose(a, np.array(value)) for a in dh.cpu_arrays.src[1:-2, -1]])
    assert all([np.allclose(a, np.array(value)) for a in dh.cpu_arrays.src[0, 1:-2]])
    assert all([np.allclose(a, np.array(value)) for a in dh.cpu_arrays.src[-1, 1:-2]])