Fix inconsistent ufl usage in demos (#3532)

* Fix inconsitent ufl import in elasticity example

* Fix inconsistent ufl import in biharmonic example

* Fix inconsistent ufl import in cahn-hilliard example

* Fix inconsistent ufl import in hdg example

* Fix inconsistent ufl import in Helmholtz example

* Fix inconsistent ufl import in Lagrange variants example

* Fix inconsistent ufl import in poisson matrix free example

* Fix inconsistent ufl import in poisson example

* Fix inconsistent ufl import in pyamg example

* Fix inconsistent ufl import in stokes example

* Fix lagrange variants

* Add 'ufl' banned import from

* Switch to consistent 'import ufl' usage in demos

* New ruff version

* Remove comment

python/demo/demo_biharmonic.py

@@ -131,7 +131,6 @@
 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
 
 # -
 
@@ -192,8 +191,8 @@
 # sides of a facet.
 
 alpha = ScalarType(8.0)
-h = CellDiameter(msh)
-n = FacetNormal(msh)
+h = ufl.CellDiameter(msh)
+n = ufl.FacetNormal(msh)
 h_avg = (h("+") + h("-")) / 2.0
 
 # After that, we can define the variational problem consisting of the bilinear
@@ -210,15 +209,15 @@
 u = ufl.TrialFunction(V)
 v = ufl.TestFunction(V)
 x = ufl.SpatialCoordinate(msh)
-f = 4.0 * pi**4 * sin(pi * x[0]) * sin(pi * x[1])
+f = 4.0 * ufl.pi**4 * ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
 
 a = (
-    inner(div(grad(u)), div(grad(v))) * dx
-    - inner(avg(div(grad(u))), jump(grad(v), n)) * dS
-    - inner(jump(grad(u), n), avg(div(grad(v)))) * dS
-    + alpha / h_avg * inner(jump(grad(u), n), jump(grad(v), n)) * dS
+    ufl.inner(ufl.div(ufl.grad(u)), ufl.div(ufl.grad(v))) * ufl.dx
+    - ufl.inner(ufl.avg(ufl.div(ufl.grad(u))), ufl.jump(ufl.grad(v), n)) * ufl.dS
+    - ufl.inner(ufl.jump(ufl.grad(u), n), ufl.avg(ufl.div(ufl.grad(v)))) * ufl.dS
+    + alpha / h_avg * ufl.inner(ufl.jump(ufl.grad(u), n), ufl.jump(ufl.grad(v), n)) * ufl.dS
 )
-L = inner(f, v) * dx
+L = ufl.inner(f, v) * ufl.dx
 # -
 
 # We create a {py:class}`LinearProblem <dolfinx.fem.petsc.LinearProblem>`

python/demo/demo_cahn-hilliard.py

@@ -145,7 +145,6 @@
 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
 
 try:
     import pyvista as pv
@@ -200,7 +199,7 @@
 c0, mu0 = ufl.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.
 #
@@ -249,8 +248,16 @@
 # which is then used in the definition of the variational forms:
 
 # Weak statement of the equations
-F0 = inner(c, q) * dx - inner(c0, q) * dx + dt * inner(grad(mu_mid), grad(q)) * dx
-F1 = inner(mu, v) * dx - inner(dfdc, v) * dx - lmbda * inner(grad(c), grad(v)) * dx
+F0 = (
+    ufl.inner(c, q) * ufl.dx
+    - ufl.inner(c0, q) * ufl.dx
+    + dt * ufl.inner(ufl.grad(mu_mid), ufl.grad(q)) * ufl.dx
+)
+F1 = (
+    ufl.inner(mu, v) * ufl.dx
+    - ufl.inner(dfdc, v) * ufl.dx
+    - lmbda * ufl.inner(ufl.grad(c), ufl.grad(v)) * ufl.dx
+)
 F = F0 + F1
 
 # This is a statement of the time-discrete equations presented as part

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

@@ -41,7 +41,6 @@
 import ufl
 from dolfinx import fem, mesh
 from dolfinx.cpp.mesh import cell_num_entities
-from ufl import div, dot, grad, inner
 
 
 def par_print(comm, string):
@@ -147,17 +146,17 @@ def u_e(x):
 x = ufl.SpatialCoordinate(msh)
 c = 1.0 + 0.1 * ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
 a = (
-    inner(c * grad(u), grad(v)) * dx_c
-    - inner(c * (u - ubar), dot(grad(v), n)) * ds_c(cell_boundaries)
-    - inner(dot(grad(u), n), c * (v - vbar)) * ds_c(cell_boundaries)
-    + gamma * inner(c * (u - ubar), v - vbar) * ds_c(cell_boundaries)
+    ufl.inner(c * ufl.grad(u), ufl.grad(v)) * dx_c
+    - ufl.inner(c * (u - ubar), ufl.dot(ufl.grad(v), n)) * ds_c(cell_boundaries)
+    - ufl.inner(ufl.dot(ufl.grad(u), n), c * (v - vbar)) * ds_c(cell_boundaries)
+    + gamma * ufl.inner(c * (u - ubar), v - vbar) * ds_c(cell_boundaries)
 )
 
 # Manufacture a source term
-f = -div(c * grad(u_e(x)))
+f = -ufl.div(c * ufl.grad(u_e(x)))
 
-L = inner(f, v) * dx_c
-L += inner(fem.Constant(facet_mesh, dtype(0.0)), vbar) * dx_f
+L = ufl.inner(f, v) * dx_c
+L += ufl.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)

python/demo/demo_helmholtz.py

@@ -51,7 +51,6 @@
 from dolfinx.fem.petsc import LinearProblem
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import create_unit_square
-from ufl import FacetNormal, dot, ds, dx, grad, inner
 
 # Wavenumber
 k0 = 4 * np.pi
@@ -79,14 +78,14 @@
 
 # Define variational problem:
 u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
-a = inner(grad(u), grad(v)) * dx - k0**2 * inner(u, v) * dx
+a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx - k0**2 * ufl.inner(u, v) * ufl.dx
 
 # solve for plane wave with mixed Dirichlet and Neumann BCs
 theta = np.pi / 4
 u_exact.interpolate(lambda x: A * np.exp(1j * k0 * (np.cos(theta) * x[0] + np.sin(theta) * x[1])))
-n = FacetNormal(msh)
-g = -dot(n, grad(u_exact))
-L = -inner(g, v) * ds
+n = ufl.FacetNormal(msh)
+g = -ufl.dot(n, ufl.grad(u_exact))
+L = -ufl.inner(g, v) * ufl.ds
 
 
 dofs_D = locate_dofs_geometrical(
@@ -122,13 +121,17 @@
 
 # H1 errors
 diff = uh - u_exact
-H1_diff = msh.comm.allreduce(assemble_scalar(form(inner(grad(diff), grad(diff)) * dx)), op=MPI.SUM)
+H1_diff = msh.comm.allreduce(
+    assemble_scalar(form(ufl.inner(ufl.grad(diff), ufl.grad(diff)) * ufl.dx)), op=MPI.SUM
+)
 H1_exact = msh.comm.allreduce(
-    assemble_scalar(form(inner(grad(u_exact), grad(u_exact)) * dx)), op=MPI.SUM
+    assemble_scalar(form(ufl.inner(ufl.grad(u_exact), ufl.grad(u_exact)) * ufl.dx)), op=MPI.SUM
 )
 print("Relative H1 error of FEM solution:", abs(np.sqrt(H1_diff) / np.sqrt(H1_exact)))
 
 # L2 errors
-L2_diff = msh.comm.allreduce(assemble_scalar(form(inner(diff, diff) * dx)), op=MPI.SUM)
-L2_exact = msh.comm.allreduce(assemble_scalar(form(inner(u_exact, u_exact) * dx)), op=MPI.SUM)
+L2_diff = msh.comm.allreduce(assemble_scalar(form(ufl.inner(diff, diff) * ufl.dx)), op=MPI.SUM)
+L2_exact = msh.comm.allreduce(
+    assemble_scalar(form(ufl.inner(u_exact, u_exact) * ufl.dx)), op=MPI.SUM
+)
 print("Relative L2 error of FEM solution:", abs(np.sqrt(L2_diff) / np.sqrt(L2_exact)))

python/demo/demo_lagrange_variants.py

@@ -24,9 +24,8 @@
 
 import basix
 import basix.ufl
-import ufl  # type: ignore
+import ufl
 from dolfinx import default_real_type, fem, mesh
-from ufl import dx
 
 # -
 
@@ -175,7 +174,7 @@ def saw_tooth(x):
     V = fem.functionspace(msh, ufl_element)
     uh = fem.Function(V)
     uh.interpolate(lambda x: saw_tooth(x[0]))
-    M = fem.form((u_exact - uh) ** 2 * dx)
+    M = fem.form((u_exact - uh) ** 2 * ufl.dx)
     error = msh.comm.allreduce(fem.assemble_scalar(M), op=MPI.SUM)
     print(f"Computed L2 interpolation error ({variant.name}):", error**0.5)
 # -

python/demo/demo_mixed-poisson.py

@@ -110,20 +110,11 @@
 import numpy as np
 
 import dolfinx.fem.petsc
+import ufl
 from basix.ufl import element
 from dolfinx import fem, la, mesh
 from dolfinx.fem.petsc import discrete_gradient, interpolation_matrix
 from dolfinx.mesh import CellType, create_unit_square
-from ufl import (
-    Measure,
-    SpatialCoordinate,
-    TestFunction,
-    TrialFunction,
-    ZeroBaseForm,
-    div,
-    exp,
-    inner,
-)
 
 # Solution scalar (e.g., float32, complex128) and geometry (float32/64)
 # types
@@ -161,16 +152,16 @@
 # $W$, with corresponding test functions $\tau$ and $v$:
 
 # +
-(sigma, u) = TrialFunction(V), TrialFunction(W)
-(tau, v) = TestFunction(V), TestFunction(W)
+(sigma, u) = ufl.TrialFunction(V), ufl.TrialFunction(W)
+(tau, v) = ufl.TestFunction(V), ufl.TestFunction(W)
 # -
 
 # The source function is set to be $f = 10\exp(-((x_{0} - 0.5)^2 +
 # (x_{1} - 0.5)^2) / 0.02)$:
 
 # +
-x = SpatialCoordinate(msh)
-f = 10 * exp(-((x[0] - 0.5) * (x[0] - 0.5) + (x[1] - 0.5) * (x[1] - 0.5)) / 0.02)
+x = ufl.SpatialCoordinate(msh)
+f = 10 * ufl.exp(-((x[0] - 0.5) * (x[0] - 0.5) + (x[1] - 0.5) * (x[1] - 0.5)) / 0.02)
 # -
 
 # We now declare the blocked bilinear and linear forms. The rows of `a`
@@ -184,9 +175,12 @@
 # corresponds to $u_{0} = 0$ on $\Gamma_{D}$.*
 
 # +
-dx = Measure("dx", msh)
-a = [[inner(sigma, tau) * dx, inner(u, div(tau)) * dx], [inner(div(sigma), v) * dx, None]]
-L = [ZeroBaseForm((tau,)), -inner(f, v) * dx]
+dx = ufl.Measure("dx", msh)
+a = [
+    [ufl.inner(sigma, tau) * dx, ufl.inner(u, ufl.div(tau)) * dx],
+    [ufl.inner(ufl.div(sigma), v) * dx, None],
+]
+L = [ufl.ZeroBaseForm((tau,)), -ufl.inner(f, v) * dx]
 # -
 
 # We now compile the abstract/symbolic forms in `a` and `L` into
@@ -263,7 +257,10 @@
 
 # +
 a_p = fem.form(
-    [[inner(sigma, tau) * dx + inner(div(sigma), div(tau)) * dx, None], [None, inner(u, v) * dx]],
+    [
+        [ufl.inner(sigma, tau) * dx + ufl.inner(ufl.div(sigma), ufl.div(tau)) * dx, None],
+        [None, ufl.inner(u, v) * dx],
+    ],
     dtype=dtype,
 )
 P = fem.petsc.assemble_matrix_nest(a_p, bcs=bcs)