Highlight usage of extract_blocks.

python/demo/demo_hdg.py

@@ -112,11 +112,11 @@ def u_e(x):
 V = fem.functionspace(msh, ("Discontinuous Lagrange", k))
 Vbar = fem.functionspace(facet_mesh, ("Discontinuous Lagrange", k))
 
-# Trial and test functions
-# Cell space
-u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
-# Facet space
-ubar, vbar = ufl.TrialFunction(Vbar), ufl.TestFunction(Vbar)
+# Trial and test functions in mixed space
+W = ufl.MixedFunctionSpace(V, Vbar)
+u, ubar = ufl.TrialFunctions(W)
+v, vbar = ufl.TestFunctions(W)
+
 
 # Define integration measures
 # Cell
@@ -146,31 +146,26 @@ 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_00 = fem.form(
+a = (
  inner(c * grad(u), grad(v)) * dx_c
- - (
- inner(c * u, dot(grad(v), n)) * ds_c(cell_boundaries)
- + inner(dot(grad(u), n), c * v) * ds_c(cell_boundaries)
- )
+ - inner(c * u, dot(grad(v), n)) * ds_c(cell_boundaries)
+ - inner(dot(grad(u), n), c * v) * ds_c(cell_boundaries)
  + gamma * inner(c * u, v) * ds_c(cell_boundaries)
 )
-a_10 = fem.form(
- inner(dot(grad(u), n) - gamma * u, c * vbar) * ds_c(cell_boundaries), entity_maps=entity_maps
-)
-a_01 = fem.form(
- inner(c * ubar, dot(grad(v), n) - gamma * v) * ds_c(cell_boundaries), entity_maps=entity_maps
-)
-a_11 = fem.form(gamma * inner(c * ubar, vbar) * ds_c(cell_boundaries), entity_maps=entity_maps)
+
+a += inner(dot(grad(u), n) - gamma * u, c * vbar) * ds_c(cell_boundaries)
+a += inner(c * ubar, dot(grad(v), n) - gamma * v) * ds_c(cell_boundaries)
+a += gamma * inner(c * ubar, vbar) * ds_c(cell_boundaries)
 
 # Manufacture a source term
 f = -div(c * grad(u_e(x)))
 
-L_0 = fem.form(inner(f, v) * dx_c)
-L_1 = fem.form(inner(fem.Constant(facet_mesh, dtype(0.0)), vbar) * dx_f)
+L = inner(f, v) * dx_c
+L += inner(fem.Constant(facet_mesh, dtype(0.0)), vbar) * dx_f
 
 # Define block structure
-a = [[a_00, a_01], [a_10, a_11]]
-L = [L_0, L_1]
+a_blocked = dolfinx.fem.form(ufl.extract_blocks(a), entity_maps=entity_maps)
+L_blocked = dolfinx.fem.form(ufl.extract_blocks(L))
 
 # Apply Dirichlet boundary conditions
 # We begin by locating the boundary facets of msh
@@ -185,9 +180,9 @@ def u_e(x):
 bc = fem.dirichletbc(dtype(0.0), dofs, Vbar)
 
 # Assemble the matrix and vector
-A = assemble_matrix_block(a, bcs=[bc])
+A = assemble_matrix_block(a_blocked, bcs=[bc])
 A.assemble()
-b = assemble_vector_block(L, a, bcs=[bc])
+b = assemble_vector_block(L_blocked, a_blocked, bcs=[bc])
 
 # Setup the solver
 ksp = PETSc.KSP().create(msh.comm)

python/demo/demo_navier-stokes.py

@@ -185,15 +185,17 @@
 from ufl import (
  CellDiameter,
  FacetNormal,
- TestFunction,
- TrialFunction,
+ MixedFunctionSpace,
+ TestFunctions,
+ TrialFunctions,
  avg,
  conditional,
  div,
  dot,
  dS,
  ds,
  dx,
+ extract_blocks,
  grad,
  gt,
  inner,
@@ -283,15 +285,16 @@ def f_expr(x):
 # Function spaces for the velocity and for the pressure
 V = fem.functionspace(msh, ("Raviart-Thomas", k + 1))
 Q = fem.functionspace(msh, ("Discontinuous Lagrange", k))
+VQ = MixedFunctionSpace(V, Q)
 
 # Funcion space for visualising the velocity field
 gdim = msh.geometry.dim
 W = fem.functionspace(msh, ("Discontinuous Lagrange", k + 1, (gdim,)))
 
 # Define trial and test functions
 
-u, v = TrialFunction(V), TestFunction(V)
-p, q = TrialFunction(Q), TestFunction(Q)
+u, p = TrialFunctions(VQ)
+v, q = TestFunctions(VQ)
 
 delta_t = fem.Constant(msh, default_real_type(t_end / num_time_steps))
 alpha = fem.Constant(msh, default_real_type(6.0 * k**2))
@@ -311,7 +314,7 @@ def jump(phi, n):
 
 
 # +
-a_00 = (1.0 / Re) * (
+a = (1.0 / Re) * (
  inner(grad(u), grad(v)) * dx
  - inner(avg(grad(u)), jump(v, n)) * dS
  - inner(jump(u, n), avg(grad(v))) * dS
@@ -320,19 +323,19 @@ def jump(phi, n):
  - inner(outer(u, n), grad(v)) * ds
  + (alpha / h) * inner(outer(u, n), outer(v, n)) * ds
 )
-a_01 = -inner(p, div(v)) * dx
-a_10 = -inner(div(u), q) * dx
+a -= inner(p, div(v)) * dx
+a -= inner(div(u), q) * dx
 
-a = fem.form([[a_00, a_01], [a_10, None]])
+a_blocked = fem.form(extract_blocks(a))
 
 f = fem.Function(W)
 u_D = fem.Function(V)
 u_D.interpolate(u_e_expr)
-L_0 = inner(f, v) * dx + (1 / Re) * (
+L = inner(f, v) * dx + (1 / Re) * (
  -inner(outer(u_D, n), grad(v)) * ds + (alpha / h) * inner(outer(u_D, n), outer(v, n)) * ds
 )
-L_1 = inner(fem.Constant(msh, default_real_type(0.0)), q) * dx
-L = fem.form([L_0, L_1])
+L += inner(fem.Constant(msh, default_real_type(0.0)), q) * dx
+L_blocked = fem.form(extract_blocks(L))
 
 # Boundary conditions
 boundary_facets = mesh.exterior_facet_indices(msh.topology)
@@ -341,9 +344,9 @@ def jump(phi, n):
 bcs = [bc_u]
 
 # Assemble Stokes problem
-A = assemble_matrix_block(a, bcs=bcs)
+A = assemble_matrix_block(a_blocked, bcs=bcs)
 A.assemble()
-b = assemble_vector_block(L, a, bcs=bcs)
+b = assemble_vector_block(L_blocked, a_blocked, bcs=bcs)
 
 # Create and configure solver
 ksp = PETSc.KSP().create(msh.comm) # type: ignore
@@ -370,7 +373,6 @@ def jump(phi, n):
  else:
  raise e
 
-
 # Split the solution
 u_h = fem.Function(V)
 p_h = fem.Function(Q)
@@ -408,29 +410,29 @@ def jump(phi, n):
 # +
 lmbda = conditional(gt(dot(u_n, n), 0), 1, 0)
 u_uw = lmbda("+") * u("+") + lmbda("-") * u("-")
-a_00 += (
+a += (
  inner(u / delta_t, v) * dx
  - inner(u, div(outer(v, u_n))) * dx
  + inner((dot(u_n, n))("+") * u_uw, v("+")) * dS
  + inner((dot(u_n, n))("-") * u_uw, v("-")) * dS
  + inner(dot(u_n, n) * lmbda * u, v) * ds
 )
-a = fem.form([[a_00, a_01], [a_10, None]])
+a_blocked = fem.form(extract_blocks(a))
 
-L_0 += inner(u_n / delta_t, v) * dx - inner(dot(u_n, n) * (1 - lmbda) * u_D, v) * ds
-L = fem.form([L_0, L_1])
+L += inner(u_n / delta_t, v) * dx - inner(dot(u_n, n) * (1 - lmbda) * u_D, v) * ds
+L_blocked = fem.form(extract_blocks(L))
 
 # Time stepping loop
 for n in range(num_time_steps):
  t += delta_t.value
 
  A.zeroEntries()
- fem.petsc.assemble_matrix_block(A, a, bcs=bcs) # type: ignore
+ fem.petsc.assemble_matrix_block(A, a_blocked, bcs=bcs) # type: ignore
  A.assemble()
 
  with b.localForm() as b_loc:
  b_loc.set(0)
- fem.petsc.assemble_vector_block(b, L, a, bcs=bcs) # type: ignore
+ fem.petsc.assemble_vector_block(b, L_blocked, a_blocked, bcs=bcs) # type: ignore
 
  # Compute solution
  ksp.solve(b, x)