1import sympy as sp

3from pystencils.boundaries.boundaryhandling import DEFAULT_FLAG_TYPE

4from pystencils.data_types import TypedSymbol, create_type

5from pystencils.field import Field

6from pystencils.integer_functions import bitwise_and

9def add_neumann_boundary(eqs, fields, flag_field, boundary_flag="neumann_flag", inverse_flag=False):

10 """

11 Replaces all neighbor accesses by flag field guarded accesses.

12 If flag in neighboring cell is set, the center value is used instead

14 Args:

15 eqs: list of equations containing field accesses to direct neighbors

16 fields: fields for which the Neumann boundary should be applied

17 flag_field: integer field marking boundary cells

18 boundary_flag: if flag field has value 'boundary_flag' (no bit operations yet)

19 the cell is assumed to be boundary

20 inverse_flag: if true, boundary cells are where flag field has not the value of boundary_flag

22 Returns:

23 list of equations with guarded field accesses

24 """

25 if not hasattr(fields, "__len__"):

26 fields = [fields]

27 fields = set(fields)

29 if type(boundary_flag) is str:

30 boundary_flag = TypedSymbol(boundary_flag, dtype=create_type(DEFAULT_FLAG_TYPE))

32 substitutions = {}

33 for eq in eqs:

34 for fa in eq.atoms(Field.Access):

35 if fa.field not in fields:

36 continue

37 if not all(offset in (-1, 0, 1) for offset in fa.offsets):

38 raise ValueError("Works only for single neighborhood stencils")

39 if all(offset == 0 for offset in fa.offsets):

40 continue

42 if inverse_flag:

43 condition = sp.Eq(bitwise_and(flag_field[tuple(fa.offsets)], boundary_flag), 0)

44 else:

45 condition = sp.Ne(bitwise_and(flag_field[tuple(fa.offsets)], boundary_flag), 0)

47 center = fa.field(*fa.index)

48 substitutions[fa] = sp.Piecewise((center, condition), (fa, True))

49 return [eq.subs(substitutions) for eq in eqs]