misc: Clean up

Author: ZoeLeibowitz

devito/petsc/solve.py

@@ -145,7 +145,9 @@ def linear_solve_args(self):
         initial_guess = InitialGuess(target, exprs, arrays, self.time_mapper)
 
         constrain_exprs = self._get_constrain_exprs(exprs)
-        constrain_bc = ConstrainBC(target, constrain_exprs, arrays) if constrain_exprs else None
+        constrain_bc = None
+        if constrain_exprs:
+            constrain_bc = ConstrainBC(target, constrain_exprs, arrays)
 
         field_data = FieldData(
             target=target,
@@ -205,11 +207,11 @@ def linear_solve_args(self):
             jacobian.target_scaler_mapper
         )
 
-        constrain_bc = {
-            t: ConstrainBC(t, self._get_constrain_exprs(as_tuple(self.target_exprs[t])), arrays)
-            for t in targets
-            if self._get_constrain_exprs(as_tuple(self.target_exprs[t]))
-        }
+        constrain_bc = {}
+        for t in targets:
+            cexprs = self._get_constrain_exprs(as_tuple(self.target_exprs[t]))
+            if cexprs:
+                constrain_bc[t] = ConstrainBC(t, cexprs, arrays)
         constrain_bc = constrain_bc if constrain_bc else None
 
         all_data = MultipleFieldData(

devito/petsc/types/array.py

@@ -125,17 +125,17 @@ def symbolic_shape(self):
         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                                                                                                                         
+    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) 
+        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

@@ -10,8 +10,8 @@ class EssentialBC(Eq):
 
     The compiler will automatically zero the corresponding rows/columns in the Jacobian
     and lift the boundary terms into the residual RHS, unless the BC is constrained
-    via a `PetscSection`. In which case, they are set once and removed from the global solver.
-    Constraining can be enabled in two ways:
+    via a `PetscSection`. In which case, they are set once and removed from the
+    global solver. Constraining can be enabled in two ways:
 
     - Globally: pass `constrain_bcs=True` to `petscsolve` to constrain all
       `EssentialBC`s in the solve.