Commit 6ad0c6dd authored by Frederik Hennig's avatar Frederik Hennig
Browse files

Some refactoring

parent b5ac8c01
from .indexing import BetweenTimestepsIndexing
from .indexing import BetweenTimestepsIndexing, NeighbourOffsetArrays
from .communication import get_communication_slices, PeriodicityHandling
from .utility import Timestep, get_accessor
from .utility import Timestep, get_accessor, is_inplace, get_timesteps, \
numeric_index, numeric_offsets, inverse_dir_index, AccessPdfValues
__all__ = ['BetweenTimestepsIndexing',
__all__ = ['BetweenTimestepsIndexing', 'NeighbourOffsetArrays',
'get_communication_slices', 'PeriodicityHandling',
'Timestep', 'get_accessor']
'Timestep', 'get_accessor', 'is_inplace', 'get_timesteps',
'numeric_index', 'numeric_offsets', 'inverse_dir_index', 'AccessPdfValues']
......@@ -6,7 +6,7 @@ from pystencils.datahandling import SerialDataHandling
from itertools import chain
def trim_slice_in_direction(slices, direction):
def _trim_slice_in_direction(slices, direction):
assert len(slices) == len(direction)
result = []
......@@ -21,15 +21,15 @@ def trim_slice_in_direction(slices, direction):
return tuple(result)
def extend_dir(direction):
def _extend_dir(direction):
if len(direction) == 0:
yield tuple()
elif direction[0] == 0:
for d in [-1, 0, 1]:
for rest in extend_dir(direction[1:]):
for rest in _extend_dir(direction[1:]):
yield (d, ) + rest
else:
for rest in extend_dir(direction[1:]):
for rest in _extend_dir(direction[1:]):
yield (direction[0], ) + rest
......@@ -77,7 +77,7 @@ def get_communication_slices(
slices_for_dir = []
for streaming_dir in set(extend_dir(comm_dir)) & set(stencil):
for streaming_dir in set(_extend_dir(comm_dir)) & set(stencil):
d = stencil.index(streaming_dir)
write_offsets = numeric_offsets(write_accesses[d])
write_index = numeric_index(write_accesses[d])[0]
......@@ -85,7 +85,7 @@ def get_communication_slices(
tangential_dir = tuple(s - c for s, c in zip(streaming_dir, comm_dir))
origin_slice = get_slice_before_ghost_layer(comm_dir, ghost_layers=ghost_layers, thickness=1)
origin_slice = _fix_length_one_slices(origin_slice)
src_slice = shift_slice(trim_slice_in_direction(origin_slice, tangential_dir), write_offsets)
src_slice = shift_slice(_trim_slice_in_direction(origin_slice, tangential_dir), write_offsets)
neighbour_transform = _get_neighbour_transform(comm_dir, ghost_layers)
dst_slice = shift_slice(src_slice, neighbour_transform)
......
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