diff --git a/finitedifferences.py b/finitedifferences.py index b81a8fa48edad54fc0df2c490de34028a0130129..171260cdc1b7dc1852dd02e8dff7db4dadfa22aa 100644 --- a/finitedifferences.py +++ b/finitedifferences.py @@ -1,6 +1,7 @@ import numpy as np import sympy as sp from pystencils.field import Field +from pystencils.transformations import fastSubs def grad(var, dim=3): @@ -131,14 +132,3 @@ def __upDownOffsets(d, dim): coord[d] = -1 down = np.array(coord, dtype=np.int) return up, down - - -def fastSubs(term, subsDict): - """Similar to sympy subs function. - This version is much faster for big substitution dictionaries than sympy version""" - def visit(expr): - if expr in subsDict: - return subsDict[expr] - paramList = [visit(a) for a in expr.args] - return expr if not paramList else expr.func(*paramList) - return visit(term) diff --git a/transformations.py b/transformations.py index e2189ace62623e1d14a255ef029bcdef6bf8ad02..e2c2e22919f12cf98727263c1c1fb2ac74a8bf61 100644 --- a/transformations.py +++ b/transformations.py @@ -9,6 +9,17 @@ from pystencils.slicing import normalizeSlice import pystencils.ast as ast +def fastSubs(term, subsDict): + """Similar to sympy subs function. + This version is much faster for big substitution dictionaries than sympy version""" + def visit(expr): + if expr in subsDict: + return subsDict[expr] + paramList = [visit(a) for a in expr.args] + return expr if not paramList else expr.func(*paramList) + return visit(term) + + def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None, loopOrder=None): """ Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.