MixedElement formulation leads to unintuitive behaviour of the API

dolfinx
4

I meant to link to: How to correctly assign values to mixed elements in FEniCS-X? - #3 by dokken
which gives you the following workflow:

from IPython import embed
import numpy as np
from mpi4py import MPI
import ufl
from dolfinx import mesh, fem, io

L, W, H = 1, 1, 1
mesh_size = 0.2

# 3D box with hexahedron cells
n_elements = [max(int(x / mesh_size), 1) for x in [L, W, H]]
domain = mesh.create_box(
    comm=MPI.COMM_WORLD,
    points=[np.array([0, 0, 0]), np.array([L, W, H])],
    n=n_elements,
    cell_type=mesh.CellType.hexahedron,
)

element_u = ufl.VectorElement(
    family="Lagrange",
    cell=domain.ufl_cell(),
    degree=2,
)
element_lm = ufl.FiniteElement(
    family="Lagrange",
    cell=domain.ufl_cell(),
    degree=1,
)

element = ufl.MixedElement(element_u, element_lm)
V = fem.FunctionSpace(domain, element)

# The trial/test functions get defined on the (mixed) function space
dudlm = ufl.TestFunction(V)
du, dlm = ufl.split(dudlm)

ulm = fem.Function(V)
(u, lm) = ulm.split()
# Fill ulm with dummy data
ulm.x.array[:] = np.arange(len(ulm.x.array))
# Create subspace and map from sub space to parent
V0, V0_to_V = V.sub(0).collapse()


u_prev = fem.Function(V0)
u_prev.x.array[:] = u.x.array[V0_to_V]
u_prev.name = "Collapsed space"
with io.XDMFFile(domain.comm, "u_prev.xdmf", "w") as xdmf:
    xdmf.write_mesh(domain)
    xdmf.write_function(u_prev)

with io.XDMFFile(domain.comm, "u0.xdmf", "w") as xdmf:
    xdmf.write_mesh(domain)
    xdmf.write_function(ulm.sub(0))

In many cases, the users do not need the map from the sub space to the parent space (for instance when defining Dirichlet conditions, as we send in both the collapsed and non-collapsed space), and therefore we do not store is as part of the Function object (as it increases memory footprint).

When you want to write variational forms, you should use ufl.split(ulm) to get each component of the function.
If you want to access data from these functions as a dolfinx.fem.Function object, you should use ulm.split() or ulm.sub(i) where i is the index of the sub space.

Here is a minimal example illustrating the use-cases:

import numpy as np
from mpi4py import MPI
import ufl
from dolfinx import mesh, fem, io, nls, log

L, W, H = 1, 1, 1
mesh_size = 0.2

# 3D box with hexahedron cells
n_elements = [max(int(x / mesh_size), 1) for x in [L, W, H]]
domain = mesh.create_box(
    comm=MPI.COMM_WORLD,
    points=[np.array([0, 0, 0]), np.array([L, W, H])],
    n=n_elements,
    cell_type=mesh.CellType.hexahedron,
)

element_u = ufl.VectorElement(
    family="Lagrange",
    cell=domain.ufl_cell(),
    degree=2,
)
element_lm = ufl.FiniteElement(
    family="Lagrange",
    cell=domain.ufl_cell(),
    degree=1,
)

element = ufl.MixedElement(element_u, element_lm)
V = fem.FunctionSpace(domain, element)

# The trial/test functions get defined on the (mixed) function space
dudlm = ufl.TestFunction(V)
du, dlm = ufl.split(dudlm)

ulm = fem.Function(V)
(u, lm) = ufl.split(ulm)

x = ufl.SpatialCoordinate(domain)
F = ufl.inner(u, du)*ufl.dx + ufl.inner(lm, dlm)*ufl.dx - \
    ufl.inner(x, du)*ufl.dx + ufl.inner(x[2], dlm)*ufl.dx

log.set_log_level(log.LogLevel.INFO)
problem = fem.petsc.NonlinearProblem(F, ulm, bcs=[])
solver = nls.petsc.NewtonSolver(domain.comm, problem)
solver.solve(ulm)


(u, lm) = ulm.split()
with io.XDMFFile(domain.comm, "u.xdmf", "w") as xdmf:
    xdmf.write_mesh(domain)
    xdmf.write_function(u)

with io.XDMFFile(domain.comm, "lm.xdmf", "w") as xdmf:
    xdmf.write_mesh(domain)
    xdmf.write_function(lm)