def _reconstruction_calls(self, split_mixed_op, split_trace_op): """This generates the reconstruction calls for the unknowns using the Lagrange multipliers. :arg split_mixed_op: a ``dict`` of split forms that make up the broken mixed operator from the original problem. :arg split_trace_op: a ``dict`` of split forms that make up the trace contribution in the hybridized mixed system. """ from firedrake.assemble import create_assembly_callable # We always eliminate the velocity block first id0, id1 = (self.vidx, self.pidx) # TODO: When PyOP2 is able to write into mixed dats, # the reconstruction expressions can simplify into # one clean expression. A = Tensor(split_mixed_op[(id0, id0)]) B = Tensor(split_mixed_op[(id0, id1)]) C = Tensor(split_mixed_op[(id1, id0)]) D = Tensor(split_mixed_op[(id1, id1)]) K_0 = Tensor(split_trace_op[(0, id0)]) K_1 = Tensor(split_trace_op[(0, id1)]) # Split functions and reconstruct each bit separately split_residual = self.broken_residual.split() split_sol = self.broken_solution.split() g = AssembledVector(split_residual[id0]) f = AssembledVector(split_residual[id1]) sigma = split_sol[id0] u = split_sol[id1] lambdar = AssembledVector(self.trace_solution) M = D - C * A.inv * B R = K_1.T - C * A.inv * K_0.T u_rec = M.solve(f - C * A.inv * g - R * lambdar, decomposition="PartialPivLU") self._sub_unknown = create_assembly_callable(u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params) sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar, decomposition="PartialPivLU") self._elim_unknown = create_assembly_callable(sigma_rec, tensor=sigma, form_compiler_parameters=self.ctx.fc_params)
def _reconstruction_calls(self, split_mixed_op, split_trace_op): """This generates the reconstruction calls for the unknowns using the Lagrange multipliers. :arg split_mixed_op: a ``dict`` of split forms that make up the broken mixed operator from the original problem. :arg split_trace_op: a ``dict`` of split forms that make up the trace contribution in the hybridized mixed system. """ from firedrake.assemble import create_assembly_callable # We always eliminate the velocity block first id0, id1 = (self.vidx, self.pidx) # TODO: When PyOP2 is able to write into mixed dats, # the reconstruction expressions can simplify into # one clean expression. A = Tensor(split_mixed_op[(id0, id0)]) B = Tensor(split_mixed_op[(id0, id1)]) C = Tensor(split_mixed_op[(id1, id0)]) D = Tensor(split_mixed_op[(id1, id1)]) K_0 = Tensor(split_trace_op[(0, id0)]) K_1 = Tensor(split_trace_op[(0, id1)]) # Split functions and reconstruct each bit separately split_residual = self.broken_residual.split() split_sol = self.broken_solution.split() g = AssembledVector(split_residual[id0]) f = AssembledVector(split_residual[id1]) sigma = split_sol[id0] u = split_sol[id1] lambdar = AssembledVector(self.trace_solution) M = D - C * A.inv * B R = K_1.T - C * A.inv * K_0.T u_rec = M.solve(f - C * A.inv * g - R * lambdar, decomposition="PartialPivLU") self._sub_unknown = create_assembly_callable( u_rec, tensor=u, form_compiler_parameters=self.ctx.fc_params) sigma_rec = A.solve(g - B * AssembledVector(u) - K_0.T * lambdar, decomposition="PartialPivLU") self._elim_unknown = create_assembly_callable( sigma_rec, tensor=sigma, form_compiler_parameters=self.ctx.fc_params)
def retrieve_user_S_approx(self, pc, usercode): """Retrieve a user-defined :class:firedrake.preconditioners.AuxiliaryOperator from the PETSc Options, which is an approximation to the Schur complement and its inverse is used to precondition the local solve in the reconstruction calls (e.g.). """ _, _, _, A11 = self.list_split_mixed_ops test, trial = A11.arguments() if usercode != "": (modname, funname) = usercode.rsplit('.', 1) mod = __import__(modname) fun = getattr(mod, funname) if isinstance(fun, type): fun = fun() return Tensor(fun.form(pc, test, trial)[0]) else: return None
def _split_mixed_operator(self): split_mixed_op = dict(split_form(self.Atilde.form)) id0, id1 = (self.vidx, self.pidx) A00 = Tensor(split_mixed_op[(id0, id0)]) A01 = Tensor(split_mixed_op[(id0, id1)]) A10 = Tensor(split_mixed_op[(id1, id0)]) A11 = Tensor(split_mixed_op[(id1, id1)]) self.list_split_mixed_ops = [A00, A01, A10, A11] split_trace_op = dict(split_form(self.K.form)) K0 = Tensor(split_trace_op[(0, id0)]) K1 = Tensor(split_trace_op[(0, id1)]) self.list_split_trace_ops = [K0, K1]
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC, assemble) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.cxt = P.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError("The python context must be an ImplicitMatrixContext") test, trial = self.cxt.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up the KSP for the hdiv residual projection hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_") # HDiv mass operator p = TrialFunction(V[self.vidx]) q = TestFunction(V[self.vidx]) mass = ufl.dot(p, q)*ufl.dx # TODO: Bcs? M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params) M.force_evaluation() Mmat = M.petscmat hdiv_mass_ksp.setOperators(Mmat) hdiv_mass_ksp.setUp() hdiv_mass_ksp.setFromOptions() self.hdiv_mass_ksp = hdiv_mass_ksp # Storing the result of A.inv * r, where A is the HDiv # mass matrix and r is the HDiv residual self._primal_r = Function(V[self.vidx]) tau = TestFunction(V_d[self.vidx]) self._assemble_broken_r = create_assembly_callable( ufl.dot(self._primal_r, tau)*ufl.dx, tensor=self.broken_residual.split()[self.vidx], form_compiler_parameters=self.cxt.fc_params) # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.cxt.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] # We zero out the contribution of the trace variables on the exterior # boundary. Extruded cells will have both horizontal and vertical # facets if mesh.cell_set._extruded: trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary"), DirichletBC(TraceSpace, Constant(0.0), "bottom"), DirichletBC(TraceSpace, Constant(0.0), "top")] K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS_h + gammar('+') * ufl.dot(sigma, n) * ufl.dS_v) else: trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")] K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS) # If boundary conditions are contained in the ImplicitMatrixContext: if self.cxt.row_bcs: raise NotImplementedError("Strong BCs not currently handled. Try imposing them weakly.") # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * self.broken_residual, tensor=self.schur_rhs, form_compiler_parameters=self.cxt.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S = create_assembly_callable(schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullspace = create_schur_nullspace(P, -K * Atilde, V, V_d, TraceSpace, pc.comm) if nullspace: Smat.setNullSpace(nullspace) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op) # NOTE: The projection stage *might* be replaced by a Fortin # operator. We may want to allow the user to specify if they # wish to use a Fortin operator over a projection, or vice-versa. # In a future add-on, we can add a switch which chooses either # the Fortin reconstruction or the usual KSP projection. # Set up the projection KSP hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_') # Reuse the mass operator from the hdiv_mass_ksp hdiv_projection_ksp.setOperators(Mmat) # Construct the RHS for the projection stage self._projection_rhs = Function(V[self.vidx]) self._assemble_projection_rhs = create_assembly_callable( ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx, tensor=self._projection_rhs, form_compiler_parameters=self.cxt.fc_params) # Finalize ksp setup hdiv_projection_ksp.setUp() hdiv_projection_ksp.setFromOptions() self.hdiv_projection_ksp = hdiv_projection_ksp
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.ctx = P.getPythonContext() if not isinstance(self.ctx, ImplicitMatrixContext): raise ValueError( "The python context must be an ImplicitMatrixContext") test, trial = self.ctx.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement( [ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) shapes = (V[self.vidx].finat_element.space_dimension(), np.prod(V[self.vidx].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(V[self.vidx]) par_loop((domain, instructions), ufl.dx, {"w": (self.weight, INC)}, is_loopy_kernel=True) instructions = """ for i, j vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j] end """ self.average_kernel = (domain, instructions) # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.ctx.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] if mesh.cell_set._extruded: Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h + gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v) else: Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS) # Here we deal with boundaries. If there are Neumann # conditions (which should be enforced strongly for # H(div)xL^2) then we need to add jump terms on the exterior # facets. If there are Dirichlet conditions (which should be # enforced weakly) then we need to zero out the trace # variables there as they are not active (otherwise the hybrid # problem is not well-posed). # If boundary conditions are contained in the ImplicitMatrixContext: if self.ctx.row_bcs: # Find all the subdomains with neumann BCS # These are Dirichlet BCs on the vidx space neumann_subdomains = set() for bc in self.ctx.row_bcs: if bc.function_space().index == self.pidx: raise NotImplementedError( "Dirichlet conditions for scalar variable not supported. Use a weak bc" ) if bc.function_space().index != self.vidx: raise NotImplementedError( "Dirichlet bc set on unsupported space.") # append the set of sub domains subdom = bc.sub_domain if isinstance(subdom, str): neumann_subdomains |= set([subdom]) else: neumann_subdomains |= set( as_tuple(subdom, numbers.Integral)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & { "top", "bottom" } neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains integrand = gammar * ufl.dot(sigma, n) measures = [] trace_subdomains = [] if mesh.cell_set._extruded: ds = ufl.ds_v for subdomain in sorted(extruded_neumann_subdomains): measures.append({ "top": ufl.ds_t, "bottom": ufl.ds_b }[subdomain]) trace_subdomains.extend( sorted({"top", "bottom"} - extruded_neumann_subdomains)) else: ds = ufl.ds if "on_boundary" in neumann_subdomains: measures.append(ds) else: measures.extend((ds(sd) for sd in sorted(neumann_subdomains))) markers = [int(x) for x in mesh.exterior_facets.unique_markers] dirichlet_subdomains = set(markers) - neumann_subdomains trace_subdomains.extend(sorted(dirichlet_subdomains)) for measure in measures: Kform += integrand * measure trace_bcs = [ DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains ] else: # No bcs were provided, we assume weak Dirichlet conditions. # We zero out the contribution of the trace variables on # the exterior boundary. Extruded cells will have both # horizontal and vertical facets trace_subdomains = ["on_boundary"] if mesh.cell_set._extruded: trace_subdomains.extend(["bottom", "top"]) trace_bcs = [ DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains ] # Make a SLATE tensor from Kform K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, form_compiler_parameters=self.ctx.fc_params) mat_type = PETSc.Options().getString(prefix + "mat_type", "aij") schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type, options_prefix=prefix) self._assemble_S = create_assembly_callable( schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.ctx.fc_params, mat_type=mat_type) with timed_region("HybridOperatorAssembly"): self._assemble_S() Smat = self.S.petscmat nullspace = self.ctx.appctx.get("trace_nullspace", None) if nullspace is not None: nsp = nullspace(TraceSpace) Smat.setNullSpace(nsp.nullspace(comm=pc.comm)) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op)
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from ufl.algorithms.map_integrands import map_integrand_dags from firedrake import (FunctionSpace, TrialFunction, TrialFunctions, TestFunction, Function, BrokenElement, MixedElement, FacetNormal, Constant, DirichletBC, Projector) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import ArgumentReplacer, split_form # Extract the problem context prefix = pc.getOptionsPrefix() _, P = pc.getOperators() context = P.getPythonContext() test, trial = context.a.arguments() V = test.function_space() if V.mesh().cell_set._extruded: # TODO: Merge FIAT branch to support TPC trace elements raise NotImplementedError("Not implemented on extruded meshes.") # Break the function spaces and define fully discontinuous spaces broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V] elem = MixedElement(broken_elements) V_d = FunctionSpace(V.mesh(), elem) arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} # Replace the problems arguments with arguments defined # on the new discontinuous spaces replacer = ArgumentReplacer(arg_map) new_form = map_integrand_dags(replacer, context.a) # Create the space of approximate traces. # The vector function space will have a non-empty value_shape W = next(v for v in V if bool(v.ufl_element().value_shape())) if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]: tdegree = W.ufl_element().degree() - 1 else: tdegree = W.ufl_element().degree() # NOTE: Once extruded is ready, we will need to be aware of this # and construct the appropriate trace space for the HDiv element TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree) # NOTE: For extruded, we will need to add "on_top" and "on_bottom" trace_conditions = [ DirichletBC(TraceSpace, Constant(0.0), "on_boundary") ] # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_rhs = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_rhs = Function(V) # Create the symbolic Schur-reduction Atilde = Tensor(new_form) gammar = TestFunction(TraceSpace) n = FacetNormal(V.mesh()) # Vector trial function will have a non-empty ufl_shape sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape)) # NOTE: Once extruded is ready, this will change slightly # to include both horizontal and vertical interior facets K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * self.broken_rhs, tensor=self.schur_rhs, form_compiler_parameters=context.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_conditions, form_compiler_parameters=context.fc_params) self._assemble_S = create_assembly_callable( schur_comp, tensor=self.S, bcs=trace_conditions, form_compiler_parameters=context.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullsp = P.getNullSpace() if nullsp.handle != 0: new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d, TraceSpace) tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs, comm=pc.comm) Smat.setNullSpace(tr_nullsp) # Set up the KSP for the system of Lagrange multipliers ksp = PETSc.KSP().create(comm=pc.comm) ksp.setOptionsPrefix(prefix + "trace_") ksp.setTolerances(rtol=1e-13) ksp.setOperators(Smat) ksp.setUp() ksp.setFromOptions() self.ksp = ksp # Now we construct the local tensors for the reconstruction stage # TODO: Add support for mixed tensors and these variables # become unnecessary split_forms = split_form(new_form) A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0))) B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0))) C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1))) trial = TrialFunction( FunctionSpace(V.mesh(), BrokenElement(W.ufl_element()))) K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS) # Split functions and reconstruct each bit separately sigma_h, u_h = self.broken_solution.split() g, f = self.broken_rhs.split() # Pressure reconstruction M = B * A.inv * B.T + C u_sol = M.inv * f + M.inv * ( B * A.inv * K_local.T * self.trace_solution - B * A.inv * g) self._assemble_pressure = create_assembly_callable( u_sol, tensor=u_h, form_compiler_parameters=context.fc_params) # Velocity reconstruction sigma_sol = A.inv * g + A.inv * (B.T * u_h - K_local.T * self.trace_solution) self._assemble_velocity = create_assembly_callable( sigma_sol, tensor=sigma_h, form_compiler_parameters=context.fc_params) # Set up the projector for projecting the broken solution # into the unbroken finite element spaces # NOTE: Tolerance here matters! sigma_b, _ = self.broken_solution.split() sigma_u, _ = self.unbroken_solution.split() self.projector = Projector(sigma_b, sigma_u, solver_parameters={ "ksp_type": "cg", "ksp_rtol": 1e-13 })
def initialize(self, pc): """Set up the problem context. This takes the incoming three-field system and constructs the static condensation operators using Slate expressions. A KSP is created for the reduced system. The eliminated variables are recovered via back-substitution. """ from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.bcs import DirichletBC from firedrake.function import Function from firedrake.functionspace import FunctionSpace from firedrake.interpolation import interpolate prefix = pc.getOptionsPrefix() + "condensed_field_" A, P = pc.getOperators() self.cxt = A.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError("Context must be an ImplicitMatrixContext") self.bilinear_form = self.cxt.a # Retrieve the mixed function space W = self.bilinear_form.arguments()[0].function_space() if len(W) > 3: raise NotImplementedError("Only supports up to three function spaces.") elim_fields = PETSc.Options().getString(pc.getOptionsPrefix() + "pc_sc_eliminate_fields", None) if elim_fields: elim_fields = [int(i) for i in elim_fields.split(',')] else: # By default, we condense down to the last field in the # mixed space. elim_fields = [i for i in range(0, len(W) - 1)] condensed_fields = list(set(range(len(W))) - set(elim_fields)) if len(condensed_fields) != 1: raise NotImplementedError("Cannot condense to more than one field") c_field, = condensed_fields # Need to duplicate a space which is NOT # associated with a subspace of a mixed space. Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element()) bcs = [] cxt_bcs = self.cxt.row_bcs for bc in cxt_bcs: if bc.function_space().index != c_field: raise NotImplementedError("Strong BC set on unsupported space") if isinstance(bc.function_arg, Function): bc_arg = interpolate(bc.function_arg, Vc) else: # Constants don't need to be interpolated bc_arg = bc.function_arg bcs.append(DirichletBC(Vc, bc_arg, bc.sub_domain)) mat_type = PETSc.Options().getString(prefix + "mat_type", "aij") self.c_field = c_field self.condensed_rhs = Function(Vc) self.residual = Function(W) self.solution = Function(W) # Get expressions for the condensed linear system A_tensor = Tensor(self.bilinear_form) reduced_sys = self.condensed_system(A_tensor, self.residual, elim_fields) S_expr = reduced_sys.lhs r_expr = reduced_sys.rhs # Construct the condensed right-hand side self._assemble_Srhs = create_assembly_callable( r_expr, tensor=self.condensed_rhs, form_compiler_parameters=self.cxt.fc_params) # Allocate and set the condensed operator self.S = allocate_matrix(S_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) self._assemble_S = create_assembly_callable( S_expr, tensor=self.S, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type) self._assemble_S() Smat = self.S.petscmat # If a different matrix is used for preconditioning, # assemble this as well if A != P: self.cxt_pc = P.getPythonContext() P_tensor = Tensor(self.cxt_pc.a) P_reduced_sys = self.condensed_system(P_tensor, self.residual, elim_fields) S_pc_expr = P_reduced_sys.lhs self.S_pc_expr = S_pc_expr # Allocate and set the condensed operator self.S_pc = allocate_matrix(S_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) self._assemble_S_pc = create_assembly_callable( S_pc_expr, tensor=self.S_pc, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type) self._assemble_S_pc() Smat_pc = self.S_pc.petscmat else: self.S_pc_expr = S_expr Smat_pc = Smat # Get nullspace for the condensed operator (if any). # This is provided as a user-specified callback which # returns the basis for the nullspace. nullspace = self.cxt.appctx.get("condensed_field_nullspace", None) if nullspace is not None: nsp = nullspace(Vc) Smat.setNullSpace(nsp.nullspace(comm=pc.comm)) # Create a SNESContext for the DM associated with the trace problem self._ctx_ref = self.new_snes_ctx(pc, S_expr, bcs, mat_type, self.cxt.fc_params, options_prefix=prefix) # Push new context onto the dm associated with the condensed problem c_dm = Vc.dm # Set up ksp for the condensed problem c_ksp = PETSc.KSP().create(comm=pc.comm) c_ksp.incrementTabLevel(1, parent=pc) # Set the dm for the condensed solver c_ksp.setDM(c_dm) c_ksp.setDMActive(False) c_ksp.setOptionsPrefix(prefix) c_ksp.setOperators(A=Smat, P=Smat_pc) self.condensed_ksp = c_ksp with dmhooks.add_hooks(c_dm, self, appctx=self._ctx_ref, save=False): c_ksp.setFromOptions() # Set up local solvers for backwards substitution self.local_solvers = self.local_solver_calls(A_tensor, self.residual, self.solution, elim_fields)
def initialize(self, pc): """Set up the problem context. Take the original mixed problem and reformulate the problem as a hybridized mixed system. A KSP is created for the Lagrange multiplier system. """ from firedrake import (FunctionSpace, Function, Constant, TrialFunction, TrialFunctions, TestFunction, DirichletBC, assemble) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace # Extract the problem context prefix = pc.getOptionsPrefix() + "hybridization_" _, P = pc.getOperators() self.cxt = P.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError("The python context must be an ImplicitMatrixContext") test, trial = self.cxt.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) != 2: raise ValueError("Expecting two function spaces.") if all(Vi.ufl_element().value_shape() for Vi in V): raise ValueError("Expecting an H(div) x L2 pair of spaces.") # Automagically determine which spaces are vector and scalar for i, Vi in enumerate(V): if Vi.ufl_element().sobolev_space().name == "HDiv": self.vidx = i else: assert Vi.ufl_element().sobolev_space().name == "L2" self.pidx = i # Create the space of approximate traces. W = V[self.vidx] if W.ufl_element().family() == "Brezzi-Douglas-Marini": tdegree = W.ufl_element().degree() else: try: # If we have a tensor product element h_deg, v_deg = W.ufl_element().degree() tdegree = (h_deg - 1, v_deg - 1) except TypeError: tdegree = W.ufl_element().degree() - 1 TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up the functions for the original, hybridized # and schur complement systems self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(TraceSpace) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up the KSP for the hdiv residual projection hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_") # HDiv mass operator p = TrialFunction(V[self.vidx]) q = TestFunction(V[self.vidx]) mass = ufl.dot(p, q)*ufl.dx # TODO: Bcs? M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params) M.force_evaluation() Mmat = M.petscmat hdiv_mass_ksp.setOperators(Mmat) hdiv_mass_ksp.setUp() hdiv_mass_ksp.setFromOptions() self.hdiv_mass_ksp = hdiv_mass_ksp # Storing the result of A.inv * r, where A is the HDiv # mass matrix and r is the HDiv residual self._primal_r = Function(V[self.vidx]) tau = TestFunction(V_d[self.vidx]) self._assemble_broken_r = create_assembly_callable( ufl.dot(self._primal_r, tau)*ufl.dx, tensor=self.broken_residual.split()[self.vidx], form_compiler_parameters=self.cxt.fc_params) # Create the symbolic Schur-reduction: # Original mixed operator replaced with "broken" # arguments arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.cxt.a, arg_map)) gammar = TestFunction(TraceSpace) n = ufl.FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] if mesh.cell_set._extruded: Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS_h + gammar('+') * ufl.dot(sigma, n) * ufl.dS_v) else: Kform = (gammar('+') * ufl.dot(sigma, n) * ufl.dS) # Here we deal with boundaries. If there are Neumann # conditions (which should be enforced strongly for # H(div)xL^2) then we need to add jump terms on the exterior # facets. If there are Dirichlet conditions (which should be # enforced weakly) then we need to zero out the trace # variables there as they are not active (otherwise the hybrid # problem is not well-posed). # If boundary conditions are contained in the ImplicitMatrixContext: if self.cxt.row_bcs: # Find all the subdomains with neumann BCS # These are Dirichlet BCs on the vidx space neumann_subdomains = set() for bc in self.cxt.row_bcs: if bc.function_space().index == self.pidx: raise NotImplementedError("Dirichlet conditions for scalar variable not supported. Use a weak bc") if bc.function_space().index != self.vidx: raise NotImplementedError("Dirichlet bc set on unsupported space.") # append the set of sub domains subdom = bc.sub_domain if isinstance(subdom, str): neumann_subdomains |= set([subdom]) else: neumann_subdomains |= set(as_tuple(subdom, int)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & {"top", "bottom"} neumann_subdomains = neumann_subdomains.difference(extruded_neumann_subdomains) integrand = gammar * ufl.dot(sigma, n) measures = [] trace_subdomains = [] if mesh.cell_set._extruded: ds = ufl.ds_v for subdomain in extruded_neumann_subdomains: measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain]) trace_subdomains.extend(sorted({"top", "bottom"} - extruded_neumann_subdomains)) else: ds = ufl.ds if "on_boundary" in neumann_subdomains: measures.append(ds) else: measures.append(ds(tuple(neumann_subdomains))) dirichlet_subdomains = set(mesh.exterior_facets.unique_markers) - neumann_subdomains trace_subdomains.append(sorted(dirichlet_subdomains)) for measure in measures: Kform += integrand*measure trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] else: # No bcs were provided, we assume weak Dirichlet conditions. # We zero out the contribution of the trace variables on # the exterior boundary. Extruded cells will have both # horizontal and vertical facets trace_subdomains = ["on_boundary"] if mesh.cell_set._extruded: trace_subdomains.extend(["bottom", "top"]) trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] # Make a SLATE tensor from Kform K = Tensor(Kform) # Assemble the Schur complement operator and right-hand side self.schur_rhs = Function(TraceSpace) self._assemble_Srhs = create_assembly_callable( K * Atilde.inv * AssembledVector(self.broken_residual), tensor=self.schur_rhs, form_compiler_parameters=self.cxt.fc_params) schur_comp = K * Atilde.inv * K.T self.S = allocate_matrix(schur_comp, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S = create_assembly_callable(schur_comp, tensor=self.S, bcs=trace_bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat # Nullspace for the multiplier problem nullspace = create_schur_nullspace(P, -K * Atilde, V, V_d, TraceSpace, pc.comm) if nullspace: Smat.setNullSpace(nullspace) # Set up the KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat) trace_ksp.setUp() trace_ksp.setFromOptions() self.trace_ksp = trace_ksp split_mixed_op = dict(split_form(Atilde.form)) split_trace_op = dict(split_form(K.form)) # Generate reconstruction calls self._reconstruction_calls(split_mixed_op, split_trace_op) # NOTE: The projection stage *might* be replaced by a Fortin # operator. We may want to allow the user to specify if they # wish to use a Fortin operator over a projection, or vice-versa. # In a future add-on, we can add a switch which chooses either # the Fortin reconstruction or the usual KSP projection. # Set up the projection KSP hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm) hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_') # Reuse the mass operator from the hdiv_mass_ksp hdiv_projection_ksp.setOperators(Mmat) # Construct the RHS for the projection stage self._projection_rhs = Function(V[self.vidx]) self._assemble_projection_rhs = create_assembly_callable( ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx, tensor=self._projection_rhs, form_compiler_parameters=self.cxt.fc_params) # Finalize ksp setup hdiv_projection_ksp.setUp() hdiv_projection_ksp.setFromOptions() self.hdiv_projection_ksp = hdiv_projection_ksp
def _slate_expressions(self): """Returns all the relevant Slate expressions for the static condensation and local recovery procedures. """ # This operator has the form: # | A B C | # | D E F | # | G H J | # NOTE: It is often the case that D = B.T, # G = C.T, H = F.T, and J = 0, but we're not making # that assumption here. _O = Tensor(self.cxt.a) O = _O.blocks # Extract sub-block: # | A B | # | D E | # which has block row indices (0, 1) and block # column indices (0, 1) as well. M = O[:2, :2] # Extract sub-block: # | C | # | F | # which has block row indices (0, 1) and block # column indices (2,) K = O[:2, 2] # Extract sub-block: # | G H | # which has block row indices (2,) and block column # indices (0, 1) L = O[2, :2] # And the final block J has block row-column # indices (2, 2) J = O[2, 2] # Schur complement for traces S = J - L * M.inv * K # Create mixed function for residual computation. # This projects the non-trace residual bits into # the trace space: # -L * M.inv * | v1 v2 |^T _R = AssembledVector(self.residual) R = _R.blocks v1v2 = R[:2] v3 = R[2] r_lambda = v3 - L * M.inv * v1v2 # Reconstruction expressions q_h, u_h, lambda_h = self.solution.split() # Local tensors needed for reconstruction A = O[0, 0] B = O[0, 1] C = O[0, 2] D = O[1, 0] E = O[1, 1] F = O[1, 2] Se = E - D * A.inv * B Sf = F - D * A.inv * C v1, v2, v3 = self.residual.split() # Solve locally using Cholesky factorizations # (Se and A are symmetric positive definite) u_h_expr = Se.solve(AssembledVector(v2) - D * A.inv * AssembledVector(v1) - Sf * AssembledVector(lambda_h), decomposition="LLT") q_h_expr = A.solve(AssembledVector(v1) - B * AssembledVector(u_h) - C * AssembledVector(lambda_h), decomposition="LLT") return (S, r_lambda, u_h_expr, q_h_expr)
def initialize(self, pc): """Set up the problem context. Take the original H1-problem and partition the spaces/functions into 'interior' and 'facet' parts. A KSP is created for the reduced system after static condensation is applied. """ from firedrake import (FunctionSpace, Function, TrialFunction, TestFunction) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from ufl.algorithms.replace import replace # Extract python context prefix = pc.getOptionsPrefix() + "static_condensation_" _, P = pc.getOperators() self.cxt = P.getPythonContext() if not isinstance(self.cxt, ImplicitMatrixContext): raise ValueError("Context must be an ImplicitMatrixContext") test, trial = self.cxt.a.arguments() V = test.function_space() mesh = V.mesh() if len(V) > 1: raise ValueError("Cannot use this PC for mixed problems.") if V.ufl_element().sobolev_space().name != "H1": raise ValueError("Expecting an H1-conforming element.") if not V.ufl_element().cell().is_simplex(): raise NotImplementedError("Only simplex meshes are implemented.") top_dim = V.finat_element._element.ref_el.get_dimension() if not V.finat_element.entity_dofs()[top_dim][0]: raise RuntimeError("There are no interior dofs to eliminate.") # We decompose the space into an interior part and facet part interior_element = V.ufl_element()["interior"] facet_element = V.ufl_element()["facet"] V_int = FunctionSpace(mesh, interior_element) V_facet = FunctionSpace(mesh, facet_element) # Get transfer kernel for moving data self._transfer_kernel = get_transfer_kernels({ 'h1-space': V, 'interior-space': V_int, 'facet-space': V_facet }) # Set up functions for the H1 functions and the interior/trace parts self.trace_solution = Function(V_facet) self.interior_solution = Function(V_int) self.h1_solution = Function(V) self.h1_residual = Function(V) self.interior_residual = Function(V_int) self.trace_residual = Function(V_facet) # TODO: Handle strong bcs in Slate if self.cxt.row_bcs: raise NotImplementedError("Strong bcs not implemented yet") self.bcs = None A00 = Tensor( replace(self.cxt.a, { test: TestFunction(V_int), trial: TrialFunction(V_int) })) A01 = Tensor( replace(self.cxt.a, { test: TestFunction(V_int), trial: TrialFunction(V_facet) })) A10 = Tensor( replace(self.cxt.a, { test: TestFunction(V_facet), trial: TrialFunction(V_int) })) A11 = Tensor( replace(self.cxt.a, { test: TestFunction(V_facet), trial: TrialFunction(V_facet) })) # Schur complement operator S = A11 - A10 * A00.inv * A01 self.S = allocate_matrix(S, bcs=self.bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S = create_assembly_callable( S, tensor=self.S, bcs=self.bcs, form_compiler_parameters=self.cxt.fc_params) self._assemble_S() Smat = self.S.petscmat # Nullspace for the reduced system nullspace = create_sc_nullspace(P, V, V_facet, pc.comm) if nullspace: Smat.setNullSpace(nullspace) # Set up KSP for the reduced problem sc_ksp = PETSc.KSP().create(comm=pc.comm) sc_ksp.setOptionsPrefix(prefix) sc_ksp.setOperators(Smat) sc_ksp.setUp() sc_ksp.setFromOptions() self.sc_ksp = sc_ksp # Set up rhs for the reduced problem F0 = AssembledVector(self.interior_residual) self.sc_rhs = Function(V_facet) self.sc_rhs_thunk = Function(V_facet) self._assemble_sc_rhs_thunk = create_assembly_callable( -A10 * A00.inv * F0, tensor=self.sc_rhs_thunk, form_compiler_parameters=self.cxt.fc_params) # Reconstruction calls u_facet = AssembledVector(self.trace_solution) self._assemble_interior_u = create_assembly_callable( A00.inv * (F0 - A01 * u_facet), tensor=self.interior_solution, form_compiler_parameters=self.cxt.fc_params)