diff --git a/pystencils/__init__.py b/pystencils/__init__.py index a402a4d9638e92c878bcc8773108ce6c8664c70d..ed08ff1327c68a1a116edb21547a86f8b27a8c81 100644 --- a/pystencils/__init__.py +++ b/pystencils/__init__.py @@ -12,6 +12,8 @@ from .kernelcreation import create_indexed_kernel, create_kernel, create_stagger from .simp import AssignmentCollection from .slicing import make_slice from .sympyextensions import SymbolCreator +from .spatial_coordinates import (x_, x_staggered, x_staggered_vector, x_vector, + y_, y_staggered, z_, z_staggered) try: import pystencils_autodiff @@ -30,5 +32,8 @@ __all__ = ['Field', 'FieldType', 'fields', 'SymbolCreator', 'create_data_handling', 'kernel', + 'x_', 'y_', 'z_', + 'x_staggered', 'y_staggered', 'z_staggered', + 'x_vector', 'x_staggered_vector', 'fd', 'stencil'] diff --git a/pystencils/spatial_coordinates.py b/pystencils/spatial_coordinates.py new file mode 100644 index 0000000000000000000000000000000000000000..6c3ba4db7d98b50bce9442619fbbce07782004a1 --- /dev/null +++ b/pystencils/spatial_coordinates.py @@ -0,0 +1,18 @@ + +import sympy + +import pystencils +import pystencils.astnodes + +x_, y_, z_ = tuple(pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(3)) +x_staggered, y_staggered, z_staggered = x_ + 0.5, y_ + 0.5, z_ + 0.5 + + +def x_vector(ndim): + return sympy.Matrix(tuple(pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(ndim))) + + +def x_staggered_vector(ndim): + return sympy.Matrix(tuple( + pystencils.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) + 0.5 for i in range(ndim) + ))