Fixes for changes to firedrake.assemble default behaviour.

.github/workflows/test.yml

@@ -61,4 +61,4 @@ jobs:
           python --version
           python -m pytest --version
           firedrake-status
-          python -m pytest -v -n 6 --durations=40 --timeout=600 --cov=asQ --cov-report=term tests/
+          python -m pytest -v -n 6 --durations=40 --timeout=420 --cov=asQ --cov-report=term tests/

asQ/allatonce/form.py

@@ -1,8 +1,7 @@
 import firedrake as fd
-from firedrake.petsc import PETSc
 
 from asQ.profiling import profiler
-from asQ.allatonce import AllAtOnceFunction, AllAtOnceCofunction
+from asQ.allatonce import AllAtOnceCofunction
 from asQ.allatonce.mixin import TimePartitionMixin
 
 from functools import partial
@@ -59,8 +58,9 @@ def __init__(self,
         for bc in self.bcs:
             bc.apply(aaofunc.function)
 
-        # function to assemble the nonlinear residual into
-        self.F = aaofunc.copy(copy_values=False).zero()
+        # cofunction to assemble the nonlinear residual into
+        self.F = AllAtOnceCofunction(self.ensemble, self.time_partition,
+                                     aaofunc.field_function_space.dual())
 
         self.form = self._construct_form()
 
@@ -137,46 +137,39 @@ def assemble(self, func=None, tensor=None):
 
         :arg func: may optionally be an AllAtOnceFunction or a global PETSc Vec.
             if not None, the form will be evaluated at the state in `func`.
-        :arg tensor: may optionally be an AllAtOnceFunction, AllAtOnceCofunction,
-            a global PETSc Vec, or a Function in AllAtOnceFunction.function_space.
-            if not None, the result will be placed into `tensor`.
+        :arg tensor: may optionally be an AllAtOnceCofunction, in which case
+            the result will be placed into `tensor`.
         """
-        # set current state
+        if tensor is not None and not isinstance(tensor, AllAtOnceCofunction):
+            raise TypeError(f"tensor must be an AllAtOnceCofunction, not {type(tensor)}")
+
+        # Set the current state
         if func is not None:
             self.aaofunc.assign(func, update_halos=False)
-        self.aaofunc.update_time_halos()
-
-        # assembly stage
-        fd.assemble(self.form, tensor=self.F.function)
 
-        # apply boundary conditions
+        # Assembly stage
+        # The residual on the DirichletBC nodes is set to zero,
+        # so we need to make sure that the function conforms
+        # with the boundary conditions.
         for bc in self.bcs:
-            bc.apply(self.F.function, u=self.aaofunc.function)
-
-        # copy into return buffer
-
-        if isinstance(tensor, AllAtOnceFunction):
-            tensor.assign(self.F)
+            bc.apply(self.aaofunc.function)
 
-        elif isinstance(tensor, fd.Function):
-            tensor.assign(self.F.function)
-
-        elif isinstance(tensor, AllAtOnceCofunction):
-            with self.F.global_vec_ro() as fvec:
-                with tensor.global_vec_wo() as tvec:
-                    fvec.copy(tvec)
-
-        elif isinstance(tensor, fd.Cofunction):
-            with self.F.function.dat.vec_ro as fvec:
-                with tensor.dat.vec_wo as tvec:
-                    fvec.copy(tvec)
+        # Update the halos after enforcing the bcs so we
+        # know they are correct. This doesn't make a
+        # difference now because we only support constant
+        # bcs, but it will be important once we support
+        # time-dependent bcs.
+        self.aaofunc.update_time_halos()
 
-        elif isinstance(tensor, PETSc.Vec):
-            with self.F.global_vec_ro() as fvec:
-                fvec.copy(tensor)
+        fd.assemble(self.form, bcs=self.bcs,
+                    tensor=self.F.cofunction)
 
-        elif tensor is not None:
-            raise TypeError(f"tensor must be AllAtOnceFunction, AllAtOnceCofunction, Function, Cofunction, or PETSc.Vec, not {type(tensor)}")
+        if tensor:
+            tensor.assign(self.F)
+            result = tensor
+        else:
+            result = self.F
+        return result
 
     def _construct_form(self):
         """

asQ/allatonce/jacobian.py

@@ -55,8 +55,11 @@ def __init__(self, aaoform,
         self.x = aaofunc.copy()
 
         # output residual, and contribution from timestep at end of previous slice
-        self.F = aaofunc.copy()
-        self.Fprev = fd.Function(aaofunc.function_space)
+        self.F = aaoform.F.copy(copy_values=False)
+        self.Fprev = fd.Cofunction(self.F.function_space)
+
+        self.bcs = aaoform.bcs
+        self.field_bcs = aaoform.field_bcs
 
         # working buffers for calculating time average when needed
         self.ureduce = fd.Function(aaofunc.field_function_space)
@@ -138,23 +141,46 @@ def mult(self, mat, X, Y):
         :arg X: a PETSc Vec to apply the action on.
         :arg Y: a PETSc Vec for the result.
         """
-
         # we could use nonblocking here and overlap comms with assembling form
         self.x.assign(X, update_halos=True, blocking=True)
 
+        # We use the same strategy as the implicit matrix context in firedrake
+        # for dealing with the boundary nodes. From the comments in that file:
+
+        # The matrix has an identity block corresponding to the Dirichlet
+        # boundary conditions.
+        # Our algorithm in this case is to save the BC values, zero them
+        # out before computing the action so that they don't pollute
+        # anything, and then set the values into the result.
+        # This has the effect of applying [ A_ii 0 ; 0 A_bb ] where A_ii
+        # is the block corresponding only to (non-fixed) internal dofs
+        # and A_bb=I is the identity block on the (fixed) boundary dofs.
+
+        # Zero the boundary nodes on the input so that A_ib = A_01 = 0
+        for bc in self.bcs:
+            bc.zero(self.x.function)
+
         # assembly stage
-        fd.assemble(self.action, tensor=self.F.function)
+        fd.assemble(self.action, bcs=self.bcs,
+                    tensor=self.F.cofunction)
+
         if self._useprev:
-            fd.assemble(self.action_prev, tensor=self.Fprev)
-            self.F.function += self.Fprev
-
-        # Apply boundary conditions
-        # For Jacobian action we should just return the values in X
-        # at boundary nodes
-        for bc in self.aaoform.bcs:
-            bc.homogenize()
-            bc.apply(self.F.function, u=self.x.function)
-            bc.restore()
+            # repeat for the halo part of the matrix action
+            for bc in self.field_bcs:
+                bc.zero(self.x.uprev)
+            fd.assemble(self.action_prev, bcs=self.bcs,
+                        tensor=self.Fprev)
+            self.F.cofunction += self.Fprev
+
+        if len(self.bcs) > 0:
+            Fbuf = self.Fprev  # just using Fprev as a working buffer
+            # Get the original values again
+            with Fbuf.dat.vec_wo as fvec:
+                X.copy(fvec)
+            # Set the output boundary nodes to the input boundary nodes.
+            # This is equivalent to setting [A_bi, A_bb] = [0 I]
+            for bc in self.bcs:
+                bc.set(self.F.cofunction, Fbuf)
 
         with self.F.global_vec_ro() as v:
             v.copy(Y)

asQ/allatonce/solver.py

@@ -90,15 +90,15 @@ def passthrough(*args, **kwargs):
 
         self.snes.setOptionsPrefix(options_prefix)
 
-        # residual vector
-        self.F = aaofunc._vec.duplicate()
-
         def assemble_function(snes, X, F):
             self.pre_function_callback(self, X)
-            self.aaoform.assemble(X, tensor=F)
+            self.aaoform.assemble(X)
+            with aaoform.F.global_vec_ro() as fvec:
+                fvec.copy(F)
             self.post_function_callback(self, X, F)
 
-        self.snes.setFunction(assemble_function, self.F)
+        self._F = aaoform.F._vec.duplicate()
+        self.snes.setFunction(assemble_function, self._F)
 
         # Jacobian
         with self.options.inserted_options():

asQ/complex_proxy/vector_impl.py

@@ -75,15 +75,11 @@ def DirichletBC(W, V, bc, function_arg=None):
     if function_arg is None:
         function_arg = bc.function_arg
 
-    sub_domain = bc.sub_domain
-
     if type(V.ufl_element()) is fd.MixedElement:
-        idx = bc.function_space().index
-        Ws = (W.sub(idx).sub(0), W.sub(idx).sub(1))
-    else:
-        Ws = (W.sub(0), W.sub(1))
+        W = W.sub(bc.function_space().index)
 
-    return tuple(fd.DirichletBC(Wsub, function_arg, sub_domain) for Wsub in Ws)
+    return tuple(fd.DirichletBC(W.sub(part), function_arg, bc.sub_domain)
+                 for part in (re, im))
 
 
 def split(u, i):

tests/allatonce/test_allatonceform.py

@@ -6,29 +6,26 @@
 from operator import mul
 
 
-def assemble(form):
-    return fd.assemble(form).riesz_representation(riesz_map='l2')
-
-
-bc_opts = ["no_bcs", "homogeneous_bcs", "inhomogeneous_bcs"]
+bc_options = ["no_bcs", "homogeneous_bcs", "inhomogeneous_bcs"]
 
 alphas = [pytest.param(None, id="alpha_None"),
           pytest.param(0, id="alpha_0"),
           pytest.param(0.1, id="alpha_0.1")]
 
 
 @pytest.mark.parallel(nprocs=4)
-@pytest.mark.parametrize("bc_opt", bc_opts)
+@pytest.mark.parametrize("bc_option", bc_options)
 @pytest.mark.parametrize("alpha", alphas)
-def test_heat_form(bc_opt, alpha):
+def test_heat_form(bc_option, alpha):
     """
     Test that assembling the AllAtOnceForm is the same as assembling the
     slice-local part of an all-at-once form for the whole timeseries.
     Diffusion coefficient is time-dependent.
     """
 
     # build the all-at-once function
-    nslices = fd.COMM_WORLD.size//2
+    nspace_ranks = 1
+    nslices = fd.COMM_WORLD.size//nspace_ranks
     slice_length = 2
 
     time_partition = tuple((slice_length for _ in range(nslices)))
@@ -61,11 +58,11 @@ def form_function(u, v, t):
     def form_mass(u, v):
         return u*v*fd.dx
 
-    if bc_opt == "inhomogeneous_bcs":
+    if bc_option == "inhomogeneous_bcs":
         bc_val = fd.sin(2*fd.pi*x)
         bc_domain = "on_boundary"
         bcs = [fd.DirichletBC(V, bc_val, bc_domain)]
-    elif bc_opt == "homogeneous_bcs":
+    elif bc_option == "homogeneous_bcs":
         bc_val = 0.
         bc_domain = 1
         bcs = [fd.DirichletBC(V, bc_val, bc_domain)]
@@ -83,12 +80,12 @@ def form_mass(u, v):
     full_function_space = reduce(mul, (V for _ in range(sum(time_partition))))
     ufull = fd.Function(full_function_space)
 
-    if bc_opt == "no_bcs":
+    if bc_option == "no_bcs":
         bcs_full = []
     else:
-        bcs_full = []
-        for i in range(sum(time_partition)):
-            bcs_full.append(fd.DirichletBC(full_function_space.sub(i), bc_val, bc_domain))
+        bcs_full = [fd.DirichletBC(full_function_space.sub(i),
+                                   bc_val, bc_domain)
+                    for i in range(sum(time_partition))]
 
     vfull = fd.TestFunction(full_function_space)
     ufulls = fd.split(ufull)
@@ -119,21 +116,25 @@ def form_mass(u, v):
 
     # assemble and compare
     aaoform.assemble()
-    Ffull = assemble(fullform)
+
+    # aaoform.assemble enforces boundary
+    # conditions, but firedrake.assemble doesn't
+    # so we need to manually enforce them here.
     for bc in bcs_full:
-        bc.apply(Ffull, u=ufull)
+        bc.apply(ufull)
+    Ffull = fd.assemble(fullform, bcs=bcs_full)
 
     for step in range(aaofunc.nlocal_timesteps):
         windx = aaofunc.transform_index(step, from_range='slice', to_range='window')
         userial = Ffull.subfunctions[windx]
         uparallel = aaoform.F[step].subfunctions[0]
-        err = fd.errornorm(userial, uparallel)
-        assert (err < 1e-12), "Each component of AllAtOnceForm residual should match component of monolithic residual calculated locally"
+        for pdat, sdat in zip(uparallel.dat, userial.dat):
+            assert np.allclose(pdat.data, sdat.data), "Each component of AllAtOnceForm residual should match component of monolithic residual calculated locally"
 
 
-@pytest.mark.parametrize("bc_opt", bc_opts)
+@pytest.mark.parametrize("bc_option", bc_options)
 @pytest.mark.parallel(nprocs=4)
-def test_mixed_heat_form(bc_opt):
+def test_mixed_heat_form(bc_option):
     """
     Test that assembling the AllAtOnceForm is the same as assembling the
     slice-local part of an all-at-once form for the whole timeseries.
@@ -172,11 +173,11 @@ def form_function(u, p, v, q, t):
     def form_mass(u, p, v, q):
         return (fd.inner(u, v) + p*q)*fd.dx
 
-    if bc_opt == "inhomogeneous_bcs":
+    if bc_option == "inhomogeneous_bcs":
         bc_val = fd.as_vector([fd.sin(2*fd.pi*x), -fd.cos(fd.pi*y)])
         bc_domain = "on_boundary"
         bcs = [fd.DirichletBC(V.sub(0), bc_val, bc_domain)]
-    elif bc_opt == "homogeneous_bcs":
+    elif bc_option == "homogeneous_bcs":
         bc_val = fd.as_vector([0., 0.])
         bc_domain = "on_boundary"
         bcs = [fd.DirichletBC(V.sub(0), bc_val, bc_domain)]
@@ -191,14 +192,12 @@ def form_mass(u, p, v, q):
     full_function_space = reduce(mul, (V for _ in range(sum(time_partition))))
     ufull = fd.Function(full_function_space)
 
-    if bc_opt == "no_bcs":
+    if bc_option == "no_bcs":
         bcs_full = []
     else:
-        bcs_full = []
-        for i in range(sum(time_partition)):
-            bcs_full.append(fd.DirichletBC(full_function_space.sub(2*i),
-                                           bc_val,
-                                           bc_domain))
+        bcs_full = [fd.DirichletBC(full_function_space.sub(2*i),
+                                   bc_val, bc_domain)
+                    for i in range(sum(time_partition))]
 
     vfull = fd.TestFunction(full_function_space)
     ufulls = fd.split(ufull)
@@ -233,18 +232,22 @@ def form_mass(u, p, v, q):
 
     # assemble and compare
     aaoform.assemble()
-    Ffull = assemble(fullform)
+
+    # aaoform.assemble enforces boundary
+    # conditions, but firedrake.assemble doesn't
+    # so we need to manually enforce them here.
     for bc in bcs_full:
-        bc.apply(Ffull, u=ufull)
+        bc.apply(ufull)
+    Ffull = fd.assemble(fullform, bcs=bcs_full)
 
     for step in range(aaofunc.nlocal_timesteps):
         windices = aaofunc._component_indices(step, to_range='window')
         for cpt in range(2):
             windx = windices[cpt]
             userial = Ffull.subfunctions[windx]
             uparallel = aaoform.F[step].subfunctions[cpt]
-            err = fd.errornorm(userial, uparallel)
-            assert (err < 1e-12), "Each component of AllAtOnceForm residual should match component of monolithic residual calculated locally"
+            for pdat, sdat in zip(uparallel.dat, userial.dat):
+                assert np.allclose(pdat.data, sdat.data), "Each component of AllAtOnceForm residual should match component of monolithic residual calculated locally"
 
 
 @pytest.mark.parallel(nprocs=4)