Update UFL element interface: move elements to FInAT (#302)

* use ufl.legacy

* cell() -> cell

* sub_elements() -> sub_elements

* value_shape() -> value_shape

* Pass in space, not element

* remove unnecessary legacy import

* FunctionSpace, not element

* Use branch of FEniCS UFL

* change README to make tests run

* Revert "change README to make tests run"

This reverts commit b5a0ff9.

* move ufl.legacy to finat.ufl

* ufl main

* flake

* Use branch on firedrake fork of UFL

* Apply suggestions from code review

---------

Co-authored-by: Jack Betteridge <43041811+JDBetteridge@users.noreply.github.com>

tests/test_create_fiat_element.py

@@ -4,6 +4,7 @@
 from FIAT.discontinuous_lagrange import HigherOrderDiscontinuousLagrange as FIAT_DiscontinuousLagrange
 
 import ufl
+import finat.ufl
 from tsfc.finatinterface import create_element as _create_element
 
 
@@ -40,7 +41,7 @@ def triangle_names(request):
 
 @pytest.fixture
 def ufl_element(triangle_names):
-    return ufl.FiniteElement(triangle_names, ufl.triangle, 2)
+    return finat.ufl.FiniteElement(triangle_names, ufl.triangle, 2)
 
 
 def test_triangle_basic(ufl_element):
@@ -58,16 +59,16 @@ def tensor_name(request):
                 ids=lambda x: x.cellname(),
                 scope="module")
 def ufl_A(request, tensor_name):
-    return ufl.FiniteElement(tensor_name, request.param, 1)
+    return finat.ufl.FiniteElement(tensor_name, request.param, 1)
 
 
 @pytest.fixture
 def ufl_B(tensor_name):
-    return ufl.FiniteElement(tensor_name, ufl.interval, 1)
+    return finat.ufl.FiniteElement(tensor_name, ufl.interval, 1)
 
 
 def test_tensor_prod_simple(ufl_A, ufl_B):
-    tensor_ufl = ufl.TensorProductElement(ufl_A, ufl_B)
+    tensor_ufl = finat.ufl.TensorProductElement(ufl_A, ufl_B)
 
     tensor = create_element(tensor_ufl)
     A = create_element(ufl_A)
@@ -84,7 +85,7 @@ def test_tensor_prod_simple(ufl_A, ufl_B):
                           ('DP', FIAT.GaussLegendre),
                           ('DP L2', FIAT.GaussLegendre)])
 def test_interval_variant_default(family, expected_cls):
-    ufl_element = ufl.FiniteElement(family, ufl.interval, 3)
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3)
     assert isinstance(create_element(ufl_element), expected_cls)
 
 
@@ -96,42 +97,42 @@ def test_interval_variant_default(family, expected_cls):
                           ('DP L2', 'equispaced', FIAT_DiscontinuousLagrange),
                           ('DP L2', 'spectral', FIAT.GaussLegendre)])
 def test_interval_variant(family, variant, expected_cls):
-    ufl_element = ufl.FiniteElement(family, ufl.interval, 3, variant=variant)
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3, variant=variant)
     assert isinstance(create_element(ufl_element), expected_cls)
 
 
 def test_triangle_variant_spectral_fail():
-    ufl_element = ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral')
+    ufl_element = finat.ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral')
     with pytest.raises(ValueError):
         create_element(ufl_element)
 
 
 def test_triangle_variant_spectral_fail_l2():
-    ufl_element = ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral')
+    ufl_element = finat.ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral')
     with pytest.raises(ValueError):
         create_element(ufl_element)
 
 
 def test_quadrilateral_variant_spectral_q():
-    element = create_element(ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral'))
     assert isinstance(element.element.A, FIAT.GaussLobattoLegendre)
     assert isinstance(element.element.B, FIAT.GaussLobattoLegendre)
 
 
 def test_quadrilateral_variant_spectral_dq():
-    element = create_element(ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
     assert isinstance(element.element.A, FIAT.GaussLegendre)
     assert isinstance(element.element.B, FIAT.GaussLegendre)
 
 
 def test_quadrilateral_variant_spectral_dq_l2():
-    element = create_element(ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral'))
     assert isinstance(element.element.A, FIAT.GaussLegendre)
     assert isinstance(element.element.B, FIAT.GaussLegendre)
 
 
 def test_quadrilateral_variant_spectral_rtcf():
-    element = create_element(ufl.FiniteElement('RTCF', ufl.quadrilateral, 2, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('RTCF', ufl.quadrilateral, 2, variant='spectral'))
     assert isinstance(element.element._elements[0].A, FIAT.GaussLobattoLegendre)
     assert isinstance(element.element._elements[0].B, FIAT.GaussLegendre)
     assert isinstance(element.element._elements[1].A, FIAT.GaussLegendre)

tests/test_create_finat_element.py

@@ -1,6 +1,7 @@
 import pytest
 
 import ufl
+import finat.ufl
 import finat
 from tsfc.finatinterface import create_element, supported_elements
 
@@ -18,7 +19,7 @@ def triangle_names(request):
 
 @pytest.fixture
 def ufl_element(triangle_names):
-    return ufl.FiniteElement(triangle_names, ufl.triangle, 2)
+    return finat.ufl.FiniteElement(triangle_names, ufl.triangle, 2)
 
 
 def test_triangle_basic(ufl_element):
@@ -28,7 +29,7 @@ def test_triangle_basic(ufl_element):
 
 @pytest.fixture
 def ufl_vector_element(triangle_names):
-    return ufl.VectorElement(triangle_names, ufl.triangle, 2)
+    return finat.ufl.VectorElement(triangle_names, ufl.triangle, 2)
 
 
 def test_triangle_vector(ufl_element, ufl_vector_element):
@@ -48,16 +49,16 @@ def tensor_name(request):
                         ufl.quadrilateral],
                 ids=lambda x: x.cellname())
 def ufl_A(request, tensor_name):
-    return ufl.FiniteElement(tensor_name, request.param, 1)
+    return finat.ufl.FiniteElement(tensor_name, request.param, 1)
 
 
 @pytest.fixture
 def ufl_B(tensor_name):
-    return ufl.FiniteElement(tensor_name, ufl.interval, 1)
+    return finat.ufl.FiniteElement(tensor_name, ufl.interval, 1)
 
 
 def test_tensor_prod_simple(ufl_A, ufl_B):
-    tensor_ufl = ufl.TensorProductElement(ufl_A, ufl_B)
+    tensor_ufl = finat.ufl.TensorProductElement(ufl_A, ufl_B)
 
     tensor = create_element(tensor_ufl)
     A = create_element(ufl_A)
@@ -73,7 +74,7 @@ def test_tensor_prod_simple(ufl_A, ufl_B):
                           ('DP', finat.GaussLegendre),
                           ('DP L2', finat.GaussLegendre)])
 def test_interval_variant_default(family, expected_cls):
-    ufl_element = ufl.FiniteElement(family, ufl.interval, 3)
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3)
     assert isinstance(create_element(ufl_element), expected_cls)
 
 
@@ -85,36 +86,36 @@ def test_interval_variant_default(family, expected_cls):
                           ('DP L2', 'equispaced', finat.DiscontinuousLagrange),
                           ('DP L2', 'spectral', finat.GaussLegendre)])
 def test_interval_variant(family, variant, expected_cls):
-    ufl_element = ufl.FiniteElement(family, ufl.interval, 3, variant=variant)
+    ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3, variant=variant)
     assert isinstance(create_element(ufl_element), expected_cls)
 
 
 def test_triangle_variant_spectral_fail():
-    ufl_element = ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral')
+    ufl_element = finat.ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral')
     with pytest.raises(ValueError):
         create_element(ufl_element)
 
 
 def test_triangle_variant_spectral_fail_l2():
-    ufl_element = ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral')
+    ufl_element = finat.ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral')
     with pytest.raises(ValueError):
         create_element(ufl_element)
 
 
 def test_quadrilateral_variant_spectral_q():
-    element = create_element(ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral'))
     assert isinstance(element.product.factors[0], finat.GaussLobattoLegendre)
     assert isinstance(element.product.factors[1], finat.GaussLobattoLegendre)
 
 
 def test_quadrilateral_variant_spectral_dq():
-    element = create_element(ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral'))
     assert isinstance(element.product.factors[0], finat.GaussLegendre)
     assert isinstance(element.product.factors[1], finat.GaussLegendre)
 
 
 def test_quadrilateral_variant_spectral_dq_l2():
-    element = create_element(ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral'))
+    element = create_element(finat.ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral'))
     assert isinstance(element.product.factors[0], finat.GaussLegendre)
     assert isinstance(element.product.factors[1], finat.GaussLegendre)
 

tests/test_dual_evaluation.py

@@ -1,12 +1,13 @@
 import pytest
 import ufl
+import finat.ufl
 from tsfc.finatinterface import create_element
 from tsfc import compile_expression_dual_evaluation
 
 
 def test_ufl_only_simple():
-    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2))
     v = ufl.Coefficient(V)
     expr = ufl.inner(v, v)
     W = V
@@ -16,8 +17,8 @@ def test_ufl_only_simple():
 
 
 def test_ufl_only_spatialcoordinate():
-    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2))
     x, y = ufl.SpatialCoordinate(mesh)
     expr = x*y - y**2 + x
     W = V
@@ -27,30 +28,30 @@ def test_ufl_only_spatialcoordinate():
 
 
 def test_ufl_only_from_contravariant_piola():
-    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1))
     v = ufl.Coefficient(V)
     expr = ufl.inner(v, v)
-    W = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    W = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2))
     to_element = create_element(W.ufl_element())
     kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
     assert kernel.needs_external_coords is True
 
 
 def test_ufl_only_to_contravariant_piola():
-    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2))
     v = ufl.Coefficient(V)
     expr = ufl.as_vector([v, v])
-    W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    W = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1))
     to_element = create_element(W.ufl_element())
     kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
     assert kernel.needs_external_coords is True
 
 
 def test_ufl_only_shape_mismatch():
-    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1))
     v = ufl.Coefficient(V)
     expr = ufl.inner(v, v)
     assert expr.ufl_shape == ()

tests/test_estimated_degree.py

@@ -3,6 +3,7 @@
 import pytest
 
 import ufl
+import finat.ufl
 from tsfc import compile_form
 from tsfc.logging import logger
 
@@ -14,8 +15,8 @@ def emit(self, record):
 
 def test_estimated_degree():
     cell = ufl.tetrahedron
-    mesh = ufl.Mesh(ufl.VectorElement('P', cell, 1))
-    V = ufl.FunctionSpace(mesh, ufl.FiniteElement('P', cell, 1))
+    mesh = ufl.Mesh(finat.ufl.VectorElement('P', cell, 1))
+    V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement('P', cell, 1))
     f = ufl.Coefficient(V)
     u = ufl.TrialFunction(V)
     v = ufl.TestFunction(V)

tests/test_firedrake_972.py

@@ -1,8 +1,9 @@
 import numpy
 import pytest
 
-from ufl import (Mesh, FunctionSpace, VectorElement, TensorElement,
+from ufl import (Mesh, FunctionSpace,
                  Coefficient, TestFunction, interval, indices, dx)
+from finat.ufl import VectorElement, TensorElement
 from ufl.classes import IndexSum, Product, MultiIndex
 
 from tsfc import compile_form

tests/test_gem_failure.py

@@ -1,5 +1,6 @@
-from ufl import (triangle, tetrahedron, FiniteElement,
+from ufl import (triangle, tetrahedron, FunctionSpace, Mesh,
                  TrialFunction, TestFunction, inner, grad, dx, dS)
+from finat.ufl import FiniteElement, VectorElement
 from tsfc import compile_form
 from FIAT.hdiv_trace import TraceError
 import pytest
@@ -12,8 +13,10 @@ def test_cell_error(cell, degree):
     cell triggers `gem.Failure` to raise the TraceError exception.
     """
     trace_element = FiniteElement("HDiv Trace", cell, degree)
-    lambdar = TrialFunction(trace_element)
-    gammar = TestFunction(trace_element)
+    domain = Mesh(VectorElement("Lagrange", cell, 1))
+    space = FunctionSpace(domain, trace_element)
+    lambdar = TrialFunction(space)
+    gammar = TestFunction(space)
 
     with pytest.raises(TraceError):
         compile_form(lambdar * gammar * dx)
@@ -27,8 +30,10 @@ def test_gradient_error(cell, degree):
     exception.
     """
     trace_element = FiniteElement("HDiv Trace", cell, degree)
-    lambdar = TrialFunction(trace_element)
-    gammar = TestFunction(trace_element)
+    domain = Mesh(VectorElement("Lagrange", cell, 1))
+    space = FunctionSpace(domain, trace_element)
+    lambdar = TrialFunction(space)
+    gammar = TestFunction(space)
 
     with pytest.raises(TraceError):
         compile_form(inner(grad(lambdar('+')), grad(gammar('+'))) * dS)

tests/test_idempotency.py

@@ -1,4 +1,5 @@
 import ufl
+import finat.ufl
 from tsfc import compile_form
 import loopy
 import pytest
@@ -21,13 +22,13 @@ def coord_degree(request):
 
 @pytest.fixture
 def mesh(cell, coord_degree):
-    c = ufl.VectorElement("CG", cell, coord_degree)
+    c = finat.ufl.VectorElement("CG", cell, coord_degree)
     return ufl.Mesh(c)
 
 
-@pytest.fixture(params=[ufl.FiniteElement,
-                        ufl.VectorElement,
-                        ufl.TensorElement],
+@pytest.fixture(params=[finat.ufl.FiniteElement,
+                        finat.ufl.VectorElement,
+                        finat.ufl.TensorElement],
                 ids=["FE", "VE", "TE"])
 def V(request, mesh):
     return ufl.FunctionSpace(mesh, request.param("CG", mesh.ufl_cell(), 2))

tests/test_impero_loopy_flop_counts.py

@@ -5,10 +5,11 @@
 import numpy
 import loopy
 from tsfc import compile_form
-from ufl import (FiniteElement, FunctionSpace, Mesh, TestFunction,
-                 TrialFunction, VectorElement, dx, grad, inner,
+from ufl import (FunctionSpace, Mesh, TestFunction,
+                 TrialFunction, dx, grad, inner,
                  interval, triangle, quadrilateral,
                  TensorProductCell)
+from finat.ufl import FiniteElement, VectorElement
 from tsfc.parameters import target
 
 

tests/test_interpolation_factorisation.py

@@ -2,9 +2,9 @@
 import numpy
 import pytest
 
-from ufl import (Mesh, FunctionSpace, FiniteElement, VectorElement,
-                 TensorElement, Coefficient,
+from ufl import (Mesh, FunctionSpace, Coefficient,
                  interval, quadrilateral, hexahedron)
+from finat.ufl import FiniteElement, VectorElement, TensorElement
 
 from tsfc import compile_expression_dual_evaluation
 from tsfc.finatinterface import create_element
@@ -54,11 +54,11 @@ def test_sum_factorisation_scalar_tensor(mesh, element):
     source = element(degree - 1)
     target = element(degree)
     tensor_flops = flop_count(mesh, source, target)
-    expect = numpy.prod(target.value_shape())
+    expect = numpy.prod(target.value_shape)
     if isinstance(target, FiniteElement):
         scalar_flops = tensor_flops
     else:
-        target = target.sub_elements()[0]
-        source = source.sub_elements()[0]
+        target = target.sub_elements[0]
+        source = source.sub_elements[0]
         scalar_flops = flop_count(mesh, source, target)
     assert numpy.allclose(tensor_flops / scalar_flops, expect, rtol=1e-2)