progress with unsafe index global to local

Author: ZoeLeibowitz

devito/data/decomposition.py

@@ -450,6 +450,39 @@ def index_loc_to_glb(self, *args):
         else:
             raise TypeError("Expected 1 arguments, found %d" % len(args))
 
+    def index_glb_to_loc_unsafe(self, glb_idx, rel=True):
+        """
+        Convert a global index to a local index even if not owned.
+        WARNING: Must not be used to index data as there are no guard rails against returning out
+        of bound indices.
+        """
+        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
+            glb_max = self.glb_max - self.glb_min
+        else:
+            loc_abs_min = self.loc_abs_min
+            loc_abs_max = self.loc_abs_max
+            glb_max = self.glb_max
+
+        base = loc_abs_min if rel else 0
+
+        # index_glb_to_loc(index)
+        # -> Base case, empty local subdomain
+        if self.loc_empty:
+            return None
+        # -> Handle negative index
+        if glb_idx < 0:
+            glb_idx = glb_max + glb_idx + 1
+        # -> 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
+        else:
+            # This should raise an exception when used to access a numpy.array
+            return glb_idx
+
     def reshape(self, *args):
         """
         Create a new Decomposition with extended or reduced boundary subdomains.

devito/petsc/equations.py

@@ -61,6 +61,7 @@ def lower_exprs_petsc(expressions, **kwargs):
         additional_exprs.append(new_e)
 
     return expressions + additional_exprs
+    # return expressions
 
 
 

devito/petsc/types/dimension.py

@@ -15,18 +15,43 @@ def __init_finalize__(self, *args, **kwargs):
     @property
     def subdim(self):
         return self._subdim
-    
-    def _arg_defaults(self, alias=None):
-        key = alias or self
-        # from IPython import embed; embed()
-        return {key.name: self.data}
+
 
     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.name, self.value)
+        grtkn = kwargs.get(self.subdim.rtkn.name, self.subdim.rtkn.value)
+
+        # decomposition info
+        decomp = dist.decomposition[self.subdim.parent]
+        g_x_M = decomp.glb_max
+        val = decomp.index_glb_to_loc_unsafe(g_x_M - grtkn)
+
+
+        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
+        )
 
-        g_x_M = grid.distributor.decomposition[self.subdim.parent].glb_max
+        if val is None:
+            return {}
 
+        return {self.name: int(val)}
 
-        return {self.name: g_x_M}