slicing.py

import sympy as sp
from pystencils.field import create_numpy_array_with_layout, get_layout_of_array


class SliceMaker(object):
    def __getitem__(self, item):
        return item


make_slice = SliceMaker()


class SlicedGetter(object):
    def __init__(self, function_returning_array):
        self._functionReturningArray = function_returning_array

    def __getitem__(self, item):
        return self._functionReturningArray(item)


class SlicedGetterDataHandling:
    def __init__(self, data_handling, name):
        self.dh = data_handling
        self.name = name

    def __getitem__(self, slice_obj):
        if slice_obj is None:
            slice_obj = make_slice[:, :] if self.data_handling.dim == 2 else make_slice[:, :, 0.5]
        return self.dh.gather_array(self.name, slice_obj).squeeze()


def normalize_slice(slices, sizes):
    """Converts slices with floating point and/or negative entries to integer slices"""

    if len(slices) != len(sizes):
        raise ValueError("Slice dimension does not match sizes")

    result = []

    for s, size in zip(slices, sizes):
        if type(s) is int:
            if s < 0:
                s = size + s
            result.append(s)
            continue
        if type(s) is float:
            result.append(int(s * size))
            continue

        assert (type(s) is slice)

        if s.start is None:
            new_start = 0
        elif type(s.start) is float:
            new_start = int(s.start * size)
        elif not isinstance(s.start, sp.Basic) and s.start < 0:
            new_start = size + s.start
        else:
            new_start = s.start

        if s.stop is None:
            new_stop = size
        elif type(s.stop) is float:
            new_stop = int(s.stop * size)
        elif not isinstance(s.stop, sp.Basic) and s.stop < 0:
            new_stop = size + s.stop
        else:
            new_stop = s.stop

        result.append(slice(new_start, new_stop, s.step if s.step is not None else 1))

    return tuple(result)


def shift_slice(slices, offset):
    def shift_slice_component(slice_comp, shift_offset):
        if slice_comp is None:
            return None
        elif isinstance(slice_comp, int):
            return slice_comp + shift_offset
        elif isinstance(slice_comp, float):
            return slice_comp  # relative entries are not shifted
        elif isinstance(slice_comp, slice):
            return slice(shift_slice_component(slice_comp.start, shift_offset),
                         shift_slice_component(slice_comp.stop, shift_offset),
                         slice_comp.step)
        else:
            raise ValueError()

    if hasattr(offset, '__len__'):
        return [shift_slice_component(k, off) for k, off in zip(slices, offset)]
    else:
        return [shift_slice_component(k, offset) for k in slices]


def slice_from_direction(direction_name, dim, normal_offset=0, tangential_offset=0):
    """
    Create a slice from a direction named by compass scheme:
        i.e. 'N' for north returns same as make_slice[:, -1]
        the naming is:
            - x: W, E (west, east)
            - y: S, N (south, north)
            - z: B, T (bottom, top)
    Also combinations are allowed like north-east 'NE'

    :param direction_name: name of direction as explained above
    :param dim: dimension of the returned slice (should be 2 or 3)
    :param normal_offset: the offset in 'normal' direction: e.g. slice_from_direction('N',2, normal_offset=2)
                         would return make_slice[:, -3]
    :param tangential_offset: offset in the other directions: e.g. slice_from_direction('N',2, tangential_offset=2)
                         would return make_slice[2:-2, -1]
    """
    if tangential_offset == 0:
        result = [slice(None, None, None)] * dim
    else:
        result = [slice(tangential_offset, -tangential_offset, None)] * dim

    normal_slice_high, normal_slice_low = -1 - normal_offset, normal_offset

    for dim_idx, (low_name, high_name) in enumerate([('W', 'E'), ('S', 'N'), ('B', 'T')]):
        if low_name in direction_name:
            assert high_name not in direction_name, "Invalid direction name"
            result[dim_idx] = normal_slice_low
        if high_name in direction_name:
            assert low_name not in direction_name, "Invalid direction name"
            result[dim_idx] = normal_slice_high
    return tuple(result)


def remove_ghost_layers(arr, index_dimensions=0, ghost_layers=1):
    if ghost_layers <= 0:
        return arr
    dimensions = len(arr.shape)
    spatial_dimensions = dimensions - index_dimensions
    indexing = [slice(ghost_layers, -ghost_layers, None), ] * spatial_dimensions
    indexing += [slice(None, None, None)] * index_dimensions
    return arr[tuple(indexing)]


def add_ghost_layers(arr, index_dimensions=0, ghost_layers=1, layout=None):
    dimensions = len(arr.shape)
    spatial_dimensions = dimensions - index_dimensions
    new_shape = [e + 2 * ghost_layers for e in arr.shape[:spatial_dimensions]] + list(arr.shape[spatial_dimensions:])
    if layout is None:
        layout = get_layout_of_array(arr)
    result = create_numpy_array_with_layout(new_shape, layout)
    result.fill(0.0)
    indexing = [slice(ghost_layers, -ghost_layers, None), ] * spatial_dimensions
    indexing += [slice(None, None, None)] * index_dimensions
    result[tuple(indexing)] = arr
    return result


def get_slice_before_ghost_layer(direction, ghost_layers=1, thickness=None, full_slice=False):
    """
    Returns slicing expression for region before ghost layer
    :param direction: tuple specifying direction of slice
    :param ghost_layers: number of ghost layers
    :param thickness: thickness of the slice, defaults to number of ghost layers
    :param full_slice:  if true also the ghost cells in directions orthogonal to direction are contained in the
                       returned slice. Example (d=W ): if full_slice then also the ghost layer in N-S and T-B
                       are included, otherwise only inner cells are returned
    """
    if not thickness:
        thickness = ghost_layers
    full_slice_inc = ghost_layers if not full_slice else 0
    slices = []
    for dir_component in direction:
        if dir_component == -1:
            s = slice(ghost_layers, thickness + ghost_layers)
        elif dir_component == 0:
            end = -full_slice_inc
            s = slice(full_slice_inc, end if end != 0 else None)
        elif dir_component == 1:
            start = -thickness - ghost_layers
            end = -ghost_layers
            s = slice(start if start != 0 else None, end if end != 0 else None)
        else:
            raise ValueError("Invalid direction: only -1, 0, 1 components are allowed")
        slices.append(s)
    return tuple(slices)


def get_ghost_region_slice(direction, ghost_layers=1, thickness=None, full_slice=False):
    """
    Returns slice of ghost region. For parameters see :func:`get_slice_before_ghost_layer`
    """
    if not thickness:
        thickness = ghost_layers
    assert thickness > 0
    assert thickness <= ghost_layers
    full_slice_inc = ghost_layers if not full_slice else 0
    slices = []
    for dir_component in direction:
        if dir_component == -1:
            s = slice(ghost_layers - thickness, ghost_layers)
        elif dir_component == 0:
            end = -full_slice_inc
            s = slice(full_slice_inc, end if end != 0 else None)
        elif dir_component == 1:
            start = -ghost_layers
            end = - ghost_layers + thickness
            s = slice(start if start != 0 else None, end if end != 0 else None)
        else:
            raise ValueError("Invalid direction: only -1, 0, 1 components are allowed")
        slices.append(s)
    return tuple(slices)


def get_periodic_boundary_src_dst_slices(stencil, ghost_layers=1, thickness=None):
    src_dst_slice_tuples = []

    for d in stencil:
        if sum([abs(e) for e in d]) == 0:
            continue
        inv_dir = (-e for e in d)
        src = get_slice_before_ghost_layer(inv_dir, ghost_layers, thickness=thickness, full_slice=False)
        dst = get_ghost_region_slice(d, ghost_layers, thickness=thickness, full_slice=False)
        src_dst_slice_tuples.append((src, dst))
    return src_dst_slice_tuples


def get_periodic_boundary_functor(stencil, ghost_layers=1, thickness=None):
    """
    Returns a function that applies periodic boundary conditions
    :param stencil: sequence of directions e.g. ( [0,1], [0,-1] ) for y periodicity
    :param ghost_layers: how many ghost layers the array has
    :param thickness: how many of the ghost layers to copy, None means 'all'
    :return: function that takes a single array and applies the periodic copy operation
    """
    src_dst_slice_tuples = get_periodic_boundary_src_dst_slices(stencil, ghost_layers, thickness)

    def functor(pdfs, **_):
        for src_slice, dst_slice in src_dst_slice_tuples:
            pdfs[dst_slice] = pdfs[src_slice]

    return functor


def slice_intersection(slice1, slice2):
    slice1 = [s if not isinstance(s, int) else slice(s, s + 1, None) for s in slice1]
    slice2 = [s if not isinstance(s, int) else slice(s, s + 1, None) for s in slice2]

    new_min = [max(s1.start, s2.start) for s1, s2 in zip(slice1, slice2)]
    new_max = [min(s1.stop, s2.stop) for s1, s2 in zip(slice1, slice2)]
    if any(max_p - min_p < 0 for min_p, max_p in zip(new_min, new_max)):
        return None

    return [slice(min_p, max_p, None) for min_p, max_p in zip(new_min, new_max)]