def initialize(self, pc): # Make a new DM. # Hook up a (new) coarsen routine on that DM. # Make a new PC, of type MG. # Assign the DM to that PC. odm = pc.getDM() ctx = get_appctx(odm) test, trial = ctx.J.arguments() if test.function_space() != trial.function_space(): raise NotImplementedError("test and trial spaces must be the same") prefix = pc.getOptionsPrefix() options_prefix = prefix + "pmg_" pdm = PETSc.DMShell().create(comm=pc.comm) pdm.setOptionsPrefix(options_prefix) # Get the coarse degree from PETSc options self.coarse_degree = PETSc.Options(options_prefix).getInt("mg_coarse_degree", default=1) # Construct a list with the elements we'll be using V = test.function_space() ele = V.ufl_element() elements = [ele] while True: try: ele_ = self.coarsen_element(ele) assert ele_.value_shape() == ele.value_shape() ele = ele_ except ValueError: break elements.append(ele) sf = odm.getPointSF() section = odm.getDefaultSection() attach_hooks(pdm, level=len(elements)-1, sf=sf, section=section) # Now overwrite some routines on the DM pdm.setRefine(None) pdm.setCoarsen(self.coarsen) pdm.setCreateInterpolation(self.create_interpolation) # We need this for p-FAS pdm.setCreateInjection(self.create_injection) pdm.setSNESJacobian(_SNESContext.form_jacobian) pdm.setSNESFunction(_SNESContext.form_function) pdm.setKSPComputeOperators(_SNESContext.compute_operators) set_function_space(pdm, get_function_space(odm)) parent = get_parent(odm) assert parent is not None add_hook(parent, setup=partial(push_parent, pdm, parent), teardown=partial(pop_parent, pdm, parent), call_setup=True) add_hook(parent, setup=partial(push_appctx, pdm, ctx), teardown=partial(pop_appctx, pdm, ctx), call_setup=True) self.ppc = self.configure_pmg(pc, pdm) self.ppc.setFromOptions() self.ppc.setUp()
def coarsen_function_space(V, self, coefficient_mapping=None): if hasattr(V, "_coarse"): return V._coarse from firedrake.dmhooks import (get_transfer_operators, get_parent, push_transfer_operators, pop_transfer_operators, push_parent, pop_parent, add_hook) fine = V indices = [] while True: if V.index is not None: indices.append(V.index) if V.component is not None: indices.append(V.component) if V.parent is not None: V = V.parent else: break mesh = self(V.mesh(), self) V = firedrake.FunctionSpace(mesh, V.ufl_element()) for i in reversed(indices): V = V.sub(i) V._fine = fine fine._coarse = V # FIXME: This replicates some code from dmhooks.coarsen, but we # can't do things there because that code calls this code. # We need to move these operators over here because if we have # fieldsplits + MG with auxiliary coefficients on spaces other # than which we do the MG, dm.coarsen is never called, so the # hooks are not attached. Instead we just call (say) inject which # coarsens the functionspace. cdm = V.dm transfer = get_transfer_operators(fine.dm) parent = get_parent(fine.dm) try: add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent), call_setup=True) add_hook(parent, setup=partial(push_transfer_operators, cdm, transfer), teardown=partial(pop_transfer_operators, cdm, transfer), call_setup=True) except ValueError: # Not in an add_hooks context pass return V
def initialize(self, pc): # Make a new DM. # Hook up a (new) coarsen routine on that DM. # Make a new PC, of type MG. # Assign the DM to that PC. odm = pc.getDM() ctx = get_appctx(odm) test, trial = ctx.J.arguments() if test.function_space() != trial.function_space(): raise NotImplementedError("test and trial spaces must be the same") # Construct a list with the elements we'll be using V = test.function_space() ele = V.ufl_element() elements = [ele] while True: try: ele_ = self.coarsen_element(ele) assert ele_.value_shape() == ele.value_shape() ele = ele_ except ValueError: break elements.append(ele) pdm = PETSc.DMShell().create(comm=pc.comm) sf = odm.getPointSF() section = odm.getDefaultSection() attach_hooks(pdm, level=len(elements) - 1, sf=sf, section=section) # Now overwrite some routines on the DM pdm.setRefine(None) pdm.setCoarsen(self.coarsen) pdm.setCreateInterpolation(self.create_interpolation) pdm.setOptionsPrefix(pc.getOptionsPrefix() + "pmg_") set_function_space(pdm, get_function_space(odm)) parent = get_parent(odm) assert parent is not None add_hook(parent, setup=partial(push_parent, pdm, parent), teardown=partial(pop_parent, pdm, parent), call_setup=True) add_hook(parent, setup=partial(push_appctx, pdm, ctx), teardown=partial(pop_appctx, pdm, ctx), call_setup=True) ppc = PETSc.PC().create(comm=pc.comm) ppc.setOptionsPrefix(pc.getOptionsPrefix() + "pmg_") ppc.setType("mg") ppc.setOperators(*pc.getOperators()) ppc.setDM(pdm) ppc.incrementTabLevel(1, parent=pc) # PETSc unfortunately requires us to make an ugly hack. # We would like to use GMG for the coarse solve, at least # sometimes. But PETSc will use this p-DM's getRefineLevels() # instead of the getRefineLevels() of the MeshHierarchy to # decide how many levels it should use for PCMG applied to # the p-MG's coarse problem. So we need to set an option # for the user, if they haven't already; I don't know any # other way to get PETSc to know this at the right time. opts = PETSc.Options(pc.getOptionsPrefix() + "pmg_") if "mg_coarse_pc_mg_levels" not in opts: opts["mg_coarse_pc_mg_levels"] = odm.getRefineLevel() + 1 ppc.setFromOptions() ppc.setUp() self.ppc = ppc
def coarsen(self, fdm, comm): fctx = get_appctx(fdm) test, trial = fctx.J.arguments() fV = test.function_space() fu = fctx._problem.u cele = self.coarsen_element(fV.ufl_element()) cV = firedrake.FunctionSpace(fV.mesh(), cele) cdm = cV.dm cu = firedrake.Function(cV) interpolators = tuple( firedrake.Interpolator(fus, cus) for fus, cus in zip(fu.split(), cu.split())) def inject_state(interpolators): for interpolator in interpolators: interpolator.interpolate() parent = get_parent(fdm) assert parent is not None add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent), call_setup=True) replace_d = { fu: cu, test: firedrake.TestFunction(cV), trial: firedrake.TrialFunction(cV) } cJ = replace(fctx.J, replace_d) cF = replace(fctx.F, replace_d) if fctx.Jp is not None: cJp = replace(fctx.Jp, replace_d) else: cJp = None cbcs = [] for bc in fctx._problem.bcs: # Don't actually need the value, since it's only used for # killing parts of the matrix. This should be generalised # for p-FAS, if anyone ever wants to do that cV_ = cV for index in bc._indices: cV_ = cV_.sub(index) cbcs.append( firedrake.DirichletBC(cV_, firedrake.zero(cV_.shape), bc.sub_domain, method=bc.method)) fcp = fctx._problem.form_compiler_parameters cproblem = firedrake.NonlinearVariationalProblem( cF, cu, cbcs, cJ, Jp=cJp, form_compiler_parameters=fcp, is_linear=fctx._problem.is_linear) cctx = _SNESContext( cproblem, fctx.mat_type, fctx.pmat_type, appctx=fctx.appctx, pre_jacobian_callback=fctx._pre_jacobian_callback, pre_function_callback=fctx._pre_function_callback, post_jacobian_callback=fctx._post_jacobian_callback, post_function_callback=fctx._post_function_callback, options_prefix=fctx.options_prefix, transfer_manager=fctx.transfer_manager) add_hook(parent, setup=partial(push_appctx, cdm, cctx), teardown=partial(pop_appctx, cdm, cctx), call_setup=True) add_hook(parent, setup=partial(inject_state, interpolators), call_setup=True) cdm.setKSPComputeOperators(_SNESContext.compute_operators) cdm.setCreateInterpolation(self.create_interpolation) cdm.setOptionsPrefix(fdm.getOptionsPrefix()) # If we're the coarsest grid of the p-hierarchy, don't # overwrite the coarsen routine; this is so that you can # use geometric multigrid for the p-coarse problem try: self.coarsen_element(cele) cdm.setCoarsen(self.coarsen) except ValueError: pass return cdm
def coarsen(self, fdm, comm): # Coarsen the _SNESContext of a DM fdm # return the coarse DM cdm of the coarse _SNESContext fctx = get_appctx(fdm) parent = get_parent(fdm) assert parent is not None test, trial = fctx.J.arguments() fV = test.function_space() cele = self.coarsen_element(fV.ufl_element()) # Have we already done this? cctx = fctx._coarse if cctx is not None: cV = cctx.J.arguments()[0].function_space() if (cV.ufl_element() == cele) and (cV.mesh() == fV.mesh()): return cV.dm cV = firedrake.FunctionSpace(fV.mesh(), cele) cdm = cV.dm fproblem = fctx._problem fu = fproblem.u cu = firedrake.Function(cV) def coarsen_quadrature(df, Nf, Nc): # Coarsen the quadrature degree in a dictionary # such that the ratio of quadrature nodes to interpolation nodes (Nq+1)/(Nf+1) is preserved if isinstance(df, dict): Nq = df.get("quadrature_degree", None) if Nq is not None: dc = dict(df) dc["quadrature_degree"] = max( 2 * Nc + 1, ((Nq + 1) * (Nc + 1) + Nf) // (Nf + 1) - 1) return dc return df def coarsen_form(form, Nf, Nc, replace_d): # Coarsen a form, by replacing the solution, test and trial functions, and # reconstructing each integral with a coarsened quadrature degree. # If form is not a Form, then return form. return Form([ f.reconstruct( metadata=coarsen_quadrature(f.metadata(), Nf, Nc)) for f in replace(form, replace_d).integrals() ]) if isinstance(form, Form) else form def coarsen_bcs(fbcs): cbcs = [] for bc in fbcs: cV_ = cV for index in bc._indices: cV_ = cV_.sub(index) cbc_value = self.coarsen_bc_value(bc, cV_) if type(bc) == firedrake.DirichletBC: cbcs.append( firedrake.DirichletBC(cV_, cbc_value, bc.sub_domain)) else: raise NotImplementedError( "Unsupported BC type, please get in touch if you need this" ) return cbcs Nf = PMGBase.max_degree(fV.ufl_element()) Nc = PMGBase.max_degree(cV.ufl_element()) # Replace dictionary with coarse state, test and trial functions replace_d = { fu: cu, test: firedrake.TestFunction(cV), trial: firedrake.TrialFunction(cV) } cF = coarsen_form(fctx.F, Nf, Nc, replace_d) cJ = coarsen_form(fctx.J, Nf, Nc, replace_d) cJp = coarsen_form(fctx.Jp, Nf, Nc, replace_d) fcp = coarsen_quadrature(fproblem.form_compiler_parameters, Nf, Nc) cbcs = coarsen_bcs(fproblem.bcs) # Coarsen the appctx: the user might want to provide solution-dependant expressions and forms cappctx = dict(fctx.appctx) for key in cappctx: val = cappctx[key] if isinstance(val, dict): cappctx[key] = coarsen_quadrature(val, Nf, Nc) elif isinstance(val, Expr): cappctx[key] = replace(val, replace_d) elif isinstance(val, Form): cappctx[key] = coarsen_form(val, Nf, Nc, replace_d) cmat_type = fctx.mat_type cpmat_type = fctx.pmat_type if Nc == self.coarse_degree: cmat_type = self.coarse_mat_type cpmat_type = self.coarse_mat_type if fcp is None: fcp = dict() fcp["mode"] = self.coarse_form_compiler_mode # Coarsen the problem and the _SNESContext cproblem = firedrake.NonlinearVariationalProblem( cF, cu, bcs=cbcs, J=cJ, Jp=cJp, form_compiler_parameters=fcp, is_linear=fproblem.is_linear) cctx = type(fctx)(cproblem, cmat_type, cpmat_type, appctx=cappctx, pre_jacobian_callback=fctx._pre_jacobian_callback, pre_function_callback=fctx._pre_function_callback, post_jacobian_callback=fctx._post_jacobian_callback, post_function_callback=fctx._post_function_callback, options_prefix=fctx.options_prefix, transfer_manager=fctx.transfer_manager) # FIXME setting up the _fine attribute triggers gmg injection. # cctx._fine = fctx fctx._coarse = cctx add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent), call_setup=True) add_hook(parent, setup=partial(push_appctx, cdm, cctx), teardown=partial(pop_appctx, cdm, cctx), call_setup=True) cdm.setOptionsPrefix(fdm.getOptionsPrefix()) cdm.setKSPComputeOperators(_SNESContext.compute_operators) cdm.setCreateInterpolation(self.create_interpolation) cdm.setCreateInjection(self.create_injection) # If we're the coarsest grid of the p-hierarchy, don't # overwrite the coarsen routine; this is so that you can # use geometric multigrid for the p-coarse problem try: self.coarsen_element(cele) cdm.setCoarsen(self.coarsen) except ValueError: pass # injection of the initial state def inject_state(mat): with cu.dat.vec_wo as xc, fu.dat.vec_ro as xf: mat.multTranspose(xf, xc) injection = self.create_injection(cdm, fdm) add_hook(parent, setup=partial(inject_state, injection), call_setup=True) # restrict the nullspace basis def coarsen_nullspace(coarse_V, mat, fine_nullspace): if isinstance(fine_nullspace, MixedVectorSpaceBasis): if mat.type == 'python': mat = mat.getPythonContext() submats = [ mat.getNestSubMatrix(i, i) for i in range(len(coarse_V)) ] coarse_bases = [] for fs, submat, basis in zip(coarse_V, submats, fine_nullspace._bases): if isinstance(basis, VectorSpaceBasis): coarse_bases.append( coarsen_nullspace(fs, submat, basis)) else: coarse_bases.append(coarse_V.sub(basis.index)) return MixedVectorSpaceBasis(coarse_V, coarse_bases) elif isinstance(fine_nullspace, VectorSpaceBasis): coarse_vecs = [] for xf in fine_nullspace._petsc_vecs: wc = firedrake.Function(coarse_V) with wc.dat.vec_wo as xc: mat.multTranspose(xf, xc) coarse_vecs.append(wc) vsb = VectorSpaceBasis(coarse_vecs, constant=fine_nullspace._constant) vsb.orthonormalize() return vsb else: return fine_nullspace I, _ = self.create_interpolation(cdm, fdm) ises = cV._ises cctx._nullspace = coarsen_nullspace(cV, I, fctx._nullspace) cctx.set_nullspace(cctx._nullspace, ises, transpose=False, near=False) cctx._nullspace_T = coarsen_nullspace(cV, I, fctx._nullspace_T) cctx.set_nullspace(cctx._nullspace_T, ises, transpose=True, near=False) cctx._near_nullspace = coarsen_nullspace(cV, I, fctx._near_nullspace) cctx.set_nullspace(cctx._near_nullspace, ises, transpose=False, near=True) return cdm
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
def coarsen(self, fdm, comm): # Coarsen the _SNESContext of a DM fdm # return the coarse DM cdm of the coarse _SNESContext fctx = get_appctx(fdm) # Have we already done this? cctx = fctx._coarse if cctx is not None: return cctx.J.arguments()[0].function_space().dm parent = get_parent(fdm) assert parent is not None test, trial = fctx.J.arguments() fV = test.function_space() cele = self.coarsen_element(fV.ufl_element()) cV = firedrake.FunctionSpace(fV.mesh(), cele) cdm = cV.dm def get_max_degree(ele): if isinstance(ele, MixedElement): return max(get_max_degree(sub) for sub in ele.sub_elements()) else: N = ele.degree() try: return max(N) except TypeError: return N def coarsen_quadrature(df, Nf, Nc): # Coarsen the quadrature degree in a dictionary # such that the ratio of quadrature nodes to interpolation nodes (Nq+1)/(Nf+1) is preserved if isinstance(df, dict): Nq = df.get("quadrature_degree", None) if Nq is not None: dc = dict(df) dc["quadrature_degree"] = max( 2 * Nc + 1, ((Nq + 1) * (Nc + 1) + Nf) // (Nf + 1) - 1) return dc return df def coarsen_form(form, Nf, Nc, replace_d): # Coarsen a form, by replacing the solution, test and trial functions, and # reconstructing each integral with a corsened quadrature degree. # If form is not a Form, then return form. return Form([ f.reconstruct( metadata=coarsen_quadrature(f.metadata(), Nf, Nc)) for f in replace(form, replace_d).integrals() ]) if isinstance(form, Form) else form def coarsen_bcs(fbcs): cbcs = [] for bc in fbcs: cV_ = cV for index in bc._indices: cV_ = cV_.sub(index) cbc_value = self.coarsen_bc_value(bc, cV_) if type(bc) == firedrake.DirichletBC: cbcs.append( firedrake.DirichletBC(cV_, cbc_value, bc.sub_domain, method=bc.method)) else: raise NotImplementedError( "Unsupported BC type, please get in touch if you need this" ) return cbcs Nf = get_max_degree(fV.ufl_element()) Nc = get_max_degree(cV.ufl_element()) fproblem = fctx._problem fu = fproblem.u cu = firedrake.Function(cV) # Replace dictionary with coarse state, test and trial functions replace_d = { fu: cu, test: firedrake.TestFunction(cV), trial: firedrake.TrialFunction(cV) } cF = coarsen_form(fctx.F, Nf, Nc, replace_d) cJ = coarsen_form(fctx.J, Nf, Nc, replace_d) cJp = coarsen_form(fctx.Jp, Nf, Nc, replace_d) fcp = coarsen_quadrature(fproblem.form_compiler_parameters, Nf, Nc) cbcs = coarsen_bcs(fproblem.bcs) # Coarsen the appctx: the user might want to provide solution-dependant expressions and forms cappctx = dict(fctx.appctx) for key in cappctx: val = cappctx[key] if isinstance(val, dict): cappctx[key] = coarsen_quadrature(val, Nf, Nc) elif isinstance(val, Expr): cappctx[key] = replace(val, replace_d) elif isinstance(val, Form): cappctx[key] = coarsen_form(val, Nf, Nc, replace_d) # Coarsen the problem and the _SNESContext cproblem = firedrake.NonlinearVariationalProblem( cF, cu, bcs=cbcs, J=cJ, Jp=cJp, form_compiler_parameters=fcp, is_linear=fproblem.is_linear) cctx = type(fctx)(cproblem, fctx.mat_type, fctx.pmat_type, appctx=cappctx, pre_jacobian_callback=fctx._pre_jacobian_callback, pre_function_callback=fctx._pre_function_callback, post_jacobian_callback=fctx._post_jacobian_callback, post_function_callback=fctx._post_function_callback, options_prefix=fctx.options_prefix, transfer_manager=fctx.transfer_manager) # FIXME setting up the _fine attribute triggers gmg injection. # cctx._fine = fctx fctx._coarse = cctx add_hook(parent, setup=partial(push_parent, cdm, parent), teardown=partial(pop_parent, cdm, parent), call_setup=True) add_hook(parent, setup=partial(push_appctx, cdm, cctx), teardown=partial(pop_appctx, cdm, cctx), call_setup=True) cdm.setOptionsPrefix(fdm.getOptionsPrefix()) cdm.setKSPComputeOperators(_SNESContext.compute_operators) cdm.setCreateInterpolation(self.create_interpolation) cdm.setCreateInjection(self.create_injection) # If we're the coarsest grid of the p-hierarchy, don't # overwrite the coarsen routine; this is so that you can # use geometric multigrid for the p-coarse problem try: self.coarsen_element(cele) cdm.setCoarsen(self.coarsen) except ValueError: pass # Injection of the initial state def inject_state(mat): with cu.dat.vec_wo as xc, fu.dat.vec_ro as xf: mat.multTranspose(xf, xc) injection = self.create_injection(cdm, fdm) add_hook(parent, setup=partial(inject_state, injection), call_setup=True) return cdm