Fix inconsistent ufl usage in demos

python/demo/demo_biharmonic.py

@@ -127,11 +127,25 @@
 
 # +
 import dolfinx
-import ufl
 from dolfinx import fem, io, mesh, plot
 from dolfinx.fem.petsc import LinearProblem
 from dolfinx.mesh import CellType, GhostMode
-from ufl import CellDiameter, FacetNormal, avg, div, dS, dx, grad, inner, jump, pi, sin
+from ufl import (
+    CellDiameter,
+    FacetNormal,
+    SpatialCoordinate,
+    TestFunction,
+    TrialFunction,
+    avg,
+    div,
+    dS,
+    dx,
+    grad,
+    inner,
+    jump,
+    pi,
+    sin,
+)
 
 # -
 
@@ -207,9 +221,9 @@
 
 # +
 # Define variational problem
-u = ufl.TrialFunction(V)
-v = ufl.TestFunction(V)
-x = ufl.SpatialCoordinate(msh)
+u = TrialFunction(V)
+v = TestFunction(V)
+x = SpatialCoordinate(msh)
 f = 4.0 * pi**4 * sin(pi * x[0]) * sin(pi * x[1])
 
 a = (

python/demo/demo_cahn-hilliard.py

@@ -137,15 +137,14 @@
 
 import numpy as np
 
-import ufl
 from basix.ufl import element, mixed_element
 from dolfinx import default_real_type, log, plot
 from dolfinx.fem import Function, functionspace
 from dolfinx.fem.petsc import NonlinearProblem
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import CellType, create_unit_square
 from dolfinx.nls.petsc import NewtonSolver
-from ufl import dx, grad, inner
+from ufl import TestFunctions, diff, dx, grad, inner, split, variable
 
 try:
     import pyvista as pv
@@ -179,7 +178,7 @@
 
 # Trial and test functions of the space `ME` are now defined:
 
-q, v = ufl.TestFunctions(ME)
+q, v = TestFunctions(ME)
 
 # ```{index} split functions
 # ```
@@ -196,11 +195,11 @@
 u0 = Function(ME)  # solution from previous converged step
 
 # Split mixed functions
-c, mu = ufl.split(u)
-c0, mu0 = ufl.split(u0)
+c, mu = split(u)
+c0, mu0 = split(u0)
 # -
 
-# The line `c, mu = ufl.split(u)` permits direct access to the
+# The line `c, mu = split(u)` permits direct access to the
 # components of a mixed function. Note that `c` and `mu` are references
 # for components of `u`, and not copies.
 #
@@ -232,9 +231,9 @@
 # differentiation:
 
 # Compute the chemical potential df/dc
-c = ufl.variable(c)
+c = variable(c)
 f = 100 * c**2 * (1 - c) ** 2
-dfdc = ufl.diff(f, c)
+dfdc = diff(f, c)
 
 # The first line declares that `c` is a variable that some function can
 # be differentiated with respect to. The next line is the function $f$

python/demo/demo_elasticity.py

@@ -53,7 +53,6 @@
 from dolfinx.fem.petsc import apply_lifting, assemble_matrix, assemble_vector
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import CellType, GhostMode, create_box, locate_entities_boundary
-from ufl import dx, grad, inner
 
 dtype = PETSc.ScalarType  # type: ignore
 # -
@@ -135,7 +134,7 @@ def build_nullspace(V: FunctionSpace):
 
 def σ(v):
     """Return an expression for the stress σ given a displacement field"""
-    return 2.0 * μ * ufl.sym(grad(v)) + λ * ufl.tr(ufl.sym(grad(v))) * ufl.Identity(len(v))
+    return 2.0 * μ * ufl.sym(ufl.grad(v)) + λ * ufl.tr(ufl.sym(ufl.grad(v))) * ufl.Identity(len(v))
 
 
 # -
@@ -146,8 +145,8 @@ def σ(v):
 
 V = functionspace(msh, ("Lagrange", 1, (msh.geometry.dim,)))
 u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
-a = form(inner(σ(u), grad(v)) * dx)
-L = form(inner(f, v) * dx)
+a = form(ufl.inner(σ(u), ufl.grad(v)) * ufl.dx)
+L = form(ufl.inner(f, v) * ufl.dx)
 
 # A homogeneous (zero) boundary condition is created on $x_0 = 0$ and
 # $x_1 = 1$ by finding all facets on these boundaries, and then creating
@@ -240,7 +239,7 @@ def σ(v):
 
 # +
 sigma_dev = σ(uh) - (1 / 3) * ufl.tr(σ(uh)) * ufl.Identity(len(uh))
-sigma_vm = ufl.sqrt((3 / 2) * inner(sigma_dev, sigma_dev))
+sigma_vm = ufl.sqrt((3 / 2) * ufl.inner(sigma_dev, sigma_dev))
 # -
 
 # Next, the Von Mises stress is interpolated in a piecewise-constant

python/demo/demo_hdg.py

@@ -38,10 +38,25 @@
 
 import numpy as np
 
-import ufl
 from dolfinx import fem, mesh
 from dolfinx.cpp.mesh import cell_num_entities
-from ufl import div, dot, grad, inner
+from ufl import (
+    CellDiameter,
+    FacetNormal,
+    Measure,
+    MixedFunctionSpace,
+    SpatialCoordinate,
+    TestFunctions,
+    TrialFunctions,
+    div,
+    dot,
+    dx,
+    extract_blocks,
+    grad,
+    inner,
+    pi,
+    sin,
+)
 
 
 def par_print(comm, string):
@@ -50,10 +65,8 @@ def par_print(comm, string):
         sys.stdout.flush()
 
 
-def norm_L2(comm, v, measure=ufl.dx):
-    return np.sqrt(
-        comm.allreduce(fem.assemble_scalar(fem.form(ufl.inner(v, v) * measure)), op=MPI.SUM)
-    )
+def norm_L2(comm, v, measure=dx):
+    return np.sqrt(comm.allreduce(fem.assemble_scalar(fem.form(inner(v, v) * measure)), op=MPI.SUM))
 
 
 def compute_cell_boundary_facets(msh):
@@ -78,7 +91,7 @@ def u_e(x):
     """Exact solution."""
     u_e = 1
     for i in range(tdim):
-        u_e *= ufl.sin(ufl.pi * x[i])
+        u_e *= sin(pi * x[i])
     return u_e
 
 
@@ -113,24 +126,24 @@ def u_e(x):
 Vbar = fem.functionspace(facet_mesh, ("Discontinuous Lagrange", k))
 
 # Trial and test functions in mixed space
-W = ufl.MixedFunctionSpace(V, Vbar)
-u, ubar = ufl.TrialFunctions(W)
-v, vbar = ufl.TestFunctions(W)
+W = MixedFunctionSpace(V, Vbar)
+u, ubar = TrialFunctions(W)
+v, vbar = TestFunctions(W)
 
 
 # Define integration measures
 # Cell
-dx_c = ufl.Measure("dx", domain=msh)
+dx_c = Measure("dx", domain=msh)
 # Cell boundaries
 # We need to define an integration measure to integrate around the
 # boundary of each cell. The integration entities can be computed
 # using the following convenience function.
 cell_boundary_facets = compute_cell_boundary_facets(msh)
 cell_boundaries = 1  # A tag
 # Create the measure
-ds_c = ufl.Measure("ds", subdomain_data=[(cell_boundaries, cell_boundary_facets)], domain=msh)
+ds_c = Measure("ds", subdomain_data=[(cell_boundaries, cell_boundary_facets)], domain=msh)
 # Create a cell integral measure over the facet mesh
-dx_f = ufl.Measure("dx", domain=facet_mesh)
+dx_f = Measure("dx", domain=facet_mesh)
 
 # We write the mixed domain forms as integrals over msh. Hence, we must
 # provide a map from facets in msh to cells in facet_mesh. This is the
@@ -140,12 +153,12 @@ def u_e(x):
 entity_maps = {facet_mesh: mesh_to_facet_mesh}
 
 # Define forms
-h = ufl.CellDiameter(msh)
-n = ufl.FacetNormal(msh)
+h = CellDiameter(msh)
+n = FacetNormal(msh)
 gamma = 16.0 * k**2 / h  # Scaled penalty parameter
 
-x = ufl.SpatialCoordinate(msh)
-c = 1.0 + 0.1 * ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
+x = SpatialCoordinate(msh)
+c = 1.0 + 0.1 * sin(pi * x[0]) * sin(pi * x[1])
 a = (
     inner(c * grad(u), grad(v)) * dx_c
     - inner(c * (u - ubar), dot(grad(v), n)) * ds_c(cell_boundaries)
@@ -160,8 +173,8 @@ def u_e(x):
 L += inner(fem.Constant(facet_mesh, dtype(0.0)), vbar) * dx_f
 
 # Define block structure
-a_blocked = dolfinx.fem.form(ufl.extract_blocks(a), entity_maps=entity_maps)
-L_blocked = dolfinx.fem.form(ufl.extract_blocks(L))
+a_blocked = dolfinx.fem.form(extract_blocks(a), entity_maps=entity_maps)
+L_blocked = dolfinx.fem.form(extract_blocks(L))
 
 # Apply Dirichlet boundary conditions
 # We begin by locating the boundary facets of msh
@@ -220,9 +233,9 @@ def u_e(x):
 
 
 # Compute errors
-x = ufl.SpatialCoordinate(msh)
+x = SpatialCoordinate(msh)
 e_u = norm_L2(msh.comm, u - u_e(x))
-x_bar = ufl.SpatialCoordinate(facet_mesh)
+x_bar = SpatialCoordinate(facet_mesh)
 e_ubar = norm_L2(msh.comm, ubar - u_e(x_bar))
 par_print(comm, f"e_u = {e_u}")
 par_print(comm, f"e_ubar = {e_ubar}")

python/demo/demo_helmholtz.py

@@ -37,12 +37,11 @@
 import numpy as np
 
 import dolfinx
-import ufl
 from dolfinx.fem import Function, assemble_scalar, form, functionspace
 from dolfinx.fem.petsc import LinearProblem
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import create_unit_square
-from ufl import dx, grad, inner
+from ufl import TestFunction, TrialFunction, dx, grad, inner
 
 # Wavenumber
 k0 = 4 * np.pi
@@ -65,7 +64,7 @@
 V = functionspace(msh, ("Lagrange", deg))
 
 # Define variational problem
-u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
+u, v = TrialFunction(V), TestFunction(V)
 f = Function(V)
 f.interpolate(lambda x: A * k0**2 * np.cos(k0 * x[0]) * np.cos(k0 * x[1]))
 a = inner(grad(u), grad(v)) * dx - k0**2 * inner(u, v) * dx

python/demo/demo_lagrange_variants.py

@@ -24,9 +24,8 @@
 
 import basix
 import basix.ufl
-import ufl  # type: ignore
 from dolfinx import default_real_type, fem, mesh
-from ufl import dx
+from ufl import SpatialCoordinate, dx
 
 # -
 
@@ -124,7 +123,7 @@ def saw_tooth(x):
 # +
 msh = mesh.create_unit_interval(MPI.COMM_WORLD, N)
 
-x = ufl.SpatialCoordinate(msh)
+x = SpatialCoordinate(msh)
 u_exact = saw_tooth(x[0])
 
 for variant in [basix.LagrangeVariant.equispaced, basix.LagrangeVariant.gll_warped]:

python/demo/demo_poisson.py

@@ -82,10 +82,9 @@
 # +
 import numpy as np
 
-import ufl
 from dolfinx import fem, io, mesh, plot
 from dolfinx.fem.petsc import LinearProblem
-from ufl import ds, dx, grad, inner
+from ufl import SpatialCoordinate, TestFunction, TrialFunction, ds, dx, exp, grad, inner, sin
 
 # -
 
@@ -141,11 +140,11 @@
 # Next, the variational problem is defined:
 
 # +
-u = ufl.TrialFunction(V)
-v = ufl.TestFunction(V)
-x = ufl.SpatialCoordinate(msh)
-f = 10 * ufl.exp(-((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) / 0.02)
-g = ufl.sin(5 * x[0])
+u = TrialFunction(V)
+v = TestFunction(V)
+x = SpatialCoordinate(msh)
+f = 10 * exp(-((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) / 0.02)
+g = sin(5 * x[0])
 a = inner(grad(u), grad(v)) * dx
 L = inner(f, v) * dx + inner(g, v) * ds
 # -

python/demo/demo_poisson_matrix_free.py

@@ -80,9 +80,8 @@
 import numpy as np
 
 import dolfinx
-import ufl
 from dolfinx import fem, la
-from ufl import action, dx, grad, inner
+from ufl import SpatialCoordinate, TestFunction, TrialFunction, action, dx, grad, inner
 
 # We begin by using {py:func}`create_rectangle
 # <dolfinx.mesh.create_rectangle>` to create a rectangular
@@ -133,9 +132,9 @@
 
 # Next, we express the variational problem using UFL.
 
-x = ufl.SpatialCoordinate(mesh)
-u = ufl.TrialFunction(V)
-v = ufl.TestFunction(V)
+x = SpatialCoordinate(mesh)
+u = TrialFunction(V)
+v = TestFunction(V)
 f = fem.Constant(mesh, dtype(-6.0))
 a = inner(grad(u), grad(v)) * dx
 L = inner(f, v) * dx