Entropic to new lbmpy

backends/cbackend.py

 import textwrap
 from sympy.utilities.codegen import CCodePrinter
-from pystencils.ast import Node
+from pystencils.astnodes import Node
 
 
 def generateC(astNode):

backends/dot.py

@@ -37,7 +37,7 @@ class DotPrinter(Printer):
 
 
 def __shortened(node):
-    from pystencils.ast import LoopOverCoordinate, KernelFunction, SympyAssignment
+    from pystencils.astnodes import LoopOverCoordinate, KernelFunction, SympyAssignment
     if isinstance(node, LoopOverCoordinate):
         return "Loop over dim %d" % (node.coordinateToLoopOver,)
     elif isinstance(node, KernelFunction):

cpu/kerncraft.py

 from pystencils.transformations import makeLoopOverDomain, typingFromSympyInspection, \
     typeAllEquations, moveConstantsBeforeLoop, getOptimalLoopOrdering
-import pystencils.ast as ast
+import pystencils.astnodes as ast
 from pystencils.backends.cbackend import CBackend, CustomSympyPrinter
 from pystencils import TypedSymbol
 

cpu/kernelcreation.py

@@ -3,7 +3,7 @@ from pystencils.transformations import resolveFieldAccesses, makeLoopOverDomain,
     typeAllEquations, getOptimalLoopOrdering, parseBasePointerInfo, moveConstantsBeforeLoop, splitInnerLoop
 from pystencils.types import TypedSymbol, DataType
 from pystencils.field import Field
-import pystencils.ast as ast
+import pystencils.astnodes as ast
 
 
 def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(),

display_utils.py

@@ -2,7 +2,7 @@
 
 def toDot(expr, graphStyle={}):
     """Show a sympy or pystencils AST as dot graph"""
-    from pystencils.ast import Node
+    from pystencils.astnodes import Node
     import graphviz
     if isinstance(expr, Node):
         from pystencils.backends.dot import dotprint

gpucuda/kernelcreation.py

@@ -2,7 +2,7 @@ import sympy as sp
 
 from pystencils.transformations import resolveFieldAccesses, typeAllEquations, \
     parseBasePointerInfo, typingFromSympyInspection
-from pystencils.ast import Block, KernelFunction
+from pystencils.astnodes import Block, KernelFunction
 from pystencils import Field
 
 BLOCK_IDX = list(sp.symbols("blockIdx.x blockIdx.y blockIdx.z"))

llvm/kernelcreation.py

@@ -3,7 +3,7 @@ from pystencils.transformations import resolveFieldAccesses, makeLoopOverDomain,
     typeAllEquations, getOptimalLoopOrdering, parseBasePointerInfo, moveConstantsBeforeLoop, splitInnerLoop
 from pystencils.types import TypedSymbol, DataType
 from pystencils.field import Field
-import pystencils.ast as ast
+import pystencils.astnodes as ast
 
 
 def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(),

sympyextensions.py

-import sympy as sp
 import operator
 from collections import defaultdict, Sequence
 import warnings
+import sympy as sp
 
 
 def fastSubs(term, subsDict):
@@ -155,9 +155,13 @@ def replaceSecondOrderProducts(expr, searchSymbols, positive=None, replaceMixed=
 
 def removeHigherOrderTerms(term, order=3, symbols=None):
     """
-    Remove all terms from a sum that contain 'order' or more factors of given 'symbols'
-    Example: symbols = ['u_x', 'u_y'] and order =2
-             removes terms u_x**2, u_x*u_y, u_y**2, u_x**3, ....
+    Removes all terms that that contain more than 'order' factors of given 'symbols'
+
+    Example:
+        >>> x, y = sp.symbols("x y")
+        >>> term = x**2 * y + y**2 * x + y**3 + x + y ** 2
+        >>> removeHigherOrderTerms(term, order=2, symbols=[x, y])
+        x + y**2
     """
     from sympy.core.power import Pow
     from sympy.core.add import Add, Mul
@@ -171,15 +175,19 @@ def removeHigherOrderTerms(term, order=3, symbols=None):
 
     def velocityFactorsInProduct(product):
         uFactorCount = 0
-        for factor in product.args:
-            if type(factor) == Pow:
-                if factor.args[0] in symbols:
-                    uFactorCount += factor.args[1]
-            if factor in symbols:
-                uFactorCount += 1
+        if type(product) is Mul:
+            for factor in product.args:
+                if type(factor) == Pow:
+                    if factor.args[0] in symbols:
+                        uFactorCount += factor.args[1]
+                if factor in symbols:
+                    uFactorCount += 1
+        elif type(product) is Pow:
+            if product.args[0] in symbols:
+                uFactorCount += product.args[1]
         return uFactorCount
 
-    if type(term) == Mul:
+    if type(term) == Mul or type(term) == Pow:
         if velocityFactorsInProduct(term) <= order:
             return term
         else:

transformations.py

@@ -6,7 +6,7 @@ from sympy.tensor import IndexedBase
 from pystencils.field import Field, offsetComponentToDirectionString
 from pystencils.types import TypedSymbol, DataType
 from pystencils.slicing import normalizeSlice
-import pystencils.ast as ast
+import pystencils.astnodes as ast
 
 
 def fastSubs(term, subsDict):