From a3cc3a8d5a215adfb8f5689c153bbc4d82c5dcf0 Mon Sep 17 00:00:00 2001
From: Markus Holzer <markus.holzer@fau.de>
Date: Wed, 12 Jul 2023 18:13:29 +0200
Subject: [PATCH] Add adjacent direcitons to stencil module
---
pystencils/stencil.py | 34 ++++++++++++++++++++++++++++++++++
pystencils/utils.py | 15 +++++++++++++++
2 files changed, 49 insertions(+)
diff --git a/pystencils/stencil.py b/pystencils/stencil.py
index 73ecb1571..aa86427fc 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 22d61d0ba..f872ae48a 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.
--
GitLab