Skip to content
Snippets Groups Projects
Commit 1edb192f authored by Jan Hönig's avatar Jan Hönig
Browse files

Introduced a converting function form AssignmentCollection to NodeCollection.

parent 1cbb3d17
No related branches found
No related tags found
No related merge requests found
...@@ -6,9 +6,10 @@ import sympy as sp ...@@ -6,9 +6,10 @@ import sympy as sp
from sympy.codegen import Assignment from sympy.codegen import Assignment
from sympy.codegen.rewriting import ReplaceOptim, optimize from sympy.codegen.rewriting import ReplaceOptim, optimize
from pystencils.astnodes import Block, Node from pystencils.astnodes import Block, Node, SympyAssignment
from pystencils.backends.cbackend import CustomCodeNode from pystencils.backends.cbackend import CustomCodeNode
from pystencils.functions import DivFunc from pystencils.functions import DivFunc
from pystencils.simp import AssignmentCollection
class NodeCollection: class NodeCollection:
...@@ -28,6 +29,10 @@ class NodeCollection: ...@@ -28,6 +29,10 @@ class NodeCollection:
self.simplification_hints = {} self.simplification_hints = {}
@staticmethod
def from_assignment_collection(assignment_collection: AssignmentCollection):
return NodeCollection([SympyAssignment(a.lhs, a.rhs) for a in assignment_collection.all_assignments])
def evaluate_terms(self): def evaluate_terms(self):
evaluate_constant_terms = ReplaceOptim( evaluate_constant_terms = ReplaceOptim(
lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer, lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
......
import sympy as sp
from pystencils import AssignmentCollection, Assignment
from pystencils.node_collection import NodeCollection
from pystencils.astnodes import SympyAssignment
def test_node_collection_from_assignment_collection():
x = sp.symbols('x')
assignment_collection = AssignmentCollection([Assignment(x, 2)])
node_collection = NodeCollection.from_assignment_collection(assignment_collection)
assert node_collection.all_assignments[0] == SympyAssignment(x, 2)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment