def step(self, snes, x, f, y): push_appctx(self.plex, self.ctx) x.copy(y) self.patch.solve(snes.vec_rhs or self.dummy, y) y.axpy(-1, x) y.scale(-1) snes.setConvergedReason(self.patch.getConvergedReason()) pop_appctx(self.plex)
def step(self, snes, x, f, y): ctx = get_appctx(snes.dm) push_appctx(self.ppc.dm, ctx) x.copy(y) self.ppc.solve(snes.vec_rhs or self.dummy, y) y.aypx(-1, x) snes.setConvergedReason(self.ppc.getConvergedReason()) pop_appctx(self.ppc.dm)
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() self.P.force_evaluation() # 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) nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx) push_appctx(dm, nctx) self._ctx_ref = nctx pc.setDM(dm) pc.setOptionsPrefix(options_prefix) pc.setOperators(Pmat, Pmat) pc.setFromOptions() pc.setUp() self.pc = pc pop_appctx(dm)
def applyTranspose(self, pc, x, y): dm = pc.getDM() push_appctx(dm, self._ctx_ref) self.pc.applyTranspose(x, y) pop_appctx(dm)
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() self.P.force_evaluation() # 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) nctx = _SNESContext(nproblem, mat_type, mat_type, octx.appctx) push_appctx(dm, nctx) pc.setDM(dm) pc.setOptionsPrefix(options_prefix) pc.setOperators(Pmat, Pmat) pc.setFromOptions() pc.setUp() self.pc = pc pop_appctx(dm)
def coarsen_snescontext(context, self, coefficient_mapping=None): if coefficient_mapping is None: coefficient_mapping = {} # Have we already done this? coarse = context._coarse if coarse is not None: return coarse problem = self(context._problem, self, coefficient_mapping=coefficient_mapping) appctx = context.appctx new_appctx = {} for k in sorted(appctx.keys()): v = appctx[k] if k != "state": # Constructor makes this one. try: new_appctx[k] = self(v, self, coefficient_mapping=coefficient_mapping) except CoarseningError: # Assume not something that needs coarsening (e.g. float) new_appctx[k] = v coarse = type(context)(problem, mat_type=context.mat_type, pmat_type=context.pmat_type, appctx=new_appctx, transfer_manager=context.transfer_manager) coarse._fine = context context._coarse = coarse # Now that we have the coarse snescontext, push it to the coarsened DMs # Otherwise they won't have the right transfer manager when they are # coarsened in turn from firedrake.dmhooks import get_appctx, push_appctx, pop_appctx from firedrake.dmhooks import add_hook, get_parent from itertools import chain for val in chain(coefficient_mapping.values(), (bc.function_arg for bc in problem.bcs)): if isinstance(val, firedrake.function.Function): V = val.function_space() coarseneddm = V.dm parentdm = get_parent(context._problem.u.function_space().dm) # Now attach the hook to the parent DM if get_appctx(coarseneddm) is None: push_appctx(coarseneddm, coarse) teardown = partial(pop_appctx, coarseneddm, coarse) add_hook(parentdm, teardown=teardown) ises = problem.J.arguments()[0].function_space()._ises coarse._nullspace = self(context._nullspace, self, coefficient_mapping=coefficient_mapping) coarse.set_nullspace(coarse._nullspace, ises, transpose=False, near=False) coarse._nullspace_T = self(context._nullspace_T, self, coefficient_mapping=coefficient_mapping) coarse.set_nullspace(coarse._nullspace_T, ises, transpose=True, near=False) coarse._near_nullspace = self(context._near_nullspace, self, coefficient_mapping=coefficient_mapping) coarse.set_nullspace(coarse._near_nullspace, ises, transpose=False, near=True) return coarse