diff --git a/pystencils/datahandling/graph_datahandling.py b/pystencils/datahandling/graph_datahandling.py
index 41d16a4c3ddea7c2893712275606c6448b4eb5ba..4c9b635b75998ddcab299bfbb005ac3e0ef5af8f 100644
--- a/pystencils/datahandling/graph_datahandling.py
+++ b/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:
diff --git a/pystencils_tests/test_graph_datahandling.py b/pystencils_tests/test_graph_datahandling.py
index d494576baa612613d0e5e7f3890b9934431a994a..441b9ff38f4cb4d28a5baee5302f48fc96b68385 100644
--- a/pystencils_tests/test_graph_datahandling.py
+++ b/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)