def visit_node(node, substitution_dict, default_type='double'):
substitution_dict = substitution_dict.copy()
for arg in node.args:
if isinstance(arg, ast.SympyAssignment):
assignment = arg
subs_expr = fast_subs(assignment.rhs, substitution_dict,
skip=lambda e: isinstance(e, ast.ResolvedFieldAccess))
assignment.rhs = visit_expr(subs_expr, default_type)
rhs_type = get_type_of_expression(assignment.rhs)
if isinstance(assignment.lhs, TypedSymbol):
lhs_type = assignment.lhs.dtype
if type(rhs_type) is VectorType and type(lhs_type) is not VectorType:
new_lhs_type = VectorType(lhs_type, rhs_type.width)
new_lhs = TypedSymbol(assignment.lhs.name, new_lhs_type)
substitution_dict[assignment.lhs] = new_lhs
assignment.lhs = new_lhs
elif isinstance(assignment.lhs, vector_memory_access):
assignment.lhs = visit_expr(assignment.lhs, default_type)
elif isinstance(arg, ast.Conditional):
arg.condition_expr = fast_subs(arg.condition_expr, substitution_dict,
skip=lambda e: isinstance(e, ast.ResolvedFieldAccess))
arg.condition_expr = visit_expr(arg.condition_expr, default_type)
visit_node(arg, substitution_dict, default_type)
else:
visit_node(arg, substitution_dict, default_type)
def _simplify_lower_order_moments(self, ac, moment_base, search_in_main_assignments):
if self.cqe is None:
return ac
moment_symbols = [sq_sym(moment_base, e) for e in moments_up_to_order(1, dim=self.dim)]
if search_in_main_assignments:
f_to_cm_dict = ac.main_assignments_dict
f_to_cm_dict_reduced = ac.new_without_subexpressions().main_assignments_dict
else:
f_to_cm_dict = ac.subexpressions_dict
f_to_cm_dict_reduced = ac.new_without_subexpressions(moment_symbols).subexpressions_dict
cqe_subs = self.cqe.new_without_subexpressions().main_assignments_dict
for m in moment_symbols:
m_eq = fast_subs(fast_subs(f_to_cm_dict_reduced[m], cqe_subs), cqe_subs)
m_eq = m_eq.expand().cancel()
for cqe_sym, cqe_exp in cqe_subs.items():
m_eq = subs_additive(m_eq, cqe_sym, cqe_exp)
f_to_cm_dict[m] = m_eq
if search_in_main_assignments:
main_assignments = [Assignment(lhs, rhs) for lhs, rhs in f_to_cm_dict.items()]
return ac.copy(main_assignments=main_assignments)
else:
subexpressions = [Assignment(lhs, rhs) for lhs, rhs in f_to_cm_dict.items()]
return ac.copy(subexpressions=subexpressions)
def __cumulant_raw_moment_transform(index, dependent_var_dict, outer_function,
default_prefix, centralized):
"""Function to express cumulants as function of moments and vice versa.
Uses multivariate version of Faa di Bruno's formula.
Args:
index: tuple describing the index of the cumulant/moment to express as function of moments/cumulants
dependent_var_dict: a dictionary from index tuple to moments/cumulants symbols, or None to use default symbols
outer_function: logarithm to transform from moments->cumulants, exp for inverse direction
default_prefix: if dependent_var_dict is None, this is used to construct symbols of the form prefix_i_j_k
centralized: if True the first order moments/cumulants are set to zero
"""
dim = len(index)
subs_dict = {}
def create_moment_symbol(idx):
idx = tuple(idx)
result_symbol = sp.Symbol(default_prefix + "_" +
"_".join(["%d"] * len(idx)) % idx)
if dependent_var_dict is not None and idx in dependent_var_dict:
subs_dict[result_symbol] = dependent_var_dict[idx]
return result_symbol
zeroth_moment = create_moment_symbol((0, ) * dim)
def outer_function_derivative(n):
x = zeroth_moment
return sp.diff(outer_function(x), *tuple([x] * n))
# index (2,1,0) means differentiate twice w.r.t to first variable, and once w.r.t to second variable
# this is transformed here into representation [0,0,1] such that each entry is one diff operation
partition_list = []
for i, index_component in enumerate(index):
for j in range(index_component):
partition_list.append(i)
if len(partition_list) == 0: # special case for zero index
return fast_subs(outer_function(zeroth_moment), subs_dict)
# implementation of Faa di Bruno's formula:
result = 0
for partition in __partition(partition_list):
factor = outer_function_derivative(len(partition))
for elements in partition:
moment_index = [
0,
] * dim
for i in elements:
moment_index[i] += 1
factor *= create_moment_symbol(moment_index)
result += factor
if centralized:
for i in range(dim):
index = [0] * dim
index[i] = 1
result = result.subs(create_moment_symbol(index), 0)
return fast_subs(result, subs_dict)
def new_merged(self, other: 'AssignmentCollection') -> 'AssignmentCollection':
"""Returns a new collection which contains self and other. Subexpressions are renamed if they clash."""
own_definitions = set([e.lhs for e in self.main_assignments])
other_definitions = set([e.lhs for e in other.main_assignments])
assert len(own_definitions.intersection(other_definitions)) == 0, \
"Cannot merge collections, since both define the same symbols"
own_subexpression_symbols = {e.lhs: e.rhs for e in self.subexpressions}
substitution_dict = {}
processed_other_subexpression_equations = []
for other_subexpression_eq in other.subexpressions:
if other_subexpression_eq.lhs in own_subexpression_symbols:
if other_subexpression_eq.rhs == own_subexpression_symbols[other_subexpression_eq.lhs]:
continue # exact the same subexpression equation exists already
else:
# different definition - a new name has to be introduced
new_lhs = next(self.subexpression_symbol_generator)
new_eq = Assignment(new_lhs, fast_subs(other_subexpression_eq.rhs, substitution_dict))
processed_other_subexpression_equations.append(new_eq)
substitution_dict[other_subexpression_eq.lhs] = new_lhs
else:
processed_other_subexpression_equations.append(fast_subs(other_subexpression_eq, substitution_dict))
processed_other_main_assignments = [fast_subs(eq, substitution_dict) for eq in other.main_assignments]
return self.copy(self.main_assignments + processed_other_main_assignments,
self.subexpressions + processed_other_subexpression_equations)
def new_without_subexpressions(
self, subexpressions_to_keep: Set[sp.Symbol] = set()
) -> 'AssignmentCollection':
"""Returns a new collection where all subexpressions have been inserted."""
if len(self.subexpressions) == 0:
return self.copy()
subexpressions_to_keep = set(subexpressions_to_keep)
kept_subexpressions = []
if self.subexpressions[0].lhs in subexpressions_to_keep:
substitution_dict = {}
kept_subexpressions = self.subexpressions[0]
else:
substitution_dict = {
self.subexpressions[0].lhs: self.subexpressions[0].rhs
}
subexpression = [e for e in self.subexpressions]
for i in range(1, len(subexpression)):
subexpression[i] = fast_subs(subexpression[i], substitution_dict)
if subexpression[i].lhs in subexpressions_to_keep:
kept_subexpressions.append(subexpression[i])
else:
substitution_dict[subexpression[i].lhs] = subexpression[i].rhs
new_assignment = [
fast_subs(eq, substitution_dict) for eq in self.main_assignments
]
return self.copy(new_assignment, kept_subexpressions)
def fast_subs(self, subs_dict, skip=None):
self.body = fast_subs(self.body, subs_dict, skip)
if isinstance(self.start, sp.Basic):
self.start = fast_subs(self.start, subs_dict, skip)
if isinstance(self.stop, sp.Basic):
self.stop = fast_subs(self.stop, subs_dict, skip)
if isinstance(self.step, sp.Basic):
self.step = fast_subs(self.step, subs_dict, skip)
return self
def new_with_inserted_subexpression(self, symbol: sp.Symbol) -> 'AssignmentCollection':
"""Eliminates the subexpression with the given symbol on its left hand side, by substituting it everywhere."""
new_subexpressions = []
subs_dict = None
for se in self.subexpressions:
if se.lhs == symbol:
subs_dict = {se.lhs: se.rhs}
else:
new_subexpressions.append(se)
if subs_dict is None:
return self
new_subexpressions = [Assignment(eq.lhs, fast_subs(eq.rhs, subs_dict)) for eq in new_subexpressions]
new_eqs = [Assignment(eq.lhs, fast_subs(eq.rhs, subs_dict)) for eq in self.main_assignments]
return self.copy(new_eqs, new_subexpressions)
def create_lbm_kernel(collision_rule,
src_field,
dst_field=None,
accessor=StreamPullTwoFieldsAccessor()):
"""Replaces the pre- and post collision symbols in the collision rule by field accesses.
Args:
collision_rule: instance of LbmCollisionRule, defining the collision step
src_field: field used for reading pdf values
dst_field: field used for writing pdf values if accessor.is_inplace this parameter is ignored
accessor: instance of PdfFieldAccessor, defining where to read and write values
to create e.g. a fused stream-collide kernel See 'fieldaccess.PdfFieldAccessor'
Returns:
LbmCollisionRule where pre- and post collision symbols have been replaced
"""
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")
method = collision_rule.method
pre_collision_symbols = method.pre_collision_pdf_symbols
post_collision_symbols = method.post_collision_pdf_symbols
substitutions = {}
input_accesses = accessor.read(src_field, method.stencil)
output_accesses = accessor.write(dst_field, method.stencil)
for (idx,
offset), input_access, output_access in zip(enumerate(method.stencil),
input_accesses,
output_accesses):
substitutions[pre_collision_symbols[idx]] = input_access
substitutions[post_collision_symbols[idx]] = output_access
result = collision_rule.new_with_substitutions(substitutions)
if 'split_groups' in result.simplification_hints:
new_split_groups = []
for split_group in result.simplification_hints['split_groups']:
new_split_groups.append(
[fast_subs(e, substitutions) for e in split_group])
result.simplification_hints['split_groups'] = new_split_groups
if accessor.is_inplace:
result = add_subexpressions_for_field_reads(result,
subexpressions=True,
main_assignments=True)
return result
def discretize_center(term, symbols_to_field_dict, dx, dim=3):
"""
Expects term that contains given symbols and gradient components of these symbols and replaces them
by field accesses. Gradients are replaced by centralized approximations:
``(upper neighbor - lower neighbor ) / ( 2*dx)``
Args:
term: term where symbols and gradient(symbol) should be replaced
symbols_to_field_dict: mapping of symbols to Field
dx: width and height of one cell
dim: dimension
Example:
>>> x = sp.Symbol("x")
>>> grad_x = grad(x, dim=3)
>>> term = x * grad_x[0]
>>> term
x*x^Delta^0
>>> f = Field.create_generic('f', spatial_dimensions=3)
>>> expected_output = f[0, 0, 0] * (-f[-1, 0, 0]/2 + f[1, 0, 0]/2)
>>> sp.simplify(discretize_center(term, { x: f }, dx=1, dim=3) - expected_output)
0
"""
substitutions = {}
for symbols, field in symbols_to_field_dict.items():
if not hasattr(symbols, "__getitem__"):
symbols = [symbols]
g = grad(symbols, dim)
substitutions.update(
{symbol: field(i)
for i, symbol in enumerate(symbols)})
for d in range(dim):
up, down = __up_down_offsets(d, dim)
substitutions.update({
g[d][i]: (field[up](i) - field[down](i)) / dx / 2
for i in range(len(symbols))
})
return fast_subs(term, substitutions)
def update_rule_with_push_boundaries(collision_rule,
field,
boundary_spec,
streaming_pattern='pull',
timestep=Timestep.BOTH):
method = collision_rule.method
accessor = get_accessor(streaming_pattern, timestep)
loads = [
Assignment(a, b) for a, b in zip(method.pre_collision_pdf_symbols,
accessor.read(field, method.stencil))
]
stores = [
Assignment(a, b) for a, b in zip(accessor.write(field, method.stencil),
method.post_collision_pdf_symbols)
]
result = collision_rule.copy()
result.subexpressions = loads + result.subexpressions
result.main_assignments += stores
for direction, boundary in boundary_spec.items():
cond = boundary_conditional(boundary, direction, streaming_pattern,
timestep, method, field)
result.main_assignments.append(cond)
if 'split_groups' in result.simplification_hints:
substitutions = {
b: a
for a, b in zip(accessor.write(field, method.stencil),
method.post_collision_pdf_symbols)
}
new_split_groups = []
for split_group in result.simplification_hints['split_groups']:
new_split_groups.append(
[fast_subs(e, substitutions) for e in split_group])
result.simplification_hints['split_groups'] = new_split_groups
return result
def vectorize_inner_loops_and_adapt_load_stores(ast_node, vector_width, assume_aligned, nontemporal_fields,
strided, keep_loop_stop, assume_sufficient_line_padding):
"""Goes over all innermost loops, changes increment to vector width and replaces field accesses by vector type."""
all_loops = filtered_tree_iteration(ast_node, ast.LoopOverCoordinate, stop_type=ast.SympyAssignment)
inner_loops = [n for n in all_loops if n.is_innermost_loop]
zero_loop_counters = {l.loop_counter_symbol: 0 for l in all_loops}
for loop_node in inner_loops:
loop_range = loop_node.stop - loop_node.start
# cut off loop tail, that is not a multiple of four
if keep_loop_stop:
pass
elif assume_aligned and assume_sufficient_line_padding:
loop_range = loop_node.stop - loop_node.start
new_stop = loop_node.start + modulo_ceil(loop_range, vector_width)
loop_node.stop = new_stop
else:
cutting_point = modulo_floor(loop_range, vector_width) + loop_node.start
loop_nodes = [l for l in cut_loop(loop_node, [cutting_point]).args if isinstance(l, ast.LoopOverCoordinate)]
assert len(loop_nodes) in (0, 1, 2) # 2 for main and tail loop, 1 if loop range divisible by vector width
if len(loop_nodes) == 0:
continue
loop_node = loop_nodes[0]
# Find all array accesses (indexed) that depend on the loop counter as offset
loop_counter_symbol = ast.LoopOverCoordinate.get_loop_counter_symbol(loop_node.coordinate_to_loop_over)
substitutions = {}
successful = True
for indexed in loop_node.atoms(sp.Indexed):
base, index = indexed.args
if loop_counter_symbol in index.atoms(sp.Symbol):
loop_counter_is_offset = loop_counter_symbol not in (index - loop_counter_symbol).atoms()
aligned_access = (index - loop_counter_symbol).subs(zero_loop_counters) == 0
stride = sp.simplify(index.subs({loop_counter_symbol: loop_counter_symbol + 1}) - index)
if not loop_counter_is_offset and (not strided or loop_counter_symbol in stride.atoms()):
successful = False
break
typed_symbol = base.label
assert type(typed_symbol.dtype) is PointerType, \
f"Type of access is {typed_symbol.dtype}, {indexed}"
vec_type = VectorType(typed_symbol.dtype.base_type, vector_width)
use_aligned_access = aligned_access and assume_aligned
nontemporal = False
if hasattr(indexed, 'field'):
nontemporal = (indexed.field in nontemporal_fields) or (indexed.field.name in nontemporal_fields)
substitutions[indexed] = vector_memory_access(indexed, vec_type, use_aligned_access, nontemporal, True,
stride if strided else 1)
if nontemporal:
# insert NontemporalFence after the outermost loop
parent = loop_node.parent
while type(parent.parent.parent) is not ast.KernelFunction:
parent = parent.parent
parent.parent.insert_after(NontemporalFence(), parent, if_not_exists=True)
# insert CachelineSize at the beginning of the kernel
parent.parent.insert_front(CachelineSize(), if_not_exists=True)
if not successful:
warnings.warn("Could not vectorize loop because of non-consecutive memory access")
continue
loop_node.step = vector_width
loop_node.subs(substitutions)
vector_int_width = ast_node.instruction_set['intwidth']
vector_loop_counter = cast_func(loop_counter_symbol, VectorType(loop_counter_symbol.dtype, vector_int_width)) \
+ cast_func(tuple(range(vector_int_width if type(vector_int_width) is int else 2)),
VectorType(loop_counter_symbol.dtype, vector_int_width))
fast_subs(loop_node, {loop_counter_symbol: vector_loop_counter},
skip=lambda e: isinstance(e, ast.ResolvedFieldAccess) or isinstance(e, vector_memory_access))
mask_conditionals(loop_node)
from pystencils.rng import RNGBase
substitutions = {}
for rng in loop_node.atoms(RNGBase):
new_result_symbols = [TypedSymbol(s.name, VectorType(s.dtype, width=vector_width))
for s in rng.result_symbols]
substitutions.update({s[0]: s[1] for s in zip(rng.result_symbols, new_result_symbols)})
rng._symbols_defined = set(new_result_symbols)
fast_subs(loop_node, substitutions, skip=lambda e: isinstance(e, RNGBase))
def fast_subs(self, subs_dict, skip):
rng = copy.deepcopy(self)
rng._args = [fast_subs(a, subs_dict, skip) for a in rng._args]
return rng
def subs(self, subs_dict):
self.lhs = fast_subs(self.lhs, subs_dict)
self.rhs = fast_subs(self.rhs, subs_dict)
def fast_subs(self, subs_dict, skip=None):
self._nodes = [fast_subs(a, subs_dict, skip) for a in self._nodes]
return self
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 discretize_staggered(term,
symbols_to_field_dict,
coordinate,
coordinate_offset,
dx,
dim=3):
"""
Expects term that contains given symbols and gradient components of these symbols and replaces them
by field accesses. Gradients in coordinate direction are replaced by staggered version at cell boundary.
Symbols themselves and gradients in other directions are replaced by interpolated version at cell face.
Args:
term: input term where symbols and gradients are replaced
symbols_to_field_dict: mapping of symbols to Field
coordinate: id for coordinate (0 for x, 1 for y, ... ) defining cell boundary.
Only gradients in this direction are replaced e.g. if symbol^Delta^coordinate
coordinate_offset: either +1 or -1 for upper or lower face in coordinate direction
dx: width and height of one cell
dim: dimension
Examples:
Discretizing at right/east face of cell i.e. coordinate=0, offset=1)
>>> x, dx = sp.symbols("x dx")
>>> grad_x = grad(x, dim=3)
>>> term = x * grad_x[0]
>>> term
x*x^Delta^0
>>> f = Field.create_generic('f', spatial_dimensions=3)
>>> discretize_staggered(term, symbols_to_field_dict={ x: f}, dx=dx, coordinate=0, coordinate_offset=1, dim=3)
(-f_C + f_E)*(f_C/2 + f_E/2)/dx
"""
assert coordinate_offset == 1 or coordinate_offset == -1
assert 0 <= coordinate < dim
substitutions = {}
for symbols, field in symbols_to_field_dict.items():
if not hasattr(symbols, "__getitem__"):
symbols = [symbols]
offset = [0] * dim
offset[coordinate] = coordinate_offset
offset = np.array(offset, dtype=np.int)
gradient = grad(symbols)[coordinate]
substitutions.update({
s: (field[offset](i) + field(i)) / 2
for i, s in enumerate(symbols)
})
substitutions.update({
g: (field[offset](i) - field(i)) / dx * coordinate_offset
for i, g in enumerate(gradient)
})
for d in range(dim):
if d == coordinate:
continue
up, down = __up_down_offsets(d, dim)
for i, s in enumerate(symbols):
center_grad = (field[up](i) - field[down](i)) / (2 * dx)
neighbor_grad = (field[up + offset](i) -
field[down + offset](i)) / (2 * dx)
substitutions[grad(s)[d]] = (center_grad + neighbor_grad) / 2
return fast_subs(term, substitutions)
def equilibrium_exprs(self):
subs_dict = {rr: 1 for rr in self._free_relaxation_rates}
subs_dict.update({rr: 1 for rr in self._fixed_relaxation_rates})
update_equations = self._collisionRule.main_assignments
return [fast_subs(eq.rhs, subs_dict) for eq in update_equations]
