compiler: Start SubDimMax

Author: ZoeLeibowitz

devito/petsc/equations.py

@@ -0,0 +1,66 @@
+from devito.symbolics import retrieve_indexed, retrieve_dimensions
+from devito.petsc import EssentialBC
+from devito.types.dimension import CustomBoundSubDimension
+from devito import Min, Max
+
+
+def lower_exprs_petsc(expressions, **kwargs):
+    mapper = {}
+
+    additional_exprs = []
+
+    # build mapper
+    for e in expressions:
+        if not isinstance(e, EssentialBC):
+            continue
+        indexeds = retrieve_indexed(e)
+        dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
+
+        dims = [d for d in dims if d.is_Sub and not d.local]
+
+        for d in dims:
+            # replace the dim with a new one that has a different symbolic_min and symbolic_max
+
+            # obvs shouldn't be obtained from indexeds[0], but how should it be obtained?
+            # USE e.lhs function -> the one that the BC is being applied to
+            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
+
+            # TODO: change name..
+
+            # global_rtkn = kwargs.get(d.rtkn.name, d.rtkn.value)
+            # in theory this class shoulod just take in d directly 
+            subdim_max = SubDimMax(d.name + '_max', subdim=d, thickness=d.thickness)
+
+            # from IPython import embed; embed()
+
+            new_dim = CustomBoundSubDimension(
+            name=d.name,
+            parent=d.parent,
+            thickness=d.thickness,
+            local=d.local,
+            custom_left=Max(d.ltkn, d.parent.symbolic_min - halo_size_left),
+            custom_right=Min(subdim_max, d.parent.symbolic_max + halo_size_right)
+            )
+            mapper[d] = new_dim
+
+            # from IPython import embed; embed()
+
+    # build new expressions
+    for e in expressions:
+        if not isinstance(e, EssentialBC):
+            continue
+
+        # build new expression
+        new_e = e.subs(mapper)
+
+        additional_exprs.append(new_e)
+
+    return expressions + additional_exprs
+    
+
+

devito/petsc/types/dimension.py

@@ -0,0 +1,32 @@
+from devito.types.dimension import Thickness
+
+
+
+class SubDimMax(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_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.name, self.value)
+
+        g_x_M = grid.distributor.decomposition[self.subdim.parent].glb_max
+
+
+        return {self.name: g_x_M}