Merge pull request #2436 from firedrakeproject/pbrubeck/fix/scpc-weig…

…ht-bcs

Fix weights and BCs when assembling the condensed residual in SCPC

firedrake/slate/static_condensation/scpc.py

@@ -1,4 +1,5 @@
 import firedrake.dmhooks as dmhooks
+import numpy as np
 from firedrake.slate.static_condensation.sc_base import SCBase
 from firedrake.matrix_free.operators import ImplicitMatrixContext
 from firedrake.petsc import PETSc
@@ -30,6 +31,8 @@ def initialize(self, pc):
         from firedrake.bcs import DirichletBC
         from firedrake.function import Function
         from firedrake.functionspace import FunctionSpace
+        from firedrake.parloops import par_loop, INC
+        from ufl import dx
 
         prefix = pc.getOptionsPrefix() + "condensed_field_"
         A, P = pc.getOperators()
@@ -68,7 +71,7 @@ def initialize(self, pc):
         for bc in cxt_bcs:
             if bc.function_space().index != c_field:
                 raise NotImplementedError("Strong BC set on unsupported space")
-            bcs.append(DirichletBC(Vc, bc.function_arg, bc.sub_domain))
+            bcs.append(DirichletBC(Vc, 0, bc.sub_domain))
 
         mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
 
@@ -77,6 +80,20 @@ def initialize(self, pc):
         self.residual = Function(W)
         self.solution = Function(W)
 
+        shapes = (Vc.finat_element.space_dimension(),
+                  np.prod(Vc.shape))
+        domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
+        instructions = """
+        for i, j
+            w[i,j] = w[i,j] + 1
+        end
+        """
+        self.weight = Function(Vc)
+        par_loop((domain, instructions), dx, {"w": (self.weight, INC)},
+                 is_loopy_kernel=True)
+        with self.weight.dat.vec as wc:
+            wc.reciprocal()
+
         # Get expressions for the condensed linear system
         A_tensor = Tensor(self.bilinear_form)
         reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields)
@@ -85,6 +102,7 @@ def initialize(self, pc):
 
         # Construct the condensed right-hand side
         self._assemble_Srhs = OneFormAssembler(r_expr, tensor=self.condensed_rhs,
+                                               bcs=bcs, zero_bc_nodes=True,
                                                form_compiler_parameters=self.cxt.fc_params).assemble
 
         # Allocate and set the condensed operator
@@ -237,6 +255,10 @@ def forward_elimination(self, pc, x):
         with self.residual.dat.vec_wo as v:
             x.copy(v)
 
+        # Disassemble the incoming right-hand side
+        with self.residual.split()[self.c_field].dat.vec as vc, self.weight.dat.vec_ro as wc:
+            vc.pointwiseMult(vc, wc)
+
         # Now assemble residual for the reduced problem
         self._assemble_Srhs()
 

tests/slate/test_cg_poisson.py

@@ -0,0 +1,63 @@
+import pytest
+from firedrake import *
+
+
+def run_CG_problem(r, degree, quads=False):
+    """
+    Solves the Dirichlet problem for the elliptic equation:
+
+    -div(grad(u)) = f in [0, 1]^2, u = g on the domain boundary.
+
+    The source function f and g are chosen such that the analytic
+    solution is:
+
+    u(x, y) = sin(x*pi)*sin(y*pi).
+
+    This test uses a CG discretization splitting interior and facet DOFs
+    and Slate to perform the static condensation and local recovery.
+    """
+
+    # Set up problem domain
+    mesh = UnitSquareMesh(2**r, 2**r, quadrilateral=quads)
+    x = SpatialCoordinate(mesh)
+    u_exact = sin(x[0]*pi)*sin(x[1]*pi)
+    f = -div(grad(u_exact))
+
+    # Set up function spaces
+    e = FiniteElement("Lagrange", cell=mesh.ufl_cell(), degree=degree)
+    Z = FunctionSpace(mesh, MixedElement(InteriorElement(e), FacetElement(e)))
+    z = Function(Z)
+    u = sum(split(z))
+
+    # Formulate the CG method in UFL
+    U = (1/2)*inner(grad(u), grad(u))*dx - inner(u, f)*dx
+    F = derivative(U, z, TestFunction(Z))
+
+    params = {'snes_type': 'ksponly',
+              'mat_type': 'matfree',
+              'pmat_type': 'matfree',
+              'ksp_type': 'preonly',
+              'pc_type': 'python',
+              'pc_python_type': 'firedrake.SCPC',
+              'pc_sc_eliminate_fields': '0',
+              'condensed_field': {'ksp_type': 'preonly',
+                                  'pc_type': 'redundant',
+                                  "redundant_pc_type": "lu",
+                                  "redundant_pc_factor_mat_solver_type": "mumps"}}
+
+    bcs = DirichletBC(Z.sub(1), zero(), "on_boundary")
+    problem = NonlinearVariationalProblem(F, z, bcs=bcs)
+    solver = NonlinearVariationalSolver(problem, solver_parameters=params)
+    solver.solve()
+    return norm(u_exact-u, norm_type="L2")
+
+
+@pytest.mark.parallel
+@pytest.mark.parametrize(('degree', 'quads', 'rate'),
+                         [(3, False, 3.75),
+                          (5, True, 5.75)])
+def test_cg_convergence(degree, quads, rate):
+    import numpy as np
+    diff = np.array([run_CG_problem(r, degree, quads) for r in range(2, 5)])
+    conv = np.log2(diff[:-1] / diff[1:])
+    assert (np.array(conv) > rate).all()