Replace UFL elements with Basix elements in tests and demos (#2381)

* Make Basix elements directly

* more direct basix

* flake

* fixing some tests

* cellname()

* equispaced variant

* UFL branch

* from basix.ufl_wrapper import ...

* flake

* UFL changes are merged

* flake

* correct test

* don't use V * W syntax. MixedElement([V, W]) instead

* ufl.VectorElement -> create_vector_element

* dim not gdim

* test against UFL branch with elements deleted

* update poisson.py

* update c++ demos

* update Python demos

* isort

* explicitly create mesh

* demos, basix branch

* basix main

* []

* ufl main

* branches

* rename ufl_wrapper -> ufl

* dlake

* one more

* remove create_ from function names

* missed a few

* update basix function name

* import enriched_element

* vector_element -> element(..., rank=1)

* more element function names

* vector_element -> element

* dim=k -> shape=(k, )

* correct two tests

* fix more tests

* fix

* element -> e

* shape not dim

* VectorElement -> BlockedElement

* Update comments

* update name

* use convenience functions not initialisers

* flake

* simplify test

* basix and ffc main

---------

Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>

.github/workflows/oneapi.yml

@@ -37,6 +37,7 @@ jobs:
         with:
           path: ./ffcx
           repository: FEniCS/ffcx
+          ref: main
       - name: Install FFCx C interface
         run: |
           . /opt/intel/oneapi/setvars.sh

cpp/cmake/scripts/copy-test-demo-data.py

@@ -29,7 +29,7 @@
 # Copy all files with the following suffixes
 suffix_patterns = ["txt", "h", "hpp", "c", "cpp", "py", "xdmf", "h5"]
 
-suffix_pattern = re.compile("(%s)," % ("|".join("[\w-]+\.%s" % pattern
+suffix_pattern = re.compile("(%s)," % ("|".join("[\\w-]+\\.%s" % pattern
                                                 for pattern in suffix_patterns)))
 
 script_rel_path = os.sep.join(__file__.split(os.sep)[:-1])
@@ -57,7 +57,7 @@ def copy_data(top_destdir, complex_mode):
     for subdir in sub_directories:
         top_dir = os.path.join(dolfinx_dir, subdir)
         for dirpath, dirnames, filenames in os.walk(top_dir):
-            if not dirpath in skip:
+            if dirpath not in skip:
                 destdir = dirpath.replace(dolfinx_dir, abs_destdir)
                 if not os.path.isdir(destdir):
                     os.makedirs(destdir)

cpp/demo/biharmonic/biharmonic.py

@@ -4,12 +4,12 @@
 # The first step is to define the variational problem at hand. We define
 # the variational problem in UFL terms in a separate form file
 # :download:`biharmonic.py`.  We begin by defining the finite element::
-from basix.ufl_wrapper import create_vector_element
-from ufl import (Coefficient, Constant, FiniteElement, FunctionSpace, Mesh,
-                 TestFunction, TrialFunction, dS, dx, grad, inner, triangle,
-                 FacetNormal, CellDiameter, jump, div, avg)
+from basix.ufl import element
+from ufl import (CellDiameter, Coefficient, Constant, FacetNormal,
+                 FunctionSpace, Mesh, TestFunction, TrialFunction, avg, div,
+                 dS, dx, grad, inner, jump)
 
-element = FiniteElement("Lagrange", triangle, 2)
+e = element("Lagrange", "triangle", 2)
 
 # The first argument to :py:class:`FiniteElement` is the finite element
 # family, the second argument specifies the domain, while the third
@@ -21,10 +21,10 @@
 # Next, we use this element to initialize the trial and test functions
 # (:math:`u` and :math:`v`) and the coefficient function :math:`f`::
 
-coord_element = create_vector_element("Lagrange", "triangle", 1)
+coord_element = element("Lagrange", "triangle", 1, rank=1)
 mesh = Mesh(coord_element)
 
-V = FunctionSpace(mesh, element)
+V = FunctionSpace(mesh, e)
 
 u = TrialFunction(V)
 v = TestFunction(V)

cpp/demo/hyperelasticity/hyperelasticity.py

@@ -8,18 +8,22 @@
 # We are interested in solving for a discrete vector field in three
 # dimensions, so first we need the appropriate finite element space and
 # trial and test functions on this space::
-from ufl import (Coefficient, Identity, TestFunction, TrialFunction,
-                 VectorElement, derivative, det, diff, dx, grad, ln,
-                 tetrahedron, tr, variable)
+from basix.ufl import element
+from ufl import (Coefficient, FunctionSpace, Identity, Mesh, TestFunction,
+                 TrialFunction, derivative, det, diff, dx, grad, ln, tr,
+                 variable)
 
 # Function spaces
-element = VectorElement("Lagrange", tetrahedron, 1)
+coord_element = element("Lagrange", "tetrahedron", 1, rank=1)
+mesh = Mesh(coord_element)
+e = element("Lagrange", "tetrahedron", 1, rank=1)
+V = FunctionSpace(mesh, e)
 
 # Trial and test functions
-du = TrialFunction(element)     # Incremental displacement
-v = TestFunction(element)      # Test function
+du = TrialFunction(V)     # Incremental displacement
+v = TestFunction(V)      # Test function
 
-# Note that ``VectorElement`` creates a finite element space of vector
+# Note that ``element`` with `rank=1` creates a finite element space of vector
 # fields. The dimension of the vector field (the number of components)
 # is assumed to be the same as the spatial dimension (in this case 3),
 # unless otherwise specified.
@@ -28,7 +32,7 @@
 # traction ``T`` and the displacement solution itself ``u``::
 
 # Functions
-u = Coefficient(element)        # Displacement from previous iteration
+u = Coefficient(V)        # Displacement from previous iteration
 # B = Coefficient(element)        # Body force per unit volume
 # T = Coefficient(element)        # Traction force on the boundary
 
@@ -74,5 +78,5 @@
 sigma = (1/J)*diff(psi, F)*F.T
 
 forms = [F_form, J_form]
-elements = [(element)]
+elements = [e]
 expressions = [(sigma, [[0.25, 0.25, 0.25]])]

cpp/demo/poisson/poisson.py

@@ -4,11 +4,11 @@
 # The first step is to define the variational problem at hand. We define
 # the variational problem in UFL terms in a separate form file
 # :download:`poisson.py`.  We begin by defining the finite element::
-from basix.ufl_wrapper import create_vector_element
-from ufl import (Coefficient, Constant, FiniteElement, FunctionSpace, Mesh,
-                 TestFunction, TrialFunction, ds, dx, grad, inner, triangle)
+from basix.ufl import element
+from ufl import (Coefficient, Constant, FunctionSpace, Mesh, TestFunction,
+                 TrialFunction, ds, dx, grad, inner)
 
-element = FiniteElement("Lagrange", triangle, 1)
+e = element("Lagrange", "triangle", 1)
 
 # The first argument to :py:class:`FiniteElement` is the finite element
 # family, the second argument specifies the domain, while the third
@@ -21,10 +21,10 @@
 # (:math:`u` and :math:`v`) and the coefficient functions (:math:`f` and
 # :math:`g`)::
 
-coord_element = create_vector_element("Lagrange", "triangle", 1)
+coord_element = element("Lagrange", "triangle", 1, rank=1)
 mesh = Mesh(coord_element)
 
-V = FunctionSpace(mesh, element)
+V = FunctionSpace(mesh, e)
 
 u = TrialFunction(V)
 v = TestFunction(V)

cpp/demo/poisson_matrix_free/poisson.py

@@ -1,16 +1,15 @@
 # UFL input for the Matrix-free Poisson Demo
 # ==================================
-from basix.ufl_wrapper import create_vector_element
-from ufl import (Coefficient, Constant, FiniteElement, FunctionSpace, Mesh,
-                 TestFunction, TrialFunction, action, dx, grad, inner,
-                 triangle)
+from basix.ufl import element
+from ufl import (Coefficient, Constant, FunctionSpace, Mesh, TestFunction,
+                 TrialFunction, action, dx, grad, inner)
 
-coord_element = create_vector_element("Lagrange", "triangle", 1)
+coord_element = element("Lagrange", "triangle", 1, rank=1)
 mesh = Mesh(coord_element)
 
 # Function Space
-element = FiniteElement("Lagrange", triangle, 2)
-V = FunctionSpace(mesh, element)
+e = element("Lagrange", "triangle", 2)
+V = FunctionSpace(mesh, e)
 
 # Trial and test functions
 u = TrialFunction(V)

cpp/dolfinx/fem/FiniteElement.h

@@ -78,7 +78,7 @@ class FiniteElement
   int space_dimension() const noexcept;
 
   /// Block size of the finite element function space. For
-  /// VectorElements and TensorElements, this is the number of DOFs
+  /// BlockedElements, this is the number of DOFs
   /// colocated at each DOF point. For other elements, this is always 1.
   /// @return Block size of the finite element space
   int block_size() const noexcept;
@@ -673,7 +673,7 @@ class FiniteElement
   // Dimension of each value space
   std::vector<std::size_t> _value_shape;
 
-  // Block size for VectorElements and TensorElements. This gives the
+  // Block size for BlockedElements. This gives the
   // number of DOFs co-located at each dof 'point'.
   int _bs;
 

cpp/test/poisson.py

@@ -1,14 +1,12 @@
-from ufl import (Coefficient, Constant, FiniteElement, FunctionSpace, Mesh,
-                 TestFunction, TrialFunction, VectorElement, dx, grad, inner,
-                 tetrahedron)
-import basix
-from basix.ufl_wrapper import create_vector_element
+from ufl import (Coefficient, Constant, FunctionSpace, Mesh,
+                 TestFunction, TrialFunction, dx, grad, inner)
+from basix.ufl import element
 
-element = FiniteElement("Lagrange", tetrahedron, 2)
-coord_element = create_vector_element("Lagrange", "tetrahedron", 1)
+e = element("Lagrange", "tetrahedron", 2)
+coord_element = element("Lagrange", "tetrahedron", 1, rank=1)
 mesh = Mesh(coord_element)
 
-V = FunctionSpace(mesh, element)
+V = FunctionSpace(mesh, e)
 
 u = TrialFunction(V)
 v = TestFunction(V)

python/demo/demo_axis/demo_axis.py

@@ -22,6 +22,7 @@
 from scipy.special import jv, jvp
 
 import ufl
+from basix.ufl import mixed_element, element
 from dolfinx import fem, io, mesh, plot
 
 from mpi4py import MPI
@@ -368,9 +369,9 @@ def create_eps_mu(pml, rho, eps_bkg, mu_bkg):
 # will use Lagrange elements:
 
 degree = 3
-curl_el = ufl.FiniteElement("N1curl", msh.ufl_cell(), degree)
-lagr_el = ufl.FiniteElement("Lagrange", msh.ufl_cell(), degree)
-V = fem.FunctionSpace(msh, ufl.MixedElement([curl_el, lagr_el]))
+curl_el = element("N1curl", msh.ufl_cell().cellname(), degree)
+lagr_el = element("Lagrange", msh.ufl_cell().cellname(), degree)
+V = fem.FunctionSpace(msh, mixed_element([curl_el, lagr_el]))
 
 # The integration domains of our problem are the following:
 
@@ -624,7 +625,7 @@ def create_eps_mu(pml, rho, eps_bkg, mu_bkg):
 assert err_ext < 0.01
 
 if has_vtx:
-    v_dg_el = ufl.VectorElement("DG", msh.ufl_cell(), degree, dim=3)
+    v_dg_el = element("DG", msh.ufl_cell().cellname(), degree, shape=(3, ))
     W = fem.FunctionSpace(msh, v_dg_el)
     Es_dg = fem.Function(W)
     Es_expr = fem.Expression(Esh, W.element.interpolation_points())

python/demo/demo_cahn-hilliard.py

@@ -123,6 +123,7 @@
 import numpy as np
 
 import ufl
+from basix.ufl import mixed_element, element
 from dolfinx import log, plot
 from dolfinx.fem import Function, FunctionSpace
 from dolfinx.fem.petsc import NonlinearProblem
@@ -160,8 +161,8 @@
 # using a pair of linear Lagrange elements.
 
 msh = create_unit_square(MPI.COMM_WORLD, 96, 96, CellType.triangle)
-P1 = ufl.FiniteElement("Lagrange", msh.ufl_cell(), 1)
-ME = FunctionSpace(msh, P1 * P1)
+P1 = element("Lagrange", msh.ufl_cell().cellname(), 1)
+ME = FunctionSpace(msh, mixed_element([P1, P1]))
 
 # Trial and test functions of the space `ME` are now defined:
 

python/demo/demo_lagrange_variants.py

@@ -22,7 +22,7 @@
 import numpy as np
 
 import basix
-import basix.ufl_wrapper
+import basix.ufl
 import ufl
 from dolfinx import fem, mesh
 from ufl import ds, dx, grad, inner
@@ -107,9 +107,8 @@
 # Elements created using Basix can be used directly with UFL via Basix's
 # UFL wrapper.
 
-element = basix.create_element(basix.ElementFamily.P, basix.CellType.triangle, 3,
-                               basix.LagrangeVariant.gll_warped)
-ufl_element = basix.ufl_wrapper.BasixElement(element)
+ufl_element = basix.ufl.element(basix.ElementFamily.P, basix.CellType.triangle, 3,
+                                basix.LagrangeVariant.gll_warped)
 
 # The UFL element `ufl_element` can be used in the same way as an
 # element created directly in UFL. For example, we could [solve a
@@ -163,8 +162,7 @@ def saw_tooth(x):
 u_exact = saw_tooth(x[0])
 
 for variant in [basix.LagrangeVariant.equispaced, basix.LagrangeVariant.gll_warped]:
-    element = basix.create_element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
-    ufl_element = basix.ufl_wrapper.BasixElement(element)
+    ufl_element = basix.ufl.element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
     V = fem.FunctionSpace(msh, ufl_element)
     uh = fem.Function(V)
     uh.interpolate(lambda x: saw_tooth(x[0]))
@@ -203,8 +201,7 @@ def saw_tooth(x):
 
 # +
 for variant in [basix.LagrangeVariant.equispaced, basix.LagrangeVariant.gll_warped]:
-    element = basix.create_element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
-    ufl_element = basix.ufl_wrapper.BasixElement(element)
+    ufl_element = basix.ufl.element(basix.ElementFamily.P, basix.CellType.interval, 10, variant)
     V = fem.FunctionSpace(msh, ufl_element)
     uh = fem.Function(V)
     uh.interpolate(lambda x: saw_tooth(x[0]))

python/demo/demo_mixed-poisson.py

@@ -86,9 +86,10 @@
 
 import numpy as np
 
+from basix.ufl import mixed_element, element
 from dolfinx import fem, io, mesh
-from ufl import (FiniteElement, Measure, MixedElement, SpatialCoordinate,
-                 TestFunctions, TrialFunctions, div, exp, inner)
+from ufl import (Measure, SpatialCoordinate, TestFunctions, TrialFunctions,
+                 div, exp, inner)
 
 from mpi4py import MPI
 from petsc4py import PETSc
@@ -100,9 +101,9 @@
 )
 
 k = 1
-Q_el = FiniteElement("BDMCF", domain.ufl_cell(), k)
-P_el = FiniteElement("DG", domain.ufl_cell(), k - 1)
-V_el = MixedElement([Q_el, P_el])
+Q_el = element("BDMCF", domain.ufl_cell().cellname(), k)
+P_el = element("DG", domain.ufl_cell().cellname(), k - 1)
+V_el = mixed_element([Q_el, P_el])
 V = fem.FunctionSpace(domain, V_el)
 
 (sigma, u) = TrialFunctions(V)

python/demo/demo_pml/demo_pml.py

@@ -37,6 +37,7 @@
 from mesh_wire_pml import generate_mesh_wire
 
 import ufl
+from basix.ufl import element
 from dolfinx import fem, mesh, plot
 from dolfinx.io import VTXWriter, gmshio
 
@@ -229,7 +230,7 @@ def pml_coordinates(x: ufl.indexed.Indexed, alpha: float, k0: complex, l_dom: fl
 # element to represent the electric field:
 
 degree = 3
-curl_el = ufl.FiniteElement("N1curl", msh.ufl_cell(), degree)
+curl_el = element("N1curl", msh.ufl_cell().cellname(), degree)
 V = fem.FunctionSpace(msh, curl_el)
 
 # Next, we interpolate $\mathbf{E}_b$ into the function space $V$,

python/demo/demo_scattering_boundary_conditions/demo_scattering_boundary_conditions.py

@@ -51,6 +51,7 @@
 from mesh_wire import generate_mesh_wire
 
 import ufl
+from basix.ufl import element
 from dolfinx import fem, io, plot
 
 from mpi4py import MPI
@@ -271,7 +272,7 @@ def curl_2d(f: fem.Function):
 # represent the electric field
 
 degree = 3
-curl_el = ufl.FiniteElement("N1curl", domain.ufl_cell(), degree)
+curl_el = element("N1curl", domain.ufl_cell().cellname(), degree)
 V = fem.FunctionSpace(domain, curl_el)
 
 # Next, we can interpolate $\mathbf{E}_b$ into the function space $V$:

python/demo/demo_static-condensation.py

@@ -28,6 +28,7 @@
 import numpy as np
 
 import ufl
+from basix.ufl import element
 from dolfinx import geometry
 from dolfinx.cpp.fem import Form_complex128, Form_float64
 from dolfinx.fem import (Function, FunctionSpace, IntegralType, dirichletbc,
@@ -47,8 +48,8 @@
 infile.close()
 
 # Stress (Se) and displacement (Ue) elements
-Se = ufl.TensorElement("DG", msh.ufl_cell(), 1, symmetry=True)
-Ue = ufl.VectorElement("Lagrange", msh.ufl_cell(), 2)
+Se = element("DG", msh.ufl_cell().cellname(), 1, rank=2, symmetry=True)
+Ue = element("Lagrange", msh.ufl_cell().cellname(), 2, rank=1)
 
 S = FunctionSpace(msh, Se)
 U = FunctionSpace(msh, Ue)

python/demo/demo_stokes.py

@@ -86,6 +86,7 @@
 import numpy as np
 
 import ufl
+from basix.ufl import mixed_element, element
 from dolfinx import fem, la
 from dolfinx.fem import (Constant, Function, FunctionSpace, dirichletbc,
                          extract_function_spaces, form,
@@ -129,8 +130,8 @@ def lid_velocity_expression(x):
 # linear basis (scalar).
 
 
-P2 = ufl.VectorElement("Lagrange", msh.ufl_cell(), 2)
-P1 = ufl.FiniteElement("Lagrange", msh.ufl_cell(), 1)
+P2 = element("Lagrange", msh.ufl_cell().cellname(), 2, rank=1)
+P1 = element("Lagrange", msh.ufl_cell().cellname(), 1)
 V, Q = FunctionSpace(msh, P2), FunctionSpace(msh, P1)
 
 # Boundary conditions for the velocity field are defined:
@@ -445,7 +446,7 @@ def block_direct_solver():
 def mixed_direct():
 
     # Create the Taylot-Hood function space
-    TH = P2 * P1
+    TH = mixed_element([P2, P1])
     W = FunctionSpace(msh, TH)
 
     # No slip boundary condition