test_timeloop.py 2.5 KB
Newer Older
Markus Holzer's avatar
Markus Holzer committed
1
2
3
4
5
6
7
8
9
10
import time
import numpy as np

from pystencils import Assignment
from pystencils import create_kernel
from pystencils.datahandling import create_data_handling
from pystencils.timeloop import TimeLoop


def test_timeloop():
Markus Holzer's avatar
Markus Holzer committed
11
    dh = create_data_handling(domain_size=(2, 2), periodicity=True)
Markus Holzer's avatar
Markus Holzer committed
12
13
14
15
16
17
18

    pre = dh.add_array('pre_run_field', values_per_cell=1)
    dh.fill("pre_run_field", 0.0, ghost_layers=True)
    f = dh.add_array('field', values_per_cell=1)
    dh.fill("field", 0.0, ghost_layers=True)
    post = dh.add_array('post_run_field', values_per_cell=1)
    dh.fill("post_run_field", 0.0, ghost_layers=True)
Markus Holzer's avatar
Markus Holzer committed
19
20
    single_step = dh.add_array('single_step_field', values_per_cell=1)
    dh.fill("single_step_field", 0.0, ghost_layers=True)
Markus Holzer's avatar
Markus Holzer committed
21
22
23
24
25
26
27

    pre_assignments = Assignment(pre.center, pre.center + 1)
    pre_kernel = create_kernel(pre_assignments).compile()
    assignments = Assignment(f.center, f.center + 1)
    kernel = create_kernel(assignments).compile()
    post_assignments = Assignment(post.center, post.center + 1)
    post_kernel = create_kernel(post_assignments).compile()
Markus Holzer's avatar
Markus Holzer committed
28
29
    single_step_assignments = Assignment(single_step.center, single_step.center + 1)
    single_step_kernel = create_kernel(single_step_assignments).compile()
Markus Holzer's avatar
Markus Holzer committed
30

Markus Holzer's avatar
Markus Holzer committed
31
32
33
    fixed_steps = 2
    timeloop = TimeLoop(steps=fixed_steps)
    assert timeloop.fixed_steps == fixed_steps
Markus Holzer's avatar
Markus Holzer committed
34
35
36
37
38
39
40

    def pre_run():
        dh.run_kernel(pre_kernel)

    def post_run():
        dh.run_kernel(post_kernel)

Markus Holzer's avatar
Markus Holzer committed
41
42
43
    def single_step_run():
        dh.run_kernel(single_step_kernel)

Markus Holzer's avatar
Markus Holzer committed
44
45
    timeloop.add_pre_run_function(pre_run)
    timeloop.add_post_run_function(post_run)
Markus Holzer's avatar
Markus Holzer committed
46
    timeloop.add_single_step_function(single_step_run)
Markus Holzer's avatar
Markus Holzer committed
47
48
    timeloop.add_call(kernel, {'field': dh.cpu_arrays["field"]})

Markus Holzer's avatar
Markus Holzer committed
49
50
    # the timeloop is initialised with 2 steps. This means a single time step consists of two steps.
    # Therefore, we have 2 main iterations and one single step iteration in this configuration
Markus Holzer's avatar
Markus Holzer committed
51
52
    timeloop.run(time_steps=5)
    assert np.all(dh.cpu_arrays["pre_run_field"] == 1.0)
Markus Holzer's avatar
Markus Holzer committed
53
54
    assert np.all(dh.cpu_arrays["field"] == 2.0)
    assert np.all(dh.cpu_arrays["single_step_field"] == 1.0)
Markus Holzer's avatar
Markus Holzer committed
55
56
57
58
59
60
61
    assert np.all(dh.cpu_arrays["post_run_field"] == 1.0)

    seconds = 2
    start = time.perf_counter()
    timeloop.run_time_span(seconds=seconds)
    end = time.perf_counter()

Markus Holzer's avatar
Markus Holzer committed
62
63
64
    # This test case fails often due to time measurements. It is not a good idea to assert here
    # np.testing.assert_almost_equal(seconds, end - start, decimal=2)
    print("timeloop: ", seconds, "  own meassurement: ", end - start)