introduce field.FieldType.STAGGERED

pystencils/datahandling/datahandling_interface.py

@@ -3,7 +3,7 @@ from typing import Callable, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
-from pystencils.field import Field
+from pystencils.field import Field, FieldType
 
 
 class DataHandling(ABC):
@@ -36,7 +36,7 @@ class DataHandling(ABC):
     @abstractmethod
     def add_array(self, name: str, values_per_cell, dtype=np.float64,
                   latex_name: Optional[str] = None, ghost_layers: Optional[int] = None, layout: Optional[str] = None,
-                  cpu: bool = True, gpu: Optional[bool] = None, alignment=False, staggered=False) -> Field:
+                  cpu: bool = True, gpu: Optional[bool] = None, alignment=False, field_type=FieldType.GENERIC) -> Field:
         """Adds a (possibly distributed) array to the handling that can be accessed using the given name.
 
         For each array a symbolic field is available via the 'fields' dictionary

pystencils/datahandling/parallel_datahandling.py

@@ -7,7 +7,7 @@ import waLBerla as wlb
 
 from pystencils.datahandling.blockiteration import block_iteration, sliced_block_iteration
 from pystencils.datahandling.datahandling_interface import DataHandling
-from pystencils.field import Field
+from pystencils.field import Field, FieldType
 from pystencils.kernelparameters import FieldPointerSymbol
 from pystencils.utils import DotDict
 
@@ -90,7 +90,7 @@ class ParallelDataHandling(DataHandling):
         self._custom_data_names.append(name)
 
     def add_array(self, name, values_per_cell=1, dtype=np.float64, latex_name=None, ghost_layers=None,
-                  layout=None, cpu=True, gpu=None, alignment=False, staggered=False):
+                  layout=None, cpu=True, gpu=None, alignment=False, field_type=FieldType.GENERIC):
         if ghost_layers is None:
             ghost_layers = self.default_ghost_layers
         if gpu is None:
@@ -141,7 +141,7 @@ class ParallelDataHandling(DataHandling):
 
         self.fields[name] = Field.create_generic(name, self.dim, dtype, index_dimensions, layout,
                                                  index_shape=(values_per_cell,) if index_dimensions > 0 else None,
-                                                 staggered=staggered)
+                                                 field_type=field_type)
         self.fields[name].latex_name = latex_name
         self._field_name_to_cpu_data_name[name] = name
         if gpu:

pystencils/datahandling/serial_datahandling.py

@@ -7,7 +7,7 @@ import numpy as np
 from pystencils.datahandling.blockiteration import SerialBlock
 from pystencils.datahandling.datahandling_interface import DataHandling
 from pystencils.field import (
-    Field, create_numpy_array_with_layout, layout_string_to_tuple, spatial_layout_string_to_tuple)
+    Field, FieldType, create_numpy_array_with_layout, layout_string_to_tuple, spatial_layout_string_to_tuple)
 from pystencils.slicing import normalize_slice, remove_ghost_layers
 from pystencils.utils import DotDict
 
@@ -76,7 +76,7 @@ class SerialDataHandling(DataHandling):
         return self._field_information[name]['values_per_cell']
 
     def add_array(self, name, values_per_cell=1, dtype=np.float64, latex_name=None, ghost_layers=None, layout=None,
-                  cpu=True, gpu=None, alignment=False, staggered=False):
+                  cpu=True, gpu=None, alignment=False, field_type=FieldType.GENERIC):
         if ghost_layers is None:
             ghost_layers = self.default_ghost_layers
         if layout is None:
@@ -129,7 +129,7 @@ class SerialDataHandling(DataHandling):
 
         assert all(f.name != name for f in self.fields.values()), "Symbolic field with this name already exists"
         self.fields[name] = Field.create_from_numpy_array(name, cpu_arr, index_dimensions=index_dimensions,
-                                                          staggered=staggered)
+                                                          field_type=field_type)
         self.fields[name].latex_name = latex_name
         return self.fields[name]
 

pystencils/field.py

@@ -21,7 +21,47 @@ from pystencils.sympyextensions import is_integer_sequence
 __all__ = ['Field', 'fields', 'FieldType', 'AbstractField']
 
 
-def fields(description=None, index_dimensions=0, layout=None, staggered=False, **kwargs):
+class FieldType(Enum):
+    # generic fields
+    GENERIC = 0
+    # index fields are currently only used for boundary handling
+    # the coordinates are not the loop counters in that case, but are read from this index field
+    INDEXED = 1
+    # communication buffer, used for (un)packing data in communication.
+    BUFFER = 2
+    # unsafe fields may be accessed in an absolute fashion - the index depends on the data
+    # and thus may lead to out-of-bounds accesses
+    CUSTOM = 3
+    # staggered field
+    STAGGERED = 4
+
+    @staticmethod
+    def is_generic(field):
+        assert isinstance(field, Field)
+        return field.field_type == FieldType.GENERIC
+
+    @staticmethod
+    def is_indexed(field):
+        assert isinstance(field, Field)
+        return field.field_type == FieldType.INDEXED
+
+    @staticmethod
+    def is_buffer(field):
+        assert isinstance(field, Field)
+        return field.field_type == FieldType.BUFFER
+
+    @staticmethod
+    def is_custom(field):
+        assert isinstance(field, Field)
+        return field.field_type == FieldType.CUSTOM
+
+    @staticmethod
+    def is_staggered(field):
+        assert isinstance(field, Field)
+        return field.field_type == FieldType.STAGGERED
+
+
+def fields(description=None, index_dimensions=0, layout=None, field_type=FieldType.GENERIC, **kwargs):
     """Creates pystencils fields from a string description.
 
     Examples:
@@ -62,16 +102,16 @@ def fields(description=None, index_dimensions=0, layout=None, staggered=False, *
                 idx_shape_of_arr = () if not len(idx_shape) else arr.shape[-len(idx_shape):]
                 assert idx_shape_of_arr == idx_shape
                 f = Field.create_from_numpy_array(field_name, kwargs[field_name], index_dimensions=len(idx_shape),
-                                                  staggered=staggered)
+                                                  field_type=field_type)
             elif isinstance(shape, tuple):
                 f = Field.create_fixed_size(field_name, shape + idx_shape, dtype=dtype,
-                                            index_dimensions=len(idx_shape), layout=layout, staggered=staggered)
+                                            index_dimensions=len(idx_shape), layout=layout, field_type=field_type)
             elif isinstance(shape, int):
                 f = Field.create_generic(field_name, spatial_dimensions=shape, dtype=dtype,
-                                         index_shape=idx_shape, layout=layout, staggered=staggered)
+                                         index_shape=idx_shape, layout=layout, field_type=field_type)
             elif shape is None:
                 f = Field.create_generic(field_name, spatial_dimensions=2, dtype=dtype,
-                                         index_shape=idx_shape, layout=layout, staggered=staggered)
+                                         index_shape=idx_shape, layout=layout, field_type=field_type)
             else:
                 assert False
             result.append(f)
@@ -79,7 +119,7 @@ def fields(description=None, index_dimensions=0, layout=None, staggered=False, *
         assert layout is None, "Layout can not be specified when creating Field from numpy array"
         for field_name, arr in kwargs.items():
             result.append(Field.create_from_numpy_array(field_name, arr, index_dimensions=index_dimensions,
-                                                        staggered=staggered))
+                                                        field_type=field_type))
 
     if len(result) == 0:
         return None
@@ -89,39 +129,6 @@ def fields(description=None, index_dimensions=0, layout=None, staggered=False, *
         return result
 
 
-class FieldType(Enum):
-    # generic fields
-    GENERIC = 0
-    # index fields are currently only used for boundary handling
-    # the coordinates are not the loop counters in that case, but are read from this index field
-    INDEXED = 1
-    # communication buffer, used for (un)packing data in communication.
-    BUFFER = 2
-    # unsafe fields may be accessed in an absolute fashion - the index depends on the data
-    # and thus may lead to out-of-bounds accesses
-    CUSTOM = 3
-
-    @staticmethod
-    def is_generic(field):
-        assert isinstance(field, Field)
-        return field.field_type == FieldType.GENERIC
-
-    @staticmethod
-    def is_indexed(field):
-        assert isinstance(field, Field)
-        return field.field_type == FieldType.INDEXED
-
-    @staticmethod
-    def is_buffer(field):
-        assert isinstance(field, Field)
-        return field.field_type == FieldType.BUFFER
-
-    @staticmethod
-    def is_custom(field):
-        assert isinstance(field, Field)
-        return field.field_type == FieldType.CUSTOM
-
-
 class AbstractField:
 
     class AbstractAccess:
@@ -182,7 +189,7 @@ class Field(AbstractField):
 
     @staticmethod
     def create_generic(field_name, spatial_dimensions, dtype=np.float64, index_dimensions=0, layout='numpy',
-                       index_shape=None, field_type=FieldType.GENERIC, staggered=False) -> 'Field':
+                       index_shape=None, field_type=FieldType.GENERIC) -> 'Field':
         """
         Creates a generic field where the field size is not fixed i.e. can be called with arrays of different sizes
 
@@ -197,9 +204,9 @@ class Field(AbstractField):
             index_shape: optional shape of the index dimensions i.e. maximum values allowed for each index dimension,
                         has to be a list or tuple
             field_type: besides the normal GENERIC fields, there are INDEXED fields that store indices of the domain
-                        that should be iterated over, and BUFFER fields that are used to generate
-                        communication packing/unpacking kernels
-            staggered: enables staggered access (with half-integer offsets) and corresponding printing
+                        that should be iterated over, BUFFER fields that are used to generate communication
+                        packing/unpacking kernels, and STAGGERED fields, which store values half-way to the next
+                        cell
         """
         if index_shape is not None:
             assert index_dimensions == 0 or index_dimensions == len(index_shape)
@@ -221,14 +228,14 @@ class Field(AbstractField):
                 raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
             shape += (1,)
             strides += (1,)
-        if staggered and index_dimensions == 0:
+        if field_type == FieldType.STAGGERED and index_dimensions == 0:
             raise ValueError("A staggered field needs at least one index dimension")
 
-        return Field(field_name, field_type, dtype, layout, shape, strides, staggered)
+        return Field(field_name, field_type, dtype, layout, shape, strides)
 
     @staticmethod
     def create_from_numpy_array(field_name: str, array: np.ndarray, index_dimensions: int = 0,
-                                staggered=False) -> 'Field':
+                                field_type=FieldType.GENERIC) -> 'Field':
         """Creates a field based on the layout, data type, and shape of a given numpy array.
 
         Kernels created for these kind of fields can only be called with arrays of the same layout, shape and type.
@@ -237,7 +244,7 @@ class Field(AbstractField):
             field_name: symbolic name for the field
             array: numpy array
             index_dimensions: see documentation of Field
-            staggered: enables staggered access (with half-integer offsets) and corresponding printing
+            field_type: kind of field
         """
         spatial_dimensions = len(array.shape) - index_dimensions
         if spatial_dimensions < 1:
@@ -256,15 +263,15 @@ class Field(AbstractField):
                 raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
             shape += (1,)
             strides += (1,)
-        if staggered and index_dimensions == 0:
+        if field_type == FieldType.STAGGERED and index_dimensions == 0:
             raise ValueError("A staggered field needs at least one index dimension")
 
-        return Field(field_name, FieldType.GENERIC, array.dtype, spatial_layout, shape, strides, staggered)
+        return Field(field_name, field_type, array.dtype, spatial_layout, shape, strides)
 
     @staticmethod
     def create_fixed_size(field_name: str, shape: Tuple[int, ...], index_dimensions: int = 0,
                           dtype=np.float64, layout: str = 'numpy', strides: Optional[Sequence[int]] = None,
-                          staggered=False) -> 'Field':
+                          field_type=FieldType.GENERIC) -> 'Field':
         """
         Creates a field with fixed sizes i.e. can be called only with arrays of the same size and layout
 
@@ -275,7 +282,7 @@ class Field(AbstractField):
             dtype: numpy data type of the array the kernel is called with later
             layout: full layout of array, not only spatial dimensions
             strides: strides in bytes or None to automatically compute them from shape (assuming no padding)
-            staggered: enables staggered access (with half-integer offsets) and corresponding printing
+            field_type: kind of field
         """
         spatial_dimensions = len(shape) - index_dimensions
         assert spatial_dimensions >= 1
@@ -296,15 +303,15 @@ class Field(AbstractField):
                 raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
             shape += (1,)
             strides += (1,)
-        if staggered and index_dimensions == 0:
+        if field_type == FieldType.STAGGERED and index_dimensions == 0:
             raise ValueError("A staggered field needs at least one index dimension")
 
         spatial_layout = list(layout)
         for i in range(spatial_dimensions, len(layout)):
             spatial_layout.remove(i)
-        return Field(field_name, FieldType.GENERIC, dtype, tuple(spatial_layout), shape, strides, staggered)
+        return Field(field_name, field_type, dtype, tuple(spatial_layout), shape, strides)
 
-    def __init__(self, field_name, field_type, dtype, layout, shape, strides, staggered):
+    def __init__(self, field_name, field_type, dtype, layout, shape, strides):
         """Do not use directly. Use static create* methods"""
         self._field_name = field_name
         assert isinstance(field_type, FieldType)
@@ -319,7 +326,6 @@ class Field(AbstractField):
             0 for _ in range(self.spatial_dimensions)
         ))  # type: tuple[float,sp.Symbol]
         self.coordinate_transform = sp.eye(self.spatial_dimensions)
-        self.is_staggered = staggered
 
     def new_field_with_different_name(self, new_name):
         if self.has_fixed_shape:
@@ -456,7 +462,7 @@ class Field(AbstractField):
         If the field stores more than one value per staggered point (e.g. a vector or a tensor), the index (integer or
         tuple of integers) refers to which of these values to access.
         """
-        assert self.is_staggered
+        assert FieldType.is_staggered(self)
 
         if type(offset) is np.ndarray:
             offset = tuple(offset)
@@ -611,7 +617,6 @@ class Field(AbstractField):
                     obj._indirect_addressing_fields.update(a.field for a in e.atoms(Field.Access))
 
             obj._is_absolute_access = is_absolute_access
-            obj.is_staggered = field.is_staggered
             return obj
 
         def __getnewargs__(self):
@@ -744,7 +749,7 @@ class Field(AbstractField):
         def _latex(self, _):
             n = self._field.latex_name if self._field.latex_name else self._field.name
             offset_str = ",".join([sp.latex(o) for o in self.offsets])
-            if self.is_staggered:
+            if FieldType.is_staggered(self._field):
                 offset_str = ",".join([sp.latex(o - sp.Rational(int(i == self.index[0]), 2))
                                        for i, o in enumerate(self.offsets)])
             if self.is_absolute_access:
@@ -752,7 +757,7 @@ class Field(AbstractField):
             elif self.field.spatial_dimensions > 1:
                 offset_str = "({})".format(offset_str)
 
-            if self.is_staggered:
+            if FieldType.is_staggered(self._field):
                 if self.index and self.field.index_dimensions > 1:
                     return "{{%s}_{%s}^{%s}}" % (n, offset_str, self.index[1:]
                                                  if len(self.index) > 2 else self.index[1])
@@ -767,13 +772,13 @@ class Field(AbstractField):
         def __str__(self):
             n = self._field.latex_name if self._field.latex_name else self._field.name
             offset_str = ",".join([sp.latex(o) for o in self.offsets])
-            if self.is_staggered:
+            if FieldType.is_staggered(self._field):
                 offset_str = ",".join([sp.latex(o - sp.Rational(int(i == self.index[0]), 2))
                                        for i, o in enumerate(self.offsets)])
             if self.is_absolute_access:
                 offset_str = "[abs]{}".format(offset_str)
 
-            if self.is_staggered:
+            if FieldType.is_staggered(self._field):
                 if self.index and self.field.index_dimensions > 1:
                     return "%s[%s](%s)" % (n, offset_str, self.index[1:] if len(self.index) > 2 else self.index[1])
                 else:

pystencils_tests/test_blocking_staggered.py

@@ -5,7 +5,7 @@ import pystencils as ps
 
 def test_blocking_staggered():
     f = ps.fields("f: double[3D]")
-    stag = ps.fields("stag(3): double[3D]", staggered=True)
+    stag = ps.fields("stag(3): double[3D]", field_type=ps.FieldType.STAGGERED)
     terms = [
        f[0, 0, 0] - f[-1, 0, 0],
        f[0, 0, 0] - f[0, -1, 0],

pystencils_tests/test_field.py

@@ -131,7 +131,7 @@ def test_itemsize():
 
 def test_staggered():
 
-    j1, j2, j3 = ps.fields('j1(2), j2(2,2), j3(2,2,2) : double[2D]', staggered=True)
+    j1, j2, j3 = ps.fields('j1(2), j2(2,2), j3(2,2,2) : double[2D]', field_type=FieldType.STAGGERED)
 
     assert j1[0, 1](1) == j1.staggered_access((0, sp.Rational(1, 2)))
     assert j1[0, 1](1) == j1.staggered_access("N")

pystencils_tests/test_loop_cutting.py

@@ -3,7 +3,7 @@ import sympy as sp
 
 import pystencils as ps
 import pystencils.astnodes as ast
-from pystencils.field import Field
+from pystencils.field import Field, FieldType
 from pystencils.astnodes import Conditional, LoopOverCoordinate, SympyAssignment
 from pystencils.cpu import create_kernel, make_python_function
 from pystencils.kernelcreation import create_staggered_kernel
@@ -34,9 +34,9 @@ def test_staggered_iteration():
     s_arr_ref = s_arr.copy()
 
     fields_fixed = (Field.create_from_numpy_array('f', f_arr),
-                    Field.create_from_numpy_array('s', s_arr, index_dimensions=1, staggered=True))
+                    Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED))
     fields_var = (Field.create_generic('f', 2),
-                  Field.create_generic('s', 2, index_dimensions=1, staggered=True))
+                  Field.create_generic('s', 2, index_dimensions=1, field_type=FieldType.STAGGERED))
 
     for f, s in [fields_var, fields_fixed]:
         # --- Manual
@@ -70,7 +70,7 @@ def test_staggered_iteration_manual():
     s_arr_ref = s_arr.copy()
 
     f = Field.create_from_numpy_array('f', f_arr)
-    s = Field.create_from_numpy_array('s', s_arr, index_dimensions=1, staggered=True)
+    s = Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED)
 
     eqs = []
 
@@ -108,7 +108,7 @@ def test_staggered_iteration_manual():
 def test_staggered_gpu():
     dim = 2
     f = ps.fields("f: double[{dim}D]".format(dim=dim))
-    s = ps.fields("s({dim}): double[{dim}D]".format(dim=dim), staggered=True)
+    s = ps.fields("s({dim}): double[{dim}D]".format(dim=dim), field_type=FieldType.STAGGERED)
     expressions = [(f[0, 0] + f[-1, 0]) / 2,
                    (f[0, 0] + f[0, -1]) / 2]
     kernel_ast = ps.create_staggered_kernel(s, expressions, target='gpu', gpu_exclusive_conditions=True)