Finish draft for graph_datahandling

pystencils/datahandling/graph_datahandling.py

@@ -15,7 +15,6 @@ import numpy as np
 import pystencils.datahandling
 import pystencils.kernel_wrapper
 from pystencils.field import FieldType
-from pystencils.integer_functions import modulo_ceil
 
 
 class DataTransferKind(str, Enum):
@@ -67,13 +66,29 @@ class KernelCall:
         return "Call " + str(self.kernel.ast.function_name)
 
 
+class TimeloopRun:
+    def __init__(self, timeloop, time_steps):
+        self.timeloop = timeloop
+        self.time_steps = time_steps
+
+    def __str__(self):
+        return (f'Timeloop:'
+                + '\nPre:\n'
+                + '\n   '.join(str(f) for f in self.timeloop._pre_run_functions)
+                + f'\n{self.time_steps} time steps:\n'
+                + '\n   '.join(str(f) for f in self.timeloop._single_step_asts)
+                + '\nPost:\n'
+                + '\n   '.join(str(f) for f in self.timeloop._post_run_functions))
+
+
 class GraphDataHandling(pystencils.datahandling.SerialDataHandling):
 
     """Docstring for GraphDataHandling. """
 
-    class TimeLoop(pystencils.TimeLoop):
+    class TimeLoop(pystencils.timeloop.TimeLoop):
         def __init__(self, parent, *args, **kwargs):
             self.parent = parent
+            self._single_step_asts = []
             super().__init__(*args, **kwargs)
 
         def add_pre_run_function(self, f):
@@ -86,13 +101,16 @@ class GraphDataHandling(pystencils.datahandling.SerialDataHandling):
             self._single_step_functions.append(f)
 
         def add_call(self, functor, argument_list):
+            for argument_dict in argument_list:
+                self._single_step_asts.append((functor, argument_dict) if not hasattr(functor, 'ast') else functor.ast)
+
             if hasattr(functor, 'kernel'):
                 functor = functor.kernel
             if not isinstance(argument_list, list):
                 argument_list = [argument_list]
 
-            for argument_dict in argument_list:
-                self._call_data.append((functor, argument_dict))
+        def run(self, time_steps=1):
+            super().run(time_steps)
 
     def __init__(self, *args, **kwargs):
 
@@ -136,12 +154,15 @@ class GraphDataHandling(pystencils.datahandling.SerialDataHandling):
 
     def run_kernel(self, kernel_function, **kwargs):
         self.call_queue.append(KernelCall(kernel_function, kwargs))
+        super().run_kernel(kernel_function, **kwargs)
         # skip calling super
 
     def to_cpu(self, name):
+        super().to_cpu(name)
         self.call_queue.append(DataTransfer(self._fields[name], DataTransferKind.HOST_TO_DEVICE))
 
     def to_gpu(self, name):
+        super().to_gpu(name)
         if name in self._custom_data_transfer_functions:
             self.call_queue.append('Custom Tranfer Function')
         else:
@@ -151,17 +172,18 @@ class GraphDataHandling(pystencils.datahandling.SerialDataHandling):
         for name in names:
             gpu = target == 'cpu'
             self.call_queue.append(Communication(self._fields[name], stencil, gpu))
+        super().synchronization_function(names, stencil=None, target=None, **_)
 
     def __str__(self):
-        return '\n'.join(str(self.call_queue))
+        return '\n'.join(str(c) for c in self.call_queue)
 
     def create_timeloop(self, *args, **kwargs):
         return self.TimeLoop(self, *args, **kwargs)
 
-    def fill(self, array_name: str, val, value_idx,
+    def fill(self, array_name: str, val, value_idx=None,
              slice_obj=None, ghost_layers=False, inner_ghost_layers=False) -> None:
-        self.call_queue('Fill ' + array_name)
-        super().fill(self, array_name, val, value_idx, slice_obj, ghost_layers, inner_ghost_layers)
+        self.call_queue.append('Fill ' + array_name)
+        super().fill(array_name, val, value_idx, slice_obj, ghost_layers, inner_ghost_layers)
 
     # TODO
     # def reduce_float_sequence(self, sequence, operation, all_reduce=False) -> np.array:

pystencils_tests/test_graph_datahandling.py

@@ -20,7 +20,7 @@ from pystencils.slicing import slice_from_direction
 pytest.importorskip('lbmpy')
 
 
-def create_lid_driven_cavity(domain_size=None, lid_velocity=0.005, lbm_kernel=None, parallel=False,
+def create_lid_driven_cavity(domain_size=None, lid_velocity=0.005, lbm_kernel=None,
                              data_handling=None, **kwargs):
     """Creates a lid driven cavity scenario.
 
@@ -42,7 +42,11 @@ def create_lid_driven_cavity(domain_size=None, lid_velocity=0.005, lbm_kernel=No
                                           periodicity=False,
                                           default_ghost_layers=1,
                                           default_target=target)
-    step = LatticeBoltzmannStep(data_handling=data_handling, lbm_kernel=lbm_kernel, name="ldc", **kwargs)
+    step = LatticeBoltzmannStep(data_handling=data_handling,
+                                lbm_kernel=lbm_kernel,
+                                name="ldc",
+                                timeloop_creation_function=data_handling.create_timeloop,
+                                **kwargs)
 
     my_ubb = UBB(velocity=[lid_velocity, 0, 0][:step.method.dim])
     step.boundary_handling.set_boundary(my_ubb, slice_from_direction('N', step.dim))
@@ -55,7 +59,7 @@ def create_lid_driven_cavity(domain_size=None, lid_velocity=0.005, lbm_kernel=No
 def ldc_setup(**kwargs):
     ldc = create_lid_driven_cavity(relaxation_rate=1.7, **kwargs)
     ldc.run(50)
-    return ldc.density_slice()
+    return ldc
 
 
 def test_graph_datahandling():
@@ -64,5 +68,5 @@ def test_graph_datahandling():
     print("--- LDC 2D test ---")
 
     opt_params = {'target': 'gpu', 'gpu_indexing_params': {'block_size': (8, 4, 2)}}
-    lbm_step: LatticeBoltzmannStep = ldc_setup(domain_size=(10, 15), parallel=False, optimization=opt_params)
+    lbm_step: LatticeBoltzmannStep = ldc_setup(domain_size=(10, 15), optimization=opt_params)
     print(lbm_step._data_handling)