How to correctly assign values to mixed elements in FEniCS-X?

Consider the following:

import dolfinx
from mpi4py import MPI
import dolfinx.io
import ufl
import numpy as np

mesh = dolfinx.UnitSquareMesh(MPI.COMM_WORLD, 10, 10)
el = ufl.FiniteElement("CG", mesh.ufl_cell(), 1)
mel = ufl.MixedElement([el, el])
V = dolfinx.FunctionSpace(mesh, mel)
num_subs = V.num_sub_spaces()
spaces = []
maps = []
for i in range(num_subs):
    space_i, map_i = V.sub(i).collapse(collapsed_dofs=True)
    spaces.append(space_i)
    maps.append(map_i)


u = dolfinx.Function(V)
u0 = dolfinx.Function(spaces[0])
u1 = dolfinx.Function(spaces[1])

u0.vector.array[:] = np.arange(len(maps[0]))
u1.vector.array[:] = np.arange(len(maps[1]))
u.vector.array[maps[0]] = u0.vector.array
u.vector.array[maps[1]] = u1.vector.array

with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "splitted_space.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(u.sub(0))
    xdmf.write_function(u.sub(1))
4 Likes