.gitattributes

-pystencils/_version.py export-subst
+src/pystencils/_version.py export-subst

pyproject.toml

@@ -12,7 +12,7 @@ authors = [
 ]
 license = { file = "COPYING.txt" }
 requires-python = ">=3.10"
-dependencies = ["sympy>=1.6,<=1.11.1", "numpy>=1.8.0", "appdirs", "joblib", "pyyaml"]
+dependencies = ["sympy>=1.9,<=1.12.1", "numpy>=1.8.0", "appdirs", "joblib", "pyyaml"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Framework :: Jupyter",
@@ -71,8 +71,7 @@ tests = [
 [build-system]
 requires = [
     "setuptools>=61",
-    "versioneer>=0.29",
-    "tomli; python_version < '3.11'",
+    "versioneer[toml]>=0.29",
     # 'Cython'
 ]
 build-backend = "setuptools.build_meta"

src/pystencils/__init__.py

@@ -15,7 +15,7 @@ from .config import (
     OpenMpConfig,
 )
 from .kernel_decorator import kernel, kernel_config
-from .kernelcreation import create_kernel
+from .kernelcreation import create_kernel, create_staggered_kernel
 from .backend.kernelfunction import KernelFunction
 from .slicing import make_slice
 from .spatial_coordinates import (
@@ -28,12 +28,13 @@ from .spatial_coordinates import (
     z_,
     z_staggered,
 )
-from .sympyextensions import Assignment, AssignmentCollection, AddAugmentedAssignment
-from .sympyextensions.astnodes import assignment_from_stencil
+from .assignment import Assignment, AddAugmentedAssignment, assignment_from_stencil
+from .simp import AssignmentCollection
 from .sympyextensions.typed_sympy import TypedSymbol
-from .sympyextensions.math import SymbolCreator
+from .sympyextensions import SymbolCreator
 from .datahandling import create_data_handling
 
+
 __all__ = [
     "Field",
     "FieldType",
@@ -48,6 +49,7 @@ __all__ = [
     "VectorizationConfig",
     "OpenMpConfig",
     "create_kernel",
+    "create_staggered_kernel",
     "KernelFunction",
     "Target",
     "show_code",
@@ -75,7 +77,5 @@ __all__ = [
     "stencil",
 ]
 
-from ._version import get_versions
-
-__version__ = get_versions()["version"]
-del get_versions
+from . import _version
+__version__ = _version.get_versions()['version']

src/pystencils/assignment.py

+import numpy as np
+import sympy as sp
+from sympy.codegen.ast import Assignment, AugmentedAssignment
+from sympy.codegen.ast import AddAugmentedAssignment as SpAddAugAssignment
+from sympy.printing.latex import LatexPrinter
+
+__all__ = ['Assignment', 'AugmentedAssignment', 'AddAugmentedAssignment', 'assignment_from_stencil']
+
+
+def print_assignment_latex(printer, expr):
+    binop = f"{expr.binop}=" if isinstance(expr, AugmentedAssignment) else ''
+    """sympy cannot print Assignments as Latex. Thus, this function is added to the sympy Latex printer"""
+    printed_lhs = printer.doprint(expr.lhs)
+    printed_rhs = printer.doprint(expr.rhs)
+    return fr"{printed_lhs} \leftarrow_{{{binop}}} {printed_rhs}"
+
+
+def assignment_str(assignment):
+    op = f"{assignment.binop}=" if isinstance(assignment, AugmentedAssignment) else '←'
+    return fr"{assignment.lhs} {op} {assignment.rhs}"
+
+
+_old_new = sp.codegen.ast.Assignment.__new__
+
+
+# TODO Typing Part2 add default type, defult_float_type, default_int_type and use sane defaults
+def _Assignment__new__(cls, lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, (list, tuple, sp.Matrix)) and isinstance(rhs, (list, tuple, sp.Matrix)):
+        assert len(lhs) == len(rhs), f'{lhs} and {rhs} must have same length when performing vector assignment!'
+        return tuple(_old_new(cls, a, b, *args, **kwargs) for a, b in zip(lhs, rhs))
+    return _old_new(cls, lhs, rhs, *args, **kwargs)
+
+
+Assignment.__str__ = assignment_str
+Assignment.__new__ = _Assignment__new__
+LatexPrinter._print_Assignment = print_assignment_latex
+
+AugmentedAssignment.__str__ = assignment_str
+LatexPrinter._print_AugmentedAssignment = print_assignment_latex
+
+sp.MutableDenseMatrix.__hash__ = lambda self: hash(tuple(self))
+
+#   Re-Export
+AddAugmentedAssignment = SpAddAugAssignment
+
+
+def assignment_from_stencil(stencil_array, input_field, output_field,
+                            normalization_factor=None, order='visual') -> Assignment:
+    """Creates an assignment
+
+    Args:
+        stencil_array: nested list of numpy array defining the stencil weights
+        input_field: field or field access, defining where the stencil should be applied to
+        output_field: field or field access where the result is written to
+        normalization_factor: optional normalization factor for the stencil
+        order: defines how the stencil_array is interpreted. Possible values are 'visual' and 'numpy'.
+               For details see examples
+
+    Returns:
+        Assignment that can be used to create a kernel
+
+    Examples:
+        >>> import pystencils as ps
+        >>> f, g = ps.fields("f, g: [2D]")
+        >>> stencil = [[0, 2, 0],
+        ...            [3, 4, 5],
+        ...            [0, 6, 0]]
+
+        By default 'visual ordering is used - i.e. the stencil is applied as the nested lists are written down
+        >>> expected_output = Assignment(g[0, 0], 3*f[-1, 0] + 6*f[0, -1] + 4*f[0, 0] + 2*f[0, 1] + 5*f[1, 0])
+        >>> assignment_from_stencil(stencil, f, g, order='visual') == expected_output
+        True
+
+        'numpy' ordering uses the first coordinate of the stencil array for x offset, second for y offset etc.
+        >>> expected_output = Assignment(g[0, 0], 2*f[-1, 0] + 3*f[0, -1] + 4*f[0, 0] + 5*f[0, 1] + 6*f[1, 0])
+        >>> assignment_from_stencil(stencil, f, g, order='numpy') == expected_output
+        True
+
+        You can also pass field accesses to apply the stencil at an already shifted position:
+        >>> expected_output = Assignment(g[2, 0], 3*f[0, 0] + 6*f[1, -1] + 4*f[1, 0] + 2*f[1, 1] + 5*f[2, 0])
+        >>> assignment_from_stencil(stencil, f[1, 0], g[2, 0]) == expected_output
+        True
+    """
+    from pystencils.field import Field
+
+    stencil_array = np.array(stencil_array)
+    if order == 'visual':
+        stencil_array = np.swapaxes(stencil_array, 0, 1)
+        stencil_array = np.flip(stencil_array, axis=1)
+    elif order == 'numpy':
+        pass
+    else:
+        raise ValueError("'order' has to be either 'visual' or 'numpy'")
+
+    if isinstance(input_field, Field):
+        input_field = input_field.center
+    if isinstance(output_field, Field):
+        output_field = output_field.center
+
+    rhs = 0
+    offset = tuple(s // 2 for s in stencil_array.shape)
+
+    for index, factor in np.ndenumerate(stencil_array):
+        shift = tuple(i - o for i, o in zip(index, offset))
+        rhs += factor * input_field.get_shifted(*shift)
+
+    if normalization_factor:
+        rhs *= normalization_factor
+
+    return Assignment(output_field, rhs)

src/pystencils/backend/kernelcreation/analysis.py

@@ -9,7 +9,8 @@ import sympy as sp
 from .context import KernelCreationContext
 
 from ...field import Field
-from ...sympyextensions import Assignment, AssignmentCollection
+from ...assignment import Assignment
+from ...simp import AssignmentCollection
 
 from ..exceptions import PsInternalCompilerError, KernelConstraintsError
 

src/pystencils/backend/kernelcreation/freeze.py

@@ -7,7 +7,8 @@ import sympy.core.relational
 import sympy.logic.boolalg
 from sympy.codegen.ast import AssignmentBase, AugmentedAssignment
 
-from ...sympyextensions.astnodes import Assignment, AssignmentCollection
+from ...assignment import Assignment
+from ...simp import AssignmentCollection
 from ...sympyextensions import (
     integer_functions,
     ConditionalFieldAccess,

src/pystencils/backend/kernelcreation/iteration_space.py

@@ -6,7 +6,7 @@ from functools import reduce
 from operator import mul
 
 from ...defaults import DEFAULTS
-from ...sympyextensions import AssignmentCollection
+from ...simp import AssignmentCollection
 from ...field import Field, FieldType
 
 from ..symbols import PsSymbol

src/pystencils/boundaries/boundaryconditions.py

 from typing import Any, List, Tuple, Sequence
 
-from pystencils.sympyextensions import Assignment
+from pystencils.assignment import Assignment
 from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
 from pystencils.types import create_type
 

src/pystencils/boundaries/boundaryhandling.py

@@ -4,7 +4,7 @@ import numpy as np
 import sympy as sp
 
 from pystencils import create_kernel, CreateKernelConfig, Target
-from pystencils.sympyextensions import Assignment
+from pystencils.assignment import Assignment
 from pystencils.boundaries.createindexlist import (
     create_boundary_index_array, numpy_data_type_for_boundary_object)
 from pystencils.sympyextensions import TypedSymbol

src/pystencils/fd/derivation.py

@@ -6,7 +6,7 @@ import sympy as sp
 
 from pystencils.field import Field
 from pystencils.stencil import direction_string_to_offset
-from pystencils.sympyextensions.math import multidimensional_sum, prod
+from pystencils.sympyextensions import multidimensional_sum, prod
 from pystencils.utils import LinearEquationSystem, fully_contains
 
 

src/pystencils/fd/derivative.py

@@ -3,7 +3,7 @@ from collections import defaultdict, namedtuple
 import sympy as sp
 
 from pystencils.field import Field
-from pystencils.sympyextensions.math import normalize_product, prod
+from pystencils.sympyextensions import normalize_product, prod
 
 
 def _default_diff_sort_key(d):

src/pystencils/fd/finitedifferences.py

@@ -7,8 +7,8 @@ from pystencils.fd import Diff
 from pystencils.fd.derivative import diff_args
 from pystencils.fd.spatial import fd_stencils_standard
 from pystencils.field import Field
-from pystencils.sympyextensions import AssignmentCollection
-from pystencils.sympyextensions.math import fast_subs
+from pystencils.simp import AssignmentCollection
+from pystencils.sympyextensions import fast_subs
 
 FieldOrFieldAccess = Union[Field, Field.Access]
 

src/pystencils/field.py

@@ -15,8 +15,8 @@ from pystencils.alignedarray import aligned_empty
 from pystencils.spatial_coordinates import x_staggered_vector, x_vector
 from pystencils.stencil import direction_string_to_offset, inverse_direction, offset_to_direction_string
 from pystencils.types import PsType, PsStructType, create_type
-from pystencils.sympyextensions.typed_sympy import (FieldShapeSymbol, FieldStrideSymbol, TypedSymbol)
-from pystencils.sympyextensions.math import is_integer_sequence
+from pystencils.sympyextensions.typed_sympy import FieldShapeSymbol, FieldStrideSymbol, TypedSymbol
+from pystencils.sympyextensions import is_integer_sequence
 from pystencils.types import UserTypeSpec
 
 

src/pystencils/kernel_decorator.py

@@ -5,8 +5,8 @@ from typing import Callable, Union, List, Dict, Tuple
 
 import sympy as sp
 
-from .sympyextensions import Assignment
-from .sympyextensions.math import SymbolCreator
+from .assignment import Assignment
+from .sympyextensions import SymbolCreator
 from pystencils.config import CreateKernelConfig
 
 __all__ = ['kernel', 'kernel_config']

src/pystencils/kernelcreation.py

@@ -32,7 +32,8 @@ from .backend.transformations import (
     SelectFunctions,
 )
 
-from .sympyextensions import AssignmentCollection, Assignment
+from .simp import AssignmentCollection
+from .assignment import Assignment
 
 
 __all__ = ["create_kernel"]
@@ -152,3 +153,7 @@ def create_kernel_function(
     return KernelFunction(
         body, target_spec, function_name, params, req_headers, ctx.constraints, jit
     )
+
+
+def create_staggered_kernel(assignments, target: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs):
+    raise NotImplementedError("Staggered kernels are not yet implemented for pystencils 2.0")

src/pystencils/placeholder_function.py

@@ -2,7 +2,7 @@ from typing import List
 
 import sympy as sp
 
-from pystencils.sympyextensions import Assignment
+from .assignment import Assignment
 from pystencils.sympyextensions import is_constant
 from pystencils.sympyextensions.astnodes import generic_visit
 

src/pystencils/sympyextensions/rng.py → src/pystencils/rng.py

@@ -2,10 +2,9 @@ import copy
 import numpy as np
 import sympy as sp
 
-from pystencils.sympyextensions import TypedSymbol, CastFunc
+from .sympyextensions import TypedSymbol, CastFunc, fast_subs
 # from pystencils.sympyextensions.astnodes import LoopOverCoordinate # TODO nbackend: replace
 # from pystencils.backends.cbackend import CustomCodeNode # TODO nbackend: replace
-from pystencils.sympyextensions import fast_subs
 
 
 # class RNGBase(CustomCodeNode): TODO nbackend: replace

src/pystencils/simp/__init__.py

+from .assignment_collection import AssignmentCollection
+from .simplifications import (
+    add_subexpressions_for_constants,
+    add_subexpressions_for_divisions,
+    add_subexpressions_for_field_reads,
+    add_subexpressions_for_sums,
+    apply_on_all_subexpressions,
+    apply_to_all_assignments,
+    subexpression_substitution_in_existing_subexpressions,
+    subexpression_substitution_in_main_assignments,
+    sympy_cse,
+    sympy_cse_on_assignment_list,
+)
+from .subexpression_insertion import (
+    insert_aliases,
+    insert_zeros,
+    insert_constants,
+    insert_constant_additions,
+    insert_constant_multiples,
+    insert_squares,
+    insert_symbol_times_minus_one,
+)
+from .simplificationstrategy import SimplificationStrategy
+
+__all__ = [
+    "AssignmentCollection",
+    "SimplificationStrategy",
+    "sympy_cse",
+    "sympy_cse_on_assignment_list",
+    "apply_to_all_assignments",
+    "apply_on_all_subexpressions",
+    "subexpression_substitution_in_existing_subexpressions",
+    "subexpression_substitution_in_main_assignments",
+    "add_subexpressions_for_constants",
+    "add_subexpressions_for_divisions",
+    "add_subexpressions_for_sums",
+    "add_subexpressions_for_field_reads",
+    "insert_aliases",
+    "insert_zeros",
+    "insert_constants",
+    "insert_constant_additions",
+    "insert_constant_multiples",
+    "insert_squares",
+    "insert_symbol_times_minus_one",
+]