Commit e83dae0d authored by Markus Holzer's avatar Markus Holzer
Browse files

Fix sparse lbm

parent cf17a9b9
Pipeline #32732 failed with stage
in 11 minutes and 7 seconds
...@@ -173,12 +173,11 @@ class UBB(LbBoundary): ...@@ -173,12 +173,11 @@ class UBB(LbBoundary):
def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field): def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
vel_from_idx_field = callable(self._velocity) vel_from_idx_field = callable(self._velocity)
vel = [index_field(f'vel_{i}') for i in range(self.dim)] if vel_from_idx_field else self._velocity vel = [index_field(f'vel_{i}') for i in range(self.dim)] if vel_from_idx_field else self._velocity
direction = dir_symbol
assert self.dim == lb_method.dim, \ assert self.dim == lb_method.dim, \
f"Dimension of UBB ({self.dim}) does not match dimension of method ({lb_method.dim})" f"Dimension of UBB ({self.dim}) does not match dimension of method ({lb_method.dim})"
neighbor_offset = NeighbourOffsetArrays.neighbour_offset(direction, lb_method.stencil) neighbor_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
velocity = tuple(v_i.get_shifted(*neighbor_offset) velocity = tuple(v_i.get_shifted(*neighbor_offset)
if isinstance(v_i, Field.Access) and not vel_from_idx_field if isinstance(v_i, Field.Access) and not vel_from_idx_field
...@@ -193,7 +192,7 @@ class UBB(LbBoundary): ...@@ -193,7 +192,7 @@ class UBB(LbBoundary):
c_s_sq = sp.Rational(1, 3) c_s_sq = sp.Rational(1, 3)
weight_of_direction = LbmWeightInfo.weight_of_direction weight_of_direction = LbmWeightInfo.weight_of_direction
vel_term = 2 / c_s_sq * sum([d_i * v_i for d_i, v_i in zip(neighbor_offset, velocity)]) * weight_of_direction( vel_term = 2 / c_s_sq * sum([d_i * v_i for d_i, v_i in zip(neighbor_offset, velocity)]) * weight_of_direction(
direction, lb_method) dir_symbol, lb_method)
# Better alternative: in conserved value computation # Better alternative: in conserved value computation
# rename what is currently called density to "virtual_density" # rename what is currently called density to "virtual_density"
...@@ -205,12 +204,12 @@ class UBB(LbBoundary): ...@@ -205,12 +204,12 @@ class UBB(LbBoundary):
density_equations = cqc.output_equations_from_pdfs(pdf_field_accesses, {'density': density_symbol}) density_equations = cqc.output_equations_from_pdfs(pdf_field_accesses, {'density': density_symbol})
density_symbol = lb_method.conserved_quantity_computation.defined_symbols()['density'] density_symbol = lb_method.conserved_quantity_computation.defined_symbols()['density']
result = density_equations.all_assignments result = density_equations.all_assignments
result += [Assignment(f_in(inv_dir[direction]), result += [Assignment(f_in(inv_dir[dir_symbol]),
f_out(direction) - vel_term * density_symbol)] f_out(dir_symbol) - vel_term * density_symbol)]
return result return result
else: else:
return [Assignment(f_in(inv_dir[direction]), return [Assignment(f_in(inv_dir[dir_symbol]),
f_out(direction) - vel_term)] f_out(dir_symbol) - vel_term)]
# end class UBB # end class UBB
......
from typing import Tuple from typing import Tuple
import operator
import numpy as np import numpy as np
import sympy as sp import sympy as sp
from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing, NeighbourOffsetArrays
from lbmpy.advanced_streaming.utility import Timestep
from lbmpy.boundaries.boundaryhandling import LbmWeightInfo from lbmpy.boundaries.boundaryhandling import LbmWeightInfo
from pystencils import Assignment, Field, TypedSymbol from pystencils import Assignment, Field, TypedSymbol, create_kernel
from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
from pystencils.boundaries.createindexlist import ( from pystencils.boundaries.createindexlist import (
boundary_index_array_coordinate_names, create_boundary_index_list, direction_member_name) boundary_index_array_coordinate_names, create_boundary_index_list, direction_member_name)
from pystencils.field import FieldType
from pystencils import Assignment, AssignmentCollection
def pdf_index(cell_index, direction_index, domain_size):
return cell_index + direction_index * domain_size
def inverse_idx(stencil, idx):
return stencil.index(tuple(-d_i for d_i in stencil[idx]))
class SparseLbMapper: class SparseLbMapper:
"""Manages the mapping of cell coordinates to indices and back. """Manages the mapping of cell coordinates to indices and back.
...@@ -16,6 +30,7 @@ class SparseLbMapper: ...@@ -16,6 +30,7 @@ class SparseLbMapper:
Args: Args:
flag_arr: integer array where each bit corresponds to a boundary or 'fluid' flag_arr: integer array where each bit corresponds to a boundary or 'fluid'
""" """
def __init__(self, stencil, flag_arr, fluid_flag, no_slip_flag, other_boundary_mask): def __init__(self, stencil, flag_arr, fluid_flag, no_slip_flag, other_boundary_mask):
self._flag_arr = flag_arr self._flag_arr = flag_arr
self._coordinate_arr = None self._coordinate_arr = None
...@@ -106,13 +121,6 @@ class SparseLbMapper: ...@@ -106,13 +121,6 @@ class SparseLbMapper:
def create_index_array(self, ghost_layers=1): def create_index_array(self, ghost_layers=1):
# TODO support different layouts here # TODO support different layouts here
stencil = self.stencil stencil = self.stencil
def pdf_index(cell_index, direction_index):
return cell_index + direction_index * len(self)
def inverse_idx(idx):
return stencil.index(tuple(-d_i for d_i in stencil[idx]))
flag_arr = self.flag_array flag_arr = self.flag_array
no_slip_flag = self.no_slip_flag no_slip_flag = self.no_slip_flag
fluid_boundary_mask = self.other_boundary_mask | self.fluid_flag fluid_boundary_mask = self.other_boundary_mask | self.fluid_flag
...@@ -126,9 +134,9 @@ class SparseLbMapper: ...@@ -126,9 +134,9 @@ class SparseLbMapper:
inv_neighbor_cell = np.array([cell_i - dir_i for cell_i, dir_i in zip(cell, direction)]) inv_neighbor_cell = np.array([cell_i - dir_i for cell_i, dir_i in zip(cell, direction)])
if flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask: if flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask:
neighbor_cell_idx = self.cell_idx(tuple(inv_neighbor_cell)) neighbor_cell_idx = self.cell_idx(tuple(inv_neighbor_cell))
result.append(pdf_index(neighbor_cell_idx, direction_idx)) result.append(pdf_index(neighbor_cell_idx, direction_idx, len(self)))
elif flag_arr[tuple(inv_neighbor_cell)] & no_slip_flag: # no-slip before periodicity! elif flag_arr[tuple(inv_neighbor_cell)] & no_slip_flag: # no-slip before periodicity!
result.append(pdf_index(own_cell_idx, inverse_idx(direction_idx))) result.append(pdf_index(own_cell_idx, inverse_idx(stencil, direction_idx), len(self)))
else: else:
at_border = False at_border = False
for i, x_i in enumerate(inv_neighbor_cell): for i, x_i in enumerate(inv_neighbor_cell):
...@@ -141,7 +149,7 @@ class SparseLbMapper: ...@@ -141,7 +149,7 @@ class SparseLbMapper:
if at_border: if at_border:
assert flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask assert flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask
neighbor_cell_idx = self.cell_idx(tuple(inv_neighbor_cell)) neighbor_cell_idx = self.cell_idx(tuple(inv_neighbor_cell))
result.append(pdf_index(neighbor_cell_idx, direction_idx)) result.append(pdf_index(neighbor_cell_idx, direction_idx, len(self)))
else: else:
raise ValueError("Could not find neighbor for {} direction {}".format(cell, direction)) raise ValueError("Could not find neighbor for {} direction {}".format(cell, direction))
...@@ -153,55 +161,76 @@ class SparseLbMapper: ...@@ -153,55 +161,76 @@ class SparseLbMapper:
class SparseLbBoundaryMapper: class SparseLbBoundaryMapper:
NEIGHBOR_IDX_NAME = 'nidx{}' NEIGHBOR_IDX_NAME = 'nidx{}'
DIR_SYMBOL = TypedSymbol("dir", np.uint32) DIR_SYMBOL = TypedSymbol("dir", np.int64)
def __init__(self, boundary, method, pdf_field_sparse): def __init__(self, boundary, method, pdf_field_sparse):
full_pdf_field = Field.create_generic('pdfFull', spatial_dimensions=method.dim, index_dimensions=1) full_pdf_field = Field.create_generic('pdfFull', spatial_dimensions=method.dim, index_dimensions=1)
full_pdf_field.field_type = FieldType.CUSTOM
prev_timestep = Timestep.BOTH
streaming_pattern = 'pull'
additional_data_field = Field.create_generic('additionalData', spatial_dimensions=1, additional_data_field = Field.create_generic('additionalData', spatial_dimensions=1,
dtype=boundary.additional_data) dtype=boundary.additional_data)
boundary_eqs = boundary(full_pdf_field, self.DIR_SYMBOL, method, additional_data_field)
indexing = BetweenTimestepsIndexing(
full_pdf_field, method.stencil, prev_timestep, streaming_pattern, np.int64, np.int64)
f_out, f_in = indexing.proxy_fields
dir_symbol = indexing.dir_symbol
inv_dir = indexing.inverse_dir_symbol
boundary_eqs = boundary(f_out, f_in, dir_symbol, inv_dir, method, additional_data_field)
neighbor_offsets = {fa.offsets for eq in boundary_eqs for fa in eq.atoms(Field.Access)} neighbor_offsets = {fa.offsets for eq in boundary_eqs for fa in eq.atoms(Field.Access)}
neighbor_offsets = list(neighbor_offsets) neighbor_offsets = list(neighbor_offsets)
neighbor_offsets_dtype = [(self.NEIGHBOR_IDX_NAME.format(i), np.uint32) neighbor_offsets_dtype = [(self.NEIGHBOR_IDX_NAME.format(i), np.int64)
for i in range(len(neighbor_offsets))] for i in range(method.dim)]
index_field_dtype = np.dtype([('dir', np.uint32), index_field_dtype = np.dtype([('dir', np.int64),
*neighbor_offsets_dtype, *neighbor_offsets_dtype,
*boundary.additional_data]) *boundary.additional_data])
index_field = Field.create_generic('indexField', spatial_dimensions=1, dtype=index_field_dtype) index_field = Field.create_generic("idx", spatial_dimensions=1, index_dimensions=1, dtype=np.int64)
boundary_eqs = boundary(full_pdf_field, self.DIR_SYMBOL, method, index_field)
if isinstance(boundary_eqs, Assignment):
offset_subs = {off: sp.Symbol(self.NEIGHBOR_IDX_NAME.format(i)) for i, off in enumerate(neighbor_offsets)} assignments = [boundary_eqs]
new_boundary_eqs = [] if not isinstance(boundary_eqs, AssignmentCollection):
for eq in boundary_eqs: assignments = AssignmentCollection(boundary_eqs)
substitutions = {
fa: pdf_field_sparse.absolute_access([index_field(offset_subs[fa.offsets].name)], fa.index) accessor_subs = dict()
for fa in eq.atoms(Field.Access) for fa in assignments.atoms(Field.Access):
if fa.field == full_pdf_field if fa.field == f_out:
} f_dir = 'out'
new_boundary_eqs.append(eq.subs(substitutions)) elif fa.field == f_in:
f_dir = 'in'
self.boundary_eqs = new_boundary_eqs else:
self.boundary_eqs_orig = boundary_eqs continue
i = 1 if f_dir == "in" else 2
accessor_subs[fa] = pdf_field_sparse.absolute_access((index_field(i),), ())
self.boundary_eqs = assignments.new_with_substitutions(accessor_subs)
self.boundary_eqs_orig = assignments
self.method = method self.method = method
self.index_field_dtype = index_field_dtype self.index_field_dtype = index_field_dtype
self.neighbor_offsets = neighbor_offsets self.neighbor_offsets = neighbor_offsets
self.index_field = index_field self.index_field = index_field
self.timesptep_indexing = indexing
self.boundary_functor = boundary
def _build_substitutions(self): def _build_substitutions(self):
dim = self.method.dim dim = self.method.dim
stencil = self.method.stencil stencil = self.method.stencil
neighbour_offset = NeighbourOffsetArrays(stencil)
result = [{BoundaryOffsetInfo.offset_from_dir(self.DIR_SYMBOL, dim)[current_dim]: offset[current_dim] result = [{self.neighbor_offsets[0][current_dim]: offset[current_dim]
for current_dim in range(dim)} for current_dim in range(dim)}
for i, offset in enumerate(self.method.stencil) for i, offset in enumerate(self.method.stencil)
] ]
for dir_idx, subs_dict in enumerate(result): for dir_idx, subs_dict in enumerate(result):
inv_idx = stencil.index(tuple(-e for e in stencil[dir_idx])) inv_idx = stencil.index(tuple(-e for e in stencil[dir_idx]))
subs_dict[BoundaryOffsetInfo.inv_dir(self.DIR_SYMBOL)] = inv_idx subs_dict[self.timesptep_indexing._index_array_symbol("in", True)] = inv_idx
return result return result
def create_index_arr(self, mapping: SparseLbMapper, boundary_mask, nr_of_ghost_layers=1): def create_index_arr(self, mapping: SparseLbMapper, boundary_mask, nr_of_ghost_layers=1):
...@@ -211,28 +240,36 @@ class SparseLbBoundaryMapper: ...@@ -211,28 +240,36 @@ class SparseLbBoundaryMapper:
flag_dtype(boundary_mask), flag_dtype(mapping.fluid_flag), flag_dtype(boundary_mask), flag_dtype(mapping.fluid_flag),
nr_of_ghost_layers) nr_of_ghost_layers)
result = np.empty(idx_arr.shape, dtype=self.index_field_dtype) result = []
dim = self.method.dim dim = self.method.dim
coord_names = boundary_index_array_coordinate_names[:dim] coord_names = boundary_index_array_coordinate_names[:dim]
center_coordinates = idx_arr[coord_names] center_coordinates = idx_arr[coord_names]
substitutions = self._build_substitutions() result = []
for j, neighbor_offset in enumerate(self.neighbor_offsets): for direction_idx, direction in enumerate(stencil):
neighbor_coordinates = center_coordinates.copy() if all(d_i == 0 for d_i in direction):
offsets = np.array([tuple(int(sp.sympify(e).subs(substitution)) for e in neighbor_offset) assert direction_idx == 0
for substitution in substitutions]) continue
for i, coord_name in enumerate(coord_names): for own_cell_idx, cell in enumerate(mapping.fluid_coordinates):
neighbor_coordinates[coord_name] += offsets[:, i][idx_arr['dir']] inv_neighbor_cell = np.array([cell_i - dir_i for cell_i, dir_i in zip(cell, direction)])
result[self.NEIGHBOR_IDX_NAME.format(j)] = mapping.cell_idx_bulk(neighbor_coordinates) if mapping.flag_array[tuple(inv_neighbor_cell)] & boundary_mask:
write = pdf_index(own_cell_idx, direction_idx, len(mapping))
result[direction_member_name] = idx_arr[direction_member_name] read = pdf_index(own_cell_idx, inverse_idx(stencil, direction_idx), len(mapping))
# Find neighbor indices result.append([direction_idx, write, read])
return result
return np.array(result, dtype=np.int64)
def assignments(self): def assignments(self):
return [BoundaryOffsetInfo(self.method.stencil), return [Assignment(self.DIR_SYMBOL, self.index_field(0)),
LbmWeightInfo(self.method),
Assignment(self.DIR_SYMBOL, self.index_field(self.DIR_SYMBOL.name)),
*self.boundary_eqs] *self.boundary_eqs]
def get_kernel(self):
kernel = create_kernel(self.assignments(), ghost_layers=0)
index_arrs_node = self.timesptep_indexing.create_code_node()
for node in self.boundary_functor.get_additional_code_nodes(self.method)[::-1]:
kernel.body.insert_front(node)
kernel.body.insert_front(index_arrs_node)
return kernel
...@@ -127,7 +127,7 @@ ...@@ -127,7 +127,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.6.5" "version": "3.8.2"
} }
}, },
"nbformat": 4, "nbformat": 4,
......
This source diff could not be displayed because it is too large. You can view the blob instead.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment