Thin Logging & MPI abstraction for data handling

datahandling/datahandling_interface.py

@@ -316,6 +316,20 @@ class DataHandling(ABC):
             result += row_format.format(arr_name, inner_min_max, with_gl_min_max)
         return result
 
+    def log(self, *args, level='INFO'):
+        """Similar to print with additional information (time, rank)."""
+
+    def log_on_root(self, *args, level='INFO'):
+        """Logs only on root process. For serial setups this is equivalent to log"""
+
+    @property
+    def is_root(self):
+        """Returns True for exactly one process in the simulation"""
+
+    @property
+    def world_rank(self):
+        """Number of current process"""
+
 
 class Block:
     """Represents locally stored part of domain.

datahandling/parallel_datahandling.py

@@ -341,3 +341,29 @@ class ParallelDataHandling(DataHandling):
             w = wlb.field.createBinarizationVTKWriter(self.blocks, data_name, mask, name)
             output.addCellDataWriter(w)
         return output
+
+    @staticmethod
+    def log(*args, level='INFO'):
+        level = level.upper()
+        message = " ".join(str(e) for e in args)
+        ParallelDataHandling._log_map[level](message)
+
+    def log_on_root(self, *args, level='INFO'):
+        if self.is_root:
+            ParallelDataHandling.log(*args, level=level)
+
+    @property
+    def is_root(self):
+        return wlb.mpi.worldRank() == 0
+
+    @property
+    def world_rank(self):
+        return wlb.mpi.worldRank()
+
+    _log_map = {
+        'DEVEL': wlb.log_devel,
+        'RESULT': wlb.log_result,
+        'INFO': wlb.log_info,
+        'WARNING': wlb.log_warning,
+        'PROGRESS': wlb.log_progress,
+    }

datahandling/serial_datahandling.py

 import itertools
 from typing import Sequence, Union
 import numpy as np
+import time
 from pystencils import Field
 from pystencils.datahandling.datahandling_interface import DataHandling
 from pystencils.field import layout_string_to_tuple, spatial_layout_string_to_tuple, create_numpy_array_with_layout
@@ -48,6 +49,7 @@ class SerialDataHandling(DataHandling):
         self._periodicity = periodicity
         self._field_information = {}
         self.default_target = default_target
+        self._start_time = time.perf_counter()
 
     @property
     def dim(self):
@@ -356,3 +358,23 @@ class SerialDataHandling(DataHandling):
         gl_to_remove = actual_ghost_layers - ghost_layers
         ind_dims = 1 if self._field_information[name]['values_per_cell'] > 1 else 0
         return remove_ghost_layers(self.cpu_arrays[name], ind_dims, gl_to_remove)
+
+    def log(self, *args, level='INFO'):
+        level = level.upper()
+        message = " ".join(str(e) for e in args)
+
+        time_running = time.perf_counter() - self._start_time
+        spacing = 7 - len(str(int(time_running)))
+        message = "[{: <8}]{}({:.3f} sec) {} ".format(level, spacing * '-', time_running, message)
+        print(message, flush=True)
+
+    def log_on_root(self, *args, level='INFO'):
+        self.log(*args, level=level)
+
+    @property
+    def is_root(self):
+        return True
+
+    @property
+    def world_rank(self):
+        return 0

timeloop.py

@@ -63,16 +63,17 @@ class TimeLoop:
     def benchmark_run(self, time_steps=0, init_time_steps=0):
         init_time_steps_rounded = modulo_ceil(init_time_steps, self._fixed_steps)
         time_steps_rounded = modulo_ceil(time_steps, self._fixed_steps)
+        call_data = self._call_data
 
         self.pre_run()
         for i in range(init_time_steps_rounded // self._fixed_steps):
-            for func, kwargs in self._call_data:
+            for func, kwargs in call_data:
                 func(**kwargs)
         self.time_steps_run += init_time_steps_rounded
 
         start = time.perf_counter()
         for i in range(time_steps_rounded // self._fixed_steps):
-            for func, kwargs in self._call_data:
+            for func, kwargs in call_data:
                 func(**kwargs)
         end = time.perf_counter()
         self.time_steps_run += time_steps_rounded