test_blocking_staggered.py 880 Bytes
Newer Older
1
import numpy as np
Martin Bauer's avatar
Martin Bauer committed
2

3
4
5
6
import pystencils as ps


def test_blocking_staggered():
Michael Kuron's avatar
Michael Kuron committed
7
    f = ps.fields("f: double[3D]")
8
    stag = ps.fields("stag(3): double[3D]", field_type=ps.FieldType.STAGGERED)
9
10
11
12
13
    terms = [
       f[0, 0, 0] - f[-1, 0, 0],
       f[0, 0, 0] - f[0, -1, 0],
       f[0, 0, 0] - f[0, 0, -1],
    ]
14
15
16
    assignments = [ps.Assignment(stag.staggered_access(d), terms[i]) for i, d in enumerate(stag.staggered_stencil)]
    kernel = ps.create_staggered_kernel(assignments, cpu_blocking=(3, 16, 8)).compile()
    reference_kernel = ps.create_staggered_kernel(assignments).compile()
17
18
19
20
21
22
23
24
    print(ps.show_code(kernel.ast))

    f_arr = np.random.rand(80, 33, 19)
    stag_arr = np.zeros((80, 33, 19, 3))
    stag_ref = np.zeros((80, 33, 19, 3))
    kernel(f=f_arr, stag=stag_arr)
    reference_kernel(f=f_arr, stag=stag_ref)
    np.testing.assert_almost_equal(stag_arr, stag_ref)