def __init__(self, normal_direction, stencil, name=None):
if isinstance(normal_direction, str):
normal_direction = direction_string_to_offset(normal_direction, dim=len(stencil[0]))
if name is None:
name = f"Simple Outflow: {offset_to_direction_string(normal_direction)}"
self.normal_direction = normal_direction
super(SimpleExtrapolationOutflow, self).__init__(name)
def __getitem__(self, offset):
if type(offset) is np.ndarray:
offset = tuple(offset)
if type(offset) is str:
offset = tuple(
direction_string_to_offset(offset, self.spatial_dimensions))
if type(offset) is not tuple:
offset = (offset, )
if len(offset) != self.spatial_dimensions:
raise ValueError("Wrong number of spatial indices: "
"Got %d, expected %d" %
(len(offset), self.spatial_dimensions))
return Field.Access(self, offset)
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 condition(direction):
"""exclude those staggered points that correspond to fluxes between ghost cells"""
exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
for elementary_direction in direction:
exclusions.remove(inverse_direction_string(elementary_direction))
conditions = []
for e in exclusions:
if e in common_exclusions:
continue
offset = direction_string_to_offset(e)
for i, o in enumerate(offset):
if o == 1:
conditions.append(counters[i] < shape[i] - 1)
elif o == -1:
conditions.append(counters[i] > 0)
return sp.And(*conditions)
def __init__(self, normal_direction, lb_method, dt=1, dx=1, name=None,
streaming_pattern='pull', zeroth_timestep=Timestep.BOTH,
initial_density=None, initial_velocity=None, data_type='double'):
self.lb_method = lb_method
self.stencil = lb_method.stencil
self.dim = len(self.stencil[0])
if isinstance(normal_direction, str):
normal_direction = direction_string_to_offset(normal_direction, dim=self.dim)
if name is None:
name = f"Outflow: {offset_to_direction_string(normal_direction)}"
self.normal_direction = normal_direction
self.streaming_pattern = streaming_pattern
self.zeroth_timestep = zeroth_timestep
self.dx = sp.Number(dx)
self.dt = sp.Number(dt)
self.c = sp.sqrt(sp.Rational(1, 3)) * (self.dx / self.dt)
self.initial_density = initial_density
self.initial_velocity = initial_velocity
self.equilibrium_calculation = None
self.data_type = data_type
if initial_density and initial_velocity:
equilibrium = lb_method.get_equilibrium(conserved_quantity_equations=AssignmentCollection([]))
rho = lb_method.zeroth_order_equilibrium_moment_symbol
u_vec = lb_method.first_order_equilibrium_moment_symbols
eq_lambda = equilibrium.lambdify((rho,) + u_vec)
post_pdf_symbols = lb_method.post_collision_pdf_symbols
def calc_eq_pdfs(density, velocity, j):
return eq_lambda(density, *velocity)[post_pdf_symbols[j]]
self.equilibrium_calculation = calc_eq_pdfs
super(ExtrapolationOutflow, self).__init__(name)
def __init__(self,
neighbor,
dim,
derivative=tuple(),
free_weights_prefix=None):
if type(neighbor) is str:
neighbor = direction_string_to_offset(neighbor)
if dim == 2:
assert neighbor[dim:] == 0
assert derivative is tuple() or max(derivative) < dim
neighbor = sp.Matrix(neighbor[:dim])
pos = neighbor / 2
def unitvec(i):
"""return the `i`-th unit vector in three dimensions"""
a = np.zeros(dim, dtype=int)
a[i] = 1
return a
def flipped(a, i):
"""return `a` with its `i`-th element's sign flipped"""
a = a.copy()
a[i] *= -1
return a
# determine the points to use, coordinates are relative to position
points = []
if np.linalg.norm(neighbor, 1) == 1:
main_points = [neighbor / 2, neighbor / -2]
elif np.linalg.norm(neighbor, 1) == 2:
nonzero_indices = [
i for i, v in enumerate(neighbor) if v != 0 and i < dim
]
main_points = [
neighbor / 2, neighbor / -2,
flipped(neighbor / 2, nonzero_indices[0]),
flipped(neighbor / -2, nonzero_indices[0])
]
else:
main_points = [
sp.Matrix(np.multiply(neighbor,
sp.Matrix(c) / 2))
for c in itertools.product([-1, 1], repeat=3)
]
points += main_points
zero_indices = [
i for i, v in enumerate(neighbor) if v == 0 and i < dim
]
for i in zero_indices:
points += [point + sp.Matrix(unitvec(i)) for point in main_points]
points += [point - sp.Matrix(unitvec(i)) for point in main_points]
points_tuple = tuple([tuple(p) for p in points])
self._stencil = points_tuple
# determine the stencil weights
if len(derivative) == 0:
weights = None
else:
derivation = FiniteDifferenceStencilDerivation(
derivative, points_tuple).get_stencil()
if not derivation.accuracy:
raise Exception(
'the requested derivative cannot be performed with the available neighbors'
)
weights = derivation.weights
# if the weights are underdefined, we can choose the free symbols to find the sparsest stencil
free_weights = set(
itertools.chain(*[w.free_symbols for w in weights]))
if free_weights_prefix is not None:
weights = [
w.subs({
fw: sp.Symbol(f"{free_weights_prefix}_{i}")
for i, fw in enumerate(free_weights)
}) for w in weights
]
elif len(free_weights) > 0:
zero_counts = defaultdict(list)
for values in itertools.product(
[-1, -sp.Rational(1, 2), 0, 1,
sp.Rational(1, 2)],
repeat=len(free_weights)):
subs = {
free_weight: value
for free_weight, value in zip(free_weights, values)
}
weights = [w.subs(subs) for w in derivation.weights]
if not all(a == 0 for a in weights):
zero_count = sum([1 for w in weights if w == 0])
zero_counts[zero_count].append(weights)
best = zero_counts[max(zero_counts.keys())]
if len(
best
) > 1: # if there are multiple, pick the one that contains a nonzero center weight
center = [tuple(p + pos) for p in points].index(
(0, 0, 0)[:dim])
best = [b for b in best if b[center] != 0]
if len(
best
) > 1: # if there are still multiple, they are equivalent, so we average
weights = [
sum([b[i] for b in best]) / len(best)
for i in range(len(weights))
]
else:
weights = best[0]
assert weights
points_tuple = tuple([tuple(p + pos) for p in points])
self._points = points_tuple
self._weights = weights
def create_staggered_kernel(assignments,
target: Target = Target.CPU,
gpu_exclusive_conditions=False,
**kwargs):
"""Kernel that updates a staggered field.
.. image:: /img/staggered_grid.svg
For a staggered field, the first index coordinate defines the location of the staggered value.
Further index coordinates can be used to store vectors/tensors at each point.
Args:
assignments: a sequence of assignments or an AssignmentCollection.
Assignments to staggered field are processed specially, while subexpressions and assignments to
regular fields are passed through to `create_kernel`. Multiple different staggered fields can be
used, but they all need to use the same stencil (i.e. the same number of staggered points) and
shape.
target: 'CPU' or 'GPU'
gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then
handled in an else branch.
kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed
Returns:
AST, see `create_kernel`
"""
if 'ghost_layers' in kwargs:
assert kwargs['ghost_layers'] is None
del kwargs['ghost_layers']
if 'iteration_slice' in kwargs:
assert kwargs['iteration_slice'] is None
del kwargs['iteration_slice']
if 'omp_single_loop' in kwargs:
assert kwargs['omp_single_loop'] is False
del kwargs['omp_single_loop']
if isinstance(assignments, AssignmentCollection):
subexpressions = assignments.subexpressions + [
a for a in assignments.main_assignments
if not hasattr(a, 'lhs') or type(a.lhs) is not Field.Access
or not FieldType.is_staggered(a.lhs.field)
]
assignments = [
a for a in assignments.main_assignments
if hasattr(a, 'lhs') and type(a.lhs) is Field.Access
and FieldType.is_staggered(a.lhs.field)
]
else:
subexpressions = [
a for a in assignments
if not hasattr(a, 'lhs') or type(a.lhs) is not Field.Access
or not FieldType.is_staggered(a.lhs.field)
]
assignments = [
a for a in assignments
if hasattr(a, 'lhs') and type(a.lhs) is Field.Access
and FieldType.is_staggered(a.lhs.field)
]
if len(set([tuple(a.lhs.field.staggered_stencil)
for a in assignments])) != 1:
raise ValueError(
"All assignments need to be made to staggered fields with the same stencil"
)
if len(set([a.lhs.field.shape for a in assignments])) != 1:
raise ValueError(
"All assignments need to be made to staggered fields with the same shape"
)
staggered_field = assignments[0].lhs.field
stencil = staggered_field.staggered_stencil
dim = staggered_field.spatial_dimensions
shape = staggered_field.shape
counters = [
LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)
]
final_assignments = []
# find out whether any of the ghost layers is not needed
common_exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
for direction in stencil:
exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
for elementary_direction in direction:
exclusions.remove(inverse_direction_string(elementary_direction))
common_exclusions.intersection_update(exclusions)
ghost_layers = [[0, 0] for d in range(dim)]
for direction in common_exclusions:
direction = direction_string_to_offset(direction)
for d, s in enumerate(direction):
if s == 1:
ghost_layers[d][1] = 1
elif s == -1:
ghost_layers[d][0] = 1
def condition(direction):
"""exclude those staggered points that correspond to fluxes between ghost cells"""
exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
for elementary_direction in direction:
exclusions.remove(inverse_direction_string(elementary_direction))
conditions = []
for e in exclusions:
if e in common_exclusions:
continue
offset = direction_string_to_offset(e)
for i, o in enumerate(offset):
if o == 1:
conditions.append(counters[i] < shape[i] - 1)
elif o == -1:
conditions.append(counters[i] > 0)
return sp.And(*conditions)
if gpu_exclusive_conditions:
outer_assignment = None
conditions = {direction: condition(direction) for direction in stencil}
for num_conditions in range(len(stencil)):
for combination in itertools.combinations(conditions.values(),
num_conditions):
for assignment in assignments:
direction = stencil[assignment.lhs.index[0]]
if conditions[direction] in combination:
assignment = SympyAssignment(assignment.lhs,
assignment.rhs)
outer_assignment = Conditional(sp.And(*combination),
Block([assignment]),
outer_assignment)
inner_assignment = []
for assignment in assignments:
inner_assignment.append(
SympyAssignment(assignment.lhs, assignment.rhs))
last_conditional = Conditional(
sp.And(*[condition(d) for d in stencil]), Block(inner_assignment),
outer_assignment)
final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
[SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
[last_conditional]
if target == Target.CPU:
from pystencils.cpu import create_kernel as create_kernel_cpu
ast = create_kernel_cpu(final_assignments,
ghost_layers=ghost_layers,
omp_single_loop=False,
**kwargs)
else:
ast = create_kernel(final_assignments,
ghost_layers=ghost_layers,
target=target,
**kwargs)
return ast
for assignment in assignments:
direction = stencil[assignment.lhs.index[0]]
sp_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
[SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
[SympyAssignment(assignment.lhs, assignment.rhs)]
last_conditional = Conditional(condition(direction),
Block(sp_assignments))
final_assignments.append(last_conditional)
remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers])
prepend_optimizations = [
lambda ast: remove_conditionals_in_staggered_kernel(
ast, remove_start_conditional), move_constants_before_loop
]
if 'cpu_prepend_optimizations' in kwargs:
prepend_optimizations += kwargs['cpu_prepend_optimizations']
del kwargs['cpu_prepend_optimizations']
ast = create_kernel(final_assignments,
ghost_layers=ghost_layers,
target=target,
omp_single_loop=False,
cpu_prepend_optimizations=prepend_optimizations,
**kwargs)
return ast