compiler: Fix constrain_essential_bcs with implicit_dims

Author: ZoeLeibowitz

devito/ir/equations/equation.py

@@ -107,10 +107,14 @@ def detect(cls, expr):
             ReduceMin: OpMin,
             PetscEq: OpPetsc
         }
-        try:
-            return reduction_mapper[type(expr)]
-        except KeyError:
-            pass
+        # try:
+        #     return reduction_mapper[type(expr)]
+        # except KeyError:
+        #     pass
+
+        for expr_type, op in reduction_mapper.items():
+            if isinstance(expr, expr_type):
+                return op
 
         # NOTE: in the future we might want to track down other kinds
         # of operations here (e.g., memcpy). However, we don't care for

devito/petsc/equations.py

@@ -22,25 +22,34 @@ def constrain_essential_bcs(expressions, **kwargs):
 
     # build mapper
     for e in expressions:
-        # from IPython import embed; embed()
         if not isinstance(e, ConstrainBC):
+            new_exprs.append(e)
             continue
 
         indexeds = retrieve_indexed(e)
         dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
-
+        # implicit_dims = set(e.implicit_dims)
+        dims.update(e.implicit_dims)
+        # from IPython import embed; embed()
         dims = [d for d in dims if d.is_Sub and not d.local]
 
+        if not dims:
+            new_exprs.append(e)
+            continue
+
         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
             # from IPython import embed; embed()
             # f._size_nodomain.left
-            halo_size_left = indexeds[0].function._size_halo[d].left
-            halo_size_right = indexeds[0].function._size_halo[d].right
+            # from IPython import embed; embed()
+            # halo_size_left = indexeds[0].function._size_halo[d].left
+            # halo_size_right = indexeds[0].function._size_halo[d].right
 
+            halo_size_left = 2
+            halo_size_right = 2
 
             from devito.petsc.types.dimension import SubDimMax, SubDimMin
 
@@ -52,7 +61,7 @@ def constrain_essential_bcs(expressions, **kwargs):
             subdim_max = SubDimMax(sregistry.make_name(prefix=d.name + '_max'), subdim=d, thickness=d.thickness)
             subdim_min = SubDimMin(sregistry.make_name(prefix=d.name + '_min'), subdim=d, thickness=d.thickness)
 
-            # from IPython import embed; embed()
+            # unique_name 
             new_dim = CustomBoundSubDimension(
             name=d.name,
             parent=d.parent,
@@ -63,13 +72,12 @@ def constrain_essential_bcs(expressions, **kwargs):
             )
             mapper[d] = new_dim
 
-    # build new expressions
-    for e in expressions:
-        if isinstance(e, ConstrainBC):
-            new_e = e.subs(mapper)
-            new_exprs.append(new_e)
-        
-        else:
-            new_exprs.append(e)
-
+        new_e = e.subs(mapper)
+        if e.implicit_dims:
+            # from devito.symbolics import uxreplace
+            implicit_dims_new = tuple(mapper.get(d, d) for d in e.implicit_dims)
+            # from IPython import embed; embed()
+            new_e = new_e._rebuild(implicit_dims=implicit_dims_new)
+        new_exprs.append(new_e)
+    
     return new_exprs

devito/petsc/types/equation.py

@@ -8,10 +8,9 @@ class EssentialBC(Eq):
     """
     Represents an essential boundary condition for use with `petscsolve`.
 
-    Due to ongoing work on PetscSection and DMDA integration (WIP),
-    these conditions are imposed as trivial equations. The compiler
-    will automatically zero the corresponding rows/columns in the Jacobian
-    and lift the boundary terms into the residual RHS.
+    The compiler will automatically zero the corresponding rows/columns in the Jacobian
+    and lift the boundary terms into the residual RHS, unless the user
+    specifies `constrain_bcs=True` to `petscsolve`.
 
     Note:
         - To define an essential boundary condition, use:
@@ -47,37 +46,14 @@ class ZeroColumn(EssentialBC):
 
 class ConstrainBC(EssentialBC):
     pass
-
-
-# class NoOfEssentialBC(ConstrainBC, Inc):
-#     """
-#     Equation used count essential boundary condition nodes.
-#     This type of equation is generated inside
-#     petscsolve if the user sets `constrain_bcs=True`.
-#     """
-#     def __new__(cls, *args, **kwargs):
-#         return Inc.__new__(Inc, *args, **kwargs)
 
 
-class NoOfEssentialBC(ConstrainBC):
+class NoOfEssentialBC(Inc, ConstrainBC):
     """Equation used count essential boundary condition nodes.
     This type of equation is generated inside
     petscsolve if the user sets `constrain_bcs=True`."""
-    # def __new__(cls, *args, **kwargs):
-    #     return Inc.__new__(Inc, *args, **kwargs)
     pass
-    
-
-# class NoOfEssentialBC(Inc, ConstrainBC):
-#     """
-#     Equation used to count essential boundary condition nodes.
-#     """
-
-#     def __new__(cls, *args, **kwargs):
-#         obj = super().__new__(cls, *args, **kwargs)
-#         return obj
 
 
 class PointEssentialBC(ConstrainBC):
     pass
-

devito/petsc/types/metadata.py

@@ -762,12 +762,14 @@ def _make_increment_expr(self, expr):
         """
         if isinstance(expr, EssentialBC):
             assert expr.lhs == self.target
-            from math import prod
-
-            rhs = prod(expr.rhs.dimensions)
+            # return NoOfEssentialBC(
+            #     TempSymb, expr.rhs,
+            #     subdomain=expr.subdomain,
+            # )
             return NoOfEssentialBC(
-                TempSymb, rhs,
-                subdomain=expr.subdomain
+                TempSymb, 1,
+                subdomain=expr.subdomain,
+                implicit_dims=expr.subdomain.dimensions
             )
         else:
             return None
@@ -791,7 +793,7 @@ def _make_point_bc_expr(self, expr):
         numBC = PetscInt(name='numBC2')
         if isinstance(expr, EssentialBC):
             assert expr.lhs == self.target
-            return NoOfEssentialBC(
+            return PointEssentialBC(
                 TempSymb, expr.rhs,
                 subdomain=expr.subdomain
             )

examples/petsc/Poisson/ed_bueler_2d.py

@@ -99,6 +99,8 @@ def exact(x, y):
 bcs += [EssentialBC(u, bc, subdomain=sub3)]
 bcs += [EssentialBC(u, bc, subdomain=sub4)]
 
+
+
 exprs = [eqn] + bcs
 petsc = petscsolve(
     exprs, target=u,
@@ -109,11 +111,11 @@ def exact(x, y):
 
 with switchconfig(log_level='DEBUG'):
     op = Operator(petsc, language='petsc')
-    summary = op.apply()
+    # summary = op.apply()
     # print(op.arguments())
 
 
-# print(op.ccode)
+print(op.ccode)
 # iters = summary.petsc[('section0', 'poisson_2d')].KSPGetIterationNumber
 
 u_exact = Function(name='u_exact', grid=grid, space_order=2)