Merge pull request #1726 from firedrakeproject/fix_equation_bc

fix direct assembly with equation_bcs

Author: wence-

firedrake/assemble.py

@@ -10,7 +10,7 @@
 from firedrake import (assemble_expressions, matrix, parameters, solving,
                        tsfc_interface, utils)
 from firedrake.adjoint import annotate_assemble
-from firedrake.bcs import DirichletBC, EquationBCSplit
+from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit
 from firedrake.slate import slac, slate
 from firedrake.utils import ScalarType
 from pyop2 import op2
@@ -162,7 +162,7 @@ def create_assembly_callable(expr, tensor=None, bcs=None, form_compiler_paramete
         raise ValueError("Have to provide tensor to write to")
     if mat_type == "matfree":
         return tensor.assemble
-    loops = _assemble(expr, tensor=tensor, bcs=bcs,
+    loops = _assemble(expr, tensor=tensor, bcs=solving._extract_bcs(bcs),
                       form_compiler_parameters=form_compiler_parameters,
                       mat_type=mat_type,
                       sub_mat_type=sub_mat_type,
@@ -200,6 +200,10 @@ def get_matrix(expr, mat_type, sub_mat_type, *, bcs=None,
     arguments = expr.arguments()
     if bcs is None:
         bcs = ()
+    else:
+        if any(isinstance(bc, EquationBC) for bc in bcs):
+            raise TypeError("EquationBC objects not expected here. "
+                            "Preprocess by extracting the appropriate form with bc.extract_form('Jp') or bc.extract_form('J')")
     if tensor is not None and tensor.a.arguments() != arguments:
         raise ValueError("Form's arguments do not match provided result tensor")
     if matfree:
@@ -425,9 +429,6 @@ def apply_bcs(tensor, bcs, *, assembly_rank=None, form_compiler_parameters=None,
     """
     dirichletbcs = tuple(bc for bc in bcs if isinstance(bc, DirichletBC))
     equationbcs = tuple(bc for bc in bcs if isinstance(bc, EquationBCSplit))
-    if any(not isinstance(bc, (DirichletBC, EquationBCSplit)) for bc in bcs):
-        raise NotImplementedError("Unhandled type of bc object")
-
     if assembly_rank == AssemblyRank.MATRIX:
         op2tensor = tensor.M
         shape = tuple(len(a.function_space()) for a in tensor.a.arguments())
@@ -669,11 +670,6 @@ def _assemble(expr, tensor=None, bcs=None, form_compiler_parameters=None,
         # building a functionspace (e.g. if integrating a constant)).
         m.init()
 
-    if bcs is None:
-        bcs = ()
-    else:
-        bcs = tuple(bcs)
-
     for o in chain(expr.arguments(), expr.coefficients()):
         domain = o.ufl_domain()
         if domain is not None and domain.topology != topology:
@@ -689,6 +685,16 @@ def _assemble(expr, tensor=None, bcs=None, form_compiler_parameters=None,
     else:
         assembly_rank = AssemblyRank.SCALAR
 
+    if not isinstance(bcs, (tuple, list)):
+        raise RuntimeError("Expecting bcs to be a tuple or a list by this stage.")
+    if assembly_rank == AssemblyRank.MATRIX:
+        # Checks will take place in get_matrix.
+        pass
+    elif assembly_rank == AssemblyRank.VECTOR:
+        # Might have gotten here without `EquationBC` objects preprocessed.
+        if any(isinstance(bc, EquationBC) for bc in bcs):
+            bcs = tuple(bc.extract_form('F') for bc in bcs)
+
     if assembly_rank == AssemblyRank.MATRIX:
         test, trial = expr.arguments()
         tensor, zeros, result = get_matrix(expr, mat_type, sub_mat_type,

firedrake/bcs.py

@@ -240,6 +240,10 @@ def increment_bc_depth(self):
         for bc in itertools.chain(*self.bcs):
             bc._bc_depth += 1
 
+    def extract_forms(self, form_type):
+        # Return boundary condition objects actually used in assembly.
+        raise NotImplementedError("Method to extract form objects not implemented.")
+
 
 class DirichletBC(BCBase, DirichletBCMixin):
     r'''Implementation of a strong Dirichlet boundary condition.
@@ -437,6 +441,10 @@ def apply(self, r, u=None):
     def integrals(self):
         return []
 
+    def extract_form(self, form_type):
+        # DirichletBC is directly used in assembly.
+        return self
+
 
 class EquationBC(object):
     r'''Construct and store EquationBCSplit objects (for `F`, `J`, and `Jp`).
@@ -516,6 +524,16 @@ def dirichlet_bcs(self):
         # _F, _J, and _Jp all have the same DirichletBCs
         yield from self._F.dirichlet_bcs()
 
+    def extract_form(self, form_type):
+        r"""Return :class:`.EquationBCSplit` associated with the given 'form_type'.
+
+        :arg form_type: Form to extract; 'F', 'J', or 'Jp'.
+        """
+        if form_type not in {"F", "J", "Jp"}:
+            raise ValueError("Unknown form_type: 'form_type' must be 'F', 'J', or 'Jp'.")
+        else:
+            return getattr(self, f"_{form_type}")
+
     def reconstruct(self, V, subu, u, field):
         _F = self._F.reconstruct(field=field, V=V, subu=subu, u=u)
         _J = self._J.reconstruct(field=field, V=V, subu=subu, u=u)

firedrake/solving.py

@@ -318,10 +318,11 @@ def _extract_bcs(bcs):
     from firedrake.bcs import BCBase, EquationBC
     if bcs is None:
         return ()
-    try:
-        bcs = tuple(bcs)
-    except TypeError:
-        bcs = (bcs,)
+    if isinstance(bcs, (BCBase, EquationBC)):
+        return (bcs, )
+    else:
+        if not isinstance(bcs, (tuple, list)):
+            raise TypeError("bcs must be BCBase, EquationBC, tuple, or list, not '%s'." % type(bcs).__name__)
     for bc in bcs:
         if not isinstance(bc, (BCBase, EquationBC)):
             raise TypeError("Provided boundary condition is a '%s', not a BCBase" % type(bc).__name__)

firedrake/solving_utils.py

@@ -82,7 +82,6 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None,
                  options_prefix=None,
                  transfer_manager=None):
         from firedrake.assemble import create_assembly_callable
-        from firedrake.bcs import DirichletBC
         if pmat_type is None:
             pmat_type = mat_type
         self.mat_type = mat_type
@@ -132,9 +131,9 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None,
             # pmat_type == mat_type and Jp_eq_J
             self.Jp = None
 
-        self.bcs_F = [bc if isinstance(bc, DirichletBC) else bc._F for bc in problem.bcs]
-        self.bcs_J = [bc if isinstance(bc, DirichletBC) else bc._J for bc in problem.bcs]
-        self.bcs_Jp = [bc if isinstance(bc, DirichletBC) else bc._Jp for bc in problem.bcs]
+        self.bcs_F = tuple(bc.extract_form('F') for bc in problem.bcs)
+        self.bcs_J = tuple(bc.extract_form('J') for bc in problem.bcs)
+        self.bcs_Jp = tuple(bc.extract_form('Jp') for bc in problem.bcs)
         self._assemble_residual = create_assembly_callable(self.F,
                                                            tensor=self._F,
                                                            bcs=self.bcs_F,

tests/equation_bcs/test_equation_bcs_assemble.py

@@ -0,0 +1,50 @@
+import pytest
+
+from firedrake import *
+import numpy as np
+
+
+def test_equation_bcs_direct_assemble_one_form():
+    mesh = UnitSquareMesh(1, 1, quadrilateral=True)
+    V = FunctionSpace(mesh, "CG", 1)
+    u = Function(V).assign(1.)
+    v = TestFunction(V)
+    F = inner(grad(u), grad(v)) * dx
+    F1 = inner(u, v) * ds(1)
+    bc = EquationBC(F1 == 0, u, 1)
+
+    g = assemble(F, bcs=bc.extract_form('F'))
+    assert(np.allclose(g.dat.data, [0.5, 0.5, 0, 0]))
+    g = assemble(F, bcs=bc)
+    assert(np.allclose(g.dat.data, [0.5, 0.5, 0, 0]))
+
+
+def test_equation_bcs_direct_assemble_two_form():
+    mesh = UnitSquareMesh(1, 1, quadrilateral=True)
+    V = FunctionSpace(mesh, "CG", 1)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    a = inner(grad(u), grad(v)) * dx
+    a1 = inner(u, v) * ds(1)
+    L1 = inner(Constant(0), v) * ds(1)
+    sol = Function(V)
+    bc = EquationBC(a1 == L1, sol, 1, Jp=2 * inner(u, v) * ds(1))
+
+    # Must preprocess bc to extract appropriate
+    # `EquationBCSplit` object.
+    A = assemble(a, bcs=bc.extract_form('J'))
+    assert(np.allclose(A.M.values, [[1 / 3, 1 / 6, 0, 0],
+                                    [1 / 6, 1 / 3, 0, 0],
+                                    [-1 / 3, -1 / 6, 2 / 3, -1 / 6],
+                                    [-1 / 6, -1 / 3, -1 / 6, 2 / 3]]))
+    A = assemble(a, bcs=bc.extract_form('Jp'))
+    assert(np.allclose(A.M.values, [[2 / 3, 2 / 6, 0, 0],
+                                    [2 / 6, 2 / 3, 0, 0],
+                                    [-1 / 3, -1 / 6, 2 / 3, -1 / 6],
+                                    [-1 / 6, -1 / 3, -1 / 6, 2 / 3]]))
+    with pytest.raises(TypeError) as excinfo:
+        # Unable to use raw `EquationBC` object, as
+        # assembler can not infer merely from the rank
+        # which form ('J' or 'Jp') should be assembled.
+        assemble(a, bcs=bc)
+    assert "EquationBC objects not expected here" in str(excinfo.value)