diff --git a/pystencils/datahandling/datahandling_interface.py b/pystencils/datahandling/datahandling_interface.py index ce153ff7a95e6a3e2875aa5b02566874c3238662..ba960edc17b2e3409dc465c257e23d9e9394dc5c 100644 --- a/pystencils/datahandling/datahandling_interface.py +++ b/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) -> 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 diff --git a/pystencils/datahandling/parallel_datahandling.py b/pystencils/datahandling/parallel_datahandling.py index 4e8d965d4335725e5e0a00b842bbffeab9497924..54f26806be318f6ef91a5ca11a9888a59524fb0c 100644 --- a/pystencils/datahandling/parallel_datahandling.py +++ b/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): + 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: @@ -140,7 +140,8 @@ class ParallelDataHandling(DataHandling): assert all(f.name != name for f in self.fields.values()), "Symbolic field with this name already exists" self.fields[name] = Field.create_generic(name, self.dim, dtype, index_dimensions, layout, - index_shape=(values_per_cell,) if index_dimensions > 0 else None) + index_shape=(values_per_cell,) if index_dimensions > 0 else None, + field_type=field_type) self.fields[name].latex_name = latex_name self._field_name_to_cpu_data_name[name] = name if gpu: diff --git a/pystencils/datahandling/serial_datahandling.py b/pystencils/datahandling/serial_datahandling.py index 8c79a0931f06fad7c7ab58143cf7058a3f52dbac..f8b0a4a1d8b56e8533eb99c6c7f915c00f307419 100644 --- a/pystencils/datahandling/serial_datahandling.py +++ b/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): + 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: @@ -128,7 +128,8 @@ class SerialDataHandling(DataHandling): self.gpu_arrays[name] = gpuarray.to_gpu(cpu_arr) 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) + self.fields[name] = Field.create_from_numpy_array(name, cpu_arr, index_dimensions=index_dimensions, + field_type=field_type) self.fields[name].latex_name = latex_name return self.fields[name] diff --git a/pystencils/field.py b/pystencils/field.py index aa0d1d36318c334a2bda3eb6ba197a4213a1c374..0c196f4ba61858f4aa03e8cde0260601d1e68c0d 100644 --- a/pystencils/field.py +++ b/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, **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: @@ -61,23 +101,25 @@ def fields(description=None, index_dimensions=0, layout=None, **kwargs): arr = kwargs[field_name] 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)) + f = Field.create_from_numpy_array(field_name, kwargs[field_name], index_dimensions=len(idx_shape), + 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) + 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) + 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) + index_shape=idx_shape, layout=layout, field_type=field_type) else: assert False result.append(f) else: 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)) + result.append(Field.create_from_numpy_array(field_name, arr, index_dimensions=index_dimensions, + field_type=field_type)) if len(result) == 0: return None @@ -87,39 +129,6 @@ def fields(description=None, index_dimensions=0, layout=None, **kwargs): 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: @@ -158,6 +167,14 @@ class Field(AbstractField): First specify the spatial offsets in [], then in case index_dimension>0 the indices in () e.g. ``f[-1,0,0](7)`` + Staggered Fields: + Staggered fields are used to store a value on a second grid shifted by half a cell with respect to the usual + grid. + + The first index dimension is used to specify the position on the staggered grid (e.g. 0 means half-way to the + eastern neighbor, 1 is half-way to the northern neighbor, etc.), while additional indices can be used to store + multiple values at each position. + Example using no index dimensions: >>> a = np.zeros([10, 10]) >>> f = Field.create_from_numpy_array("f", a, index_dimensions=0) @@ -187,8 +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 + 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) @@ -210,11 +228,14 @@ class Field(AbstractField): raise ValueError("Structured arrays/fields are not allowed to have an index dimension") shape += (1,) strides += (1,) + 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) @staticmethod - def create_from_numpy_array(field_name: str, array: np.ndarray, index_dimensions: int = 0) -> 'Field': + def create_from_numpy_array(field_name: str, array: np.ndarray, index_dimensions: int = 0, + 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. @@ -223,6 +244,7 @@ class Field(AbstractField): field_name: symbolic name for the field array: numpy array index_dimensions: see documentation of Field + field_type: kind of field """ spatial_dimensions = len(array.shape) - index_dimensions if spatial_dimensions < 1: @@ -241,12 +263,15 @@ class Field(AbstractField): raise ValueError("Structured arrays/fields are not allowed to have an index dimension") shape += (1,) strides += (1,) + 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) + 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) -> 'Field': + dtype=np.float64, layout: str = 'numpy', strides: Optional[Sequence[int]] = None, + 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 @@ -257,6 +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) + field_type: kind of field """ spatial_dimensions = len(shape) - index_dimensions assert spatial_dimensions >= 1 @@ -277,11 +303,13 @@ class Field(AbstractField): raise ValueError("Structured arrays/fields are not allowed to have an index dimension") shape += (1,) strides += (1,) + 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) + return Field(field_name, field_type, dtype, tuple(spatial_layout), shape, strides) def __init__(self, field_name, field_type, dtype, layout, shape, strides): """Do not use directly. Use static create* methods""" @@ -298,6 +326,8 @@ class Field(AbstractField): 0 for _ in range(self.spatial_dimensions) )) # type: tuple[float,sp.Symbol] self.coordinate_transform = sp.eye(self.spatial_dimensions) + if field_type == FieldType.STAGGERED: + assert self.staggered_stencil def new_field_with_different_name(self, new_name): if self.has_fixed_shape: @@ -427,6 +457,73 @@ class Field(AbstractField): address_mode, allow_textures=allow_textures).at(offset) + def staggered_access(self, offset, index=None): + """If this field is a staggered field, it can be accessed using half-integer offsets. + For example, an offset of ``(0, sp.Rational(1,2))`` or ``"E"`` corresponds to the staggered point to the east + of the cell center, i.e. half-way to the eastern-next cell. + 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 FieldType.is_staggered(self) + + if type(offset) is np.ndarray: + offset = tuple(offset) + if type(offset) is str: + offset = tuple(direction_string_to_offset(offset, self.spatial_dimensions)) + offset = tuple([o * sp.Rational(1, 2) for o in offset]) + if type(offset) is not tuple: + offset = (offset,) + if len(offset) != self.spatial_dimensions: + raise ValueError("Wrong number of spatial indices: " + "Got %d, expected %d" % (len(offset), self.spatial_dimensions)) + + offset = list(offset) + neighbor = [0] * len(offset) + for i, o in enumerate(offset): + if (o + sp.Rational(1, 2)).is_Integer: + offset[i] += sp.Rational(1, 2) + neighbor[i] = 1 + neighbor = offset_to_direction_string(neighbor) + idx = self.staggered_stencil.index(neighbor) + offset = tuple(offset) + + if self.index_dimensions == 1: # this field stores a scalar value at each staggered position + if index is not None: + raise ValueError("Cannot specify an index for a scalar staggered field") + return Field.Access(self, offset, (idx,)) + else: # this field stores a vector or tensor at each staggered position + if index is None: + raise ValueError("Wrong number of indices: " + "Got %d, expected %d" % (0, self.index_dimensions - 1)) + if type(index) is np.ndarray: + index = tuple(index) + if type(index) is not tuple: + index = (index,) + if self.index_dimensions != len(index) + 1: + raise ValueError("Wrong number of indices: " + "Got %d, expected %d" % (len(index), self.index_dimensions - 1)) + + return Field.Access(self, offset, (idx, *index)) + + @property + def staggered_stencil(self): + assert FieldType.is_staggered(self) + stencils = { + 2: { + 2: ["E", "N"], # D2Q5 + 4: ["E", "N", "NE", "SE"] # D2Q9 + }, + 3: { + 3: ["E", "N", "T"], # D3Q7 + 7: ["E", "N", "T", "TNE", "BNE", "TSE", "BSE "], # D3Q15 + 9: ["E", "N", "T", "NE", "SE", "TE", "BE", "TN", "BN"], # D3Q19 + 13: ["E", "N", "T", "NE", "SE", "TE", "BE", "TN", "BN", "TNE", "BNE", "TSE", "BSE"] # D3Q27 + } + } + if not self.index_shape[0] in stencils[self.spatial_dimensions]: + raise ValueError("No known stencil has {} staggered points".format(self.index_shape[0])) + return stencils[self.spatial_dimensions][self.index_shape[0]] + def __call__(self, *args, **kwargs): center = tuple([0] * self.spatial_dimensions) return Field.Access(self, center)(*args, **kwargs) @@ -673,28 +770,54 @@ class Field(AbstractField): super_class_contents = super(Field.Access, self)._hashable_content() return (super_class_contents, self._field.hashable_contents(), *self._index, *self._offsets) + def _staggered_offset(self, offsets, index): + assert FieldType.is_staggered(self._field) + neighbor = self._field.staggered_stencil[index] + neighbor = direction_string_to_offset(neighbor, self._field.spatial_dimensions) + return [(o - sp.Rational(neighbor[i], 2)) for i, o in enumerate(offsets)] + 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 FieldType.is_staggered(self._field): + offset_str = ",".join([sp.latex(self._staggered_offset(self.offsets, self.index[0])[i]) + for i in range(len(self.offsets))]) if self.is_absolute_access: offset_str = "\\mathbf{}".format(offset_str) elif self.field.spatial_dimensions > 1: offset_str = "({})".format(offset_str) - if self.index and self.index != (0,): - return "{{%s}_{%s}^{%s}}" % (n, offset_str, self.index if len(self.index) > 1 else self.index[0]) + 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: + return "{{%s}_{%s}}" % (n, offset_str) else: - return "{{%s}_{%s}}" % (n, offset_str) + if self.index and self.field.index_dimensions > 0: + return "{{%s}_{%s}^{%s}}" % (n, offset_str, self.index if len(self.index) > 1 else self.index[0]) + else: + return "{{%s}_{%s}}" % (n, offset_str) 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 FieldType.is_staggered(self._field): + offset_str = ",".join([sp.latex(self._staggered_offset(self.offsets, self.index[0])[i]) + for i in range(len(self.offsets))]) if self.is_absolute_access: offset_str = "[abs]{}".format(offset_str) - if self.index and self.index != (0,): - return "%s[%s](%s)" % (n, offset_str, self.index if len(self.index) > 1 else self.index[0]) + + 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: + return "%s[%s]" % (n, offset_str) else: - return "%s[%s]" % (n, offset_str) + if self.index and self.field.index_dimensions > 0: + return "%s[%s](%s)" % (n, offset_str, self.index if len(self.index) > 1 else self.index[0]) + else: + return "%s[%s]" % (n, offset_str) def get_layout_from_strides(strides: Sequence[int], index_dimension_ids: Optional[List[int]] = None): diff --git a/pystencils_tests/test_blocking_staggered.py b/pystencils_tests/test_blocking_staggered.py index 207d05a280f790c5c0b74a50e7ecfb252bd04276..76ec8abf0e3b76d30415571bb43e975f317e9a3f 100644 --- a/pystencils_tests/test_blocking_staggered.py +++ b/pystencils_tests/test_blocking_staggered.py @@ -4,7 +4,8 @@ import pystencils as ps def test_blocking_staggered(): - f, stag = ps.fields("f, stag(3): double[3D]") + f = ps.fields("f: double[3D]") + 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], diff --git a/pystencils_tests/test_field.py b/pystencils_tests/test_field.py index 49bd4486474b91caf90c7e405f1a2a2ac7d7e33d..227b332cd1dbf1f68bab43d1b7042b89cc44b8ad 100644 --- a/pystencils_tests/test_field.py +++ b/pystencils_tests/test_field.py @@ -127,3 +127,25 @@ def test_itemsize(): assert x.itemsize == 4 assert y.itemsize == 8 assert i.itemsize == 2 + + +def test_staggered(): + + # D2Q5 + 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") + assert j1[0, 0](1) == j1.staggered_access("S") + + assert j2[0, 1](1, 1) == j2.staggered_access((0, sp.Rational(1, 2)), 1) + assert j2[0, 1](1, 1) == j2.staggered_access("N", 1) + + assert j3[0, 1](1, 1, 1) == j3.staggered_access((0, sp.Rational(1, 2)), (1, 1)) + assert j3[0, 1](1, 1, 1) == j3.staggered_access("N", (1, 1)) + + # D2Q9 + k = ps.fields('k(4) : double[2D]', field_type=FieldType.STAGGERED) + + assert k[1, 1](2) == k.staggered_access("NE") + assert k[0, 0](2) == k.staggered_access("SW") diff --git a/pystencils_tests/test_loop_cutting.py b/pystencils_tests/test_loop_cutting.py index 58126928daea4e7e6df5e0e71b3e57ac1491cff2..999e7b52a8b40111243c09aca1aa3fc1549a0cc2 100644 --- a/pystencils_tests/test_loop_cutting.py +++ b/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)) + 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)) + Field.create_generic('s', 2, index_dimensions=1, index_shape=(dim,), 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) + s = Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED) eqs = [] @@ -107,7 +107,8 @@ def test_staggered_iteration_manual(): def test_staggered_gpu(): dim = 2 - f, s = ps.fields("f, s({dim}): double[{dim}D]".format(dim=dim)) + f = ps.fields("f: double[{dim}D]".format(dim=dim)) + 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)