test_walberla_codegen.py 3.5 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
2
import unittest

3
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
4

5
import pystencils as ps
Martin Bauer's avatar
Martin Bauer committed
6
from lbmpy.boundaries import UBB, NoSlip
Martin Bauer's avatar
Martin Bauer committed
7
from lbmpy.creationfunctions import create_lb_method, create_lb_update_rule, create_lb_collision_rule
Martin Bauer's avatar
Martin Bauer committed
8
from lbmpy_walberla import RefinementScaling, generate_boundary, generate_lattice_model
9
10
from lbmpy_walberla.sparse import ListLbGenerator
from pystencils_walberla import generate_pack_info_for_field, generate_pack_info_from_kernel
Martin Bauer's avatar
Martin Bauer committed
11
from pystencils_walberla.cmake_integration import ManualCodeGenerationContext
Martin Bauer's avatar
Martin Bauer committed
12
13
14
15
16
17
18
19
20
21


class WalberlaLbmpyCodegenTest(unittest.TestCase):

    @staticmethod
    def test_lattice_model():
        with ManualCodeGenerationContext() as ctx:
            force_field = ps.fields("force(3): [3D]", layout='fzyx')
            omega = sp.Symbol("omega")

Martin Bauer's avatar
Martin Bauer committed
22
23
            cr = create_lb_collision_rule(stencil='D3Q19', method='srt', relaxation_rates=[omega], compressible=True,
                                          force_model='guo', force=force_field.center_vector)
Martin Bauer's avatar
Martin Bauer committed
24
25
26
27
28

            scaling = RefinementScaling()
            scaling.add_standard_relaxation_rate_scaling(omega)
            scaling.add_force_scaling(force_field)

Martin Bauer's avatar
Martin Bauer committed
29
30
31
            generate_lattice_model(ctx, 'SrtWithForceFieldModel', cr, refinement_scaling=scaling)
            generate_boundary(ctx, 'MyUBB', UBB([0.05, 0, 0]), cr.method)
            generate_boundary(ctx, 'MyNoSlip', NoSlip(), cr.method)
32
            assert 'static const bool compressible = true;' in ctx.files['SrtWithForceFieldModel.h']
Martin Bauer's avatar
Martin Bauer committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

    @staticmethod
    def test_sparse():
        from lbmpy.creationfunctions import create_lb_collision_rule
        from pystencils import show_code
        g = ListLbGenerator(create_lb_collision_rule())
        kernel_code = str(show_code(g.kernel()))
        assert 'num_cells' in kernel_code
        setter_code = str(show_code(g.setter_ast()))
        assert 'num_cells' in setter_code
        getter_code = str(show_code(g.getter_ast()))
        assert 'num_cells' in getter_code

    @staticmethod
    def test_pack_info():
        with ManualCodeGenerationContext() as ctx:
            f = ps.fields("f(9): [3D]")
            generate_pack_info_for_field(ctx, 'MyPackInfo1', f)

            lb_assignments = create_lb_update_rule(stencil='D3Q19', method='srt').main_assignments
            generate_pack_info_from_kernel(ctx, 'MyPackInfo2', lb_assignments)
54
55
56
57
58
59

    @staticmethod
    def test_incompressible():
        with ManualCodeGenerationContext() as ctx:
            omega = sp.Symbol("omega")

Martin Bauer's avatar
Martin Bauer committed
60
61
            cr = create_lb_collision_rule(stencil='D3Q19', method='srt', relaxation_rates=[omega], compressible=False)
            generate_lattice_model(ctx, 'Model', cr)
62
            assert 'static const bool compressible = false;' in ctx.files['Model.h']
63
64
65
66
67
68

    @staticmethod
    def test_output_field():
        with ManualCodeGenerationContext(openmp=True, double_accuracy=True) as ctx:
            omega_field = ps.fields("omega_out: [3D]", layout='fzyx')
            parameters = {
69
70
71
72
                'stencil': 'D3Q27',
                'method': 'trt-kbc-n1',
                'compressible': True,
                'entropic': True,
73
74
                'omega_output_field': omega_field,
            }
Martin Bauer's avatar
Martin Bauer committed
75
76
            cr = create_lb_collision_rule(**parameters)
            generate_lattice_model(ctx, 'Model', cr)
77
78
79
80
81
82

    @staticmethod
    def test_boundary():
        with ManualCodeGenerationContext(openmp=True, double_accuracy=True) as ctx:
            lb_method = create_lb_method(stencil='D3Q19', method='srt')
            generate_boundary(ctx, 'Boundary', NoSlip(), lb_method, target='gpu')