dsl/misc/compiler: Clean up, fix __indices_setup__ in PETScArray and allow constraining individual EssentialBC equations

Author: ZoeLeibowitz

devito/petsc/solve.py

@@ -79,10 +79,12 @@ def petscsolve(target_exprs, target=None, solver_parameters=None,
            'kspgettype', 'kspgetnormtype', 'snesgetiterationnumber']
 
     constrain_bcs : bool, optional
-        If `True`, essential boundary conditions specifed by `EssentialBC` equations
-        are constrained through a `PetscSection`. As a result, the corresponding degrees
-        of freedom are excluded from the global solver and are not imposed using
-        trivial equations.
+        If `True`, ALL `EssentialBC` equations are constrained via a `PetscSection`,
+        excluding the corresponding degrees of freedom from the global solver.
+        Individual `EssentialBC`s can also be constrained independently by passing
+        ``constrain=True`` to their constructor (e.g.
+        ``EssentialBC(lhs, rhs, subdomain=..., constrain=True)``), regardless of
+        this flag.
 
     Returns
     -------
@@ -142,11 +144,8 @@ def linear_solve_args(self):
         residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag)
         initial_guess = InitialGuess(target, exprs, arrays, self.time_mapper)
 
-        constrain_bc = (
-            ConstrainBC(target, exprs, arrays)
-            if self.constrain_bcs
-            else None
-        )
+        constrain_exprs = self._get_constrain_exprs(exprs)
+        constrain_bc = ConstrainBC(target, constrain_exprs, arrays) if constrain_exprs else None
 
         field_data = FieldData(
             target=target,
@@ -159,6 +158,16 @@ def linear_solve_args(self):
 
         return target, funcs, field_data
 
+    def _get_constrain_exprs(self, exprs):
+        """
+        Return the subset of `exprs` to be constrained via PetscSection.
+        If `constrain_bcs=True`, all `EssentialBC` exprs are returned.
+        Otherwise, only those individually marked with ``constrain=True``.
+        """
+        if self.constrain_bcs:
+            return tuple(e for e in exprs if isinstance(e, EssentialBC))
+        return tuple(e for e in exprs if isinstance(e, EssentialBC) and e.constrain)
+
     def generate_arrays(self, *targets):
         return {
             t: {
@@ -197,9 +206,11 @@ def linear_solve_args(self):
         )
 
         constrain_bc = {
-            t: ConstrainBC(t, as_tuple(self.target_exprs[t]), arrays)
+            t: ConstrainBC(t, self._get_constrain_exprs(as_tuple(self.target_exprs[t])), arrays)
             for t in targets
-        } if self.constrain_bcs else None
+            if self._get_constrain_exprs(as_tuple(self.target_exprs[t]))
+        }
+        constrain_bc = constrain_bc if constrain_bc else None
 
         all_data = MultipleFieldData(
             targets=targets,

devito/petsc/types/array.py

@@ -51,11 +51,12 @@ def __dtype_setup__(cls, **kwargs):
     @classmethod
     def __indices_setup__(cls, *args, **kwargs):
         target = kwargs['target']
-        dimensions = tuple(target.indices[d] for d in target.space_dimensions)
+        dimensions = tuple(target.space_dimensions)
+        target_indices = tuple(target.indices[d] for d in target.space_dimensions)
         if args:
             indices = args
         else:
-            indices = dimensions
+            indices = target_indices
         return as_tuple(dimensions), as_tuple(indices)
 
     def __halo_setup__(self, **kwargs):
@@ -123,6 +124,19 @@ def symbolic_shape(self):
         # Reverse it since DMDA is setup backwards to Devito dimensions.
         return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions)
 
+    # TODO: Is this necessary? Taken directly from `Function`.
+    def _eval_at(self, func):                                                                                                                   
+        if self.staggered == func.staggered:                                                                                                  
+            return self                                                                                                                         
+        mapper = {}
+        for d in self.dimensions:                                                                                                               
+            try:                                                                                                                              
+                if self.indices_ref[d] is not func.indices_ref[d]:                                                                              
+                    mapper[self.indices_ref[d]] = func.indices_ref[d]                                                                           
+            except KeyError:                                                                                                                    
+                pass                                                                                                                            
+        return self.subs(mapper) 
+
 
 class PetscBundle(Bundle):
     """

devito/petsc/types/equation.py

@@ -9,30 +9,40 @@ class EssentialBC(Eq):
     Represents an essential boundary condition for use with `petscsolve`.
 
     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`.
+    and lift the boundary terms into the residual RHS, unless the BC is constrained
+    via a `PetscSection`. Constraining can be enabled in two ways:
+
+    - Globally: pass `constrain_bcs=True` to `petscsolve` to constrain all
+      `EssentialBC`s in the solve.
+    - Individually: pass `constrain=True` to a specific `EssentialBC` constructor,
+      e.g. ``EssentialBC(lhs, rhs, subdomain=..., constrain=True)``.
 
     Note:
         - To define an essential boundary condition, use:
             Eq(target, boundary_value, subdomain=...),
           where `target` is the Function-like object passed to `petscsolve`.
         - SubDomains used for multiple `EssentialBC`s must not overlap.
     """
-    __rkwargs__ = Eq.__rkwargs__ + ("target",)
+    __rkwargs__ = Eq.__rkwargs__ + ("target", "constrain")
 
-    def __new__(cls, *args, target=None, **kwargs):
+    def __new__(cls, *args, target=None, constrain=False, **kwargs):
         obj = super().__new__(cls, *args, **kwargs)
 
         if target is None:
             target = obj.lhs.function
 
         obj._target = target
+        obj._constrain = constrain
         return obj
 
     @property
     def target(self):
         return self._target
 
+    @property
+    def constrain(self):
+        return self._constrain
+
 
 class ZeroRow(EssentialBC):
     """

tests/test_petsc.py

@@ -1481,9 +1481,9 @@ class TestMPI:
     def test_laplacian_1d(self, nx, unorm, mode):
         """
         """
-        configuration['compiler'] = 'custom'
-        os.environ['CC'] = 'mpicc'
-        PetscInitialize()
+        # configuration['compiler'] = 'custom'
+        # os.environ['CC'] = 'mpicc'
+        # PetscInitialize()
 
         class SubSide(SubDomain):
             def __init__(self, side='left', grid=None):