extended timeloop test case

pystencils_tests/test_timeloop.py

@@ -8,7 +8,7 @@ from pystencils.timeloop import TimeLoop
 
 
 def test_timeloop():
-    dh = create_data_handling(domain_size=(10, 10), periodicity=True)
+    dh = create_data_handling(domain_size=(2, 2), periodicity=True)
 
     pre = dh.add_array('pre_run_field', values_per_cell=1)
     dh.fill("pre_run_field", 0.0, ghost_layers=True)
@@ -16,6 +16,8 @@ def test_timeloop():
     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)
+    single_step = dh.add_array('single_step_field', values_per_cell=1)
+    dh.fill("single_step_field", 0.0, ghost_layers=True)
 
     pre_assignments = Assignment(pre.center, pre.center + 1)
     pre_kernel = create_kernel(pre_assignments).compile()
@@ -23,8 +25,12 @@ def test_timeloop():
     kernel = create_kernel(assignments).compile()
     post_assignments = Assignment(post.center, post.center + 1)
     post_kernel = create_kernel(post_assignments).compile()
+    single_step_assignments = Assignment(single_step.center, single_step.center + 1)
+    single_step_kernel = create_kernel(single_step_assignments).compile()
 
-    timeloop = TimeLoop(steps=1)
+    fixed_steps = 2
+    timeloop = TimeLoop(steps=fixed_steps)
+    assert timeloop.fixed_steps == fixed_steps
 
     def pre_run():
         dh.run_kernel(pre_kernel)
@@ -32,13 +38,20 @@ def test_timeloop():
     def post_run():
         dh.run_kernel(post_kernel)
 
+    def single_step_run():
+        dh.run_kernel(single_step_kernel)
+
     timeloop.add_pre_run_function(pre_run)
     timeloop.add_post_run_function(post_run)
+    timeloop.add_single_step_function(single_step_run)
     timeloop.add_call(kernel, {'field': dh.cpu_arrays["field"]})
 
+    # 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
     timeloop.run(time_steps=5)
     assert np.all(dh.cpu_arrays["pre_run_field"] == 1.0)
-    assert np.all(dh.cpu_arrays["field"] == 5.0)
+    assert np.all(dh.cpu_arrays["field"] == 2.0)
+    assert np.all(dh.cpu_arrays["single_step_field"] == 1.0)
     assert np.all(dh.cpu_arrays["post_run_field"] == 1.0)
 
     seconds = 2