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