Expunge Expr.ufl_domain (#3259)

---------

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

docs/source/interpolation.rst

@@ -251,7 +251,7 @@ proceeds as follows:
 .. code-block:: python3
 
    # First, grab the mesh.
-   m = V.ufl_domain()
+   m = V.mesh()
 
    # Now make the VectorFunctionSpace corresponding to V.
    W = VectorFunctionSpace(m, V.ufl_element())

firedrake/assemble.py

@@ -1338,7 +1338,7 @@ def _as_global_kernel_arg_cell_facet(_, self):
 @_as_global_kernel_arg.register(kernel_args.CellOrientationsKernelArg)
 def _as_global_kernel_arg_cell_orientations(_, self):
     # this mirrors firedrake.mesh.MeshGeometry.init_cell_orientations
-    ufl_element = finat.ufl.FiniteElement("DG", cell=self._form.ufl_domain().ufl_cell(), degree=0)
+    ufl_element = finat.ufl.FiniteElement("DG", cell=self._mesh.ufl_cell(), degree=0)
     finat_element = create_element(ufl_element)
     return self._make_dat_global_kernel_arg(finat_element)
 

firedrake/function.py

@@ -616,16 +616,16 @@ def at(self, arg, *args, **kwargs):
         dont_raise = kwargs.get('dont_raise', False)
 
         tolerance = kwargs.get('tolerance', None)
+        mesh = self.function_space().mesh()
         if tolerance is None:
-            tolerance = self.ufl_domain().tolerance
+            tolerance = mesh.tolerance
         else:
-            self.ufl_domain().tolerance = tolerance
+            mesh.tolerance = tolerance
 
         # Handle f.at(0.3)
         if not arg.shape:
             arg = arg.reshape(-1)
 
-        mesh = self.function_space().mesh()
         if mesh.variable_layers:
             raise NotImplementedError("Point evaluation not implemented for variable layers")
 

firedrake/functionspaceimpl.py

@@ -168,7 +168,7 @@ def ufl_function_space(self):
 
     def ufl_cell(self):
         r"""The :class:`~ufl.classes.Cell` this FunctionSpace is defined on."""
-        return self.ufl_domain().ufl_cell()
+        return self.mesh().ufl_cell()
 
     @PETSc.Log.EventDecorator()
     def split(self):

firedrake/interpolation.py

@@ -9,7 +9,7 @@
 import finat.ufl
 from ufl.algorithms import extract_arguments, extract_coefficients
 from ufl.algorithms.signature import compute_expression_signature
-from ufl.domain import extract_unique_domain
+from ufl.domain import as_domain, extract_unique_domain
 
 from pyop2 import op2
 from pyop2.caching import disk_cached
@@ -191,7 +191,7 @@ class Interpolator(abc.ABC):
     """
 
     def __new__(cls, expr, V, **kwargs):
-        target_mesh = V.ufl_domain()
+        target_mesh = as_domain(V)
         source_mesh = extract_unique_domain(expr) or target_mesh
         if target_mesh is not source_mesh:
             if isinstance(target_mesh.topology, firedrake.mesh.VertexOnlyMeshTopology):
@@ -323,7 +323,7 @@ def __init__(
         # setup
         V_dest = V
         src_mesh = extract_unique_domain(expr)
-        dest_mesh = V_dest.ufl_domain()
+        dest_mesh = as_domain(V_dest)
         src_mesh_gdim = src_mesh.geometric_dimension()
         dest_mesh_gdim = dest_mesh.geometric_dimension()
         if src_mesh_gdim != dest_mesh_gdim:
@@ -702,7 +702,7 @@ def make_interpolator(expr, V, subset, access, bcs=None):
     assert isinstance(expr, ufl.classes.Expr)
 
     arguments = extract_arguments(expr)
-    target_mesh = V.ufl_domain()
+    target_mesh = as_domain(V)
     if len(arguments) == 0:
         source_mesh = extract_unique_domain(expr) or target_mesh
         vom_onto_other_vom = (
@@ -852,7 +852,7 @@ def _interpolator(V, tensor, expr, subset, arguments, access, bcs=None):
                            % (expr.ufl_shape, V.ufl_element().value_shape))
 
     # NOTE: The par_loop is always over the target mesh cells.
-    target_mesh = V.ufl_domain()
+    target_mesh = as_domain(V)
     source_mesh = extract_unique_domain(expr) or target_mesh
 
     if target_mesh is not source_mesh:

firedrake/mg/interface.py

@@ -69,7 +69,7 @@ def prolong(coarse, fine):
             Vf = firedrake.FunctionSpace(meshes[next_level], element)
             next = firedrake.Function(Vf)
 
-        coarse_coords = Vc.ufl_domain().coordinates
+        coarse_coords = Vc.mesh().coordinates
         fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc)
         fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space())
         kernel = kernels.prolong_kernel(coarse)
@@ -136,7 +136,7 @@ def restrict(fine_dual, coarse_dual):
         # x = \sum_i c_i \phi_i(x_hat)
         node_locations = utils.physical_node_locations(Vf)
 
-        coarse_coords = Vc.ufl_domain().coordinates
+        coarse_coords = Vc.mesh().coordinates
         fine_to_coarse = utils.fine_node_to_coarse_node_map(Vf, Vc)
         fine_to_coarse_coords = utils.fine_node_to_coarse_node_map(Vf, coarse_coords.function_space())
         # Have to do this, because the node set core size is not right for
@@ -210,7 +210,7 @@ def inject(fine, coarse):
         if not dg:
             node_locations = utils.physical_node_locations(Vc)
 
-            fine_coords = Vf.ufl_domain().coordinates
+            fine_coords = Vf.mesh().coordinates
             coarse_node_to_fine_nodes = utils.coarse_node_to_fine_node_map(Vc, Vf)
             coarse_node_to_fine_coords = utils.coarse_node_to_fine_node_map(Vc, fine_coords.function_space())
 

firedrake/mg/kernels.py

@@ -292,10 +292,10 @@ def prolong_kernel(expression):
 
 
 def restrict_kernel(Vf, Vc):
-    hierarchy, level = utils.get_level(Vc.ufl_domain())
+    hierarchy, level = utils.get_level(Vc.mesh())
     levelf = level + Fraction(1, hierarchy.refinements_per_level)
     cache = hierarchy._shared_data_cache["transfer_kernels"]
-    coordinates = Vc.ufl_domain().coordinates
+    coordinates = Vc.mesh().coordinates
     if Vf.extruded:
         assert Vc.extruded
     key = (("restrict",)
@@ -376,9 +376,9 @@ def restrict_kernel(Vf, Vc):
 
 
 def inject_kernel(Vf, Vc):
-    hierarchy, level = utils.get_level(Vc.ufl_domain())
+    hierarchy, level = utils.get_level(Vc.mesh())
     cache = hierarchy._shared_data_cache["transfer_kernels"]
-    coordinates = Vf.ufl_domain().coordinates
+    coordinates = Vf.mesh().coordinates
     if Vf.extruded:
         assert Vc.extruded
         level_ratio = (Vf.mesh().layers - 1) // (Vc.mesh().layers - 1)
@@ -397,7 +397,7 @@ def inject_kernel(Vf, Vc):
         if Vc.finat_element.entity_dofs() == Vc.finat_element.entity_closure_dofs():
             return cache.setdefault(key, (dg_injection_kernel(Vf, Vc, ncandidate), True))
 
-        coordinates = Vf.ufl_domain().coordinates
+        coordinates = Vf.mesh().coordinates
         evaluate_code = compile_element(ufl.Coefficient(Vf))
         to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
 
@@ -456,7 +456,7 @@ def inject_kernel(Vf, Vc):
             "inside_cell": inside_check(Vc.finat_element.cell, eps=1e-8, X="Xref"),
             "spacedim": Vc.finat_element.cell.get_spatial_dimension(),
             "celldist_l1_c_expr": celldist_l1_c_expr(Vc.finat_element.cell, X="Xref"),
-            "tdim": Vc.ufl_domain().topological_dimension(),
+            "tdim": Vc.mesh().topological_dimension(),
             "ncandidate": ncandidate,
             "Rdim": numpy.prod(Vf_element.value_shape),
             "Xf_cell_inc": coords_element.space_dimension(),

firedrake/mg/ufl_utils.py

@@ -1,7 +1,7 @@
 import ufl
 from ufl.corealg.map_dag import map_expr_dag
 from ufl.corealg.multifunction import MultiFunction
-from ufl.domain import extract_unique_domain
+from ufl.domain import as_domain, extract_unique_domain
 from ufl.duals import is_dual
 
 from functools import singledispatch, partial
@@ -95,7 +95,7 @@ def coarsen_form(form, self, coefficient_mapping=None):
     integrals = []
     for it in form.integrals():
         integrand = map_expr_dag(mapper, it.integrand())
-        mesh = it.ufl_domain()
+        mesh = as_domain(it)
         hierarchy, level = utils.get_level(mesh)
         new_mesh = hierarchy[level-1]
         if isinstance(integrand, ufl.classes.Zero):

firedrake/mg/utils.py

@@ -14,8 +14,8 @@ def fine_node_to_coarse_node_map(Vf, Vc):
         return op2.MixedMap(fine_node_to_coarse_node_map(f, c) for f, c in zip(Vf, Vc))
     mesh = Vf.mesh()
     assert hasattr(mesh, "_shared_data_cache")
-    hierarchyf, levelf = get_level(Vf.ufl_domain())
-    hierarchyc, levelc = get_level(Vc.ufl_domain())
+    hierarchyf, levelf = get_level(Vf.mesh())
+    hierarchyc, levelc = get_level(Vc.mesh())
 
     if hierarchyc != hierarchyf:
         raise ValueError("Can't map across hierarchies")
@@ -52,8 +52,8 @@ def coarse_node_to_fine_node_map(Vc, Vf):
         return op2.MixedMap(coarse_node_to_fine_node_map(f, c) for f, c in zip(Vf, Vc))
     mesh = Vc.mesh()
     assert hasattr(mesh, "_shared_data_cache")
-    hierarchyf, levelf = get_level(Vf.ufl_domain())
-    hierarchyc, levelc = get_level(Vc.ufl_domain())
+    hierarchyf, levelf = get_level(Vf.mesh())
+    hierarchyc, levelc = get_level(Vc.mesh())
 
     if hierarchyc != hierarchyf:
         raise ValueError("Can't map across hierarchies")
@@ -90,8 +90,8 @@ def coarse_cell_to_fine_node_map(Vc, Vf):
         return op2.MixedMap(coarse_cell_to_fine_node_map(f, c) for f, c in zip(Vf, Vc))
     mesh = Vc.mesh()
     assert hasattr(mesh, "_shared_data_cache")
-    hierarchyf, levelf = get_level(Vf.ufl_domain())
-    hierarchyc, levelc = get_level(Vc.ufl_domain())
+    hierarchyf, levelf = get_level(Vf.mesh())
+    hierarchyc, levelc = get_level(Vc.mesh())
 
     if hierarchyc != hierarchyf:
         raise ValueError("Can't map across hierarchies")

firedrake/output.py

@@ -73,7 +73,7 @@ def is_cg(V):
 
     :arg V: A FunctionSpace.
     """
-    nvertex = V.ufl_domain().ufl_cell().num_vertices()
+    nvertex = V.mesh().ufl_cell().num_vertices()
     entity_dofs = V.finat_element.entity_dofs()
     # If there are as many dofs on vertices as there are vertices,
     # assume a continuous space.
@@ -96,7 +96,7 @@ def is_linear(V):
 
     :arg V: A FunctionSpace.
     """
-    nvertex = V.ufl_domain().ufl_cell().num_vertices()
+    nvertex = V.mesh().ufl_cell().num_vertices()
     return V.finat_element.space_dimension() == nvertex
 
 
@@ -135,7 +135,7 @@ def get_topology(coordinates):
     V = coordinates.function_space()
 
     nonLinear = not is_linear(V)
-    mesh = V.ufl_domain().topology
+    mesh = V.mesh().topology
     cell = mesh.ufl_cell()
     values = V.cell_node_map().values
     value_shape = values.shape

firedrake/plot.py

@@ -1000,7 +1000,7 @@ def _pgfplot_create_patches(f, coords, complex_component):
     fiat_cell = V.finat_element.cell
     degree = elem.degree()
     coordV = coords.function_space()
-    mesh = V.ufl_domain()
+    mesh = V.mesh()
     cdata = coords.dat.data_ro.real
     fdata = f.dat.data_ro.real if complex_component == 'real' else f.dat.data_ro.imag
     map_facet_dofs, patch_type = _pgfplot_make_perms(fiat_cell, degree)
@@ -1068,7 +1068,7 @@ def pgfplot(f, filename, degree=1, complex_component='real', print_latex_example
         raise NotImplementedError(f"complex_component must be {'real', 'imag'}: got {complex_component}")
     V = f.function_space()
     elem = V.ufl_element()
-    mesh = V.ufl_domain()
+    mesh = V.mesh()
     dim = mesh.geometric_dimension()
     if dim not in (2, 3):
         raise NotImplementedError(f"Not yet implemented for functions in spatial dimension {dim}")

firedrake/preconditioners/patch.py

@@ -192,11 +192,11 @@ def matrix_funptr(form, state):
         arg = mesh.coordinates.dat(op2.READ, get_map(mesh.coordinates))
         args.append(arg)
         if kinfo.oriented:
-            c = form.ufl_domain().cell_orientations()
+            c = mesh.cell_orientations()
             arg = c.dat(op2.READ, get_map(c))
             args.append(arg)
         if kinfo.needs_cell_sizes:
-            c = form.ufl_domain().cell_sizes
+            c = mesh.cell_sizes
             arg = c.dat(op2.READ, get_map(c))
             args.append(arg)
         for n, indices in kinfo.coefficient_numbers:
@@ -217,7 +217,7 @@ def matrix_funptr(form, state):
             args.append(all_constants[constant_index].dat(op2.READ))
 
         if kinfo.integral_type == "interior_facet":
-            arg = test.ufl_domain().interior_facets.local_facet_dat(op2.READ)
+            arg = mesh.interior_facets.local_facet_dat(op2.READ)
             args.append(arg)
         iterset = op2.Subset(iterset, [])
 
@@ -286,11 +286,11 @@ def residual_funptr(form, state):
         args.append(arg)
 
         if kinfo.oriented:
-            c = form.ufl_domain().cell_orientations()
+            c = mesh.cell_orientations()
             arg = c.dat(op2.READ, get_map(c))
             args.append(arg)
         if kinfo.needs_cell_sizes:
-            c = form.ufl_domain().cell_sizes
+            c = mesh.cell_sizes
             arg = c.dat(op2.READ, get_map(c))
             args.append(arg)
         for n, indices in kinfo.coefficient_numbers:
@@ -514,11 +514,12 @@ def load_c_function(code, name, comm):
 
 def make_c_arguments(form, kernel, state, get_map, require_state=False,
                      require_facet_number=False):
-    coeffs = [form.ufl_domain().coordinates]
+    mesh = form.ufl_domains()[kernel.kinfo.domain_number]
+    coeffs = [mesh.coordinates]
     if kernel.kinfo.oriented:
-        coeffs.append(form.ufl_domain().cell_orientations())
+        coeffs.append(mesh.cell_orientations())
     if kernel.kinfo.needs_cell_sizes:
-        coeffs.append(form.ufl_domain().cell_sizes)
+        coeffs.append(mesh.cell_sizes)
     for n, indices in kernel.kinfo.coefficient_numbers:
         c = form.coefficients()[n]
         if c is state:
@@ -550,7 +551,7 @@ def make_c_arguments(form, kernel, state, get_map, require_state=False,
         data_args.extend(all_constants[constant_index].dat._kernel_args_)
 
     if require_facet_number:
-        data_args.extend(form.ufl_domain().interior_facets.local_facet_dat._kernel_args_)
+        data_args.extend(mesh.interior_facets.local_facet_dat._kernel_args_)
     return data_args, map_args
 
 
@@ -778,7 +779,8 @@ def initialize(self, obj):
             J = ctx.Jp or ctx.J
             bcs = ctx._problem.bcs
 
-        mesh = J.ufl_domain()
+        V = J.arguments()[0].function_space()
+        mesh = V.mesh()
         self.plex = mesh.topology_dm
         # We need to attach the mesh and appctx to the plex, so that
         # PlaneSmoothers (and any other user-customised patch
@@ -810,8 +812,6 @@ def initialize(self, obj):
             Jstate = None
             is_snes = False
 
-        V, _ = map(operator.methodcaller("function_space"), J.arguments())
-
         if len(bcs) > 0:
             ghost_bc_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=True)
                                                              for bc in bcs]))
@@ -838,7 +838,7 @@ def initialize(self, obj):
                                                                        require_facet_number=True)
             code, Struct = make_jacobian_wrapper(facet_Jop_data_args, facet_Jop_map_args)
             facet_Jop_function = load_c_function(code, "ComputeJacobian", obj.comm)
-            point2facet = J.ufl_domain().interior_facets.point2facetnumber.ctypes.data
+            point2facet = mesh.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)

firedrake/ufl_expr.py

@@ -3,6 +3,7 @@
 from ufl.duals import is_dual
 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
@@ -230,7 +231,7 @@ def derivative(form, u, du=None, coefficient_derivatives=None):
         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 = form.ufl_domain()
+    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?")

tests/extrusion/test_real_tensorproduct.py

@@ -28,7 +28,7 @@ def variant(request):
 
 @pytest.fixture
 def expr(variant, V, fs_kind):
-    x, y, z = SpatialCoordinate(V.ufl_domain())
+    x, y, z = SpatialCoordinate(V.mesh())
     val = {"linear": z, "sin": sin(pi*z)}[variant]
     ret = {
         "scalar": val,

tests/multigrid/test_non_nested.py

@@ -24,7 +24,7 @@ def coarsen(expr, self, coefficient_mapping=None):
 
     @coarsen.register(functionspaceimpl.WithGeometryBase)
     def coarsen_fs(V, self, coefficient_mapping=None):
-        mesh = self(V.ufl_domain(), self)
+        mesh = self(V.mesh(), self)
         return FunctionSpace(mesh, "CG", 1)
 
     uh = Function(V)

tests/multigrid/test_transfer_manager.py

@@ -49,7 +49,7 @@ def test_transfer_manager_inside_coarsen(sub, mesh):
 
     bc, = cctx._problem.bcs
     V = bc.function_space()
-    mesh = V.ufl_domain()
+    mesh = V.mesh()
     x, y = SpatialCoordinate(mesh)
     expect = project(as_vector([-y, x]), V)
     assert numpy.allclose(bc.function_arg.dat.data_ro, expect.dat.data_ro)