compiler: Expose separate attribute in SubDomain API

Author: FabioLuporini

devito/types/grid.py

@@ -583,6 +583,14 @@ class SubDomain(AbstractSubDomain):
           region of ``d_size - (N + M)`` points starting at ``N`` and finishing
           at ``d_sizeM - M``.
 
+    Attributes
+    ----------
+    separated : bool, default=True
+        Require a stencil-safe gap between opposite left/right regions, as in
+        SubDimension. Set to False to allow touching or overlapping regions.
+        Applies to SubDimensions generated from tuple definitions; explicit
+        Dimensions returned by :meth:`define` retain their own settings.
+
     Examples
     --------
     An "Inner" SubDomain, which spans the entire domain except for an exterior
@@ -614,6 +622,8 @@ class SubDomain(AbstractSubDomain):
     especially when defining BCs.
     """
 
+    separated = True
+
     def __subdomain_finalize__(self):
         self.__subdomain_finalize_legacy__(self.grid)
         self._distributor = SubDistributor(self)
@@ -655,7 +665,8 @@ def __subdomain_finalize_legacy__(self, grid):
                             f"Maximum thickness of dimension {k.name} "
                             f"is {s}, not {thickness}"
                         ) from None
-                    sub_dimensions.append(constructor(f'i{k.name}', k, thickness))
+                    sub_dimensions.append(constructor(f'i{k.name}', k, thickness,
+                                                      separated=self.separated))
                     sdshape.append(thickness)
                 else:
                     if side != 'middle':
@@ -671,7 +682,8 @@ def __subdomain_finalize_legacy__(self, grid):
                         )
 
                     sub_dimensions.append(
-                        SubDimension.middle(f'i{k.name}', k, ltkn, rtkn)
+                        SubDimension.middle(f'i{k.name}', k, ltkn, rtkn,
+                                            separated=self.separated)
                     )
                     sdshape.append(thickness)
 

tests/test_pickle.py

@@ -49,6 +49,10 @@ def define(self, dimensions):
         return {x: x, y: ('middle', 1, 1), z: ('right', 2)}
 
 
+class OverlappingSD(SD):
+    separated = False
+
+
 @pytest.mark.parametrize('pickle', [pickle0, pickle1])
 class TestBasic:
 
@@ -112,6 +116,16 @@ def test_enrichedtuple_rebuild(self, pickle):
         assert new_t.left == tup.left
         assert new_t.right == tup.right
 
+    def test_subdomain(self, pickle):
+        grid = Grid(shape=(5, 5, 5))
+        sd = OverlappingSD(grid=grid)
+
+        new_sd = pickle.loads(pickle.dumps(sd))
+
+        assert not new_sd.separated
+        assert new_sd.shape == sd.shape
+        assert all(not d.separated for d in new_sd.dimensions if d.is_Sub)
+
     @pytest.mark.parametrize('on_sd', [False, True])
     def test_function(self, pickle, on_sd):
         grid = Grid(shape=(3, 3, 3))

tests/test_subdomains.py

@@ -145,6 +145,31 @@ def make_region():
             expected = size - 16 if spec[0] == 'middle' else 16
             assert make_region().shape == (expected,)
 
+    @pytest.mark.parametrize('legacy', [False, True])
+    @pytest.mark.parametrize('separated', [False, True])
+    def test_separated(self, legacy, separated):
+        class Region(SubDomain):
+            name = 'region'
+
+            def define(self, dimensions):
+                x, y, z = dimensions
+                return {x: ('left', 4), y: ('middle', 1, 1), z: ('right', 4)}
+
+        class OverlappingRegion(Region):
+            separated = False
+
+        cls = Region if separated else OverlappingRegion
+        if legacy:
+            region = cls()
+            Grid(shape=(7, 7, 7), subdomains=(region,))
+        else:
+            region = cls(grid=Grid(shape=(7, 7, 7)))
+
+        assert region.separated is separated
+        assert all(d.separated is separated for d in region.dimensions)
+        assert all(t.separated is separated for d in region.dimensions
+                   for t in d.thickness)
+
     def test_definitions(self):
 
         class sd0(SubDomain):
@@ -289,6 +314,26 @@ def define(self, dimensions):
 
         assert_structure(op, ['t', 'txyz', 'txyz'], 'txyzyz')
 
+    def test_overlapping(self):
+        class Overlapping(ReducedDomain):
+            separated = False
+
+        grid = Grid(shape=(7, 5))
+        left = Overlapping(('left', 4), ('middle', 1, 1), grid=grid)
+        right = Overlapping(('right', 4), None, grid=grid)
+
+        f = Function(name='f', grid=grid)
+
+        eqs = [Eq(f, f + 1, subdomain=left), Eq(f, 2*f + 2, subdomain=right)]
+
+        op = Operator(eqs, name='overlapping_subdomains')
+        op.apply()
+
+        expected = np.zeros(grid.shape)
+        expected[:4, 1:-1] += 1
+        expected[3:, :] = 2*expected[3:, :] + 2
+        assert np.array_equal(f.data, expected)
+
 
 class TestSubDomainScheduling:
     """Tests scheduling across different SubDomains."""