1import sympy as sp 

2 

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 

7 

8 

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 

13 

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 

21 

22 Returns: 

23 list of equations with guarded field accesses 

24 """ 

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

26 fields = [fields] 

27 fields = set(fields) 

28 

29 if type(boundary_flag) is str: 

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

31 

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 

41 

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) 

46 

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

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

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