Staggered fields with face/edge neighbor points

pystencils/field.py

@@ -326,6 +326,8 @@ class Field(AbstractField):
             0 for _ in range(self.spatial_dimensions)
         ))  # type: tuple[float,sp.Symbol]
         self.coordinate_transform = sp.eye(self.spatial_dimensions)
+        if field_type == FieldType.STAGGERED:
+            assert self.staggered_stencil
 
     def new_field_with_different_name(self, new_name):
         if self.has_fixed_shape:
@@ -476,10 +478,13 @@ class Field(AbstractField):
                              "Got %d, expected %d" % (len(offset), self.spatial_dimensions))
 
         offset = list(offset)
+        neighbor = [0] * len(offset)
         for i, o in enumerate(offset):
             if (o + sp.Rational(1, 2)).is_Integer:
                 offset[i] += sp.Rational(1, 2)
-                idx = i
+                neighbor[i] = 1
+        neighbor = offset_to_direction_string(neighbor)
+        idx = self.staggered_stencil.index(neighbor)
         offset = tuple(offset)
 
         if self.index_dimensions == 1:  # this field stores a scalar value at each staggered position
@@ -500,6 +505,25 @@ class Field(AbstractField):
 
             return Field.Access(self, offset, (idx, *index))
 
+    @property
+    def staggered_stencil(self):
+        assert FieldType.is_staggered(self)
+        stencils = {
+            2: {
+                2: ["E", "N"],  # D2Q5
+                4: ["E", "N", "NE", "SE"]  # D2Q9
+            },
+            3: {
+                3: ["E", "N", "T"],  # D3Q7
+                7: ["E", "N", "T", "TNE", "BNE", "TSE", "BSE "],  # D3Q15
+                9: ["E", "N", "T", "NE", "SE", "TE", "BE", "TN", "BN"],  # D3Q19
+                13: ["E", "N", "T", "NE", "SE", "TE", "BE", "TN", "BN", "TNE", "BNE", "TSE", "BSE"]  # D3Q27
+            }
+        }
+        if not self.index_shape[0] in stencils[self.spatial_dimensions]:
+            raise ValueError("No known stencil has {} staggered points".format(self.index_shape[0]))
+        return stencils[self.spatial_dimensions][self.index_shape[0]]
+
     def __call__(self, *args, **kwargs):
         center = tuple([0] * self.spatial_dimensions)
         return Field.Access(self, center)(*args, **kwargs)
@@ -746,12 +770,18 @@ class Field(AbstractField):
             super_class_contents = super(Field.Access, self)._hashable_content()
             return (super_class_contents, self._field.hashable_contents(), *self._index, *self._offsets)
 
+        def _staggered_offset(self, offsets, index):
+            assert FieldType.is_staggered(self._field)
+            neighbor = self._field.staggered_stencil[index]
+            neighbor = direction_string_to_offset(neighbor, self._field.spatial_dimensions)
+            return [(o - sp.Rational(neighbor[i], 2)) for i, o in enumerate(offsets)]
+
         def _latex(self, _):
             n = self._field.latex_name if self._field.latex_name else self._field.name
             offset_str = ",".join([sp.latex(o) for o in self.offsets])
             if FieldType.is_staggered(self._field):
-                offset_str = ",".join([sp.latex(o - sp.Rational(int(i == self.index[0]), 2))
-                                       for i, o in enumerate(self.offsets)])
+                offset_str = ",".join([sp.latex(self._staggered_offset(self.offsets, self.index[0])[i])
+                                       for i in range(len(self.offsets))])
             if self.is_absolute_access:
                 offset_str = "\\mathbf{}".format(offset_str)
             elif self.field.spatial_dimensions > 1:
@@ -773,8 +803,8 @@ class Field(AbstractField):
             n = self._field.latex_name if self._field.latex_name else self._field.name
             offset_str = ",".join([sp.latex(o) for o in self.offsets])
             if FieldType.is_staggered(self._field):
-                offset_str = ",".join([sp.latex(o - sp.Rational(int(i == self.index[0]), 2))
-                                       for i, o in enumerate(self.offsets)])
+                offset_str = ",".join([sp.latex(self._staggered_offset(self.offsets, self.index[0])[i])
+                                       for i in range(len(self.offsets))])
             if self.is_absolute_access:
                 offset_str = "[abs]{}".format(offset_str)
 

pystencils_tests/test_field.py

@@ -131,13 +131,21 @@ def test_itemsize():
 
 def test_staggered():
 
+    # D2Q5
     j1, j2, j3 = ps.fields('j1(2), j2(2,2), j3(2,2,2) : double[2D]', field_type=FieldType.STAGGERED)
 
     assert j1[0, 1](1) == j1.staggered_access((0, sp.Rational(1, 2)))
     assert j1[0, 1](1) == j1.staggered_access("N")
+    assert j1[0, 0](1) == j1.staggered_access("S")
 
     assert j2[0, 1](1, 1) == j2.staggered_access((0, sp.Rational(1, 2)), 1)
     assert j2[0, 1](1, 1) == j2.staggered_access("N", 1)
 
     assert j3[0, 1](1, 1, 1) == j3.staggered_access((0, sp.Rational(1, 2)), (1, 1))
     assert j3[0, 1](1, 1, 1) == j3.staggered_access("N", (1, 1))
+
+    # D2Q9
+    k = ps.fields('k(4) : double[2D]', field_type=FieldType.STAGGERED)
+
+    assert k[1, 1](2) == k.staggered_access("NE")
+    assert k[0, 0](2) == k.staggered_access("SW")

pystencils_tests/test_loop_cutting.py

@@ -36,7 +36,7 @@ def test_staggered_iteration():
     fields_fixed = (Field.create_from_numpy_array('f', f_arr),
                     Field.create_from_numpy_array('s', s_arr, index_dimensions=1, field_type=FieldType.STAGGERED))
     fields_var = (Field.create_generic('f', 2),
-                  Field.create_generic('s', 2, index_dimensions=1, field_type=FieldType.STAGGERED))
+                  Field.create_generic('s', 2, index_dimensions=1, index_shape=(dim,), field_type=FieldType.STAGGERED))
 
     for f, s in [fields_var, fields_fixed]:
         # --- Manual