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 test_matrix_subblocks(mesh): if mesh.ufl_cell() == quadrilateral: U = FunctionSpace(mesh, "RTCF", 1) else: U = FunctionSpace(mesh, "RT", 1) V = FunctionSpace(mesh, "DG", 0) T = FunctionSpace(mesh, "HDiv Trace", 0) n = FacetNormal(mesh) W = U * V * T u, p, lambdar = TrialFunctions(W) w, q, gammar = TestFunctions(W) A = Tensor(inner(u, w)*dx + p*q*dx - div(w)*p*dx + q*div(u)*dx + lambdar('+')*jump(w, n=n)*dS + gammar('+')*jump(u, n=n)*dS + lambdar*gammar*ds) # Test individual blocks indices = [(0, 0), (0, 1), (1, 0), (1, 1), (1, 2), (2, 1), (2, 2)] refs = dict(split_form(A.form)) _A = A.blocks for x, y in indices: ref = assemble(refs[x, y]).M.values block = _A[x, y] assert np.allclose(assemble(block).M.values, ref, rtol=1e-14) # Mixed blocks A0101 = _A[:2, :2] A1212 = _A[1:3, 1:3] _A0101 = A0101.blocks _A1212 = A1212.blocks # Block of blocks A0101_00 = _A0101[0, 0] A0101_11 = _A0101[1, 1] A0101_01 = _A0101[0, 1] A0101_10 = _A0101[1, 0] A1212_00 = _A1212[0, 0] A1212_11 = _A1212[1, 1] A1212_01 = _A1212[0, 1] A1212_10 = _A1212[1, 0] items = [(A0101_00, refs[(0, 0)]), (A0101_11, refs[(1, 1)]), (A0101_01, refs[(0, 1)]), (A0101_10, refs[(1, 0)]), (A1212_00, refs[(1, 1)]), (A1212_11, refs[(2, 2)]), (A1212_01, refs[(1, 2)]), (A1212_10, refs[(2, 1)])] # Test assembly of blocks of mixed blocks for tensor, form in items: ref = assemble(form).M.values assert np.allclose(assemble(tensor).M.values, ref, rtol=1e-14)
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 compile_form(form, name, parameters=None, inverse=False, split=True): """Compile a form using TSFC. :arg form: the :class:`~ufl.classes.Form` to compile. :arg name: a prefix for the generated kernel functions. :arg parameters: optional dict of parameters to pass to the form compiler. If not provided, parameters are read from the ``form_compiler`` slot of the Firedrake :data:`~.parameters` dictionary (which see). :arg inverse: If True then assemble the inverse of the local tensor. :arg split: If ``False``, then don't split mixed forms. Returns a tuple of tuples of (index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`). ``needs_orientations`` indicates whether the form requires cell orientation information (for correctly pulling back to reference elements on embedded manifolds). The coordinates are extracted from the domain of the integral (a :func:`~.Mesh`) """ # Check that we get a Form if not isinstance(form, Form): raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form)) if parameters is None: parameters = default_parameters["form_compiler"].copy() else: # Override defaults with user-specified values _ = parameters parameters = default_parameters["form_compiler"].copy() parameters.update(_) # We stash the compiled kernels on the form so we don't have to recompile # if we assemble the same form again with the same optimisations cache = form._cache.setdefault("firedrake_kernels", {}) def tuplify(params): return tuple((k, params[k]) for k in sorted(params)) key = (tuplify(default_parameters["coffee"]), name, tuplify(parameters), split) try: return cache[key] except KeyError: pass kernels = [] # A map from all form coefficients to their number. coefficient_numbers = dict( (c, n) for (n, c) in enumerate(form.coefficients())) if split: iterable = split_form(form) else: iterable = ([(0, ) * len(form.arguments()), form], ) for idx, f in iterable: f = _real_mangle(f) # Map local coefficient numbers (as seen inside the # compiler) to the global coefficient numbers number_map = dict((n, coefficient_numbers[c]) for (n, c) in enumerate(f.coefficients())) kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters, number_map).kernels for kinfo in kinfos: kernels.append(SplitKernel(idx, kinfo)) kernels = tuple(kernels) return cache.setdefault(key, kernels)
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 compile_form(form, name, parameters=None, inverse=False): """Compile a form using TSFC. :arg form: the :class:`~ufl.classes.Form` to compile. :arg name: a prefix for the generated kernel functions. :arg parameters: optional dict of parameters to pass to the form compiler. If not provided, parameters are read from the ``form_compiler`` slot of the Firedrake :data:`~.parameters` dictionary (which see). :arg inverse: If True then assemble the inverse of the local tensor. Returns a tuple of tuples of (index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`). ``needs_orientations`` indicates whether the form requires cell orientation information (for correctly pulling back to reference elements on embedded manifolds). The coordinates are extracted from the domain of the integral (a :func:`~.Mesh`) """ # Check that we get a Form if not isinstance(form, Form): raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form)) if parameters is None: parameters = default_parameters["form_compiler"].copy() else: # Override defaults with user-specified values _ = parameters parameters = default_parameters["form_compiler"].copy() parameters.update(_) # We stash the compiled kernels on the form so we don't have to recompile # if we assemble the same form again with the same optimisations if "firedrake_kernels" in form._cache: # Save both kernels and TSFC params so we can tell if this # cached version is valid (the TSFC parameters might have changed) kernels, coffee_params, old_name, params = form._cache[ "firedrake_kernels"] if coffee_params == default_parameters["coffee"] and \ name == old_name and \ params == parameters: return kernels kernels = [] # A map from all form coefficients to their number. coefficient_numbers = dict( (c, n) for (n, c) in enumerate(form.coefficients())) for idx, f in split_form(form): f = _real_mangle(f) # Map local coefficient numbers (as seen inside the # compiler) to the global coefficient numbers number_map = dict((n, coefficient_numbers[c]) for (n, c) in enumerate(f.coefficients())) kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters, number_map).kernels for kinfo in kinfos: kernels.append(SplitKernel(idx, kinfo)) kernels = tuple(kernels) form._cache["firedrake_kernels"] = (kernels, default_parameters["coffee"].copy(), name, parameters) return kernels
def compile_form(form, name, parameters=None, inverse=False): """Compile a form using TSFC. :arg form: the :class:`~ufl.classes.Form` to compile. :arg name: a prefix for the generated kernel functions. :arg parameters: optional dict of parameters to pass to the form compiler. If not provided, parameters are read from the ``form_compiler`` slot of the Firedrake :data:`~.parameters` dictionary (which see). :arg inverse: If True then assemble the inverse of the local tensor. Returns a tuple of tuples of (index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`). ``needs_orientations`` indicates whether the form requires cell orientation information (for correctly pulling back to reference elements on embedded manifolds). The coordinates are extracted from the domain of the integral (a :func:`~.Mesh`) """ # Check that we get a Form if not isinstance(form, Form): raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form)) if parameters is None: parameters = default_parameters["form_compiler"].copy() else: # Override defaults with user-specified values _ = parameters parameters = default_parameters["form_compiler"].copy() parameters.update(_) # We stash the compiled kernels on the form so we don't have to recompile # if we assemble the same form again with the same optimisations if "firedrake_kernels" in form._cache: # Save both kernels and TSFC params so we can tell if this # cached version is valid (the TSFC parameters might have changed) kernels, coffee_params, old_name, params = form._cache["firedrake_kernels"] if coffee_params == default_parameters["coffee"] and \ name == old_name and \ params == parameters: return kernels kernels = [] # A map from all form coefficients to their number. coefficient_numbers = dict((c, n) for (n, c) in enumerate(form.coefficients())) for idx, f in split_form(form): # Map local coefficient numbers (as seen inside the # compiler) to the global coefficient numbers number_map = dict((n, coefficient_numbers[c]) for (n, c) in enumerate(f.coefficients())) kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters, number_map).kernels for kinfo in kinfos: kernels.append(SplitKernel(idx, kinfo)) kernels = tuple(kernels) form._cache["firedrake_kernels"] = (kernels, default_parameters["coffee"].copy(), name, parameters) return kernels
def compile_form(form, name, parameters=None, inverse=False, split=True, interface=None, coffee=False): """Compile a form using TSFC. :arg form: the :class:`~ufl.classes.Form` to compile. :arg name: a prefix for the generated kernel functions. :arg parameters: optional dict of parameters to pass to the form compiler. If not provided, parameters are read from the ``form_compiler`` slot of the Firedrake :data:`~.parameters` dictionary (which see). :arg inverse: If True then assemble the inverse of the local tensor. :arg split: If ``False``, then don't split mixed forms. :arg coffee: compile coffee kernel instead of loopy kernel Returns a tuple of tuples of (index, integral type, subdomain id, coordinates, coefficients, needs_orientations, :class:`Kernels <pyop2.op2.Kernel>`). ``needs_orientations`` indicates whether the form requires cell orientation information (for correctly pulling back to reference elements on embedded manifolds). The coordinates are extracted from the domain of the integral (a :func:`~.Mesh`) """ # Check that we get a Form if not isinstance(form, Form): raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form)) if parameters is None: parameters = default_parameters["form_compiler"].copy() else: # Override defaults with user-specified values _ = parameters parameters = default_parameters["form_compiler"].copy() parameters.update(_) # We stash the compiled kernels on the form so we don't have to recompile # if we assemble the same form again with the same optimisations cache = form._cache.setdefault("firedrake_kernels", {}) def tuplify(params): return tuple((k, params[k]) for k in sorted(params)) key = (tuplify(default_parameters["coffee"]), name, tuplify(parameters), split) try: return cache[key] except KeyError: pass kernels = [] # A map from all form coefficients to their number. coefficient_numbers = dict((c, n) for (n, c) in enumerate(form.coefficients())) if split: iterable = split_form(form) else: iterable = ([(0, )*len(form.arguments()), form], ) for idx, f in iterable: f = _real_mangle(f) # Map local coefficient numbers (as seen inside the # compiler) to the global coefficient numbers number_map = dict((n, coefficient_numbers[c]) for (n, c) in enumerate(f.coefficients())) kinfos = TSFCKernel(f, name + "".join(map(str, idx)), parameters, number_map, interface, coffee).kernels for kinfo in kinfos: kernels.append(SplitKernel(idx, kinfo)) kernels = tuple(kernels) return cache.setdefault(key, kernels)
def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.residual = equation.residual if self.field_name is not None: self.idx = equation.field_names.index(self.field_name) self.fs = self.state.fields(self.field_name).function_space() self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( split_form(t.form)[self.idx].form, t.labels), drop) bcs = equation.bcs[self.field_name] else: self.field_name = equation.field_name self.fs = equation.function_space self.idx = None if type(self.fs.ufl_element()) is MixedElement: bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs] else: bcs = equation.bcs[self.field_name] if len(active_labels) > 0: self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, *active_labels)), map_if_false=drop) options = self.options # -------------------------------------------------------------------- # # Routines relating to transport # -------------------------------------------------------------------- # if hasattr(self.options, 'ibp'): self.replace_transport_term() self.replace_transporting_velocity(uadv) # -------------------------------------------------------------------- # # Wrappers for embedded / recovery methods # -------------------------------------------------------------------- # if self.discretisation_option in ["embedded_dg", "recovered"]: # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) self.fs = FunctionSpace(self.state.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) if self.idx is None: self.x_projected = Function(equation.function_space) else: self.x_projected = Function(self.state.fields(self.field_name).function_space()) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} # -------------------------------------------------------------------- # # Make boundary conditions # -------------------------------------------------------------------- # if not apply_bcs: self.bcs = None elif self.discretisation_option in ["embedded_dg", "recovered"]: # Transfer boundary conditions onto test function space self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs] else: self.bcs = bcs # -------------------------------------------------------------------- # # Modify test function for SUPG methods # -------------------------------------------------------------------- # if self.discretisation_option == "supg": # construct tau, if it is not specified dim = self.state.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" else: # create tuple of default values of size dim default_vals = [options.default*self.dt]*dim # check for directions is which the space is discontinuous # so that we don't apply supg in that direction if is_cg(self.fs): vals = default_vals else: space = self.fs.ufl_element().sobolev_space() if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] else: raise ValueError("I don't know what to do with space %s" % space) tau = Constant(tuple([ tuple( [vals[j] if i == j else 0. for i, v in enumerate(vals)] ) for j in range(dim)]) ) self.solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} test = TestFunction(self.fs) new_test = test + dot(dot(uadv, tau), grad(test)) if self.discretisation_option is not None: # replace the original test function with one defined on # the embedding space, as this is the space where the # the problem will be solved self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) if self.discretisation_option == "embedded_dg": if self.limiter is None: self.x_out_projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) else: self.x_out_projector = Recoverer(self.xdg_out, self.x_projected) if self.discretisation_option == "recovered": # set up the necessary functions self.x_in = Function(self.state.fields(self.field_name).function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) else: self.x_out_projector = Projector(self.xdg_out, self.x_projected) # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs)
def initialize(self, pc): """ Set up the problem context. Takes the original mixed problem and transforms it into the equivalent hybrid-mixed system. A KSP object is created for the Lagrange multipliers on the top/bottom faces of the mesh cells. """ from firedrake import (FunctionSpace, Function, Constant, FiniteElement, TensorProductElement, TrialFunction, TrialFunctions, TestFunction, DirichletBC, interval, MixedElement, BrokenElement) from firedrake.assemble import (allocate_matrix, create_assembly_callable) from firedrake.formmanipulation import split_form from ufl.algorithms.replace import replace from ufl.cell import TensorProductCell # Extract PC context prefix = pc.getOptionsPrefix() + "vert_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() # Magically determine which spaces are vector and scalar valued for i, Vi in enumerate(V): # Vector-valued spaces will have a non-empty value_shape if Vi.ufl_element().value_shape(): self.vidx = i else: self.pidx = i Vv = V[self.vidx] Vp = V[self.pidx] # Create the space of approximate traces in the vertical. # NOTE: Technically a hack since the resulting space is technically # defined in cell interiors, however the degrees of freedom will only # be geometrically defined on edges. Arguments will only be used in # surface integrals deg, _ = Vv.ufl_element().degree() # Assumes a tensor product cell (quads, triangular-prisms, cubes) if not isinstance(Vp.ufl_element().cell(), TensorProductCell): raise NotImplementedError( "Currently only implemented for tensor product discretizations" ) # Only want the horizontal cell cell, _ = Vp.ufl_element().cell()._cells DG = FiniteElement("DG", cell, deg) CG = FiniteElement("CG", interval, 1) Vv_tr_element = TensorProductElement(DG, CG) Vv_tr = FunctionSpace(mesh, Vv_tr_element) # Break the spaces broken_elements = MixedElement( [BrokenElement(Vi.ufl_element()) for Vi in V]) V_d = FunctionSpace(mesh, broken_elements) # Set up relevant functions self.broken_solution = Function(V_d) self.broken_residual = Function(V_d) self.trace_solution = Function(Vv_tr) self.unbroken_solution = Function(V) self.unbroken_residual = Function(V) # Set up transfer kernels to and from the broken velocity space # NOTE: Since this snippet of code is used in a couple places in # in Gusto, might be worth creating a utility function that is # is importable and just called where needed. shapes = { "i": Vv.finat_element.space_dimension(), "j": np.prod(Vv.shape, dtype=int) } weight_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) w[i*{j} + j] += 1.0; """.format(**shapes) self.weight = Function(Vv) par_loop(weight_kernel, dx, {"w": (self.weight, INC)}) # Averaging kernel self.average_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j]; """.format(**shapes) # 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(Vv_tr) n = FacetNormal(mesh) sigma = TrialFunctions(V_d)[self.vidx] # Again, assumes tensor product structure. Why use this if you # don't have some form of vertical extrusion? Kform = gammar('+') * jump(sigma, n=n) * dS_h # Here we deal with boundary conditions 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, int)) # separate out the top and bottom bcs extruded_neumann_subdomains = neumann_subdomains & { "top", "bottom" } neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains integrand = gammar * dot(sigma, n) measures = [] trace_subdomains = [] for subdomain in sorted(extruded_neumann_subdomains): measures.append({"top": ds_t, "bottom": ds_b}[subdomain]) trace_subdomains.extend( sorted({"top", "bottom"} - extruded_neumann_subdomains)) 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 else: trace_subdomains = ["top", "bottom"] trace_bcs = [ DirichletBC(Vv_tr, 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(Vv_tr) 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) self._assemble_S() self.S.force_evaluation() Smat = self.S.petscmat nullspace = self.ctx.appctx.get("vert_trace_nullspace", None) if nullspace is not None: nsp = nullspace(Vv_tr) 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 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 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, int)) # 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) self._assemble_S() self.S.force_evaluation() 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)