fix

fix

lint

Author: leo-collins

firedrake/assemble.py

@@ -387,6 +387,16 @@ def allocation_integral_types(self):
         else:
             return self._allocation_integral_types
 
+    @staticmethod
+    def _as_pyop2_type(tensor, indices=None):
+        if isinstance(tensor, (firedrake.Cofunction, firedrake.Function)):
+            return OneFormAssembler._as_pyop2_type(tensor, indices=indices)
+        elif isinstance(tensor, ufl.Matrix):
+            return ExplicitMatrixAssembler._as_pyop2_type(tensor, indices=indices)
+        else:
+            assert indices is None
+            return tensor
+
     def assemble(self, tensor=None, current_state=None):
         """Assemble the form.
 
@@ -411,21 +421,22 @@ def assemble(self, tensor=None, current_state=None):
         """
         def visitor(e, *operands):
             t = tensor if e is self._form else None
-            return self.base_form_assembly_visitor(e, t, *operands)
+            # Deal with 2-form bcs inside the visitor
+            bcs = self._bcs if isinstance(e, ufl.BaseForm) and len(e.arguments()) == 2 else ()
+            return self.base_form_assembly_visitor(e, t, bcs, *operands)
 
         # DAG assembly: traverse the DAG in a post-order fashion and evaluate the node on the fly.
         visited = {}
         result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited)
 
-        # Apply BCs after assembly
+        # Deal with 1-form bcs outside the visitor
         rank = len(self._form.arguments())
         if rank == 1 and not isinstance(result, ufl.ZeroBaseForm):
             for bc in self._bcs:
                 OneFormAssembler._apply_bc(self, result, bc, u=current_state)
-
         return result
 
-    def base_form_assembly_visitor(self, expr, tensor, *args):
+    def base_form_assembly_visitor(self, expr, tensor, bcs, *args):
         r"""Assemble a :class:`~ufl.classes.BaseForm` object given its assembled operands.
 
             This functions contains the assembly handlers corresponding to the different nodes that
@@ -446,7 +457,7 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 assembler = OneFormAssembler(form, form_compiler_parameters=self._form_compiler_params,
                                              zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight)
             elif rank == 2:
-                assembler = TwoFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params,
+                assembler = TwoFormAssembler(form, bcs=bcs, form_compiler_parameters=self._form_compiler_params,
                                              mat_type=self._mat_type, sub_mat_type=self._sub_mat_type,
                                              options_prefix=self._options_prefix, appctx=self._appctx, weight=self._weight,
                                              allocation_integral_types=self.allocation_integral_types)
@@ -457,13 +468,12 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
             if len(args) != 1:
                 raise TypeError("Not enough operands for Adjoint")
             mat, = args
-            res = tensor.petscmat if tensor else PETSc.Mat()
-            petsc_mat = mat.petscmat
+            result = tensor.petscmat if tensor else PETSc.Mat()
             # Out-of-place Hermitian transpose
-            petsc_mat.hermitianTranspose(out=res)
-            (row, col) = mat.arguments()
-            return matrix.AssembledMatrix((col, row), self._bcs, res,
-                                          options_prefix=self._options_prefix)
+            mat.petscmat.hermitianTranspose(out=result)
+            if tensor is None:
+                tensor = self.assembled_matrix(expr, bcs, result)
+            return tensor
         elif isinstance(expr, ufl.Action):
             if len(args) != 2:
                 raise TypeError("Not enough operands for Action")
@@ -481,7 +491,7 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                     result = tensor.petscmat if tensor else PETSc.Mat()
                     lhs.petscmat.matMult(rhs.petscmat, result=result)
                     if tensor is None:
-                        tensor = self.assembled_matrix(expr, result)
+                        tensor = self.assembled_matrix(expr, bcs, result)
                     return tensor
                 else:
                     raise TypeError("Incompatible RHS for Action.")
@@ -500,9 +510,6 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 raise TypeError("Mismatching weights and operands in FormSum")
             if len(args) == 0:
                 raise TypeError("Empty FormSum")
-            if tensor:
-                tensor.zero()
-
             # Assemble weights
             weights = []
             for w in expr.weights():
@@ -520,27 +527,54 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                     raise ValueError("Expecting a scalar weight expression")
                 weights.append(w)
 
+            # Scalar FormSum
             if all(isinstance(op, numbers.Complex) for op in args):
-                result = sum(weight * arg for weight, arg in zip(weights, args))
-                return tensor.assign(result) if tensor else result
-            elif (all(isinstance(op, firedrake.Cofunction) for op in args)
+                result = numpy.dot(weights, args)
+                return tensor.assign(result) if tensor else result.item()
+
+            # Accumulate coefficients in a dictionary for each unique Dat/Mat
+            terms = defaultdict(PETSc.ScalarType)
+            for arg, weight in zip(args, weights):
+                t = self._as_pyop2_type(arg)
+                terms[t] += weight
+
+            # Zero the output tensor, or rescale it if it appears in the sum
+            tensor_scale = terms.pop(self._as_pyop2_type(tensor), 0)
+            if tensor is None or tensor_scale == 1:
+                pass
+            elif tensor_scale == 0:
+                tensor.zero()
+            elif isinstance(tensor, (firedrake.Cofunction, firedrake.Function)):
+                with tensor.dat.vec as v:
+                    v.scale(tensor_scale)
+            elif isinstance(tensor, ufl.Matrix):
+                tensor.petscmat.scale(tensor_scale)
+            else:
+                raise ValueError("Expecting tensor to be None, Function, Cofunction, or Matrix")
+
+            # Compute the linear combination
+            if (all(isinstance(op, firedrake.Cofunction) for op in args)
                     or all(isinstance(op, firedrake.Function) for op in args)):
+                # Vector FormSum
                 V, = set(a.function_space() for a in args)
                 result = tensor if tensor else firedrake.Function(V)
-                result.dat.maxpy(weights, [a.dat for a in args])
+                weights = terms.values()
+                dats = terms.keys()
+                result.dat.maxpy(weights, dats)
                 return result
             elif all(isinstance(op, ufl.Matrix) for op in args):
+                # Matrix FormSum
                 result = tensor.petscmat if tensor else PETSc.Mat()
-                for (op, w) in zip(args, weights):
+                for (op, w) in terms.items():
                     if result:
                         # If result is not void, then accumulate on it
-                        result.axpy(w, op.petscmat)
+                        result.axpy(w, op.handle)
                     else:
                         # If result is void, then allocate it with first term
-                        op.petscmat.copy(result=result)
+                        op.handle.copy(result=result)
                         result.scale(w)
                 if tensor is None:
-                    tensor = self.assembled_matrix(expr, result)
+                    tensor = self.assembled_matrix(expr, bcs, result)
                 return tensor
             else:
                 raise TypeError("Mismatching FormSum shapes")
@@ -572,9 +606,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
                 # Occur in situations such as Interpolate composition
                 operand = assembled_operand[0]
 
-            reconstruct_interp = expr._ufl_expr_reconstruct_
             if (v, operand) != expr.argument_slots():
-                expr = reconstruct_interp(operand, v=v)
+                expr = expr._ufl_expr_reconstruct_(operand, v=v)
 
             rank = len(expr.arguments())
             if rank > 2:
@@ -591,8 +624,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args):
         else:
             raise TypeError(f"Unrecognised BaseForm instance: {expr}")
 
-    def assembled_matrix(self, expr, petscmat):
-        return matrix.AssembledMatrix(expr.arguments(), self._bcs, petscmat,
+    def assembled_matrix(self, expr, bcs, petscmat):
+        return matrix.AssembledMatrix(expr.arguments(), bcs, petscmat,
                                       options_prefix=self._options_prefix)
 
     @staticmethod
@@ -1441,10 +1474,11 @@ def _apply_bc(self, tensor, bc, u=None):
         index = 0 if V.index is None else V.index
         space = V if V.parent is None else V.parent
         if isinstance(bc, DirichletBC):
-            if space != spaces[0]:
-                raise TypeError("bc space does not match the test function space")
-            elif space != spaces[1]:
-                raise TypeError("bc space does not match the trial function space")
+            if not any(space == fs for fs in spaces):
+                raise TypeError("bc space does not match the test or trial function space")
+            if spaces[0] != spaces[1]:
+                # Not on a diagonal block, we cannot set diagonal entries
+                return
 
             # Set diagonal entries on bc nodes to 1 if the current
             # block is on the matrix diagonal and its index matches the

tests/firedrake/regression/test_interpolate.py

@@ -599,3 +599,24 @@ def test_interpolator_reuse(family, degree, mode):
 
         # Test for correctness
         assert np.allclose(result.dat.data, expected)
+
+
+def test_mixed_space_bcs():
+    mesh = UnitSquareMesh(2, 2)
+    V = FunctionSpace(mesh, "CG", 1)
+    W = V * V
+    rg = RandomGenerator(PCG64(seed=123456789))
+    w = rg.uniform(W)
+
+    bcs = [DirichletBC(W.sub(0), 0, 1),
+           DirichletBC(W.sub(1), 0, 2),
+           DirichletBC(V, 0, (3, 4))]
+
+    I = assemble(interpolate(sum(TrialFunction(W)), V), bcs=bcs)
+    result = assemble(action(I, w))
+
+    for bc in bcs[:-1]:
+        bc.zero(w)
+    expected = assemble(interpolate(sum(w), V), bcs=bcs[-1:])
+
+    assert np.allclose(result.dat.data, expected.dat.data)

tests/firedrake/submesh/test_submesh_interpolate.py

@@ -10,14 +10,14 @@
 
 def _get_expr(V):
     m = V.ufl_domain()
-    if m.geometric_dimension() == 1:
+    if m.geometric_dimension == 1:
         x, = SpatialCoordinate(m)
         y = x * x
         z = x + y
-    elif m.geometric_dimension() == 2:
+    elif m.geometric_dimension == 2:
         x, y = SpatialCoordinate(m)
         z = x + y
-    elif m.geometric_dimension() == 3:
+    elif m.geometric_dimension == 3:
         x, y, z = SpatialCoordinate(m)
     else:
         raise NotImplementedError("Not implemented")
@@ -28,7 +28,7 @@ def _get_expr(V):
 
 
 def make_submesh(mesh, subdomain_cond, label_value):
-    dim = mesh.topological_dimension()
+    dim = mesh.topological_dimension
     DG0 = FunctionSpace(mesh, "DG", 0)
     indicator_function = Function(DG0).interpolate(subdomain_cond)
     mesh.mark_entities(indicator_function, label_value)
@@ -145,7 +145,7 @@ def test_submesh_interpolate_subcell_subcell_2_processes():
     mesh = RectangleMesh(
         3, 1, 3., 1., quadrilateral=True, distribution_parameters={"partitioner_type": "simple"},
     )
-    dim = mesh.topological_dimension()
+    dim = mesh.topological_dimension
     x, _ = SpatialCoordinate(mesh)
     DG0 = FunctionSpace(mesh, "DG", 0)
     f_l = Function(DG0).interpolate(conditional(x < 2.0, 1, 0))
@@ -210,7 +210,7 @@ def expr(m):
     )
     facet_value = 999
     mesh = RelabeledMesh(mesh, [facet_function], [facet_value])
-    subm = Submesh(mesh, mesh.topological_dimension() - 1, facet_value)
+    subm = Submesh(mesh, mesh.topological_dimension - 1, facet_value)
     DG3d = FunctionSpace(mesh, "DG", degree)
     dg3d = Function(DG3d).interpolate(expr(mesh))
     DG2d = FunctionSpace(subm, "DG", degree)
@@ -258,7 +258,7 @@ def expr(m):
     facet_function = Function(V).interpolate(Constant(1.))
     facet_value = 999
     mesh = RelabeledMesh(mesh, [facet_function], [facet_value])
-    subm = Submesh(mesh, mesh.topological_dimension() - 1, facet_value)
+    subm = Submesh(mesh, mesh.topological_dimension - 1, facet_value)
     HDivT3d = FunctionSpace(mesh, "HDiv Trace", degree)
     hdivt3d = Function(HDivT3d).interpolate(expr(mesh))
     DG2d = FunctionSpace(subm, "DG", degree)