mpi: Fix index_glb_to_loc_unsafe using base

Author: ZoeLeibowitz

devito/data/decomposition.py

@@ -466,6 +466,7 @@ def index_glb_to_loc_unsafe(self, glb_idx, rel=True):
             glb_max = self.glb_max
 
         base = loc_abs_min if rel else 0
+        glb_min = 0
 
         # index_glb_to_loc(index)
         # -> Base case, empty local subdomain
@@ -477,8 +478,8 @@ def index_glb_to_loc_unsafe(self, glb_idx, rel=True):
         # -> Do the actual conversion
         if loc_abs_min <= glb_idx <= loc_abs_max:
             return glb_idx - base
-        # elif glb_min <= glb_idx <= glb_max:
-        #     return base
+        elif glb_min <= glb_idx <= glb_max:
+            return glb_idx - base
         else:
             # This should raise an exception when used to access a numpy.array
             return glb_idx

devito/petsc/equations.py

@@ -1,3 +1,4 @@
+from sympy import Eq
 from devito.symbolics import retrieve_indexed, retrieve_dimensions
 from devito.petsc import EssentialBC
 from devito.types.dimension import CustomBoundSubDimension
@@ -26,15 +27,17 @@ def lower_exprs_petsc(expressions, **kwargs):
             halo_size_left = indexeds[0].function._size_halo[d].left
             halo_size_right = indexeds[0].function._size_halo[d].right
 
-            grid = indexeds[0].function.grid
 
-            from devito.petsc.types.dimension import SubDimMax
+            from devito.petsc.types.dimension import SubDimMax, SubDimMin
+
+
 
             # TODO: change name..
 
-            # global_rtkn = kwargs.get(d.rtkn.name, d.rtkn.value)
-            # in theory this class shoulod just take in d directly 
+            # in theory this class shoulod just take in d
+            # TODO: use unique name
             subdim_max = SubDimMax(d.name + '_max', subdim=d, thickness=d.thickness)
+            subdim_min = SubDimMin(d.name + '_min', subdim=d, thickness=d.thickness)
 
             # from IPython import embed; embed()
 
@@ -43,7 +46,7 @@ def lower_exprs_petsc(expressions, **kwargs):
             parent=d.parent,
             thickness=d.thickness,
             local=d.local,
-            custom_left=Max(d.ltkn, d.parent.symbolic_min - halo_size_left),
+            custom_left=Max(subdim_min, d.parent.symbolic_min - halo_size_left),
             custom_right=Min(subdim_max, d.parent.symbolic_max + halo_size_right)
             )
             mapper[d] = new_dim
@@ -60,8 +63,8 @@ def lower_exprs_petsc(expressions, **kwargs):
 
         additional_exprs.append(new_e)
 
+    # return expressions + additional_exprs
     return expressions + additional_exprs
-    # return expressions
-    
+
 
 

devito/petsc/types/dimension.py

@@ -19,39 +19,48 @@ def subdim(self):
 
     def _arg_values(self, grid=None, **kwargs):
 
-        # # global rtkn
-        # grtkn = kwargs.get(self.subdim.rtkn.name, self.subdim.rtkn.value)
-
-        # g_x_M = grid.distributor.decomposition[self.subdim.parent].glb_max
-        # val = grid.distributor.decomposition[self.subdim.parent].index_glb_to_loc(g_x_M - grtkn)
-
-        # return {self.name: int(val)}
-    
-
         dist = grid.distributor
-        rank = dist.myrank
-        comm = dist.comm
 
         # global rtkn
         grtkn = kwargs.get(self.subdim.rtkn.name, self.subdim.rtkn.value)
-
+        # print(g_x_M)
         # decomposition info
         decomp = dist.decomposition[self.subdim.parent]
         g_x_M = decomp.glb_max
+        # print(g_x_M)
         val = decomp.index_glb_to_loc_unsafe(g_x_M - grtkn)
+        print(val)
+
+
+        return {self.name: int(val)}
+    
+
+
+class SubDimMin(Thickness):
+    """
+    """
+
+    def __init_finalize__(self, *args, **kwargs):
+        self._subdim = kwargs.pop('subdim')
+        self._dtype = self._subdim.dtype
+
+        super().__init_finalize__(*args, **kwargs)
+
+    @property
+    def subdim(self):
+        return self._subdim
 
+    def _arg_values(self, grid=None, **kwargs):
 
-        print(
-            f"[Rank {rank}] "
-            f"grtkn={grtkn}, "
-            f"g_x_M={g_x_M}, "
-            f"glb_idx={g_x_M - grtkn}, "
-            f"loc_val={val}",
-            flush=True
-        )
+        dist = grid.distributor
 
-        if val is None:
-            return {}
+        # global ltkn
+        gltkn = kwargs.get(self.subdim.ltkn.name, self.subdim.ltkn.value)
+
+        # decomposition info
+        decomp = dist.decomposition[self.subdim.parent]
+        g_x_m = decomp.glb_min
+        val = decomp.index_glb_to_loc_unsafe(g_x_m + gltkn)
 
         return {self.name: int(val)}