From e4f9dc790a250c8a4b46f5dfff054eaff93c5867 Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Wed, 4 Sep 2024 00:57:41 +0100 Subject: [PATCH 1/9] DROP BEFORE MERGE --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0fff428f43..cdf1857304 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,6 +83,8 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch tsfc ksagiyam/introduce_mixed_map \ + --package-branch ufl ksagiyam/introduce_mixed_map \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From 90d50ba709d9d2ad27a16c1c531973a41b8b950a Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Fri, 12 Apr 2024 02:34:36 +0100 Subject: [PATCH 2/9] introduce MixedMesh --- .../saddle_point_systems.py.rst | 12 +- firedrake/assemble.py | 173 +++++-- firedrake/checkpointing.py | 12 +- firedrake/dmhooks.py | 17 +- firedrake/function.py | 10 +- firedrake/functionspace.py | 21 +- firedrake/functionspaceimpl.py | 79 ++- firedrake/mesh.py | 199 ++++++++ firedrake/mg/kernels.py | 10 +- firedrake/mg/ufl_utils.py | 5 +- firedrake/pointeval_utils.py | 1 + firedrake/pointquery_utils.py | 2 +- firedrake/preconditioners/asm.py | 34 +- firedrake/preconditioners/fdm.py | 2 +- firedrake/preconditioners/patch.py | 142 +++--- firedrake/preconditioners/pmg.py | 12 +- firedrake/slate/slac/compiler.py | 8 +- firedrake/slate/slac/kernel_builder.py | 39 +- firedrake/slate/slate.py | 15 +- .../static_condensation/hybridization.py | 16 +- firedrake/slate/static_condensation/scpc.py | 2 +- firedrake/tsfc_interface.py | 53 +- firedrake/ufl_expr.py | 85 ++-- tests/regression/test_function_spaces.py | 4 +- tests/submesh/test_submesh_assemble.py | 258 ++++++++++ tests/submesh/test_submesh_base.py | 289 +++++++++++ tests/submesh/test_submesh_solve.py | 459 ++++++++++++++++++ 27 files changed, 1699 insertions(+), 260 deletions(-) create mode 100644 tests/submesh/test_submesh_assemble.py create mode 100644 tests/submesh/test_submesh_base.py create mode 100644 tests/submesh/test_submesh_solve.py diff --git a/demos/saddle_point_pc/saddle_point_systems.py.rst b/demos/saddle_point_pc/saddle_point_systems.py.rst index 60d7133ebd..ac023687b6 100644 --- a/demos/saddle_point_pc/saddle_point_systems.py.rst +++ b/demos/saddle_point_pc/saddle_point_systems.py.rst @@ -187,7 +187,7 @@ Finally, at each mesh size, we print out the number of cells in the mesh and the number of iterations the solver took to converge :: # - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) The resulting convergence is unimpressive: @@ -289,7 +289,7 @@ applying the action of blocks, so we can use a block matrix format. :: for n in range(8): solver, w = build_problem(n, parameters, block_matrix=True) solver.solve() - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) The resulting convergence is algorithmically good, however, the larger problems still take a long time. @@ -374,7 +374,7 @@ Let's see what happens. :: for n in range(8): solver, w = build_problem(n, parameters, block_matrix=True) solver.solve() - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) This is much better, the problem takes much less time to solve and when observing the iteration counts for inverting :math:`S` we can see @@ -429,7 +429,7 @@ and so we no longer need a flexible Krylov method. :: for n in range(8): solver, w = build_problem(n, parameters, block_matrix=True) solver.solve() - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) This results in the following GMRES iteration counts @@ -494,7 +494,7 @@ variable. We can provide it as an :class:`~.AuxiliaryOperatorPC` via a python pr for n in range(8): solver, w = build_problem(n, parameters, aP=None, block_matrix=False) solver.solve() - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) This actually results in slightly worse convergence than the diagonal approximation we used above. @@ -578,7 +578,7 @@ Let's see what the iteration count looks like now. :: for n in range(8): solver, w = build_problem(n, parameters, aP=riesz, block_matrix=True) solver.solve() - print(w.function_space().mesh().num_cells(), solver.snes.ksp.getIterationNumber()) + print(w.function_space().mesh().unique().num_cells(), solver.snes.ksp.getIterationNumber()) ============== ================== Mesh elements GMRES iterations diff --git a/firedrake/assemble.py b/firedrake/assemble.py index d3f4cfb8ba..3bdcbc97d6 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -19,7 +19,7 @@ from firedrake import (extrusion_utils as eutils, matrix, parameters, solving, tsfc_interface, utils) from firedrake.adjoint_utils import annotate_assemble -from firedrake.ufl_expr import extract_unique_domain +from firedrake.ufl_expr import extract_domains from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace from firedrake.functionspacedata import entity_dofs_key, entity_permutations_key @@ -1018,7 +1018,7 @@ def parloops(self, tensor): self._bcs, local_kernel, subdomain_id, - self.all_integer_subdomain_ids, + self.all_integer_subdomain_ids[local_kernel.kinfo.domain_number], diagonal=self.diagonal, ) pyop2_tensor = self._as_pyop2_type(tensor, local_kernel.indices) @@ -1040,14 +1040,15 @@ def local_kernels(self): """ try: - topology, = set(d.topology for d in self._form.ufl_domains()) + topology, = set(d.topology.submesh_ancesters[-1] for d in self._form.ufl_domains()) except ValueError: raise NotImplementedError("All integration domains must share a mesh topology") for o in itertools.chain(self._form.arguments(), self._form.coefficients()): - domain = extract_unique_domain(o) - if domain is not None and domain.topology != topology: - raise NotImplementedError("Assembly with multiple meshes is not supported") + domains = extract_domains(o) + for domain in domains: + if domain is not None and domain.topology.submesh_ancesters[-1] != topology: + raise NotImplementedError("Assembly with multiple meshes is not supported") if isinstance(self._form, ufl.Form): kernels = tsfc_interface.compile_form( @@ -1351,12 +1352,12 @@ def _make_maps_and_regions(self): else: maps_and_regions = defaultdict(lambda: defaultdict(set)) for assembler in self._all_assemblers: - all_meshes = assembler._form.ufl_domains() + all_meshes = extract_domains(assembler._form) for local_kernel, subdomain_id in assembler.local_kernels: i, j = local_kernel.indices mesh = all_meshes[local_kernel.kinfo.domain_number] # integration domain integral_type = local_kernel.kinfo.integral_type - all_subdomain_ids = assembler.all_integer_subdomain_ids + all_subdomain_ids = assembler.all_integer_subdomain_ids[local_kernel.kinfo.domain_number] # Make Sparsity independent of the subdomain of integration for better reusability; # subdomain_id is passed here only to determine the integration_type on the target domain # (see ``entity_node_map``). @@ -1531,6 +1532,10 @@ def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomai # N.B. Generating the global kernel is not a collective operation so the # communicator does not need to be a part of this cache key. + # Maps in the cached global kernel depend on concrete mesh data. + all_meshes = extract_domains(form) + domain_ids = tuple(mesh.ufl_id() for mesh in all_meshes) + if isinstance(form, ufl.Form): sig = form.signature() elif isinstance(form, slate.TensorBase): @@ -1550,7 +1555,8 @@ def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomai else: subdomain_key.append((k, i)) - return ((sig, subdomain_id) + return (domain_ids + + (sig, subdomain_id) + tuple(subdomain_key) + tuplify(all_integer_subdomain_ids) + cachetools.keys.hashkey(local_knl, **kwargs)) @@ -1587,8 +1593,13 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia self._diagonal = diagonal self._unroll = unroll + self._active_coordinates = _FormHandler.iter_active_coordinates(form, local_knl.kinfo) + self._active_cell_orientations = _FormHandler.iter_active_cell_orientations(form, local_knl.kinfo) + self._active_cell_sizes = _FormHandler.iter_active_cell_sizes(form, local_knl.kinfo) self._active_coefficients = _FormHandler.iter_active_coefficients(form, local_knl.kinfo) self._constants = _FormHandler.iter_constants(form, local_knl.kinfo) + self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo) + self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo) self._map_arg_cache = {} # Cache for holding :class:`op2.MapKernelArg` instances. @@ -1602,8 +1613,13 @@ def build(self): for arg in self._kinfo.arguments] # we should use up all of the coefficients and constants + assert_empty(self._active_coordinates) + assert_empty(self._active_cell_orientations) + assert_empty(self._active_cell_sizes) assert_empty(self._active_coefficients) assert_empty(self._constants) + assert_empty(self._active_exterior_facets) + assert_empty(self._active_interior_facets) iteration_regions = {"exterior_facet_top": op2.ON_TOP, "exterior_facet_bottom": op2.ON_BOTTOM, @@ -1628,7 +1644,8 @@ def _integral_type(self): @cached_property def _mesh(self): - return self._form.ufl_domains()[self._kinfo.domain_number] + all_meshes = extract_domains(self._form) + return all_meshes[self._kinfo.domain_number] @cached_property def _needs_subset(self): @@ -1733,7 +1750,22 @@ def _as_global_kernel_arg_output(_, self): @_as_global_kernel_arg.register(kernel_args.CoordinatesKernelArg) def _as_global_kernel_arg_coordinates(_, self): - V = self._mesh.coordinates.function_space() + coord = next(self._active_coordinates) + V = coord.function_space() + return self._make_dat_global_kernel_arg(V) + + +@_as_global_kernel_arg.register(kernel_args.CellOrientationsKernelArg) +def _as_global_kernel_arg_cell_orientations(_, self): + c = next(self._active_cell_orientations) + V = c.function_space() + return self._make_dat_global_kernel_arg(V) + + +@_as_global_kernel_arg.register(kernel_args.CellSizesKernelArg) +def _as_global_kernel_arg_cell_sizes(_, self): + c = next(self._active_cell_sizes) + V = c.function_space() return self._make_dat_global_kernel_arg(V) @@ -1761,20 +1793,26 @@ def _as_global_kernel_arg_constant(_, self): return op2.GlobalKernelArg((value_size,)) -@_as_global_kernel_arg.register(kernel_args.CellSizesKernelArg) -def _as_global_kernel_arg_cell_sizes(_, self): - V = self._mesh.cell_sizes.function_space() - return self._make_dat_global_kernel_arg(V) - - @_as_global_kernel_arg.register(kernel_args.ExteriorFacetKernelArg) def _as_global_kernel_arg_exterior_facet(_, self): - return op2.DatKernelArg((1,)) + mesh = next(self._active_exterior_facets) + if mesh is self._mesh: + return op2.DatKernelArg((1,)) + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "exterior_facet" + return op2.DatKernelArg((1,), m._global_kernel_arg) @_as_global_kernel_arg.register(kernel_args.InteriorFacetKernelArg) def _as_global_kernel_arg_interior_facet(_, self): - return op2.DatKernelArg((2,)) + mesh = next(self._active_interior_facets) + if mesh is self._mesh: + return op2.DatKernelArg((2,)) + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "interior_facet" + return op2.DatKernelArg((2,), m._global_kernel_arg) @_as_global_kernel_arg.register(kernel_args.ExteriorFacetOrientationKernelArg) @@ -1796,12 +1834,6 @@ def _as_global_kernel_arg_cell_facet(_, self): return op2.DatKernelArg((num_facets, 2)) -@_as_global_kernel_arg.register(kernel_args.CellOrientationsKernelArg) -def _as_global_kernel_arg_cell_orientations(_, self): - V = self._mesh.cell_orientations().function_space() - return self._make_dat_global_kernel_arg(V) - - @_as_global_kernel_arg.register(LayerCountKernelArg) def _as_global_kernel_arg_layer_count(_, self): return op2.GlobalKernelArg((1,)) @@ -1835,8 +1867,13 @@ def __init__(self, form, bcs, local_knl, subdomain_id, self._diagonal = diagonal self._bcs = bcs + self._active_coordinates = _FormHandler.iter_active_coordinates(form, local_knl.kinfo) + self._active_cell_orientations = _FormHandler.iter_active_cell_orientations(form, local_knl.kinfo) + self._active_cell_sizes = _FormHandler.iter_active_cell_sizes(form, local_knl.kinfo) self._active_coefficients = _FormHandler.iter_active_coefficients(form, local_knl.kinfo) self._constants = _FormHandler.iter_constants(form, local_knl.kinfo) + self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo) + self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo) def build(self, tensor: op2.Global | op2.Dat | op2.Mat) -> op2.Parloop: """Construct the parloop. @@ -1970,7 +2007,8 @@ def _indexed_function_spaces(self): @cached_property def _mesh(self): - return self._form.ufl_domains()[self._kinfo.domain_number] + all_meshes = extract_domains(self._form) + return all_meshes[self._kinfo.domain_number] @cached_property def _iterset(self): @@ -2042,7 +2080,21 @@ def _as_parloop_arg_output(_, self): @_as_parloop_arg.register(kernel_args.CoordinatesKernelArg) def _as_parloop_arg_coordinates(_, self): - func = self._mesh.coordinates + func = next(self._active_coordinates) + map_ = self._get_map(func.function_space()) + return op2.DatParloopArg(func.dat, map_) + + +@_as_parloop_arg.register(kernel_args.CellOrientationsKernelArg) +def _as_parloop_arg_cell_orientations(_, self): + func = next(self._active_cell_orientations) + map_ = self._get_map(func.function_space()) + return op2.DatParloopArg(func.dat, map_) + + +@_as_parloop_arg.register(kernel_args.CellSizesKernelArg) +def _as_parloop_arg_cell_sizes(_, self): + func = next(self._active_cell_sizes) map_ = self._get_map(func.function_space()) return op2.DatParloopArg(func.dat, map_) @@ -2063,28 +2115,26 @@ def _as_parloop_arg_constant(arg, self): return op2.GlobalParloopArg(const.dat) -@_as_parloop_arg.register(kernel_args.CellOrientationsKernelArg) -def _as_parloop_arg_cell_orientations(_, self): - func = self._mesh.cell_orientations() - m = self._get_map(func.function_space()) - return op2.DatParloopArg(func.dat, m) - - -@_as_parloop_arg.register(kernel_args.CellSizesKernelArg) -def _as_parloop_arg_cell_sizes(_, self): - func = self._mesh.cell_sizes - m = self._get_map(func.function_space()) - return op2.DatParloopArg(func.dat, m) - - @_as_parloop_arg.register(kernel_args.ExteriorFacetKernelArg) def _as_parloop_arg_exterior_facet(_, self): - return op2.DatParloopArg(self._mesh.exterior_facets.local_facet_dat) + mesh = next(self._active_exterior_facets) + if mesh is self._mesh: + m = None + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "exterior_facet" + return op2.DatParloopArg(mesh.exterior_facets.local_facet_dat, m) @_as_parloop_arg.register(kernel_args.InteriorFacetKernelArg) def _as_parloop_arg_interior_facet(_, self): - return op2.DatParloopArg(self._mesh.interior_facets.local_facet_dat) + mesh = next(self._active_interior_facets) + if mesh is self._mesh: + m = None + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "interior_facet" + return op2.DatParloopArg(mesh.interior_facets.local_facet_dat, m) @_as_parloop_arg.register(kernel_args.ExteriorFacetOrientationKernelArg) @@ -2116,6 +2166,27 @@ def _as_parloop_arg_layer_count(_, self): class _FormHandler: """Utility class for inspecting forms and local kernels.""" + @staticmethod + def iter_active_coordinates(form, kinfo): + """Yield the form coordinates referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.coordinates: + yield all_meshes[i].coordinates + + @staticmethod + def iter_active_cell_orientations(form, kinfo): + """Yield the form cell orientations referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.cell_orientations: + yield all_meshes[i].cell_orientations() + + @staticmethod + def iter_active_cell_sizes(form, kinfo): + """Yield the form cell sizes referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.cell_sizes: + yield all_meshes[i].cell_sizes + @staticmethod def iter_active_coefficients(form, kinfo): """Yield the form coefficients referenced in ``kinfo``.""" @@ -2135,6 +2206,22 @@ def iter_constants(form, kinfo): for constant_index in kinfo.constant_numbers: yield all_constants[constant_index] + @staticmethod + def iter_active_exterior_facets(form, kinfo): + """Yield the form exterior facets referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.exterior_facets: + mesh = all_meshes[i] + yield mesh + + @staticmethod + def iter_active_interior_facets(form, kinfo): + """Yield the form interior facets referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.interior_facets: + mesh = all_meshes[i] + yield mesh + @staticmethod def index_function_spaces(form, indices): """Return the function spaces of the form's arguments, indexed diff --git a/firedrake/checkpointing.py b/firedrake/checkpointing.py index 5070a96315..affdf92ad6 100644 --- a/firedrake/checkpointing.py +++ b/firedrake/checkpointing.py @@ -565,6 +565,8 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None): :kwarg distribution_name: the name under which distribution is saved; if `None`, auto-generated name will be used. :kwarg permutation_name: the name under which permutation is saved; if `None`, auto-generated name will be used. """ + # TODO: Add general MixedMesh support. + mesh = mesh.unique() mesh.init() # Handle extruded mesh tmesh = mesh.topology @@ -836,6 +838,8 @@ def get_timestepping_history(self, mesh, name): @PETSc.Log.EventDecorator("SaveFunctionSpace") def _save_function_space(self, V): mesh = V.mesh() + # TODO: Add general MixedMesh support. + mesh = mesh.unique() if isinstance(V.topological, impl.MixedFunctionSpace): V_name = self._generate_function_space_name(V) base_path = self._path_to_mixed_function_space(mesh.name, V_name) @@ -911,10 +915,12 @@ def save_function(self, f, idx=None, name=None, timestepping_info={}): each index. """ V = f.function_space() - mesh = V.mesh() if name: g = Function(V, val=f.dat, name=name) return self.save_function(g, idx=idx, timestepping_info=timestepping_info) + mesh = V.mesh() + # TODO: Add general MixedMesh support. + mesh = mesh.unique() # -- Save function space -- self._save_function_space(V) # -- Save function -- @@ -1233,6 +1239,8 @@ def _load_mesh_topology(self, tmesh_name, reorder, distribution_parameters): @PETSc.Log.EventDecorator("LoadFunctionSpace") def _load_function_space(self, mesh, name): + # TODO: Add general MixedMesh support. + mesh = mesh.unique() mesh.init() mesh_key = self._generate_mesh_key_from_names(mesh.name, mesh.topology._distribution_name, @@ -1310,6 +1318,8 @@ def load_function(self, mesh, name, idx=None): be loaded with idx only when it was saved with idx. :returns: the loaded :class:`~.Function`. """ + # TODO: Add general MixedMesh support. + mesh = mesh.unique() tmesh = mesh.topology if name in self._get_mixed_function_name_mixed_function_space_name_map(mesh.name): V_name = self._get_mixed_function_name_mixed_function_space_name_map(mesh.name)[name] diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index c1cca65b56..c3b61bcf08 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -43,6 +43,7 @@ import firedrake from firedrake.petsc import PETSc +from firedrake.mesh import MixedMeshGeometry @PETSc.Log.EventDecorator() @@ -53,8 +54,11 @@ def get_function_space(dm): :raises RuntimeError: if no function space was found. """ info = dm.getAttr("__fs_info__") - meshref, element, indices, (name, names) = info - mesh = meshref() + meshref_tuple, element, indices, (name, names) = info + if len(meshref_tuple) == 1: + mesh = meshref_tuple[0]() + else: + mesh = MixedMeshGeometry(*(meshref() for meshref in meshref_tuple)) if mesh is None: raise RuntimeError("Somehow your mesh was collected, this should never happen") V = firedrake.FunctionSpace(mesh, element, name=name) @@ -78,8 +82,6 @@ def set_function_space(dm, V): This stores the information necessary to make a function space given a DM. """ - mesh = V.mesh() - indices = [] names = [] while V.parent is not None: @@ -90,11 +92,12 @@ def set_function_space(dm, V): assert V.index is None indices.append(V.component) V = V.parent + mesh = V.mesh() if len(V) > 1: names = tuple(V_.name for V_ in V) element = V.ufl_element() - info = (weakref.ref(mesh), element, tuple(reversed(indices)), (V.name, names)) + info = (tuple(weakref.ref(m) for m in mesh), element, tuple(reversed(indices)), (V.name, names)) dm.setAttr("__fs_info__", info) @@ -412,7 +415,9 @@ def coarsen(dm, comm): """ from firedrake.mg.utils import get_level V = get_function_space(dm) - hierarchy, level = get_level(V.mesh()) + # TODO: Think harder. + m, = set(m_ for m_ in V.mesh()) + hierarchy, level = get_level(m) if level < 1: raise RuntimeError("Cannot coarsen coarsest DM") coarsen = get_ctx_coarsener(dm) diff --git a/firedrake/function.py b/firedrake/function.py index 83136407ea..77f9ca5668 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -620,16 +620,20 @@ def at(self, arg, *args, **kwargs): tolerance = kwargs.get('tolerance', None) mesh = self.function_space().mesh() + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") if tolerance is None: - tolerance = mesh.tolerance + tolerance = mesh_unique.tolerance else: - mesh.tolerance = tolerance + mesh_unique.tolerance = tolerance # Handle f.at(0.3) if not arg.shape: arg = arg.reshape(-1) - if mesh.variable_layers: + if mesh_unique.variable_layers: raise NotImplementedError("Point evaluation not implemented for variable layers") # Validate geometric dimension diff --git a/firedrake/functionspace.py b/firedrake/functionspace.py index 74dab04703..8d6ae617b2 100644 --- a/firedrake/functionspace.py +++ b/firedrake/functionspace.py @@ -4,6 +4,7 @@ API is functional, rather than object-based, to allow for simple backwards-compatibility, argument checking, and dispatch. """ +import itertools import ufl import finat.ufl @@ -258,6 +259,8 @@ def MixedFunctionSpace(spaces, name=None, mesh=None): :class:`finat.ufl.mixedelement.MixedElement`, ignored otherwise. """ + from firedrake.mesh import MixedMeshGeometry + if isinstance(spaces, finat.ufl.FiniteElementBase): # Build the spaces if we got a mixed element assert type(spaces) is finat.ufl.MixedElement and mesh is not None @@ -272,13 +275,8 @@ def rec(eles): sub_elements.append(ele) rec(spaces.sub_elements) spaces = [FunctionSpace(mesh, element) for element in sub_elements] - - # Check that function spaces are on the same mesh - meshes = [space.mesh() for space in spaces] - for i in range(1, len(meshes)): - if meshes[i] is not meshes[0]: - raise ValueError("All function spaces must be defined on the same mesh!") - + # Flatten MixedMeshes. + meshes = list(itertools.chain(*[space.mesh() for space in spaces])) try: cls, = set(type(s) for s in spaces) except ValueError: @@ -286,8 +284,6 @@ def rec(eles): # We had not implemented something in between, so let's make it primal cls = impl.WithGeometry - # Select mesh - mesh = meshes[0] # Get topological spaces spaces = tuple(s.topological for s in flatten(spaces)) # Error checking @@ -301,10 +297,9 @@ def rec(eles): else: raise ValueError("Can't make mixed space with %s" % type(space)) - new = impl.MixedFunctionSpace(spaces, name=name) - if mesh is not mesh.topology: - new = cls.create(new, mesh) - return new + mixed_mesh_geometry = MixedMeshGeometry(*meshes) + new = impl.MixedFunctionSpace(spaces, mixed_mesh_geometry.topology, name=name) + return cls.create(new, mixed_mesh_geometry) @PETSc.Log.EventDecorator("CreateFunctionSpace") diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 53f69e92ce..6f16399c62 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -16,6 +16,7 @@ from pyop2 import op2, mpi from firedrake import dmhooks, utils +from firedrake.mesh import MeshGeometry, MixedMeshTopology, MixedMeshGeometry from firedrake.functionspacedata import get_shared_data, create_element from firedrake.petsc import PETSc @@ -92,9 +93,11 @@ class WithGeometryBase(object): generation. """ def __init__(self, mesh, element, component=None, cargo=None): + if type(element) is finat.ufl.MixedElement: + if not isinstance(mesh, MixedMeshGeometry): + raise TypeError(f"Can only use MixedElement with MixedMeshGeometry: got {type(mesh)}") assert component is None or isinstance(component, int) assert cargo is None or isinstance(cargo, FunctionSpaceCargo) - super().__init__(mesh, element, label=cargo.topological._label or "") self.component = component self.cargo = cargo @@ -102,13 +105,22 @@ def __init__(self, mesh, element, component=None, cargo=None): self._comm = mpi.internal_comm(mesh.comm, self) @classmethod - def create(cls, function_space, mesh): + def create(cls, function_space, mesh, parent=None): """Create a :class:`WithGeometry`. - :arg function_space: The topological function space to attach - geometry to. - :arg mesh: The mesh with geometric information to use. + Parameters + ---------- + function_space : FunctionSpace or MixedFunctionSpace + Topological function space to attach geometry to. + mesh : MeshGeometry + Mesh with geometric information to use. + parent : WithGeometry + Parent geometric function space if exists. + """ + if isinstance(function_space, MixedFunctionSpace): + if not isinstance(mesh, MixedMeshGeometry): + raise TypeError(f"Can only use MixedFunctionSpace with MixedMeshGeometry: got {type(mesh)}") function_space = function_space.topological assert mesh.topology is function_space.mesh() assert mesh.topology is not mesh @@ -119,7 +131,8 @@ def create(cls, function_space, mesh): component = function_space.component if function_space.parent is not None: - parent = cls.create(function_space.parent, mesh) + if parent is None: + raise ValueError("Must pass parent if function_space.parent is not None") else: parent = None @@ -149,8 +162,11 @@ def topological(self, val): @utils.cached_property def subfunctions(self): r"""Split into a tuple of constituent spaces.""" - return tuple(type(self).create(subspace, self.mesh()) - for subspace in self.topological.subfunctions) + if len(self) == 1: + return (self, ) + else: + return tuple(type(self).create(subspace, mesh, parent=self) + for mesh, subspace in zip(self.mesh(), self.topological.subfunctions, strict=True)) mesh = ufl.FunctionSpace.ufl_domain @@ -174,18 +190,18 @@ def split(self): @utils.cached_property def _components(self): - if len(self) == 1: - return tuple(type(self).create(self.topological.sub(i), self.mesh()) - for i in range(self.block_size)) - else: - return self.subfunctions + return tuple(type(self).create(self.topological.sub(i), self.mesh(), parent=self) + for i in range(self.block_size)) @PETSc.Log.EventDecorator() def sub(self, i): bound = len(self._components) if i < 0 or i >= bound: raise IndexError("Invalid component %d, not in [0, %d)" % (i, bound)) - return self._components[i] + if len(self) == 1: + return self._components[i] + else: + return self.subfunctions[i] @utils.cached_property def dm(self): @@ -295,7 +311,7 @@ def restore_work_function(self, function): def __eq__(self, other): try: return self.topological == other.topological and \ - self.mesh() is other.mesh() + self.mesh() == other.mesh() except AttributeError: return False @@ -358,9 +374,17 @@ def make_function_space(cls, mesh, element, name=None): topology = mesh.topology # Create a new abstract (Mixed/Real)FunctionSpace, these are neither primal nor dual. if type(element) is finat.ufl.MixedElement: - spaces = [cls.make_function_space(topology, e) for e in element.sub_elements] - new = MixedFunctionSpace(spaces, name=name) + if isinstance(mesh, MeshGeometry): + mesh = MixedMeshGeometry(*[mesh for _ in element.sub_elements]) + topology = mesh.topology + else: + if not isinstance(mesh, MixedMeshGeometry): + raise TypeError(f"mesh must be MixedMeshGeometry: got {mesh}") + spaces = [cls.make_function_space(topo, e) for topo, e in zip(topology, element.sub_elements, strict=True)] + new = MixedFunctionSpace(spaces, topology, name=name) else: + if isinstance(mesh, MixedMeshGeometry): + raise TypeError(f"mesh must not be MixedMeshGeometry: got {mesh}") # Check that any Vector/Tensor/Mixed modifiers are outermost. check_element(element) if element.family() == "Real": @@ -418,7 +442,8 @@ def __init__(self, mesh, element, component=None, cargo=None): cargo=cargo) def dual(self): - return FiredrakeDualSpace.create(self.topological, self.mesh()) + parent = None if self.parent is None else self.parent.dual() + return FiredrakeDualSpace.create(self.topological, self.mesh(), parent=parent) class FiredrakeDualSpace(WithGeometryBase, ufl.functionspace.DualSpace): @@ -429,7 +454,8 @@ def __init__(self, mesh, element, component=None, cargo=None): cargo=cargo) def dual(self): - return WithGeometry.create(self.topological, self.mesh()) + parent = None if self.parent is None else self.parent.dual() + return WithGeometry.create(self.topological, self.mesh(), parent=parent) class FunctionSpace(object): @@ -560,7 +586,7 @@ def __eq__(self, other): if not isinstance(other, FunctionSpace): return False # FIXME: Think harder about equality - return self.mesh() is other.mesh() and \ + return self.mesh() == other.mesh() and \ self.dof_dset is other.dof_dset and \ self.ufl_element() == other.ufl_element() and \ self.component == other.component @@ -945,11 +971,14 @@ class MixedFunctionSpace(object): but should instead use the functional interface provided by :func:`.MixedFunctionSpace`. """ - def __init__(self, spaces, name=None): + def __init__(self, spaces, mesh, name=None): super(MixedFunctionSpace, self).__init__() + if not isinstance(mesh, MixedMeshTopology): + raise TypeError(f"mesh must be MixedMeshTopology: got {mesh}") + if len(mesh) != len(spaces): + raise RuntimeError(f"len(mesh) ({len(mesh)}) != len(spaces) ({len(spaces)})") self._spaces = tuple(IndexedFunctionSpace(i, s, self) for i, s in enumerate(spaces)) - mesh, = set(s.mesh() for s in spaces) self._ufl_function_space = ufl.FunctionSpace(mesh.ufl_mesh(), finat.ufl.MixedElement(*[s.ufl_element() for s in spaces])) self.name = name or "_".join(str(s.name) for s in spaces) @@ -1151,7 +1180,9 @@ def dm(self): def _dm(self): from firedrake.mg.utils import get_level dm = self.dof_dset.dm - _, level = get_level(self.mesh()) + # TODO: Think harder. + m = self.mesh()[0] + _, level = get_level(m) dmhooks.attach_hooks(dm, level=level) return dm @@ -1328,7 +1359,7 @@ def __eq__(self, other): if not isinstance(other, RealFunctionSpace): return False # FIXME: Think harder about equality - return self.mesh() is other.mesh() and \ + return self.mesh() == other.mesh() and \ self.ufl_element() == other.ufl_element() def __ne__(self, other): diff --git a/firedrake/mesh.py b/firedrake/mesh.py index fd37e7b575..62f3803aab 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -990,6 +990,12 @@ def mark_entities(self, tf, label_value, label_name=None): def extruded_periodic(self): return self.cell_set._extruded_periodic + def __iter__(self): + yield self + + def unique(self): + return self + # submesh @utils.cached_property @@ -2819,6 +2825,12 @@ def mark_entities(self, f, label_value, label_name=None): """ self.topology.mark_entities(f.topological, label_value, label_name) + def __iter__(self): + yield self + + def unique(self): + return self + @PETSc.Log.EventDecorator() def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): @@ -4668,3 +4680,190 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): submesh.submesh_parent = mesh submesh.init() return submesh + + +class MixedMeshGeometry(ufl.MixedMesh): + """A representation of mixed mesh geometry.""" + + def __init__(self, *meshes, set_hierarchy=True): + """Initialise. + + Parameters + ---------- + meshes : tuple or list + `MeshGeometry`s to make `MixedMeshGeometry` with. + set_hierarchy : bool + Flag for making hierarchy. + + """ + for m in meshes: + if not isinstance(m, MeshGeometry): + raise ValueError(f"Got {type(m)}") + super().__init__(*meshes) + self.comm = meshes[0].comm + self._comm = internal_comm(self.comm, self) + # Only set hierarchy at top level. + if set_hierarchy: + self.set_hierarchy() + + @utils.cached_property + def topology(self): + return MixedMeshTopology(*[m.topology for m in self._meshes]) + + @property + def topological(self): + """Alias of topology. + + This is to ensure consistent naming for some multigrid codes.""" + return self.topology + + def init(self): + for m in self._meshes: + m.init() + + def __eq__(self, other): + if type(other) != type(self): + return False + if len(other) != len(self): + return False + for o, s in zip(other, self): + if o is not s: + return False + return True + + def __ne__(self, other): + return not self.__eq__(other) + + def __hash__(self): + return hash(self._meshes) + + def __len__(self): + return len(self._meshes) + + def __iter__(self): + return iter(self._meshes) + + def __getitem__(self, i): + return self._meshes[i] + + @utils.cached_property + def extruded(self): + m = self.unique() + return m.extruded + + def unique(self): + """Return a single component or raise exception.""" + if len(set(self._meshes)) > 1: + raise RuntimeError(f"Found multiple meshes in {self} where a single mesh is expected") + m, = set(self._meshes) + return m + + def set_hierarchy(self): + """Set mesh hierarchy if needed.""" + from firedrake.mg.utils import set_level, get_level, has_level + + # TODO: Think harder on how mesh hierarchy should work with mixed meshes. + if all(not has_level(m) for m in self._meshes): + return + else: + if not all(has_level(m) for m in self._meshes): + raise RuntimeError("Found inconsistent component meshes") + hierarchy_list = [] + level_list = [] + for m in self: + hierarchy, level = get_level(m) + hierarchy_list.append(hierarchy) + level_list.append(level) + nlevels, = set(len(hierarchy) for hierarchy in hierarchy_list) + level, = set(level_list) + result = [] + for ilevel in range(nlevels): + if ilevel == level: + result.append(self) + else: + result.append(MixedMeshGeometry(*[hierarchy[ilevel] for hierarchy in hierarchy_list], set_hierarchy=False)) + result = tuple(result) + for i, m in enumerate(result): + set_level(m, result, i) + + +class MixedMeshTopology(object): + """A representation of mixed mesh topology.""" + + def __init__(self, *meshes): + """Initialise. + + Parameters + ---------- + meshes : tuple or list + `MeshTopology`s to make `MixedMeshTopology` with. + + """ + for m in meshes: + if not isinstance(m, AbstractMeshTopology): + raise ValueError(f"Got {type(m)}") + self._meshes = tuple(meshes) + self.comm = meshes[0].comm + self._comm = internal_comm(self.comm, self) + + def init(self): + for m in self._meshes: + m.init() + + @property + def topology(self): + """The underlying mesh topology object.""" + return self + + @property + def topological(self): + """Alias of topology. + + This is to ensure consistent naming for some multigrid codes.""" + return self + + def ufl_cell(self): + cell, = set(m.ufl_cell() for m in self._meshes) + return cell + + def ufl_mesh(self): + cell = self.ufl_cell() + return ufl.MixedMesh(*[ufl.Mesh(finat.ufl.VectorElement("Lagrange", cell, 1, dim=cell.topological_dimension())) + for _ in self._meshes]) + + def __eq__(self, other): + if type(other) != type(self): + return False + if len(other) != len(self): + return False + for o, s in zip(other, self): + if o is not s: + return False + return True + + def __ne__(self, other): + return not self.__eq__(other) + + def __hash__(self): + return hash(self._meshes) + + def __len__(self): + return len(self._meshes) + + def __iter__(self): + return iter(self._meshes) + + def __getitem__(self, i): + return self._meshes[i] + + @utils.cached_property + def extruded(self): + m = self.unique() + return m.extruded + + def unique(self): + """Return a single component or raise exception.""" + if len(set(self._meshes)) > 1: + raise RuntimeError(f"Found multiple meshes in {self} where a single mesh is expected") + m, = set(self._meshes) + return m diff --git a/firedrake/mg/kernels.py b/firedrake/mg/kernels.py index 5405b6d726..09d7448ef8 100644 --- a/firedrake/mg/kernels.py +++ b/firedrake/mg/kernels.py @@ -135,6 +135,7 @@ def compile_element(expression, dual_space=None, parameters=None, # Replace coordinates (if any) builder = firedrake_interface.KernelBuilderBase(scalar_type=ScalarType) domain = extract_unique_domain(expression) + builder._domain_integral_type_map = {domain: "cell"} # Translate to GEM cell = domain.ufl_cell() dim = cell.topological_dimension() @@ -520,6 +521,7 @@ def dg_injection_kernel(Vf, Vc, ncell): if complex_mode: raise NotImplementedError("In complex mode we are waiting for Slate") macro_builder = MacroKernelBuilder(ScalarType, ncell) + macro_builder._domain_integral_type_map = {Vf.mesh(): "cell"} f = ufl.Coefficient(Vf) macro_builder.set_coefficients([f]) macro_builder.set_coordinates(Vf.mesh()) @@ -542,7 +544,8 @@ def dg_injection_kernel(Vf, Vc, ncell): entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=macro_quadrature_rule, - scalar_type=parameters["scalar_type"]) + scalar_type=parameters["scalar_type"], + domain_integral_type_map={Vf.mesh(): "cell"}) macro_context = fem.PointSetContext(**macro_cfg) fexpr, = fem.compile_ufl(f, macro_context) @@ -558,8 +561,10 @@ def dg_injection_kernel(Vf, Vc, ncell): integral_type="cell", subdomain_id=("otherwise",), domain_number=0, + domain_integral_type_map={Vc.mesh(): "cell"}, arguments=(ufl.TestFunction(Vc), ), coefficients=(), + coefficient_split={}, coefficient_numbers=()) coarse_builder = firedrake_interface.KernelBuilder(info, parameters["scalar_type"]) @@ -579,7 +584,8 @@ def dg_injection_kernel(Vf, Vc, ncell): entity_ids=entity_ids, index_cache=index_cache, quadrature_rule=quadrature_rule, - scalar_type=parameters["scalar_type"]) + scalar_type=parameters["scalar_type"], + domain_integral_type_map={Vc.mesh(): "cell"}) X = ufl.SpatialCoordinate(Vc.mesh()) K = ufl_utils.preprocess_expression(ufl.JacobianInverse(Vc.mesh()), diff --git a/firedrake/mg/ufl_utils.py b/firedrake/mg/ufl_utils.py index 0408a6d2a6..11b2b4fb51 100644 --- a/firedrake/mg/ufl_utils.py +++ b/firedrake/mg/ufl_utils.py @@ -62,6 +62,7 @@ def coarsen(expr, self, coefficient_mapping=None): @coarsen.register(ufl.Mesh) +@coarsen.register(ufl.MixedMesh) def coarsen_mesh(mesh, self, coefficient_mapping=None): hierarchy, level = utils.get_level(mesh) if hierarchy is None: @@ -132,7 +133,9 @@ def coarsen_function_space(V, self, coefficient_mapping=None): return V._coarse V_fine = V - mesh_coarse = self(V_fine.mesh(), self) + # Handle MixedFunctionSpace : V_fine.reconstruct requires MixedMesh. + fine_mesh = V_fine.mesh() if V_fine.index is None else V_fine.parent.mesh() + mesh_coarse = self(fine_mesh, self) name = f"coarse_{V.name}" if V.name else None V_coarse = V_fine.reconstruct(mesh=mesh_coarse, name=name) V_coarse._fine = V_fine diff --git a/firedrake/pointeval_utils.py b/firedrake/pointeval_utils.py index 75ed2b4564..fc056738a9 100644 --- a/firedrake/pointeval_utils.py +++ b/firedrake/pointeval_utils.py @@ -54,6 +54,7 @@ def compile_element(expression, coordinates, parameters=None): # Initialise kernel builder builder = firedrake_interface.KernelBuilderBase(utils.ScalarType) + builder._domain_integral_type_map = {domain: "cell"} builder.domain_coordinate[domain] = coordinates builder._coefficient(coordinates, "x") x_arg = builder.generate_arg_from_expression(builder.coefficient_map[coordinates]) diff --git a/firedrake/pointquery_utils.py b/firedrake/pointquery_utils.py index 4793af0bfe..58ceeb10c7 100644 --- a/firedrake/pointquery_utils.py +++ b/firedrake/pointquery_utils.py @@ -143,8 +143,8 @@ def to_reference_coords_newton_step(ufl_coordinate_element, parameters, x0_dtype expr = ufl_utils.simplify_abs(expr, complex_mode) builder = firedrake_interface.KernelBuilderBase(ScalarType) + builder._domain_integral_type_map = {domain: "cell"} builder.domain_coordinate[domain] = C - Cexpr = builder._coefficient(C, "C") x0_expr = builder._coefficient(x0, "x0") loopy_args = [ diff --git a/firedrake/preconditioners/asm.py b/firedrake/preconditioners/asm.py index 19738fe438..7dbbdc2341 100644 --- a/firedrake/preconditioners/asm.py +++ b/firedrake/preconditioners/asm.py @@ -134,8 +134,12 @@ class ASMStarPC(ASMPatchPC): def get_patches(self, V): mesh = V._mesh - mesh_dm = mesh.topology_dm - if mesh.cell_set._extruded: + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + mesh_dm = mesh_unique.topology_dm + if mesh_unique.cell_set._extruded: warning("applying ASMStarPC on an extruded mesh") # Obtain the topological entities to use to construct the stars @@ -189,8 +193,12 @@ class ASMVankaPC(ASMPatchPC): def get_patches(self, V): mesh = V._mesh - mesh_dm = mesh.topology_dm - if mesh.layers: + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + mesh_dm = mesh_unique.topology_dm + if mesh_unique.layers: warning("applying ASMVankaPC on an extruded mesh") # Obtain the topological entities to use to construct the stars @@ -278,8 +286,12 @@ class ASMLinesmoothPC(ASMPatchPC): def get_patches(self, V): mesh = V._mesh - assert mesh.cell_set._extruded - dm = mesh.topology_dm + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + assert mesh_unique.cell_set._extruded + dm = mesh_unique.topology_dm section = V.dm.getDefaultSection() # Obtain the codimensions to loop over from options, if present opts = PETSc.Options(self.prefix) @@ -377,9 +389,13 @@ class ASMExtrudedStarPC(ASMStarPC): def get_patches(self, V): mesh = V.mesh() - mesh_dm = mesh.topology_dm - nlayers = mesh.layers - if not mesh.cell_set._extruded: + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + mesh_dm = mesh_unique.topology_dm + nlayers = mesh_unique.layers + if not mesh_unique.cell_set._extruded: return super(ASMExtrudedStarPC, self).get_patches(V) # Obtain the topological entities to use to construct the stars diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 82e56057e7..c2f621a63e 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -216,7 +216,7 @@ def allocate_matrix(self, Amat, V, J, bcs, fcp, pmat_type, use_static_condensati elif len(ifacet) == 1: Vfacet = V[ifacet[0]] ebig, = set(unrestrict_element(Vsub.ufl_element()) for Vsub in V) - Vbig = FunctionSpace(V.mesh(), ebig) + Vbig = FunctionSpace(V.mesh().unique(), ebig) if len(V) > 1: dims = [Vsub.finat_element.space_dimension() for Vsub in V] assert sum(dims) == Vbig.finat_element.space_dimension() diff --git a/firedrake/preconditioners/patch.py b/firedrake/preconditioners/patch.py index 0a7bad5575..68d89ccf03 100644 --- a/firedrake/preconditioners/patch.py +++ b/firedrake/preconditioners/patch.py @@ -6,6 +6,7 @@ from firedrake.dmhooks import get_appctx, push_appctx, pop_appctx from firedrake.functionspace import FunctionSpace from firedrake.interpolation import Interpolate +from firedrake.ufl_expr import extract_domains from collections import namedtuple import operator @@ -14,7 +15,6 @@ import numpy from finat.ufl import VectorElement, MixedElement from ufl.domain import extract_unique_domain -from tsfc.kernel_interface.firedrake_loopy import make_builder from tsfc.ufl_utils import extract_firedrake_constants import weakref @@ -133,6 +133,10 @@ def __call__(self, access, map_=None): CompiledKernel = namedtuple('CompiledKernel', ["funptr", "kinfo"]) +def get_map(V, base_mesh, base_integral_type): + return V.topological.entity_node_map(base_mesh.topology, base_integral_type, None, None) + + def matrix_funptr(form, state): from firedrake.tsfc_interface import compile_form test, trial = map(operator.methodcaller("function_space"), form.arguments()) @@ -140,16 +144,19 @@ def matrix_funptr(form, state): raise NotImplementedError("Only for matching test and trial spaces") if state is not None: - interface = make_builder(dont_split=(state, )) + dont_split = (state, ) else: - interface = None + dont_split = None - kernels = compile_form(form, "subspace_form", split=False, interface=interface) + kernels = compile_form(form, "subspace_form", split=False, dont_split=dont_split) + all_meshes = extract_domains(form) cell_kernels = [] int_facet_kernels = [] for kernel in kernels: kinfo = kernel.kinfo + mesh = all_meshes[kinfo.domain_number] # integration domain + integral_type = kinfo.integral_type if kinfo.subdomain_id != ("otherwise",): raise NotImplementedError("Only for full domain integrals") @@ -159,21 +166,17 @@ def matrix_funptr(form, state): # OK, now we've validated the kernel, let's build the callback args = [] - if kinfo.integral_type == "cell": - get_map = operator.methodcaller("cell_node_map") + if integral_type == "cell": kernels = cell_kernels - elif kinfo.integral_type == "interior_facet": - get_map = operator.methodcaller("interior_facet_node_map") + elif integral_type == "interior_facet": kernels = int_facet_kernels - else: - get_map = None toset = op2.Set(1, comm=test.comm) dofset = op2.DataSet(toset, 1) arity = sum(m.arity*s.cdim - for m, s in zip(get_map(test), + for m, s in zip(get_map(test, mesh, integral_type), test.dof_dset)) - iterset = get_map(test).iterset + iterset = get_map(test, mesh, integral_type).iterset entity_node_map = op2.Map(iterset, toset, arity, values=numpy.zeros(iterset.total_size*arity, dtype=IntType)) @@ -187,16 +190,17 @@ def matrix_funptr(form, state): values=numpy.zeros(iterset.total_size*arity, dtype=IntType)) statearg = statedat(op2.READ, state_entity_node_map) - mesh = form.ufl_domains()[kinfo.domain_number] - arg = mesh.coordinates.dat(op2.READ, get_map(mesh.coordinates)) - args.append(arg) - if kinfo.oriented: - c = mesh.cell_orientations() - arg = c.dat(op2.READ, get_map(c)) + for i in kinfo.active_domain_numbers.coordinates: + c = all_meshes[i].coordinates + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) args.append(arg) - if kinfo.needs_cell_sizes: - c = mesh.cell_sizes - arg = c.dat(op2.READ, get_map(c)) + for i in kinfo.active_domain_numbers.cell_orientations: + c = all_meshes[i].cell_orientations() + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) + args.append(arg) + for i in kinfo.active_domain_numbers.cell_sizes: + c = all_meshes[i].cell_sizes + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) args.append(arg) for n, indices in kinfo.coefficient_numbers: c = form.coefficients()[n] @@ -207,7 +211,7 @@ def matrix_funptr(form, state): continue for ind in indices: c_ = c.subfunctions[ind] - map_ = get_map(c_) + map_ = get_map(c_.function_space(), mesh, integral_type) arg = c_.dat(op2.READ, map_) args.append(arg) @@ -215,7 +219,7 @@ def matrix_funptr(form, state): for constant_index in kinfo.constant_numbers: args.append(all_constants[constant_index].dat(op2.READ)) - if kinfo.integral_type == "interior_facet": + if integral_type == "interior_facet": arg = mesh.interior_facets.local_facet_dat(op2.READ) args.append(arg) iterset = op2.Subset(iterset, []) @@ -234,16 +238,19 @@ def residual_funptr(form, state): raise NotImplementedError("State and test space must be dual to one-another") if state is not None: - interface = make_builder(dont_split=(state, )) + dont_split = (state, ) else: - interface = None + dont_split = None - kernels = compile_form(form, "subspace_form", split=False, interface=interface) + kernels = compile_form(form, "subspace_form", split=False, dont_split=dont_split) + all_meshes = extract_domains(form) cell_kernels = [] int_facet_kernels = [] for kernel in kernels: kinfo = kernel.kinfo + mesh = all_meshes[kinfo.domain_number] # integration domain + integral_type = kinfo.integral_type if kinfo.subdomain_id != ("otherwise",): raise NotImplementedError("Only for full domain integrals") @@ -252,20 +259,16 @@ def residual_funptr(form, state): args = [] if kinfo.integral_type == "cell": - get_map = operator.methodcaller("cell_node_map") kernels = cell_kernels elif kinfo.integral_type == "interior_facet": - get_map = operator.methodcaller("interior_facet_node_map") kernels = int_facet_kernels - else: - get_map = None toset = op2.Set(1, comm=test.comm) dofset = op2.DataSet(toset, 1) arity = sum(m.arity*s.cdim - for m, s in zip(get_map(test), + for m, s in zip(get_map(test, mesh, integral_type), test.dof_dset)) - iterset = get_map(test).iterset + iterset = get_map(test, mesh, integral_type).iterset entity_node_map = op2.Map(iterset, toset, arity, values=numpy.zeros(iterset.total_size*arity, dtype=IntType)) @@ -280,17 +283,17 @@ def residual_funptr(form, state): arg = dat(op2.INC, entity_node_map) args.append(arg) - mesh = form.ufl_domains()[kinfo.domain_number] - arg = mesh.coordinates.dat(op2.READ, get_map(mesh.coordinates)) - args.append(arg) - - if kinfo.oriented: - c = mesh.cell_orientations() - arg = c.dat(op2.READ, get_map(c)) + for i in kinfo.active_domain_numbers.coordinates: + c = all_meshes[i].coordinates + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) + args.append(arg) + for i in kinfo.active_domain_numbers.cell_orientations: + c = all_meshes[i].cell_orientations() + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) args.append(arg) - if kinfo.needs_cell_sizes: - c = mesh.cell_sizes - arg = c.dat(op2.READ, get_map(c)) + for i in kinfo.active_domain_numbers.cell_sizes: + c = all_meshes[i].cell_sizes + arg = c.dat(op2.READ, get_map(c.function_space(), mesh, integral_type)) args.append(arg) for n, indices in kinfo.coefficient_numbers: c = form.coefficients()[n] @@ -301,7 +304,7 @@ def residual_funptr(form, state): continue for ind in indices: c_ = c.subfunctions[ind] - map_ = get_map(c_) + map_ = get_map(c_.function_space(), mesh, integral_type) arg = c_.dat(op2.READ, map_) args.append(arg) @@ -513,14 +516,14 @@ def load_c_function(code, name, comm): comm=comm) -def make_c_arguments(form, kernel, state, get_map, require_state=False, +def make_c_arguments(form, kernel, state, integral_type, require_state=False, require_facet_number=False): - mesh = form.ufl_domains()[kernel.kinfo.domain_number] - coeffs = [mesh.coordinates] - if kernel.kinfo.oriented: - coeffs.append(mesh.cell_orientations()) - if kernel.kinfo.needs_cell_sizes: - coeffs.append(mesh.cell_sizes) + all_meshes = extract_domains(form) + mesh = all_meshes[kernel.kinfo.domain_number] + coeffs = [] + coeffs.extend([all_meshes[i].coordinates for i in kernel.kinfo.active_domain_numbers.coordinates]) + coeffs.extend([all_meshes[i].cell_orientations() for i in kernel.kinfo.active_domain_numbers.cell_orientations]) + coeffs.extend([all_meshes[i].cell_sizes for i in kernel.kinfo.active_domain_numbers.cell_sizes]) for n, indices in kernel.kinfo.coefficient_numbers: c = form.coefficients()[n] if c is state: @@ -540,7 +543,7 @@ def make_c_arguments(form, kernel, state, get_map, require_state=False, map_args.append(None) else: data_args.extend(c.dat._kernel_args_) - map_ = get_map(c) + map_ = get_map(c.function_space(), mesh, integral_type) if map_ is not None: for k in map_._kernel_args_: if k not in seen: @@ -649,8 +652,12 @@ def sort_entities(self, dm, axis, dir, ndiv=None, divisions=None): raise RuntimeError("Must either set ndiv or divisions for PlaneSmoother!") mesh = dm.getAttr("__firedrake_mesh__") - ele = mesh.coordinates.function_space().ufl_element() - V = mesh.coordinates.function_space() + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + ele = mesh_unique.coordinates.function_space().ufl_element() + V = mesh_unique.coordinates.function_space() if V.finat_element.entity_dofs() == V.finat_element.entity_closure_dofs(): # We're using DG or DQ for our coordinates, so we got # a periodic mesh. We need to interpolate to CGk @@ -660,11 +667,11 @@ def sort_entities(self, dm, axis, dir, ndiv=None, divisions=None): # Need to supply the actual mesh to the FunctionSpace constructor, # not its weakref proxy (the variable `mesh`) # as interpolation fails because they are not hashable - CGk = FunctionSpace(mesh.coordinates.function_space().mesh(), CGkele) - coordinates = Interpolate(mesh.coordinates, CGk, access=op2.MAX) + CGk = FunctionSpace(mesh_unique.coordinates.function_space().mesh(), CGkele) + coordinates = Interpolate(mesh_unique.coordinates, CGk, access=op2.MAX) coordinates = assemble(coordinates) else: - coordinates = mesh.coordinates + coordinates = mesh_unique.coordinates select = partial(select_entity, dm=dm, exclude="pyop2_ghost") entities = [(p, self.coords(dm, p, coordinates)) for p in @@ -764,7 +771,11 @@ def initialize(self, obj): J, bcs = self.form(obj) V = J.arguments()[0].function_space() mesh = V.mesh() - self.plex = mesh.topology_dm + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") + self.plex = mesh_unique.topology_dm # We need to attach the mesh and appctx to the plex, so that # PlaneSmoothers (and any other user-customised patch # constructors) can use firedrake's opinion of what @@ -773,10 +784,10 @@ def initialize(self, obj): self.ctx = ctx self.plex.setAttr("__firedrake_ctx__", weakref.proxy(ctx)) - if mesh.cell_set._extruded: + if mesh_unique.cell_set._extruded: raise NotImplementedError("Not implemented on extruded meshes") - if "overlap_type" not in mesh._distribution_parameters: + if "overlap_type" not in mesh_unique._distribution_parameters: if mesh.comm.size > 1: # Want to do # warnings.warn("You almost surely want to set an overlap_type in your mesh's distribution_parameters.") @@ -810,8 +821,7 @@ def initialize(self, obj): Jcell_kernels, Jint_facet_kernels = matrix_funptr(J, Jstate) Jcell_kernel, = Jcell_kernels Jcell_flops = Jcell_kernel.kinfo.kernel.num_flops - Jop_data_args, Jop_map_args = make_c_arguments(J, Jcell_kernel, Jstate, - operator.methodcaller("cell_node_map")) + Jop_data_args, Jop_map_args = make_c_arguments(J, Jcell_kernel, Jstate, "cell") code, Struct = make_jacobian_wrapper(Jop_data_args, Jop_map_args, Jcell_flops) Jop_function = load_c_function(code, "ComputeJacobian", mesh.comm) Jop_struct = make_c_struct(Jop_data_args, Jop_map_args, Jcell_kernel.funptr, Struct) @@ -822,11 +832,11 @@ def initialize(self, obj): Jhas_int_facet_kernel = True Jint_facet_flops = Jint_facet_kernel.kinfo.kernel.num_flops facet_Jop_data_args, facet_Jop_map_args = make_c_arguments(J, Jint_facet_kernel, Jstate, - operator.methodcaller("interior_facet_node_map"), + "interior_facet", require_facet_number=True) code, Struct = make_jacobian_wrapper(facet_Jop_data_args, facet_Jop_map_args, Jint_facet_flops) facet_Jop_function = load_c_function(code, "ComputeJacobian", mesh.comm) - point2facet = mesh.interior_facets.point2facetnumber.ctypes.data + point2facet = mesh_unique.interior_facets.point2facetnumber.ctypes.data facet_Jop_struct = make_c_struct(facet_Jop_data_args, facet_Jop_map_args, Jint_facet_kernel.funptr, Struct, point2facet=point2facet) @@ -840,7 +850,7 @@ def initialize(self, obj): Fcell_kernel, = Fcell_kernels Fcell_flops = Fcell_kernel.kinfo.kernel.num_flops Fop_data_args, Fop_map_args = make_c_arguments(F, Fcell_kernel, Fstate, - operator.methodcaller("cell_node_map"), + "cell", require_state=True) code, Struct = make_residual_wrapper(Fop_data_args, Fop_map_args, Fcell_flops) Fop_function = load_c_function(code, "ComputeResidual", mesh.comm) @@ -852,7 +862,7 @@ def initialize(self, obj): Fhas_int_facet_kernel = True Fint_facet_flops = Fint_facet_kernel.kinfo.kernel.num_flops facet_Fop_data_args, facet_Fop_map_args = make_c_arguments(F, Fint_facet_kernel, Fstate, - operator.methodcaller("interior_facet_node_map"), + "interior_facet", require_state=True, require_facet_number=True) code, Struct = make_jacobian_wrapper(facet_Fop_data_args, facet_Fop_map_args, Fint_facet_flops) @@ -863,7 +873,7 @@ def initialize(self, obj): point2facet=point2facet) patch.setDM(self.plex) - patch.setPatchCellNumbering(mesh._cell_numbering) + patch.setPatchCellNumbering(mesh_unique._cell_numbering) offsets = numpy.append([0], numpy.cumsum([W.dof_count for W in V])).astype(PETSc.IntType) diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 251f585cbe..77dff36b3f 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -1229,7 +1229,7 @@ def _weight(self): }} """ kernel = op2.Kernel(kernel_code, "weight", requires_zeroed_output_arguments=True) - op2.par_loop(kernel, weight.cell_set, weight.dat(op2.INC, weight.cell_node_map())) + op2.par_loop(kernel, weight.function_space().mesh().topology.unique().cell_set, weight.dat(op2.INC, weight.cell_node_map())) with weight.dat.vec as w: w.reciprocal() return weight @@ -1243,7 +1243,7 @@ def _kernels(self): uf_map = get_permuted_map(self.Vf) uc_map = get_permuted_map(self.Vc) prolong_kernel, restrict_kernel, coefficients = self.make_blas_kernels(self.Vf, self.Vc) - prolong_args = [prolong_kernel, self.uf.cell_set, + prolong_args = [prolong_kernel, self.uf.function_space().mesh().topology.unique().cell_set, self.uf.dat(op2.INC, uf_map), self.uc.dat(op2.READ, uc_map), self._weight.dat(op2.READ, uf_map)] @@ -1253,11 +1253,11 @@ def _kernels(self): uf_map = self.Vf.cell_node_map() uc_map = self.Vc.cell_node_map() prolong_kernel, restrict_kernel, coefficients = self.make_kernels(self.Vf, self.Vc) - prolong_args = [prolong_kernel, self.uf.cell_set, + prolong_args = [prolong_kernel, self.uf.function_space().mesh().topology.unique().cell_set, self.uf.dat(op2.WRITE, uf_map), self.uc.dat(op2.READ, uc_map)] - restrict_args = [restrict_kernel, self.uf.cell_set, + restrict_args = [restrict_kernel, self.uf.function_space().mesh().topology.unique().cell_set, self.uc.dat(op2.INC, uc_map), self.uf.dat(op2.READ, uf_map), self._weight.dat(op2.READ, uf_map)] @@ -1598,7 +1598,7 @@ def prolongation_matrix_aij(P1, Pk, P1_bcs=[], Pk_bcs=[]): lgmaps=((rlgmap, clgmap), ), unroll_map=unroll) expr = firedrake.TestFunction(P1.sub(i)) kernel, coefficients = prolongation_transfer_kernel_action(Pk.sub(i), expr) - parloop_args = [kernel, mesh.cell_set, matarg] + parloop_args = [kernel, mesh.topology.unique().cell_set, matarg] for coefficient in coefficients: m_ = coefficient.cell_node_map() parloop_args.append(coefficient.dat(op2.READ, m_)) @@ -1615,7 +1615,7 @@ def prolongation_matrix_aij(P1, Pk, P1_bcs=[], Pk_bcs=[]): lgmaps=((rlgmap, clgmap), ), unroll_map=unroll) expr = firedrake.TestFunction(P1) kernel, coefficients = prolongation_transfer_kernel_action(Pk, expr) - parloop_args = [kernel, mesh.cell_set, matarg] + parloop_args = [kernel, mesh.topology.unique().cell_set, matarg] for coefficient in coefficients: m_ = coefficient.cell_node_map() parloop_args.append(coefficient.dat(op2.READ, m_)) diff --git a/firedrake/slate/slac/compiler.py b/firedrake/slate/slac/compiler.py index 7e1b14281c..a49d3ea11a 100644 --- a/firedrake/slate/slac/compiler.py +++ b/firedrake/slate/slac/compiler.py @@ -36,6 +36,7 @@ from gem import indices as make_indices from tsfc.kernel_args import OutputKernelArg, CoefficientKernelArg from tsfc.loopy import generate as generate_loopy +from tsfc.kernel_interface.firedrake_loopy import ActiveDomainNumbers import copy from petsc4py import PETSc @@ -190,14 +191,17 @@ def generate_loopy_kernel(slate_expr, compiler_parameters=None): kinfo = KernelInfo(kernel=loopykernel, integral_type="cell", # slate can only do things as contributions to the cell integrals - oriented=builder.bag.needs_cell_orientations, subdomain_id=("otherwise",), domain_number=0, + active_domain_numbers=ActiveDomainNumbers(coordinates=(0, ) if builder.bag.needs_coordinates else (), + cell_orientations=(0, ) if builder.bag.needs_cell_orientations else (), + cell_sizes=(0, ) if builder.bag.needs_cell_sizes else (), + exterior_facets=(), + interior_facets=()), coefficient_numbers=coefficient_numbers, constant_numbers=constant_numbers, needs_cell_facets=builder.bag.needs_cell_facets, pass_layer_arg=builder.bag.needs_mesh_layers, - needs_cell_sizes=builder.bag.needs_cell_sizes, arguments=arguments, events=events) diff --git a/firedrake/slate/slac/kernel_builder.py b/firedrake/slate/slac/kernel_builder.py index cbc9b6fed2..7f1c44e128 100644 --- a/firedrake/slate/slac/kernel_builder.py +++ b/firedrake/slate/slac/kernel_builder.py @@ -132,13 +132,22 @@ def collect_tsfc_kernel_data(self, mesh, tsfc_coefficients, tsfc_constants, wrap that are coordinates, orientations, cell sizes and cofficients. """ - kernel_data = [(mesh.coordinates, self.coordinates_arg_name)] - - if kinfo.oriented: + kernel_data = [] + for coord_domain_number in kinfo.active_domain_numbers.coordinates: + if coord_domain_number != 0: + raise ValueError("Slate currently only supports single domain") + self.bag.needs_coordinates = True + kernel_data.append((mesh.coordinates, self.coordinates_arg_name)) + + for cell_orientation_domain_number in kinfo.active_domain_numbers.cell_orientations: + if cell_orientation_domain_number != 0: + raise ValueError("Slate currently only supports single domain") self.bag.needs_cell_orientations = True kernel_data.append((mesh.cell_orientations(), self.cell_orientations_arg_name)) - if kinfo.needs_cell_sizes: + for cell_size_domain_number in kinfo.active_domain_numbers.cell_sizes: + if cell_size_domain_number != 0: + raise ValueError("Slate currently only supports single domain") self.bag.needs_cell_sizes = True kernel_data.append((mesh.cell_sizes, self.cell_sizes_arg_name)) @@ -339,10 +348,11 @@ def generate_wrapper_kernel_args(self, tensor2temp): args = [] tmp_args = [] - coords_extent = self.extent(self.expression.ufl_domain().coordinates) - coords_loopy_arg = loopy.GlobalArg(self.coordinates_arg_name, shape=coords_extent, - dtype=self.tsfc_parameters["scalar_type"]) - args.append(kernel_args.CoordinatesKernelArg(coords_loopy_arg)) + if self.bag.needs_coordinates: + coords_extent = self.extent(self.expression.ufl_domain().coordinates) + coords_loopy_arg = loopy.GlobalArg(self.coordinates_arg_name, shape=coords_extent, + dtype=self.tsfc_parameters["scalar_type"]) + args.append(kernel_args.CoordinatesKernelArg(coords_loopy_arg)) if self.bag.needs_cell_orientations: ori_extent = self.extent(self.expression.ufl_domain().cell_orientations()) @@ -442,9 +452,15 @@ def generate_tsfc_calls(self, terminal, loopy_tensor): if subdomain_id != "otherwise": raise NotImplementedError("No subdomain markers for cells yet") elif self.is_integral_type(integral_type, "facet_integral"): - predicates, fidx, facet_arg = self.facet_integral_predicates(mesh, integral_type, kinfo, subdomain_id) - reads.append(facet_arg) - inames_dep.append(fidx[0].name) + if kinfo.active_domain_numbers._asdict()[{"exterior_facet": "exterior_facets", + "exterior_facet_vert": "exterior_facets", + "interior_facet": "interior_facets", + "interior_facet_vert": "interior_facets"}[kinfo.integral_type]] != (): + predicates, fidx, facet_arg = self.facet_integral_predicates(mesh, integral_type, kinfo, subdomain_id) + reads.append(facet_arg) + inames_dep.append(fidx[0].name) + else: + predicates = None elif self.is_integral_type(integral_type, "layer_integral"): predicates = self.layer_integral_predicates(slate_tensor, integral_type) else: @@ -471,6 +487,7 @@ def __init__(self, coeffs, constants): self.coefficients = coeffs self.constants = constants self.inames = OrderedDict() + self.needs_coordinates = False self.needs_cell_orientations = False self.needs_cell_sizes = False self.needs_cell_facets = False diff --git a/firedrake/slate/slate.py b/firedrake/slate/slate.py index 85b7af3635..9abbc5980f 100644 --- a/firedrake/slate/slate.py +++ b/firedrake/slate/slate.py @@ -15,7 +15,7 @@ functions to be executed within the Firedrake architecture. """ from abc import ABCMeta, abstractproperty, abstractmethod - +import functools from collections import OrderedDict, namedtuple, defaultdict from ufl import Constant @@ -253,6 +253,13 @@ def ufl_domain(self): raise ValueError("All integrals must share the same domain of integration.") return domain + @staticmethod + def _expand_mixed_meshes(func): + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + return join_domains(func(self, *args, **kwargs)) + return wrapper + @abstractmethod def ufl_domains(self): """Returns the integration domains of the integrals associated with @@ -477,6 +484,7 @@ def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. @@ -553,6 +561,7 @@ def slate_coefficients(self): """Returns a BlockFunction in a tuple which carries all information to generate the right coefficients and maps.""" return (BlockFunction(self._function, self._indices, self._original_function),) + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. """ @@ -726,6 +735,7 @@ def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. @@ -821,6 +831,7 @@ def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. @@ -924,6 +935,7 @@ def slate_coefficients(self): """Returns a tuple of coefficients associated with the tensor.""" return self.coefficients() + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. @@ -978,6 +990,7 @@ def slate_coefficients(self): coeffs = [op.slate_coefficients() for op in self.operands] return tuple(OrderedDict.fromkeys(chain(*coeffs))) + @TensorBase._expand_mixed_meshes def ufl_domains(self): """Returns the integration domains of the integrals associated with the tensor. diff --git a/firedrake/slate/static_condensation/hybridization.py b/firedrake/slate/static_condensation/hybridization.py index 997f969684..191823b149 100644 --- a/firedrake/slate/static_condensation/hybridization.py +++ b/firedrake/slate/static_condensation/hybridization.py @@ -57,6 +57,10 @@ def initialize(self, pc): V = test.function_space() mesh = V.mesh() + if len(set(mesh)) == 1: + mesh_unique = mesh.unique() + else: + raise NotImplementedError("Not implemented for general mixed meshes") if len(V) != 2: raise ValueError("Expecting two function spaces.") @@ -86,7 +90,7 @@ def initialize(self, pc): except TypeError: tdegree = W.ufl_element().degree() - 1 - TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree) + TraceSpace = FunctionSpace(mesh[self.vidx], "HDiv Trace", tdegree) # Break the function spaces and define fully discontinuous spaces broken_elements = finat.ufl.MixedElement([finat.ufl.BrokenElement(Vi.ufl_element()) for Vi in V]) @@ -125,10 +129,10 @@ def initialize(self, pc): trial: TrialFunction(V_d)} Atilde = Tensor(replace(self.ctx.a, arg_map)) gammar = TestFunction(TraceSpace) - n = ufl.FacetNormal(mesh) + n = ufl.FacetNormal(mesh_unique) sigma = TrialFunctions(V_d)[self.vidx] - if mesh.cell_set._extruded: + if mesh_unique.cell_set._extruded: Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h + gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v) else: @@ -166,7 +170,7 @@ def initialize(self, pc): integrand = gammar * ufl.dot(sigma, n) measures = [] trace_subdomains = [] - if mesh.cell_set._extruded: + if mesh_unique.cell_set._extruded: ds = ufl.ds_v for subdomain in sorted(extruded_neumann_subdomains): measures.append({"top": ufl.ds_t, "bottom": ufl.ds_b}[subdomain]) @@ -177,7 +181,7 @@ def initialize(self, pc): measures.append(ds) else: measures.extend((ds(sd) for sd in sorted(neumann_subdomains))) - markers = [int(x) for x in mesh.exterior_facets.unique_markers] + markers = [int(x) for x in mesh_unique.exterior_facets.unique_markers] dirichlet_subdomains = set(markers) - neumann_subdomains trace_subdomains.extend(sorted(dirichlet_subdomains)) @@ -192,7 +196,7 @@ def initialize(self, pc): # the exterior boundary. Extruded cells will have both # horizontal and vertical facets trace_subdomains = ["on_boundary"] - if mesh.cell_set._extruded: + if mesh_unique.cell_set._extruded: trace_subdomains.extend(["bottom", "top"]) trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), subdomain) for subdomain in trace_subdomains] diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index 35fa4742eb..c4122ba81a 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/firedrake/slate/static_condensation/scpc.py @@ -67,7 +67,7 @@ def initialize(self, pc): # Need to duplicate a space which is NOT # associated with a subspace of a mixed space. - Vc = FunctionSpace(W.mesh(), W[c_field].ufl_element()) + Vc = FunctionSpace(W.mesh()[c_field], W[c_field].ufl_element()) bcs = [] cxt_bcs = self.cxt.row_bcs for bc in cxt_bcs: diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index a4a57ae0cb..e015f5d087 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -12,11 +12,12 @@ import ufl import finat.ufl from ufl import Form, conj -from .ufl_expr import TestFunction +from .ufl_expr import TestFunction, extract_domains from tsfc import compile_form as original_tsfc_compile_form from tsfc.parameters import PARAMETERS as tsfc_default_parameters from tsfc.ufl_utils import extract_firedrake_constants +from tsfc.kernel_interface.firedrake_loopy import ActiveDomainNumbers from pyop2 import op2 from pyop2.caching import memory_and_disk_cache, default_parallel_hashkey @@ -35,14 +36,13 @@ KernelInfo = collections.namedtuple("KernelInfo", ["kernel", "integral_type", - "oriented", "subdomain_id", "domain_number", + "active_domain_numbers", "coefficient_numbers", "constant_numbers", "needs_cell_facets", "pass_layer_arg", - "needs_cell_sizes", "arguments", "events"]) @@ -53,9 +53,9 @@ ) -def tsfc_compile_form_hashkey(form, prefix, parameters, interface, diagonal, log): +def tsfc_compile_form_hashkey(form, prefix, parameters, dont_split_numbers, diagonal, log): # Drop prefix as it's only used for naming and log - return default_parallel_hashkey(form.signature(), prefix, parameters, interface, diagonal) + return default_parallel_hashkey(form.signature(), prefix, parameters, dont_split_numbers, diagonal) def tsfc_compile_form_comm_fetcher(*args, **kwargs): @@ -77,9 +77,10 @@ def __init__( form, name, parameters, + domain_number_map, coefficient_numbers, constant_numbers, - interface, + dont_split_numbers, diagonal=False ): """A wrapper object for one or more TSFC kernels compiled from a given :class:`~ufl.classes.Form`. @@ -87,16 +88,19 @@ def __init__( :arg form: the :class:`~ufl.classes.Form` from which to compile the kernels. :arg name: a prefix to be applied to the compiled kernel names. This is primarily useful for debugging. :arg parameters: a dict of parameters to pass to the form compiler. + :arg domain_number_map: Map from domain numbers in the provided (split) form to domain numbers in the original form. :arg coefficient_numbers: Map from coefficient numbers in the provided (split) form to coefficient numbers in the original form. :arg constant_numbers: Map from local constant numbers in the provided (split) form to constant numbers in the original form. - :arg interface: the KernelBuilder interface for TSFC (may be None) + :arg dont_split_numbers: Block-local numbers of coefficients that are not to be split. :arg diagonal: If assembling a matrix is it diagonal? """ tree = tsfc_compile_form(form, prefix=name, parameters=parameters, - interface=interface, + dont_split_numbers=dont_split_numbers, diagonal=diagonal, log=PETSc.Log.isActive()) kernels = [] for kernel in tree: + domain_number = domain_number_map[kernel.domain_number] + active_domain_numbers = ActiveDomainNumbers(*(tuple(domain_number_map[dn] for dn in dn_tuple) for dn_tuple in kernel.active_domain_numbers)) # Individual kernels do not have to use all of the coefficients # provided by the (split) form. Here we combine the numberings # of (kernel coefficients -> split form coefficients) and @@ -117,14 +121,13 @@ def __init__( events=events) kernels.append(KernelInfo(kernel=pyop2_kernel, integral_type=kernel.integral_type, - oriented=kernel.oriented, subdomain_id=kernel.subdomain_id, - domain_number=kernel.domain_number, + domain_number=domain_number, + active_domain_numbers=active_domain_numbers, coefficient_numbers=coefficient_numbers_per_kernel, constant_numbers=constant_numbers_per_kernel, needs_cell_facets=False, pass_layer_arg=False, - needs_cell_sizes=kernel.needs_cell_sizes, arguments=kernel.arguments, events=events)) self.kernels = tuple(kernels) @@ -156,7 +159,7 @@ def _compile_form_comm(*args, **kwargs): cachedir=_cachedir ) @PETSc.Log.EventDecorator() -def compile_form(form, name, parameters=None, split=True, interface=None, diagonal=False): +def compile_form(form, name, parameters=None, split=True, dont_split=None, diagonal=False): """Compile a form using TSFC. :arg form: the :class:`~ufl.classes.Form` to compile. @@ -193,6 +196,7 @@ def compile_form(form, name, parameters=None, split=True, interface=None, diagon kernels = [] numbering = form.terminal_numbering() + all_meshes = extract_domains(form) if split: iterable = split_form(form, diagonal=diagonal) else: @@ -208,22 +212,32 @@ def compile_form(form, name, parameters=None, split=True, interface=None, diagon # and that component doesn't actually appear in the form then we # have an empty form, which we should not attempt to assemble. continue - # Map local coefficient/constant numbers (as seen inside the + # Map local domain/coefficient/constant numbers (as seen inside the # compiler) to the global coefficient/constant numbers + meshes = extract_domains(f) + domain_number_map = tuple(all_meshes.index(m) for m in meshes) coefficient_numbers = tuple( numbering[c] for c in f.coefficients() ) constant_numbers = tuple( numbering[c] for c in extract_firedrake_constants(f) ) + if dont_split is not None: + dont_split_numbers = tuple(f.coefficients().index(c) # block-local numbering + for c in dont_split + if c in f.coefficients()) + else: + dont_split_numbers = () prefix = name + "".join(map(str, (i for i in idx if i is not None))) tsfc_kernel = TSFCKernel( f, prefix, parameters, + domain_number_map, coefficient_numbers, constant_numbers, - interface, diagonal + dont_split_numbers, + diagonal, ) for kinfo in tsfc_kernel.kernels: kernels.append(SplitKernel(idx, kinfo)) @@ -266,20 +280,21 @@ def _ensure_cachedir(comm=None): def gather_integer_subdomain_ids(knls): - """Gather a dict of all integer subdomain IDs per integral type. + """Gather a dict of all integer subdomain IDs per integral type per domain. This is needed to correctly interpret the ``"otherwise"`` subdomain ID. :arg knls: Iterable of :class:`SplitKernel` objects. """ - all_integer_subdomain_ids = collections.defaultdict(list) + all_integer_subdomain_ids = collections.defaultdict(lambda: collections.defaultdict(set)) for _, kinfo in knls: for subdomain_id in kinfo.subdomain_id: if subdomain_id != "otherwise": - all_integer_subdomain_ids[kinfo.integral_type].append(subdomain_id) + all_integer_subdomain_ids[kinfo.domain_number][kinfo.integral_type].add(subdomain_id) - for k, v in all_integer_subdomain_ids.items(): - all_integer_subdomain_ids[k] = tuple(sorted(v)) + for domain_number, integral_type_subdomain_ids_dict in all_integer_subdomain_ids.items(): + for integral_type, subdomain_ids in integral_type_subdomain_ids_dict.items(): + all_integer_subdomain_ids[domain_number][integral_type] = tuple(sorted(subdomain_ids)) return all_integer_subdomain_ids diff --git a/firedrake/ufl_expr.py b/firedrake/ufl_expr.py index 1f9a66df75..a569fdd2a9 100644 --- a/firedrake/ufl_expr.py +++ b/firedrake/ufl_expr.py @@ -5,9 +5,9 @@ from ufl.split_functions import split from ufl.algorithms import extract_arguments, extract_coefficients from ufl.domain import as_domain - import firedrake from firedrake import utils, function, cofunction +from firedrake.mesh import MixedMeshGeometry from firedrake.constant import Constant from firedrake.petsc import PETSc @@ -227,63 +227,64 @@ def derivative(form, u, du=None, coefficient_derivatives=None): raise TypeError( f"Cannot take the derivative of a {type(form).__name__}" ) - u_is_x = isinstance(u, ufl.SpatialCoordinate) - if u_is_x or isinstance(u, (Constant, BaseFormOperator)): - uc = u - else: - uc, = extract_coefficients(u) - if not (u_is_x or isinstance(u, BaseFormOperator)) and len(uc.subfunctions) > 1 and set(extract_coefficients(form)) & set(uc.subfunctions): - raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.subfunctions." - "\nYou probably meant to write split(u) when defining your form.") - - mesh = as_domain(form) - if not mesh: - raise ValueError("Expression to be differentiated has no ufl domain." - "\nDo you need to add a domain to your Constant?") - is_dX = u_is_x or u is mesh.coordinates - try: args = form.arguments() except AttributeError: args = extract_arguments(form) # UFL arguments need unique indices within a form n = max(a.number() for a in args) if args else -1 - - if is_dX: - coords = mesh.coordinates - u = ufl.SpatialCoordinate(mesh) + set_internal_coord_derivatives = False + all_meshes = extract_domains(form) + if isinstance(u, ufl.SpatialCoordinate): + uc = u + coords_mesh, = extract_unique_domain(u) + coords = coords_mesh.coordinates V = coords.function_space() - elif isinstance(uc, (firedrake.Function, firedrake.Cofunction, BaseFormOperator)): + set_internal_coord_derivatives = True + elif any(u is m.coordinates for m in all_meshes): + uc = u + coords = u + coord_mesh = u.function_space().mesh() + u = ufl.SpatialCoordinate(coord_mesh) + V = coords.function_space() + set_internal_coord_derivatives = True + elif isinstance(u, BaseFormOperator): + uc = u V = uc.function_space() - elif isinstance(uc, firedrake.Constant): + elif isinstance(u, Constant): + uc = u if uc.ufl_shape != (): raise ValueError("Real function space of vector elements not supported") # Replace instances of the constant with a new argument ``x`` # and differentiate wrt ``x``. + mesh = as_domain(form) # integration domain V = firedrake.FunctionSpace(mesh, "Real", 0) x = ufl.Coefficient(V) # TODO: Update this line when https://github.com/FEniCS/ufl/issues/171 is fixed form = ufl.replace(form, {u: x}) u_orig, u = u, x else: - raise RuntimeError("Can't compute derivative for form") - + uc, = extract_coefficients(u) + if not isinstance(uc, (firedrake.Function, firedrake.Cofunction)): + raise RuntimeError(f"Can't compute derivative for form w.r.t {u}") + if len(uc.subfunctions) > 1 and set(extract_coefficients(form)) & set(uc.subfunctions): + raise ValueError("Taking derivative of form wrt u, but form contains coefficients from u.subfunctions." + "\nYou probably meant to write split(u) when defining your form.") + V = uc.function_space() if du is None: du = Argument(V, n + 1) - - if is_dX: + if set_internal_coord_derivatives: internal_coefficient_derivatives = {coords: du} else: internal_coefficient_derivatives = {} if coefficient_derivatives: internal_coefficient_derivatives.update(coefficient_derivatives) - if u.ufl_shape != du.ufl_shape: raise ValueError("Shapes of u and du do not match.\n" "If you passed an indexed part of split(u) into " "derivative, you need to provide an appropriate du as well.") dform = ufl.derivative(form, u, du, internal_coefficient_derivatives) - if isinstance(uc, firedrake.Constant): + if isinstance(uc, Constant): # If we replaced constants with ``x`` to differentiate, # replace them back to the original symbolic constant dform = ufl.replace(dform, {u: u_orig}) @@ -366,23 +367,35 @@ def FacetNormal(mesh): return ufl.FacetNormal(mesh) -def extract_domains(func): - """Extract the domain from `func`. +def extract_domains(f): + """Extract the domain from `f`. Parameters ---------- - x : firedrake.function.Function, firedrake.cofunction.Cofunction, or firedrake.constant.Constant - The function to extract the domain from. + f : ufl.form.Form or firedrake.slate.TensorBase or firedrake.function.Function or firedrake.cofunction.Cofunction or firedrake.constant.Constant + The form, tensor, or function to extract the domain from. Returns ------- list of firedrake.mesh.MeshGeometry Extracted domains. """ - if isinstance(func, (function.Function, cofunction.Cofunction)): - return [func.function_space().mesh()] + if isinstance(f, firedrake.slate.TensorBase): + return f.ufl_domains() + elif isinstance(f, cofunction.Cofunction): + # ufl.domain.extract_domains does not work. + mesh = f.function_space().mesh() + if isinstance(mesh, MixedMeshGeometry): + return list(set(mesh._meshes)) + else: + return [mesh] + elif isinstance(f, (ufl.form.FormSum, ufl.Action)): + # ufl.domain.extract_domains does not work. + if f._domains is None: + f._analyze_domains() + return f._domains else: - return ufl.domain.extract_domains(func) + return ufl.domain.extract_domains(f) def extract_unique_domain(func): @@ -399,6 +412,6 @@ def extract_unique_domain(func): Extracted domains. """ if isinstance(func, (function.Function, cofunction.Cofunction)): - return func.function_space().mesh() + return func.function_space().mesh().unique() else: return ufl.domain.extract_unique_domain(func) diff --git a/tests/regression/test_function_spaces.py b/tests/regression/test_function_spaces.py index dbb81ec700..3bb20197fa 100644 --- a/tests/regression/test_function_spaces.py +++ b/tests/regression/test_function_spaces.py @@ -243,8 +243,8 @@ def test_reconstruct_variant(family, dual): def test_reconstruct_mixed(fs, mesh, mesh2, dual): W1 = fs.dual() if dual else fs W2 = W1.reconstruct(mesh=mesh2) - assert W1.mesh() == mesh - assert W2.mesh() == mesh2 + assert W1.mesh().unique() == mesh + assert W2.mesh().unique() == mesh2 assert W1.ufl_element() == W2.ufl_element() for index, V in enumerate(W1): V1 = W1.sub(index) diff --git a/tests/submesh/test_submesh_assemble.py b/tests/submesh/test_submesh_assemble.py new file mode 100644 index 0000000000..dd99076aa8 --- /dev/null +++ b/tests/submesh/test_submesh_assemble.py @@ -0,0 +1,258 @@ +import numpy as np +from firedrake import * + + +def test_submesh_assemble_cell_cell_integral_cell(): + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True) + x, y = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(conditional(x > 1., 1, 0)) + mesh.mark_entities(indicator_function, 999) + subm = Submesh(mesh, dim, 999) + V0 = FunctionSpace(mesh, "CG", 1) + V1 = FunctionSpace(subm, "CG", 1) + V = V0 * V1 + u = TrialFunction(V) + v = TestFunction(V) + u0, u1 = split(u) + v0, v1 = split(v) + dx0 = Measure("cell", domain=mesh) + dx1 = Measure("cell", domain=subm) + a = inner(u1, v0) * dx0(999) + inner(u0, v1) * dx1 + A = assemble(a, mat_type="nest") + assert np.allclose(A.M.sparsity[0][0].nnz, [1, 1, 1, 1, 1, 1]) # bc nodes + assert np.allclose(A.M.sparsity[0][1].nnz, [4, 4, 4, 4, 0, 0]) + assert np.allclose(A.M.sparsity[1][0].nnz, [4, 4, 4, 4]) + assert np.allclose(A.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + M10 = np.array([[1./9. , 1./18., 1./36., 1./18., 0., 0.], # noqa: E203 + [1./18., 1./9. , 1./18., 1./36., 0., 0.], # noqa: E203 + [1./36., 1./18., 1./9. , 1./18., 0., 0.], # noqa: E203 + [1./18., 1./36., 1./18., 1./9. , 0., 0.]]) # noqa: E203 + assert np.allclose(A.M[0][1].values, np.transpose(M10)) + assert np.allclose(A.M[1][0].values, M10) + + +def test_submesh_assemble_cell_cell_integral_facet(): + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True) + x, y = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(conditional(x > 1., 1, 0)) + mesh.mark_entities(indicator_function, 999) + subm = Submesh(mesh, dim, 999) + V0 = FunctionSpace(mesh, "DQ", 1, variant="equispaced") + V1 = FunctionSpace(subm, "DQ", 1, variant="equispaced") + V = V0 * V1 + u = TrialFunction(V) + v = TestFunction(V) + u0, u1 = split(u) + v0, v1 = split(v) + dS0 = Measure("dS", domain=mesh) + ds1 = Measure("ds", domain=subm) + a = inner(u1('|'), v0('+')) * dS0 + inner(u0('+'), v1('|')) * ds1(5) + A = assemble(a, mat_type="nest") + assert np.allclose(A.M.sparsity[0][0].nnz, [1, 1, 1, 1, 1, 1, 1, 1]) # bc nodes + assert np.allclose(A.M.sparsity[0][1].nnz, [4, 4, 4, 4, 4, 4, 4, 4]) + assert np.allclose(A.M.sparsity[1][0].nnz, [8, 8, 8, 8]) + assert np.allclose(A.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + M10 = np.array([[0., 0., 0., 0., 0., 0., 1. / 3., 1. / 6.], + [0., 0., 0., 0., 0., 0., 1. / 6., 1. / 3.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.]]) + assert np.allclose(A.M[0][1].values, np.transpose(M10)) + assert np.allclose(A.M[1][0].values, M10) + b = inner(u1('|'), v0('+')) * ds1(5) + inner(u0('+'), v1('|')) * dS0 + B = assemble(b, mat_type="nest") + assert np.allclose(B.M.sparsity[0][0].nnz, [1, 1, 1, 1, 1, 1, 1, 1]) # bc nodes + assert np.allclose(B.M.sparsity[0][1].nnz, [4, 4, 4, 4, 4, 4, 4, 4]) + assert np.allclose(B.M.sparsity[1][0].nnz, [8, 8, 8, 8]) + assert np.allclose(B.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + assert np.allclose(B.M[0][1].values, A.M[0][1].values) + assert np.allclose(B.M[1][0].values, A.M[1][0].values) + + +def test_submesh_assemble_cell_cell_cell_cell_integral_various(): + # +-------+-------+-------+-------+ + # | | | | | + # | | 555 | | mesh + # | | | | | + # +-------+-------+-------+-------+ + # +-------+-------+ + # | | | + # | | 555 mesh_l + # | | | + # +-------+-------+ + # +-------+-------+ + # | | | + # 555 | | mesh_r + # | | | + # +-------+-------+ + # +-------+ + # | | + # 555 | mesh_rl + # | | + # +-------+ + dim = 2 + mesh = RectangleMesh(4, 1, 4., 1., quadrilateral=True) + x, y = SpatialCoordinate(mesh) + label_int = 555 + label_l = 81100 + label_r = 80011 + label_rl = 80010 + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + DG0 = FunctionSpace(mesh, "DG", 0) + f_int = Function(HDivTrace0).interpolate(conditional(And(x > 1.9, x < 2.1), 1, 0)) + f_l = Function(DG0).interpolate(conditional(x < 2., 1, 0)) + f_r = Function(DG0).interpolate(conditional(x > 2., 1, 0)) + f_rl = Function(DG0).interpolate(conditional(And(x > 2., x < 3.), 1, 0)) + mesh = RelabeledMesh(mesh, [f_int, f_l, f_r, f_rl], [label_int, label_l, label_r, label_rl]) + x, y = SpatialCoordinate(mesh) + mesh_l = Submesh(mesh, dim, label_l) + mesh_r = Submesh(mesh, dim, label_r) + mesh_rl = Submesh(mesh_r, dim, label_rl) + dS = Measure("dS", domain=mesh) + ds_l = Measure("ds", domain=mesh_l) + ds_r = Measure("ds", domain=mesh_r) + ds_rl = Measure("ds", domain=mesh_rl) + n_l = FacetNormal(mesh_l) + n_rl = FacetNormal(mesh_rl) + assert assemble(dot(n_rl + n_l, n_rl + n_l) * ds_rl(label_int)) < 1.e-32 + assert assemble(dot(n_rl + n_l, n_rl + n_l) * ds_r(label_int)) < 1.e-32 + assert assemble(dot(n_rl + n_l, n_rl + n_l) * ds_l(label_int)) < 1.e-32 + assert assemble(dot(n_rl + n_l, n_rl + n_l) * dS(label_int)) < 1.e-32 + V_l = FunctionSpace(mesh_l, "DQ", 1, variant='equispaced') + V_rl = FunctionSpace(mesh_rl, "DQ", 1, variant='equispaced') + V = V_l * V_rl + u_l, u_rl = TrialFunctions(V) + v_l, v_rl = TestFunctions(V) + a = inner(u_rl('|'), v_l('|')) * ds_l(label_int) + inner(u_l('|'), v_rl('|')) * ds_rl(label_int) + A = assemble(a, mat_type="nest") + assert np.allclose(A.M.sparsity[0][0].nnz, [1, 1, 1, 1, 1, 1, 1, 1]) # bc nodes + assert np.allclose(A.M.sparsity[0][1].nnz, [4, 4, 4, 4, 0, 0, 0, 0]) + assert np.allclose(A.M.sparsity[1][0].nnz, [4, 4, 4, 4]) + assert np.allclose(A.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + M10 = np.array([[0., 0., 1. / 3., 1. / 6., 0., 0., 0., 0.], + [0., 0., 1. / 6., 1. / 3., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.]]) + assert np.allclose(A.M[0][1].values, np.transpose(M10)) + assert np.allclose(A.M[1][0].values, M10) + b = inner(u_rl('|'), v_l('|')) * dS(label_int) + inner(u_l('|'), v_rl('|')) * dS(label_int) + B = assemble(b, mat_type="nest") + assert np.allclose(B.M.sparsity[0][0].nnz, [1, 1, 1, 1, 1, 1, 1, 1]) # bc nodes + assert np.allclose(B.M.sparsity[0][1].nnz, [4, 4, 4, 4, 0, 0, 0, 0]) + assert np.allclose(B.M.sparsity[1][0].nnz, [4, 4, 4, 4]) + assert np.allclose(B.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + assert np.allclose(B.M[0][1].values, A.M[0][1].values) + assert np.allclose(B.M[1][0].values, A.M[1][0].values) + + +def test_submesh_assemble_cell_cell_cell_cell_integral_avg(): + # +-------+-------+-------+-------+ + # | | | | | + # | | 555 | | mesh + # | | | | | + # +-------+-------+-------+-------+ + # +-------+-------+-------+ + # | | | | + # | | 555 | mesh_l + # | | | | + # +-------+-------+-------+ + # +-------+-------+ + # | | | + # 555 | | mesh_r + # | | | + # +-------+-------+ + # +-------+ + # | | + # 555 | mesh_rl + # | | + # +-------+ + dim = 2 + mesh = RectangleMesh(4, 1, 4., 1., quadrilateral=True) + x, y = SpatialCoordinate(mesh) + label_int = 555 + label_l = 81110 + label_r = 80011 + label_rl = 80010 + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + DG0 = FunctionSpace(mesh, "DG", 0) + f_int = Function(HDivTrace0).interpolate(conditional(And(x > 1.9, x < 2.1), 1, 0)) + f_l = Function(DG0).interpolate(conditional(x < 3., 1, 0)) + f_r = Function(DG0).interpolate(conditional(x > 2., 1, 0)) + f_rl = Function(DG0).interpolate(conditional(And(x > 2., x < 3.), 1, 0)) + mesh = RelabeledMesh(mesh, [f_int, f_l, f_r, f_rl], [label_int, label_l, label_r, label_rl]) + x, y = SpatialCoordinate(mesh) + mesh_l = Submesh(mesh, dim, label_l) + x_l, y_l = SpatialCoordinate(mesh_l) + mesh_r = Submesh(mesh, dim, label_r) + x_r, y_r = SpatialCoordinate(mesh_r) + mesh_rl = Submesh(mesh_r, dim, label_rl) + x_rl, y_rl = SpatialCoordinate(mesh_rl) + dx = Measure("dx", domain=mesh) + dx_l = Measure("dx", domain=mesh_l) + dx_rl = Measure("dx", domain=mesh_rl) + dS = Measure("dS", domain=mesh) + dS_l = Measure("dS", domain=mesh_l) + ds_rl = Measure("ds", domain=mesh_rl) + assert abs(assemble(cell_avg(x) * dx(label_rl)) - 2.5) < 5.e-16 + assert abs(assemble(cell_avg(x) * dx_rl) - 2.5) < 5.e-16 + assert abs(assemble(cell_avg(x_rl) * dx(label_rl)) - 2.5) < 5.e-16 + assert abs(assemble(cell_avg(x_rl) * dx_l(label_rl)) - 2.5) < 5.e-16 + assert abs(assemble(cell_avg(x_l) * dx_rl) - 2.5) < 5.e-16 + assert abs(assemble(facet_avg(y * y) * dS(label_int)) - 1. / 3.) < 5.e-16 + assert abs(assemble(facet_avg(y('+') * y('-')) * ds_rl(label_int)) - 1. / 3.) < 5.e-16 + assert abs(assemble(facet_avg(y_rl('|') * y_rl('|')) * dS(label_int)) - 1. / 3.) < 5.e-16 + assert abs(assemble(facet_avg(y_rl('|') * y_rl('|')) * dS_l(label_int)) - 1. / 3.) < 5.e-16 + assert abs(assemble(facet_avg(y_l('+') * y_l('-')) * ds_rl(label_int)) - 1. / 3.) < 5.e-16 + + +def test_submesh_assemble_cell_cell_equation_bc(): + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True) + x, y = SpatialCoordinate(mesh) + label_int = 555 + label_l = 810 + label_r = 801 + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + DQ0 = FunctionSpace(mesh, "DQ", 0) + f_int = Function(HDivTrace0).interpolate(conditional(And(x > 0.9, x < 1.1), 1, 0)) + f_l = Function(DQ0).interpolate(conditional(x < 1., 1, 0)) + f_r = Function(DQ0).interpolate(conditional(x > 1., 1, 0)) + mesh = RelabeledMesh(mesh, [f_int, f_l, f_r], [label_int, label_l, label_r]) + mesh_l = Submesh(mesh, dim, label_l) + mesh_r = Submesh(mesh, dim, label_r) + V_l = FunctionSpace(mesh_l, "CG", 1) + V_r = FunctionSpace(mesh_r, "CG", 1) + V = V_l * V_r + u = TrialFunction(V) + v = TestFunction(V) + u_l, u_r = split(u) + v_l, v_r = split(v) + dx_l = Measure("dx", domain=mesh_l) + ds_l = Measure("ds", domain=mesh_l) + a = inner(u_l, v_l) * dx_l + a_int = inner(u_l - u_r, v_l) * ds_l(label_int) + L_int = inner(Constant(0), v_l) * ds_l(label_int) + sol = Function(V) + bc = EquationBC(a_int == L_int, sol, label_int, V=V.sub(0)) + A = assemble(a, bcs=bc.extract_form('J'), mat_type="nest") + assert np.allclose(Function(V_l).interpolate(SpatialCoordinate(mesh_l)[0]).dat.data, [0., 0., 1., 1.]) + assert np.allclose(Function(V_l).interpolate(SpatialCoordinate(mesh_l)[1]).dat.data, [0., 1., 1., 0.]) + assert np.allclose(Function(V_r).interpolate(SpatialCoordinate(mesh_r)[0]).dat.data, [1., 1., 2., 2.]) + assert np.allclose(Function(V_r).interpolate(SpatialCoordinate(mesh_r)[1]).dat.data, [0., 1., 1., 0.]) + assert np.allclose(A.M.sparsity[0][0].nnz, [4, 4, 4, 4]) + assert np.allclose(A.M.sparsity[0][1].nnz, [4, 4, 4, 4]) + assert np.allclose(A.M.sparsity[1][0].nnz, [0, 0, 0, 0]) + assert np.allclose(A.M.sparsity[1][1].nnz, [1, 1, 1, 1]) # bc nodes + M00 = np.array([[1. / 9. , 1. / 18., 1. / 36., 1. / 18.], # noqa: E203 + [1. / 18., 1. / 9. , 1. / 18., 1. / 36.], # noqa: E203 + [0., 0., 1. / 3., 1. / 6.], + [0., 0., 1. / 6., 1. / 3.]]) + M01 = np.array([[0., 0., 0., 0.], + [0., 0., 0., 0.], + [- 1. / 6., - 1. / 3., 0., 0.], + [- 1. / 3., - 1. / 6., 0., 0.]]) + assert np.allclose(A.M[0][0].values, M00) + assert np.allclose(A.M[0][1].values, M01) diff --git a/tests/submesh/test_submesh_base.py b/tests/submesh/test_submesh_base.py new file mode 100644 index 0000000000..1c125e8e05 --- /dev/null +++ b/tests/submesh/test_submesh_base.py @@ -0,0 +1,289 @@ +import pytest +from firedrake import * + + +def _get_expr(m): + if m.geometric_dimension() == 1: + x, = SpatialCoordinate(m) + y = x * x + z = x + y + elif m.geometric_dimension() == 2: + x, y = SpatialCoordinate(m) + z = x + y + elif m.geometric_dimension() == 3: + x, y, z = SpatialCoordinate(m) + else: + raise NotImplementedError("Not implemented") + return exp(x + y * y + z * z * z) + + +def _test_submesh_base_cell_integral_quad(family_degree, nelem): + dim = 2 + family, degree = family_degree + mesh = UnitSquareMesh(nelem, nelem, quadrilateral=True) + V = FunctionSpace(mesh, family, degree) + f = Function(V).interpolate(_get_expr(mesh)) + x, y = SpatialCoordinate(mesh) + cond = conditional(x > .5, 1, + conditional(y > .5, 1, 0)) # noqa: E128 + target = assemble(f * cond * dx) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(cond) + label_value = 999 + mesh.mark_entities(indicator_function, label_value) + msub = Submesh(mesh, dim, label_value) + Vsub = FunctionSpace(msub, family, degree) + fsub = Function(Vsub).interpolate(_get_expr(msub)) + result = assemble(fsub * dx) + assert abs(result - target) < 1e-12 + + +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_cell_integral_quad_1_process(family_degree, nelem): + _test_submesh_base_cell_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_cell_integral_quad_2_processes(family_degree, nelem): + _test_submesh_base_cell_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_cell_integral_quad_3_processes(family_degree, nelem): + _test_submesh_base_cell_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_cell_integral_quad_4_processes(family_degree, nelem): + _test_submesh_base_cell_integral_quad(family_degree, nelem) + + +def _test_submesh_base_facet_integral_quad(family_degree, nelem): + dim = 2 + family, degree = family_degree + mesh = UnitSquareMesh(nelem, nelem, quadrilateral=True) + x, y = SpatialCoordinate(mesh) + cond = conditional(x > .5, 1, + conditional(y > .5, 1, 0)) # noqa: E128 + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(cond) + label_value = 999 + mesh.mark_entities(indicator_function, label_value) + subm = Submesh(mesh, dim, label_value) + for i in [1, 2, 3, 4]: + target = assemble(cond * _get_expr(mesh) * ds(i)) + result = assemble(_get_expr(subm) * ds(i)) + assert abs(result - target) < 2e-12 + # Check new boundary. + assert abs(assemble(Constant(1.) * ds(subdomain_id=5, domain=subm)) - 1.0) < 1e-12 + x, y = SpatialCoordinate(subm) + assert abs(assemble(x**4 * ds(5)) - (.5**5 / 5 + .5**4 * .5)) < 1e-12 + assert abs(assemble(y**4 * ds(5)) - (.5**5 / 5 + .5**4 * .5)) < 1e-12 + + +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_facet_integral_quad_1_process(family_degree, nelem): + _test_submesh_base_facet_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_facet_integral_quad_2_processes(family_degree, nelem): + _test_submesh_base_facet_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_facet_integral_quad_3_processes(family_degree, nelem): + _test_submesh_base_facet_integral_quad(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8, 16]) +def test_submesh_base_facet_integral_quad_4_processes(family_degree, nelem): + _test_submesh_base_facet_integral_quad(family_degree, nelem) + + +def _test_submesh_base_cell_integral_hex(family_degree, nelem): + dim = 3 + family, degree = family_degree + mesh = UnitCubeMesh(nelem, nelem, nelem, hexahedral=True) + V = FunctionSpace(mesh, family, degree) + f = Function(V).interpolate(_get_expr(mesh)) + x, y, z = SpatialCoordinate(mesh) + cond = conditional(x > .5, 1, + conditional(y > .5, 1, # noqa: E128 + conditional(z > .5, 1, 0))) # noqa: E128 + target = assemble(f * cond * dx) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(cond) + label_value = 999 + mesh.mark_entities(indicator_function, label_value) + msub = Submesh(mesh, dim, label_value) + Vsub = FunctionSpace(msub, family, degree) + fsub = Function(Vsub).interpolate(_get_expr(msub)) + result = assemble(fsub * dx) + assert abs(result - target) < 1e-12 + + +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_cell_integral_hex_1_process(family_degree, nelem): + _test_submesh_base_cell_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_cell_integral_hex_2_processes(family_degree, nelem): + _test_submesh_base_cell_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_cell_integral_hex_3_processes(family_degree, nelem): + _test_submesh_base_cell_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('family_degree', [("Q", 4), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_cell_integral_hex_4_processes(family_degree, nelem): + _test_submesh_base_cell_integral_hex(family_degree, nelem) + + +def _test_submesh_base_facet_integral_hex(family_degree, nelem): + dim = 3 + family, degree = family_degree + mesh = UnitCubeMesh(nelem, nelem, nelem, hexahedral=True) + x, y, z = SpatialCoordinate(mesh) + cond = conditional(x > .5, 1, + conditional(y > .5, 1, # noqa: E128 + conditional(z > .5, 1, 0))) # noqa: E128 + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(cond) + label_value = 999 + mesh.mark_entities(indicator_function, label_value) + subm = Submesh(mesh, dim, label_value) + for i in [1, 2, 3, 4, 5, 6]: + target = assemble(cond * _get_expr(mesh) * ds(i)) + result = assemble(_get_expr(subm) * ds(i)) + assert abs(result - target) < 2e-12 + # Check new boundary. + assert abs(assemble(Constant(1) * ds(subdomain_id=7, domain=subm)) - .75) < 1e-12 + x, y, z = SpatialCoordinate(subm) + assert abs(assemble(x**4 * ds(7)) - (.5**5 / 5 * .5 * 2 + .5**4 * .5**2)) < 1e-12 + assert abs(assemble(y**4 * ds(7)) - (.5**5 / 5 * .5 * 2 + .5**4 * .5**2)) < 1e-12 + assert abs(assemble(z**4 * ds(7)) - (.5**5 / 5 * .5 * 2 + .5**4 * .5**2)) < 1e-12 + + +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_facet_integral_hex_1_process(family_degree, nelem): + _test_submesh_base_facet_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=2) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_facet_integral_hex_2_processes(family_degree, nelem): + _test_submesh_base_facet_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_facet_integral_hex_3_processes(family_degree, nelem): + _test_submesh_base_facet_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('family_degree', [("Q", 3), ]) +@pytest.mark.parametrize('nelem', [2, 4, 8]) +def test_submesh_base_facet_integral_hex_4_processes(family_degree, nelem): + _test_submesh_base_facet_integral_hex(family_degree, nelem) + + +@pytest.mark.parallel(nprocs=2) +def test_submesh_base_entity_maps(): + + # 3---9--(5)-(12)(7) (7)-(13)-3---9---5 + # | | | | | | + # 8 0 (11) (1) (13) (12) (1) 8 0 10 mesh + # | | | | | | + # 2--10--(4)(14)-(6) (6)-(14)-2--11---4 + # + # 2---6---4 (4)-(7)-(2) + # | | | | + # 5 0 8 (6) (0) (5) submesh + # | | | | + # 1---7---3 (3)-(8)-(1) + # + # rank 0 rank 1 + + dim = 2 + mesh = RectangleMesh(2, 1, 2., 1., quadrilateral=True, distribution_parameters={"partitioner_type": "simple"}) + assert mesh.comm.size == 2 + rank = mesh.comm.rank + x, y = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(conditional(x < 1., 1, 0)) + label_value = 999 + mesh.mark_entities(indicator_function, label_value) + submesh = Submesh(mesh, dim, label_value) + submesh.topology_dm.viewFromOptions("-dm_view") + subdm = submesh.topology.topology_dm + if rank == 0: + assert subdm.getLabel("pyop2_core").getStratumSize(1) == 0 + assert subdm.getLabel("pyop2_owned").getStratumSize(1) == 9 + assert subdm.getLabel("pyop2_ghost").getStratumSize(1) == 0 + assert (subdm.getLabel("pyop2_owned").getStratumIS(1).getIndices() == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])).all() + assert (mesh.interior_facets.facets == np.array([11])).all + assert (mesh.exterior_facets.facets == np.array([8, 9, 10, 12, 13, 14])).all + assert (submesh.interior_facets.facets == np.array([])).all + assert (submesh.exterior_facets.facets == np.array([5, 6, 8, 7])).all() + else: + assert subdm.getLabel("pyop2_core").getStratumSize(1) == 0 + assert subdm.getLabel("pyop2_owned").getStratumSize(1) == 0 + assert subdm.getLabel("pyop2_ghost").getStratumSize(1) == 9 + assert (subdm.getLabel("pyop2_ghost").getStratumIS(1).getIndices() == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])).all() + assert (mesh.interior_facets.facets == np.array([8])).all + assert (mesh.exterior_facets.facets == np.array([9, 10, 11, 12, 13, 14])).all + assert (submesh.interior_facets.facets == np.array([])).all + assert (submesh.exterior_facets.facets == np.array([6, 7, 5, 8])).all() + composed_map, integral_type = mesh.topology.trans_mesh_entity_map(submesh.topology, "cell", None, None) + assert integral_type == "cell" + if rank == 0: + assert (composed_map.maps_[0].values_with_halo == np.array([0])).all() + else: + assert (composed_map.maps_[0].values_with_halo == np.array([1])).all() + composed_map, integral_type = mesh.topology.trans_mesh_entity_map(submesh.topology, "exterior_facet", 5, None) + assert integral_type == "interior_facet" + if rank == 0: + assert (composed_map.maps_[0].values_with_halo == np.array([-1, -1, 0, -1]).reshape((-1, 1))).all() # entire exterior-interior map + else: + assert (composed_map.maps_[0].values_with_halo == np.array([-1, -1, 0, -1]).reshape((-1, 1))).all() # entire exterior-interior map + composed_map, integral_type = mesh.topology.trans_mesh_entity_map(submesh.topology, "exterior_facet", 4, None) + assert integral_type == "exterior_facet" + if rank == 0: + assert (composed_map.maps_[0].values_with_halo == np.array([0, 1, -1, 2]).reshape((-1, 1))).all() # entire exterior-exterior map + else: + assert (composed_map.maps_[0].values_with_halo == np.array([3, 4, -1, 5]).reshape((-1, 1))).all() # entire exterior-exterior map + composed_map, integral_type = submesh.topology.trans_mesh_entity_map(mesh.topology, "exterior_facet", 1, None) + assert integral_type == "exterior_facet" + if rank == 0: + assert (composed_map.maps_[0].values_with_halo == np.array([0, 1, 3, -1, -1, -1]).reshape((-1, 1))).all() + else: + assert (composed_map.maps_[0].values_with_halo == np.array([-1, -1, -1, 0, 1, 3]).reshape((-1, 1))).all() diff --git a/tests/submesh/test_submesh_solve.py b/tests/submesh/test_submesh_solve.py new file mode 100644 index 0000000000..361df30910 --- /dev/null +++ b/tests/submesh/test_submesh_solve.py @@ -0,0 +1,459 @@ +import pytest +from os.path import abspath, dirname, join +import numpy as np +from firedrake import * + + +cwd = abspath(dirname(__file__)) + + +def _solve_helmholtz(mesh): + V = FunctionSpace(mesh, "CG", 1) + u = TrialFunction(V) + v = TestFunction(V) + x = SpatialCoordinate(mesh) + u_exact = sin(x[0]) * sin(x[1]) + f = Function(V).interpolate(2 * u_exact) + a = (inner(grad(u), grad(v)) + inner(u, v)) * dx + L = inner(f, v) * dx + bc = DirichletBC(V, u_exact, "on_boundary") + sol = Function(V) + solve(a == L, sol, bcs=[bc], solver_parameters={'ksp_type': 'preonly', + 'pc_type': 'lu'}) + return sqrt(assemble((sol - u_exact)**2 * dx)) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('nelem', [2, 4]) +@pytest.mark.parametrize('distribution_parameters', [None, {"overlap_type": (DistributedMeshOverlapType.NONE, 0)}]) +def test_submesh_solve_simple(nelem, distribution_parameters): + dim = 2 + # Compute reference error. + mesh = RectangleMesh(nelem, nelem * 2, 1., 1., quadrilateral=True, distribution_parameters=distribution_parameters) + error = _solve_helmholtz(mesh) + # Compute submesh error. + mesh = RectangleMesh(nelem * 2, nelem * 2, 2., 1., quadrilateral=True, distribution_parameters=distribution_parameters) + x, y = SpatialCoordinate(mesh) + DQ0 = FunctionSpace(mesh, "DQ", 0) + indicator_function = Function(DQ0).interpolate(conditional(x < 1., 1, 0)) + mesh.mark_entities(indicator_function, 999) + mesh = Submesh(mesh, dim, 999) + suberror = _solve_helmholtz(mesh) + assert abs(error - suberror) < 1e-15 + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('dim', [2, 3]) +@pytest.mark.parametrize('simplex', [True, False]) +def test_submesh_solve_cell_cell_mixed_scalar(dim, simplex): + if dim == 2: + if simplex: + mesh = Mesh("./docs/notebooks/stokes-control.msh") + bid = (1, 2, 3, 4, 5) + submesh_expr = lambda x: conditional(x[0] < 10., 1, 0) + solution_expr = lambda x: x[0] + x[1] + else: + mesh = Mesh(join(cwd, "..", "meshes", "unitsquare_unstructured_quadrilaterals.msh")) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + x, y = SpatialCoordinate(mesh) + hdivtrace0x = Function(HDivTrace0).interpolate(conditional(And(x > .001, x < .999), 0, 1)) + hdivtrace0y = Function(HDivTrace0).interpolate(conditional(And(y > .001, y < .999), 0, 1)) + mesh = RelabeledMesh(mesh, [hdivtrace0x, hdivtrace0y], [111, 222]) + bid = (111, 222) + submesh_expr = lambda x: conditional(x[0] < .5, 1, 0) + solution_expr = lambda x: x[0] + x[1] + elif dim == 3: + if simplex: + nref = 3 + mesh = BoxMesh(2 ** nref, 2 ** nref, 2 ** nref, 1., 1., 1., hexahedral=False) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + else: + mesh = Mesh(join(cwd, "..", "meshes", "cube_hex.msh")) + HDivTrace0 = FunctionSpace(mesh, "Q", 2) + x, y, z = SpatialCoordinate(mesh) + hdivtrace0x = Function(HDivTrace0).interpolate(conditional(And(x > .001, x < .999), 0, 1)) + hdivtrace0y = Function(HDivTrace0).interpolate(conditional(And(y > .001, y < .999), 0, 1)) + hdivtrace0z = Function(HDivTrace0).interpolate(conditional(And(z > .001, z < .999), 0, 1)) + mesh = RelabeledMesh(mesh, [hdivtrace0x, hdivtrace0y, hdivtrace0z], [111, 222, 333]) + bid = (111, 222, 333) + submesh_expr = lambda x: conditional(x[0] > .5, 1, 0) + solution_expr = lambda x: x[0] + x[1] + x[2] + else: + raise NotImplementedError + DG0 = FunctionSpace(mesh, "DG", 0) + submesh_function = Function(DG0).interpolate(submesh_expr(SpatialCoordinate(mesh))) + submesh_label = 999 + mesh.mark_entities(submesh_function, submesh_label) + subm = Submesh(mesh, dim, submesh_label) + V0 = FunctionSpace(mesh, "CG", 2) + V1 = FunctionSpace(subm, "CG", 3) + V = V0 * V1 + u = TrialFunction(V) + v = TestFunction(V) + u0, u1 = split(u) + v0, v1 = split(v) + dx0 = Measure("cell", domain=mesh) + dx1 = Measure("cell", domain=subm) + a = inner(grad(u0), grad(v0)) * dx0 + inner(u0 - u1, v1) * dx1 + L = inner(Constant(0.), v1) * dx1 + g = Function(V0).interpolate(solution_expr(SpatialCoordinate(mesh))) + bc = DirichletBC(V.sub(0), g, bid) + solution = Function(V) + solve(a == L, solution, bcs=[bc]) + target = Function(V1).interpolate(solution_expr(SpatialCoordinate(subm))) + assert np.allclose(solution.subfunctions[1].dat.data_ro_with_halos, target.dat.data_ro_with_halos) + + +@pytest.mark.parallel(nprocs=3) +@pytest.mark.parametrize('dim', [2, 3]) +@pytest.mark.parametrize('simplex', [True, False]) +def test_submesh_solve_cell_cell_mixed_vector(dim, simplex): + if dim == 2: + if simplex: + mesh = Mesh("./docs/notebooks/stokes-control.msh") + submesh_expr = lambda x: conditional(x[0] < 10., 1, 0) + elem0 = FiniteElement("RT", "triangle", 3) + elem1 = VectorElement("P", "triangle", 3) + else: + mesh = Mesh(join(cwd, "..", "meshes", "unitsquare_unstructured_quadrilaterals.msh")) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + x, y = SpatialCoordinate(mesh) + hdivtrace0x = Function(HDivTrace0).interpolate(conditional(And(x > .001, x < .999), 0, 1)) + hdivtrace0y = Function(HDivTrace0).interpolate(conditional(And(y > .001, y < .999), 0, 1)) + mesh = RelabeledMesh(mesh, [hdivtrace0x, hdivtrace0y], [111, 222]) + submesh_expr = lambda x: conditional(x[0] < .5, 1, 0) + elem0 = FiniteElement("RTCF", "quadrilateral", 2) + elem1 = VectorElement("Q", "quadrilateral", 3) + elif dim == 3: + if simplex: + nref = 3 + mesh = BoxMesh(2 ** nref, 2 ** nref, 2 ** nref, 1., 1., 1., hexahedral=False) + x, y, z = SpatialCoordinate(mesh) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + hdivtrace0x = Function(HDivTrace0).interpolate(conditional(And(x > .001, x < .999), 0, 1)) + hdivtrace0y = Function(HDivTrace0).interpolate(conditional(And(y > .001, y < .999), 0, 1)) + hdivtrace0z = Function(HDivTrace0).interpolate(conditional(And(z > .001, z < .999), 0, 1)) + mesh = RelabeledMesh(mesh, [hdivtrace0x, hdivtrace0y, hdivtrace0z], [111, 222, 333]) + submesh_expr = lambda x: conditional(x[0] > .5, 1, 0) + elem0 = FiniteElement("N1F", "tetrahedron", 3) + elem1 = VectorElement("P", "tetrahedron", 3) + else: + mesh = Mesh(join(cwd, "..", "meshes", "cube_hex.msh")) + HDivTrace0 = FunctionSpace(mesh, "Q", 2) + x, y, z = SpatialCoordinate(mesh) + hdivtrace0x = Function(HDivTrace0).interpolate(conditional(And(x > .001, x < .999), 0, 1)) + hdivtrace0y = Function(HDivTrace0).interpolate(conditional(And(y > .001, y < .999), 0, 1)) + hdivtrace0z = Function(HDivTrace0).interpolate(conditional(And(z > .001, z < .999), 0, 1)) + mesh = RelabeledMesh(mesh, [hdivtrace0x, hdivtrace0y, hdivtrace0z], [111, 222, 333]) + submesh_expr = lambda x: conditional(x[0] > .5, 1, 0) + elem0 = FiniteElement("NCF", "hexahedron", 2) + elem1 = VectorElement("Q", "hexahedron", 3) + with pytest.raises(NotImplementedError): + _ = FunctionSpace(mesh, elem0) + return + else: + raise NotImplementedError + DG0 = FunctionSpace(mesh, "DG", 0) + submesh_function = Function(DG0).interpolate(submesh_expr(SpatialCoordinate(mesh))) + submesh_label = 999 + mesh.mark_entities(submesh_function, submesh_label) + subm = Submesh(mesh, dim, submesh_label) + V0 = FunctionSpace(mesh, elem0) + V1 = FunctionSpace(subm, elem1) + V = V0 * V1 + u = TrialFunction(V) + v = TestFunction(V) + u0, u1 = split(u) + v0, v1 = split(v) + dx0 = Measure("cell", domain=mesh) + dx1 = Measure("cell", domain=subm) + a = inner(u0, v0) * dx0 + inner(u0 - u1, v1) * dx1 + L = inner(SpatialCoordinate(mesh), v0) * dx0 + solution = Function(V) + solve(a == L, solution) + s0, s1 = split(solution) + x = SpatialCoordinate(subm) + assert assemble(inner(s1 - x, s1 - x) * dx1) < 1.e-20 + + +def _mixed_poisson_create_mesh_2d(nref, quadrilateral, submesh_region, label_submesh, label_submesh_compl): + # y + # | + # | + # 1.0 +--17---+--18---+ + # | | | + # 12 20 14 + # | | | + # 0.5 +--21---+--22---+ + # | | | + # 11 19 13 + # | | | + # 0.0 +--15---+--16---+----x + # + # 0.0 0.5 1.0 + mesh = UnitSquareMesh(2 ** nref, 2 ** nref, quadrilateral=quadrilateral) + eps = 1. / (2 ** nref) / 100. + x, y = SpatialCoordinate(mesh) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + f11 = Function(HDivTrace0).interpolate(conditional(And(x < eps, y < .5), 1, 0)) + f12 = Function(HDivTrace0).interpolate(conditional(And(x < eps, y > .5), 1, 0)) + f13 = Function(HDivTrace0).interpolate(conditional(And(x > 1 - eps, y < .5), 1, 0)) + f14 = Function(HDivTrace0).interpolate(conditional(And(x > 1 - eps, y > .5), 1, 0)) + f15 = Function(HDivTrace0).interpolate(conditional(And(x < .5, y < eps), 1, 0)) + f16 = Function(HDivTrace0).interpolate(conditional(And(x > .5, y < eps), 1, 0)) + f17 = Function(HDivTrace0).interpolate(conditional(And(x < .5, y > 1 - eps), 1, 0)) + f18 = Function(HDivTrace0).interpolate(conditional(And(x > .5, y > 1 - eps), 1, 0)) + f19 = Function(HDivTrace0).interpolate(conditional(And(And(x > .5 - eps, x < .5 + eps), y < .5), 1, 0)) + f20 = Function(HDivTrace0).interpolate(conditional(And(And(x > .5 - eps, x < .5 + eps), y > .5), 1, 0)) + f21 = Function(HDivTrace0).interpolate(conditional(And(x < .5, And(y > .5 - eps, y < .5 + eps)), 1, 0)) + f22 = Function(HDivTrace0).interpolate(conditional(And(x > .5, And(y > .5 - eps, y < .5 + eps)), 1, 0)) + DG0 = FunctionSpace(mesh, "DG", 0) + if submesh_region == "left": + submesh_function = Function(DG0).interpolate(conditional(x < .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(x > .5, 1, 0)) + elif submesh_region == "right": + submesh_function = Function(DG0).interpolate(conditional(x > .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(x < .5, 1, 0)) + elif submesh_region == "bottom": + submesh_function = Function(DG0).interpolate(conditional(y < .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(y > .5, 1, 0)) + elif submesh_region == "top": + submesh_function = Function(DG0).interpolate(conditional(y > .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(y < .5, 1, 0)) + else: + raise NotImplementedError(f"Unknown submesh_region: {submesh_region}") + return RelabeledMesh(mesh, [f11, f12, f13, f14, f15, f16, f17, f18, f19, f20, f21, f22, submesh_function, submesh_function_compl], + [11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, label_submesh, label_submesh_compl]) + + +def _mixed_poisson_solve_2d(nref, degree, quadrilateral, submesh_region): + dim = 2 + label_submesh = 999 + label_submesh_compl = 888 + mesh = _mixed_poisson_create_mesh_2d(nref, quadrilateral, submesh_region, label_submesh, label_submesh_compl) + x, y = SpatialCoordinate(mesh) + subm = Submesh(mesh, dim, label_submesh) + subx, suby = SpatialCoordinate(subm) + if submesh_region == "left": + boun_ext = (11, 12) + boun_int = (19, 20) + boun_dirichlet = (15, 17) + elif submesh_region == "right": + boun_ext = (13, 14) + boun_int = (19, 20) + boun_dirichlet = (16, 18) + elif submesh_region == "bottom": + boun_ext = (15, 16) + boun_int = (21, 22) + boun_dirichlet = (11, 13) + elif submesh_region == "top": + boun_ext = (17, 18) + boun_int = (21, 22) + boun_dirichlet = (12, 14) + else: + raise NotImplementedError(f"Unknown submesh_region: {submesh_region}") + BDM = FunctionSpace(subm, "RTCF" if quadrilateral else "BDM", degree) + DG = FunctionSpace(mesh, "DG", degree - 1) + W = BDM * DG + tau, v = TestFunctions(W) + nsub = FacetNormal(subm) + u_exact = Function(DG).interpolate(cos(2 * pi * x) * cos(2 * pi * y)) + sigma_exact = Function(BDM).project(as_vector([- 2 * pi * sin(2 * pi * subx) * cos(2 * pi * suby), - 2 * pi * cos(2 * pi * subx) * sin(2 * pi * suby)]), + solver_parameters={"ksp_type": "cg", "ksp_rtol": 1.e-16}) + f = Function(DG).interpolate(- 8 * pi * pi * cos(2 * pi * x) * cos(2 * pi * y)) + dx0 = Measure("dx", domain=mesh) + dx1 = Measure("dx", domain=subm) + ds0 = Measure("ds", domain=mesh) + ds1 = Measure("ds", domain=subm) + dS0 = Measure("dS", domain=mesh) + bc = DirichletBC(W.sub(0), sigma_exact, boun_dirichlet) + # Do the base case. + w = Function(W) + sigma, u = split(w) + a = (inner(sigma, tau) + inner(u, div(tau)) + inner(div(sigma), v)) * dx1 + inner(u - u_exact, v) * dx0(label_submesh_compl) + L = inner(f, v) * dx1 + inner((u('+') + u('-')) / 2., dot(tau('|'), nsub('|'))) * dS0(boun_int) + inner(u_exact('|'), dot(tau('|'), nsub('|'))) * ds0(boun_ext) + solve(a - L == 0, w, bcs=[bc]) + # Change domains of integration. + w_ = Function(W) + sigma_, u_ = split(w_) + a_ = (inner(sigma_, tau) + inner(u_, div(tau)) + inner(div(sigma_), v)) * dx1 + inner(u_ - u_exact, v) * dx0(label_submesh_compl) + L_ = inner(f, v) * dx0(label_submesh) + inner((u_('+') + u_('-')) / 2., dot(tau('|'), nsub('|'))) * ds1(boun_int) + inner(u_exact('|'), dot(tau('|'), nsub('|'))) * ds1(boun_ext) + solve(a_ - L_ == 0, w_, bcs=[bc]) + assert assemble(inner(sigma_ - sigma, sigma_ - sigma) * dx1) < 1.e-20 + assert assemble(inner(u_ - u, u_ - u) * dx0(label_submesh)) < 1.e-20 + sigma_error = sqrt(assemble(inner(sigma - sigma_exact, sigma - sigma_exact) * dx1)) + u_error = sqrt(assemble(inner(u - u_exact, u - u_exact) * dx0(label_submesh))) + return sigma_error, u_error + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('nref', [1, 2, 3, 4]) +@pytest.mark.parametrize('degree', [1]) +@pytest.mark.parametrize('quadrilateral', [False, True]) +@pytest.mark.parametrize('submesh_region', ["left", "right", "bottom", "top"]) +def test_submesh_solve_mixed_poisson_check_sanity_2d(nref, degree, quadrilateral, submesh_region): + _, _ = _mixed_poisson_solve_2d(nref, degree, quadrilateral, submesh_region) + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('quadrilateral', [True]) +@pytest.mark.parametrize('degree', [3]) +@pytest.mark.parametrize('submesh_region', ["left", "right"]) +def test_submesh_solve_mixed_poisson_check_convergence_2d(quadrilateral, degree, submesh_region): + nrefs = [5, 6, 7] + start = nrefs[0] + s_error_array = np.zeros(len(nrefs)) + u_error_array = np.zeros(len(nrefs)) + for nref in nrefs: + i = nref - start + s_error_array[i], u_error_array[i] = _mixed_poisson_solve_2d(nref, degree, quadrilateral, submesh_region) + assert (np.log2(s_error_array[:-1] / s_error_array[1:]) > degree + .95).all() + assert (np.log2(u_error_array[:-1] / u_error_array[1:]) > degree + .95).all() + + +def _mixed_poisson_create_mesh_3d(hexahedral, submesh_region, label_submesh, label_submesh_compl): + if hexahedral: + mesh = Mesh(join(cwd, "..", "meshes", "cube_hex.msh")) + DG0 = FunctionSpace(mesh, "DQ", 0) + HDivTrace0 = FunctionSpace(mesh, "Q", 2) + else: + mesh = BoxMesh(4, 4, 4, 1., 1., 1., hexahedral=False) + DG0 = FunctionSpace(mesh, "DP", 0) + HDivTrace0 = FunctionSpace(mesh, "HDiv Trace", 0) + x, y, z = SpatialCoordinate(mesh) + eps = 1.e-6 + f101 = Function(HDivTrace0).interpolate(conditional(x < eps, 1, 0)) + f102 = Function(HDivTrace0).interpolate(conditional(x > 1. - eps, 1, 0)) + f103 = Function(HDivTrace0).interpolate(conditional(y < eps, 1, 0)) + f104 = Function(HDivTrace0).interpolate(conditional(y > 1. - eps, 1, 0)) + f105 = Function(HDivTrace0).interpolate(conditional(z < eps, 1, 0)) + f106 = Function(HDivTrace0).interpolate(conditional(z > 1. - eps, 1, 0)) + if submesh_region == "left": + submesh_function = Function(DG0).interpolate(conditional(x < .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(x > .5, 1, 0)) + elif submesh_region == "right": + submesh_function = Function(DG0).interpolate(conditional(x > .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(x < .5, 1, 0)) + elif submesh_region == "front": + submesh_function = Function(DG0).interpolate(conditional(y < .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(y > .5, 1, 0)) + elif submesh_region == "back": + submesh_function = Function(DG0).interpolate(conditional(y > .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(y < .5, 1, 0)) + elif submesh_region == "bottom": + submesh_function = Function(DG0).interpolate(conditional(z < .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(z > .5, 1, 0)) + elif submesh_region == "top": + submesh_function = Function(DG0).interpolate(conditional(z > .5, 1, 0)) + submesh_function_compl = Function(DG0).interpolate(conditional(z < .5, 1, 0)) + else: + raise NotImplementedError(f"Unknown submesh_region: {submesh_region}") + return RelabeledMesh(mesh, [f101, f102, f103, f104, f105, f106, submesh_function, submesh_function_compl], + [101, 102, 103, 104, 105, 106, label_submesh, label_submesh_compl]) + + +def _mixed_poisson_solve_3d(hexahedral, degree, submesh_region): + dim = 3 + label_submesh = 999 + label_submesh_compl = 888 + mesh = _mixed_poisson_create_mesh_3d(hexahedral, submesh_region, label_submesh, label_submesh_compl) + x, y, z = SpatialCoordinate(mesh) + subm = Submesh(mesh, dim, label_submesh) + subx, suby, subz = SpatialCoordinate(subm) + if submesh_region == "left": + boun_ext = (101, ) + boun_dirichlet = (103, 104, 105, 106) + elif submesh_region == "right": + boun_ext = (102, ) + boun_dirichlet = (103, 104, 105, 106) + elif submesh_region == "front": + boun_ext = (103, ) + boun_dirichlet = (101, 102, 105, 106) + elif submesh_region == "back": + boun_ext = (104, ) + boun_dirichlet = (101, 102, 105, 106) + elif submesh_region == "bottom": + boun_ext = (105, ) + boun_dirichlet = (101, 102, 103, 104) + elif submesh_region == "top": + boun_ext = (106, ) + boun_dirichlet = (101, 102, 103, 104) + else: + raise NotImplementedError(f"Unknown submesh_region: {submesh_region}") + boun_int = (107, ) # labeled automatically. + NCF = FunctionSpace(subm, "NCF" if hexahedral else "N2F", degree) + DG = FunctionSpace(mesh, "DG", degree - 1) + W = NCF * DG + tau, v = TestFunctions(W) + nsub = FacetNormal(subm) + u_exact = Function(DG).interpolate(cos(2 * pi * x) * cos(2 * pi * y) * cos(2 * pi * z)) + sigma_exact = Function(NCF).project(as_vector([- 2 * pi * sin(2 * pi * subx) * cos(2 * pi * suby) * cos(2 * pi * subz), + - 2 * pi * cos(2 * pi * subx) * sin(2 * pi * suby) * cos(2 * pi * subz), + - 2 * pi * cos(2 * pi * subx) * cos(2 * pi * suby) * sin(2 * pi * subz)]), + solver_parameters={"ksp_type": "cg", "ksp_rtol": 1.e-16}) + f = Function(DG).interpolate(- 12 * pi * pi * cos(2 * pi * x) * cos(2 * pi * y) * cos(2 * pi * z)) + dx0 = Measure("dx", domain=mesh) + dx1 = Measure("dx", domain=subm) + ds0 = Measure("ds", domain=mesh) + ds1 = Measure("ds", domain=subm) + bc = DirichletBC(W.sub(0), sigma_exact, boun_dirichlet) + # Do the base case. + w = Function(W) + sigma, u = split(w) + a = (inner(sigma, tau) + inner(u, div(tau)) + inner(div(sigma), v)) * dx1 + inner(u - u_exact, v) * dx0(label_submesh_compl) + L = inner(f, v) * dx1 + inner((u('+') + u('-')) / 2., dot(tau('|'), nsub('|'))) * ds1(boun_int) + inner(u_exact('|'), dot(tau('|'), nsub('|'))) * ds0(boun_ext) + solve(a - L == 0, w, bcs=[bc]) + sigma_error = sqrt(assemble(inner(sigma - sigma_exact, sigma - sigma_exact) * dx1)) + u_error = sqrt(assemble(inner(u - u_exact, u - u_exact) * dx0(label_submesh))) + return sigma_error, u_error + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('hexahedral', [False]) +@pytest.mark.parametrize('degree', [4]) +@pytest.mark.parametrize('submesh_region', ["left", "right", "front", "back", "bottom", "top"]) +def test_submesh_solve_mixed_poisson_check_sanity_3d(hexahedral, degree, submesh_region): + sigma_error, u_error = _mixed_poisson_solve_3d(hexahedral, degree, submesh_region) + assert sigma_error < 0.07 + assert u_error < 0.003 + + +@pytest.mark.parallel(nprocs=4) +@pytest.mark.parametrize('simplex', [True, False]) +@pytest.mark.parametrize('nref', [1, 3]) +@pytest.mark.parametrize('degree', [2, 4]) +def test_submesh_solve_cell_cell_equation_bc(nref, degree, simplex): + dim = 2 + mesh = RectangleMesh(3 ** nref, 2 ** nref, 3., 2., quadrilateral=not simplex) + x, y = SpatialCoordinate(mesh) + label_outer = 101 + label_inner = 100 + label_interface = 5 # automatically labeled by Submesh + DG0 = FunctionSpace(mesh, "DG", 0) + f_outer = Function(DG0).interpolate(conditional(Or(Or(x < 1., x > 2.), y > 1.), 1, 0)) + f_inner = Function(DG0).interpolate(conditional(And(And(x > 1., x < 2.), y < 1.), 1, 0)) + mesh = RelabeledMesh(mesh, [f_outer, f_inner], [label_outer, label_inner]) + x, y = SpatialCoordinate(mesh) + mesh_outer = Submesh(mesh, dim, label_outer) + x_outer, y_outer = SpatialCoordinate(mesh_outer) + mesh_inner = Submesh(mesh, dim, label_inner) + x_inner, y_inner = SpatialCoordinate(mesh_inner) + V_outer = FunctionSpace(mesh_outer, "CG", degree) + V_inner = FunctionSpace(mesh_inner, "CG", degree) + V = V_outer * V_inner + u = TrialFunction(V) + v = TestFunction(V) + sol = Function(V) + u_outer, u_inner = split(u) + v_outer, v_inner = split(v) + dx_outer = Measure("dx", domain=mesh_outer) + dx_inner = Measure("dx", domain=mesh_inner) + ds_outer = Measure("ds", domain=mesh_outer) + a = inner(grad(u_outer), grad(v_outer)) * dx_outer + \ + inner(u_inner, v_inner) * dx_inner + L = inner(x * y, v_inner) * dx_inner + dbc = DirichletBC(V.sub(0), x_outer * y_outer, (1, 2, 3, 4)) + ebc = EquationBC(inner(u_outer - u_inner, v_outer) * ds_outer(label_interface) == inner(Constant(0.), v_outer) * ds_outer(label_interface), sol, label_interface, V=V.sub(0)) + solve(a == L, sol, bcs=[dbc, ebc]) + assert sqrt(assemble(inner(sol[0] - x * y, sol[0] - x * y) * dx_outer)) < 1.e-12 + assert sqrt(assemble(inner(sol[1] - x * y, sol[1] - x * y) * dx_inner)) < 1.e-12 From 096de4220ebc03da6f2ab4315baa2b5075b3d2cb Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Thu, 30 May 2024 17:01:03 +0100 Subject: [PATCH 3/9] equation bc: remove bbc check --- firedrake/bcs.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 272b74ddbf..ad229e95f0 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -162,8 +162,6 @@ def hermite_stride(bcnodes): # take intersection of facet nodes, and add it to bcnodes # i, j, k can also be strings. bcnodes1 = [] - if len(s) > 1 and not isinstance(self._function_space.finat_element, (finat.Lagrange, finat.GaussLobattoLegendre)): - raise TypeError("Currently, edge conditions have only been tested with CG Lagrange elements") for ss in s: # intersection of facets # Edge conditions have only been tested with Lagrange elements. From 5f36dc1b0f0bdf4f341de123ada69bf248818897 Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Thu, 12 Sep 2024 02:41:49 +0100 Subject: [PATCH 4/9] ufl update: MixedMesh([list]) --- firedrake/dmhooks.py | 2 +- firedrake/functionspace.py | 2 +- firedrake/functionspaceimpl.py | 2 +- firedrake/mesh.py | 14 +++++++------- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/firedrake/dmhooks.py b/firedrake/dmhooks.py index c3b61bcf08..fa97a8565a 100644 --- a/firedrake/dmhooks.py +++ b/firedrake/dmhooks.py @@ -58,7 +58,7 @@ def get_function_space(dm): if len(meshref_tuple) == 1: mesh = meshref_tuple[0]() else: - mesh = MixedMeshGeometry(*(meshref() for meshref in meshref_tuple)) + mesh = MixedMeshGeometry([meshref() for meshref in meshref_tuple]) if mesh is None: raise RuntimeError("Somehow your mesh was collected, this should never happen") V = firedrake.FunctionSpace(mesh, element, name=name) diff --git a/firedrake/functionspace.py b/firedrake/functionspace.py index 8d6ae617b2..844be0dfe6 100644 --- a/firedrake/functionspace.py +++ b/firedrake/functionspace.py @@ -297,7 +297,7 @@ def rec(eles): else: raise ValueError("Can't make mixed space with %s" % type(space)) - mixed_mesh_geometry = MixedMeshGeometry(*meshes) + mixed_mesh_geometry = MixedMeshGeometry(meshes) new = impl.MixedFunctionSpace(spaces, mixed_mesh_geometry.topology, name=name) return cls.create(new, mixed_mesh_geometry) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 6f16399c62..f20526ab50 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -375,7 +375,7 @@ def make_function_space(cls, mesh, element, name=None): # Create a new abstract (Mixed/Real)FunctionSpace, these are neither primal nor dual. if type(element) is finat.ufl.MixedElement: if isinstance(mesh, MeshGeometry): - mesh = MixedMeshGeometry(*[mesh for _ in element.sub_elements]) + mesh = MixedMeshGeometry([mesh for _ in element.sub_elements]) topology = mesh.topology else: if not isinstance(mesh, MixedMeshGeometry): diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 62f3803aab..d500687442 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4685,7 +4685,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): class MixedMeshGeometry(ufl.MixedMesh): """A representation of mixed mesh geometry.""" - def __init__(self, *meshes, set_hierarchy=True): + def __init__(self, meshes, set_hierarchy=True): """Initialise. Parameters @@ -4699,7 +4699,7 @@ def __init__(self, *meshes, set_hierarchy=True): for m in meshes: if not isinstance(m, MeshGeometry): raise ValueError(f"Got {type(m)}") - super().__init__(*meshes) + super().__init__(meshes) self.comm = meshes[0].comm self._comm = internal_comm(self.comm, self) # Only set hierarchy at top level. @@ -4708,7 +4708,7 @@ def __init__(self, *meshes, set_hierarchy=True): @utils.cached_property def topology(self): - return MixedMeshTopology(*[m.topology for m in self._meshes]) + return MixedMeshTopology([m.topology for m in self._meshes]) @property def topological(self): @@ -4781,7 +4781,7 @@ def set_hierarchy(self): if ilevel == level: result.append(self) else: - result.append(MixedMeshGeometry(*[hierarchy[ilevel] for hierarchy in hierarchy_list], set_hierarchy=False)) + result.append(MixedMeshGeometry([hierarchy[ilevel] for hierarchy in hierarchy_list], set_hierarchy=False)) result = tuple(result) for i, m in enumerate(result): set_level(m, result, i) @@ -4790,7 +4790,7 @@ def set_hierarchy(self): class MixedMeshTopology(object): """A representation of mixed mesh topology.""" - def __init__(self, *meshes): + def __init__(self, meshes): """Initialise. Parameters @@ -4828,8 +4828,8 @@ def ufl_cell(self): def ufl_mesh(self): cell = self.ufl_cell() - return ufl.MixedMesh(*[ufl.Mesh(finat.ufl.VectorElement("Lagrange", cell, 1, dim=cell.topological_dimension())) - for _ in self._meshes]) + return ufl.MixedMesh([ufl.Mesh(finat.ufl.VectorElement("Lagrange", cell, 1, dim=cell.topological_dimension())) + for _ in self._meshes]) def __eq__(self, other): if type(other) != type(self): From e2fd7f64f69ad17995615a503907928ac24e00cf Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Fri, 13 Sep 2024 02:59:46 +0100 Subject: [PATCH 5/9] rebase fix --- firedrake/assemble.py | 54 +++++++++++++++++++++++++++++--- firedrake/mg/kernels.py | 4 +-- firedrake/pointeval_utils.py | 2 +- firedrake/pointquery_utils.py | 2 +- firedrake/slate/slac/compiler.py | 5 ++- 5 files changed, 57 insertions(+), 10 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 3bdcbc97d6..27f86cbd00 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1600,6 +1600,8 @@ def __init__(self, form, local_knl, subdomain_id, all_integer_subdomain_ids, dia self._constants = _FormHandler.iter_constants(form, local_knl.kinfo) self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo) self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo) + self._active_exterior_facet_orientations = _FormHandler.iter_active_exterior_facet_orientations(form, local_knl.kinfo) + self._active_interior_facet_orientations = _FormHandler.iter_active_interior_facet_orientations(form, local_knl.kinfo) self._map_arg_cache = {} # Cache for holding :class:`op2.MapKernelArg` instances. @@ -1620,6 +1622,8 @@ def build(self): assert_empty(self._constants) assert_empty(self._active_exterior_facets) assert_empty(self._active_interior_facets) + assert_empty(self._active_exterior_facet_orientations) + assert_empty(self._active_interior_facet_orientations) iteration_regions = {"exterior_facet_top": op2.ON_TOP, "exterior_facet_bottom": op2.ON_BOTTOM, @@ -1817,12 +1821,24 @@ def _as_global_kernel_arg_interior_facet(_, self): @_as_global_kernel_arg.register(kernel_args.ExteriorFacetOrientationKernelArg) def _as_global_kernel_arg_exterior_facet_orientation(_, self): - return op2.DatKernelArg((1,)) + mesh = next(self._active_exterior_facet_orientations) + if mesh is self._mesh: + return op2.DatKernelArg((1,)) + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "exterior_facet" + return op2.DatKernelArg((1,), m._global_kernel_arg) @_as_global_kernel_arg.register(kernel_args.InteriorFacetOrientationKernelArg) def _as_global_kernel_arg_interior_facet_orientation(_, self): - return op2.DatKernelArg((2,)) + mesh = next(self._active_interior_facet_orientations) + if mesh is self._mesh: + return op2.DatKernelArg((2,)) + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "interior_facet" + return op2.DatKernelArg((2,), m._global_kernel_arg) @_as_global_kernel_arg.register(CellFacetKernelArg) @@ -1874,6 +1890,8 @@ def __init__(self, form, bcs, local_knl, subdomain_id, self._constants = _FormHandler.iter_constants(form, local_knl.kinfo) self._active_exterior_facets = _FormHandler.iter_active_exterior_facets(form, local_knl.kinfo) self._active_interior_facets = _FormHandler.iter_active_interior_facets(form, local_knl.kinfo) + self._active_exterior_facet_orientations = _FormHandler.iter_active_exterior_facet_orientations(form, local_knl.kinfo) + self._active_interior_facet_orientations = _FormHandler.iter_active_interior_facet_orientations(form, local_knl.kinfo) def build(self, tensor: op2.Global | op2.Dat | op2.Mat) -> op2.Parloop: """Construct the parloop. @@ -2139,12 +2157,24 @@ def _as_parloop_arg_interior_facet(_, self): @_as_parloop_arg.register(kernel_args.ExteriorFacetOrientationKernelArg) def _as_parloop_arg_exterior_facet_orientation(_, self): - return op2.DatParloopArg(self._mesh.exterior_facets.local_facet_orientation_dat) + mesh = next(self._active_exterior_facet_orientations) + if mesh is self._mesh: + m = None + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "exterior_facet" + return op2.DatParloopArg(mesh.exterior_facets.local_facet_orientation_dat, m) @_as_parloop_arg.register(kernel_args.InteriorFacetOrientationKernelArg) def _as_parloop_arg_interior_facet_orientation(_, self): - return op2.DatParloopArg(self._mesh.interior_facets.local_facet_orientation_dat) + mesh = next(self._active_interior_facet_orientations) + if mesh is self._mesh: + m = None + else: + m, integral_type = mesh.topology.trans_mesh_entity_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + assert integral_type == "interior_facet" + return op2.DatParloopArg(mesh.interior_facets.local_facet_orientation_dat, m) @_as_parloop_arg.register(CellFacetKernelArg) @@ -2222,6 +2252,22 @@ def iter_active_interior_facets(form, kinfo): mesh = all_meshes[i] yield mesh + @staticmethod + def iter_active_exterior_facet_orientations(form, kinfo): + """Yield the form exterior facet orientations referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.exterior_facet_orientations: + mesh = all_meshes[i] + yield mesh + + @staticmethod + def iter_active_interior_facet_orientations(form, kinfo): + """Yield the form interior facet orientations referenced in ``kinfo``.""" + all_meshes = extract_domains(form) + for i in kinfo.active_domain_numbers.interior_facet_orientations: + mesh = all_meshes[i] + yield mesh + @staticmethod def index_function_spaces(form, indices): """Return the function spaces of the form's arguments, indexed diff --git a/firedrake/mg/kernels.py b/firedrake/mg/kernels.py index 09d7448ef8..a1d2a80e78 100644 --- a/firedrake/mg/kernels.py +++ b/firedrake/mg/kernels.py @@ -144,7 +144,7 @@ def compile_element(expression, dual_space=None, parameters=None, config = dict(interface=builder, ufl_cell=cell, - integral_type="cell", + domain_integral_type_map={domain: "cell"}, point_indices=(), point_expr=point, argument_multiindices=argument_multiindices, @@ -539,7 +539,6 @@ def dg_injection_kernel(Vf, Vc, ncell): integration_dim, entity_ids = lower_integral_type(Vfe.cell, "cell") macro_cfg = dict(interface=macro_builder, ufl_cell=Vf.ufl_cell(), - integral_type="cell", integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, @@ -579,7 +578,6 @@ def dg_injection_kernel(Vf, Vc, ncell): coarse_cfg = dict(interface=coarse_builder, ufl_cell=Vc.ufl_cell(), - integral_type="cell", integration_dim=integration_dim, entity_ids=entity_ids, index_cache=index_cache, diff --git a/firedrake/pointeval_utils.py b/firedrake/pointeval_utils.py index fc056738a9..17f9832b8e 100644 --- a/firedrake/pointeval_utils.py +++ b/firedrake/pointeval_utils.py @@ -72,7 +72,7 @@ def compile_element(expression, coordinates, parameters=None): config = dict(interface=builder, ufl_cell=extract_unique_domain(coordinates).ufl_cell(), - integral_type="cell", + domain_integral_type_map={domain: "cell"}, point_indices=(), point_expr=point, scalar_type=utils.ScalarType) diff --git a/firedrake/pointquery_utils.py b/firedrake/pointquery_utils.py index 58ceeb10c7..4fcfec1084 100644 --- a/firedrake/pointquery_utils.py +++ b/firedrake/pointquery_utils.py @@ -162,7 +162,7 @@ def to_reference_coords_newton_step(ufl_coordinate_element, parameters, x0_dtype context = tsfc.fem.GemPointContext( interface=builder, ufl_cell=cell, - integral_type="cell", + domain_integral_type_map={domain: "cell"}, point_indices=(), point_expr=point, scalar_type=parameters["scalar_type"] diff --git a/firedrake/slate/slac/compiler.py b/firedrake/slate/slac/compiler.py index a49d3ea11a..bbb2acb666 100644 --- a/firedrake/slate/slac/compiler.py +++ b/firedrake/slate/slac/compiler.py @@ -197,7 +197,10 @@ def generate_loopy_kernel(slate_expr, compiler_parameters=None): cell_orientations=(0, ) if builder.bag.needs_cell_orientations else (), cell_sizes=(0, ) if builder.bag.needs_cell_sizes else (), exterior_facets=(), - interior_facets=()), + interior_facets=(), + exterior_facet_orientations=(), + interior_facet_orientations=(), + ), coefficient_numbers=coefficient_numbers, constant_numbers=constant_numbers, needs_cell_facets=builder.bag.needs_cell_facets, From d06c980a0399bbab5418859a9aa2e696e833c31f Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Fri, 13 Sep 2024 09:43:57 +0100 Subject: [PATCH 6/9] k --- firedrake/slate/slac/compiler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/firedrake/slate/slac/compiler.py b/firedrake/slate/slac/compiler.py index bbb2acb666..787d7364c3 100644 --- a/firedrake/slate/slac/compiler.py +++ b/firedrake/slate/slac/compiler.py @@ -199,8 +199,7 @@ def generate_loopy_kernel(slate_expr, compiler_parameters=None): exterior_facets=(), interior_facets=(), exterior_facet_orientations=(), - interior_facet_orientations=(), - ), + interior_facet_orientations=(),), coefficient_numbers=coefficient_numbers, constant_numbers=constant_numbers, needs_cell_facets=builder.bag.needs_cell_facets, From ed0e3b11481bf36711bf2ad0c20dde22e7671fa9 Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Sat, 16 Nov 2024 00:40:06 +0000 Subject: [PATCH 7/9] todo --- firedrake/functionspaceimpl.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index f20526ab50..1d3bd2bf24 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -195,9 +195,9 @@ def _components(self): @PETSc.Log.EventDecorator() def sub(self, i): - bound = len(self._components) - if i < 0 or i >= bound: - raise IndexError("Invalid component %d, not in [0, %d)" % (i, bound)) + #bound = len(self._components) + #if i < 0 or i >= bound: + # raise IndexError("Invalid component %d, not in [0, %d)" % (i, bound)) if len(self) == 1: return self._components[i] else: From cad0cd2564ed2f19a50e7bdc010cbc19133ed90f Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Sat, 16 Nov 2024 14:03:58 +0000 Subject: [PATCH 8/9] k --- firedrake/slate/slate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/slate/slate.py b/firedrake/slate/slate.py index 9abbc5980f..7706b1eb1d 100644 --- a/firedrake/slate/slate.py +++ b/firedrake/slate/slate.py @@ -257,7 +257,7 @@ def ufl_domain(self): def _expand_mixed_meshes(func): @functools.wraps(func) def wrapper(self, *args, **kwargs): - return join_domains(func(self, *args, **kwargs)) + return sort_domains(join_domains(func(self, *args, **kwargs))) return wrapper @abstractmethod From 8417b41657d69bd030cfcaacd61dac33a79f9fc4 Mon Sep 17 00:00:00 2001 From: ksagiyam Date: Sat, 16 Nov 2024 20:15:21 +0000 Subject: [PATCH 9/9] k --- firedrake/functionspaceimpl.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 1d3bd2bf24..ac98149549 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -195,12 +195,15 @@ def _components(self): @PETSc.Log.EventDecorator() def sub(self, i): - #bound = len(self._components) - #if i < 0 or i >= bound: - # raise IndexError("Invalid component %d, not in [0, %d)" % (i, bound)) if len(self) == 1: + bound = len(self._components) + if i < 0 or i >= bound: + raise IndexError("Invalid component %d, not in [0, %d)" % (i, bound)) return self._components[i] else: + bound = len(self) + if i < 0 or i >= bound: + raise IndexError("Invalid index %d, not in [0, %d)" % (i, bound)) return self.subfunctions[i] @utils.cached_property