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):
def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
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
direction = dir_symbol
assert self.dim == 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)
if isinstance(v_i, Field.Access) and not vel_from_idx_field
......@@ -193,7 +192,7 @@ class UBB(LbBoundary):
c_s_sq = sp.Rational(1, 3)
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(
direction, lb_method)
dir_symbol, lb_method)
# Better alternative: in conserved value computation
# rename what is currently called density to "virtual_density"
......@@ -205,12 +204,12 @@ class UBB(LbBoundary):
density_equations = cqc.output_equations_from_pdfs(pdf_field_accesses, {'density': density_symbol})
density_symbol = lb_method.conserved_quantity_computation.defined_symbols()['density']
result = density_equations.all_assignments
result += [Assignment(f_in(inv_dir[direction]),
f_out(direction) - vel_term * density_symbol)]
result += [Assignment(f_in(inv_dir[dir_symbol]),
f_out(dir_symbol) - vel_term * density_symbol)]
return result
else:
return [Assignment(f_in(inv_dir[direction]),
f_out(direction) - vel_term)]
return [Assignment(f_in(inv_dir[dir_symbol]),
f_out(dir_symbol) - vel_term)]
# end class UBB
......
from typing import Tuple
import operator
import numpy as np
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 pystencils import Assignment, Field, TypedSymbol
from pystencils import Assignment, Field, TypedSymbol, create_kernel
from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
from pystencils.boundaries.createindexlist import (
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:
"""Manages the mapping of cell coordinates to indices and back.
......@@ -16,6 +30,7 @@ class SparseLbMapper:
Args:
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):
self._flag_arr = flag_arr
self._coordinate_arr = None
......@@ -106,13 +121,6 @@ class SparseLbMapper:
def create_index_array(self, ghost_layers=1):
# TODO support different layouts here
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
no_slip_flag = self.no_slip_flag
fluid_boundary_mask = self.other_boundary_mask | self.fluid_flag
......@@ -126,9 +134,9 @@ class SparseLbMapper:
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:
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!
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:
at_border = False
for i, x_i in enumerate(inv_neighbor_cell):
......@@ -141,7 +149,7 @@ class SparseLbMapper:
if at_border:
assert flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask
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:
raise ValueError("Could not find neighbor for {} direction {}".format(cell, direction))
......@@ -153,55 +161,76 @@ class SparseLbMapper:
class SparseLbBoundaryMapper:
NEIGHBOR_IDX_NAME = 'nidx{}'
DIR_SYMBOL = TypedSymbol("dir", np.uint32)
DIR_SYMBOL = TypedSymbol("dir", np.int64)
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_type = FieldType.CUSTOM
prev_timestep = Timestep.BOTH
streaming_pattern = 'pull'
additional_data_field = Field.create_generic('additionalData', spatial_dimensions=1,
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 = list(neighbor_offsets)
neighbor_offsets_dtype = [(self.NEIGHBOR_IDX_NAME.format(i), np.uint32)
for i in range(len(neighbor_offsets))]
neighbor_offsets_dtype = [(self.NEIGHBOR_IDX_NAME.format(i), np.int64)
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,
*boundary.additional_data])
index_field = Field.create_generic('indexField', spatial_dimensions=1, dtype=index_field_dtype)
boundary_eqs = boundary(full_pdf_field, self.DIR_SYMBOL, method, index_field)
offset_subs = {off: sp.Symbol(self.NEIGHBOR_IDX_NAME.format(i)) for i, off in enumerate(neighbor_offsets)}
new_boundary_eqs = []
for eq in boundary_eqs:
substitutions = {
fa: pdf_field_sparse.absolute_access([index_field(offset_subs[fa.offsets].name)], fa.index)
for fa in eq.atoms(Field.Access)
if fa.field == full_pdf_field
}
new_boundary_eqs.append(eq.subs(substitutions))
self.boundary_eqs = new_boundary_eqs
self.boundary_eqs_orig = boundary_eqs
index_field = Field.create_generic("idx", spatial_dimensions=1, index_dimensions=1, dtype=np.int64)
if isinstance(boundary_eqs, Assignment):
assignments = [boundary_eqs]
if not isinstance(boundary_eqs, AssignmentCollection):
assignments = AssignmentCollection(boundary_eqs)
accessor_subs = dict()
for fa in assignments.atoms(Field.Access):
if fa.field == f_out:
f_dir = 'out'
elif fa.field == f_in:
f_dir = 'in'
else:
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.index_field_dtype = index_field_dtype
self.neighbor_offsets = neighbor_offsets
self.index_field = index_field
self.timesptep_indexing = indexing
self.boundary_functor = boundary
def _build_substitutions(self):
dim = self.method.dim
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 i, offset in enumerate(self.method.stencil)
]
for dir_idx, subs_dict in enumerate(result):
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
def create_index_arr(self, mapping: SparseLbMapper, boundary_mask, nr_of_ghost_layers=1):
......@@ -211,28 +240,36 @@ class SparseLbBoundaryMapper:
flag_dtype(boundary_mask), flag_dtype(mapping.fluid_flag),
nr_of_ghost_layers)
result = np.empty(idx_arr.shape, dtype=self.index_field_dtype)
result = []
dim = self.method.dim
coord_names = boundary_index_array_coordinate_names[:dim]
center_coordinates = idx_arr[coord_names]
substitutions = self._build_substitutions()
for j, neighbor_offset in enumerate(self.neighbor_offsets):
neighbor_coordinates = center_coordinates.copy()
offsets = np.array([tuple(int(sp.sympify(e).subs(substitution)) for e in neighbor_offset)
for substitution in substitutions])
for i, coord_name in enumerate(coord_names):
neighbor_coordinates[coord_name] += offsets[:, i][idx_arr['dir']]
result[self.NEIGHBOR_IDX_NAME.format(j)] = mapping.cell_idx_bulk(neighbor_coordinates)
result[direction_member_name] = idx_arr[direction_member_name]
# Find neighbor indices
return result
result = []
for direction_idx, direction in enumerate(stencil):
if all(d_i == 0 for d_i in direction):
assert direction_idx == 0
continue
for own_cell_idx, cell in enumerate(mapping.fluid_coordinates):
inv_neighbor_cell = np.array([cell_i - dir_i for cell_i, dir_i in zip(cell, direction)])
if mapping.flag_array[tuple(inv_neighbor_cell)] & boundary_mask:
write = pdf_index(own_cell_idx, direction_idx, len(mapping))
read = pdf_index(own_cell_idx, inverse_idx(stencil, direction_idx), len(mapping))
result.append([direction_idx, write, read])
return np.array(result, dtype=np.int64)
def assignments(self):
return [BoundaryOffsetInfo(self.method.stencil),
LbmWeightInfo(self.method),
Assignment(self.DIR_SYMBOL, self.index_field(self.DIR_SYMBOL.name)),
return [Assignment(self.DIR_SYMBOL, self.index_field(0)),
*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 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.5"
"version": "3.8.2"
}
},
"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