Skip to content
Snippets Groups Projects
field.py 30.1 KiB
Newer Older
from itertools import chain
Martin Bauer's avatar
Martin Bauer committed
from typing import Tuple, Sequence, Optional, List
import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
from sympy.tensor import IndexedBase
from pystencils.alignedarray import aligned_empty
Martin Bauer's avatar
Martin Bauer committed
from pystencils.data_types import TypedSymbol, create_type, create_composite_type_from_string, StructType
Martin Bauer's avatar
Martin Bauer committed
from pystencils.sympyextensions import is_integer_sequence
def fields(description=None, index_dimensions=0, layout=None, **kwargs):
    """Creates pystencils fields from a string description.

    Examples:
        Create a 2D scalar and vector field
        >>> s, v = fields("s, v(2): double[2D]")
        >>> assert s.spatial_dimensions == 2 and s.index_dimensions == 0
        >>> assert v.spatial_dimensions == 2 and v.index_dimensions == 1 and v.index_shape == (2,)

        Create an integer field of shape (10, 20)
        >>> f = fields("f : int32[10, 20]")
        >>> f.has_fixed_shape, f.shape
        (True, (10, 20))

        Numpy arrays can be used as template for shape and data type of field
        >>> arr_s, arr_v = np.zeros([20, 20]), np.zeros([20, 20, 2])
        >>> s, v = fields("s, v(2)", s=arr_s, v=arr_v)
        >>> assert s.index_dimensions == 0 and v.index_shape == (2,) and s.dtype.numpy_dtype == arr_s.dtype

        Format string can be left out, field names are taken from keyword arguments.
        >>> fields(f1=arr_s, f2=arr_s)
        [f1, f2]

        The keyword names 'index_dimension' and 'layout' have special meaning and thus can not be used to pass
        numpy arrays:
        >>> f = fields(f=arr_v, index_dimensions=1)
        >>> assert f.index_dimensions == 1

        >>> f = fields("pdfs(19) : float32[3D]", layout='fzyx')
        >>> f.layout
        (2, 1, 0)
    """
    result = []
    if description:
        field_descriptions, dtype, shape = _parse_description(description)
        layout = 'numpy' if layout is None else layout
        for field_name, idx_shape in field_descriptions:
            if field_name in 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))
            elif isinstance(shape, tuple):
                f = Field.create_fixed_size(field_name, shape + idx_shape, dtype=dtype,
                                            index_dimensions=len(idx_shape), layout=layout)
            elif isinstance(shape, int):
                f = Field.create_generic(field_name, spatial_dimensions=shape, dtype=dtype,
                                         index_shape=idx_shape, layout=layout)
            elif shape is None:
                f = Field.create_generic(field_name, spatial_dimensions=2, dtype=dtype,
                                         index_shape=idx_shape, layout=layout)
            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))

    if len(result) == 0:
        return None
    elif len(result) == 1:
        return result[0]
    else:
        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

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
    def is_generic(field):
        assert isinstance(field, Field)
Martin Bauer's avatar
Martin Bauer committed
        return field.field_type == FieldType.GENERIC
Martin Bauer's avatar
Martin Bauer committed
    def is_indexed(field):
        assert isinstance(field, Field)
Martin Bauer's avatar
Martin Bauer committed
        return field.field_type == FieldType.INDEXED
Martin Bauer's avatar
Martin Bauer committed
    def is_buffer(field):
        assert isinstance(field, Field)
Martin Bauer's avatar
Martin Bauer committed
        return field.field_type == FieldType.BUFFER
    """
    With fields one can formulate stencil-like update rules on structured grids.
    This Field class knows about the dimension, memory layout (strides) and optionally about the size of an array.

    Creating Fields:

        To create a field use one of the static create* members. There are two options:

        1. create a kernel with fixed loop sizes i.e. the shape of the array is already known. This is usually the
Martin Bauer's avatar
Martin Bauer committed
           case if just-in-time compilation directly from Python is done. (see :func:`Field.create_from_numpy_array`)
        2. create a more general kernel that works for variable array sizes. This can be used to create kernels
Martin Bauer's avatar
Martin Bauer committed
           beforehand for a library. (see :func:`Field.create_generic`)

    Dimensions:
        A field has spatial and index dimensions, where the spatial dimensions come first.
        The interpretation is that the field has multiple cells in (usually) two or three dimensional space which are
Martin Bauer's avatar
Martin Bauer committed
        looped over. Additionally  N values are stored per cell. In this case spatial_dimensions is two or three,
        and index_dimensions equals N. If you want to store a matrix on each point in a two dimensional grid, there
Martin Bauer's avatar
Martin Bauer committed
        are four dimensions, two spatial and two index dimensions: ``len(arr.shape) == spatial_dims + index_dims``

    Indexing:
        When accessing (indexing) a field the result is a FieldAccess which is derived from sympy Symbol.
Martin Bauer's avatar
Martin Bauer committed
        First specify the spatial offsets in [], then in case index_dimension>0 the indices in ()
        e.g. ``f[-1,0,0](7)``

    Example without index dimensions:
        >>> a = np.zeros([10, 10])
Martin Bauer's avatar
Martin Bauer committed
        >>> f = Field.create_from_numpy_array("f", a, index_dimensions=0)
        >>> jacobi = ( f[-1,0] + f[1,0] + f[0,-1] + f[0,1] ) / 4

    Example with index dimensions: LBM D2Q9 stream pull
Martin Bauer's avatar
Martin Bauer committed
        >>> from pystencils import Assignment
        >>> stencil = np.array([[0,0], [0,1], [0,-1]])
Martin Bauer's avatar
Martin Bauer committed
        >>> src = Field.create_generic("src", spatial_dimensions=2, index_dimensions=1)
        >>> dst = Field.create_generic("dst", spatial_dimensions=2, index_dimensions=1)
        >>> for i, offset in enumerate(stencil):
        ...     Assignment(dst[0,0](i), src[-offset](i))
        Assignment(dst_C^0, src_C^0)
        Assignment(dst_C^1, src_S^1)
        Assignment(dst_C^2, src_N^2)
Martin Bauer's avatar
Martin Bauer committed
    def create_generic(field_name, spatial_dimensions, dtype=np.float64, index_dimensions=0, layout='numpy',
                       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

Martin Bauer's avatar
Martin Bauer committed
        Args:
            field_name: symbolic name for the field
            dtype: numpy data type of the array the kernel is called with later
            spatial_dimensions: see documentation of Field
            index_dimensions: see documentation of Field
            layout: tuple specifying the loop ordering of the spatial dimensions e.g. (2, 1, 0 ) means that
                    the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
Martin Bauer's avatar
Martin Bauer committed
                    over dimension 0. Also allowed: the strings 'numpy' (0,1,..d) or 'reverse_numpy' (d, ..., 1, 0)
Martin Bauer's avatar
Martin Bauer committed
            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
        if index_shape is not None:
            assert index_dimensions == 0 or index_dimensions == len(index_shape)
            index_dimensions = len(index_shape)
        if isinstance(layout, str):
Martin Bauer's avatar
Martin Bauer committed
            layout = spatial_layout_string_to_tuple(layout, dim=spatial_dimensions)
        shape_symbol = IndexedBase(TypedSymbol(Field.SHAPE_PREFIX + field_name, Field.SHAPE_DTYPE), shape=(1,))
        stride_symbol = IndexedBase(TypedSymbol(Field.STRIDE_PREFIX + field_name, Field.STRIDE_DTYPE), shape=(1,))
        total_dimensions = spatial_dimensions + index_dimensions
        if index_shape is None or len(index_shape) == 0:
            shape = tuple([shape_symbol[i] for i in range(total_dimensions)])
Martin Bauer's avatar
Martin Bauer committed
            shape = tuple([shape_symbol[i] for i in range(spatial_dimensions)] + list(index_shape))
Martin Bauer's avatar
Martin Bauer committed
        strides = tuple([stride_symbol[i] for i in range(total_dimensions)])
Martin Bauer's avatar
Martin Bauer committed
        np_data_type = np.dtype(dtype)
        if np_data_type.fields is not None:
            if index_dimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

Martin Bauer's avatar
Martin Bauer committed
        return Field(field_name, field_type, dtype, layout, shape, strides)
    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
    def create_from_numpy_array(field_name: str, array: np.ndarray, index_dimensions: int = 0) -> '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.
Martin Bauer's avatar
Martin Bauer committed

        Args:
            field_name: symbolic name for the field
            array: numpy array
            index_dimensions: see documentation of Field
Martin Bauer's avatar
Martin Bauer committed
        spatial_dimensions = len(array.shape) - index_dimensions
        if spatial_dimensions < 1:
            raise ValueError("Too many index dimensions. At least one spatial dimension required")

Martin Bauer's avatar
Martin Bauer committed
        full_layout = get_layout_of_array(array)
        spatial_layout = tuple([i for i in full_layout if i < spatial_dimensions])
        assert len(spatial_layout) == spatial_dimensions
Martin Bauer's avatar
Martin Bauer committed
        strides = tuple([s // np.dtype(array.dtype).itemsize for s in array.strides])
        shape = tuple(int(s) for s in array.shape)
Martin Bauer's avatar
Martin Bauer committed
        numpy_dtype = np.dtype(array.dtype)
        if numpy_dtype.fields is not None:
            if index_dimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

Martin Bauer's avatar
Martin Bauer committed
        return Field(field_name, FieldType.GENERIC, array.dtype, spatial_layout, shape, strides)
Martin Bauer's avatar
Martin Bauer committed
    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':
        Creates a field with fixed sizes i.e. can be called only with arrays of the same size and layout
Martin Bauer's avatar
Martin Bauer committed
        Args:
            field_name: symbolic name for the field
            shape: overall shape of the array
            index_dimensions: how many of the trailing dimensions are interpreted as index (as opposed to spatial)
            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)
Martin Bauer's avatar
Martin Bauer committed
        spatial_dimensions = len(shape) - index_dimensions
        assert spatial_dimensions >= 1
Martin Bauer's avatar
Martin Bauer committed
            layout = layout_string_to_tuple(layout, spatial_dimensions + index_dimensions)

        shape = tuple(int(s) for s in shape)
        if strides is None:
Martin Bauer's avatar
Martin Bauer committed
            strides = compute_strides(shape, layout)
        else:
            assert len(strides) == len(shape)
            strides = tuple([s // np.dtype(dtype).itemsize for s in strides])
Martin Bauer's avatar
Martin Bauer committed
        numpy_dtype = np.dtype(dtype)
        if numpy_dtype.fields is not None:
            if index_dimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

Martin Bauer's avatar
Martin Bauer committed
        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)
Martin Bauer's avatar
Martin Bauer committed
    def __init__(self, field_name, field_type, dtype, layout, shape, strides):
        """Do not use directly. Use static create* methods"""
        self._field_name = field_name
Martin Bauer's avatar
Martin Bauer committed
        assert isinstance(field_type, FieldType)
        assert len(shape) == len(strides)
Martin Bauer's avatar
Martin Bauer committed
        self.field_type = field_type
Martin Bauer's avatar
Martin Bauer committed
        self._dtype = create_type(dtype)
Martin Bauer's avatar
Martin Bauer committed
        self._layout = normalize_layout(layout)
        self.shape = shape
        self.strides = strides
        self.latex_name = None  # type: Optional[str]
Martin Bauer's avatar
Martin Bauer committed
    def new_field_with_different_name(self, new_name):
Martin Bauer's avatar
Martin Bauer committed
        return Field(new_name, self.field_type, self._dtype, self._layout, self.shape, self.strides)
    @property
Martin Bauer's avatar
Martin Bauer committed
    def spatial_dimensions(self) -> int:
        return len(self._layout)

    @property
Martin Bauer's avatar
Martin Bauer committed
    def index_dimensions(self) -> int:
        return len(self.shape) - len(self._layout)

    @property
    def layout(self):
        return self._layout

    @property
Martin Bauer's avatar
Martin Bauer committed
    def name(self) -> str:
        return self._field_name
Martin Bauer's avatar
Martin Bauer committed
    def spatial_shape(self) -> Tuple[int, ...]:
        return self.shape[:self.spatial_dimensions]
Martin Bauer's avatar
Martin Bauer committed
    def has_fixed_shape(self):
Martin Bauer's avatar
Martin Bauer committed
        return is_integer_sequence(self.shape)
    @property
Martin Bauer's avatar
Martin Bauer committed
    def index_shape(self):
        return self.shape[self.spatial_dimensions:]
Martin Bauer's avatar
Martin Bauer committed
    def has_fixed_index_shape(self):
        return is_integer_sequence(self.index_shape)
    @property
Martin Bauer's avatar
Martin Bauer committed
    def spatial_strides(self):
        return self.strides[:self.spatial_dimensions]
Martin Bauer's avatar
Martin Bauer committed
    def index_strides(self):
        return self.strides[self.spatial_dimensions:]

    @property
    def dtype(self):
        return self._dtype

    def __repr__(self):
        return self._field_name
Martin Bauer's avatar
Martin Bauer committed
    def neighbor(self, coord_id, offset):
        offset_list = [0] * self.spatial_dimensions
        offset_list[coord_id] = offset
        return Field.Access(self, tuple(offset_list))
    def neighbors(self, stencil):
        return [self.__getitem__(s) for s in stencil]
Martin Bauer's avatar
Martin Bauer committed
    def center_vector(self):
        index_shape = self.index_shape
        if len(index_shape) == 0:
Martin Bauer's avatar
Martin Bauer committed
        elif len(index_shape) == 1:
            return sp.Matrix([self(i) for i in range(index_shape[0])])
        elif len(index_shape) == 2:
            def cb(*args):
                r = self.__call__(*args)
                return r
Martin Bauer's avatar
Martin Bauer committed
            return sp.Matrix(*index_shape, cb)
    def center(self):
Martin Bauer's avatar
Martin Bauer committed
        center = tuple([0] * self.spatial_dimensions)
        return Field.Access(self, center)

    def __getitem__(self, offset):
        if type(offset) is np.ndarray:
            offset = tuple(offset)
        if type(offset) is str:
Martin Bauer's avatar
Martin Bauer committed
            offset = tuple(direction_string_to_offset(offset, self.spatial_dimensions))
        if type(offset) is not tuple:
            offset = (offset,)
Martin Bauer's avatar
Martin Bauer committed
        if len(offset) != self.spatial_dimensions:
            raise ValueError("Wrong number of spatial indices: "
Martin Bauer's avatar
Martin Bauer committed
                             "Got %d, expected %d" % (len(offset), self.spatial_dimensions))
        return Field.Access(self, offset)

    def __call__(self, *args, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
        center = tuple([0] * self.spatial_dimensions)
        return Field.Access(self, center)(*args, **kwargs)

    def __hash__(self):
        return hash((self._layout, self.shape, self.strides, self._dtype, self.field_type, self._field_name))

    def __eq__(self, other):
Martin Bauer's avatar
Martin Bauer committed
        self_tuple = (self.shape, self.strides, self.name, self.dtype, self.field_type)
        other_tuple = (other.shape, other.strides, other.name, other.dtype, other.field_type)
Martin Bauer's avatar
Martin Bauer committed
        return self_tuple == other_tuple
    PREFIX = "f"
    STRIDE_PREFIX = PREFIX + "stride_"
    SHAPE_PREFIX = PREFIX + "shape_"
Martin Bauer's avatar
Martin Bauer committed
    STRIDE_DTYPE = create_composite_type_from_string("const int *")
    SHAPE_DTYPE = create_composite_type_from_string("const int *")
    DATA_PREFIX = PREFIX + "d_"
Martin Bauer's avatar
Martin Bauer committed
    # noinspection PyAttributeOutsideInit,PyUnresolvedReferences
    class Access(sp.Symbol):
        def __new__(cls, name, *args, **kwargs):
            obj = Field.Access.__xnew_cached_(cls, name, *args, **kwargs)
            return obj

        def __new_stage2__(self, field, offsets=(0, 0, 0), idx=None):
Martin Bauer's avatar
Martin Bauer committed
            field_name = field.name
            offsets_and_index = chain(offsets, idx) if idx is not None else offsets
            constant_offsets = not any([isinstance(o, sp.Basic) and not o.is_Integer for o in offsets_and_index])
Martin Bauer's avatar
Martin Bauer committed
                idx = tuple([0] * field.index_dimensions)
Martin Bauer's avatar
Martin Bauer committed
            if constant_offsets:
                offset_name = offset_to_direction_string(offsets)
                if field.index_dimensions == 0:
Martin Bauer's avatar
Martin Bauer committed
                elif field.index_dimensions == 1:
                    superscript = str(idx[0])
Martin Bauer's avatar
Martin Bauer committed
                    idx_str = ",".join([str(e) for e in idx])
                    superscript = idx_str
                if field.has_fixed_index_shape and not isinstance(field.dtype, StructType):
                    for i, bound in zip(idx, field.index_shape):
                        if i >= bound:
                            raise ValueError("Field index out of bounds")
Martin Bauer's avatar
Martin Bauer committed
                offset_name = "%0.10X" % (abs(hash(tuple(offsets_and_index))))
Martin Bauer's avatar
Martin Bauer committed
            symbol_name = "%s_%s" % (field_name, offset_name)
            if superscript is not None:
Martin Bauer's avatar
Martin Bauer committed
                symbol_name += "^" + superscript
Martin Bauer's avatar
Martin Bauer committed
            obj = super(Field.Access, self).__xnew__(self, symbol_name)
            obj._field = field
            obj._offsets = []
            for o in offsets:
                if isinstance(o, sp.Basic):
                    obj._offsets.append(o)
                else:
                    obj._offsets.append(int(o))
Martin Bauer's avatar
Martin Bauer committed
            obj._offsetName = offset_name
            obj._superscript = superscript
            obj._index = idx

            return obj

            return self.field, self.offsets, self.index
Martin Bauer's avatar
Martin Bauer committed
        # noinspection SpellCheckingInspection
        __xnew__ = staticmethod(__new_stage2__)
Martin Bauer's avatar
Martin Bauer committed
        # noinspection SpellCheckingInspection
        __xnew_cached_ = staticmethod(cacheit(__new_stage2__))

        def __call__(self, *idx):
Martin Bauer's avatar
Martin Bauer committed
            if self._index != tuple([0] * self.field.index_dimensions):
                raise ValueError("Indexing an already indexed Field.Access")

            idx = tuple(idx)
Martin Bauer's avatar
Martin Bauer committed
            if self.field.index_dimensions == 0 and idx == (0,):
Martin Bauer's avatar
Martin Bauer committed
            if len(idx) != self.field.index_dimensions:
                raise ValueError("Wrong number of indices: "
Martin Bauer's avatar
Martin Bauer committed
                                 "Got %d, expected %d" % (len(idx), self.field.index_dimensions))
            return Field.Access(self.field, self._offsets, idx)

        def __getitem__(self, *idx):
            return self.__call__(*idx)

Martin Bauer's avatar
Martin Bauer committed
        def __iter__(self):
            """This is necessary to work with parts of sympy that test if an object is iterable (e.g. simplify).
            The __getitem__ would make it iterable"""
            raise TypeError("Field access is not iterable")

        @property
        def field(self):
            return self._field

        @property
        def offsets(self):
            return tuple(self._offsets)
        @property
Martin Bauer's avatar
Martin Bauer committed
        def required_ghost_layers(self):
            return int(np.max(np.abs(self._offsets)))

        @property
Martin Bauer's avatar
Martin Bauer committed
        def nr_of_coordinates(self):
            return len(self._offsets)

        @property
Martin Bauer's avatar
Martin Bauer committed
        def offset_name(self) -> str:
            return self._offsetName

        @property
        def index(self):
            return self._index

Martin Bauer's avatar
Martin Bauer committed
        def neighbor(self, coord_id: int, offset: Sequence[int]) -> 'Field.Access':
            offset_list = list(self.offsets)
            offset_list[coord_id] += offset
            return Field.Access(self.field, tuple(offset_list), self.index)
Martin Bauer's avatar
Martin Bauer committed
        def get_shifted(self, *shift)-> 'Field.Access':
            return Field.Access(self.field, tuple(a + b for a, b in zip(shift, self.offsets)), self.index)

        def at_index(self, *idx_tuple):
            return Field.Access(self.field, self.offsets, idx_tuple)

        def _hashable_content(self):
Martin Bauer's avatar
Martin Bauer committed
            super_class_contents = list(super(Field.Access, self)._hashable_content())
            t = tuple(super_class_contents + [hash(self._field), self._index] + self._offsets)
Martin Bauer's avatar
Martin Bauer committed
        def _latex(self, _):
            n = self._field.latex_name if self._field.latex_name else self._field.name
Martin Bauer's avatar
Martin Bauer committed
            if self._superscript:
                return "{{%s}_{%s}^{%s}}" % (n, self._offsetName, self._superscript)
            else:
                return "{{%s}_{%s}}" % (n, self._offsetName)
Martin Bauer's avatar
Martin Bauer committed
def get_layout_from_strides(strides: Sequence[int], index_dimension_ids: Optional[List[int]] = None):
    index_dimension_ids = [] if index_dimension_ids is None else index_dimension_ids
    coordinates = list(range(len(strides)))
Martin Bauer's avatar
Martin Bauer committed
    relevant_strides = [stride for i, stride in enumerate(strides) if i not in index_dimension_ids]
    result = [x for (y, x) in sorted(zip(relevant_strides, coordinates), key=lambda pair: pair[0], reverse=True)]
Martin Bauer's avatar
Martin Bauer committed
    return normalize_layout(result)
Martin Bauer's avatar
Martin Bauer committed
def get_layout_of_array(arr: np.ndarray, index_dimension_ids: Optional[List[int]] = None):
    """ Returns a list indicating the memory layout (linearization order) of the numpy array.

    Examples:
        >>> get_layout_of_array(np.zeros([3,3,3]))
        (0, 1, 2)

    In this example the loop over the zeroth coordinate should be the outermost loop,
    followed by the first and second. Elements arr[x,y,0] and arr[x,y,1] are adjacent in memory.
    Normally constructed numpy arrays have this order, however by stride tricks or other frameworks, arrays
    with different memory layout can be created.
Martin Bauer's avatar
Martin Bauer committed
    The index_dimension_ids parameter leaves specifies which coordinates should not be
Martin Bauer's avatar
Martin Bauer committed
    index_dimension_ids = [] if index_dimension_ids is None else index_dimension_ids
    return get_layout_from_strides(arr.strides, index_dimension_ids)
Martin Bauer's avatar
Martin Bauer committed
def create_numpy_array_with_layout(shape, layout, alignment=False, byte_offset=0, **kwargs):
    """Creates numpy array with given memory layout.

    Args:
        shape: shape of the resulting array
        layout: layout as tuple, where the coordinates are ordered from slow to fast
        alignment: number of bytes to align the beginning and the innermost coordinate to, or False for no alignment
        byte_offset: only used when alignment is specified, align not beginning but address at this offset
                     mostly used to align first inner cell, not ghost cells

    Example:
        >>> res = create_numpy_array_with_layout(shape=(2, 3, 4, 5), layout=(3, 2, 0, 1))
        >>> res.shape
        (2, 3, 4, 5)
        >>> get_layout_of_array(res)
        (3, 2, 0, 1)
    """
    assert set(layout) == set(range(len(shape))), "Wrong layout descriptor"
Martin Bauer's avatar
Martin Bauer committed
    cur_layout = list(range(len(shape)))
Martin Bauer's avatar
Martin Bauer committed
        if cur_layout[i] != layout[i]:
            index_to_swap_with = cur_layout.index(layout[i])
            swaps.append((i, index_to_swap_with))
            cur_layout[i], cur_layout[index_to_swap_with] = cur_layout[index_to_swap_with], cur_layout[i]
    assert tuple(cur_layout) == tuple(layout)

    shape = list(shape)
    for a, b in swaps:
        shape[a], shape[b] = shape[b], shape[a]

    if not alignment:
        res = np.empty(shape, order='c', **kwargs)
    else:
        if alignment is True:
            alignment = 8 * 4
Martin Bauer's avatar
Martin Bauer committed
        res = aligned_empty(shape, alignment, byte_offset=byte_offset, **kwargs)
    for a, b in reversed(swaps):
        res = res.swapaxes(a, b)
    return res


Martin Bauer's avatar
Martin Bauer committed
def spatial_layout_string_to_tuple(layout_str: str, dim: int) -> Tuple[int, ...]:
    if layout_str in ('fzyx', 'zyxf'):
        assert dim <= 3
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
    if layout_str in ('fzyx', 'f', 'reverse_numpy', 'SoA'):
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
    elif layout_str in ('c', 'numpy', 'AoS'):
Martin Bauer's avatar
Martin Bauer committed
    raise ValueError("Unknown layout descriptor " + layout_str)
Martin Bauer's avatar
Martin Bauer committed
def layout_string_to_tuple(layout_str, dim):
    layout_str = layout_str.lower()
    if layout_str == 'fzyx' or layout_str == 'soa':
        assert dim <= 4
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
    elif layout_str == 'zyxf' or layout_str == 'aos':
        assert dim <= 4
Martin Bauer's avatar
Martin Bauer committed
        return tuple(reversed(range(dim - 1))) + (dim - 1,)
Martin Bauer's avatar
Martin Bauer committed
    elif layout_str == 'f' or layout_str == 'reverse_numpy':
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
    elif layout_str == 'c' or layout_str == 'numpy':
        return tuple(range(dim))
Martin Bauer's avatar
Martin Bauer committed
    raise ValueError("Unknown layout descriptor " + layout_str)
Martin Bauer's avatar
Martin Bauer committed
def normalize_layout(layout):
    """Takes a layout tuple and subtracts the minimum from all entries"""
Martin Bauer's avatar
Martin Bauer committed
    min_entry = min(layout)
    return tuple(i - min_entry for i in layout)
Martin Bauer's avatar
Martin Bauer committed
def compute_strides(shape, layout):
    """
    Computes strides assuming no padding exists
Martin Bauer's avatar
Martin Bauer committed

    Args:
        shape: shape (size) of array
        layout: layout specification as tuple

    Returns:
        strides in elements, not in bytes
Martin Bauer's avatar
Martin Bauer committed
    dim = len(shape)
    assert len(layout) == dim
    assert len(set(layout)) == dim
    strides = [0] * dim
        strides[j] = product
        product *= shape[j]
    return tuple(strides)
Martin Bauer's avatar
Martin Bauer committed
def offset_component_to_direction_string(coordinate_id: int, value: int) -> str:
    """Translates numerical offset to string notation.

    x offsets are labeled with east 'E' and 'W',
    y offsets with north 'N' and 'S' and
    z offsets with top 'T' and bottom 'B'
    If the absolute value of the offset is bigger than 1, this number is prefixed.

Martin Bauer's avatar
Martin Bauer committed
    Args:
        coordinate_id: integer 0, 1 or 2 standing for x,y and z
        value: integer offset

    Examples:
        >>> offset_component_to_direction_string(0, 1)
        'E'
        >>> offset_component_to_direction_string(1, 2)
        '2N'
    assert 0 <= coordinate_id < 3, "Works only for at most 3D arrays"
Martin Bauer's avatar
Martin Bauer committed
    name_components = (('W', 'E'),  # west, east
                       ('S', 'N'),  # south, north
                       ('B', 'T'))  # bottom, top
    if value == 0:
        result = ""
    elif value < 0:
Martin Bauer's avatar
Martin Bauer committed
        result = name_components[coordinate_id][0]
Martin Bauer's avatar
Martin Bauer committed
        result = name_components[coordinate_id][1]
    if abs(value) > 1:
        result = "%d%s" % (abs(value), result)
    return result


Martin Bauer's avatar
Martin Bauer committed
def offset_to_direction_string(offsets: Sequence[int]) -> str:
    """
    Translates numerical offset to string notation.
Martin Bauer's avatar
Martin Bauer committed
    For details see :func:`offset_component_to_direction_string`
    Args:
        offsets: 3-tuple with x,y,z offset

    Examples:
        >>> offset_to_direction_string([1, -1, 0])
        'SE'
        >>> offset_to_direction_string(([-3, 0, -2]))
        '2B3W'
    if len(offsets) > 3:
        return str(offsets)
    names = ["", "", ""]
Martin Bauer's avatar
Martin Bauer committed
    for i in range(len(offsets)):
        names[i] = offset_component_to_direction_string(i, offsets[i])
    name = "".join(reversed(names))
    if name == "":
        name = "C"
    return name


Martin Bauer's avatar
Martin Bauer committed
def direction_string_to_offset(direction: str, dim: int = 3):
Martin Bauer's avatar
Martin Bauer committed
    Reverse mapping of :func:`offset_to_direction_string`
Martin Bauer's avatar
Martin Bauer committed

    Args:
        direction: string representation of offset
        dim: dimension of offset, i.e the length of the returned list

    Examples:
        >>> direction_string_to_offset('NW', dim=3)
        array([-1,  1,  0])
        >>> direction_string_to_offset('NW', dim=2)
        array([-1,  1])
        >>> direction_string_to_offset(offset_to_direction_string((3,-2,1)))
        array([ 3, -2,  1])
Martin Bauer's avatar
Martin Bauer committed
    offset_dict = {
        'C': np.array([0, 0, 0]),

        'W': np.array([-1, 0, 0]),
        'E': np.array([1, 0, 0]),

        'S': np.array([0, -1, 0]),
        'N': np.array([0, 1, 0]),

        'B': np.array([0, 0, -1]),
        'T': np.array([0, 0, 1]),
    }
    offset = np.array([0, 0, 0])

Martin Bauer's avatar
Martin Bauer committed
    while len(direction) > 0:
        factor = 1
Martin Bauer's avatar
Martin Bauer committed
        first_non_digit = 0
        while direction[first_non_digit].isdigit():
            first_non_digit += 1
        if first_non_digit > 0:
            factor = int(direction[:first_non_digit])
            direction = direction[first_non_digit:]
        cur_offset = offset_dict[direction[0]]
        offset += factor * cur_offset
        direction = direction[1:]
    return offset[:dim]

def _parse_type_description(type_description):
    if not type_description:
        return np.float64, None
    elif '[' in type_description:
        assert type_description[-1] == ']'
        splitted = type_description[:-1].split("[", )
        type_part, size_part = type_description[:-1].split("[", )
        if not type_part:
            type_part = "float64"
        if size_part.lower()[-1] == 'd':
            size_part = int(size_part[:-1])
        else:
            size_part = tuple(int(i) for i in size_part.split(','))
    else:
        type_part, size_part = type_description, None

    dtype = np.dtype(type_part).type
    return dtype, size_part


def _parse_field_description(description):
    if '(' not in description:
        return description, ()
    assert description[-1] == ')'
    name, index_shape = description[:-1].split('(')
    index_shape = tuple(int(i) for i in index_shape.split(','))
    return name, index_shape


def _parse_description(description):
    description = description.replace(' ', '')
    if ':' in description:
        name_descr, type_descr = description.split(':')
    else:
        name_descr, type_descr = description, ''

    # correct ',' splits inside brackets
    field_names = name_descr.split(',')
    cleaned_field_names = []
    prefix = ''
    for field_name in field_names:
        full_field_name = prefix + field_name
        if '(' in full_field_name and ')' not in full_field_name:
            prefix += field_name + ','
        else:
            prefix = ''
            cleaned_field_names.append(full_field_name)

    dtype, size = _parse_type_description(type_descr)
    fields_info = tuple(_parse_field_description(fd) for fd in cleaned_field_names)
    return fields_info, dtype, size