diff --git a/pystencils/stencil.py b/pystencils/stencil.py index 73ecb15713a54a5a5da3bbe68a8368699bb9145d..aa86427fcd59113e502c91e7444a6815a5307f96 100644 --- a/pystencils/stencil.py +++ b/pystencils/stencil.py @@ -5,6 +5,8 @@ from typing import Sequence import numpy as np import sympy as sp +from pystencils.utils import binary_numbers + def inverse_direction(direction): """Returns inverse i.e. negative of given direction tuple @@ -293,6 +295,38 @@ def direction_string_to_offset(direction: str, dim: int = 3): return offset[:dim] +def adjacent_directions(direction): + """ + Returns all adjacent directions for a direction as tuple of tuples. This is useful for exmple to find all directions + relevant for neighbour communication. + + Args: + direction: tuple representing a direction. For example (0, 1, 0) for the northern side + + Examples: + >>> adjacent_directions((0, 0, 0)) + ((0, 0, 0),) + >>> adjacent_directions((0, 1, 0)) + ((0, 1, 0),) + >>> adjacent_directions((0, 1, 1)) + ((0, 0, 1), (0, 1, 0), (0, 1, 1)) + >>> adjacent_directions((-1, -1)) + ((-1, -1), (-1, 0), (0, -1)) + """ + result = set() + if all(e == 0 for e in direction): + result.add(direction) + return tuple(result) + binary_numbers_list = binary_numbers(len(direction)) + for adjacent_direction in binary_numbers_list: + for i, entry in enumerate(direction): + if entry == 0: + adjacent_direction[i] = 0 + if entry == -1 and adjacent_direction[i] == 1: + adjacent_direction[i] = -1 + if not all(e == 0 for e in adjacent_direction): + result.add(tuple(adjacent_direction)) + return tuple(sorted(result)) # -------------------------------------- Visualization ----------------------------------------------------------------- diff --git a/pystencils/utils.py b/pystencils/utils.py index 22d61d0bac6c402e10a7f48a07a55264ec4ddf27..f872ae48a54a2ae8c9437ec826be9b51c061f52e 100644 --- a/pystencils/utils.py +++ b/pystencils/utils.py @@ -96,6 +96,21 @@ def boolean_array_bounding_box(boolean_array): return bounds +def binary_numbers(n): + """Returns all binary numbers up to 2^n - 1 + + Example: + >>> binary_numbers(2) + [[0, 0], [0, 1], [1, 0], [1, 1]] + """ + result = list() + for i in range(1 << n): + binary_number = bin(i)[2:] + binary_number = '0' * (n - len(binary_number)) + binary_number + result.append((list(map(int, binary_number)))) + return result + + class LinearEquationSystem: """Symbolic linear system of equations - consisting of matrix and right hand side.