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 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 get_restriction_weights(coarse, fine): mesh = coarse.mesh() assert hasattr(mesh, "_shared_data_cache") cache = mesh._shared_data_cache["hierarchy_restriction_weights"] key = entity_dofs_key(coarse.fiat_element.entity_dofs()) try: return cache[key] except KeyError: # We hit each fine dof more than once since we loop # elementwise over the coarse cells. So we need a count of # how many times we did this to weight the final contribution # appropriately. if not (coarse.ufl_element() == fine.ufl_element()): raise ValueError("Can't transfer between different spaces") if coarse.fiat_element.entity_dofs() == coarse.fiat_element.entity_closure_dofs(): return cache.setdefault(key, None) ele = coarse.ufl_element() if isinstance(ele, ufl.VectorElement): ele = ele.sub_elements()[0] weights = firedrake.Function(firedrake.FunctionSpace(fine.mesh(), ele)) else: weights = firedrake.Function(fine) c2f_map = coarse_to_fine_node_map(coarse, fine) kernel = get_count_kernel(c2f_map.arity) op2.par_loop(kernel, op2.LocalSet(mesh.cell_set), weights.dat(op2.INC, c2f_map[op2.i[0]])) weights.assign(1/weights) return cache.setdefault(key, weights)
def coarse_to_fine_node_map(coarse, fine): if len(coarse) > 1: assert len(fine) == len(coarse) return op2.MixedMap( coarse_to_fine_node_map(c, f) for c, f in zip(coarse, fine)) mesh = coarse.mesh() assert hasattr(mesh, "_shared_data_cache") if not (coarse.ufl_element() == fine.ufl_element()): raise ValueError("Can't transfer between different spaces") ch, level = get_level(mesh) fh, fine_level = get_level(fine.mesh()) if ch is not fh: raise ValueError("Can't map between different hierarchies") refinements_per_level = ch.refinements_per_level if refinements_per_level * level + 1 != refinements_per_level * fine_level: raise ValueError("Can't map between level %s and level %s" % (level, fine_level)) c2f, vperm = ch._cells_vperm[int(level * refinements_per_level)] key = entity_dofs_key(coarse.finat_element.entity_dofs()) + (level, ) cache = mesh._shared_data_cache["hierarchy_cell_node_map"] try: return cache[key] except KeyError: from .impl import create_cell_node_map map_vals, offset = create_cell_node_map(coarse, fine, c2f, vperm) return cache.setdefault( key, op2.Map(mesh.cell_set, fine.node_set, map_vals.shape[1], map_vals, offset=offset))
def get_restriction_weights(coarse, fine): mesh = coarse.mesh() assert hasattr(mesh, "_shared_data_cache") cache = mesh._shared_data_cache["hierarchy_restriction_weights"] key = entity_dofs_key(coarse.finat_element.entity_dofs()) try: return cache[key] except KeyError: # We hit each fine dof more than once since we loop # elementwise over the coarse cells. So we need a count of # how many times we did this to weight the final contribution # appropriately. if not (coarse.ufl_element() == fine.ufl_element()): raise ValueError("Can't transfer between different spaces") if coarse.finat_element.entity_dofs( ) == coarse.finat_element.entity_closure_dofs(): return cache.setdefault(key, None) ele = coarse.ufl_element() if isinstance(ele, ufl.VectorElement): ele = ele.sub_elements()[0] weights = firedrake.Function( firedrake.FunctionSpace(fine.mesh(), ele)) else: weights = firedrake.Function(fine) c2f_map = coarse_to_fine_node_map(coarse, fine) kernel = get_count_kernel(c2f_map.arity) op2.par_loop(kernel, c2f_map.iterset, weights.dat(op2.INC, c2f_map[op2.i[0]])) weights.assign(1 / weights) return cache.setdefault(key, weights)
def get_transfer_kernel(coarse, fine, typ=None): ch, level = get_level(coarse.mesh()) mesh = ch[0] assert hasattr(mesh, "_shared_data_cache") key = entity_dofs_key(coarse.fiat_element.entity_dofs()) + (typ, ) cache = mesh._shared_data_cache["hierarchy_transfer_kernel"] try: return cache[key] except KeyError: if not (coarse.ufl_element() == fine.ufl_element()): raise ValueError("Can't transfer between different spaces") ch, level = get_level(coarse.mesh()) fh, fine_level = get_level(fine.mesh()) refinements_per_level = ch.refinements_per_level assert ch is fh assert refinements_per_level*level + 1 == refinements_per_level*fine_level dim = coarse.dim element = coarse.fiat_element omap = fine.cell_node_map().values c2f, vperm = ch._cells_vperm[int(level*refinements_per_level)] indices, _ = get_unique_indices(element, omap[c2f[0, :], ...].reshape(-1), vperm[0, :], offset=None) if typ == "prolong": kernel = get_prolongation_kernel(element, indices, dim) elif typ == "inject": kernel = get_injection_kernel(element, indices, dim) elif typ == "restrict": discontinuous = element.entity_dofs() == element.entity_closure_dofs() kernel = get_restriction_kernel(element, indices, dim, no_weights=discontinuous) else: raise ValueError("Unknown transfer type '%s'" % typ) return cache.setdefault(key, kernel)
def coarse_to_fine_node_map(coarse, fine): if len(coarse) > 1: assert len(fine) == len(coarse) return op2.MixedMap(coarse_to_fine_node_map(c, f) for c, f in zip(coarse, fine)) mesh = coarse.mesh() assert hasattr(mesh, "_shared_data_cache") if not (coarse.ufl_element() == fine.ufl_element()): raise ValueError("Can't transfer between different spaces") ch, level = get_level(mesh) fh, fine_level = get_level(fine.mesh()) if ch is not fh: raise ValueError("Can't map between different hierarchies") refinements_per_level = ch.refinements_per_level if refinements_per_level*level + 1 != refinements_per_level*fine_level: raise ValueError("Can't map between level %s and level %s" % (level, fine_level)) c2f, vperm = ch._cells_vperm[int(level*refinements_per_level)] key = entity_dofs_key(coarse.fiat_element.entity_dofs()) + (level, ) cache = mesh._shared_data_cache["hierarchy_cell_node_map"] try: return cache[key] except KeyError: from .impl import create_cell_node_map map_vals, offset = create_cell_node_map(coarse, fine, c2f, vperm) return cache.setdefault(key, op2.Map(op2.LocalSet(mesh.cell_set), fine.node_set, map_vals.shape[1], map_vals, offset=offset))
def coarse_node_to_fine_node_map(Vc, Vf): if len(Vf) > 1: assert len(Vf) == len(Vc) return op2.MixedMap( coarse_node_to_fine_node_map(f, c) for f, c in zip(Vf, Vc)) mesh = Vc.mesh() assert hasattr(mesh, "_shared_data_cache") hierarchyf, levelf = get_level(Vf.ufl_domain()) hierarchyc, levelc = get_level(Vc.ufl_domain()) if hierarchyc != hierarchyf: raise ValueError("Can't map across hierarchies") hierarchy = hierarchyf increment = Fraction(1, hierarchyf.refinements_per_level) if levelc + increment != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) key = (entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf)) cache = mesh._shared_data_cache["hierarchy_coarse_node_to_fine_node_map"] try: return cache[key] except KeyError: assert Vc.extruded == Vf.extruded if Vc.mesh().variable_layers or Vf.mesh().variable_layers: raise NotImplementedError( "Not implemented for variable layers, sorry") if Vc.extruded and not ((Vf.mesh().layers - 1) / (Vc.mesh().layers - 1)).is_integer(): raise ValueError( "Coarse and fine meshes must have an integer ratio of layers") coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc] coarse_to_fine_nodes = impl.coarse_to_fine_nodes( Vc, Vf, coarse_to_fine) return cache.setdefault( key, op2.Map(Vc.node_set, Vf.node_set, coarse_to_fine_nodes.shape[1], values=coarse_to_fine_nodes))
def fine_node_to_coarse_node_map(Vf, Vc): if len(Vf) > 1: assert len(Vf) == len(Vc) return op2.MixedMap( fine_node_to_coarse_node_map(f, c) for f, c in zip(Vf, Vc)) mesh = Vf.mesh() assert hasattr(mesh, "_shared_data_cache") hierarchyf, levelf = get_level(Vf.ufl_domain()) hierarchyc, levelc = get_level(Vc.ufl_domain()) if hierarchyc != hierarchyf: raise ValueError("Can't map across hierarchies") hierarchy = hierarchyf if levelc + 1 != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) key = (entity_dofs_key(Vc.finat_element.entity_dofs()) + entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf)) cache = mesh._shared_data_cache["hierarchy_fine_node_to_coarse_node_map"] try: return cache[key] except KeyError: assert Vc.extruded == Vf.extruded if Vc.mesh().variable_layers or Vf.mesh().variable_layers: raise NotImplementedError( "Not implemented for variable layers, sorry") if Vc.extruded and Vc.mesh().layers != Vf.mesh().layers: raise ValueError( "Coarse and fine meshes must have same number of layers") fine_to_coarse = hierarchy.fine_to_coarse_cells[levelc + 1] fine_to_coarse_nodes = impl.fine_to_coarse_nodes( Vf, Vc, fine_to_coarse) return cache.setdefault( key, op2.Map(Vf.node_set, Vc.node_set, fine_to_coarse_nodes.shape[1], values=fine_to_coarse_nodes))
def coarse_cell_to_fine_node_map(Vc, Vf): if len(Vf) > 1: assert len(Vf) == len(Vc) return op2.MixedMap( coarse_cell_to_fine_node_map(f, c) for f, c in zip(Vf, Vc)) mesh = Vc.mesh() assert hasattr(mesh, "_shared_data_cache") hierarchyf, levelf = get_level(Vf.ufl_domain()) hierarchyc, levelc = get_level(Vc.ufl_domain()) if hierarchyc != hierarchyf: raise ValueError("Can't map across hierarchies") hierarchy = hierarchyf increment = Fraction(1, hierarchyf.refinements_per_level) if levelc + increment != levelf: raise ValueError("Can't map between level %s and level %s" % (levelc, levelf)) key = (entity_dofs_key(Vf.finat_element.entity_dofs()) + (levelc, levelf)) cache = mesh._shared_data_cache["hierarchy_coarse_cell_to_fine_node_map"] try: return cache[key] except KeyError: assert Vc.extruded == Vf.extruded if Vc.mesh().variable_layers or Vf.mesh().variable_layers: raise NotImplementedError( "Not implemented for variable layers, sorry") if Vc.extruded and Vc.mesh().layers != Vf.mesh().layers: raise ValueError( "Coarse and fine meshes must have same number of layers") coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc] _, ncell = coarse_to_fine.shape iterset = Vc.mesh().cell_set arity = Vf.finat_element.space_dimension() * ncell coarse_to_fine_nodes = numpy.full((iterset.total_size, arity), -1, dtype=IntType) values = Vf.cell_node_map().values[coarse_to_fine, :].reshape( iterset.size, arity) coarse_to_fine_nodes[:Vc.mesh().cell_set.size, :] = values offset = Vf.offset if offset is not None: offset = numpy.tile(offset, ncell) return cache.setdefault( key, op2.Map(iterset, Vf.node_set, arity=arity, values=coarse_to_fine_nodes, offset=offset))
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 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 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 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 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"))
def prolong_kernel(expression): meshc = expression.ufl_domain() 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 if meshc.cell_set._extruded: idx = levelf * hierarchy.refinements_per_level assert idx == int(idx) level_ratio = (hierarchy._meshes[int(idx)].layers - 1) // (meshc.layers - 1) else: level_ratio = 1 key = (("prolong", level_ratio) + 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 = """#include <petsc.h> %(to_reference)s %(evaluate)s __attribute__((noinline)) /* Clang bug */ static void pyop2_kernel_prolong(PetscScalar *R, %(args)s, const PetscScalar *X, const PetscScalar *Xc) { PetscScalar Xref[%(tdim)d]; int cell = -1; int bestcell = -1; double bestdist = 1e10; for (int i = 0; i < %(ncandidate)d; i++) { const PetscScalar *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 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 PetscScalar *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, "spacedim": element.cell.get_spatial_dimension(), "coarse": coarse, "ncandidate": hierarchy.fine_to_coarse_cells[levelf].shape[1] * level_ratio, "Rdim": numpy.prod(element.value_shape), "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(), "tdim": mesh.topological_dimension() } return cache.setdefault( key, op2.Kernel(my_kernel, name="pyop2_kernel_prolong"))