def sc_solve(self, pc): """Solve the condensed linear system for the condensed field. :arg pc: a Preconditioner instance. """ dm = self.condensed_ksp.getDM() with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): with self.condensed_rhs.dat.vec_ro as rhs: if self.condensed_ksp.getInitialGuessNonzero(): acc = self.solution.split()[self.c_field].dat.vec else: acc = self.solution.split()[self.c_field].dat.vec_wo with acc as sol: self.condensed_ksp.solve(rhs, sol)
def adjoint_solve(self): r"""Solve the adjoint problem. """ # Make sure appcontext is attached to the DM before the adjoint solve. dm = self.ts.getDM() with ExitStack() as stack: for ctx in chain( ( self.inserted_options(), dmhooks.add_hooks(dm, self, appctx=self._ctx), ), self._transfer_operators, ): stack.enter_context(ctx) self.ts.adjointSolve() check_ts_convergence(self.ts)
def sc_solve(self, pc): """Solve the condensed linear system for the condensed field. :arg pc: a Preconditioner instance. """ dm = self.trace_ksp.getDM() with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): # Solve the system for the Lagrange multipliers with self.schur_rhs.dat.vec_ro as b: if self.trace_ksp.getInitialGuessNonzero(): acc = self.trace_solution.dat.vec else: acc = self.trace_solution.dat.vec_wo with acc as x_trace: self.trace_ksp.solve(b, x_trace)
def solve(self, bounds=None): r"""Solve the variational problem. :arg bounds: Optional bounds on the solution (lower, upper). ``lower`` and ``upper`` must both be :class:`~.Function`\s. or :class:`~.Vector`\s. .. note:: If bounds are provided the ``snes_type`` must be set to ``vinewtonssls`` or ``vinewtonrsls``. """ # Make sure the DM has this solver's callback functions self._ctx.set_function(self.snes) self._ctx.set_jacobian(self.snes) # Make sure appcontext is attached to the DM before we solve. dm = self.snes.getDM() for dbc in self._problem.dirichlet_bcs(): dbc.apply(self._problem.u) if bounds is not None: lower, upper = bounds with lower.dat.vec_ro as lb, upper.dat.vec_ro as ub: self.snes.setVariableBounds(lb, ub) work = self._work with self._problem.u.dat.vec as u: u.copy(work) with ExitStack() as stack: # Ensure options database has full set of options (so monitors # work right) for ctx in chain( (self.inserted_options(), dmhooks.add_hooks(dm, self, appctx=self._ctx)), self._transfer_operators): stack.enter_context(ctx) self.snes.solve(None, work) work.copy(u) self._setup = True solving_utils.check_snes_convergence(self.snes)
def solve(self, x, b): if not isinstance(x, (function.Function, vector.Vector)): raise TypeError( "Provided solution is a '%s', not a Function or Vector" % type(x).__name__) if isinstance(b, vector.Vector): b = b.function if not isinstance(b, function.Function): raise TypeError("Provided RHS is a '%s', not a Function" % type(b).__name__) if len(self.trial_space) > 1 and self.nullspace is not None: self.nullspace._apply(self.trial_space.dof_dset.field_ises) if len(self.test_space) > 1 and self.transpose_nullspace is not None: self.transpose_nullspace._apply( self.test_space.dof_dset.field_ises, transpose=True) if len(self.trial_space) > 1 and self.near_nullspace is not None: self.near_nullspace._apply(self.trial_space.dof_dset.field_ises, near=True) if self.A.has_bcs: b = self._lifted(b) if self.ksp.getInitialGuessNonzero(): acc = x.dat.vec else: acc = x.dat.vec_wo with self.inserted_options( ), b.dat.vec_ro as rhs, acc as solution, dmhooks.add_hooks( self.ksp.dm, self): self.ksp.solve(rhs, solution) r = self.ksp.getConvergedReason() if r < 0: raise ConvergenceError( "LinearSolver failed to converge after %d iterations with reason: %s", self.ksp.getIterationNumber(), solving_utils.KSPReasons[r])
def _la_solve(A, x, b, **kwargs): r"""Solve a linear algebra problem. :arg A: the assembled bilinear form, a :class:`.Matrix`. :arg x: the :class:`.Function` to write the solution into. :arg b: the :class:`.Function` defining the right hand side values. :kwarg solver_parameters: optional solver parameters. :kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or :class:`.MixedVectorSpaceBasis`) spanning the null space of the operator. :kwarg transpose_nullspace: as for the nullspace, but used to make the right hand side consistent. :kwarg near_nullspace: as for the nullspace, but used to add the near nullspace. :kwarg options_prefix: an optional prefix used to distinguish PETSc options. If not provided a unique prefix will be created. Use this option if you want to pass options to the solver from the command line in addition to through the ``solver_parameters`` dict. .. note:: This function no longer accepts :class:`.DirichletBC`\s or :class:`.EquationBC`\s as arguments. Any boundary conditions must be applied when assembling the bilinear form as: .. code-block:: python A = assemble(a, bcs=[bc1]) solve(A, x, b) Example usage: .. code-block:: python _la_solve(A, x, b, solver_parameters=parameters_dict).""" bcs, solver_parameters, nullspace, nullspace_T, near_nullspace, \ options_prefix = _extract_linear_solver_args(A, x, b, **kwargs) if bcs is not None: raise RuntimeError( "It is no longer possible to apply or change boundary conditions after assembling the matrix `A`; pass any necessary boundary conditions to `assemble` when assembling `A`." ) solver = ls.LinearSolver(A, solver_parameters=solver_parameters, nullspace=nullspace, transpose_nullspace=nullspace_T, near_nullspace=near_nullspace, options_prefix=options_prefix) if isinstance(x, firedrake.Vector): x = x.function # linear MG doesn't need RHS, supply zero. lvp = vs.LinearVariationalProblem(a=A.a, L=0, u=x, bcs=A.bcs) mat_type = A.mat_type appctx = solver_parameters.get("appctx", {}) ctx = solving_utils._SNESContext(lvp, mat_type=mat_type, pmat_type=mat_type, appctx=appctx, options_prefix=options_prefix) dm = solver.ksp.dm with dmhooks.add_hooks(dm, solver, appctx=ctx): solver.solve(x, b)
def __init__(self, problem, **kwargs): r""" :arg problem: A :class:`NonlinearVariationalProblem` to solve. :kwarg nullspace: an optional :class:`.VectorSpaceBasis` (or :class:`.MixedVectorSpaceBasis`) spanning the null space of the operator. :kwarg transpose_nullspace: as for the nullspace, but used to make the right hand side consistent. :kwarg near_nullspace: as for the nullspace, but used to specify the near nullspace (for multigrid solvers). :kwarg solver_parameters: Solver parameters to pass to PETSc. This should be a dict mapping PETSc options to values. :kwarg appctx: A dictionary containing application context that is passed to the preconditioner if matrix-free. :kwarg options_prefix: an optional prefix used to distinguish PETSc options. If not provided a unique prefix will be created. Use this option if you want to pass options to the solver from the command line in addition to through the ``solver_parameters`` dict. :kwarg pre_jacobian_callback: A user-defined function that will be called immediately before Jacobian assembly. This can be used, for example, to update a coefficient function that has a complicated dependence on the unknown solution. :kwarg pre_function_callback: As above, but called immediately before residual assembly Example usage of the ``solver_parameters`` option: to set the nonlinear solver type to just use a linear solver, use .. code-block:: python {'snes_type': 'ksponly'} PETSc flag options (where the presence of the option means something) should be specified with ``None``. For example: .. code-block:: python {'snes_monitor': None} To use the ``pre_jacobian_callback`` or ``pre_function_callback`` functionality, the user-defined function must accept the current solution as a petsc4py Vec. Example usage is given below: .. code-block:: python def update_diffusivity(current_solution): with cursol.dat.vec_wo as v: current_solution.copy(v) solve(trial*test*dx == dot(grad(cursol), grad(test))*dx, diffusivity) solver = NonlinearVariationalSolver(problem, pre_jacobian_callback=update_diffusivity) """ assert isinstance(problem, NonlinearVariationalProblem) parameters = kwargs.get("solver_parameters") if "parameters" in kwargs: raise TypeError("Use solver_parameters, not parameters") nullspace = kwargs.get("nullspace") nullspace_T = kwargs.get("transpose_nullspace") near_nullspace = kwargs.get("near_nullspace") options_prefix = kwargs.get("options_prefix") pre_j_callback = kwargs.get("pre_jacobian_callback") pre_f_callback = kwargs.get("pre_function_callback") super(NonlinearVariationalSolver, self).__init__(parameters, options_prefix) # Allow anything, interpret "matfree" as matrix_free. mat_type = self.parameters.get("mat_type") pmat_type = self.parameters.get("pmat_type") matfree = mat_type == "matfree" pmatfree = pmat_type == "matfree" appctx = kwargs.get("appctx") ctx = solving_utils._SNESContext(problem, mat_type=mat_type, pmat_type=pmat_type, appctx=appctx, pre_jacobian_callback=pre_j_callback, pre_function_callback=pre_f_callback, options_prefix=self.options_prefix) # No preconditioner by default for matrix-free if (problem.Jp is not None and pmatfree) or matfree: self.set_default_parameter("pc_type", "none") elif ctx.is_mixed: # Mixed problem, use jacobi pc if user has not supplied # one. self.set_default_parameter("pc_type", "jacobi") self.snes = PETSc.SNES().create(comm=problem.dm.comm) self._problem = problem self._ctx = ctx self._work = problem.u.dof_dset.layout_vec.duplicate() self.snes.setDM(problem.dm) ctx.set_function(self.snes) ctx.set_jacobian(self.snes) ctx.set_nullspace(nullspace, problem.J.arguments()[0].function_space()._ises, transpose=False, near=False) ctx.set_nullspace(nullspace_T, problem.J.arguments()[1].function_space()._ises, transpose=True, near=False) ctx.set_nullspace(near_nullspace, problem.J.arguments()[0].function_space()._ises, transpose=False, near=True) ctx._nullspace = nullspace ctx._nullspace_T = nullspace_T ctx._near_nullspace = near_nullspace # Set from options now, so that people who want to noodle with # the snes object directly (mostly Patrick), can. We need the # DM with an app context in place so that if the DM is active # on a subKSP the context is available. dm = self.snes.getDM() with dmhooks.add_hooks(dm, self, appctx=self._ctx, save=False): self.set_from_options(self.snes) # Used for custom grid transfer. self._transfer_operators = () self._setup = False
def applyTranspose(self, pc, X, Y): dm = self._dm with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): self.pc.applyTranspose(X, Y)
def initialize(self, pc): from firedrake import TestFunction, parameters from firedrake.assemble import allocate_matrix, create_assembly_callable from firedrake.interpolation import Interpolator from firedrake.solving_utils import _SNESContext from firedrake.matrix_free.operators import ImplicitMatrixContext _, P = pc.getOperators() appctx = self.get_appctx(pc) fcp = appctx.get("form_compiler_parameters") if pc.getType() != "python": raise ValueError("Expecting PC type python") ctx = dmhooks.get_appctx(pc.getDM()) if ctx is None: raise ValueError("No context found.") if not isinstance(ctx, _SNESContext): raise ValueError("Don't know how to get form from %r", ctx) prefix = pc.getOptionsPrefix() options_prefix = prefix + self._prefix opts = PETSc.Options() # Handle the fine operator if type is python if P.getType() == "python": ictx = P.getPythonContext() if ictx is None: raise ValueError("No context found on matrix") if not isinstance(ictx, ImplicitMatrixContext): raise ValueError("Don't know how to get form from %r", ictx) fine_operator = ictx.a fine_bcs = ictx.row_bcs if fine_bcs != ictx.col_bcs: raise NotImplementedError("Row and column bcs must match") fine_mat_type = opts.getString(options_prefix + "mat_type", parameters["default_matrix_type"]) self.fine_op = allocate_matrix(fine_operator, bcs=fine_bcs, form_compiler_parameters=fcp, mat_type=fine_mat_type, options_prefix=options_prefix) self._assemble_fine_op = create_assembly_callable( fine_operator, tensor=self.fine_op, bcs=fine_bcs, form_compiler_parameters=fcp, mat_type=fine_mat_type) self._assemble_fine_op() fine_petscmat = self.fine_op.petscmat else: fine_petscmat = P # Transfer fine operator null space fine_petscmat.setNullSpace(P.getNullSpace()) fine_transpose_nullspace = P.getTransposeNullSpace() if fine_transpose_nullspace.handle != 0: fine_petscmat.setTransposeNullSpace(fine_transpose_nullspace) # Handle the coarse operator coarse_options_prefix = options_prefix + "mg_coarse" coarse_mat_type = opts.getString(coarse_options_prefix + "mat_type", parameters["default_matrix_type"]) get_coarse_space = appctx.get("get_coarse_space", None) if not get_coarse_space: raise ValueError( "Need to provide a callback which provides the coarse space.") coarse_space = get_coarse_space() get_coarse_operator = appctx.get("get_coarse_operator", None) if not get_coarse_operator: raise ValueError( "Need to provide a callback which provides the coarse operator." ) coarse_operator = get_coarse_operator() coarse_space_bcs = appctx.get("coarse_space_bcs", None) # These should be callbacks which return the relevant nullspaces get_coarse_nullspace = appctx.get("get_coarse_op_nullspace", None) get_coarse_transpose_nullspace = appctx.get( "get_coarse_op_transpose_nullspace", None) self.coarse_op = allocate_matrix(coarse_operator, bcs=coarse_space_bcs, form_compiler_parameters=fcp, mat_type=coarse_mat_type, options_prefix=coarse_options_prefix) self._assemble_coarse_op = create_assembly_callable( coarse_operator, tensor=self.coarse_op, bcs=coarse_space_bcs, form_compiler_parameters=fcp) self._assemble_coarse_op() coarse_opmat = self.coarse_op.petscmat # Set nullspace if provided if get_coarse_nullspace: nsp = get_coarse_nullspace() coarse_opmat.setNullSpace(nsp.nullspace(comm=pc.comm)) if get_coarse_transpose_nullspace: tnsp = get_coarse_transpose_nullspace() coarse_opmat.setTransposeNullSpace(tnsp.nullspace(comm=pc.comm)) interp_petscmat = appctx.get("interpolation_matrix", None) if interp_petscmat is None: # Create interpolation matrix from coarse space to fine space fine_space = ctx.J.arguments()[0].function_space() interpolator = Interpolator(TestFunction(coarse_space), fine_space) interpolation_matrix = interpolator.callable() interp_petscmat = interpolation_matrix.handle # We set up a PCMG object that uses the constructed interpolation # matrix to generate the restriction/prolongation operators. # This is a two-level multigrid preconditioner. pcmg = PETSc.PC().create(comm=pc.comm) pcmg.incrementTabLevel(1, parent=pc) pcmg.setType(pc.Type.MG) pcmg.setOptionsPrefix(options_prefix) pcmg.setMGLevels(2) pcmg.setMGCycleType(pc.MGCycleType.V) pcmg.setMGInterpolation(1, interp_petscmat) pcmg.setOperators(A=fine_petscmat, P=fine_petscmat) # Create new appctx self._ctx_ref = self.new_snes_ctx(pc, coarse_operator, coarse_space_bcs, coarse_mat_type, fcp) coarse_solver = pcmg.getMGCoarseSolve() coarse_solver.setOperators(A=coarse_opmat, P=coarse_opmat) # coarse space dm coarse_dm = coarse_space.dm coarse_solver.setDM(coarse_dm) coarse_solver.setDMActive(False) pcmg.setDM(coarse_dm) pcmg.setFromOptions() self.pc = pcmg self._dm = coarse_dm with dmhooks.add_hooks(coarse_dm, self, appctx=self._ctx_ref, save=False): coarse_solver.setFromOptions()
def initialize(self, pc): from firedrake.assemble import allocate_matrix, create_assembly_callable _, P = pc.getOperators() if pc.getType() != "python": raise ValueError("Expecting PC type python") opc = pc appctx = self.get_appctx(pc) fcp = appctx.get("form_compiler_parameters") V = get_function_space(pc.getDM()) if len(V) == 1: V = FunctionSpace(V.mesh(), V.ufl_element()) else: V = MixedFunctionSpace([V_ for V_ in V]) test = TestFunction(V) trial = TrialFunction(V) if P.type == "python": context = P.getPythonContext() # It only makes sense to preconditioner/invert a diagonal # block in general. That's all we're going to allow. if not context.on_diag: raise ValueError("Only makes sense to invert diagonal block") prefix = pc.getOptionsPrefix() options_prefix = prefix + self._prefix mat_type = PETSc.Options().getString(options_prefix + "mat_type", "aij") (a, bcs) = self.form(pc, test, trial) self.P = allocate_matrix(a, bcs=bcs, form_compiler_parameters=fcp, mat_type=mat_type, options_prefix=options_prefix) self._assemble_P = create_assembly_callable( a, tensor=self.P, bcs=bcs, form_compiler_parameters=fcp, mat_type=mat_type) self._assemble_P() # Transfer nullspace over Pmat = self.P.petscmat Pmat.setNullSpace(P.getNullSpace()) tnullsp = P.getTransposeNullSpace() if tnullsp.handle != 0: Pmat.setTransposeNullSpace(tnullsp) # Internally, we just set up a PC object that the user can configure # however from the PETSc command line. Since PC allows the user to specify # a KSP, we can do iterative by -assembled_pc_type ksp. pc = PETSc.PC().create(comm=opc.comm) pc.incrementTabLevel(1, parent=opc) # We set a DM and an appropriate SNESContext on the constructed PC so one # can do e.g. multigrid or patch solves. from firedrake.variational_solver import NonlinearVariationalProblem from firedrake.solving_utils import _SNESContext dm = opc.getDM() octx = get_appctx(dm) oproblem = octx._problem nproblem = NonlinearVariationalProblem(oproblem.F, oproblem.u, bcs, J=a, form_compiler_parameters=fcp) self._ctx_ref = _SNESContext(nproblem, mat_type, mat_type, octx.appctx, options_prefix=options_prefix) pc.setDM(dm) pc.setOptionsPrefix(options_prefix) pc.setOperators(Pmat, Pmat) self.pc = pc with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref, save=False): pc.setFromOptions()
def applyTranspose(self, pc, x, y): dm = pc.getDM() with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref): self.pc.applyTranspose(x, y)
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) 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, appctx=self.get_appctx(pc)) 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)) # Create a SNESContext for the DM associated with the trace problem self._ctx_ref = self.new_snes_ctx(pc, schur_comp, trace_bcs, mat_type, self.ctx.fc_params) # dm associated with the trace problem trace_dm = TraceSpace.dm # KSP for the system of Lagrange multipliers trace_ksp = PETSc.KSP().create(comm=pc.comm) trace_ksp.incrementTabLevel(1, parent=pc) # Set the dm for the trace solver trace_ksp.setDM(trace_dm) trace_ksp.setDMActive(False) trace_ksp.setOptionsPrefix(prefix) trace_ksp.setOperators(Smat, Smat) self.trace_ksp = trace_ksp with dmhooks.add_hooks(trace_dm, self, appctx=self._ctx_ref, save=False): trace_ksp.setFromOptions() 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)