def _split_backward_equations(self, backward_assignments, subexp_symgen):
all_subexpressions = []
split_main_assignments = []
known_coeffs_dict = dict()
for asm, stencil_dir in zip(backward_assignments, self.stencil):
split_asm = self._split_backward_equations_recursive(
asm, all_subexpressions, stencil_dir, subexp_symgen, known_coeffs_dict)
split_main_assignments.append(split_asm)
ac = AssignmentCollection(split_main_assignments, subexpressions=all_subexpressions,
subexpression_symbol_generator=subexp_symgen)
ac.topological_sort(sort_main_assignments=False)
return ac
def create_lb_update_rule_sparse(collision_rule, src, dst, idx, kernel_type='stream_pull_collide') -> AC:
"""Creates a update rule from a collision rule using compressed pdf storage and two (src/dst) arrays.
Args:
collision_rule: arbitrary collision rule, e.g. created with create_lb_collision_rule
src: symbolic field to read from
dst: symbolic field to write to
idx: symbolic index field
kernel_type: one of 'stream_pull_collide', 'collide_only' or 'stream_pull_only'
Returns:
update rule
"""
assert kernel_type in ('stream_pull_collide', 'collide_only', 'stream_pull_only')
method = collision_rule.method
q = len(method.stencil)
symbol_subs = _list_substitutions(method, src, idx)
if kernel_type == 'stream_pull_only':
assignments = []
for i in range(q):
lhs = dst(i)
rhs = symbol_subs[method.pre_collision_pdf_symbols[i]]
if lhs - rhs != 0:
assignments.append(Assignment(lhs, rhs))
return AssignmentCollection(assignments, subexpressions=[])
else:
write_target = src if kernel_type == 'collide_only' else dst
symbol_subs.update({sym: write_target(i) for i, sym in enumerate(method.post_collision_pdf_symbols)})
return collision_rule.new_with_substitutions(symbol_subs)
def test_complex_execution(dtype, target, with_complex_argument):
complex_dtype = f'complex{64 if dtype ==np.float32 else 128}'
x, y = pystencils.fields(f'x, y: {complex_dtype}[2d]')
x_arr = np.zeros((20, 30), complex_dtype)
y_arr = np.zeros((20, 30), complex_dtype)
if with_complex_argument:
a = pystencils.TypedSymbol('a', create_type(complex_dtype))
else:
a = (2j + 1)
assignments = AssignmentCollection({y.center: x.center + a})
if target == pystencils.Target.GPU:
pytest.importorskip('pycuda')
from pycuda.gpuarray import zeros
x_arr = zeros((20, 30), complex_dtype)
y_arr = zeros((20, 30), complex_dtype)
kernel = pystencils.create_kernel(assignments,
target=target,
data_type=dtype).compile()
if with_complex_argument:
kernel(x=x_arr, y=y_arr, a=2j + 1)
else:
kernel(x=x_arr, y=y_arr)
if target == pystencils.Target.GPU:
y_arr = y_arr.get()
assert np.allclose(y_arr, 2j + 1)
def equilibrium_input_equations_from_init_values(self,
density=1,
velocity=(0, 0, 0),
force_substitution=True):
dim = self._stencil.D
zeroth_order_moment = density if self._compressible else density - sp.Rational(
1, 1)
first_order_moments = velocity[:dim]
vel_offset = [0] * dim
if self._compressible:
if self._forceModel is not None:
vel_offset = self._forceModel.macroscopic_velocity_shift(
zeroth_order_moment)
else:
if self._forceModel is not None:
vel_offset = self._forceModel.macroscopic_velocity_shift(
sp.Rational(1, 1))
eqs = [Assignment(self._symbolOrder0, zeroth_order_moment)]
first_order_moments = [
a - b for a, b in zip(first_order_moments, vel_offset)
]
eqs += [
Assignment(l, r)
for l, r in zip(self._symbolsOrder1, first_order_moments)
]
result = AssignmentCollection(eqs, [])
if self._forceModel is not None and force_substitution:
result = add_symbolic_force_substitutions(
result, self._forceModel._subs_dict_force)
return result
def create_stream_only_kernel(stencil,
src_field,
dst_field=None,
accessor=StreamPullTwoFieldsAccessor()):
"""Creates a stream kernel, without collision.
Args:
stencil: lattice Boltzmann stencil which is used in the form of a tuple of tuples
src_field: field used for reading pdf values
dst_field: field used for writing pdf values if accessor.is_inplace this parameter is not necessary but it
is used if provided
accessor: instance of PdfFieldAccessor, defining where to read and write values
to create e.g. a fused stream-collide kernel See 'fieldaccess.PdfFieldAccessor'
Returns:
AssignmentCollection of the stream only update rule
"""
if accessor.is_inplace and dst_field is None:
dst_field = src_field
if not accessor.is_inplace and dst_field is None:
raise ValueError(
"For two field accessors a destination field has to be provided")
temporary_symbols = sp.symbols(f'streamed_:{stencil.Q}')
subexpressions = [
Assignment(tmp, acc) for tmp, acc in zip(
temporary_symbols, accessor.read(src_field, stencil))
]
main_assignments = [
Assignment(acc, tmp) for acc, tmp in zip(
accessor.write(dst_field, stencil), temporary_symbols)
]
return AssignmentCollection(main_assignments,
subexpressions=subexpressions)
def test_free_and_defined_symbols():
ac = AssignmentCollection([
Assignment(z, x + y),
Conditional(t > 0, Assignment(a, b + 1), Assignment(a, b + 2))
], [],
subexpression_symbol_generator=symbol_gen)
print(ac)
print(ac.__repr__)
def relax_central_moments(pre_collision_symbols, post_collision_symbols,
relaxation_rates, equilibrium_values,
force_terms):
equilibrium_vec = sp.Matrix(equilibrium_values)
moment_vec = sp.Matrix(pre_collision_symbols)
relaxation_matrix = sp.diag(*relaxation_rates)
moment_vec = moment_vec + relaxation_matrix * (equilibrium_vec - moment_vec) + force_terms
main_assignments = [Assignment(s, eq) for s, eq in zip(post_collision_symbols, moment_vec)]
return AssignmentCollection(main_assignments)
def test_add_subexpressions_for_field_reads():
s, v = fields("s(5), v(5): double[2D]")
subexpressions = []
main = [
Assignment(s[0, 0](0), 3 * v[0, 0](0)),
Assignment(s[0, 0](1), 10 * v[0, 0](1))
]
ac = AssignmentCollection(main, subexpressions)
assert len(ac.subexpressions) == 0
ac = add_subexpressions_for_field_reads(ac)
assert len(ac.subexpressions) == 2
def _get_collision_rule_with_relaxation_rate(
self,
relaxation_rate,
include_force_terms=True,
conserved_quantity_equations=None):
f = sp.Matrix(self.pre_collision_pdf_symbols)
rho = self._cqc.zeroth_order_moment_symbol
u = self._cqc.first_order_moment_symbols
all_subexpressions = []
if self._forceModel is not None:
all_subexpressions += AssignmentCollection(
self._forceModel.subs_dict_force).all_assignments
if conserved_quantity_equations is None:
conserved_quantity_equations = self._cqc.equilibrium_input_equations_from_pdfs(
f, False)
all_subexpressions += conserved_quantity_equations.all_assignments
eq = []
for w_i, direction in zip(self.weights, self.stencil):
f_i = rho * w_i
for u_a, e_ia in zip(u, direction):
b = sp.sqrt(1 + 3 * u_a**2)
f_i *= (2 - b) * ((2 * u_a + b) / (1 - u_a))**e_ia
eq.append(f_i)
collision_eqs = [
Assignment(lhs,
(1 - relaxation_rate) * f_i + relaxation_rate * eq_i)
for lhs, f_i, eq_i in zip(self.post_collision_pdf_symbols,
self.pre_collision_pdf_symbols, eq)
]
if (self._forceModel is not None) and include_force_terms:
force_model_terms = self._forceModel(self)
force_term_symbols = sp.symbols("forceTerm_:%d" %
(len(force_model_terms, )))
force_subexpressions = [
Assignment(sym,
force_model_term) for sym, force_model_term in zip(
force_term_symbols, force_model_terms)
]
all_subexpressions += force_subexpressions
collision_eqs = [
Assignment(eq.lhs, eq.rhs + force_term_symbol) for eq,
force_term_symbol in zip(collision_eqs, force_term_symbols)
]
cr = LbmCollisionRule(self, collision_eqs, all_subexpressions)
cr.simplification_hints['relaxation_rates'] = []
return cr
def backward_transform(self,
pdf_symbols,
simplification=True,
subexpression_base='sub_k_to_f',
start_from_monomials=False):
r"""Returns an assignment collection containing equations for post-collision populations,
expressed in terms of the post-collision polynomial moments by matrix-multiplication.
The moment transformation matrix :math:`M` provided by :func:`lbmpy.moments.moment_matrix` is
inverted and used to compute the pre-collision moments as
:math:`\mathbf{f}^{\ast} = M^{-1} \cdot \mathbf{M}_{\mathrm{post}}`, which is returned element-wise.
**Simplifications**
If simplification is enabled, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}`
of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
that
.. math::
f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
Args:
pdf_symbols: List of symbols that represent the post-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
start_from_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'backward')
if start_from_monomials:
assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
mm_inv = moment_matrix(self.moment_exponents, self.stencil).inv()
post_collision_moments = self.post_collision_monomial_symbols
else:
mm_inv = self.inv_moment_matrix
post_collision_moments = self.post_collision_symbols
m_to_f_vec = mm_inv * sp.Matrix(post_collision_moments)
main_assignments = [
Assignment(f, eq) for f, eq in zip(pdf_symbols, m_to_f_vec)
]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments,
subexpression_symbol_generator=symbol_gen)
ac.add_simplification_hint('stencil', self.stencil)
ac.add_simplification_hint('post_collision_pdf_symbols', pdf_symbols)
if simplification:
ac = simplification.apply(ac)
return ac
def test_split_inner_loop():
dst = ps.fields('dst(8): double[2D]')
s = sp.symbols('s_:8')
x = sp.symbols('x')
subexpressions = []
main = [
Assignment(dst[0, 0](0), s[0]),
Assignment(dst[0, 0](1), s[1]),
Assignment(dst[0, 0](2), s[2]),
Assignment(dst[0, 0](3), s[3]),
Assignment(dst[0, 0](4), s[4]),
Assignment(dst[0, 0](5), s[5]),
Assignment(dst[0, 0](6), s[6]),
Assignment(dst[0, 0](7), s[7]),
Assignment(x, sum(s))
]
ac = AssignmentCollection(main, subexpressions)
split_groups = [[dst[0, 0](0), dst[0, 0](1)], [dst[0, 0](2), dst[0, 0](3)],
[dst[0, 0](4), dst[0, 0](5)],
[dst[0, 0](6), dst[0, 0](7), x]]
ac.simplification_hints['split_groups'] = split_groups
ast = ps.create_kernel(ac)
code = ps.get_code_str(ast)
# we have four inner loops as indicated in split groups (4 elements) plus one outer loop
assert code.count('for') == 5
ast = ps.create_kernel(ac, target=ps.Target.GPU)
code = ps.get_code_str(ast)
# on GPUs is wouldn't be good to use loop splitting
assert code.count('for') == 0
ac = AssignmentCollection(main, subexpressions)
ast = ps.create_kernel(ac)
code = ps.get_code_str(ast)
# one inner loop and one outer loop
assert code.count('for') == 2
def _compute_weights(self):
defaults = self._conserved_quantity_computation.default_values
cqe = AssignmentCollection(
[Assignment(s, e) for s, e in defaults.items()])
eq_ac = self.get_equilibrium(cqe,
subexpressions=False,
keep_cqc_subexpressions=False)
weights = []
for eq in eq_ac.main_assignments:
value = eq.rhs.expand()
assert len(value.atoms(
sp.Symbol)) == 0, "Failed to compute weights"
weights.append(value)
return weights
def create_stream_pull_with_output_kernel(
lb_method,
src_field,
dst_field=None,
output=None,
accessor=StreamPullTwoFieldsAccessor()):
"""Creates a stream kernel, without collision but macroscopic quantaties like density or velocity can be calculated.
Args:
lb_method: lattice Boltzmann method see 'creationfunctions.create_lb_method'
src_field: field used for reading pdf values
dst_field: field used for writing pdf values if accessor.is_inplace this parameter is ignored
output: dictonary which containes macroscopic quantities as keys which should be calculated and fields as
values which should be used to write the data e.g.: {'density': density_field}
accessor: instance of PdfFieldAccessor, defining where to read and write values
to create e.g. a fused stream-collide kernel See 'fieldaccess.PdfFieldAccessor'
Returns:
AssignmentCollection of the stream only update rule
"""
if accessor.is_inplace:
dst_field = src_field
if not accessor.is_inplace and dst_field is None:
raise ValueError(
"For two field accessors a destination field has to be provided")
stencil = lb_method.stencil
cqc = lb_method.conserved_quantity_computation
streamed = sp.symbols(f"streamed_:{stencil.Q}")
stream_assignments = [
Assignment(a, b)
for a, b in zip(streamed, accessor.read(src_field, stencil))
]
output_eq_collection = cqc.output_equations_from_pdfs(streamed, output) if output\
else AssignmentCollection(main_assignments=[])
write_eqs = [
Assignment(a, b)
for a, b in zip(accessor.write(dst_field, stencil), streamed)
]
subexpressions = stream_assignments + output_eq_collection.subexpressions
main_eqs = output_eq_collection.main_assignments + write_eqs
return LbmCollisionRule(
lb_method,
main_eqs,
subexpressions,
simplification_hints=output_eq_collection.simplification_hints)
def test_subexpression_substitution_in_main_assignments():
subexpressions = [
Assignment(s0, 2 * a + 2 * b),
Assignment(s1, 2 * a + 2 * b + 2 * c),
Assignment(s2, 2 * a + 2 * b + 2 * c + 2 * d),
Assignment(s3, 2 * a + 2 * b * c),
Assignment(x, s1 + s2 + s0 + s3)
]
main = [
Assignment(f[0], s1 + s2 + s0 + s3),
Assignment(f[1], s1 + s2 + s0 + s3),
Assignment(f[2], s1 + s2 + s0 + s3),
Assignment(f[3], s1 + s2 + s0 + s3),
Assignment(f[4], s1 + s2 + s0 + s3)
]
ac = AssignmentCollection(main, subexpressions)
ac = subexpression_substitution_in_main_assignments(ac)
for i in range(0, len(ac.main_assignments)):
assert ac.main_assignments[i].rhs == x
def forward_transform(self,
pdf_symbols,
simplification=True,
subexpression_base='sub_f_to_M',
return_monomials=False):
r"""Returns an assignment collection containing equations for pre-collision polynomial
moments, expressed in terms of the pre-collision populations by matrix-multiplication.
The moment transformation matrix :math:`M` provided by :func:`lbmpy.moments.moment_matrix` is
used to compute the pre-collision moments as :math:`\mathbf{M} = M \cdot \mathbf{f}`,
which is returned element-wise.
Args:
pdf_symbols: List of symbols that represent the pre-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
return_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'forward')
if return_monomials:
assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
mm = moment_matrix(self.moment_exponents, self.stencil)
pre_collision_moments = self.pre_collision_monomial_symbols
else:
mm = self.moment_matrix
pre_collision_moments = self.pre_collision_symbols
f_to_m_vec = mm * sp.Matrix(pdf_symbols)
main_assignments = [
Assignment(m, eq)
for m, eq in zip(pre_collision_moments, f_to_m_vec)
]
symbol_gen = SymbolGen(symbol=subexpression_base)
ac = AssignmentCollection(main_assignments,
subexpression_symbol_generator=symbol_gen)
if simplification:
ac = simplification.apply(ac)
return ac
def forward_transform(self,
cumulant_base=PRE_COLLISION_CUMULANT,
central_moment_base=PRE_COLLISION_MONOMIAL_CENTRAL_MOMENT,
simplification=True,
subexpression_base='sub_k_to_C'):
simplification = self._get_simp_strategy(simplification)
main_assignments = []
for exp in self.cumulant_exponents:
eq = self.cumulant_from_central_moments(exp, central_moment_base)
c_symbol = statistical_quantity_symbol(cumulant_base, exp)
main_assignments.append(Assignment(c_symbol, eq))
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(
main_assignments, subexpression_symbol_generator=symbol_gen)
if simplification:
ac = simplification.apply(ac)
return ac
def test_assignment_collection():
ac = AssignmentCollection([Assignment(z, x + y)], [],
subexpression_symbol_generator=symbol_gen)
lhs = ac.add_subexpression(t)
assert lhs == sp.Symbol("a_0")
ac.subexpressions.append(Assignment(t, 3))
ac.topological_sort(sort_main_assignments=False, sort_subexpressions=True)
assert ac.subexpressions[0].lhs == t
assert ac.new_with_inserted_subexpression(sp.Symbol("not_defined")) == ac
ac_inserted = ac.new_with_inserted_subexpression(t)
ac_inserted2 = ac.new_without_subexpressions({lhs})
assert all(a == b for a, b in zip(ac_inserted.all_assignments,
ac_inserted2.all_assignments))
print(ac_inserted)
assert ac_inserted.subexpressions[0] == Assignment(lhs, 3)
assert 'a_0' in str(ac_inserted)
assert '<table' in ac_inserted._repr_html_()
def relax_polynomial_cumulants(monomial_exponents,
polynomials,
relaxation_rates,
equilibrium_values,
pre_simplification,
galilean_correction_terms=None,
pre_collision_base=PRE_COLLISION_CUMULANT,
post_collision_base=POST_COLLISION_CUMULANT,
subexpression_base='sub_col'):
mon_to_poly_matrix = monomial_to_polynomial_transformation_matrix(
monomial_exponents, polynomials)
mon_vec = sp.Matrix([
statistical_quantity_symbol(pre_collision_base, exp)
for exp in monomial_exponents
])
equilibrium_vec = sp.Matrix(equilibrium_values)
relaxation_matrix = sp.diag(*relaxation_rates)
subexpressions = []
poly_vec = mon_to_poly_matrix * mon_vec
relaxed_polys = poly_vec + relaxation_matrix * (equilibrium_vec - poly_vec)
if galilean_correction_terms is not None:
relaxed_polys = add_galilean_correction(relaxed_polys, polynomials,
galilean_correction_terms)
subexpressions = galilean_correction_terms.all_assignments
relaxed_monos = mon_to_poly_matrix.inv() * relaxed_polys
main_assignments = [
Assignment(statistical_quantity_symbol(post_collision_base, exp), v)
for exp, v in zip(monomial_exponents, relaxed_monos)
]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments,
subexpressions=subexpressions,
subexpression_symbol_generator=symbol_gen)
if pre_simplification == 'default_with_cse':
ac = sympy_cse(ac)
return ac
def test_subexpression_insertion():
f, g = fields('f(10), g(10) : [2D]')
xi = sp.symbols('xi_:10')
xi_set = set(xi)
subexpressions = [
Assignment(xi[0], -f(4)),
Assignment(xi[1], -(f(1) * f(2))),
Assignment(xi[2], 2.31 * f(5)),
Assignment(xi[3], 1.8 + f(5) + f(6)),
Assignment(xi[4], 5.7 + f(6)),
Assignment(xi[5], (f(4) + f(5))**2),
Assignment(xi[6],
f(3)**2),
Assignment(xi[7], f(4)),
Assignment(xi[8], 13),
Assignment(xi[9], 0),
]
assignments = [Assignment(g(i), x) for i, x in enumerate(xi)]
ac = AssignmentCollection(assignments, subexpressions=subexpressions)
ac_ins = insert_symbol_times_minus_one(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[0]})
ac_ins = insert_constant_multiples(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[0], xi[2]})
ac_ins = insert_constant_additions(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[4]})
ac_ins = insert_squares(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[6]})
ac_ins = insert_aliases(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[7]})
ac_ins = insert_zeros(ac)
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[9]})
ac_ins = insert_constants(ac, skip={xi[9]})
assert (ac_ins.bound_symbols & xi_set) == (xi_set - {xi[8]})
def relax_lower_order_central_moments(
moment_indices,
pre_collision_values,
relaxation_rates,
equilibrium_values,
post_collision_base=POST_COLLISION_MONOMIAL_CENTRAL_MOMENT):
post_collision_symbols = [
statistical_quantity_symbol(post_collision_base, i)
for i in moment_indices
]
equilibrium_vec = sp.Matrix(equilibrium_values)
moment_vec = sp.Matrix(pre_collision_values)
relaxation_matrix = sp.diag(*relaxation_rates)
moment_vec = moment_vec + relaxation_matrix * (equilibrium_vec -
moment_vec)
main_assignments = [
Assignment(s, eq) for s, eq in zip(post_collision_symbols, moment_vec)
]
return AssignmentCollection(main_assignments)
def backward_transform(self,
cumulant_base=POST_COLLISION_CUMULANT,
central_moment_base=POST_COLLISION_MONOMIAL_CENTRAL_MOMENT,
simplification=True,
omit_conserved_moments=False,
subexpression_base='sub_C_to_k'):
simplification = self._get_simp_strategy(simplification)
main_assignments = []
for exp in self.central_moment_exponents:
if omit_conserved_moments and get_order(exp) <= 1:
continue
eq = self.central_moment_from_cumulants(exp, cumulant_base)
k_symbol = statistical_quantity_symbol(central_moment_base, exp)
main_assignments.append(Assignment(k_symbol, eq))
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments, subexpression_symbol_generator=symbol_gen)
if simplification:
ac = simplification.apply(ac)
return ac
def test_add_subexpressions_for_sums():
subexpressions = [
Assignment(s0, a + b + c + d),
Assignment(s1, 3 * a * sp.sqrt(x) + 4 * b + c),
Assignment(s2, 3 * a * sp.sqrt(x) + 4 * b + c),
Assignment(s3, 3 * a * sp.sqrt(x) + 4 * b + c)
]
main = [Assignment(f[0], s1 + s2 + s0 + s3)]
ac = AssignmentCollection(main, subexpressions)
ops_before_optimisation = ac.operation_count
ac = add_subexpressions_for_sums(ac)
ops_after_optimisation = ac.operation_count
assert ops_after_optimisation["adds"] == ops_before_optimisation["adds"]
assert ops_after_optimisation["muls"] < ops_before_optimisation["muls"]
assert ops_after_optimisation["sqrts"] < ops_before_optimisation["sqrts"]
rhs = []
for i in range(len(ac.subexpressions)):
rhs.append(ac.subexpressions[i].rhs)
assert a + b + c + d in rhs
assert 3 * a * sp.sqrt(x) in rhs
def test_add_subexpressions_for_divisions():
subexpressions = [
Assignment(s0, 2 / a + 2 / b),
Assignment(s1, 2 / a + 2 / b + 2 / c),
Assignment(s2, 2 / a + 2 / b + 2 / c + 2 / d),
Assignment(s3, 2 / a + 2 / b / c),
Assignment(x, s1 + s2 + s0 + s3)
]
main = [Assignment(f[0], s1 + s2 + s0 + s3)]
ac = AssignmentCollection(main, subexpressions)
divs_before_optimisation = ac.operation_count["divs"]
ac = add_subexpressions_for_divisions(ac)
divs_after_optimisation = ac.operation_count["divs"]
assert divs_before_optimisation - divs_after_optimisation == 8
rhs = []
for i in range(len(ac.subexpressions)):
rhs.append(ac.subexpressions[i].rhs)
assert 1 / a in rhs
assert 1 / b in rhs
assert 1 / c in rhs
assert 1 / d in rhs
def backward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_k_to_f',
start_from_monomials=False):
r"""Returns an assignment collection containing equations for post-collision populations,
expressed in terms of the post-collision polynomial central moments by matrix-multiplication.
The moment transformation matrix :math:`K` provided by :func:`lbmpy.moments.moment_matrix` is
inverted and used to compute the pre-collision moments as
:math:`\mathbf{f}^{\ast} = K^{-1} \cdot \mathbf{K}_{\mathrm{post}}`, which is returned element-wise.
Args:
pdf_symbols: List of symbols that represent the post-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
start_from_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification)
if start_from_monomials:
assert len(self.moment_exponents) == self.q, "Could not derive invertible monomial transform." \
f"Expected {self.q} monomials, but got {len(self.moment_exponents)}."
mm_inv = moment_matrix(self.moment_exponents, self.stencil).inv()
shift_inv = set_up_shift_matrix(self.moment_exponents, self.stencil, self.equilibrium_velocity).inv()
km_inv = mm_inv * shift_inv
post_collision_moments = self.post_collision_monomial_symbols
else:
km_inv = self.backward_matrix
post_collision_moments = self.post_collision_symbols
m_to_f_vec = km_inv * sp.Matrix(post_collision_moments)
main_assignments = [Assignment(f, eq) for f, eq in zip(pdf_symbols, m_to_f_vec)]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments, subexpression_symbol_generator=symbol_gen)
if simplification:
ac = simplification.apply(ac)
return ac
def test_simplification_strategy():
a, b, c, d, x, y, z = sp.symbols("a b c d x y z")
s0, s1, s2, s3 = sp.symbols("s_:4")
a0, a1, a2, a3 = sp.symbols("a_:4")
subexpressions = [
Assignment(s0, 2 * a + 2 * b),
Assignment(s1, 2 * a + 2 * b + 2 * c),
Assignment(s2, 2 * a + 2 * b + 2 * c + 2 * d),
]
main = [
Assignment(a0, s0 + s1),
Assignment(a1, s0 + s2),
Assignment(a2, s1 + s2),
]
ac = AssignmentCollection(main, subexpressions)
strategy = SimplificationStrategy()
strategy.add(subexpression_substitution_in_existing_subexpressions)
strategy.add(apply_on_all_subexpressions(sp.factor))
result = strategy(ac)
assert result.operation_count['adds'] == 7
assert result.operation_count['muls'] == 5
assert result.operation_count['divs'] == 0
# Trigger display routines, such that they are at least executed
report = strategy.show_intermediate_results(ac, symbols=[s0])
assert 's_0' in str(report)
report = strategy.show_intermediate_results(ac)
assert 's_{1}' in report._repr_html_()
report = strategy.create_simplification_report(ac)
assert 'Adds' in str(report)
assert 'Adds' in report._repr_html_()
assert 'factor' in str(strategy)
def test_add_subexpressions_for_constants():
half = sp.Rational(1, 2)
sqrt_2 = sp.sqrt(2)
main = [
Assignment(f[0], half * a + half * b + half * c),
Assignment(f[1], -half * a - half * b),
Assignment(f[2], a * sqrt_2 - b * sqrt_2),
Assignment(f[3], a**2 + b**2)
]
ac = AssignmentCollection(main)
ac = add_subexpressions_for_constants(ac)
assert len(ac.subexpressions) == 2
half_subexp = None
sqrt_subexp = None
for asm in ac.subexpressions:
if asm.rhs == half:
half_subexp = asm.lhs
elif asm.rhs == sqrt_2:
sqrt_subexp = asm.lhs
else:
pytest.fail(f"An unexpected subexpression was encountered: {asm}")
assert half_subexp is not None
assert sqrt_subexp is not None
for asm in ac.main_assignments[:3]:
assert isinstance(asm.rhs, sp.Mul)
assert any(arg == half_subexp for arg in ac.main_assignments[0].rhs.args)
assert any(arg == half_subexp for arg in ac.main_assignments[1].rhs.args)
assert any(arg == sqrt_subexp for arg in ac.main_assignments[2].rhs.args)
# Do not replace exponents!
assert ac.main_assignments[3].rhs == a**2 + b**2
def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
return_monomials=False):
r"""Returns equations for polynomial central moments, computed from pre-collision populations
through a cascade of three steps.
First, the monomial raw moment vector :math:`\mathbf{m}` is computed using the raw-moment
chimera transform (see `lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform`). Then, the
monomial shift matrix :math:`N` provided by `lbmpy.moments.set_up_shift_matrix` is used to compute
the monomial central moment vector as :math:`\mathbf{\kappa} = N \mathbf{m}`. Lastly, the polynomial
central moments are computed using the polynomialization matrix as :math:`\mathbf{K} = P \mathbf{\kappa}`.
**Conserved Quantity Equations**
If given, this transform absorbs the conserved quantity equations and simplifies them
using the raw moment equations, if simplification is enabled.
**Simplification**
If simplification is enabled, the absorbed conserved quantity equations are - if possible -
rewritten using the monomial symbols. If the conserved quantities originate somewhere else
than in the lower-order moments (like from an external field), they are not affected by this
simplification.
The relations between conserved quantities and raw moments are used to simplify the equations
obtained from the shift matrix. Further, these equations are simplified by recursively inserting
lower-order moments into equations for higher-order moments.
**De-Aliasing**
If more than :math:`q` monomial moments are extracted from the polynomial set, they
are de-aliased and reduced to a set of only :math:`q` moments using the same rules
as for raw moments. For polynomialization, a special reduced matrix :math:`\tilde{P}`
is used, which is computed using `lbmpy.moments.central_moment_reduced_monomial_to_polynomial_matrix`.
Args:
pdf_symbols: List of symbols that represent the pre-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
return_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'forward')
raw_moment_base = self.raw_moment_transform.mono_base_pre
central_moment_base = self.mono_base_pre
mono_rm_symbols = self.raw_moment_transform.pre_collision_monomial_symbols
mono_cm_symbols = self.pre_collision_monomial_symbols
rm_ac = self.raw_moment_transform.forward_transform(pdf_symbols, simplification=False, return_monomials=True)
cq_symbols_to_moments = self.raw_moment_transform.get_cq_to_moment_symbols_dict(raw_moment_base)
rm_to_cm_vec = self.shift_matrix * sp.Matrix(mono_rm_symbols)
cq_subs = dict()
if simplification:
from lbmpy.methods.momentbased.momentbasedsimplifications import (
substitute_moments_in_conserved_quantity_equations)
rm_ac = substitute_moments_in_conserved_quantity_equations(rm_ac)
# Compute replacements for conserved moments in terms of the CQE
rm_asm_dict = rm_ac.main_assignments_dict
for cq_sym, moment_sym in cq_symbols_to_moments.items():
cq_eq = rm_asm_dict[cq_sym]
solutions = sp.solve(cq_eq - cq_sym, moment_sym)
if len(solutions) > 0:
cq_subs[moment_sym] = solutions[0]
rm_to_cm_vec = fast_subs(rm_to_cm_vec, cq_subs)
rm_to_cm_dict = {cm: rm for cm, rm in zip(mono_cm_symbols, rm_to_cm_vec)}
if simplification:
rm_to_cm_dict = self._simplify_raw_to_central_moments(
rm_to_cm_dict, self.moment_exponents, raw_moment_base, central_moment_base)
rm_to_cm_dict = self._undo_remaining_cq_subexpressions(rm_to_cm_dict, cq_subs)
subexpressions = rm_ac.all_assignments
if return_monomials:
main_assignments = [Assignment(lhs, rhs) for lhs, rhs in rm_to_cm_dict.items()]
else:
subexpressions += [Assignment(lhs, rhs) for lhs, rhs in rm_to_cm_dict.items()]
poly_eqs = self.mono_to_poly_matrix * sp.Matrix(mono_cm_symbols)
main_assignments = [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, poly_eqs)]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments=main_assignments, subexpressions=subexpressions,
subexpression_symbol_generator=symbol_gen)
if simplification:
ac = simplification.apply(ac)
return ac
def forward_transform(self, pdf_symbols, simplification=True, subexpression_base='sub_f_to_k',
return_monomials=False):
r"""Returns an assignment collection containing equations for pre-collision polynomial
central moments, expressed in terms of the pre-collision populations.
The monomial central moments are computed from populations through the central-moment
chimera transform:
.. math::
f_{xyz} &:= f_i \text{ such that } c_i = (x,y,z)^T \\
\kappa_{xy|\gamma} &:= \sum_{z \in \{-1, 0, 1\} } f_{xyz} \cdot (z - u_z)^{\gamma} \\
\kappa_{x|\beta \gamma} &:= \sum_{y \in \{-1, 0, 1\}} \kappa_{xy|\gamma} \cdot (y - u_y)^{\beta} \\
\kappa_{\alpha \beta \gamma} &:= \sum_{x \in \{-1, 0, 1\}} \kappa_{x|\beta \gamma} \cdot (x - u_x)^{\alpha}
The polynomial moments are afterward computed from the monomials by matrix-multiplication
using the polynomialization matrix :math:`P`.
**De-Aliasing**
If more than :math:`q` monomial moments are extracted from the polynomial set, they
are de-aliased and reduced to a set of only :math:`q` moments using the same rules
as for raw moments. For polynomialization, a special reduced matrix :math:`\tilde{P}`
is used, which is computed using `lbmpy.moments.central_moment_reduced_monomial_to_polynomial_matrix`.
Args:
pdf_symbols: List of symbols that represent the pre-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
return_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'forward')
monomial_symbol_base = self.mono_base_pre
def _partial_kappa_symbol(fixed_directions, remaining_exponents):
fixed_str = '_'.join(str(direction) for direction in fixed_directions).replace('-', 'm')
exp_str = '_'.join(str(exp) for exp in remaining_exponents).replace('-', 'm')
return sp.Symbol(f"partial_{monomial_symbol_base}_{fixed_str}_e_{exp_str}")
partial_sums_dict = dict()
monomial_moment_eqs = []
def collect_partial_sums(exponents, dimension=0, fixed_directions=tuple()):
if dimension == self.dim:
# Base Case
if fixed_directions in self.stencil:
return pdf_symbols[self.stencil.index(fixed_directions)]
else:
return 0
else:
# Recursive Case
summation = sp.sympify(0)
for d in [-1, 0, 1]:
next_partial = collect_partial_sums(
exponents, dimension=dimension + 1, fixed_directions=fixed_directions + (d,))
summation += next_partial * (d - self.equilibrium_velocity[dimension])**exponents[dimension]
if dimension == 0:
lhs_symbol = sq_sym(monomial_symbol_base, exponents)
monomial_moment_eqs.append(Assignment(lhs_symbol, summation))
else:
lhs_symbol = _partial_kappa_symbol(fixed_directions, exponents[dimension:])
partial_sums_dict[lhs_symbol] = summation
return lhs_symbol
for e in self.moment_exponents:
collect_partial_sums(e)
subexpressions = [Assignment(lhs, rhs) for lhs, rhs in partial_sums_dict.items()]
if return_monomials:
main_assignments = monomial_moment_eqs
else:
subexpressions += monomial_moment_eqs
moment_eqs = self.mono_to_poly_matrix * sp.Matrix(self.pre_collision_monomial_symbols)
main_assignments = [Assignment(m, v) for m, v in zip(self.pre_collision_symbols, moment_eqs)]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments, subexpressions=subexpressions,
subexpression_symbol_generator=symbol_gen)
if simplification:
ac = self._simplify_lower_order_moments(ac, monomial_symbol_base, return_monomials)
ac = simplification.apply(ac)
return ac
def backward_transform(self,
pdf_symbols,
simplification=True,
subexpression_base='sub_k_to_f',
start_from_monomials=False):
r"""Returns an assignment collection containing equations for post-collision populations,
expressed in terms of the post-collision polynomial moments by matrix-multiplication.
The post-collision monomial moments :math:`\mathbf{m}_{\mathrm{post}}` are first obtained from the polynomials.
Then, the monomial transformation matrix :math:`M_r` provided by :func:`lbmpy.moments.moment_matrix`
is inverted and used to compute the post-collision populations as
:math:`\mathbf{f}_{\mathrm{post}} = M_r^{-1} \cdot \mathbf{m}_{\mathrm{post}}`.
**De-Aliasing**
See `PdfsToMomentsByChimeraTransform.forward_transform`.
**Simplifications**
If simplification is enabled, the equations for populations :math:`f_i` and :math:`f_{\bar{i}}`
of opposite stencil directions :math:`\mathbf{c}_i` and :math:`\mathbf{c}_{\bar{i}} = - \mathbf{c}_i`
are split into their symmetric and antisymmetric parts :math:`f_i^{\mathrm{sym}}, f_i^{\mathrm{anti}}`, such
that
.. math::
f_i = f_i^{\mathrm{sym}} + f_i^{\mathrm{anti}}
f_{\bar{i}} = f_i^{\mathrm{sym}} - f_i^{\mathrm{anti}}
Args:
pdf_symbols: List of symbols that represent the post-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
start_from_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'backward')
post_collision_moments = self.post_collision_symbols
post_collision_monomial_moments = self.post_collision_monomial_symbols
subexpressions = []
if not start_from_monomials:
raw_moment_eqs = self.poly_to_mono_matrix * sp.Matrix(
post_collision_moments)
subexpressions += [
Assignment(rm, v) for rm, v in zip(
post_collision_monomial_moments, raw_moment_eqs)
]
rm_to_f_vec = self.inv_moment_matrix * sp.Matrix(
post_collision_monomial_moments)
main_assignments = [
Assignment(f, eq) for f, eq in zip(pdf_symbols, rm_to_f_vec)
]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments,
subexpressions=subexpressions,
subexpression_symbol_generator=symbol_gen)
ac.add_simplification_hint('stencil', self.stencil)
ac.add_simplification_hint('post_collision_pdf_symbols', pdf_symbols)
if simplification:
ac = simplification.apply(ac)
return ac
def forward_transform(self,
pdf_symbols,
simplification=True,
subexpression_base='sub_f_to_m',
return_monomials=False):
r"""Returns an assignment collection containing equations for pre-collision polynomial
moments, expressed in terms of the pre-collision populations, using the raw moment chimera transform.
The chimera transform for raw moments is given by :cite:`geier2015` :
.. math::
f_{xyz} &:= f_i \text{ such that } c_i = (x,y,z)^T \\
m_{xy|\gamma} &:= \sum_{z \in \{-1, 0, 1\} } f_{xyz} \cdot z^{\gamma} \\
m_{x|\beta \gamma} &:= \sum_{y \in \{-1, 0, 1\}} m_{xy|\gamma} \cdot y^{\beta} \\
m_{\alpha \beta \gamma} &:= \sum_{x \in \{-1, 0, 1\}} m_{x|\beta \gamma} \cdot x^{\alpha}
The obtained raw moments are afterward combined to the desired polynomial moments.
**Conserved Quantity Equations**
If given, this transform absorbs the conserved quantity equations and simplifies them
using the monomial raw moment equations, if simplification is enabled.
**De-Aliasing**
If more than :math:`q` monomial moments are extracted from the polynomial set, the polynomials
are de-aliased by eliminating aliases according to the stencil
using `lbmpy.moments.non_aliased_polynomial_raw_moments`.
**Simplification**
If simplification is enabled, the absorbed conserved quantity equations are - if possible -
rewritten using the monomial symbols. If the conserved quantities originate somewhere else
than in the lower-order moments (like from an external field), they are not affected by this
simplification.
Args:
pdf_symbols: List of symbols that represent the pre-collision populations
simplification: Simplification specification. See :class:`AbstractMomentTransform`
subexpression_base: The base name used for any subexpressions of the transformation.
return_monomials: Return equations for monomial moments. Use only when specifying
``moment_exponents`` in constructor!
"""
simplification = self._get_simp_strategy(simplification, 'forward')
monomial_symbol_base = self.mono_base_pre
def _partial_kappa_symbol(fixed_directions, remaining_exponents):
fixed_str = '_'.join(
str(direction)
for direction in fixed_directions).replace('-', 'm')
exp_str = '_'.join(
str(exp) for exp in remaining_exponents).replace('-', 'm')
return sp.Symbol(
f"partial_{monomial_symbol_base}_{fixed_str}_e_{exp_str}")
partial_sums_dict = dict()
monomial_eqs = []
def collect_partial_sums(exponents,
dimension=0,
fixed_directions=tuple()):
if dimension == self.dim:
# Base Case
if fixed_directions in self.stencil:
return pdf_symbols[self.stencil.index(fixed_directions)]
else:
return 0
else:
# Recursive Case
summation = sp.sympify(0)
for d in [-1, 0, 1]:
next_partial = collect_partial_sums(
exponents,
dimension=dimension + 1,
fixed_directions=fixed_directions + (d, ))
summation += next_partial * d**exponents[dimension]
if dimension == 0:
lhs_symbol = sq_sym(monomial_symbol_base, exponents)
monomial_eqs.append(Assignment(lhs_symbol, summation))
else:
lhs_symbol = _partial_kappa_symbol(fixed_directions,
exponents[dimension:])
partial_sums_dict[lhs_symbol] = summation
return lhs_symbol
for e in self.moment_exponents:
collect_partial_sums(e)
main_assignments = self.cqe.main_assignments.copy(
) if self.cqe is not None else []
subexpressions = self.cqe.subexpressions.copy(
) if self.cqe is not None else []
subexpressions += [
Assignment(lhs, rhs) for lhs, rhs in partial_sums_dict.items()
]
if return_monomials:
main_assignments += monomial_eqs
else:
subexpressions += monomial_eqs
moment_eqs = self.mono_to_poly_matrix * sp.Matrix(
self.pre_collision_monomial_symbols)
main_assignments += [
Assignment(m, v)
for m, v in zip(self.pre_collision_symbols, moment_eqs)
]
symbol_gen = SymbolGen(subexpression_base)
ac = AssignmentCollection(main_assignments,
subexpressions=subexpressions,
subexpression_symbol_generator=symbol_gen)
ac.add_simplification_hint(
'cq_symbols_to_moments',
self.get_cq_to_moment_symbols_dict(monomial_symbol_base))
if simplification:
ac = simplification.apply(ac)
return ac