Vectorial non homogeneous Dirichlet boundary condition

Good evening,

I am trying to solve the stokes problem by considering the force term and the velocity as a three component vectors.
I consider as boundary condition non homogeneous Dirichlet, which also coincide with the true solution. The problem is that I don’t understand how to impose the vector of non homogeneous Dirichlet condition, I don’t know if it is a problem related to the use of the function dolfinx.fem.Expression.

I attach the code and the dolfinx version I am using is 0.6.0:

######### VARIATIONAL FORM
def create_bilinear_form(V, Q):
    u, p = TrialFunction(V), TrialFunction(Q)
    v, q = TestFunction(V), TestFunction(Q)
    a_uu = inner(grad(u), grad(v)) * dx
    a_up = inner(p, div(v)) * dx
    a_pu = inner(div(u), q) * dx
    return fem.form([[a_uu, a_up], [a_pu, None]])

def create_linear_form(V, Q):
    v, q = TestFunction(V), TestFunction(Q)
    domain = V.mesh
    x = SpatialCoordinate(domain)
    f = ufl.as_vector([1-24*(-1 + x[0])*(x[0]-1)*x[0]*x[0]*(-0.5 + x[1])- 48*(0.166667 - x[0] + x[0]*x[0])*(-1 + x[1])*(-0.5 + x[1])*x[1],
                    1+24*(-0.5 + x[0])*(-1 + x[1])*(x[1]-1)*x[1]*x[1] + 8*(-1 + x[0])*(-0.5 + x[0])*x[0]*(1 - 6*x[1] + 6*x[1]*x[1]),
                    1.0])
    return fem.form([inner(f, v) * dx,
                     inner(fem.Constant(domain, 0.), q) * dx])             

########### BOUNDARY CONDITIONS
def create_velocity_bc(V):
    domain = V.mesh
    x = SpatialCoordinate(domain)
    uD_vec = ufl.as_vector([4*(x[1]-1)*x[1]*(x[0]-1)*(x[0]-1)*x[0]*x[0]*(x[1]-0.5), -4*(x[0]-0.5)*x[1]*x[1]*x[0]*(x[1]-1)*(x[1]-1)*(x[0]-1), 0.0])
    '''uD = Function(V)
                uD.interpolate(u_ex)'''
    uD = Expression(uD_vec, V.element.interpolation_points())
    import numpy
    # Create facet to cell connectivity required to determine boundary facets
    tdim = domain.topology.dim
    fdim = tdim - 1
    domain.topology.create_connectivity(fdim, tdim)
    boundary_facets = mesh.exterior_facet_indices(domain.topology)
    boundary_dofs = fem.locate_dofs_topological(V, fdim, boundary_facets)
    bc = fem.dirichletbc(uD, boundary_dofs,V)

    return bc   

def create_nullspace(rhs_form):
    null_vec = fem.petsc.create_vector_nest(rhs_form)
    null_vecs = null_vec.getNestSubVecs()
    null_vecs[0].set(0.0)
    null_vecs[1].set(1.0)
    null_vec.normalize()
    nsp = PETSc.NullSpace().create(vectors=[null_vec])
    return nsp

########### CREATE BLOCK PRECONDITIONER
def create_preconditioner(Q, a, bcs):
    p, q = TrialFunction(Q), TestFunction(Q)
    a_p11 = fem.form(inner(p, q) * dx)
    a_p = fem.form([[a[0][0], None],
                    [None, a_p11]])
    P = fem.petsc.assemble_matrix_nest(a_p, bcs)
    P.assemble()
    return P

######### ASSEMBLE THE NESTED SYSTEM
def assemble_system(lhs_form, rhs_form, bcs):
    A = fem.petsc.assemble_matrix_nest(lhs_form, bcs=bcs)
    A.assemble()

    b = fem.petsc.assemble_vector_nest(rhs_form)
    # https://docs.fenicsproject.org/dolfinx/main/python/generated/dolfinx.fem.html#dolfinx.fem.apply_lifting
	# The right-hand side vector is assembled and then modified to account for non-homogeneous Dirichlet boundary conditions
	# https://jsdokken.com/dolfinx-tutorial/chapter1/nitsche.html
    fem.petsc.apply_lifting_nest(b, lhs_form, bcs=bcs)
    for b_sub in b.getNestSubVecs():
        b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD,
                          mode=PETSc.ScatterMode.REVERSE)
    spaces = fem.extract_function_spaces(rhs_form)
    bcs0 = fem.bcs_by_block(spaces, bcs)
    fem.petsc.set_bc_nest(b, bcs0)
    return A, b

######### KRILOV SOLVER
def create_block_solver(A, b, P, comm):
    ksp = PETSc.KSP().create(comm)
    ksp.setOperators(A, P)
    ksp.setType("minres")
    ksp.setTolerances(rtol=1e-9)
    ksp.getPC().setType("fieldsplit")
    ksp.getPC().setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)

    nested_IS = P.getNestISs()
    ksp.getPC().setFieldSplitIS(("u", nested_IS[0][0]),
                                ("p", nested_IS[0][1]))

    # Set the preconditioners for each block
    ksp_u, ksp_p = ksp.getPC().getFieldSplitSubKSP()
    ksp_u.setType("preonly")
    ksp_u.getPC().setType("gamg")
    ksp_p.setType("preonly")
    ksp_p.getPC().setType("jacobi")

    # Monitor the convergence of the KSP
    ksp.setFromOptions()
    return ksp
######### COMPUTE ERROR
def assemble_scalar(J, comm: MPI.Comm):
    scalar_form = fem.form(J)
    local_J = fem.assemble_scalar(scalar_form)
    return comm.allreduce(local_J, op=MPI.SUM)

def compute_errors(u, p):
    domain = u.function_space.mesh
    x = SpatialCoordinate(domain)
    error_u = u - u_ex(x)
    H1_u = inner(error_u, error_u) * dx
    H1_u += inner(grad(error_u), grad(error_u)) * dx
    velocity_error = np.sqrt(assemble_scalar(H1_u, domain.comm))

    error_p = -p - p_ex(x)
    L2_p = fem.form(error_p * error_p * dx)
    pressure_error = np.sqrt(assemble_scalar(L2_p, domain.comm))
    return velocity_error, pressure_error
######### SOLVE THE SYSTEM
def solve_stokes(u_element, p_element, domain):
    V = fem.FunctionSpace(domain, u_element)
    Q = fem.FunctionSpace(domain, p_element)

    lhs_form = create_bilinear_form(V, Q)
    rhs_form = create_linear_form(V, Q)

    bcs = create_velocity_bc(V)
    nsp = create_nullspace(rhs_form)
    A, b = assemble_system(lhs_form, rhs_form, bcs)
    assert nsp.test(A)
    A.setNullSpace(nsp)

    P = create_preconditioner(Q, lhs_form, bcs)
    ksp = create_block_solver(A, b, P, domain.comm)

    u, p = fem.Function(V), fem.Function(Q)
    w = PETSc.Vec().createNest([u.vector, p.vector])
    ksp.solve(b, w)
    assert ksp.getConvergedReason() > 0
    u.x.scatter_forward()
    p.x.scatter_forward()

    import pyvista
    import matplotlib.pyplot as plt
    from dolfinx import plot
    pyvista.start_xvfb()
    plotter = pyvista.Plotter()
    plotter.open_gif("deformation.gif", fps=3)

    topology, cells, geometry = plot.create_vtk_mesh(u.function_space)
    function_grid = pyvista.UnstructuredGrid(topology, cells, geometry)

    values = np.zeros((geometry.shape[0], 3))
    values[:, :len(u)] = u.x.array.reshape(geometry.shape[0], len(u))
    function_grid["u"] = values
    function_grid.set_active_vectors("u")

    # Warp mesh by deformation
    warped = function_grid.warp_by_vector("u", factor=1)
    warped.set_active_vectors("u")

    # Add mesh to plotter and visualize
    actor = plotter.add_mesh(warped, show_edges=True, lighting=False, clim=[0, 10])

    
    return compute_errors(u, p)

N = 20
domain = mesh.create_box(comm=MPI.COMM_WORLD,
                            points=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), n=(N, N, N))
element_u = ufl.VectorElement("Lagrange", domain.ufl_cell(), 2)
element_p = ufl.FiniteElement("Lagrange", domain.ufl_cell(), 1)

Can someone help me in solving this issue?
Thanks in advance

Have you read: Stokes equations using Taylor-Hood elements — DOLFINx 0.7.0.0 documentation