def update(self, w): if not self.separate_wind: return if isinstance(w, Function): self.wind.assign(w) else: with self.wind.dat.vec_wo as wind_v: w.copy(wind_v) # Need to update your coarse grid mates, too. # There's probably a better way to do this. dm = self.V.dm wind = self.wind while get_level(get_function_space(dm).mesh())[1] > 0: cdm = dm.coarsen() cctx = get_appctx(cdm) if cctx is None: break cwind = cctx.F._cache['coefficient_mapping'][wind] inject(wind, cwind) dm = cdm wind = cwind
def error_of_flood_estimate(mesh, p, T): # mixed function space h, lvl = get_level(mesh) v_h = FunctionSpace(mesh, "DG", p) v_mu = FunctionSpace(mesh, "DG", p) v_mv = FunctionSpace(mesh, "DG", p) V = v_h*v_mu*v_mv # parameters if p == 0: parameters["flooddrake"].update({"eps2": 8e-3}) if p == 1: parameters["flooddrake"].update({"eps2": 3.5e-3}) # setup free surface depth g = Function(V) x = SpatialCoordinate(V.mesh()) g.sub(0).interpolate(conditional(x[0] < 20, 1.0/9.8, parameters["flooddrake"]["eps2"])) # setup bed bed = Function(V) # setup state state = State(V, g, bed) # setup source (is only a depth function) source = Function(v_h) # timestep solution = Timestepper(V, state.bed, source, 0.05) solution.stepper(0, T, state.w, 0.5) return solution
def coarsen(dm, comm): """Callback to coarsen a DM. :arg DM: The DM to coarsen. :arg comm: The communicator for the new DM (ignored) This transfers a coarse application context over to the coarsened DM (if found on the input DM). """ from firedrake.mg.utils import get_level from firedrake.mg.ufl_utils import coarsen V = get_function_space(dm) if V is None: raise RuntimeError("No functionspace found on DM") hierarchy, level = get_level(V.mesh()) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") if hasattr(V, "_coarse"): cdm = V._coarse.dm else: V._coarse = firedrake.FunctionSpace(hierarchy[level - 1], V.ufl_element()) cdm = V._coarse.dm ctx = get_appctx(dm) if ctx is not None: set_appctx(cdm, coarsen(ctx)) # Necessary for MG inside a fieldsplit in a SNES. cdm.setKSPComputeOperators( firedrake.solving_utils._SNESContext.compute_operators) return cdm
def coarsen(dm, comm): """Callback to coarsen a DM. :arg DM: The DM to coarsen. :arg comm: The communicator for the new DM (ignored) This transfers a coarse application context over to the coarsened DM (if found on the input DM). """ from firedrake.mg.utils import get_level V = get_function_space(dm) if V is None: raise RuntimeError("No functionspace found on DM") hierarchy, level = get_level(V.mesh()) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") if hasattr(V, "_coarse"): cdm = V._coarse.dm else: coarsen = get_ctx_coarsener(dm) V._coarse = coarsen(V, coarsen) cdm = V._coarse.dm transfer = get_transfer_operators(dm) push_transfer_operators(cdm, *transfer) coarsen = get_ctx_coarsener(dm) push_ctx_coarsener(cdm, coarsen) ctx = get_appctx(dm) if ctx is not None: push_appctx(cdm, coarsen(ctx, coarsen)) # Necessary for MG inside a fieldsplit in a SNES. cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators) V._coarse._fine = V return cdm
def coarsen(dm, comm): """Callback to coarsen a DM. :arg DM: The DM to coarsen. :arg comm: The communicator for the new DM (ignored) This transfers a coarse application context over to the coarsened DM (if found on the input DM). """ from firedrake.mg.utils import get_level from firedrake.mg.ufl_utils import coarsen V = get_function_space(dm) if V is None: raise RuntimeError("No functionspace found on DM") hierarchy, level = get_level(V.mesh()) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") if hasattr(V, "_coarse"): cdm = V._coarse.dm else: V._coarse = firedrake.FunctionSpace(hierarchy[level - 1], V.ufl_element()) cdm = V._coarse.dm ctx = get_appctx(dm) if ctx is not None: set_appctx(cdm, coarsen(ctx)) # Necessary for MG inside a fieldsplit in a SNES. cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators) return cdm
def form_function(cls, snes, X, F): """Form the residual for this problem :arg snes: a PETSc SNES object :arg X: the current guess (a Vec) :arg F: the residual at X (a Vec) """ from firedrake.assemble import assemble dm = snes.getDM() ctx = dm.getAppCtx() _, lvl = utils.get_level(dm) # FIXME: Think about case where DM is refined but we don't # have a hierarchy of problems better. if len(ctx._problems) == 1: lvl = -1 problem = ctx._problems[lvl] # X may not be the same vector as the vec behind self._x, so # copy guess in from X. with ctx._xs[lvl].dat.vec as v: if v != X: X.copy(v) assemble(ctx.Fs[lvl], tensor=ctx._Fs[lvl], form_compiler_parameters=problem.form_compiler_parameters, nest=problem._nest) for bc in problem.bcs: bc.zero(ctx._Fs[lvl]) # F may not be the same vector as self._F, so copy # residual out to F. with ctx._Fs[lvl].dat.vec_ro as v: v.copy(F)
def prolong(self, coarse, fine): # Rebuild without any indices V = FunctionSpace(fine.ufl_domain(), fine.function_space().ufl_element()) key = V.dim() firsttime = self.bcs.get(key, None) is None gamma = self.gamma _, level = get_level(fine.function_space().mesh()) nu = self.nus[level] if firsttime: bcs = self.fix_coarse_boundaries(V) u = TrialFunction(V) v = TestFunction(V) tildeu, rhs = Function(V), Function(V) if self.element == "sv": A = assemble(nu * inner(grad(u), grad(v)) * dx + gamma * inner(div(u), div(v)) * dx, bcs=bcs, mat_type=self.patchparams["mat_type"]) bform = nu * inner(grad(rhs), grad(v)) * dx + gamma * inner( div(rhs), div(v)) * dx else: A = assemble(nu * inner(grad(u), grad(v)) * dx + gamma * inner(cell_avg(div(u)), div(v)) * dx, bcs=bcs, mat_type=self.patchparams["mat_type"]) bform = nu * inner(grad(rhs), grad(v)) * dx + gamma * inner( cell_avg(div(rhs)), div(v)) * dx b = Function(V) solver = LinearSolver(A, solver_parameters=self.patchparams, options_prefix="prolongation") self.bcs[key] = bcs self.solver[key] = solver self.rhs[key] = tildeu, rhs self.tensors[key] = A, b, bform else: bcs = self.bcs[key] solver = self.solver[key] A, b, bform = self.tensors[key] tildeu, rhs = self.rhs[key] prolong(coarse, rhs) b = assemble(bform, bcs=bcs, tensor=b) # # Could do # #solver.solve(tildeu, b) # # but that calls a lot of SNES and KSP overhead. # # We know we just want to apply the PC: with solver.inserted_options(): with b.dat.vec_ro as rhsv: with tildeu.dat.vec_wo as x: solver.ksp.pc.apply(rhsv, x) fine.assign(rhs - tildeu)
def inject_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vf.ufl_domain().coordinates key = (("inject", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1] if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs(): return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True)) coordinates = Vf.ufl_domain().coordinates evaluate_kernel = compile_element(ufl.Coefficient(Vf)) to_reference_kernel = to_reference_coordinates(coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) Vf_element = create_element(Vf.ufl_element()) kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf) { double Xref[%(tdim)d]; int cell = -1; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xfi = Xf + i*%(Xf_cell_inc)d; to_reference_coords_kernel(Xref, X, Xfi); if (%(inside_cell)s) { cell = i; break; } } if (cell == -1) { abort(); } const double *fi = f + cell*%(f_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { R[i] = 0; } pyop2_kernel_evaluate(R, fi, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"), "tdim": Vc.ufl_domain().topological_dimension(), "ncandidate": ncandidate, "Rdim": numpy.prod(Vf_element.value_shape), "Xf_cell_inc": coords_element.space_dimension(), "f_cell_inc": Vf_element.space_dimension() } return cache.setdefault(key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def test_get_level(): mesh = UnitSquareMesh(5, 5) L = 2 refinements_per_level = 2 MH = MeshHierarchy(mesh, L, refinements_per_level=refinements_per_level) for i, mesh in enumerate(MH): assert get_level(mesh)[1] == i
def test_get_skip_level(): mesh = UnitSquareMesh(5, 5) L = 2 refinements_per_level = 2 MH = MeshHierarchy(mesh, L, refinements_per_level=refinements_per_level) for i, mesh in enumerate(MH._unskipped_hierarchy): assert get_level(mesh)[1] == Fraction(i, refinements_per_level)
def inject_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vf.ufl_domain().coordinates key = (("inject", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1] if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs(): return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True)) coordinates = Vf.ufl_domain().coordinates evaluate_kernel = compile_element(ufl.Coefficient(Vf)) to_reference_kernel = to_reference_coordinates(coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) Vf_element = create_element(Vf.ufl_element()) kernel = """ %(to_reference)s %(evaluate)s void inject_kernel(double *R, const double *X, const double *f, const double *Xf) { double Xref[%(tdim)d]; int cell = -1; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xfi = Xf + i*%(Xf_cell_inc)d; to_reference_coords_kernel(Xref, X, Xfi); if (%(inside_cell)s) { cell = i; break; } } if (cell == -1) { abort(); } const double *fi = f + cell*%(f_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { R[i] = 0; } evaluate_kernel(R, fi, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"), "tdim": Vc.ufl_domain().topological_dimension(), "ncandidate": ncandidate, "Rdim": numpy.prod(Vf_element.value_shape), "Xf_cell_inc": coords_element.space_dimension(), "f_cell_inc": Vf_element.space_dimension() } return cache.setdefault(key, (op2.Kernel(kernel, name="inject_kernel"), False))
def prolong_kernel(expression): hierarchy, level = utils.get_level(expression.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = expression.ufl_domain().coordinates key = (("prolong", ) + expression.ufl_element().value_shape() + entity_dofs_key(expression.function_space().finat_element.entity_dofs()) + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(expression) to_reference_kernel = to_reference_coordinates(coordinates.ufl_element()) element = create_element(expression.ufl_element()) eval_args = evaluate_kernel.args[:-1] coords_element = create_element(coordinates.ufl_element()) args = eval_args[-1].gencode(not_scope=True) R, coarse = (a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_prolong(double *R, %(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; int cell = -1; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xci = Xc + i*%(Xc_cell_inc)d; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; break; } } if (cell == -1) abort(); const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { %(R)s[i] = 0; } pyop2_kernel_evaluate(%(R)s, coarsei, Xref); } """ % {"to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "args": args, "R": R, "coarse": coarse, "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1], "Rdim": numpy.prod(element.value_shape), "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "tdim": mesh.topological_dimension()} return cache.setdefault(key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))
def restrict_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vc.ufl_domain().coordinates key = (("restrict", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) element = create_element(Vc.ufl_element()) eval_args = evaluate_kernel.args[:-1] args = eval_args[-1].gencode(not_scope=True) R, fine = (a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; int cell = -1; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xci = Xc + i*%(Xc_cell_inc)d; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; const double *Ri = %(R)s + cell*%(coarse_cell_inc)d; pyop2_kernel_evaluate(Ri, %(fine)s, Xref); break; } } } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1], "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "args": args, "R": R, "fine": fine, "tdim": mesh.topological_dimension() } return cache.setdefault( key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def _dm(self): from firedrake.mg.utils import get_level dm = self.dof_dset.dm _, level = get_level(self.mesh()) dmhooks.attach_hooks(dm, level=level, sf=self.mesh()._plex.getPointSF(), section=None) # Remember the function space so we can get from DM back to FunctionSpace. dmhooks.set_function_space(dm, self) return dm
def _dm(self): from firedrake.mg.utils import get_level dm = self.dof_dset.dm _, level = get_level(self.mesh()) dmhooks.attach_hooks(dm, level=level, sf=self.mesh().topology_dm.getPointSF(), section=self._shared_data.global_numbering) # Remember the function space so we can get from DM back to FunctionSpace. dmhooks.set_function_space(dm, self) return dm
def _dm(self): from firedrake.mg.utils import get_level dm = self.dof_dset.dm _, level = get_level(self.mesh()) dmhooks.attach_hooks(dm, level=level, sf=self.mesh()._plex.getPointSF(), section=self._shared_data.global_numbering) # Remember the function space so we can get from DM back to FunctionSpace. dmhooks.set_function_space(dm, self) return dm
def coarsen(dm, comm): """Callback to coarsen a DM. :arg DM: The DM to coarsen. :arg comm: The communicator for the new DM (ignored) This transfers a coarse application context over to the coarsened DM (if found on the input DM). """ from firedrake.mg.utils import get_level V = get_function_space(dm) hierarchy, level = get_level(V.mesh()) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") coarsen = get_ctx_coarsener(dm) Vc = coarsen(V, coarsen) cdm = Vc.dm transfer = get_transfer_operators(dm) parent = get_parent(dm) 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) if len(V) > 1: for V_, Vc_ in zip(V, Vc): transfer = get_transfer_operators(V_.dm) add_hook(parent, setup=partial(push_parent, Vc_.dm, parent), teardown=partial(pop_parent, Vc_.dm, parent), call_setup=True) add_hook(parent, setup=partial(push_transfer_operators, Vc_.dm, transfer), teardown=partial(pop_transfer_operators, Vc_.dm, transfer), call_setup=True) add_hook(parent, setup=partial(push_ctx_coarsener, cdm, coarsen), teardown=partial(pop_ctx_coarsener, cdm, coarsen), call_setup=True) ctx = get_appctx(dm) if ctx is not None: cctx = coarsen(ctx, coarsen) add_hook(parent, setup=partial(push_appctx, cdm, cctx), teardown=partial(pop_appctx, cdm, cctx), call_setup=True) # Necessary for MG inside a fieldsplit in a SNES. cdm.setKSPComputeOperators( firedrake.solving_utils._SNESContext.compute_operators) return cdm
def aug_jacobian(X, J, ctx): mh, level = get_level(ctx._x.ufl_domain()) if case == 4 or case == 5: BTWBlevel = BTWB_dict[level] if level == nref: Jsub = J.getNestSubMatrix(0, 0) rmap, cmap = Jsub.getLGMap() Jsub.axpy(1, BTWBlevel, Jsub.Structure.SUBSET_NONZERO_PATTERN) Jsub.setLGMap(rmap, cmap) else: rmap, cmap = J.getLGMap() J.axpy(1, BTWBlevel, J.Structure.SUBSET_NONZERO_PATTERN) J.setLGMap(rmap, cmap)
def __call__(self, pc): from firedrake.mg.utils import get_level from firedrake.cython.mgimpl import get_entity_renumbering dmf = pc.getDM() ctx = pc.getAttr("ctx") mf = ctx._x.ufl_domain() (mh, level) = get_level(mf) coarse_to_fine_cell_map = mh.coarse_to_fine_cells[level - 1] (_, firedrake_to_plex) = get_entity_renumbering(dmf, mf._cell_numbering, "cell") mc = mh[level - 1] (_, coarse_firedrake_to_plex) = get_entity_renumbering( mc._topology_dm, mc._cell_numbering, "cell") patches = [] tdim = mf.topological_dimension() for i, fine_firedrake in enumerate(coarse_to_fine_cell_map): # there are d+1 many coarse cells that all map to the same fine cells. # We only want to build the patch once, so skip repitions if coarse_firedrake_to_plex[i] % (tdim + 1) != 0: continue # we need to convert firedrake cell numbering to plex cell numbering fine_plex = [firedrake_to_plex[ff] for ff in fine_firedrake] entities = [] for fp in fine_plex: (pts, _) = dmf.getTransitiveClosure(fp, True) for pt in pts: value = dmf.getLabelValue("prolongation", pt) if not (value > -1 and value <= level): entities.append(pt) iset = PETSc.IS().createGeneral(unique(entities), comm=PETSc.COMM_SELF) patches.append(iset) piterset = PETSc.IS().createStride(size=len(patches), first=0, step=1, comm=PETSc.COMM_SELF) return (patches, piterset)
def test_inject_to_any_level(): M = MeshHierarchy(UnitSquareMesh(10, 10), 3) V = FunctionSpaceHierarchy(M, 'DG', 0) F = FunctionHierarchy(V) # check for prolonging to top level F[-1].interpolate(Expression("3")) F[0].interpolate(Expression("3")) H = FunctionHierarchy(V) A = InjectDownToAnyLevel(0, F[-1], H) assert get_level(A)[1] == 0 assert norm(assemble(A - F[0])) <= 0
def test_prolong_to_any_level(): M = MeshHierarchy(UnitSquareMesh(10, 10), 3) V = FunctionSpaceHierarchy(M, 'DG', 0) F = FunctionHierarchy(V) # check for prolonging to top level F[0].interpolate(Expression("3")) F[2].interpolate(Expression("3")) H = FunctionHierarchy(V) A = ProlongUpToAnyLevel(2, F[0], H) assert get_level(A)[1] == 2 assert norm(assemble(A - F[2])) <= 0
def refine(dm, comm): """Callback to refine a DM. :arg DM: The DM to refine. :arg comm: The communicator for the new DM (ignored) """ from firedrake.mg.utils import get_level V = get_function_space(dm) if V is None: raise RuntimeError("No functionspace found on DM") hierarchy, level = get_level(V.mesh()) if level >= len(hierarchy) - 1: raise RuntimeError("Cannot refine finest DM") if hasattr(V, "_fine"): fdm = V._fine.dm else: V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element()) fdm = V._fine.dm return fdm
def form_jacobian(cls, snes, X, J, P): """Form the Jacobian for this problem :arg snes: a PETSc SNES object :arg X: the current guess (a Vec) :arg J: the Jacobian (a Mat) :arg P: the preconditioner matrix (a Mat) """ from firedrake.assemble import assemble dm = snes.getDM() ctx = dm.getAppCtx() _, lvl = utils.get_level(dm) # FIXME: Think about case where DM is refined but we don't # have a hierarchy of problems better. if len(ctx._problems) == 1: lvl = -1 problem = ctx._problems[lvl] if problem._constant_jacobian and ctx._jacobians_assembled[lvl]: # Don't need to do any work with a constant jacobian # that's already assembled return ctx._jacobians_assembled[lvl] = True # X may not be the same vector as the vec behind self._x, so # copy guess in from X. with ctx._xs[lvl].dat.vec as v: X.copy(v) assemble(ctx.Js[lvl], tensor=ctx._jacs[lvl], bcs=problem.bcs, form_compiler_parameters=problem.form_compiler_parameters, nest=problem._nest) ctx._jacs[lvl].M._force_evaluation() if ctx.Jps[lvl] is not None: assemble(ctx.Jps[lvl], tensor=ctx._pjacs[lvl], bcs=problem.bcs, form_compiler_parameters=problem.form_compiler_parameters, nest=problem._nest) ctx._pjacs[lvl].M._force_evaluation()
def refine(dm, comm): """Callback to refine a DM. :arg DM: The DM to refine. :arg comm: The communicator for the new DM (ignored) """ from firedrake.mg.utils import get_level V = get_function_space(dm) if V is None: raise RuntimeError("No functionspace found on DM") hierarchy, level = get_level(V.mesh()) if level >= len(hierarchy) - 1: raise RuntimeError("Cannot refine finest DM") if hasattr(V, "_fine"): fdm = V._fine.dm else: V._fine = firedrake.FunctionSpace(hierarchy[level + 1], V.ufl_element()) fdm = V._fine.dm V._fine._coarse = V return fdm
def prolong_kernel(expression): hierarchy, _ = utils.get_level(expression.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = expression.ufl_domain().coordinates key = (("prolong", ) + expression.ufl_element().value_shape() + entity_dofs_key(expression.function_space().finat_element.entity_dofs()) + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(expression) to_reference_kernel = to_reference_coordinates(coordinates.ufl_element()) element = create_element(expression.ufl_element()) eval_args = evaluate_kernel.args[:-1] args = ", ".join(a.gencode(not_scope=True) for a in eval_args) arg_names = ", ".join(a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s void prolong_kernel(%(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; to_reference_coords_kernel(Xref, X, Xc); for ( int i = 0; i < %(Rdim)d; i++ ) { %(R)s[i] = 0; } evaluate_kernel(%(arg_names)s, Xref); } """ % {"to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "args": args, "R": arg_names[0], "Rdim": numpy.prod(element.value_shape), "arg_names": arg_names, "tdim": mesh.topological_dimension()} return cache.setdefault(key, op2.Kernel(my_kernel, name="prolong_kernel"))
def fix_coarse_boundaries(V): hierarchy, level = get_level(V.mesh()) dm = V.mesh()._topology_dm section = V.dm.getDefaultSection() indices = [] fStart, fEnd = dm.getHeightStratum(1) # Spin over faces, if the face is marked with a magic label # value, it means it was in the coarse mesh. for p in range(fStart, fEnd): value = dm.getLabelValue("prolongation", p) if value > -1 and value <= level: # OK, so this is a coarse mesh face. # Grab all the points in the closure. closure, _ = dm.getTransitiveClosure(p) for c in closure: # Now add all the dofs on that point to the list # of boundary nodes. dof = section.getDof(c) off = section.getOffset(c) for d in range(dof): indices.append(off + d) nodelist = unique(indices).astype(IntType) class FixedDirichletBC(DirichletBC): def __init__(self, V, g, nodelist): self.nodelist = nodelist DirichletBC.__init__(self, V, g, "on_boundary") @utils.cached_property def nodes(self): return self.nodelist dim = V.mesh().topological_dimension() bc = FixedDirichletBC(V, ufl.zero(V.ufl_element().value_shape()), nodelist) return bc
def standard_transfer(self, source, target, mode): if not (source.ufl_shape[0] == 3 and "CG1" in source.ufl_element().shortstr()): return super().standard_transfer(source, target, mode) if mode == "prolong": coarse = source fine = target elif mode == "restrict": fine = source coarse = target else: raise NotImplementedError (mh, level) = get_level(coarse.ufl_domain()) if level not in self.transfers: self.transfers[level] = BubbleTransfer(coarse.function_space(), fine.function_space()) if mode == "prolong": self.transfers[level].prolong(coarse, fine) elif mode == "restrict": self.transfers[level].restrict(fine, coarse) else: raise NotImplementedError
def __call__(self, pc): from firedrake.mg.utils import get_level from firedrake.cython.mgimpl import get_entity_renumbering dmf = pc.getDM() ctx = pc.getAttr("ctx") mf = ctx._x.ufl_domain() (mh, level) = get_level(mf) coarse_to_fine_cell_map = mh.coarse_to_fine_cells[level - 1] (_, firedrake_to_plex) = get_entity_renumbering(dmf, mf._cell_numbering, "cell") patches = [] for fine_firedrake in coarse_to_fine_cell_map: # we need to convert firedrake cell numbering to plex cell numbering fine_plex = [firedrake_to_plex[ff] for ff in fine_firedrake] entities = [] for fp in fine_plex: (pts, _) = dmf.getTransitiveClosure(fp, True) for pt in pts: value = dmf.getLabelValue("prolongation", pt) if not (value > -1 and value <= level): entities.append(pt) iset = PETSc.IS().createGeneral(unique(entities), comm=PETSc.COMM_SELF) patches.append(iset) piterset = PETSc.IS().createStride(size=len(patches), first=0, step=1, comm=PETSc.COMM_SELF) return (patches, piterset)
def coarsen(dm, comm): """Callback to coarsen a DM. :arg DM: The DM to coarsen. :arg comm: The communicator for the new DM (ignored) This transfers a coarse application context over to the coarsened DM (if found on the input DM). """ from firedrake.mg.utils import get_level V = get_function_space(dm) hierarchy, level = get_level(V.mesh()) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") coarsen = get_ctx_coarsener(dm) Vc = coarsen(V, coarsen) cdm = Vc.dm push_ctx_coarsener(cdm, coarsen) ctx = get_appctx(dm) if ctx is not None: push_appctx(cdm, coarsen(ctx, coarsen)) # Necessary for MG inside a fieldsplit in a SNES. cdm.setKSPComputeOperators(firedrake.solving_utils._SNESContext.compute_operators) return cdm
def error_of_flood_estimate(mesh, p, T): # mixed function space h, lvl = get_level(mesh) v_h = FunctionSpace(mesh, "DG", p) v_mu = FunctionSpace(mesh, "DG", p) v_mv = FunctionSpace(mesh, "DG", p) V = v_h * v_mu * v_mv # parameters if p == 0: parameters["flooddrake"].update({"eps2": 8e-3}) if p == 1: parameters["flooddrake"].update({"eps2": 3.5e-3}) # setup free surface depth g = Function(V) x = SpatialCoordinate(V.mesh()) g.sub(0).interpolate( conditional(x[0] < 20, 1.0 / 9.8, parameters["flooddrake"]["eps2"])) # setup bed bed = Function(V) # setup state state = State(V, g, bed) # setup source (is only a depth function) source = Function(v_h) # timestep solution = Timestepper(V, state.bed, source, 0.05) solution.stepper(0, T, state.w, 0.5) return solution
def restrict_kernel(Vf, Vc): hierarchy, _ = utils.get_level(Vc.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vc.ufl_domain().coordinates key = (("restrict", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf) to_reference_kernel = to_reference_coordinates(coordinates.ufl_element()) eval_args = evaluate_kernel.args[:-1] args = ", ".join(a.gencode(not_scope=True) for a in eval_args) arg_names = ", ".join(a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s void restrict_kernel(%(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; to_reference_coords_kernel(Xref, X, Xc); evaluate_kernel(%(arg_names)s, Xref); } """ % {"to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "args": args, "arg_names": arg_names, "tdim": mesh.topological_dimension()} return cache.setdefault(key, op2.Kernel(my_kernel, name="restrict_kernel"))
def _dm(self): from firedrake.mg.utils import get_level dm = self.dof_dset.dm _, level = get_level(self.mesh()) dmhooks.attach_hooks(dm, level=level) return dm
def create_matrix(cls, dm): ctx = dm.getAppCtx() _, lvl = utils.get_level(dm) return ctx._jacs[lvl]._M.handle
def aug_jacobian(X, J, ctx): mh, level = get_level(ctx._x.ufl_domain()) levelmesh = mh[level] if case == 4 or case == 5: if args.quad: Vlevel = FunctionSpace(levelmesh, "RTCF", k) Qlevel = FunctionSpace(levelmesh, "DQ", k - 1) else: if args.discretisation == "hdiv": Vlevel = FunctionSpace(levelmesh, "BDM", k) Qlevel = FunctionSpace(levelmesh, "DG", k - 1) elif args.discretisation == "cg": if dim == 2: Vlevel = VectorFunctionSpace(levelmesh, "CG", k) Qlevel = FunctionSpace(levelmesh, "DG", 0) elif dim == 3: Pklevel = FiniteElement("Lagrange", levelmesh.ufl_cell(), 2) FBlevel = FiniteElement("FacetBubble", levelmesh.ufl_cell(), 3) eleulevel = VectorElement( NodalEnrichedElement(Pklevel, FBlevel)) Vlevel = FunctionSpace(levelmesh, eleulevel) Qlevel = FunctionSpace(levelmesh, "DG", 0) else: raise NotImplementedError("Only implemented for dim=2,3") else: raise ValueError( "please specify hdiv or cg for --discretisation") Zlevel = Vlevel * Qlevel # Get B tmpu, tmpp = TrialFunctions(Zlevel) tmpv, tmpq = TestFunctions(Zlevel) tmpF = -tmpq * div(tmpu) * dx if dim == 2: tmpbcs = [ DirichletBC(Zlevel.sub(0), Constant((0., 0.)), "on_boundary") ] elif dim == 3: tmpbcs = [ DirichletBC(Zlevel.sub(0), Constant((0., 0., 0.)), "on_boundary") ] else: raise NotImplementedError("Only implemented for dim=2,3") tmpa = lhs(tmpF) M = assemble(tmpa, bcs=tmpbcs) Blevel = M.M[1, 0].handle # Get W ptrial = TrialFunction(Qlevel) ptest = TestFunction(Qlevel) if case == 4: Wlevel = assemble(Tensor(inner(ptrial, ptest) * dx).inv).M[0, 0].handle if case == 5: Wlevel = assemble(Tensor(1.0/mu(levelmesh)*inner(ptrial, ptest)*\ dx).inv).M[0,0].handle # Form BTWB BTWlevel = Blevel.transposeMatMult(Wlevel) BTWlevel *= args.gamma BTWBlevel = BTWlevel.matMult(Blevel) J.axpy(1, BTWBlevel, structure=J.Structure.SUBSET_NONZERO_PATTERN)
def assemble_mixed_mass_matrix(V_A, V_B): """ Construct the mixed mass matrix of two function spaces, using the TrialFunction from V_A and the TestFunction from V_B. """ if len(V_A) > 1 or len(V_B) > 1: raise NotImplementedError( "Sorry, only implemented for non-mixed spaces") if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping( ) != "identity": msg = """ Sorry, only implemented for affine maps for now. To do non-affine, we'd need to import much more of the assembly engine of UFL/TSFC/etc to do the assembly on each supermesh cell. """ raise NotImplementedError(msg) mesh_A = V_A.mesh() mesh_B = V_B.mesh() dim = mesh_A.geometric_dimension() assert dim == mesh_B.geometric_dimension() assert dim == mesh_A.topological_dimension() assert dim == mesh_B.topological_dimension() (mh_A, level_A) = get_level(mesh_A) (mh_B, level_B) = get_level(mesh_B) if mesh_A is mesh_B: def likely(cell_A): return [cell_A] else: if (mh_A is None or mh_B is None) or (mh_A is not mh_B): # No mesh hierarchy structure, call libsupermesh for # intersection finding intersections = intersection_finder(mesh_A, mesh_B) likely = intersections.__getitem__ else: # We do have a mesh hierarchy, use it if abs(level_A - level_B) > 1: raise NotImplementedError( "Only works for transferring between adjacent levels for now." ) # What are the cells of B that (probably) intersect with a given cell in A? if level_A > level_B: cell_map = mh_A.fine_to_coarse_cells[level_A] def likely(cell_A): return cell_map[cell_A] elif level_A < level_B: cell_map = mh_A.coarse_to_fine_cells[level_A] def likely(cell_A): return cell_map[cell_A] assert V_A.value_size == V_B.value_size orig_value_size = V_A.value_size if V_A.value_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements()[0]) if V_B.value_size > 1: V_B = firedrake.FunctionSpace(mesh_B, V_B.ufl_element().sub_elements()[0]) assert V_A.value_size == 1 assert V_B.value_size == 1 preallocator = PETSc.Mat().create(comm=mesh_A.comm) preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) rset = V_B.dof_dset cset = V_A.dof_dset nrows = rset.layout_vec.getSizes() ncols = cset.layout_vec.getSizes() preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap) preallocator.setSizes(size=(nrows, ncols), bsize=1) preallocator.setUp() zeros = numpy.zeros((V_B.cell_node_map().arity, V_A.cell_node_map().arity), dtype=ScalarType) for cell_A, dofs_A in enumerate(V_A.cell_node_map().values): for cell_B in likely(cell_A): dofs_B = V_B.cell_node_map().values_with_halo[cell_B, :] preallocator.setValuesLocal(dofs_B, dofs_A, zeros) preallocator.assemble() dnnz, onnz = get_preallocation(preallocator, nrows[0]) # Unroll from block to AIJ dnnz = dnnz * cset.cdim dnnz = numpy.repeat(dnnz, rset.cdim) onnz = onnz * cset.cdim onnz = numpy.repeat(onnz, cset.cdim) preallocator.destroy() assert V_A.value_size == V_B.value_size rdim = V_B.dof_dset.cdim cdim = V_A.dof_dset.cdim # # Preallocate M_AB. # mat = PETSc.Mat().create(comm=mesh_A.comm) mat.setType(PETSc.Mat.Type.AIJ) rsizes = tuple(n * rdim for n in nrows) csizes = tuple(c * cdim for c in ncols) mat.setSizes(size=(rsizes, csizes), bsize=(rdim, cdim)) mat.setPreallocationNNZ((dnnz, onnz)) mat.setLGMap(rmap=rset.lgmap, cmap=cset.lgmap) # TODO: Boundary conditions not handled. mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True) mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, False) mat.setOption(mat.Option.IGNORE_ZERO_ENTRIES, True) mat.setUp() evaluate_kernel_A = compile_element(ufl.Coefficient(V_A), name="evaluate_kernel_A") evaluate_kernel_B = compile_element(ufl.Coefficient(V_B), name="evaluate_kernel_B") # We only need one of these since we assume that the two meshes both have CG1 coordinates to_reference_kernel = to_reference_coordinates( mesh_A.coordinates.ufl_element()) if dim == 2: reference_mesh = UnitTriangleMesh(comm=COMM_SELF) else: reference_mesh = UnitTetrahedronMesh(comm=COMM_SELF) evaluate_kernel_S = compile_element(ufl.Coefficient( reference_mesh.coordinates.function_space()), name="evaluate_kernel_S") V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element()) V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element()) M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx) M_SS = M_SS.M.handle[:, :] node_locations_A = utils.physical_node_locations( V_S_A).dat.data_ro_with_halos node_locations_B = utils.physical_node_locations( V_S_B).dat.data_ro_with_halos num_nodes_A = node_locations_A.shape[0] num_nodes_B = node_locations_B.shape[0] to_reference_kernel = to_reference_coordinates( mesh_A.coordinates.ufl_element()) supermesh_kernel_str = """ #include "libsupermesh-c.h" #include <petsc.h> %(to_reference)s %(evaluate_S)s %(evaluate_A)s %(evaluate_B)s #define complex_mode %(complex_mode)s #define PrintInfo(...) do { if (PetscLogPrintInfo) printf(__VA_ARGS__); } while (0) static void print_array(PetscScalar *arr, int d) { for(int j=0; j<d; j++) PrintInfo(stderr, "%%+.2f ", arr[j]); } static void print_coordinates(PetscScalar *simplex, int d) { for(int i=0; i<d+1; i++) { PrintInfo("\t"); print_array(&simplex[d*i], d); PrintInfo("\\n"); } } #if complex_mode static void seperate_real_and_imag(PetscScalar *simplex, double *real_simplex, double *imag_simplex, int d) { for(int i=0; i<d+1; i++) { for(int j=0; j<d; j++) { real_simplex[d*i+j] = creal(simplex[d*i+j]); imag_simplex[d*i+j] = cimag(simplex[d*i+j]); } } } static void merge_back_to_simplex(PetscScalar* simplex, double* real_simplex, double* imag_simplex, int d) { print_coordinates(simplex,d); for(int i=0; i<d+1; i++) { for(int j=0; j<d; j++) { simplex[d*i+j] = real_simplex[d*i+j]+imag_simplex[d*i+j]*_Complex_I; } } } #endif int supermesh_kernel(PetscScalar* simplex_A, PetscScalar* simplex_B, PetscScalar* simplices_C, PetscScalar* nodes_A, PetscScalar* nodes_B, PetscScalar* M_SS, PetscScalar* outptr, int num_ele) { #define d %(dim)s #define num_nodes_A %(num_nodes_A)s #define num_nodes_B %(num_nodes_B)s double simplex_ref_measure; PrintInfo("simplex_A coordinates\\n"); print_coordinates(simplex_A, d); PrintInfo("simplex_B coordinates\\n"); print_coordinates(simplex_B, d); int num_elements = num_ele; if (d == 2) simplex_ref_measure = 0.5; else if (d == 3) simplex_ref_measure = 1.0/6; PetscScalar R_AS[num_nodes_A][num_nodes_A]; PetscScalar R_BS[num_nodes_B][num_nodes_B]; PetscScalar coeffs_A[%(num_nodes_A)s] = {0.}; PetscScalar coeffs_B[%(num_nodes_B)s] = {0.}; PetscScalar reference_nodes_A[num_nodes_A][d]; PetscScalar reference_nodes_B[num_nodes_B][d]; #if complex_mode double real_simplex_A[d*(d+1)]; double imag_simplex_A[d*(d+1)]; seperate_real_and_imag(simplex_A, real_simplex_A, imag_simplex_A, d); double real_simplex_B[d*(d+1)]; double imag_simplex_B[d*(d+1)]; seperate_real_and_imag(simplex_B, real_simplex_B, imag_simplex_B, d); double real_simplices_C[num_elements*d*(d+1)]; double imag_simplices_C[num_elements*d*(d+1)]; for (int ii=0; ii<num_elements*d*(d+1); ++ii) imag_simplices_C[ii] = 0.; %(libsupermesh_intersect_simplices)s(real_simplex_A, real_simplex_B, real_simplices_C, &num_elements); merge_back_to_simplex(simplex_A, real_simplex_A, imag_simplex_A, d); merge_back_to_simplex(simplex_B, real_simplex_B, imag_simplex_B, d); for(int s=0; s<num_elements; s++) { PetscScalar* simplex_C = &simplices_C[s * d * (d+1)]; double* real_simplex_C = &real_simplices_C[s * d * (d+1)]; double* imag_simplex_C = &imag_simplices_C[s * d * (d+1)]; merge_back_to_simplex(simplex_C, real_simplex_C, imag_simplex_C, d); } #else %(libsupermesh_intersect_simplices)s(simplex_A, simplex_B, simplices_C, &num_elements); #endif PrintInfo("Supermesh consists of %%i elements\\n", num_elements); // would like to do this //PetscScalar MAB[%(num_nodes_A)s][%(num_nodes_B)s] = (PetscScalar (*)[%(num_nodes_B)s])outptr; // but have to do this instead because we don't grok C PetscScalar (*MAB)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])outptr; PetscScalar (*MSS)[num_nodes_A] = (PetscScalar (*)[num_nodes_A])M_SS; // note the underscore for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { MAB[i][j] = 0.0; } } for(int s=0; s<num_elements; s++) { PetscScalar* simplex_S = &simplices_C[s * d * (d+1)]; double simplex_S_measure; #if complex_mode double real_simplex_S[d*(d+1)]; double imag_simplex_S[d*(d+1)]; seperate_real_and_imag(simplex_S, real_simplex_S, imag_simplex_S, d); %(libsupermesh_simplex_measure)s(real_simplex_S, &simplex_S_measure); merge_back_to_simplex(simplex_S, real_simplex_S, imag_simplex_S, d); #else %(libsupermesh_simplex_measure)s(simplex_S, &simplex_S_measure); #endif PrintInfo("simplex_S coordinates with measure %%f\\n", simplex_S_measure); print_coordinates(simplex_S, d); PrintInfo("Start mapping nodes for V_A\\n"); PetscScalar physical_nodes_A[num_nodes_A][d]; for(int n=0; n < num_nodes_A; n++) { PetscScalar* reference_node_location = &nodes_A[n*d]; PetscScalar* physical_node_location = physical_nodes_A[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("Start mapping nodes for V_B\\n"); PetscScalar physical_nodes_B[num_nodes_B][d]; for(int n=0; n < num_nodes_B; n++) { PetscScalar* reference_node_location = &nodes_B[n*d]; PetscScalar* physical_node_location = physical_nodes_B[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("==========================================================\\n"); PrintInfo("Start pulling back dof from S into reference space for A.\\n"); for(int n=0; n < num_nodes_A; n++) { for(int i=0; i<d; i++) reference_nodes_A[n][i] = 0.; to_reference_coords_kernel(reference_nodes_A[n], physical_nodes_A[n], simplex_A); PrintInfo("Pulling back "); print_array(physical_nodes_A[n], d); PrintInfo(" to "); print_array(reference_nodes_A[n], d); PrintInfo("\\n"); } PrintInfo("Start pulling back dof from S into reference space for B.\\n"); for(int n=0; n < num_nodes_B; n++) { for(int i=0; i<d; i++) reference_nodes_B[n][i] = 0.; to_reference_coords_kernel(reference_nodes_B[n], physical_nodes_B[n], simplex_B); PrintInfo("Pulling back "); print_array(physical_nodes_B[n], d); PrintInfo(" to "); print_array(reference_nodes_B[n], d); PrintInfo("\\n"); } PrintInfo("Start evaluating basis functions of V_A at dofs for V_A on S\\n"); for(int i=0; i<num_nodes_A; i++) { coeffs_A[i] = 1.; for(int j=0; j<num_nodes_A; j++) { R_AS[i][j] = 0.; pyop2_kernel_evaluate_kernel_A(&R_AS[i][j], coeffs_A, reference_nodes_A[j]); } print_array(R_AS[i], num_nodes_A); PrintInfo("\\n"); coeffs_A[i] = 0.; } PrintInfo("Start evaluating basis functions of V_B at dofs for V_B on S\\n"); for(int i=0; i<num_nodes_B; i++) { coeffs_B[i] = 1.; for(int j=0; j<num_nodes_B; j++) { R_BS[i][j] = 0.; pyop2_kernel_evaluate_kernel_B(&R_BS[i][j], coeffs_B, reference_nodes_B[j]); } print_array(R_BS[i], num_nodes_B); PrintInfo("\\n"); coeffs_B[i] = 0.; } PrintInfo("Start doing the matmatmat mult\\n"); for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { for ( int k = 0; k < num_nodes_B; k++) { for ( int l = 0; l < num_nodes_A; l++) { MAB[i][j] += (simplex_S_measure/simplex_ref_measure) * R_BS[i][k] * MSS[k][l] * R_AS[j][l]; } } } } } return num_elements; } """ % { "evaluate_S": str(evaluate_kernel_S), "evaluate_A": str(evaluate_kernel_A), "evaluate_B": str(evaluate_kernel_B), "to_reference": str(to_reference_kernel), "num_nodes_A": num_nodes_A, "num_nodes_B": num_nodes_B, "libsupermesh_simplex_measure": "libsupermesh_triangle_area" if dim == 2 else "libsupermesh_tetrahedron_volume", "libsupermesh_intersect_simplices": "libsupermesh_intersect_tris_real" if dim == 2 else "libsupermesh_intersect_tets_real", "dim": dim, "complex_mode": 1 if complex_mode else 0 } dirs = get_petsc_dir() + (sys.prefix, ) includes = ["-I%s/include" % d for d in dirs] libs = ["-L%s/lib" % d for d in dirs] libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"] lib = load(supermesh_kernel_str, "c", "supermesh_kernel", cppargs=includes, ldargs=libs, argtypes=[ ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp ], restype=ctypes.c_int) ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat) if orig_value_size == 1: return mat else: (lrows, grows), (lcols, gcols) = mat.getSizes() lrows *= orig_value_size grows *= orig_value_size lcols *= orig_value_size gcols *= orig_value_size size = ((lrows, grows), (lcols, gcols)) context = BlockMatrix(mat, orig_value_size) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat
def test_convergence_decerasing_nl(): mesh = UnitSquareMesh(4, 4) L = 4 # Generate the Mesh Hierarchy and the FunctionSpaceHierarchies Mesh_Hierarchy = GenerateMeshHierarchy(mesh, L) FunctionSpaceHierarchies = FunctionSpaceHierarchy(Mesh_Hierarchy, 'CG', 1) # Create Discretization Object ensemble_hierarchy = EnsembleHierarchy() level_to_prolong_to = 5 eps = 4e-2 n = 16 converge = 0 i = 0 Nl = [] while converge == 0: lvlc = i lvlf = i + 1 for j in range(n): u_h = FunctionHierarchy(FunctionSpaceHierarchies) u = state(u_h[lvlc], u_h[lvlf]) xc = np.random.normal(0.0, 0.05, 1) u.state[0].interpolate( Expression( "exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))", x_center=xc)) u.state[1].interpolate( Expression( "exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))", x_center=xc)) u.state[0].assign( Solve( lvlc, FunctionSpaceHierarchies, u.state[0])) u.state[1].assign( Solve( lvlf, FunctionSpaceHierarchies, u.state[1])) u_new = state( ProlongUpToFinestLevel( u.state[0], FunctionHierarchy(FunctionSpaceHierarchies)), ProlongUpToFinestLevel( u.state[1], FunctionHierarchy(FunctionSpaceHierarchies))) ensemble_hierarchy.AppendToEnsemble(u_new.state, i) assert get_level(u_new.state[1])[1] == len(Mesh_Hierarchy) - 1 Nl.append(n) # Calculate sample statistics SampleStatistics(ensemble_hierarchy) Ns = OptimalNl(ensemble_hierarchy, eps) for j in range(i + 1): if Ns[j] > Nl[j]: dN = Ns[j] - Nl[j] for k in range(dN): u_h = FunctionHierarchy(FunctionSpaceHierarchies) u = state(u_h[lvlc], u_h[lvlf]) xc = np.random.normal(0.0, 0.05, 1) u.state[0].interpolate( Expression( "exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))", x_center=xc)) u.state[1].interpolate( Expression( "exp(-(pow((x[0]-0.5+x_center),2)/0.25)-(pow(x[1]-0.5,2)/0.25))", x_center=xc)) u.state[0].assign( Solve( lvlc, FunctionSpaceHierarchies, u.state[0])) u.state[1].assign( Solve( lvlf, FunctionSpaceHierarchies, u.state[1])) u_new = state( ProlongUpToFinestLevel( u.state[0], FunctionHierarchy(FunctionSpaceHierarchies)), ProlongUpToFinestLevel( u.state[1], FunctionHierarchy(FunctionSpaceHierarchies))) ensemble_hierarchy.AppendToEnsemble(u_new.state, i) Nl[j] = Ns[j] # Recalculate sample statistics SampleStatistics(ensemble_hierarchy) converge = Convergence(ensemble_hierarchy, eps) i += 1 assert np.all(np.diff(np.array(Nl)) <= 0) == 1 return ensemble_hierarchy
def aug_jacobian(X, J, ctx): mh, level = get_level(ctx._x.ufl_domain()) if args.galerkin: rmap, cmap = J.getLGMap() levelOps[nref-level].copy(J,structure=J.Structure.DIFFERENT_NONZERO_PATTERN) J.setLGMap(rmap, cmap)
def assemble_mixed_mass_matrix(V_A, V_B): """ Construct the mixed mass matrix of two function spaces, using the TrialFunction from V_A and the TestFunction from V_B. """ if len(V_A) > 1 or len(V_B) > 1: raise NotImplementedError("Sorry, only implemented for non-mixed spaces") if V_A.ufl_element().mapping() != "identity" or V_B.ufl_element().mapping() != "identity": msg = """ Sorry, only implemented for affine maps for now. To do non-affine, we'd need to import much more of the assembly engine of UFL/TSFC/etc to do the assembly on each supermesh cell. """ raise NotImplementedError(msg) mesh_A = V_A.mesh() mesh_B = V_B.mesh() dim = mesh_A.geometric_dimension() assert dim == mesh_B.geometric_dimension() assert dim == mesh_A.topological_dimension() assert dim == mesh_B.topological_dimension() (mh_A, level_A) = get_level(mesh_A) (mh_B, level_B) = get_level(mesh_B) if mesh_A is mesh_B: def likely(cell_A): return [cell_A] else: if (mh_A is None or mh_B is None) or (mh_A is not mh_B): # No mesh hierarchy structure, call libsupermesh for # intersection finding intersections = intersection_finder(mesh_A, mesh_B) likely = intersections.__getitem__ else: # We do have a mesh hierarchy, use it if abs(level_A - level_B) > 1: raise NotImplementedError("Only works for transferring between adjacent levels for now.") # What are the cells of B that (probably) intersect with a given cell in A? if level_A > level_B: cell_map = mh_A.fine_to_coarse_cells[level_A] def likely(cell_A): return cell_map[cell_A] elif level_A < level_B: cell_map = mh_A.coarse_to_fine_cells[level_A] def likely(cell_A): return cell_map[cell_A] assert V_A.value_size == V_B.value_size orig_value_size = V_A.value_size if V_A.value_size > 1: V_A = firedrake.FunctionSpace(mesh_A, V_A.ufl_element().sub_elements()[0]) if V_B.value_size > 1: V_B = firedrake.FunctionSpace(mesh_B, V_B.ufl_element().sub_elements()[0]) assert V_A.value_size == 1 assert V_B.value_size == 1 preallocator = PETSc.Mat().create(comm=mesh_A.comm) preallocator.setType(PETSc.Mat.Type.PREALLOCATOR) rset = V_B.dof_dset cset = V_A.dof_dset nrows = rset.layout_vec.getSizes() ncols = cset.layout_vec.getSizes() preallocator.setLGMap(rmap=rset.scalar_lgmap, cmap=cset.scalar_lgmap) preallocator.setSizes(size=(nrows, ncols), bsize=1) preallocator.setUp() zeros = numpy.zeros((V_B.cell_node_map().arity, V_A.cell_node_map().arity), dtype=ScalarType) for cell_A, dofs_A in enumerate(V_A.cell_node_map().values): for cell_B in likely(cell_A): dofs_B = V_B.cell_node_map().values_with_halo[cell_B, :] preallocator.setValuesLocal(dofs_B, dofs_A, zeros) preallocator.assemble() dnnz, onnz = get_preallocation(preallocator, nrows[0]) # Unroll from block to AIJ dnnz = dnnz * cset.cdim dnnz = numpy.repeat(dnnz, rset.cdim) onnz = onnz * cset.cdim onnz = numpy.repeat(onnz, cset.cdim) preallocator.destroy() assert V_A.value_size == V_B.value_size rdim = V_B.dof_dset.cdim cdim = V_A.dof_dset.cdim # # Preallocate M_AB. # mat = PETSc.Mat().create(comm=mesh_A.comm) mat.setType(PETSc.Mat.Type.AIJ) rsizes = tuple(n * rdim for n in nrows) csizes = tuple(c * cdim for c in ncols) mat.setSizes(size=(rsizes, csizes), bsize=(rdim, cdim)) mat.setPreallocationNNZ((dnnz, onnz)) mat.setLGMap(rmap=rset.lgmap, cmap=cset.lgmap) # TODO: Boundary conditions not handled. mat.setOption(mat.Option.IGNORE_OFF_PROC_ENTRIES, False) mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, True) mat.setOption(mat.Option.KEEP_NONZERO_PATTERN, True) mat.setOption(mat.Option.UNUSED_NONZERO_LOCATION_ERR, False) mat.setOption(mat.Option.IGNORE_ZERO_ENTRIES, True) mat.setUp() evaluate_kernel_A = compile_element(ufl.Coefficient(V_A), name="evaluate_kernel_A") evaluate_kernel_B = compile_element(ufl.Coefficient(V_B), name="evaluate_kernel_B") # We only need one of these since we assume that the two meshes both have CG1 coordinates to_reference_kernel = to_reference_coordinates(mesh_A.coordinates.ufl_element()) if dim == 2: reference_mesh = UnitTriangleMesh(comm=COMM_SELF) else: reference_mesh = UnitTetrahedronMesh(comm=COMM_SELF) evaluate_kernel_S = compile_element(ufl.Coefficient(reference_mesh.coordinates.function_space()), name="evaluate_kernel_S") V_S_A = FunctionSpace(reference_mesh, V_A.ufl_element()) V_S_B = FunctionSpace(reference_mesh, V_B.ufl_element()) M_SS = assemble(inner(TrialFunction(V_S_A), TestFunction(V_S_B)) * dx) M_SS.force_evaluation() M_SS = M_SS.M.handle[:, :] node_locations_A = utils.physical_node_locations(V_S_A).dat.data_ro_with_halos node_locations_B = utils.physical_node_locations(V_S_B).dat.data_ro_with_halos num_nodes_A = node_locations_A.shape[0] num_nodes_B = node_locations_B.shape[0] to_reference_kernel = to_reference_coordinates(mesh_A.coordinates.ufl_element()) supermesh_kernel_str = """ #include "libsupermesh-c.h" #include <petsc.h> %(to_reference)s %(evaluate_S)s %(evaluate_A)s %(evaluate_B)s #define PrintInfo(...) do { if (PetscLogPrintInfo) printf(__VA_ARGS__); } while (0) static void print_array(double *arr, int d) { for(int j=0; j<d; j++) PrintInfo("%%+.2f ", arr[j]); } static void print_coordinates(double *simplex, int d) { for(int i=0; i<d+1; i++) { PrintInfo("\t"); print_array(&simplex[d*i], d); PrintInfo("\\n"); } } int supermesh_kernel(double* simplex_A, double* simplex_B, double* simplices_C, double* nodes_A, double* nodes_B, double* M_SS, double* outptr) { #define d %(dim)s #define num_nodes_A %(num_nodes_A)s #define num_nodes_B %(num_nodes_B)s double simplex_ref_measure; PrintInfo("simplex_A coordinates\\n"); print_coordinates(simplex_A, d); PrintInfo("simplex_B coordinates\\n"); print_coordinates(simplex_B, d); int num_elements; if (d == 2) simplex_ref_measure = 0.5; else if (d == 3) simplex_ref_measure = 1.0/6; double R_AS[num_nodes_A][num_nodes_A]; double R_BS[num_nodes_B][num_nodes_B]; double coeffs_A[%(num_nodes_A)s] = {0.}; double coeffs_B[%(num_nodes_B)s] = {0.}; double reference_nodes_A[num_nodes_A][d]; double reference_nodes_B[num_nodes_B][d]; %(libsupermesh_intersect_simplices)s(simplex_A, simplex_B, simplices_C, &num_elements); PrintInfo("Supermesh consists of %%i elements\\n", num_elements); // would like to do this //double MAB[%(num_nodes_A)s][%(num_nodes_B)s] = (double (*)[%(num_nodes_B)s])outptr; // but have to do this instead because we don't grok C double (*MAB)[num_nodes_A] = (double (*)[num_nodes_A])outptr; double (*MSS)[num_nodes_A] = (double (*)[num_nodes_A])M_SS; // note the underscore for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { MAB[i][j] = 0; } } for(int s=0; s<num_elements; s++) { double* simplex_S = &simplices_C[s * d * (d+1)]; double simplex_S_measure; %(libsupermesh_simplex_measure)s(simplex_S, &simplex_S_measure); PrintInfo("simplex_S coordinates with measure %%f\\n", simplex_S_measure); print_coordinates(simplex_S, d); PrintInfo("Start mapping nodes for V_A\\n"); double physical_nodes_A[num_nodes_A][d]; for(int n=0; n < num_nodes_A; n++) { double* reference_node_location = &nodes_A[n*d]; double* physical_node_location = physical_nodes_A[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("Start mapping nodes for V_B\\n"); double physical_nodes_B[num_nodes_B][d]; for(int n=0; n < num_nodes_B; n++) { double* reference_node_location = &nodes_B[n*d]; double* physical_node_location = physical_nodes_B[n]; for (int j=0; j < d; j++) physical_node_location[j] = 0.0; pyop2_kernel_evaluate_kernel_S(physical_node_location, simplex_S, reference_node_location); PrintInfo("\\tNode "); print_array(reference_node_location, d); PrintInfo(" mapped to "); print_array(physical_node_location, d); PrintInfo("\\n"); } PrintInfo("==========================================================\\n"); PrintInfo("Start pulling back dof from S into reference space for A.\\n"); for(int n=0; n < num_nodes_A; n++) { for(int i=0; i<d; i++) reference_nodes_A[n][i] = 0.; to_reference_coords_kernel(reference_nodes_A[n], physical_nodes_A[n], simplex_A); PrintInfo("Pulling back "); print_array(physical_nodes_A[n], d); PrintInfo(" to "); print_array(reference_nodes_A[n], d); PrintInfo("\\n"); } PrintInfo("Start pulling back dof from S into reference space for B.\\n"); for(int n=0; n < num_nodes_B; n++) { for(int i=0; i<d; i++) reference_nodes_B[n][i] = 0.; to_reference_coords_kernel(reference_nodes_B[n], physical_nodes_B[n], simplex_B); PrintInfo("Pulling back "); print_array(physical_nodes_B[n], d); PrintInfo(" to "); print_array(reference_nodes_B[n], d); PrintInfo("\\n"); } PrintInfo("Start evaluating basis functions of V_A at dofs for V_A on S\\n"); for(int i=0; i<num_nodes_A; i++) { coeffs_A[i] = 1.; for(int j=0; j<num_nodes_A; j++) { R_AS[i][j] = 0.; pyop2_kernel_evaluate_kernel_A(&R_AS[i][j], coeffs_A, reference_nodes_A[j]); } print_array(R_AS[i], num_nodes_A); PrintInfo("\\n"); coeffs_A[i] = 0.; } PrintInfo("Start evaluating basis functions of V_B at dofs for V_B on S\\n"); for(int i=0; i<num_nodes_B; i++) { coeffs_B[i] = 1.; for(int j=0; j<num_nodes_B; j++) { R_BS[i][j] = 0.; pyop2_kernel_evaluate_kernel_B(&R_BS[i][j], coeffs_B, reference_nodes_B[j]); } print_array(R_BS[i], num_nodes_B); PrintInfo("\\n"); coeffs_B[i] = 0.; } PrintInfo("Start doing the matmatmat mult\\n"); for ( int i = 0; i < num_nodes_B; i++ ) { for (int j = 0; j < num_nodes_A; j++) { for ( int k = 0; k < num_nodes_B; k++) { for ( int l = 0; l < num_nodes_A; l++) { MAB[i][j] += (simplex_S_measure/simplex_ref_measure) * R_BS[i][k] * MSS[k][l] * R_AS[j][l]; } } } } } return num_elements; } """ % { "evaluate_S": str(evaluate_kernel_S), "evaluate_A": str(evaluate_kernel_A), "evaluate_B": str(evaluate_kernel_B), "to_reference": str(to_reference_kernel), "num_nodes_A": num_nodes_A, "num_nodes_B": num_nodes_B, "value_size_A": V_A.value_size, "value_size_B": V_B.value_size, "libsupermesh_simplex_measure": "libsupermesh_triangle_area" if dim == 2 else "libsupermesh_tetrahedron_volume", "libsupermesh_intersect_simplices": "libsupermesh_intersect_tris_real" if dim == 2 else "libsupermesh_intersect_tets_real", "dim": dim } dirs = get_petsc_dir() + (os.environ["VIRTUAL_ENV"], ) includes = ["-I%s/include" % d for d in dirs] libs = ["-L%s/lib" % d for d in dirs] libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"] lib = load(supermesh_kernel_str, "c", "supermesh_kernel", cppargs=includes, ldargs=libs, argtypes=[ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp], restype=ctypes.c_int) ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat) if orig_value_size == 1: return mat else: (lrows, grows), (lcols, gcols) = mat.getSizes() lrows *= orig_value_size grows *= orig_value_size lcols *= orig_value_size gcols *= orig_value_size size = ((lrows, grows), (lcols, gcols)) context = BlockMatrix(mat, orig_value_size) blockmat = PETSc.Mat().createPython(size, context=context, comm=mat.comm) blockmat.setUp() return blockmat
def restrict_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vc.ufl_domain().coordinates key = (("restrict", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(firedrake.TestFunction(Vc), Vf) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) element = create_element(Vc.ufl_element()) eval_args = evaluate_kernel.args[:-1] args = eval_args[-1].gencode(not_scope=True) R, fine = (a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_restrict(double *R, %(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xci = Xc + i*%(Xc_cell_inc)d; double celldist = 2*bestdist; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; break; } %(compute_celldist)s /* fprintf(stderr, "cell %%d celldist: %%.14e\\n", i, celldist); fprintf(stderr, "Xref: %%.14e %%.14e %%.14e\\n", Xref[0], Xref[1], Xref[2]); */ if (celldist < bestdist) { bestdist = celldist; bestcell = i; } } if (cell == -1) { /* We didn't find a cell that contained this point exactly. Did we find one that was close enough? */ if (bestdist < 10) { cell = bestcell; } else { fprintf(stderr, "Could not identify cell in transfer operator. Point: "); for (int coord = 0; coord < %(spacedim)s; coord++) { fprintf(stderr, "%%.14e ", X[coord]); } fprintf(stderr, "\\n"); fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist); abort(); } } { const double *Ri = %(R)s + cell*%(coarse_cell_inc)d; pyop2_kernel_evaluate(Ri, %(fine)s, Xref); } } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1], "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "compute_celldist": compute_celldist(element.cell, X="Xref", celldist="celldist"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "args": args, "spacedim": element.cell.get_spatial_dimension(), "R": R, "fine": fine, "tdim": mesh.topological_dimension() } return cache.setdefault( key, op2.Kernel(my_kernel, name="pyop2_kernel_restrict"))
def inject_kernel(Vf, Vc): hierarchy, level = utils.get_level(Vc.ufl_domain()) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = Vf.ufl_domain().coordinates key = ( ("inject", ) + Vf.ufl_element().value_shape() + entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + entity_dofs_key( Vc.mesh().coordinates.function_space().finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: ncandidate = hierarchy.coarse_to_fine_cells[level].shape[1] if Vc.finat_element.entity_dofs( ) == Vc.finat_element.entity_closure_dofs(): return cache.setdefault( key, (dg_injection_kernel(Vf, Vc, ncandidate), True)) coordinates = Vf.ufl_domain().coordinates evaluate_kernel = compile_element(ufl.Coefficient(Vf)) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) coords_element = create_element(coordinates.ufl_element()) Vf_element = create_element(Vf.ufl_element()) kernel = """ %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_inject(double *R, const double *X, const double *f, const double *Xf) { double Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xfi = Xf + i*%(Xf_cell_inc)d; double celldist = 2*bestdist; to_reference_coords_kernel(Xref, X, Xfi); if (%(inside_cell)s) { cell = i; break; } %(compute_celldist)s if (celldist < bestdist) { bestdist = celldist; bestcell = i; } } if (cell == -1) { /* We didn't find a cell that contained this point exactly. Did we find one that was close enough? */ if (bestdist < 10) { cell = bestcell; } else { fprintf(stderr, "Could not identify cell in transfer operator. Point: "); for (int coord = 0; coord < %(spacedim)s; coord++) { fprintf(stderr, "%%.14e ", X[coord]); } fprintf(stderr, "\\n"); fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist); abort(); } } const double *fi = f + cell*%(f_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { R[i] = 0; } pyop2_kernel_evaluate(R, fi, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"), "spacedim": Vc.finat_element.cell.get_spatial_dimension(), "compute_celldist": compute_celldist( Vc.finat_element.cell, X="Xref", celldist="celldist"), "tdim": Vc.ufl_domain().topological_dimension(), "ncandidate": ncandidate, "Rdim": numpy.prod(Vf_element.value_shape), "Xf_cell_inc": coords_element.space_dimension(), "f_cell_inc": Vf_element.space_dimension() } return cache.setdefault( key, (op2.Kernel(kernel, name="pyop2_kernel_inject"), False))
def prolong_kernel(expression): hierarchy, level = utils.get_level(expression.ufl_domain()) levelf = level + Fraction(1 / hierarchy.refinements_per_level) cache = hierarchy._shared_data_cache["transfer_kernels"] coordinates = expression.ufl_domain().coordinates key = (("prolong", ) + expression.ufl_element().value_shape() + entity_dofs_key( expression.function_space().finat_element.entity_dofs()) + entity_dofs_key( coordinates.function_space().finat_element.entity_dofs())) try: return cache[key] except KeyError: mesh = coordinates.ufl_domain() evaluate_kernel = compile_element(expression) to_reference_kernel = to_reference_coordinates( coordinates.ufl_element()) element = create_element(expression.ufl_element()) eval_args = evaluate_kernel.args[:-1] coords_element = create_element(coordinates.ufl_element()) args = eval_args[-1].gencode(not_scope=True) R, coarse = (a.sym.symbol for a in eval_args) my_kernel = """ %(to_reference)s %(evaluate)s void prolong_kernel(double *R, %(args)s, const double *X, const double *Xc) { double Xref[%(tdim)d]; int cell = -1; for (int i = 0; i < %(ncandidate)d; i++) { const double *Xci = Xc + i*%(Xc_cell_inc)d; to_reference_coords_kernel(Xref, X, Xci); if (%(inside_cell)s) { cell = i; break; } } if (cell == -1) abort(); const double *coarsei = %(coarse)s + cell*%(coarse_cell_inc)d; for ( int i = 0; i < %(Rdim)d; i++ ) { %(R)s[i] = 0; } evaluate_kernel(%(R)s, coarsei, Xref); } """ % { "to_reference": str(to_reference_kernel), "evaluate": str(evaluate_kernel), "args": args, "R": R, "coarse": coarse, "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1], "Rdim": numpy.prod(element.value_shape), "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"), "Xc_cell_inc": coords_element.space_dimension(), "coarse_cell_inc": element.space_dimension(), "tdim": mesh.topological_dimension() } return cache.setdefault(key, op2.Kernel(my_kernel, name="prolong_kernel"))