Skip to content

GitLab

Commit 68acb242 authored by Michael Kuron's avatar Michael Kuron
Browse files

introduce field.FieldType.STAGGERED

parent 23de3441
Hide whitespace changes
Inline Side-by-side
pystencils/datahandling/datahandling_interface.py
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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
View file @ 68acb242
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment