def test_block():
assignments = [Assignment(dst[0, 0](0), s[0]), Assignment(x, dst[0, 0](2))]
bl = Block(assignments)
assert bl.symbols_defined == {dst[0, 0](0), dst[0, 0](2), s[0], x}
bl.append([Assignment(y, 10)])
assert bl.symbols_defined == {dst[0, 0](0), dst[0, 0](2), s[0], x, y}
assert len(bl.args) == 3
list_iterator = iter([Assignment(s[1], 11)])
bl.insert_front(list_iterator)
assert bl.args[0] == Assignment(s[1], 11)
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 test_subset_cell_values():
"""Tests (un)packing a subset of cell values of the a field (from)to a buffer."""
num_cell_values = 19
# Cell indices of the field to be (un)packed (from)to the buffer
cell_indices = [1, 5, 7, 8, 10, 12, 13]
fields = _generate_fields(num_directions=num_cell_values)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(buffer(buffer_idx), src_field(cell_idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(dst_field(cell_idx), buffer(buffer_idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
mask_arr = np.ma.masked_where((src_arr - dst_arr) != 0, src_arr)
np.testing.assert_equal(dst_arr, mask_arr.filled(int(0)))
def test_all_cell_values():
"""Tests (un)packing all cell values of the a field (from)to a buffer."""
num_cell_values = 7
fields = _generate_fields(stencil_directions=num_cell_values)
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field",
gpu_src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
gpu_src_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=gpu_src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs, type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
dst_arr = gpu_dst_arr.get()
np.testing.assert_equal(src_arr, dst_arr)
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_replace_second_order_products():
x, y = sympy.symbols('x y')
expr = 4 * x * y
expected_expr_positive = 2 * ((x + y)**2 - x**2 - y**2)
expected_expr_negative = 2 * (-(x - y)**2 + x**2 + y**2)
result = replace_second_order_products(expr,
search_symbols=[x, y],
positive=True)
assert result == expected_expr_positive
assert (result - expected_expr_positive).simplify() == 0
result = replace_second_order_products(expr,
search_symbols=[x, y],
positive=False)
assert result == expected_expr_negative
assert (result - expected_expr_negative).simplify() == 0
result = replace_second_order_products(expr,
search_symbols=[x, y],
positive=None)
assert result == expected_expr_positive
a = [Assignment(sympy.symbols('z'), x + y)]
replace_second_order_products(expr,
search_symbols=[x, y],
positive=True,
replace_mixed=a)
assert len(a) == 2
assert replace_second_order_products(4 + y, search_symbols=[x, y]) == y + 4
def test_jacobi_fixed_field_size():
size = (30, 20)
src_field_llvm = np.random.rand(*size)
src_field_py = np.copy(src_field_llvm)
dst_field_llvm = np.zeros(size)
dst_field_py = np.zeros(size)
f = Field.create_from_numpy_array("f", src_field_llvm)
d = Field.create_from_numpy_array("d", dst_field_llvm)
jacobi = Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
ast = create_kernel([jacobi])
for x in range(1, size[0] - 1):
for y in range(1, size[1] - 1):
dst_field_py[
x,
y] = 0.25 * (src_field_py[x - 1, y] + src_field_py[x + 1, y] +
src_field_py[x, y - 1] + src_field_py[x, y + 1])
jit = generate_and_jit(ast)
jit('kernel', dst_field_llvm, src_field_llvm)
error = np.sum(np.abs(dst_field_py - dst_field_llvm))
np.testing.assert_almost_equal(error, 0.0)
def test_jacobi_variable_field_size():
size = (3, 3, 3)
f = Field.create_generic("f", 3)
d = Field.create_generic("d", 3)
jacobi = Assignment(
d[0, 0, 0], (f[1, 0, 0] + f[-1, 0, 0] + f[0, 1, 0] + f[0, -1, 0]) / 4)
ast = create_kernel([jacobi])
src_field_llvm = np.random.rand(*size)
src_field_py = np.copy(src_field_llvm)
dst_field_llvm = np.zeros(size)
dst_field_py = np.zeros(size)
for x in range(1, size[0] - 1):
for y in range(1, size[1] - 1):
for z in range(1, size[2] - 1):
dst_field_py[x, y, z] = 0.25 * (
src_field_py[x - 1, y, z] + src_field_py[x + 1, y, z] +
src_field_py[x, y - 1, z] + src_field_py[x, y + 1, z])
kernel = make_python_function(ast, {
'f': src_field_llvm,
'd': dst_field_llvm
})
kernel()
error = np.sum(np.abs(dst_field_py - dst_field_llvm))
np.testing.assert_almost_equal(error, 0.0)
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 _compute_weights(self):
replacements = self._conservedQuantityComputation.default_values
ac = self.get_equilibrium(include_force_terms=False)
ac = ac.new_with_substitutions(
replacements,
substitute_on_lhs=False).new_without_subexpressions()
new_assignments = [
Assignment(
e.lhs,
subs_additive(e.rhs,
sp.sympify(1),
sum(self.pre_collision_pdf_symbols),
required_match_replacement=1.0))
for e in ac.main_assignments
]
ac = ac.copy(new_assignments)
weights = []
for eq in ac.main_assignments:
value = eq.rhs.expand()
assert len(value.atoms(
sp.Symbol)) == 0, "Failed to compute weights " + str(value)
weights.append(value)
return weights
def _generate_symbolic_relaxation_matrix(self):
"""
This function replaces the numbers in the relaxation matrix with symbols in this case, and returns also
the subexpressions, that assign the number to the newly introduced symbol
"""
rr = [
self.relaxation_matrix[i, i]
for i in range(self.relaxation_matrix.rows)
]
unique_relaxation_rates = set()
subexpressions = {}
for relaxation_rate in rr:
if relaxation_rate not in unique_relaxation_rates:
relaxation_rate = sp.sympify(relaxation_rate)
# special treatment for zero, sp.Zero would be an integer ..
if isinstance(relaxation_rate, Zero):
relaxation_rate = 0.0
if not isinstance(relaxation_rate, sp.Symbol):
rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
subexpressions[relaxation_rate] = rt_symbol
unique_relaxation_rates.add(relaxation_rate)
new_rr = [
subexpressions[sp.sympify(e)]
if sp.sympify(e) in subexpressions else e for e in rr
]
substitutions = [
Assignment(e[1], e[0]) for e in subexpressions.items()
]
return substitutions, sp.diag(*new_rr)
def initializer_kernel_hydro_lb(lb_velocity_field, velocity_field, mrt_method):
r"""
Returns an assignment list for initializing the velocity distribution functions
Args:
lb_velocity_field: source field of velocity distribution function
velocity_field: velocity field
mrt_method: lattice Boltzmann method of the hydrodynamic lattice Boltzmann step
"""
stencil = mrt_method.stencil
weights = get_weights(stencil, c_s_sq=sp.Rational(1, 3))
u_symp = sp.symbols(f"u_:{stencil.D}")
gamma = mrt_method.get_equilibrium_terms()
gamma = gamma.subs({sp.symbols("rho"): 1})
gamma_init = gamma.subs(
{x: y
for x, y in zip(u_symp, velocity_field.center_vector)})
g_updates = list()
for i, _ in enumerate(stencil):
g_updates.append(
Assignment(lb_velocity_field.center(i),
gamma_init[i] - weights[i]))
return g_updates
def add_density_offset(assignment_collection, offset=sp.Rational(1, 1)):
r"""
Assumes that first equation is the density (zeroth moment). Changes the density equations by adding offset to it.
"""
old_eqs = assignment_collection.main_assignments
new_density = Assignment(old_eqs[0].lhs, old_eqs[0].rhs + offset)
return assignment_collection.copy([new_density] + old_eqs[1:])
def apply_force_model_shift(shift_func,
dim,
assignment_collection,
compressible,
reverse=False):
"""
Modifies the first order moment equations in assignment collection according to the force model shift.
It is applied if force model has a method named shift_member_name. The equations 1: dim+1 of the passed
equation collection are assumed to be the velocity equations.
Args:
shift_func: shift function which is applied. See lbmpy.forcemodels.AbstractForceModel for details
dim: number of spatial dimensions
assignment_collection: assignment collection containing the conserved quantity computation
compressible: True if a compressible LB method is used. Otherwise the Helmholtz decomposition was applied
for rho
reverse: If True the sign of the shift is flipped
"""
old_eqs = assignment_collection.main_assignments
density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
old_vel_eqs = old_eqs[1:dim + 1]
vel_offsets = shift_func(density)
if reverse:
vel_offsets = [-v for v in vel_offsets]
shifted_velocity_eqs = [
Assignment(old_eq.lhs, old_eq.rhs + offset)
for old_eq, offset in zip(old_vel_eqs, vel_offsets)
]
new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
return assignment_collection.copy(new_eqs)
def test_indexed_cuda_kernel():
try:
import pycuda
except ImportError:
pycuda = None
if pycuda:
from pystencils.gpucuda import make_python_function
import pycuda.gpuarray as gpuarray
from pystencils.gpucuda.kernelcreation import created_indexed_cuda_kernel
arr = np.zeros((3, 4))
dtype = np.dtype([('x', int), ('y', int), ('value', arr.dtype)])
index_arr = np.zeros((3, ), dtype=dtype)
index_arr[0] = (0, 2, 3.0)
index_arr[1] = (1, 3, 42.0)
index_arr[2] = (2, 1, 5.0)
indexed_field = Field.create_from_numpy_array('index', index_arr)
normal_field = Field.create_from_numpy_array('f', arr)
update_rule = Assignment(normal_field[0, 0], indexed_field('value'))
ast = created_indexed_cuda_kernel([update_rule], [indexed_field])
kernel = make_python_function(ast)
gpu_arr = gpuarray.to_gpu(arr)
gpu_index_arr = gpuarray.to_gpu(index_arr)
kernel(f=gpu_arr, index=gpu_index_arr)
gpu_arr.get(arr)
for i in range(index_arr.shape[0]):
np.testing.assert_allclose(arr[index_arr[i]['x'],
index_arr[i]['y']],
index_arr[i]['value'],
atol=1e-13)
else:
print("Did not run test on GPU since no pycuda is available")
def create_copy_kernel(domain_size,
from_slice,
to_slice,
index_dimensions=0,
index_dim_shape=1,
dtype=np.float64):
"""Copies a rectangular part of an array to another non-overlapping part"""
f = Field.create_generic("pdfs",
len(domain_size),
index_dimensions=index_dimensions,
dtype=dtype)
normalized_from_slice = normalize_slice(from_slice, f.spatial_shape)
normalized_to_slice = normalize_slice(to_slice, f.spatial_shape)
offset = [
s1.start - s2.start
for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
]
assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
"Slices have to have same size"
update_eqs = []
if index_dimensions < 2:
index_dim_shape = [index_dim_shape]
for i in product(*[range(d) for d in index_dim_shape]):
eq = Assignment(f(*i), f[tuple(offset)](*i))
update_eqs.append(eq)
ast = create_cuda_kernel(update_eqs,
iteration_slice=to_slice,
skip_independence_check=True)
return ast
def test_ghost_layer():
size = (6, 5)
src_arr = np.ones(size)
dst_arr = np.zeros_like(src_arr)
src_field = Field.create_from_numpy_array('src',
src_arr,
index_dimensions=0)
dst_field = Field.create_from_numpy_array('dst',
dst_arr,
index_dimensions=0)
update_rule = Assignment(dst_field[0, 0], src_field[0, 0])
ghost_layers = [(1, 2), (2, 1)]
ast = create_cuda_kernel([update_rule],
ghost_layers=ghost_layers,
indexing_creator=LineIndexing)
kernel = make_python_function(ast)
gpu_src_arr = gpuarray.to_gpu(src_arr)
gpu_dst_arr = gpuarray.to_gpu(dst_arr)
kernel(src=gpu_src_arr, dst=gpu_dst_arr)
gpu_dst_arr.get(dst_arr)
reference = np.zeros_like(src_arr)
reference[ghost_layers[0][0]:-ghost_layers[0][1],
ghost_layers[1][0]:-ghost_layers[1][1]] = 1
np.testing.assert_equal(reference, dst_arr)
def create_copy_kernel(domain_size,
from_slice,
to_slice,
index_dimensions=0,
index_dim_shape=1,
dtype=np.float64):
"""Copies a rectangular part of an array to another non-overlapping part"""
if index_dimensions not in (0, 1):
raise NotImplementedError(
"Works only for one or zero index coordinates")
f = Field.create_generic("pdfs",
len(domain_size),
index_dimensions=index_dimensions,
dtype=dtype)
normalized_from_slice = normalize_slice(from_slice, f.spatial_shape)
normalized_to_slice = normalize_slice(to_slice, f.spatial_shape)
offset = [
s1.start - s2.start
for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
]
assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
"Slices have to have same size"
update_eqs = []
for i in range(index_dim_shape):
eq = Assignment(f(i), f[tuple(offset)](i))
update_eqs.append(eq)
ast = create_cuda_kernel(update_eqs,
iteration_slice=to_slice,
skip_independence_check=True)
return make_python_function(ast)
def __init__(self, force, rho=1):
self._force = force
self._rho = rho
self.force_symp = sp.symbols(f"F_:{len(force)}")
self.subs_terms = [
Assignment(rhs, lhs) for rhs, lhs in zip(self.force_symp, force)
]
def test_sliced_iteration_llvm():
size = (4, 4)
src_arr = np.ones(size)
dst_arr = np.zeros_like(src_arr)
src_field = Field.create_from_numpy_array('src', src_arr)
dst_field = Field.create_from_numpy_array('dst', dst_arr)
a, b = sp.symbols("a b")
update_rule = Assignment(dst_field[0, 0],
(a * src_field[0, 1] + a * src_field[0, -1] +
b * src_field[1, 0] + b * src_field[-1, 0]) / 4)
x_end = TypedSymbol("x_end", "int")
s = make_slice[1:x_end, 1]
x_end_value = size[1] - 1
import pystencils.llvm as llvm_generator
ast = llvm_generator.create_kernel(sympy_cse_on_assignment_list(
[update_rule]),
iteration_slice=s)
kernel = llvm_generator.make_python_function(ast)
kernel(src=src_arr, dst=dst_arr, a=1.0, b=1.0, x_end=x_end_value)
expected_result = np.zeros(size)
expected_result[1:x_end_value, 1] = 1
np.testing.assert_almost_equal(expected_result, dst_arr)
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
def test_variable_sized_fields():
src_field = Field.create_generic('src', spatial_dimensions=2)
dst_field = Field.create_generic('dst', spatial_dimensions=2)
update_rule = Assignment(dst_field[0, 0],
(src_field[0, 1] + src_field[0, -1] +
src_field[1, 0] + src_field[-1, 0]) / 4)
ast = create_cuda_kernel(sympy_cse_on_assignment_list([update_rule]))
kernel = make_python_function(ast)
size = (3, 3)
src_arr = np.random.rand(*size)
src_arr = add_ghost_layers(src_arr)
dst_arr = np.zeros_like(src_arr)
gpu_src_arr = gpuarray.to_gpu(src_arr)
gpu_dst_arr = gpuarray.to_gpu(dst_arr)
kernel(src=gpu_src_arr, dst=gpu_dst_arr)
gpu_dst_arr.get(dst_arr)
stencil = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]) / 4.0
reference = convolve(remove_ghost_layers(src_arr),
stencil,
mode='constant',
cval=0.0)
reference = add_ghost_layers(reference)
np.testing.assert_almost_equal(reference, dst_arr)
def test_field_layouts():
num_cell_values = 7
for layout_str in ['numpy', 'fzyx', 'zyxf', 'reverse_numpy']:
fields = _generate_fields(stencil_directions=num_cell_values,
layout=layout_str)
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field",
gpu_src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
gpu_src_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs,
type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
def test_all_cell_values():
"""Tests (un)packing all cell values of the a field (from)to a buffer."""
num_cell_values = 19
fields = _generate_fields(num_directions=num_cell_values)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
np.testing.assert_equal(src_arr, dst_arr)
def test_field_slice():
"""Tests (un)packing slices of a scalar field (from)to a buffer."""
fields = _generate_fields()
for d in ['N', 'S', 'NW', 'SW', 'TNW', 'B']:
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
# Extract slice from N direction of the field
slice_dir = direction_string_to_offset(d, dim=len(src_arr.shape))
pack_slice = get_slice_before_ghost_layer(slice_dir)
unpack_slice = get_ghost_region_slice(slice_dir)
src_field = Field.create_from_numpy_array("src_field",
src_arr[pack_slice])
dst_field = Field.create_from_numpy_array("dst_field",
src_arr[unpack_slice])
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = [Assignment(buffer.center(), src_field.center())]
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr,
src_field=gpu_src_arr[pack_slice])
# Unpack into ghost layer of dst_field in N direction
unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs,
type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr,
dst_field=gpu_dst_arr[unpack_slice])
dst_arr = gpu_dst_arr.get()
np.testing.assert_equal(src_arr[pack_slice], dst_arr[unpack_slice])
def __call__(self, field, direction_symbol, **kwargs):
neighbor = BoundaryOffsetInfo.offset_from_dir(direction_symbol,
field.spatial_dimensions)
if field.index_dimensions == 0:
return [Assignment(field.center, field[neighbor])]
else:
from itertools import product
if not field.has_fixed_index_shape:
raise NotImplementedError(
"Neumann boundary works only for fields with fixed index shape"
)
index_iter = product(*(range(i) for i in field.index_shape))
return [
Assignment(field(*idx), field[neighbor](*idx))
for idx in index_iter
]
def test_field_layouts():
num_cell_values = 27
for layout_str in ['numpy', 'fzyx', 'zyxf', 'reverse_numpy']:
fields = _generate_fields(num_directions=num_cell_values,
layout=layout_str)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
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 using the backward
fast central moment transform.
First, monomial central moments are obtained from the polynomial moments by multiplication
with :math:`P^{-1}`. Then, the elementwise equations of the matrix
multiplication :math:`K^{-1} \cdot \mathbf{K}` with the monomial central moment matrix
(see `PdfsToCentralMomentsByMatrix`) are recursively simplified by extracting certain linear
combinations of velocities, to obtain equations similar to the ones given in :cite:`geier2015`.
The backward transform is designed for D3Q27, inherently generalizes to D2Q9, and is tested
for D3Q19. It also returns correct equations for D3Q15, whose efficiency is however questionable.
**De-Aliasing**:
See `FastCentralMomentTransform.forward_transform`.
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:
monomial_eqs = self.poly_to_mono_matrix * sp.Matrix(post_collision_moments)
subexpressions += [Assignment(m, v) for m, v in zip(post_collision_monomial_moments, monomial_eqs)]
raw_equations = self.inv_monomial_matrix * sp.Matrix(post_collision_monomial_moments)
raw_equations = [Assignment(f, eq) for f, eq in zip(pdf_symbols, raw_equations)]
symbol_gen = SymbolGen(subexpression_base)
ac = self._split_backward_equations(raw_equations, symbol_gen)
ac.subexpressions = subexpressions + ac.subexpressions
if simplification:
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 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 mu_kernel(free_energy, order_parameters, phi_field, mu_field, dx=1, discretization='standard'):
"""Reads from order parameter (phi) field and updates chemical potentials"""
assert phi_field.spatial_dimensions == mu_field.spatial_dimensions
dim = phi_field.spatial_dimensions
chemical_potential = chemical_potentials_from_free_energy(free_energy, order_parameters)
chemical_potential = substitute_laplacian_by_sum(chemical_potential, dim)
chemical_potential = chemical_potential.subs({op: phi_field(i) for i, op in enumerate(order_parameters)})
return [Assignment(mu_field(i), discretize_spatial(mu_i, dx, discretization))
for i, mu_i in enumerate(chemical_potential)]
