Prototype `ufl.MixedFunctionSpace`

[jorgensd]

To simplify extraction of blocks and ease variational formulations.

Usage:

V = dolfinx.fem.functionspace(...)
Q = scifem.create_real_space(...)
W = scifem.MixedFunctionSpace(....)
u, p = ufl.TrialFunctions(W)
v, q = ufl.TestFunctions(W)

a = ufl.inner(u, v)*ufl.dx + ufl.inner(p, v)*ufl.dx + ufl.inner(v, q)*ufl.dx

and internally use

a_ij = ufl.extract_blocks(a)

[jorgensd]

Almost works out of the box:

    V = dolfinx.fem.functionspace(mesh, ("Lagrange", degree, value_shape))
    R = scifem.create_real_functionspace(mesh, value_shape)

    W = ufl.MixedFunctionSpace(V, R)
    u, c = ufl.TrialFunctions(W)
    v, d = ufl.TestFunctions(W)
    x = ufl.SpatialCoordinate(mesh)
    pol = x[0] ** degree - 2 * x[1] ** degree
    # Compute average value of polynomial to make mean 0
    C = mesh.comm.allreduce(
        dolfinx.fem.assemble_scalar(dolfinx.fem.form(pol * ufl.dx, dtype=dtype)), op=MPI.SUM
    )
    u_scalar = pol - dolfinx.fem.Constant(mesh, dtype(C))
    if tensor == 0:
        u_ex = u_scalar
        zero = dolfinx.fem.Constant(mesh, dtype(0.0))
    elif tensor == 1:
        u_ex = ufl.as_vector([u_scalar, -u_scalar])
        zero = dolfinx.fem.Constant(mesh, dtype((0.0, 0.0)))
    else:
        u_ex = ufl.as_tensor(
            [
                [u_scalar, 2 * u_scalar],
                [3 * u_scalar, -u_scalar],
                [u_scalar, 2 * u_scalar],
            ]
        )
        zero = dolfinx.fem.Constant(mesh, dtype(((0.0, 0.0), (0.0, 0.0), (0.0, 0.0))))

    dx = ufl.Measure("dx", domain=mesh)
    f = -ufl.div(ufl.grad(u_ex))
    n = ufl.FacetNormal(mesh)
    g = ufl.dot(ufl.grad(u_ex), n)
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx
    a += ufl.inner(c, v) * dx
    #a += ufl.inner(u, d) * dx
    L = ufl.inner(f, v) * dx + ufl.inner(g, v) * ufl.ds
    L += ufl.inner(zero, d) * dx

    a_blocked = [[None for _ in range(W.num_sub_spaces())] for _ in range(W.num_sub_spaces())]
    L_blocked = [None for _ in range(W.num_sub_spaces())]
    for i in range(W.num_sub_spaces()):
        for j in range(W.num_sub_spaces()):
              a_blocked[i][j] = ufl.extract_blocks(a, i, j)
        L_blocked[i] = ufl.extract_blocks(L, i)
    a = dolfinx.fem.form(a_blocked, dtype=dtype)
    L = dolfinx.fem.form(L_blocked, dtype=dtype)

Need some fixes in both dolfinx and basix.

[finsberg]

For nonlinear problems I guess we need to do

W = ufl.MixedFunctionSpace(V, R)
w = dolfinx.fem.Function(W)
u, c = ufl.split(w)

But then I get

Traceback (most recent call last):
  File "/home/shared/mwe.py", line 180, in <module>
    main()
  File "/home/shared/mwe.py", line 77, in main
    w = dolfinx.fem.Function(W)
  File "/usr/local/dolfinx-real/lib/python3.10/dist-packages/dolfinx/fem/function.py", line 340, in __init__
    V.element.dtype, np.dtype(dtype).type(0).real.dtype
AttributeError: 'MixedFunctionSpace' object has no attribute 'element'

[jorgensd]

You would have to make functions in V and R, and later call derivative wrt TrialFunctions from W.sub(i) to make up the block jacobian

[astqx]

Thanks for this. I tried it with the old FEniCS demo Poisson equation with pure Neumann boundary conditions, which can finally be translated to FEniCSx. I know is not a non-linear problem, but I also wanted to test the blocked Newton solver (#5).

# ref: https://fenicsproject.org/olddocs/dolfin/latest/python/demos/neumann-poisson/demo_neumann-poisson.py.html

domain = dolfinx.mesh.create_unit_square(comm=MPI.COMM_WORLD, nx=128, ny=128, cell_type=dolfinx.mesh.CellType.triangle)

V = dolfinx.fem.functionspace(domain, ("Lagrange", 1, ()))
R = scifem.create_real_functionspace(domain)

v = ufl.TestFunction(V)
d = ufl.TestFunction(R)
du = ufl.TrialFunction(V)
dc = ufl.TrialFunction(R)
u = dolfinx.fem.Function(V)
c = dolfinx.fem.Function(R)

x = ufl.SpatialCoordinate(domain)
f = 10 * ufl.exp(-((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) / 0.02)
g = -ufl.sin(5 * x[0])
zero = dolfinx.fem.Constant(domain, dolfinx.default_scalar_type(0.0))

a0 = (ufl.inner(ufl.grad(u), ufl.grad(v)) + ufl.inner(c, v))*ufl.dx
a1 = ufl.inner(u, d)*ufl.dx
L0 = ufl.inner(f, v)*ufl.dx + ufl.inner(g, v)*ufl.ds
L1 = ufl.inner(zero, d)*ufl.dx

F0 = a0 - L0
F1 = a1 - L1 
F = [F0, F1]

J = [
    [ufl.derivative(F0, u, du), ufl.derivative(F0, c, dc)],
    [ufl.derivative(F1, u, du), ufl.derivative(F1, c, dc)],
]

petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
solver = NewtonSolver(F, J, [u, c], max_iterations=25, petsc_options=petsc_options)
solver.solve()

[jorgensd]

@astqx I am assuming you are trying to say that this does not work? Please note that I have an open PR for a NewtonSolver for blocked problems: https://github.com/scientificcomputing/scifem/pull/35

[astqx]

@astqx I am assuming you are trying to say that this does not work?

Please note that I have an open PR for a NewtonSolver for blocked problems:

https://github.com/scientificcomputing/scifem/pull/35

Sorry, I meant this does work! I've used the solver from that PR. It could be nice to include this problem as a demo in FEniCSx too.

[jorgensd]

@astqx I am assuming you are trying to say that this does not work? Please note that I have an open PR for a NewtonSolver for blocked problems:

35

Sorry, I meant this does work! I've used the solver from that PR. It could be nice to include this problem as a demo in FEniCSx too.

We can aim for a demo on the scifem web-pages if you are interested.

[jorgensd]

This should work on the main branch of DOLFINx now.