Merge pull request #151 from firedrakeproject/physical-geometry

Make UFL -> GEM aware of physical geometry

tsfc/driver.py

@@ -14,7 +14,7 @@
 import ufl
 from ufl.algorithms import extract_arguments, extract_coefficients
 from ufl.algorithms.analysis import has_type
-from ufl.classes import Form, CellVolume
+from ufl.classes import Form, GeometricQuantity
 from ufl.log import GREEN
 from ufl.utils.sequences import max_degree
 
@@ -32,7 +32,6 @@
 from tsfc.logging import logger
 from tsfc.parameters import default_parameters
 
-from tsfc.kernel_interface import ProxyKernelInterface
 import tsfc.kernel_interface.firedrake as firedrake_interface
 
 
@@ -109,8 +108,7 @@ def compile_integral(integral_data, form_data, prefix, parameters,
                                   for arg in arguments)
     return_variables = builder.set_arguments(arguments, argument_multiindices)
 
-    coordinates = ufl_utils.coordinate_coefficient(mesh)
-    builder.set_coordinates(coordinates)
+    builder.set_coordinates(mesh)
 
     builder.set_coefficients(integral_data, form_data)
 
@@ -126,15 +124,13 @@ def compile_integral(integral_data, form_data, prefix, parameters,
 
     kernel_cfg = dict(interface=builder,
                       ufl_cell=cell,
+                      integral_type=integral_type,
                       precision=parameters["precision"],
                       integration_dim=integration_dim,
                       entity_ids=entity_ids,
                       argument_multiindices=argument_multiindices,
                       index_cache=index_cache)
 
-    kernel_cfg["facetarea"] = facetarea_generator(mesh, coordinates, kernel_cfg, integral_type)
-    kernel_cfg["cellvolume"] = cellvolume_generator(mesh, coordinates, kernel_cfg)
-
     mode_irs = collections.OrderedDict()
     for integral in integral_data.integrals:
         params = parameters.copy()
@@ -145,8 +141,7 @@ def compile_integral(integral_data, form_data, prefix, parameters,
         mode = pick_mode(params["mode"])
         mode_irs.setdefault(mode, collections.OrderedDict())
 
-        integrand = ufl_utils.replace_coordinates(integral.integrand(), coordinates)
-        integrand = ufl.replace(integrand, form_data.function_replace_map)
+        integrand = ufl.replace(integral.integrand(), form_data.function_replace_map)
         integrand = ufl_utils.split_coefficients(integrand, builder.coefficient_split)
 
         # Check if the integral has a quad degree attached, otherwise use
@@ -157,7 +152,7 @@ def compile_integral(integral_data, form_data, prefix, parameters,
             quadrature_degree = params["quadrature_degree"]
         except KeyError:
             quadrature_degree = params["estimated_polynomial_degree"]
-            functions = list(arguments) + [coordinates] + list(integral_data.integral_coefficients)
+            functions = list(arguments) + [builder.coordinate(mesh)] + list(integral_data.integral_coefficients)
             function_degrees = [f.ufl_function_space().ufl_element().degree() for f in functions]
             if all((asarray(quadrature_degree) > 10 * asarray(degree)).all()
                    for degree in function_degrees):
@@ -248,50 +243,6 @@ def name_multiindex(multiindex, name):
     return builder.construct_kernel(kernel_name, body)
 
 
-class CellVolumeKernelInterface(ProxyKernelInterface):
-    # Since CellVolume is evaluated as a cell integral, we must ensure
-    # that the right restriction is applied when it is used in an
-    # interior facet integral.  This proxy diverts coefficient
-    # translation to use a specified restriction.
-
-    def __init__(self, wrapee, restriction):
-        ProxyKernelInterface.__init__(self, wrapee)
-        self.restriction = restriction
-
-    def coefficient(self, ufl_coefficient, r):
-        assert r is None
-        return self._wrapee.coefficient(ufl_coefficient, self.restriction)
-
-
-def cellvolume_generator(domain, coordinate_coefficient, kernel_config):
-    def cellvolume(restriction):
-        from ufl import dx
-        integrand, degree = ufl_utils.one_times(dx(domain=domain))
-        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)
-        interface = CellVolumeKernelInterface(kernel_config["interface"], restriction)
-
-        config = {k: v for k, v in kernel_config.items()
-                  if k in ["ufl_cell", "precision", "index_cache"]}
-        config.update(interface=interface, quadrature_degree=degree)
-        expr, = fem.compile_ufl(integrand, point_sum=True, **config)
-        return expr
-    return cellvolume
-
-
-def facetarea_generator(domain, coordinate_coefficient, kernel_config, integral_type):
-    def facetarea():
-        from ufl import Measure
-        assert integral_type != 'cell'
-        integrand, degree = ufl_utils.one_times(Measure(integral_type, domain=domain))
-        integrand = ufl_utils.replace_coordinates(integrand, coordinate_coefficient)
-
-        config = kernel_config.copy()
-        config.update(quadrature_degree=degree)
-        expr, = fem.compile_ufl(integrand, point_sum=True, **config)
-        return expr
-    return facetarea
-
-
 def compile_expression_at_points(expression, points, coordinates, parameters=None):
     """Compiles a UFL expression to be evaluated at compile-time known
     reference points.  Useful for interpolating UFL expressions onto
@@ -318,19 +269,19 @@ def compile_expression_at_points(expression, points, coordinates, parameters=Non
     # Apply UFL preprocessing
     expression = ufl_utils.preprocess_expression(expression)
 
+    # Initialise kernel builder
+    builder = firedrake_interface.ExpressionKernelBuilder()
+
     # Replace coordinates (if any)
     domain = expression.ufl_domain()
     if domain:
         assert coordinates.ufl_domain() == domain
-        expression = ufl_utils.replace_coordinates(expression, coordinates)
+        builder.domain_coordinate[domain] = coordinates
 
     # Collect required coefficients
     coefficients = extract_coefficients(expression)
-    if coordinates not in coefficients and has_type(expression, CellVolume):
+    if has_type(expression, GeometricQuantity):
         coefficients = [coordinates] + coefficients
-
-    # Initialise kernel builder
-    builder = firedrake_interface.ExpressionKernelBuilder()
     builder.set_coefficients(coefficients)
 
     # Split mixed coefficients
@@ -342,7 +293,6 @@ def compile_expression_at_points(expression, points, coordinates, parameters=Non
                   ufl_cell=coordinates.ufl_domain().ufl_cell(),
                   precision=parameters["precision"],
                   point_set=point_set)
-    config["cellvolume"] = cellvolume_generator(coordinates.ufl_domain(), coordinates, config)
     ir, = fem.compile_ufl(expression, point_sum=False, **config)
 
     # Deal with non-scalar expressions

tsfc/fem.py

@@ -11,14 +11,16 @@
 import numpy
 from singledispatch import singledispatch
 
+import ufl
 from ufl.corealg.map_dag import map_expr_dag, map_expr_dags
 from ufl.corealg.multifunction import MultiFunction
 from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors,
                          CellFacetJacobian, CellOrientation,
-                         CellVolume, Coefficient, FacetArea,
-                         FacetCoordinate, GeometricQuantity,
-                         QuadratureWeight, ReferenceCellVolume,
-                         ReferenceFacetVolume, ReferenceNormal)
+                         CellOrigin, CellVolume, Coefficient,
+                         FacetArea, FacetCoordinate,
+                         GeometricQuantity, QuadratureWeight,
+                         ReferenceCellVolume, ReferenceFacetVolume,
+                         ReferenceNormal, SpatialCoordinate)
 
 from FIAT.reference_element import make_affine_mapping
 
@@ -28,26 +30,30 @@
 from gem.unconcatenate import unconcatenate
 from gem.utils import cached_property
 
+from finat.point_set import PointSingleton
 from finat.quadrature import make_quadrature
 
 from tsfc import ufl2gem
 from tsfc.finatinterface import as_fiat_cell
 from tsfc.kernel_interface import ProxyKernelInterface
-from tsfc.modified_terminals import analyse_modified_terminal
+from tsfc.modified_terminals import (analyse_modified_terminal,
+                                     construct_modified_terminal)
 from tsfc.parameters import NUMPY_TYPE, PARAMETERS
-from tsfc.ufl_utils import ModifiedTerminalMixin, PickRestriction, simplify_abs
+from tsfc.ufl_utils import (ModifiedTerminalMixin, PickRestriction,
+                            one_times, simplify_abs,
+                            preprocess_expression)
 
 
 class ContextBase(ProxyKernelInterface):
     """Common UFL -> GEM translation context."""
 
     keywords = ('ufl_cell',
                 'fiat_cell',
+                'integral_type',
                 'integration_dim',
                 'entity_ids',
                 'precision',
                 'argument_multiindices',
-                'cellvolume',
                 'facetarea',
                 'index_cache')
 
@@ -99,6 +105,11 @@ def entity_selector(self, callback, restriction):
     def index_cache(self):
         return {}
 
+    @cached_property
+    def translator(self):
+        # NOTE: reference cycle!
+        return Translator(self)
+
 
 class PointSetContext(ContextBase):
     """Context for compile-time known evaluation points."""
@@ -153,12 +164,17 @@ def basis_evaluation(self, finat_element, local_derivatives, entity_id):
 
 
 class Translator(MultiFunction, ModifiedTerminalMixin, ufl2gem.Mixin):
-    """Contains all the context necessary to translate UFL into GEM."""
+    """Multifunction for translating UFL -> GEM.  Incorporates ufl2gem.Mixin, and
+    dispatches on terminal type when reaching modified terminals."""
 
     def __init__(self, context):
+        # MultiFunction.__init__ does not call further __init__
+        # methods, but ufl2gem.Mixin must be initialised.
+        # (ModifiedTerminalMixin requires no initialisation.)
         MultiFunction.__init__(self)
         ufl2gem.Mixin.__init__(self)
 
+        # Need context during translation!
         self.context = context
 
     def modified_terminal(self, o):
@@ -283,14 +299,71 @@ def translate_facet_coordinate(terminal, mt, ctx):
     return ctx.point_expr
 
 
+@translate.register(SpatialCoordinate)
+def translate_spatialcoordinate(terminal, mt, ctx):
+    # Replace terminal with a Coefficient
+    terminal = ctx.coordinate(terminal.ufl_domain())
+    # Get back to reference space
+    terminal = preprocess_expression(terminal)
+    # Rebuild modified terminal
+    expr = construct_modified_terminal(mt, terminal)
+    # Translate replaced UFL snippet
+    return ctx.translator(expr)
+
+
+class CellVolumeKernelInterface(ProxyKernelInterface):
+    # Since CellVolume is evaluated as a cell integral, we must ensure
+    # that the right restriction is applied when it is used in an
+    # interior facet integral.  This proxy diverts coefficient
+    # translation to use a specified restriction.
+
+    def __init__(self, wrapee, restriction):
+        ProxyKernelInterface.__init__(self, wrapee)
+        self.restriction = restriction
+
+    def coefficient(self, ufl_coefficient, r):
+        assert r is None
+        return self._wrapee.coefficient(ufl_coefficient, self.restriction)
+
+
 @translate.register(CellVolume)
 def translate_cellvolume(terminal, mt, ctx):
-    return ctx.cellvolume(mt.restriction)
+    integrand, degree = one_times(ufl.dx(domain=terminal.ufl_domain()))
+    interface = CellVolumeKernelInterface(ctx, mt.restriction)
+
+    config = {name: getattr(ctx, name)
+              for name in ["ufl_cell", "precision", "index_cache"]}
+    config.update(interface=interface, quadrature_degree=degree)
+    expr, = compile_ufl(integrand, point_sum=True, **config)
+    return expr
 
 
 @translate.register(FacetArea)
 def translate_facetarea(terminal, mt, ctx):
-    return ctx.facetarea()
+    assert ctx.integral_type != 'cell'
+    domain = terminal.ufl_domain()
+    integrand, degree = one_times(ufl.Measure(ctx.integral_type, domain=domain))
+
+    config = {name: getattr(ctx, name)
+              for name in ["ufl_cell", "integration_dim",
+                           "entity_ids", "precision", "index_cache"]}
+    config.update(interface=ctx, quadrature_degree=degree)
+    expr, = compile_ufl(integrand, point_sum=True, **config)
+    return expr
+
+
+@translate.register(CellOrigin)
+def translate_cellorigin(terminal, mt, ctx):
+    domain = terminal.ufl_domain()
+    coords = SpatialCoordinate(domain)
+    expression = construct_modified_terminal(mt, coords)
+    point_set = PointSingleton((0.0,) * domain.topological_dimension())
+
+    config = {name: getattr(ctx, name)
+              for name in ["ufl_cell", "precision", "index_cache"]}
+    config.update(interface=ctx, point_set=point_set)
+    context = PointSetContext(**config)
+    return context.translator(expression)
 
 
 def fiat_to_ufl(fiat_dict, order):
@@ -414,8 +487,7 @@ def compile_ufl(expression, interior_facet=False, point_sum=False, **kwargs):
         expressions = [expression]
 
     # Translate UFL to GEM, lowering finite element specific nodes
-    translator = Translator(context)
-    result = map_expr_dags(translator, expressions)
+    result = map_expr_dags(context.translator, expressions)
     if point_sum:
         result = [gem.index_sum(expr, context.point_indices) for expr in result]
     return result

tsfc/kernel_interface/__init__.py

@@ -10,6 +10,11 @@ class KernelInterface(with_metaclass(ABCMeta)):
     """Abstract interface for accessing the GEM expressions corresponding
     to kernel arguments."""
 
+    @abstractmethod
+    def coordinate(self, ufl_domain):
+        """A function that maps :class:`ufl.Domain`s to coordinate
+        :class:`ufl.Coefficient`s."""
+
     @abstractmethod
     def coefficient(self, ufl_coefficient, restriction):
         """A function that maps :class:`ufl.Coefficient`s to GEM

tsfc/kernel_interface/common.py

@@ -23,9 +23,15 @@ def __init__(self, interior_facet=False):
         self.prepare = []
         self.finalise = []
 
+        # Coordinates
+        self.domain_coordinate = {}
+
         # Coefficients
         self.coefficient_map = {}
 
+    def coordinate(self, domain):
+        return self.domain_coordinate[domain]
+
     def coefficient(self, ufl_coefficient, restriction):
         """A function that maps :class:`ufl.Coefficient`s to GEM
         expressions."""

tsfc/kernel_interface/firedrake.py

@@ -181,12 +181,15 @@ def set_arguments(self, arguments, multiindices):
             arguments, multiindices, interior_facet=self.interior_facet)
         return expressions
 
-    def set_coordinates(self, coefficient):
+    def set_coordinates(self, domain):
         """Prepare the coordinate field.
 
-        :arg coefficient: :class:`ufl.Coefficient`
+        :arg domain: :class:`ufl.Domain`
         """
-        self.coordinates_arg = self._coefficient(coefficient, "coords")
+        # Create a fake coordinate coefficient for a domain.
+        f = Coefficient(FunctionSpace(domain, domain.ufl_coordinate_element()))
+        self.domain_coordinate[domain] = f
+        self.coordinates_arg = self._coefficient(f, "coords")
 
     def set_coefficients(self, integral_data, form_data):
         """Prepare the coefficients of the form.

tsfc/kernel_interface/ufc.py

@@ -12,6 +12,8 @@
 
 from finat import TensorFiniteElement
 
+import ufl
+
 from tsfc.kernel_interface.common import KernelBuilderBase
 from tsfc.finatinterface import create_element as _create_element
 from tsfc.coffee import SCALAR_TYPE
@@ -62,15 +64,17 @@ def set_arguments(self, arguments, multiindices):
         self.apply_glue(prepare)
         return expressions
 
-    def set_coordinates(self, coefficient, mode=None):
+    def set_coordinates(self, domain):
         """Prepare the coordinate field.
 
-        :arg coefficient: :class:`ufl.Coefficient`
-        :arg mode: (ignored)
+        :arg domain: :class:`ufl.Domain`
         """
+        # Create a fake coordinate coefficient for a domain.
+        f = ufl.Coefficient(ufl.FunctionSpace(domain, domain.ufl_coordinate_element()))
+        self.domain_coordinate[domain] = f
         self.coordinates_args, expression = prepare_coordinates(
-            coefficient, "coordinate_dofs", interior_facet=self.interior_facet)
-        self.coefficient_map[coefficient] = expression
+            f, "coordinate_dofs", interior_facet=self.interior_facet)
+        self.coefficient_map[f] = expression
 
     def set_coefficients(self, integral_data, form_data):
         """Prepare the coefficients of the form.