diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index a1f1925fd1..f0500c8b91 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/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() diff --git a/tests/slate/test_cg_poisson.py b/tests/slate/test_cg_poisson.py new file mode 100644 index 0000000000..f04d545ef1 --- /dev/null +++ b/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()