Mixed formulation of stokes equations

I am working towards developing a monolithic stokes equation solver using mixed pressure velocity formulation and Taylor Hood elements. Having reviewed all the different posts, examples, and forums, I am able to get my implementation to run but get wrong results. Here is my attempt at replicating an MWE for flow through a long channel (aiming to replicate Poiseuille flow). At this point, I will highly benefit from someone to look at this and let me know why I am getting wrong results.

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx as dfx
from dolfinx.fem.petsc import LinearProblem
from dolfinx.io import gmshio
import ufl as ufl
import numpy as np

def inflow(x):
    return np.stack((np.full(x.shape[1], 1.0), np.zeros(x.shape[1])))

def wall(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

meshFile = 'test.msh'

ID_L = 5
ID_T = 6
ID_B = 8

mesh, c_id, f_id = gmshio.read_from_msh(meshFile, MPI.COMM_WORLD, gdim=2)

tdim = mesh.topology.dim
fdim = tdim - 1

V = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
M = ufl.MixedElement([V,P])
W = dfx.fem.FunctionSpace(mesh, M)

V_sub, _ = W.sub(0).collapse()
P_sub, _ = W.sub(1).collapse()

b_dofs_L = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, f_id.find(ID_L))
b_dofs_T = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, f_id.find(ID_T))
b_dofs_B = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, f_id.find(ID_B))

u_in = dfx.fem.Function(V_sub)
u_in.interpolate(inflow)

u_wall_1 = dfx.fem.Function(V_sub)
u_wall_1.interpolate(wall)

u_wall_2 = dfx.fem.Function(V_sub)
u_wall_2.interpolate(wall)

bc_1 = dfx.fem.dirichletbc(u_in, b_dofs_L, V_sub)
bc_2 = dfx.fem.dirichletbc(u_wall_1, b_dofs_T, V_sub)
bc_3 = dfx.fem.dirichletbc(u_wall_2, b_dofs_B, V_sub)

bc = [bc_1, bc_2, bc_3]

(us,ps) = ufl.TrialFunctions(W)
(vt,qt) = ufl.TestFunctions(W)

mu = dfx.fem.Constant(mesh, dfx.default_scalar_type(1.0))
b = dfx.fem.Constant(mesh, PETSc.ScalarType((0.0,0.0)))

weakForm = mu * ufl.inner(ufl.grad(vt), ufl.grad(us)) * ufl.dx \
    - ps * ufl.div(vt) * ufl.dx \
    + qt * ufl.div(us) * ufl.dx \
    - ufl.dot(vt,b) * ufl.dx

a = ufl.lhs(weakForm)
L = ufl.rhs(weakForm)

uh = dfx.fem.Function(W)
problem = LinearProblem(a, L, bcs=bc, u=uh)
problem.solve()

vtkFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./mwe.pvd", "w")
vtkFile.write_function([uh.split()[0].collapse()])
vtkFile.close()

Happy to share more if needed.

1 Like

Hi,

First of all, as you haven’t disclosed the mesh, the code isn’t reproducible. Couldn’t you just use dolfinx.mesh.create_rectangle for now?

You are also not setting any solver options.

Note that you could also compare with: Bad convergence of stokes problem
that solves a similar problem

How do I attach a mesh?
Also: My issue is not with convergence. I can attach image of the solution as read into Paraview, and it is giving me incorrect solution.

Please try to use dolfinx.mesh.create_rectangle first, as it makes it easy for us to reproduce the results.
Is this the pressure or velocity that you are looking at?

I was looking at the velocity solution (that is at least per my understanding of what the code is doing).

Here I have recreated the above MWE using create_rectangle built in functionality.

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx as dfx
from dolfinx.fem.petsc import LinearProblem
from dolfinx.io import gmshio
import ufl as ufl
import numpy as np

def inflow(x):
    return np.stack((np.full(x.shape[1], 1.0), np.zeros(x.shape[1])))

def noslip(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

def inlet(x):
    return np.isclose(x[0], 0.0)

def walls(x):
    return np.logical_or(np.isclose(x[1], 0.0), np.isclose(x[1], 1.0))

def outlet(x):
    return np.isclose(x[0], 5.0)

mesh = dfx.mesh.create_rectangle(MPI.COMM_WORLD, \
    [np.array([0.0, 0.0]), np.array([5, 1])], [50, 10], dfx.mesh.CellType.triangle)

tdim = mesh.topology.dim
fdim = tdim - 1

V = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
M = ufl.MixedElement([V,P])
W = dfx.fem.FunctionSpace(mesh, M)

V_sub, _ = W.sub(0).collapse()
P_sub, _ = W.sub(1).collapse()

ID_Wall = dfx.mesh.locate_entities_boundary(mesh, fdim, walls)
b_dofs_Wall = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_Wall)
u_wall = dfx.fem.Function(V_sub)
u_wall.interpolate(noslip)
bc_wall = dfx.fem.dirichletbc(u_wall, b_dofs_Wall, V_sub)

ID_Inlet = dfx.mesh.locate_entities_boundary(mesh, fdim, inlet)
b_dofs_Inlet = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_Inlet)
u_in = dfx.fem.Function(V_sub)
u_in.interpolate(inflow)
bc_inlet = dfx.fem.dirichletbc(u_in, b_dofs_Inlet, V_sub)

bc = [bc_inlet, bc_wall]

(us,ps) = ufl.TrialFunctions(W)
(vt,qt) = ufl.TestFunctions(W)

mu = dfx.fem.Constant(mesh, dfx.default_scalar_type(1.0))
b = dfx.fem.Constant(mesh, PETSc.ScalarType((0.0,0.0)))

weakForm = mu * ufl.inner(ufl.grad(vt), ufl.grad(us)) * ufl.dx \
    - ps * ufl.div(vt) * ufl.dx \
    + qt * ufl.div(us) * ufl.dx \
    - ufl.dot(vt,b) * ufl.dx

a = ufl.lhs(weakForm)
L = ufl.rhs(weakForm)

uh = dfx.fem.Function(W)
problem = LinearProblem(a, L, bcs=bc, u=uh)
problem.solve()

vtkFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./mwe.pvd", "w")
vtkFile.write_function([uh.split()[0].collapse()])
vtkFile.close()

I still get incorrect results with this for expected channel flow - here is a screenshot:

The issue is the boundary conditions.
The last argument of dirichletbc should not be V_sub, but W.sub(0) to indicate to the assemblers that it is the space we want to apply u_in on. I also changed the sign of your divergence term for symmetry.
MWE:

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx as dfx
from dolfinx.fem.petsc import LinearProblem
import ufl as ufl
import numpy as np

def inflow(x):
    return np.stack((np.full(x.shape[1], 1.0), np.zeros(x.shape[1])))

def noslip(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

def inlet(x):
    return np.isclose(x[0], 0.0)

def walls(x):
    return np.logical_or(np.isclose(x[1], 0.0), np.isclose(x[1], 1.0))

def outlet(x):
    return np.isclose(x[0], 5.0)

mesh = dfx.mesh.create_rectangle(MPI.COMM_WORLD, \
    [np.array([0.0, 0.0]), np.array([5, 1])], [50, 10], dfx.mesh.CellType.triangle)

tdim = mesh.topology.dim
fdim = tdim - 1

V = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
M = ufl.MixedElement([V,P])
W = dfx.fem.FunctionSpace(mesh, M)

V_sub, _ = W.sub(0).collapse()
P_sub, _ = W.sub(1).collapse()

ID_Wall = dfx.mesh.locate_entities_boundary(mesh, fdim, walls)
b_dofs_Wall = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_Wall)
u_wall = dfx.fem.Function(V_sub)
u_wall.interpolate(noslip)
bc_wall = dfx.fem.dirichletbc(u_wall, b_dofs_Wall, W.sub(0))

ID_Inlet = dfx.mesh.locate_entities_boundary(mesh, fdim, inlet)
b_dofs_Inlet = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_Inlet)
u_in = dfx.fem.Function(V_sub)
u_in.interpolate(inflow)
bc_inlet = dfx.fem.dirichletbc(u_in, b_dofs_Inlet, W.sub(0))

bc = [bc_inlet, bc_wall]

(us,ps) = ufl.TrialFunctions(W)
(vt,qt) = ufl.TestFunctions(W)

mu = dfx.fem.Constant(mesh, dfx.default_scalar_type(1.0))
b = dfx.fem.Constant(mesh, PETSc.ScalarType((0.0,0.0)))

weakForm = mu * ufl.inner(ufl.grad(vt), ufl.grad(us)) * ufl.dx \
    + ps * ufl.div(vt) * ufl.dx \
    + qt * ufl.div(us) * ufl.dx \
    - ufl.dot(vt,b) * ufl.dx

a = ufl.lhs(weakForm)
L = ufl.rhs(weakForm)

uh = dfx.fem.Function(W)
problem = LinearProblem(a, L, bcs=bc, u=uh)
problem.solve()

vtkFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./mwe.pvd", "w")
vtkFile.write_function([uh.split()[0].collapse()])
vtkFile.close()
1 Like

Thanks for the help. With your pointer above, we have been able to get several of our workflows corrected into Dolfinx format from legacy Fenics. However, some issues still remain. I share with you an example as below. This is a problem that specifies a body force moving a fluid within a box with 4 no slip walls. We know the analytical solutions for velocity and pressure; and are comparing the computed solution against this. I have listed an MWE version of this below. We have extensively used this example simulation before, and we know that with Legacy Fenics this worked out perfectly. However, with the Dolfinx version - the solution seems to vary wildly with mesh. If you use this MWE, you can see that for varying meshes the solutions will either match or depart from analytical. Why would such mesh dependence arise in the Dolfinx implementation?

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx as dfx
from dolfinx.fem.petsc import LinearProblem
from dolfinx.io import gmshio
import ufl as ufl
import numpy as np
import sys

def noSlipBC(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

def pressureBC(x):
    return np.zeros(x.shape[1])

def bodyForce(x):

    vals = np.zeros((mesh.geometry.dim, x.shape[1]))

    vals[0] = (12.0 - 24.0*x[1])*(x[0]*x[0]*x[0]*x[0]) + \
            (-24.0 + 48.0*x[1])*(x[0]*x[0]*x[0]) + \
            (-48.0*x[1] + 72.0*x[1]*x[1] - 48.0*x[1]*x[1]*x[1] + 12.0)*(x[0]*x[0]) + \
            (-2.0 + 24.0*x[1] - 72.0*x[1]*x[1] + 48.0*x[1]*x[1]*x[1])*x[0] + \
            1.0 - 4.0*x[1] + 12.0*x[1]*x[1] - 8.0*x[1]*x[1]*x[1]

    vals[1] = (8.0 - 48.0*x[1] + 48.0*x[1]*x[1])*(x[0]*x[0]*x[0]) + \
            (-12.0 + 72.0*x[1] - 72.0*x[1]*x[1])*(x[0]*x[0]) + \
            (4.0 - 24.0*x[1] + 48.0*x[1]*x[1] - 48.0*x[1]*x[1]*x[1] + 24.0*x[1]*x[1]*x[1]*x[1])*x[0] \
            - 12.0*x[1]*x[1] + 24.0*x[1]*x[1]*x[1] - 12.0*x[1]*x[1]*x[1]*x[1]

    return vals

def analyticalVelocity(x):

    vals = np.zeros((mesh.geometry.dim, x.shape[1]))

    vals[0] = (x[0]**2) * (1.0 - x[0])**2 * (2.0 * x[1] - 6.0 * x[1]**2 + 4.0 * x[1]**3)
    vals[1] = -(x[1]**2) * (1.0 - x[1])**2 * (2.0 * x[0] - 6.0 * x[0]**2 + 4.0 * x[0]**3)

    return vals

def analyticalPressure(x):
    return x[0] * (1.0 - x[0])

def left(x):
    return np.isclose(x[0], 0.0)

def right(x):
    return np.isclose(x[0], 1.0)

def top(x):
    return np.isclose(x[1], 1.0)

def bottom(x):
    return np.isclose(x[1], 0.0)

viscosity = 1.0
velName = './vv-sol-v'
presName = './vv-sol-p'
Nx = 80
Ny = 80

mesh = dfx.mesh.create_rectangle(MPI.COMM_WORLD, \
    [np.array([0.0, 0.0]), np.array([1, 1])], [Nx, Ny], dfx.mesh.CellType.triangle)

tdim = mesh.topology.dim
fdim = tdim - 1

V = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
M = ufl.MixedElement([V,P])
W = dfx.fem.FunctionSpace(mesh, M)

V_sub, V_submap = W.sub(0).collapse()
P_sub, P_submap = W.sub(1).collapse()

ID_L = dfx.mesh.locate_entities_boundary(mesh, fdim, left)
ID_R = dfx.mesh.locate_entities_boundary(mesh, fdim, right)
ID_T = dfx.mesh.locate_entities_boundary(mesh, fdim, top)
ID_B = dfx.mesh.locate_entities_boundary(mesh, fdim, bottom)

b_dofs_L = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_L)
b_dofs_T = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_T)
b_dofs_B = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_B)
b_dofs_R = dfx.fem.locate_dofs_topological((W.sub(0),V_sub), fdim, ID_R)
b_dofs_P = dfx.fem.locate_dofs_topological((W.sub(1),P_sub), fdim, ID_R)

uD_Wall = dfx.fem.Function(V_sub)
uD_Wall.interpolate(noSlipBC)

uD_P = dfx.fem.Function(P_sub)
uD_P.interpolate(pressureBC)

bc_L = dfx.fem.dirichletbc(uD_Wall, b_dofs_L, W.sub(0))
bc_T = dfx.fem.dirichletbc(uD_Wall, b_dofs_T, W.sub(0))
bc_B = dfx.fem.dirichletbc(uD_Wall, b_dofs_B, W.sub(0))
bc_R = dfx.fem.dirichletbc(uD_Wall, b_dofs_R, W.sub(0))
bc_P = dfx.fem.dirichletbc(uD_P, b_dofs_P, W.sub(1))

bc = [bc_L, bc_T, bc_B, bc_R, bc_P]

(us,ps) = ufl.TrialFunctions(W)
(vt,qt) = ufl.TestFunctions(W)

mu = dfx.fem.Constant(mesh, dfx.default_scalar_type(viscosity))

bF = dfx.fem.VectorFunctionSpace(mesh, ("Lagrange", 4))
b = dfx.fem.Function(bF)
b.interpolate(bodyForce)

weakForm = mu * ufl.inner(ufl.grad(vt), ufl.grad(us)) * ufl.dx \
    - ps * ufl.div(vt) * ufl.dx - qt * ufl.div(us) * ufl.dx - ufl.dot(vt,b) * ufl.dx

a = ufl.lhs(weakForm)
L = ufl.rhs(weakForm)

uh = dfx.fem.Function(W)
problem = LinearProblem(a, L, bcs=bc, u=uh)
problem.solve()

us = uh.split()[0].collapse()
ps = uh.split()[1].collapse()

us.name='vel'
ps.name='pres'

vFile = dfx.io.VTKFile(MPI.COMM_WORLD, velName+'.pvd', "w")
vFile.write_function(us, 0.)
vFile.close()

pFile = dfx.io.VTKFile(MPI.COMM_WORLD, presName+'.pvd', "w")
pFile.write_function(ps, 0.)
pFile.close()

UF = dfx.fem.VectorFunctionSpace(mesh, ("Lagrange", 2))
u0Func = dfx.fem.Function(UF)
u0Func.interpolate(analyticalVelocity)
u0Func.name = "vel"

PF = dfx.fem.FunctionSpace(mesh, ("Lagrange", 1))
p0Func = dfx.fem.Function(PF)
p0Func.interpolate(analyticalPressure)
p0Func.name = "pres"

vFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./vv-ex-v.pvd", "w")
vFile.write_function(u0Func, 0.)
vFile.close()

pFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./vv-ex-p.pvd", "w")
pFile.write_function(p0Func, 0.)
pFile.close()

You are not specifying a solver for the linear problem. I would suggest you start by testing it with a direct solver, and ensure that this converges.
See for instance: Conspicuous speed ups in parallel computing - #2 by dokken

Is the default solver not a direct solver? Also: for the Legacy Fenics version, for this simple a problem, we never specified a solver - and the default solver calls worked totally fine. What seems to be the difference in Dolfinx?

For a quick follow up - I included the following lines of code after the solve() function call in the MWE above

p_solver = problem.solver
viewer = PETSc.Viewer().createASCII("solver_output.txt")
p_solver.view(viewer)
solver_output = open("solver_output.txt", "r")
for line in solver_output.readlines():
    print(line)

This is based on documentation provided on Dolfinx site. I realize that this default call is setting gmres with ilu preconditioning (assuming I understand the readout correctly). With this solver - for Nx = Ny = 20 the solution matches analytical, but Nx = Ny = 80 (for instance) does not. Hope this is helpful information.

Legacy dolfin does a lot of magic under the hood to determine a suitable solver.

DOLFINx relies on the defaults provided by your petsc installation, thus I would start by setting

petsc_options={"ksp_type": "preonly", "pc_type": "lu",
"pc_factor_mat_solver_type": "mumps"}

as done in the aforementioned post

Alright, the above fix with the petsc_options works but only if we use pc_factor_mat_solver_type: umfpack. If we specify this to be mumps then the code crashes with a segmentation fault. I am not sure why this is the case.

I cannot reproduce any segfault with mumps with dolfinx stable (v0.8.0, updated code below):

from mpi4py import MPI
import dolfinx as dfx
from dolfinx.fem.petsc import LinearProblem
import ufl as ufl
import numpy as np
import basix.ufl


def noSlipBC(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))


def pressureBC(x):
    return np.zeros(x.shape[1])


def bodyForce(x):
    vals = np.zeros((mesh.geometry.dim, x.shape[1]))

    vals[0] = (
        (12.0 - 24.0 * x[1]) * (x[0] * x[0] * x[0] * x[0])
        + (-24.0 + 48.0 * x[1]) * (x[0] * x[0] * x[0])
        + (-48.0 * x[1] + 72.0 * x[1] * x[1] - 48.0 * x[1] * x[1] * x[1] + 12.0)
        * (x[0] * x[0])
        + (-2.0 + 24.0 * x[1] - 72.0 * x[1] * x[1] + 48.0 * x[1] * x[1] * x[1]) * x[0]
        + 1.0
        - 4.0 * x[1]
        + 12.0 * x[1] * x[1]
        - 8.0 * x[1] * x[1] * x[1]
    )

    vals[1] = (
        (8.0 - 48.0 * x[1] + 48.0 * x[1] * x[1]) * (x[0] * x[0] * x[0])
        + (-12.0 + 72.0 * x[1] - 72.0 * x[1] * x[1]) * (x[0] * x[0])
        + (
            4.0
            - 24.0 * x[1]
            + 48.0 * x[1] * x[1]
            - 48.0 * x[1] * x[1] * x[1]
            + 24.0 * x[1] * x[1] * x[1] * x[1]
        )
        * x[0]
        - 12.0 * x[1] * x[1]
        + 24.0 * x[1] * x[1] * x[1]
        - 12.0 * x[1] * x[1] * x[1] * x[1]
    )

    return vals


def analyticalVelocity(x):
    vals = np.zeros((mesh.geometry.dim, x.shape[1]))

    vals[0] = (
        (x[0] ** 2)
        * (1.0 - x[0]) ** 2
        * (2.0 * x[1] - 6.0 * x[1] ** 2 + 4.0 * x[1] ** 3)
    )
    vals[1] = (
        -(x[1] ** 2)
        * (1.0 - x[1]) ** 2
        * (2.0 * x[0] - 6.0 * x[0] ** 2 + 4.0 * x[0] ** 3)
    )

    return vals


def analyticalPressure(x):
    return x[0] * (1.0 - x[0])


def left(x):
    return np.isclose(x[0], 0.0)


def right(x):
    return np.isclose(x[0], 1.0)


def top(x):
    return np.isclose(x[1], 1.0)


def bottom(x):
    return np.isclose(x[1], 0.0)


viscosity = 1.0
velName = "./vv-sol-v"
presName = "./vv-sol-p"
Nx = 80
Ny = 80

mesh = dfx.mesh.create_rectangle(
    MPI.COMM_WORLD,
    [np.array([0.0, 0.0]), np.array([1, 1])],
    [Nx, Ny],
    dfx.mesh.CellType.triangle,
)

tdim = mesh.topology.dim
fdim = tdim - 1

V = basix.ufl.element("Lagrange", mesh.basix_cell(), 2, shape=(mesh.geometry.dim,))
P = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
M = basix.ufl.mixed_element([V, P])
W = dfx.fem.functionspace(mesh, M)

V_sub, V_submap = W.sub(0).collapse()
P_sub, P_submap = W.sub(1).collapse()

ID_L = dfx.mesh.locate_entities_boundary(mesh, fdim, left)
ID_R = dfx.mesh.locate_entities_boundary(mesh, fdim, right)
ID_T = dfx.mesh.locate_entities_boundary(mesh, fdim, top)
ID_B = dfx.mesh.locate_entities_boundary(mesh, fdim, bottom)

b_dofs_L = dfx.fem.locate_dofs_topological((W.sub(0), V_sub), fdim, ID_L)
b_dofs_T = dfx.fem.locate_dofs_topological((W.sub(0), V_sub), fdim, ID_T)
b_dofs_B = dfx.fem.locate_dofs_topological((W.sub(0), V_sub), fdim, ID_B)
b_dofs_R = dfx.fem.locate_dofs_topological((W.sub(0), V_sub), fdim, ID_R)
b_dofs_P = dfx.fem.locate_dofs_topological((W.sub(1), P_sub), fdim, ID_R)

uD_Wall = dfx.fem.Function(V_sub)
uD_Wall.interpolate(noSlipBC)

uD_P = dfx.fem.Function(P_sub)
uD_P.interpolate(pressureBC)

bc_L = dfx.fem.dirichletbc(uD_Wall, b_dofs_L, W.sub(0))
bc_T = dfx.fem.dirichletbc(uD_Wall, b_dofs_T, W.sub(0))
bc_B = dfx.fem.dirichletbc(uD_Wall, b_dofs_B, W.sub(0))
bc_R = dfx.fem.dirichletbc(uD_Wall, b_dofs_R, W.sub(0))
bc_P = dfx.fem.dirichletbc(uD_P, b_dofs_P, W.sub(1))

bc = [bc_L, bc_T, bc_B, bc_R, bc_P]

(us, ps) = ufl.TrialFunctions(W)
(vt, qt) = ufl.TestFunctions(W)

mu = dfx.fem.Constant(mesh, dfx.default_scalar_type(viscosity))

bF = dfx.fem.functionspace(mesh, ("Lagrange", 4, (mesh.geometry.dim,)))
b = dfx.fem.Function(bF)
b.interpolate(bodyForce)

weakForm = (
    mu * ufl.inner(ufl.grad(vt), ufl.grad(us)) * ufl.dx
    - ps * ufl.div(vt) * ufl.dx
    - qt * ufl.div(us) * ufl.dx
    - ufl.dot(vt, b) * ufl.dx
)

a = ufl.lhs(weakForm)
L = ufl.rhs(weakForm)

uh = dfx.fem.Function(W)
problem = LinearProblem(
    a,
    L,
    bcs=bc,
    u=uh,
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
    },
)
problem.solve()
print(problem.solver.getConvergedReason())
us = uh.split()[0].collapse()
ps = uh.split()[1].collapse()

us.name = "vel"
ps.name = "pres"

vFile = dfx.io.VTKFile(MPI.COMM_WORLD, velName + ".pvd", "w")
vFile.write_function(us, 0.0)
vFile.close()

pFile = dfx.io.VTKFile(MPI.COMM_WORLD, presName + ".pvd", "w")
pFile.write_function(ps, 0.0)
pFile.close()

UF = dfx.fem.functionspace(mesh, ("Lagrange", 2, (mesh.geometry.dim,)))
u0Func = dfx.fem.Function(UF)
u0Func.interpolate(analyticalVelocity)
u0Func.name = "vel"

PF = dfx.fem.functionspace(mesh, ("Lagrange", 1))
p0Func = dfx.fem.Function(PF)
p0Func.interpolate(analyticalPressure)
p0Func.name = "pres"

vFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./vv-ex-v.pvd", "w")
vFile.write_function(u0Func, 0.0)
vFile.close()

pFile = dfx.io.VTKFile(MPI.COMM_WORLD, "./vv-ex-p.pvd", "w")
pFile.write_function(p0Func, 0.0)
pFile.close()