dsl: Introduce ability to define Functions on Subdomains

devito/data/data.py

@@ -237,8 +237,9 @@ def __getitem__(self, glb_idx, comm_type, gather_rank=None):
                     stop += sendcounts[i]
                     data_slice = recvbuf[slice(start, stop, step)]
                     shape = [r.stop-r.start for r in self._distributor.all_ranges[i]]
-                    idx = [slice(r.start, r.stop, r.step)
-                           for r in self._distributor.all_ranges[i]]
+                    idx = [slice(r.start - d.glb_min, r.stop - d.glb_min, r.step)
+                           for r, d in zip(self._distributor.all_ranges[i],
+                                           self._distributor.decomposition)]
                     for i in range(len(self.shape) - len(self._distributor.glb_shape)):
                         shape.insert(i, glb_shape[i])
                         idx.insert(i, slice(0, glb_shape[i]+1, 1))
@@ -328,6 +329,7 @@ def __getitem__(self, glb_idx, comm_type, gather_rank=None):
     @_check_idx
     def __setitem__(self, glb_idx, val, comm_type):
         loc_idx = self._index_glb_to_loc(glb_idx)
+
         if loc_idx is NONLOCAL:
             # no-op
             return

devito/data/decomposition.py

@@ -204,9 +204,20 @@ def index_glb_to_loc(self, *args, rel=True):
         >>> d.index_glb_to_loc((1, 6), rel=False)
         (5, 6)
         """
+        # Need to offset the loc_abs_min, loc_abs_max, glb_min, and glb_max
+        # Distributed indices may not go from 0-i_M, but array indices will
+        # Subtract glb_min from all of these
+        if not self.loc_empty:
+            loc_abs_min = self.loc_abs_min - self.glb_min
+            loc_abs_max = self.loc_abs_max - self.glb_min
+        else:
+            loc_abs_min = self.loc_abs_min
+            loc_abs_max = self.loc_abs_max
+
+        glb_max = self.glb_max - self.glb_min
 
-        base = self.loc_abs_min if rel is True else 0
-        top = self.loc_abs_max
+        base = loc_abs_min if rel is True else 0
+        top = loc_abs_max
 
         if len(args) == 1:
             glb_idx = args[0]
@@ -217,11 +228,11 @@ def index_glb_to_loc(self, *args, rel=True):
                     return None
                 # -> Handle negative index
                 if glb_idx < 0:
-                    glb_idx = self.glb_max + glb_idx + 1
+                    glb_idx = glb_max + glb_idx + 1
                 # -> Do the actual conversion
-                if self.loc_abs_min <= glb_idx <= self.loc_abs_max:
+                if loc_abs_min <= glb_idx <= loc_abs_max:
                     return glb_idx - base
-                elif self.glb_min <= glb_idx <= self.glb_max:
+                elif 0 <= glb_idx <= glb_max:
                     return None
                 else:
                     # This should raise an exception when used to access a numpy.array
@@ -239,30 +250,36 @@ def index_glb_to_loc(self, *args, rel=True):
                 elif isinstance(glb_idx, slice):
                     if self.loc_empty:
                         return slice(-1, -3)
-                    if glb_idx.step >= 0 and glb_idx.stop == self.glb_min:
-                        glb_idx_min = self.glb_min if glb_idx.start is None \
+                    if glb_idx.step >= 0 and glb_idx.stop == 0:
+                        glb_idx_min = 0 if glb_idx.start is None \
                             else glb_idx.start
-                        glb_idx_max = self.glb_min
+                        glb_idx_max = 0
                         retfunc = lambda a, b: slice(a, b, glb_idx.step)
                     elif glb_idx.step >= 0:
-                        glb_idx_min = self.glb_min if glb_idx.start is None \
+                        glb_idx_min = 0 if glb_idx.start is None \
                             else glb_idx.start
-                        glb_idx_max = self.glb_max if glb_idx.stop is None \
+                        glb_idx_max = glb_max \
+                            if glb_idx.stop is None \
+                            else glb_idx.stop-1
+                        glb_idx_min = 0 if glb_idx.start is None \
+                            else glb_idx.start
+                        glb_idx_max = glb_max if glb_idx.stop is None \
                             else glb_idx.stop-1
                         retfunc = lambda a, b: slice(a, b + 1, glb_idx.step)
                     else:
-                        glb_idx_min = self.glb_min if glb_idx.stop is None \
+                        glb_idx_min = 0 if glb_idx.stop is None \
                             else glb_idx.stop+1
-                        glb_idx_max = self.glb_max if glb_idx.start is None \
+                        glb_idx_max = glb_max if glb_idx.start is None \
                             else glb_idx.start
                         retfunc = lambda a, b: slice(b, a - 1, glb_idx.step)
                 else:
                     raise TypeError("Cannot convert index from `%s`" % type(glb_idx))
                 # -> Handle negative min/max
                 if glb_idx_min is not None and glb_idx_min < 0:
-                    glb_idx_min = self.glb_max + glb_idx_min + 1
+                    glb_idx_min = glb_max + glb_idx_min + 1
                 if glb_idx_max is not None and glb_idx_max < 0:
-                    glb_idx_max = self.glb_max + glb_idx_max + 1
+                    glb_idx_max = glb_max + glb_idx_max + 1
+
                 # -> Do the actual conversion
                 # Compute loc_min. For a slice with step > 0 this will be
                 # used to produce slice.start and for a slice with step < 0 slice.stop.
@@ -271,19 +288,19 @@ def index_glb_to_loc(self, *args, rel=True):
                 # coincide with loc_abs_min.
                 if isinstance(glb_idx, slice) and glb_idx.step is not None \
                         and glb_idx.step > 1:
-                    if glb_idx_min > self.loc_abs_max:
+                    if glb_idx_min > loc_abs_max:
                         return retfunc(-1, -3)
                     elif glb_idx.start is None:  # glb start is zero.
-                        loc_min = self.loc_abs_min - base \
+                        loc_min = loc_abs_min - base \
                             + np.mod(glb_idx.step - np.mod(base, glb_idx.step),
                                      glb_idx.step)
                     else:  # glb start is given explicitly
-                        loc_min = self.loc_abs_min - base \
+                        loc_min = loc_abs_min - base \
                             + np.mod(glb_idx.step - np.mod(base - glb_idx.start,
                                                            glb_idx.step), glb_idx.step)
-                elif glb_idx_min is None or glb_idx_min < self.loc_abs_min:
-                    loc_min = self.loc_abs_min - base
-                elif glb_idx_min > self.loc_abs_max:
+                elif glb_idx_min is None or glb_idx_min < loc_abs_min:
+                    loc_min = loc_abs_min - base
+                elif glb_idx_min > loc_abs_max:
                     return retfunc(-1, -3)
                 else:
                     loc_min = glb_idx_min - base
@@ -294,19 +311,19 @@ def index_glb_to_loc(self, *args, rel=True):
                 # coincide with loc_abs_max.
                 if isinstance(glb_idx, slice) and glb_idx.step is not None \
                         and glb_idx.step < -1:
-                    if glb_idx_max < self.loc_abs_min:
+                    if glb_idx_max < loc_abs_min:
                         return retfunc(-1, -3)
                     elif glb_idx.start is None:
                         loc_max = top - base \
-                            + np.mod(glb_idx.step - np.mod(top - self.glb_max,
+                            + np.mod(glb_idx.step - np.mod(top - glb_max,
                                                            glb_idx.step), glb_idx.step)
                     else:
                         loc_max = top - base \
                             + np.mod(glb_idx.step - np.mod(top - glb_idx.start,
                                                            glb_idx.step), glb_idx.step)
-                elif glb_idx_max is None or glb_idx_max > self.loc_abs_max:
-                    loc_max = self.loc_abs_max - base
-                elif glb_idx_max < self.loc_abs_min:
+                elif glb_idx_max is None or glb_idx_max > loc_abs_max:
+                    loc_max = loc_abs_max - base
+                elif glb_idx_max < loc_abs_min:
                     return retfunc(-1, -3)
                 else:
                     loc_max = glb_idx_max - base
@@ -321,19 +338,19 @@ def index_glb_to_loc(self, *args, rel=True):
                 return None
             abs_ofs, side = args
             if side is LEFT:
-                rel_ofs = self.glb_min + abs_ofs - base
+                rel_ofs = abs_ofs - base
                 if abs_ofs >= base and abs_ofs <= top:
                     return rel_ofs
                 elif abs_ofs > top:
                     return top + 1
                 else:
                     return None
             else:
-                rel_ofs = abs_ofs - (self.glb_max - top)
-                if abs_ofs >= self.glb_max - top and abs_ofs <= self.glb_max - base:
+                rel_ofs = abs_ofs - (glb_max - top)
+                if abs_ofs >= glb_max - top and abs_ofs <= glb_max - base:
                     return rel_ofs
-                elif abs_ofs > self.glb_max - base:
-                    return self.glb_max - base + 1
+                elif abs_ofs > glb_max - base:
+                    return glb_max - base + 1
                 else:
                     return None
         else: