def test_staggered_iteration_manual():
dim = 2
f_arr = np.arange(5**dim).reshape([5] * dim)
s_arr = np.ones([5] * dim + [dim]) * 1234
s_arr_ref = s_arr.copy()
f = Field.create_from_numpy_array('f', f_arr)
s = Field.create_from_numpy_array('s', s_arr, index_dimensions=1)
eqs = []
counters = [
LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)
]
conditions = [counters[i] < f.shape[i] - 1 for i in range(dim)]
for d in range(dim):
eq = SympyAssignment(
s(d),
sum(f[o] for o in offsets_in_plane(d, 0, dim)) -
sum(f[o] for o in offsets_in_plane(d, -1, dim)))
cond = sp.And(*[conditions[i] for i in range(dim) if d != i])
eqs.append(Conditional(cond, eq))
kernel_ast = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)])
func = make_python_function(kernel_ast)
func(f=f_arr, s=s_arr_ref)
inner_loop = [
n for n in kernel_ast.atoms(ast.LoopOverCoordinate)
if n.is_innermost_loop
][0]
cut_loop(inner_loop, [4])
outer_loop = [
n for n in kernel_ast.atoms(ast.LoopOverCoordinate)
if n.is_outermost_loop
][0]
cut_loop(outer_loop, [4])
simplify_conditionals(kernel_ast.body, loop_counter_simplification=True)
cleanup_blocks(kernel_ast.body)
move_constants_before_loop(kernel_ast.body)
cleanup_blocks(kernel_ast.body)
assert not kernel_ast.atoms(
Conditional), "Loop cutting optimization did not work"
func_optimized = make_python_function(kernel_ast)
func_optimized(f=f_arr, s=s_arr)
np.testing.assert_almost_equal(s_arr_ref, s_arr)
def test_staggered_iteration():
dim = 2
f_arr = np.arange(5**dim).reshape([5] * dim).astype(np.float64)
s_arr = np.ones([5] * dim + [dim]) * 1234
s_arr_ref = s_arr.copy()
fields_fixed = (Field.create_from_numpy_array('f', f_arr),
Field.create_from_numpy_array('s',
s_arr,
index_dimensions=1))
fields_var = (Field.create_generic('f', 2),
Field.create_generic('s', 2, index_dimensions=1))
for f, s in [fields_var, fields_fixed]:
# --- Manual
eqs = []
counters = [
LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)
]
conditions = [counters[i] < f.shape[i] - 1 for i in range(dim)]
for d in range(dim):
eq = SympyAssignment(
s(d),
sum(f[o] for o in offsets_in_plane(d, 0, dim)) -
sum(f[o] for o in offsets_in_plane(d, -1, dim)))
cond = sp.And(*[conditions[i] for i in range(dim) if d != i])
eqs.append(Conditional(cond, eq))
func = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0),
(1, 0)]).compile()
# --- Built-in optimized
expressions = []
for d in range(dim):
expressions.append(
sum(f[o] for o in offsets_in_plane(d, 0, dim)) -
sum(f[o] for o in offsets_in_plane(d, -1, dim)))
func_optimized = create_staggered_kernel(s, expressions).compile()
assert not func_optimized.ast.atoms(
Conditional), "Loop cutting optimization did not work"
func(f=f_arr, s=s_arr_ref)
func_optimized(f=f_arr, s=s_arr)
np.testing.assert_almost_equal(s_arr_ref, s_arr)
def test_jacobi_fixed_field_size():
size = (30, 20)
src_field_c = np.random.rand(*size)
src_field_py = np.copy(src_field_c)
dst_field_c = np.zeros(size)
dst_field_py = np.zeros(size)
f = Field.create_from_numpy_array("f", src_field_c)
d = Field.create_from_numpy_array("d", dst_field_c)
jacobi = SympyAssignment(d[0, 0],
(f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
body = Block([jacobi])
loop_node, gl_info = make_loop_over_domain(body)
ast_node = KernelFunction(loop_node,
'cpu',
'c',
make_python_function,
ghost_layers=gl_info)
resolve_field_accesses(ast_node)
move_constants_before_loop(ast_node)
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])
kernel = ast_node.compile()
kernel(f=src_field_c, d=dst_field_c)
error = np.sum(np.abs(dst_field_py - dst_field_c))
np.testing.assert_allclose(error, 0.0, atol=1e-13)
code_display = show_code(ast_node)
assert 'for' in str(code_display)
assert 'for' in code_display._repr_html_()
def add_array(self,
name,
values_per_cell=1,
dtype=np.float64,
latex_name=None,
ghost_layers=None,
layout=None,
cpu=True,
gpu=None,
alignment=False,
field_type=FieldType.GENERIC):
if ghost_layers is None:
ghost_layers = self.default_ghost_layers
if layout is None:
layout = self.default_layout
if gpu is None:
gpu = self.default_target in self._GPU_LIKE_TARGETS
kwargs = {
'shape': tuple(s + 2 * ghost_layers for s in self._domainSize),
'dtype': dtype,
}
if not hasattr(values_per_cell, '__len__'):
values_per_cell = (values_per_cell, )
if len(values_per_cell) == 1 and values_per_cell[0] == 1:
values_per_cell = ()
self._field_information[name] = {
'ghost_layers': ghost_layers,
'values_per_cell': values_per_cell,
'layout': layout,
'dtype': dtype,
'alignment': alignment,
'field_type': field_type,
}
index_dimensions = len(values_per_cell)
kwargs['shape'] = kwargs['shape'] + values_per_cell
if index_dimensions > 0:
layout_tuple = layout_string_to_tuple(layout,
self.dim + index_dimensions)
else:
layout_tuple = spatial_layout_string_to_tuple(layout, self.dim)
# cpu_arr is always created - since there is no create_pycuda_array_with_layout()
byte_offset = ghost_layers * np.dtype(dtype).itemsize
cpu_arr = create_numpy_array_with_layout(layout=layout_tuple,
alignment=alignment,
byte_offset=byte_offset,
**kwargs)
if alignment and gpu:
raise NotImplementedError("Alignment for GPU fields not supported")
if cpu:
if name in self.cpu_arrays:
raise ValueError("CPU Field with this name already exists")
self.cpu_arrays[name] = cpu_arr
if gpu:
if name in self.gpu_arrays:
raise ValueError("GPU Field with this name already exists")
self.gpu_arrays[name] = self.array_handler.to_gpu(cpu_arr)
assert all(f.name != name for f in self.fields.values()
), "Symbolic field with this name already exists"
self.fields[name] = Field.create_from_numpy_array(
name,
cpu_arr,
index_dimensions=index_dimensions,
field_type=field_type)
self.fields[name].latex_name = latex_name
return self.fields[name]
def compile_macroscopic_values_getter(lb_method,
output_quantities,
pdf_arr=None,
ghost_layers=1,
iteration_slice=None,
field_layout='numpy',
target=Target.CPU,
streaming_pattern='pull',
previous_timestep=Timestep.BOTH):
"""
Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
Args:
lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
output_quantities: sequence of quantities to compute e.g. ['density', 'velocity']
pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
that should be excluded from the iteration. If None, the number of ghost layers
is determined automatically and assumed to be equal for all dimensions.
iteration_slice: if not None, iteration is done only over this slice of the field
field_layout: layout for output field, also used for pdf field if pdf_arr is not given
target: `Target.CPU` or `Target.GPU`
previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
Returns:
a function to compute macroscopic values:
- pdf_array
- keyword arguments from name of conserved quantity (as in output_quantities) to numpy field
"""
if not (isinstance(output_quantities, list)
or isinstance(output_quantities, tuple)):
output_quantities = [output_quantities]
cqc = lb_method.conserved_quantity_computation
unknown_quantities = [
oq for oq in output_quantities if oq not in cqc.conserved_quantities
]
if unknown_quantities:
raise ValueError(
"No such conserved quantity: %s, conserved quantities are %s" %
(str(unknown_quantities), str(cqc.conserved_quantities.keys())))
if pdf_arr is None:
pdf_field = Field.create_generic('pdfs',
lb_method.dim,
index_dimensions=1,
layout=field_layout)
else:
pdf_field = Field.create_from_numpy_array('pdfs',
pdf_arr,
index_dimensions=1)
output_mapping = {}
for output_quantity in output_quantities:
number_of_elements = cqc.conserved_quantities[output_quantity]
assert number_of_elements >= 1
ind_dims = 0 if number_of_elements <= 1 else 1
if pdf_arr is None:
output_field = Field.create_generic(output_quantity,
lb_method.dim,
layout=field_layout,
index_dimensions=ind_dims)
else:
output_field_shape = pdf_arr.shape[:-1]
if ind_dims > 0:
output_field_shape += (number_of_elements, )
field_layout = get_layout_of_array(pdf_arr)
else:
field_layout = get_layout_of_array(
pdf_arr, index_dimension_ids=[len(pdf_field.shape) - 1])
output_field = Field.create_fixed_size(output_quantity,
output_field_shape,
ind_dims, pdf_arr.dtype,
field_layout)
output_mapping[output_quantity] = [
output_field(i) for i in range(number_of_elements)
]
if len(output_mapping[output_quantity]) == 1:
output_mapping[output_quantity] = output_mapping[output_quantity][
0]
stencil = lb_method.stencil
previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
pdf_symbols = previous_step_accessor.write(pdf_field, stencil)
eqs = cqc.output_equations_from_pdfs(pdf_symbols,
output_mapping).all_assignments
if target == Target.CPU:
import pystencils.cpu as cpu
kernel = cpu.make_python_function(
cpu.create_kernel(eqs,
ghost_layers=ghost_layers,
iteration_slice=iteration_slice))
elif target == Target.GPU:
import pystencils.gpucuda as gpu
kernel = gpu.make_python_function(
gpu.create_cuda_kernel(eqs,
ghost_layers=ghost_layers,
iteration_slice=iteration_slice))
else:
raise ValueError(
"Unknown target '%s'. Possible targets are `Target.CPU` and `Target.GPU`"
% (target, ))
def getter(pdfs, **kwargs):
if pdf_arr is not None:
assert pdfs.shape == pdf_arr.shape and pdfs.strides == pdf_arr.strides, \
"Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdf_arr.shape)
if not set(output_quantities).issubset(kwargs.keys()):
raise ValueError(
"You have to specify the output field for each of the following quantities: %s"
% (str(output_quantities), ))
kernel(pdfs=pdfs, **kwargs)
return getter
def compile_macroscopic_values_setter(lb_method,
quantities_to_set,
pdf_arr=None,
ghost_layers=1,
iteration_slice=None,
field_layout='numpy',
target=Target.CPU,
streaming_pattern='pull',
previous_timestep=Timestep.BOTH):
"""
Creates a function that sets a pdf field to specified macroscopic quantities
The returned function can be called with the pdf field to set as single argument
Args:
lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
quantities_to_set: map from conserved quantity name to fixed value or numpy array
pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
that should be excluded from the iteration. If None, the number of ghost layers
is determined automatically and assumed to be equal for all dimensions.
iteration_slice: if not None, iteration is done only over this slice of the field
field_layout: layout of the pdf field if pdf_arr was not given
target: `Target.CPU` or `Target.GPU`
previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
Returns:
function taking pdf array as single argument and which sets the field to the given values
"""
if pdf_arr is not None:
pdf_field = Field.create_from_numpy_array('pdfs',
pdf_arr,
index_dimensions=1)
else:
pdf_field = Field.create_generic('pdfs',
lb_method.dim,
index_dimensions=1,
layout=field_layout)
fixed_kernel_parameters = {}
cqc = lb_method.conserved_quantity_computation
value_map = {}
at_least_one_field_input = False
for quantity_name, value in quantities_to_set.items():
if hasattr(value, 'shape'):
fixed_kernel_parameters[quantity_name] = value
at_least_one_field_input = True
num_components = cqc.conserved_quantities[quantity_name]
field = Field.create_from_numpy_array(
quantity_name,
value,
index_dimensions=0 if num_components <= 1 else 1)
if num_components == 1:
value = field(0)
else:
value = [field(i) for i in range(num_components)]
value_map[quantity_name] = value
cq_equations = cqc.equilibrium_input_equations_from_init_values(
**value_map, force_substitution=False)
eq = lb_method.get_equilibrium(conserved_quantity_equations=cq_equations)
if at_least_one_field_input:
simplification = create_simplification_strategy(lb_method)
eq = simplification(eq)
else:
eq = eq.new_without_subexpressions()
previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
substitutions = {
sym: write_accesses[i]
for i, sym in enumerate(lb_method.post_collision_pdf_symbols)
}
eq = eq.new_with_substitutions(substitutions).all_assignments
if target == Target.CPU:
import pystencils.cpu as cpu
kernel = cpu.make_python_function(cpu.create_kernel(eq))
kernel = functools.partial(kernel, **fixed_kernel_parameters)
elif target == Target.GPU:
import pystencils.gpucuda as gpu
kernel = gpu.make_python_function(gpu.create_cuda_kernel(eq))
kernel = functools.partial(kernel, **fixed_kernel_parameters)
else:
raise ValueError(
"Unknown target '%s'. Possible targets are `Target.CPU` and `Target.GPU`"
% (target, ))
def setter(pdfs, **kwargs):
if pdf_arr is not None:
assert pdfs.shape == pdf_arr.shape and pdfs.strides == pdf_arr.strides, \
"Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdf_arr.shape)
kernel(pdfs=pdfs, **kwargs)
return setter
def test_error_handling():
struct_dtype = np.dtype([('a', np.int32), ('b', np.float64),
('c', np.uint32)])
Field.create_generic('f',
spatial_dimensions=2,
index_dimensions=0,
dtype=struct_dtype)
with pytest.raises(ValueError) as e:
Field.create_generic('f',
spatial_dimensions=2,
index_dimensions=1,
dtype=struct_dtype)
assert 'index dimension' in str(e.value)
arr = np.array([[1, 2.0, 3], [1, 2.0, 3]], dtype=struct_dtype)
Field.create_from_numpy_array('f', arr, index_dimensions=0)
with pytest.raises(ValueError) as e:
Field.create_from_numpy_array('f', arr, index_dimensions=1)
assert 'Structured arrays' in str(e.value)
arr = np.zeros([3, 3, 3])
Field.create_from_numpy_array('f', arr, index_dimensions=2)
with pytest.raises(ValueError) as e:
Field.create_from_numpy_array('f', arr, index_dimensions=3)
assert 'Too many' in str(e.value)
Field.create_fixed_size('f', (3, 2, 4),
index_dimensions=0,
dtype=struct_dtype,
layout='reverse_numpy')
with pytest.raises(ValueError) as e:
Field.create_fixed_size('f', (3, 2, 4),
index_dimensions=1,
dtype=struct_dtype,
layout='reverse_numpy')
assert 'Structured arrays' in str(e.value)
f = Field.create_fixed_size('f', (10, 10))
with pytest.raises(ValueError) as e:
f[1]
assert 'Wrong number of spatial indices' in str(e.value)
f = Field.create_generic('f', spatial_dimensions=2, index_shape=(3, ))
with pytest.raises(ValueError) as e:
f(3)
assert 'out of bounds' in str(e.value)
f = Field.create_fixed_size('f', (10, 10, 3, 4), index_dimensions=2)
with pytest.raises(ValueError) as e:
f(3, 0)
assert 'out of bounds' in str(e.value)
with pytest.raises(ValueError) as e:
f(1, 0)(1, 0)
assert 'Indexing an already indexed' in str(e.value)
with pytest.raises(ValueError) as e:
f(1)
assert 'Wrong number of indices' in str(e.value)
with pytest.raises(ValueError) as e:
Field.create_generic('f', spatial_dimensions=2, layout='wrong')
assert 'Unknown layout descriptor' in str(e.value)
assert layout_string_to_tuple('fzyx', dim=4) == (3, 2, 1, 0)
with pytest.raises(ValueError) as e:
layout_string_to_tuple('wrong', dim=4)
assert 'Unknown layout descriptor' in str(e.value)