diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py
index ed03bf7627040093167a175ec0229a17a484fe44..d2826fa850a6dbf06db192bf45509147633bfc93 100644
--- a/pystencils/astnodes.py
+++ b/pystencils/astnodes.py
@@ -294,6 +294,8 @@ class SkipIteration(Node):
class Block(Node):
def __init__(self, nodes: List[Node]):
super(Block, self).__init__()
+ if not isinstance(nodes, list):
+ nodes = [nodes]
self._nodes = nodes
self.parent = None
for n in self._nodes:
diff --git a/pystencils/backends/cbackend.py b/pystencils/backends/cbackend.py
index 8437bdb6801ecc5adc717dec2983e5d6eda0baf0..a8a3dbc8d467585ba71570d6c0df7d31a8738dcf 100644
--- a/pystencils/backends/cbackend.py
+++ b/pystencils/backends/cbackend.py
@@ -125,7 +125,7 @@ def get_headers(ast_node: Node) -> Set[str]:
# --------------------------------------- Backend Specific Nodes -------------------------------------------------------
-
+# TODO CustomCodeNode should not be backend specific
class CustomCodeNode(Node):
def __init__(self, code, symbols_read, symbols_defined, parent=None):
super(CustomCodeNode, self).__init__(parent=parent)
diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py
index c63946bf0b05dd321a5c1da1672369932c80fe1e..50eaedf4f51eea2b969673d8011bd4c6a2dfa062 100644
--- a/pystencils/kernelcreation.py
+++ b/pystencils/kernelcreation.py
@@ -19,7 +19,7 @@ from pystencils.transformations import (
loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel)
-def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCollection, List[Node]], *,
+def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCollection, List[Node], NodeCollection], *,
config: CreateKernelConfig = None, **kwargs):
"""
Creates abstract syntax tree (AST) of kernel, using a list of update equations.
diff --git a/pystencils/node_collection.py b/pystencils/node_collection.py
index 68f53f36c578fcb470cb18122b08300a89f1fe8b..0287b6fc8dbec76f03283b032d47f77eacb45516 100644
--- a/pystencils/node_collection.py
+++ b/pystencils/node_collection.py
@@ -1,12 +1,13 @@
import logging
from typing import List, Union
+import sympy
import sympy as sp
from sympy.codegen import Assignment
from sympy.codegen.rewriting import ReplaceOptim, optimize
-from pystencils.astnodes import Node
-
+from pystencils.astnodes import Block, Node
+from pystencils.backends.cbackend import CustomCodeNode
from pystencils.functions import DivFunc
@@ -28,11 +29,6 @@ class NodeCollection:
self.simplification_hints = ()
def evaluate_terms(self):
-
- # There is no visitor implemented now so working with nodes does not work
- if self.is_Nodes:
- return
-
evaluate_constant_terms = ReplaceOptim(
lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
lambda p: p.evalf())
@@ -43,8 +39,22 @@ class NodeCollection:
sp.UnevaluatedExpr(sp.Mul(*([p.base] * +p.exp), evaluate=False)) if p.exp > 0 else
DivFunc(sp.Integer(1), sp.Mul(*([p.base] * -p.exp), evaluate=False))
))
-
sympy_optimisations = [evaluate_constant_terms, evaluate_pow]
- self.all_assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
- if hasattr(a, 'lhs')
- else a for a in self.all_assignments]
+
+ if self.is_Nodes:
+ def visitor(node):
+ if isinstance(node, CustomCodeNode):
+ return node
+ elif isinstance(node, Block):
+ return node.func([visitor(child) for child in node.args])
+ elif isinstance(node, Node):
+ return node.func(*[visitor(child) for child in node.args])
+ elif isinstance(node, sympy.Basic):
+ return optimize(node, sympy_optimisations)
+ else:
+ raise NotImplementedError(f'{node} {type(node)} has no valid visitor')
+ self.all_assignments = [visitor(assignment) for assignment in self.all_assignments]
+ else:
+ self.all_assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
+ if hasattr(a, 'lhs')
+ else a for a in self.all_assignments]
diff --git a/pystencils_tests/test_fvm.py b/pystencils_tests/test_fvm.py
index 55ce1f7db2e859978a4afb5bb3a763ca5e7d2d56..c0e22f8bd497620863b185a460780eb6af95bb27 100644
--- a/pystencils_tests/test_fvm.py
+++ b/pystencils_tests/test_fvm.py
@@ -4,6 +4,8 @@ import numpy as np
import pytest
from itertools import product
from pystencils.rng import random_symbol
+from pystencils.astnodes import SympyAssignment
+from pystencils.node_collection import NodeCollection
def advection_diffusion(dim: int):
@@ -240,12 +242,10 @@ def advection_diffusion_fluctuations(dim: int):
advection_diffusion_fluctuations.runners = {}
-# @pytest.mark.parametrize("velocity", list(product([0, 0.00041], [0, -0.00031])))
-# @pytest.mark.parametrize("density", [27.0, 56.5])
-# @pytest.mark.longrun
-def test_advection_diffusion_fluctuation_2d():
- density = 27.0
- velocity = [0, 0.00041]
+@pytest.mark.parametrize("velocity", list(product([0, 0.00041], [0, -0.00031])))
+@pytest.mark.parametrize("density", [27.0, 56.5])
+@pytest.mark.longrun
+def test_advection_diffusion_fluctuation_2d(density, velocity):
if 2 not in advection_diffusion_fluctuations.runners:
advection_diffusion_fluctuations.runners[2] = advection_diffusion_fluctuations(2)
advection_diffusion_fluctuations.runners[2](density, velocity)
@@ -318,33 +318,38 @@ def diffusion_reaction(fluctuations: bool):
# add fluctuations
fluct *= 2 * (next(rng_symbol_gen) - 0.5) * sp.sqrt(3)
- flux.main_assignments[i] = ps.Assignment(flux.main_assignments[i].lhs, flux.main_assignments[i].rhs + fluct)
+ flux.main_assignments[i] = ps.Assignment(flux.main_assignments[i].lhs,
+ flux.main_assignments[i].rhs + fluct)
# Add the folding to the flux, so that the random numbers persist through the ghostlayers.
fold = {ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
flux.subs(fold)
- r_flux = ps.AssignmentCollection([ps.Assignment(j_fields[i].center, 0) for i in range(species)])
+ r_flux = NodeCollection([SympyAssignment(j_fields[i].center, 0) for i in range(species)])
reaction = r_rate_const
for i in range(species):
reaction *= sp.Pow(n_fields[i].center, r_order[i])
- if(fluctuations):
- rng_symbol_gen = random_symbol(r_flux.subexpressions, dim=dh.dim)
+ new_assignments = []
+ if fluctuations:
+ rng_symbol_gen = random_symbol(new_assignments, dim=dh.dim)
reaction_fluctuations = sp.sqrt(sp.Abs(reaction)) * 2 * (next(rng_symbol_gen) - 0.5) * sp.sqrt(3)
reaction_fluctuations *= sp.Min(1, sp.Abs(reaction**2))
else:
reaction_fluctuations = 0.0
for i in range(species):
- r_flux.main_assignments[i] = ps.Assignment(
+ r_flux.all_assignments[i] = SympyAssignment(
r_flux_fields[i].center, (reaction + reaction_fluctuations) * r_coefs[i])
+ [r_flux.all_assignments.insert(0, new) for new in new_assignments]
- continuity_assignments.append(ps.Assignment(n_fields[0].center, n_fields[0].center + r_flux_fields[0].center))
+ continuity_assignments = [SympyAssignment(*assignment.args) for assignment in continuity_assignments]
+ continuity_assignments.append(SympyAssignment(n_fields[0].center, n_fields[0].center + r_flux_fields[0].center))
flux_kernel = ps.create_staggered_kernel(flux).compile()
reaction_kernel = ps.create_kernel(r_flux).compile()
- pde_kernel = ps.create_kernel(continuity_assignments).compile()
+ config = ps.CreateKernelConfig(allow_double_writes=True)
+ pde_kernel = ps.create_kernel(continuity_assignments, config=config).compile()
sync_conc = dh.synchronization_function([n_fields[0].name, n_fields[1].name])
@@ -414,7 +419,7 @@ advection_diffusion_fluctuations.runners = {}
@pytest.mark.parametrize("density", [27.0, 56.5])
@pytest.mark.parametrize("fluctuations", [False, True])
@pytest.mark.longrun
-def test_diffusion_reaction(density, velocity, fluctuations):
+def test_diffusion_reaction(velocity, density, fluctuations):
diffusion_reaction.runner = diffusion_reaction(fluctuations)
diffusion_reaction.runner(density, velocity)