simplificationfactory.py 1.85 KB
Newer Older
1
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
2

Martin Bauer's avatar
Martin Bauer committed
3
from lbmpy.innerloopsplit import create_lbm_split_groups
Frederik Hennig's avatar
Frederik Hennig committed
4
from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
Markus Holzer's avatar
Markus Holzer committed
5
6
from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
from lbmpy.methods.centeredcumulant.simplification import insert_aliases, insert_zeros
Frederik Hennig's avatar
Frederik Hennig committed
7
from lbmpy.methods.momentbased.momentbasedsimplifications import (
Martin Bauer's avatar
Martin Bauer committed
8
9
    factor_density_after_factoring_relaxation_times, factor_relaxation_rates,
    replace_common_quadratic_and_constant_term, replace_density_and_velocity, replace_second_order_velocity_products)
Martin Bauer's avatar
Martin Bauer committed
10
from pystencils.simp import (
Martin Bauer's avatar
Martin Bauer committed
11
12
    SimplificationStrategy, add_subexpressions_for_divisions, apply_to_all_assignments,
    subexpression_substitution_in_main_assignments)
13
14


Martin Bauer's avatar
Martin Bauer committed
15
def create_simplification_strategy(lb_method, split_inner_loop=False):
16
    s = SimplificationStrategy()
17
    expand = apply_to_all_assignments(sp.expand)
18

Martin Bauer's avatar
Martin Bauer committed
19
    if isinstance(lb_method, MomentBasedLbMethod):
Martin Bauer's avatar
Martin Bauer committed
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
        if len(set(lb_method.relaxation_rates)) <= 2:
            s.add(expand)
            s.add(replace_second_order_velocity_products)
            s.add(expand)
            s.add(factor_relaxation_rates)
            s.add(replace_density_and_velocity)
            s.add(replace_common_quadratic_and_constant_term)
            s.add(factor_density_after_factoring_relaxation_times)
            s.add(subexpression_substitution_in_main_assignments)
            if split_inner_loop:
                s.add(create_lbm_split_groups)
            s.add(add_subexpressions_for_divisions)
        else:
            s.add(subexpression_substitution_in_main_assignments)
            if split_inner_loop:
                s.add(create_lbm_split_groups)
Markus Holzer's avatar
Markus Holzer committed
36
37
38
39
    elif isinstance(lb_method, CenteredCumulantBasedLbMethod):
        s.add(insert_zeros)
        s.add(insert_aliases)
        s.add(lambda ac: ac.new_without_unused_subexpressions())
40
    return s