Update UFL element interface: move elements to FInAT

.github/workflows/build.yml

@@ -87,6 +87,7 @@ jobs:
           python -m pip install \
             pytest-cov pytest-timeout pytest-xdist pytest-timeout
           python -m pip list
+
       - name: Test Firedrake
         run: |
           . ../firedrake_venv/bin/activate

AUTHORS.rst

@@ -148,6 +148,8 @@ Tomasz J. Salwa
 
 Kaho Sato
 
+Matthew Scroggs...............<https://mscroggs.co.uk>
+
 Ben Sepanski..................<https://bensepanski.github.io>
 
 Daniel R. Shapero.............<https://psc.apl.uw.edu/people/investigators/daniel-shapero/>

docs/source/team.ini

@@ -107,6 +107,7 @@ Francis P. Russell:
 Thomas Roy:
 Tomasz J. Salwa:
 Kaho Sato:
+Matthew Scroggs: https://mscroggs.co.uk
 Ben Sepanski: https://bensepanski.github.io
 Jemma Shipton:
 Joseph G. Wallwork: https://www.imperial.ac.uk/people/j.wallwork16

firedrake/__init__.py

@@ -54,6 +54,7 @@
     pass
 del ufl
 from ufl import *
+from finat.ufl import *
 # Set up the cache directories before importing PyOP2.
 firedrake_configuration.setup_cache_dirs()
 

firedrake/adjoint_utils/constant.py

@@ -95,7 +95,7 @@ def _ad_dot(self, other, options=None):
 
     @staticmethod
     def _ad_assign_numpy(dst, src, offset):
-        l = dst.ufl_element().value_size()
+        l = dst.ufl_element().value_size
         dst.assign(numpy.reshape(src[offset:offset + l], dst.ufl_shape), annotate=False)
         offset += l
         return dst, offset

firedrake/assemble.py

@@ -14,6 +14,7 @@
 from tsfc.finatinterface import create_element
 from tsfc.ufl_utils import extract_firedrake_constants
 import ufl
+import finat.ufl
 from firedrake import (extrusion_utils as eutils, matrix, parameters, solving,
                        tsfc_interface, utils)
 from firedrake.adjoint_utils import annotate_assemble
@@ -1294,7 +1295,7 @@ def _as_global_kernel_arg_coefficient(_, self):
 
     ufl_element = V.ufl_element()
     if ufl_element.family() == "Real":
-        return op2.GlobalKernelArg((ufl_element.value_size(),))
+        return op2.GlobalKernelArg((ufl_element.value_size,))
     else:
         finat_element = create_element(ufl_element)
         return self._make_dat_global_kernel_arg(finat_element, index)
@@ -1310,7 +1311,7 @@ def _as_global_kernel_arg_constant(_, self):
 @_as_global_kernel_arg.register(kernel_args.CellSizesKernelArg)
 def _as_global_kernel_arg_cell_sizes(_, self):
     # this mirrors tsfc.kernel_interface.firedrake_loopy.KernelBuilder.set_cell_sizes
-    ufl_element = ufl.FiniteElement("P", self._mesh.ufl_cell(), 1)
+    ufl_element = finat.ufl.FiniteElement("P", self._mesh.ufl_cell(), 1)
     finat_element = create_element(ufl_element)
     return self._make_dat_global_kernel_arg(finat_element)
 
@@ -1337,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 = ufl.FiniteElement("DG", cell=self._form.ufl_domain().ufl_cell(), degree=0)
+    ufl_element = finat.ufl.FiniteElement("DG", cell=self._form.ufl_domain().ufl_cell(), degree=0)
     finat_element = create_element(ufl_element)
     return self._make_dat_global_kernel_arg(finat_element)
 

firedrake/assign.py

@@ -5,7 +5,7 @@
 from pyadjoint.tape import annotate_tape
 from pyop2.utils import cached_property
 import pytools
-import ufl
+import finat.ufl
 from ufl.algorithms import extract_coefficients
 from ufl.constantvalue import as_ufl
 from ufl.corealg.map_dag import map_expr_dag
@@ -155,9 +155,9 @@ def __init__(self, assignee, expression, subset=None):
                     raise ValueError("All functions in the expression must use the same "
                                      "mesh as the assignee")
 
-        if (subset and type(assignee.ufl_element()) == ufl.MixedElement
+        if (subset and type(assignee.ufl_element()) == finat.ufl.MixedElement
                 and any(el.family() == "Real"
-                        for el in assignee.ufl_element().sub_elements())):
+                        for el in assignee.ufl_element().sub_elements)):
             raise ValueError("Subset is not a valid argument for assigning to a mixed "
                              "element including a real element")
 

firedrake/bcs.py

@@ -5,7 +5,8 @@
 import itertools
 
 import ufl
-from ufl import as_ufl, as_tensor, VectorElement
+from ufl import as_ufl, as_tensor
+from finat.ufl import VectorElement
 import finat
 
 import pyop2 as op2
@@ -329,12 +330,12 @@ def function_arg(self, g):
                 raise RuntimeError("%r is defined on incompatible FunctionSpace!" % g)
             self._function_arg = g
         elif isinstance(g, ufl.classes.Zero):
-            if g.ufl_shape and g.ufl_shape != self.function_space().ufl_element().value_shape():
+            if g.ufl_shape and g.ufl_shape != self.function_space().ufl_element().value_shape:
                 raise ValueError(f"Provided boundary value {g} does not match shape of space")
             # Special case. Scalar zero for direct Function.assign.
             self._function_arg = g
         elif isinstance(g, ufl.classes.Expr):
-            if g.ufl_shape != self.function_space().ufl_element().value_shape():
+            if g.ufl_shape != self.function_space().ufl_element().value_shape:
                 raise RuntimeError(f"Provided boundary value {g} does not match shape of space")
             try:
                 self._function_arg = firedrake.Function(self.function_space())

firedrake/checkpointing.py

@@ -1,7 +1,7 @@
 import functools
 import pickle
 from petsc4py.PETSc import ViewerHDF5
-import ufl
+import finat.ufl
 from pyop2 import op2
 from pyop2.mpi import COMM_WORLD, internal_comm, decref, MPI
 from firedrake.cython import hdf5interface as h5i
@@ -577,7 +577,7 @@ def save_mesh(self, mesh, distribution_name=None, permutation_name=None):
                     # Save tmesh.layers, which contains (start layer, stop layer)-tuple for each cell
                     # Conceptually, we project these integer pairs onto DG0 vector space of dim=2.
                     cell = base_tmesh.ufl_cell()
-                    element = ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2)
+                    element = finat.ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2)
                     layers_tV = impl.FunctionSpace(base_tmesh, element)
                     self._save_function_space_topology(layers_tV)
                     # Note that _cell_numbering coincides with DG0 section, so we can use tmesh.layers directly.
@@ -762,7 +762,7 @@ def _save_function_space_topology(self, tV):
         path = self._path_to_dms(tmesh.name)
         if dm_name not in self.require_group(path):
             if element.family() == "Real":
-                assert not isinstance(element, (ufl.VectorElement, ufl.TensorElement))
+                assert not isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement))
             else:
                 dm = self._get_dm_for_checkpointing(tV)
                 topology_dm = tmesh.topology_dm
@@ -845,7 +845,7 @@ def _save_function_topology(self, tf, idx=None):
         tmesh = tV.mesh()
         element = tV.ufl_element()
         if element.family() == "Real":
-            assert not isinstance(element, (ufl.VectorElement, ufl.TensorElement))
+            assert not isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement))
             dm_name = self._get_dm_name_for_checkpointing(tmesh, element)
             path = self._path_to_vec(tmesh.name, dm_name, tf.name())
             self.require_group(path)
@@ -894,7 +894,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
             variable_layers = self.get_attr(path, PREFIX_EXTRUDED + "_variable_layers")
             if variable_layers:
                 cell = base_tmesh.ufl_cell()
-                element = ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2)
+                element = finat.ufl.VectorElement("DP" if cell.is_simplex() else "DQ", cell, 0, dim=2)
                 _ = self._load_function_space_topology(base_tmesh, element)
                 base_tmesh_key = self._generate_mesh_key_from_names(base_tmesh.name,
                                                                     base_tmesh._distribution_name,
@@ -955,7 +955,7 @@ def load_mesh(self, name=DEFAULT_MESH_NAME, reorder=None, distribution_parameter
             path = self._path_to_mesh_immersed(tmesh.name, name)
             if path in self.h5pyfile:
                 cell = tmesh.ufl_cell()
-                element = ufl.FiniteElement("DP" if cell.is_simplex() else "DQ", cell, 0)
+                element = finat.ufl.FiniteElement("DP" if cell.is_simplex() else "DQ", cell, 0)
                 cell_orientations_tV = self._load_function_space_topology(tmesh, element)
                 tmesh_key = self._generate_mesh_key_from_names(tmesh.name,
                                                                tmesh._distribution_name,
@@ -1208,7 +1208,7 @@ def _load_function_topology(self, tmesh, element, tf_name, idx=None):
             self.viewer.pushTimestepping()
             self.viewer.setTimestep(idx)
         if element.family() == "Real":
-            assert not isinstance(element, (ufl.VectorElement, ufl.TensorElement))
+            assert not isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement))
             value = self.get_attr(path, "_".join([PREFIX, "value" if idx is None else "value_" + str(idx)]))
             tf.dat.data.itemset(value)
         else:
@@ -1247,10 +1247,10 @@ def _generate_function_space_name(self, V):
         V_names = [PREFIX + "_function_space"]
         for Vsub in V:
             elem = Vsub.ufl_element()
-            if isinstance(elem, ufl.RestrictedElement):
+            if isinstance(elem, finat.ufl.RestrictedElement):
                 # RestrictedElement.shortstr() contains '<>|{}'.
                 elem_name = "RestrictedElement(%s,%s)" % (elem.sub_element().shortstr(), elem.restriction_domain())
-            elif isinstance(elem, ufl.EnrichedElement):
+            elif isinstance(elem, finat.ufl.EnrichedElement):
                 # EnrichedElement.shortstr() contains '<>+'.
                 elem_name = "EnrichedElement(%s)" % ",".join(e.shortstr() for e in elem._elements)
             else:
@@ -1274,19 +1274,19 @@ def _get_shared_data_key_for_checkpointing(self, mesh, ufl_element):
         real_tensorproduct = eutils.is_real_tensor_product_element(finat_element)
         entity_dofs = finat_element.entity_dofs()
         nodes_per_entity = tuple(mesh.make_dofs_per_plex_entity(entity_dofs))
-        if isinstance(ufl_element, ufl.TensorElement):
-            shape = ufl_element.reference_value_shape()
+        if isinstance(ufl_element, finat.ufl.TensorElement):
+            shape = ufl_element.reference_value_shape
             block_size = np.product(shape)
-        elif isinstance(ufl_element, ufl.VectorElement):
-            shape = ufl_element.value_shape()[:1]
+        elif isinstance(ufl_element, finat.ufl.VectorElement):
+            shape = ufl_element.value_shape[:1]
             block_size = np.product(shape)
         else:
             block_size = 1
         return (nodes_per_entity, real_tensorproduct, block_size)
 
     def _get_dm_for_checkpointing(self, tV):
         sd_key = self._get_shared_data_key_for_checkpointing(tV.mesh(), tV.ufl_element())
-        if isinstance(tV.ufl_element(), (ufl.VectorElement, ufl.TensorElement)):
+        if isinstance(tV.ufl_element(), (finat.ufl.VectorElement, finat.ufl.TensorElement)):
             nodes_per_entity, real_tensorproduct, block_size = sd_key
             global_numbering = tV.mesh().create_section(nodes_per_entity, real_tensorproduct, block_size=block_size)
             topology_dm = tV.mesh().topology_dm
@@ -1421,6 +1421,7 @@ def _load_ufl_element(self, path, name):
             globals = {}
             locals = {}
             exec("from ufl import *", globals, locals)
+            exec("from finat.ufl import *", globals, locals)
             return eval(self.get_attr(path, name + "_repr"), globals, locals)
         else:
             return self._unpickle(self.get_attr(path, name))  # backward compat.

firedrake/constant.py

@@ -1,5 +1,6 @@
 import numpy as np
 import ufl
+import finat.ufl
 
 from tsfc.ufl_utils import TSFCConstantMixin
 from pyop2 import op2
@@ -67,14 +68,18 @@ def __new__(cls, value, domain=None, name=None, count=None):
 
             dat, rank, shape = _create_dat(op2.Global, value, domain._comm)
 
-            domain = ufl.as_domain(domain)
+            if not isinstance(domain, ufl.AbstractDomain):
+                cell = ufl.as_cell(domain)
+                coordinate_element = finat.ufl.VectorElement("Lagrange", cell, 1, gdim=cell.geometric_dimension)
+                domain = ufl.Mesh(coordinate_element)
+
             cell = domain.ufl_cell()
             if rank == 0:
-                element = ufl.FiniteElement("R", cell, 0)
+                element = finat.ufl.FiniteElement("R", cell, 0)
             elif rank == 1:
-                element = ufl.VectorElement("R", cell, 0, shape[0])
+                element = finat.ufl.VectorElement("R", cell, 0, shape[0])
             else:
-                element = ufl.TensorElement("R", cell, 0, shape=shape)
+                element = finat.ufl.TensorElement("R", cell, 0, shape=shape)
 
             R = FunctionSpace(domain, element, name="firedrake.Constant")
             return Function(R, val=dat).assign(value)

firedrake/embedding.py

@@ -1,32 +1,33 @@
 # -*- coding: utf-8 -*-
 """Module for utility functions for scalable HDF5 I/O."""
+import finat.ufl
 import ufl
 
 
 def get_embedding_dg_element(element):
-    cell = element.cell()
+    cell = element.cell
     degree = element.degree()
     family = lambda c: "DG" if c.is_simplex() else "DQ"
     if isinstance(cell, ufl.TensorProductCell):
         if type(degree) is int:
-            scalar_element = ufl.FiniteElement("DQ", cell=cell, degree=degree)
+            scalar_element = finat.ufl.FiniteElement("DQ", cell=cell, degree=degree)
         else:
-            scalar_element = ufl.TensorProductElement(*(ufl.FiniteElement(family(c), cell=c, degree=d)
-                                                        for (c, d) in zip(cell.sub_cells(), degree)))
+            scalar_element = finat.ufl.TensorProductElement(*(finat.ufl.FiniteElement(family(c), cell=c, degree=d)
+                                                              for (c, d) in zip(cell.sub_cells(), degree)))
     else:
-        scalar_element = ufl.FiniteElement(family(cell), cell=cell, degree=degree)
-    shape = element.value_shape()
+        scalar_element = finat.ufl.FiniteElement(family(cell), cell=cell, degree=degree)
+    shape = element.value_shape
     if len(shape) == 0:
         DG = scalar_element
     elif len(shape) == 1:
         shape, = shape
-        DG = ufl.VectorElement(scalar_element, dim=shape)
+        DG = finat.ufl.VectorElement(scalar_element, dim=shape)
     else:
-        if isinstance(element, ufl.TensorElement):
+        if isinstance(element, finat.ufl.TensorElement):
             symmetry = element.symmetry()
         else:
             symmetry = None
-        DG = ufl.TensorElement(scalar_element, shape=shape, symmetry=symmetry)
+        DG = finat.ufl.TensorElement(scalar_element, shape=shape, symmetry=symmetry)
     return DG
 
 
@@ -43,19 +44,19 @@ def get_embedding_element_for_checkpointing(element):
 
 def get_embedding_method_for_checkpointing(element):
     """Return the method used to embed element in dg space."""
-    if isinstance(element, (ufl.HDivElement, ufl.HCurlElement, ufl.WithMapping)):
+    if isinstance(element, (finat.ufl.HDivElement, finat.ufl.HCurlElement, finat.ufl.WithMapping)):
         return "project"
-    elif isinstance(element, (ufl.VectorElement, ufl.TensorElement)):
-        elem, = set(element.sub_elements())
+    elif isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement)):
+        elem, = set(element.sub_elements)
         return get_embedding_method_for_checkpointing(elem)
     elif element.family() in ['Lagrange', 'Discontinuous Lagrange',
                               'Nedelec 1st kind H(curl)', 'Raviart-Thomas',
                               'Nedelec 2nd kind H(curl)', 'Brezzi-Douglas-Marini',
                               'Q', 'DQ',
                               'S', 'DPC', 'Real']:
         return "interpolate"
-    elif isinstance(element, ufl.TensorProductElement):
-        methods = [get_embedding_method_for_checkpointing(elem) for elem in element.sub_elements()]
+    elif isinstance(element, finat.ufl.TensorProductElement):
+        methods = [get_embedding_method_for_checkpointing(elem) for elem in element.sub_elements]
         if any(method == "project" for method in methods):
             return "project"
         else:

firedrake/extrusion_utils.py

@@ -50,7 +50,7 @@ def make_extruded_coords(extruded_topology, base_coords, ext_coords,
     coordinates on the extruded cell (to write to), the fixed layer
     height, and the current cell layer.
     """
-    _, vert_space = ext_coords.function_space().ufl_element().sub_elements()[0].sub_elements()
+    _, vert_space = ext_coords.function_space().ufl_element().sub_elements[0].sub_elements
     if kernel is None and not (vert_space.degree() == 1
                                and vert_space.family() in ['Lagrange',
                                                            'Discontinuous Lagrange']):

firedrake/formmanipulation.py

@@ -125,7 +125,7 @@ def argument(self, o):
             else:
                 args += [Zero()
                          for j in numpy.ndindex(
-                         V_is[i].ufl_element().value_shape())]
+                         V_is[i].ufl_element().value_shape)]
         return self._arg_cache.setdefault(o, as_vector(args))