Demo: Assignment collections and simplification

Assignment collections

The assignment collection class helps to formulate and simplify assignments for numerical kernels.

An AssignmentCollection is an ordered collection of assignments, together with an optional ordered collection of subexpressions, that are required to evaluate the main assignments. There are various simplification rules available that operate on AssignmentCollections.

We start by defining some stencil update rule. Here we also use the pystencils Field, note however that the assignment collection module works purely on the sympy level.

a,b,c = sp.symbols("a b c")
f = ps.fields("f(2) : [2D]")

a1 = ps.Assignment(f[0,0](1), (a**2 +b) * f[0,1] + \
                  (a**2 - c) * f[1,0] + \
                  (a**2 - 2*c) * f[-1,0] + \
                  (a**2) * f[0, -1])

a2 = ps.Assignment(f[0,0](0), (c**2 +b) * f[0,1] + \
                  (c**2 - c) * f[1,0] + \
                  (c**2 - 2*c) * f[-1,0] + \
                  (c**2 - a**2) * f[0, -1])


ac = ps.AssignmentCollection([a1, a2], subexpressions=[])
ac

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \left(a^{2} - c\right) + {{f}_{N}^{0}} \left(a^{2} + b\right) + {{f}_{S}^{0}} a^{2} + {{f}_{W}^{0}} \left(a^{2} - 2 c\right)$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \left(c^{2} - c\right) + {{f}_{N}^{0}} \left(b + c^{2}\right) + {{f}_{S}^{0}} \left(- a^{2} + c^{2}\right) + {{f}_{W}^{0}} \left(c^{2} - 2 c\right)$$

sympy operations can be applied on an assignment collection: In this example we first expand the collection, then look for common subexpressions.

expand_all = ps.simp.apply_to_all_assignments(sp.expand)
expandedEc = expand_all(ac)

ac_cse = ps.simp.sympy_cse(expandedEc)
ac_cse

Subexpressions:

$$\xi_{0} \leftarrow - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b - 2 {{f}_{W}^{0}} c$$

$$\xi_{1} \leftarrow a^{2}$$

$$\xi_{2} \leftarrow {{f}_{S}^{0}} \xi_{1}$$

$$\xi_{3} \leftarrow c^{2}$$

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \xi_{3} + {{f}_{N}^{0}} \xi_{3} + {{f}_{S}^{0}} \xi_{3} + {{f}_{W}^{0}} \xi_{3} + \xi_{0} - \xi_{2}$$

Symbols occuring in assignment collections are classified into 3 categories: - free_symbols: symbols that occur in right-hand-sides but never on left-hand-sides - bound_symbols: symbols that occur on left-hand-sides - defined_symbols: symbols that occur on left-hand-sides of a main assignment

ac_cse.free_symbols

$$\left\{{{f}_{E}^{0}}, {{f}_{N}^{0}}, {{f}_{S}^{0}}, {{f}_{W}^{0}}, a, b, c\right\}$$

ac_cse.bound_symbols

$$\left\{{{f}_{C}^{0}}, {{f}_{C}^{1}}, \xi_{0}, \xi_{1}, \xi_{2}, \xi_{3}\right\}$$

ac_cse.defined_symbols

$$\left\{{{f}_{C}^{0}}, {{f}_{C}^{1}}\right\}$$

Assignment collections can be splitted up, and merged together. For splitting, a list of symbols that occur on the left-hand-side in the main assignments has to be passed. The returned assignment collection only contains these main assignments together with all necessary subexpressions.

ac_f0 = ac_cse.new_filtered([f(0)])
ac_f1 = ac_cse.new_filtered([f(1)])
ac_f1

Subexpressions:

$$\xi_{0} \leftarrow - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b - 2 {{f}_{W}^{0}} c$$

$$\xi_{1} \leftarrow a^{2}$$

$$\xi_{2} \leftarrow {{f}_{S}^{0}} \xi_{1}$$

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$

Note here that \(\xi_4\) is no longer part of the subexpressions, since it is not used in the main assignment of \(f_C^1\).

If we merge both collections together, we end up with the original collection.

ac_f0.new_merged(ac_f1)

Subexpressions:

$$\xi_{0} \leftarrow - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b - 2 {{f}_{W}^{0}} c$$

$$\xi_{1} \leftarrow a^{2}$$

$$\xi_{2} \leftarrow {{f}_{S}^{0}} \xi_{1}$$

$$\xi_{3} \leftarrow c^{2}$$

Main Assignments:

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \xi_{3} + {{f}_{N}^{0}} \xi_{3} + {{f}_{S}^{0}} \xi_{3} + {{f}_{W}^{0}} \xi_{3} + \xi_{0} - \xi_{2}$$

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$

There is also a method that inserts all subexpressions into the main assignments. This is the inverse operation of common subexpression elimination.

assert sp.simplify(ac_f0.new_without_subexpressions().main_assignments[0].rhs - a2.rhs) == 0
ac_f0.new_without_subexpressions()

Main Assignments:

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} c^{2} - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b + {{f}_{N}^{0}} c^{2} - {{f}_{S}^{0}} a^{2} + {{f}_{S}^{0}} c^{2} + {{f}_{W}^{0}} c^{2} - 2 {{f}_{W}^{0}} c$$

To evaluate an assignment collection, use the lambdify method. It is very similar to sympys lambdify function.

evalFct = ac_cse.lambdify([f[0,1], f[1,0]],  # new parameters of returned function
                          fixed_symbols={a:1, b:2, c:3, f[0,-1]: 4, f[-1,0]: 5}) # fix values of other symbols
evalFct(2,1)

$$\left \{ {{f}_{C}^{0}} : 75, \quad {{f}_{C}^{1}} : -17\right \}$$

lambdify is rather slow for evaluation. The intended way to evaluate an assignment collection is pystencils i.e. create a fast kernel, that applies the update at every site of a structured grid. The collection can be directly passed to the create_kernel function.

func = ps.create_kernel(ac_cse).compile()

Simplification Strategies

In above examples, we already applied simplification rules to assignment collections. Simplification rules are functions that take, as a single argument, an assignment collection and return an modified/simplified copy of it. The SimplificationStrategy class holds a list of simplification rules and can apply all of them in the specified order. Additionally it provides useful printing and reporting functions.

We start by creating a simplification strategy, consisting of the expand and CSE simplifications we have already applied above:

strategy = ps.simp.SimplificationStrategy()
strategy.add(ps.simp.apply_to_all_assignments(sp.expand))
strategy.add(ps.simp.sympy_cse)

This strategy can be applied to any assignment collection:

strategy(ac)

Subexpressions:

$$\xi_{0} \leftarrow - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b - 2 {{f}_{W}^{0}} c$$

$$\xi_{1} \leftarrow a^{2}$$

$$\xi_{2} \leftarrow {{f}_{S}^{0}} \xi_{1}$$

$$\xi_{3} \leftarrow c^{2}$$

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \xi_{3} + {{f}_{N}^{0}} \xi_{3} + {{f}_{S}^{0}} \xi_{3} + {{f}_{W}^{0}} \xi_{3} + \xi_{0} - \xi_{2}$$

The strategy can also print the simplification results at each stage. The report contains information about the number of operations after each simplification as well as the runtime of each simplification routine.

strategy.create_simplification_report(ac)

Name	Runtime	Adds	Muls	Divs	Total
OriginalTerm	-	13	19	0	32
expand	0.11 ms	13	26	0	39
sympy_cse	3.25 ms	11	14	0	25

The strategy can also print the full collection after each simplification…

strategy.show_intermediate_results(ac)

Initial Version

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \left(a^{2} - c\right) + {{f}_{N}^{0}} \left(a^{2} + b\right) + {{f}_{S}^{0}} a^{2} + {{f}_{W}^{0}} \left(a^{2} - 2 c\right)$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \left(c^{2} - c\right) + {{f}_{N}^{0}} \left(b + c^{2}\right) + {{f}_{S}^{0}} \left(- a^{2} + c^{2}\right) + {{f}_{W}^{0}} \left(c^{2} - 2 c\right)$$

expand

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} a^{2} - {{f}_{E}^{0}} c + {{f}_{N}^{0}} a^{2} + {{f}_{N}^{0}} b + {{f}_{S}^{0}} a^{2} + {{f}_{W}^{0}} a^{2} - 2 {{f}_{W}^{0}} c$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} c^{2} - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b + {{f}_{N}^{0}} c^{2} - {{f}_{S}^{0}} a^{2} + {{f}_{S}^{0}} c^{2} + {{f}_{W}^{0}} c^{2} - 2 {{f}_{W}^{0}} c$$

sympy_cse

Subexpressions:

$$\xi_{0} \leftarrow - {{f}_{E}^{0}} c + {{f}_{N}^{0}} b - 2 {{f}_{W}^{0}} c$$

$$\xi_{1} \leftarrow a^{2}$$

$$\xi_{2} \leftarrow {{f}_{S}^{0}} \xi_{1}$$

$$\xi_{3} \leftarrow c^{2}$$

Main Assignments:

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$

$${{f}_{C}^{0}} \leftarrow {{f}_{E}^{0}} \xi_{3} + {{f}_{N}^{0}} \xi_{3} + {{f}_{S}^{0}} \xi_{3} + {{f}_{W}^{0}} \xi_{3} + \xi_{0} - \xi_{2}$$

… or only specific assignments for better readability

strategy.show_intermediate_results(ac, symbols=[f(1)])

Initial Version

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \left(a^{2} - c\right) + {{f}_{N}^{0}} \left(a^{2} + b\right) + {{f}_{S}^{0}} a^{2} + {{f}_{W}^{0}} \left(a^{2} - 2 c\right)$$

expand

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} a^{2} - {{f}_{E}^{0}} c + {{f}_{N}^{0}} a^{2} + {{f}_{N}^{0}} b + {{f}_{S}^{0}} a^{2} + {{f}_{W}^{0}} a^{2} - 2 {{f}_{W}^{0}} c$$

sympy_cse

$${{f}_{C}^{1}} \leftarrow {{f}_{E}^{0}} \xi_{1} + {{f}_{N}^{0}} \xi_{1} + {{f}_{W}^{0}} \xi_{1} + \xi_{0} + \xi_{2}$$