Thanks a lot Francesco, that is exactly what I was looking for!
I put the modified code for people who does not have multiphenics and viskex (and thanks for making me discover viskex, this looks pretty cool, I’ll install it):
"""
Adapted from https://multiphenics.github.io/tutorials/02_navier_stokes/tutorial_navier_stokes.html
"""
import typing
from mpi4py import MPI
import numpy as np
import gmsh
from petsc4py.PETSc import ScalarType, InsertMode, ScatterMode, KSP, Vec, SNES, Mat
import dolfinx
import dolfinx.fem.petsc as fem_petsc
import ufl
import basix.ufl
pre_step_length = 4.
after_step_length = 14.
pre_step_height = 3.
after_step_height = 5.
mesh_size = 1. / 5.
nu = 0.01
def create_mesh():
gmsh.initialize()
gmsh.model.add("mesh")
p0 = gmsh.model.geo.addPoint(0.0, after_step_height - pre_step_height, 0.0, mesh_size)
p1 = gmsh.model.geo.addPoint(pre_step_length, after_step_height - pre_step_height, 0.0, mesh_size)
p2 = gmsh.model.geo.addPoint(pre_step_length, 0.0, 0.0, mesh_size)
p3 = gmsh.model.geo.addPoint(pre_step_length + after_step_length, 0.0, 0.0, mesh_size)
p4 = gmsh.model.geo.addPoint(pre_step_length + after_step_length, after_step_height, 0.0, mesh_size)
p5 = gmsh.model.geo.addPoint(0.0, after_step_height, 0.0, mesh_size)
l0 = gmsh.model.geo.addLine(p0, p1)
l1 = gmsh.model.geo.addLine(p1, p2)
l2 = gmsh.model.geo.addLine(p2, p3)
l3 = gmsh.model.geo.addLine(p3, p4)
l4 = gmsh.model.geo.addLine(p4, p5)
l5 = gmsh.model.geo.addLine(p5, p0)
line_loop = gmsh.model.geo.addCurveLoop([l0, l1, l2, l3, l4, l5])
domain = gmsh.model.geo.addPlaneSurface([line_loop])
gmsh.model.geo.synchronize()
gmsh.model.addPhysicalGroup(1, [l5], 1)
gmsh.model.addPhysicalGroup(1, [l0, l1, l2, l4], 2)
gmsh.model.addPhysicalGroup(2, [domain], 0)
gmsh.model.mesh.generate(2)
mesh, subdomains, boundaries = dolfinx.io.gmshio.model_to_mesh(
gmsh.model, comm=MPI.COMM_WORLD, rank=0, gdim=2)
gmsh.finalize()
return mesh, subdomains, boundaries
mesh, subdomains, boundaries = create_mesh()
# Create connectivities required by the rest of the code
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
boundaries_1 = boundaries.indices[boundaries.values == 1]
boundaries_2 = boundaries.indices[boundaries.values == 2]
def u_in_eval(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[ # type: ignore[no-any-unimported]
ScalarType]:
"""Return the flat velocity profile at the inlet."""
values = np.zeros((2, x.shape[1]))
values[0, :] = 1.0
return values
def u_wall_eval(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[ # type: ignore[no-any-unimported]
ScalarType]:
"""Return the zero velocity at the wall."""
return np.zeros((2, x.shape[1]))
V_element = basix.ufl.element("Lagrange", mesh.basix_cell(), 2, shape=(mesh.geometry.dim, ))
Q_element = basix.ufl.element("Lagrange", mesh.basix_cell(), 1)
def run_monolithic() -> dolfinx.fem.Function:
"""Run standard formulation using a mixed function space."""
# Function spaces
W_element = basix.ufl.mixed_element([V_element, Q_element])
W = dolfinx.fem.functionspace(mesh, W_element)
# Test and trial functions: monolithic
vq = ufl.TestFunction(W)
(v, q) = ufl.split(vq)
dup = ufl.TrialFunction(W)
up = dolfinx.fem.Function(W)
(u, p) = ufl.split(up)
# Variational forms
F = (nu * ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
+ ufl.inner(ufl.grad(u) * u, v) * ufl.dx
- ufl.inner(p, ufl.div(v)) * ufl.dx
+ ufl.inner(ufl.div(u), q) * ufl.dx)
J = ufl.derivative(F, up, dup)
# Boundary conditions
W0 = W.sub(0)
V, _ = W0.collapse()
u_in = dolfinx.fem.Function(V)
u_in.interpolate(u_in_eval)
u_wall = dolfinx.fem.Function(V)
u_wall.interpolate(u_wall_eval)
bdofs_V_1 = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundaries_1)
bdofs_V_2 = dolfinx.fem.locate_dofs_topological((W0, V), mesh.topology.dim - 1, boundaries_2)
inlet_bc = dolfinx.fem.dirichletbc(u_in, bdofs_V_1, W0)
wall_bc = dolfinx.fem.dirichletbc(u_wall, bdofs_V_2, W0)
bc = [inlet_bc, wall_bc]
# Class for interfacing with SNES
class NavierStokesProblem:
"""Define a nonlinear problem, interfacing with SNES."""
def __init__( # type: ignore[no-any-unimported]
self, F: ufl.Form, J: ufl.Form, solution: dolfinx.fem.Function,
bcs: list[dolfinx.fem.DirichletBC], P: typing.Optional[ufl.Form] = None
) -> None:
self._F = dolfinx.fem.form(F)
self._J = dolfinx.fem.form(J)
self._obj_vec = fem_petsc.create_vector(self._F)
self._solution = solution
self._bcs = bcs
self._P = P
def create_snes_solution(self) -> Vec: # type: ignore[no-any-unimported]
"""
Create a petsc4py.PETSc.Vec to be passed to petsc4py.PETSc.SNES.solve.
The returned vector will be initialized with the initial guess provided in `self._solution`.
"""
x = self._solution.vector.copy()
with x.localForm() as _x, self._solution.vector.localForm() as _solution:
_x[:] = _solution
return x
def update_solution(self, x: Vec) -> None: # type: ignore[no-any-unimported]
"""Update `self._solution` with data in `x`."""
x.ghostUpdate(addv=InsertMode.INSERT, mode=ScatterMode.FORWARD)
with x.localForm() as _x, self._solution.vector.localForm() as _solution:
_solution[:] = _x
def obj( # type: ignore[no-any-unimported]
self, snes: SNES, x: Vec
) -> np.float64:
"""Compute the norm of the residual."""
self.F(snes, x, self._obj_vec)
return self._obj_vec.norm() # type: ignore[no-any-return]
def F( # type: ignore[no-any-unimported]
self, snes: SNES, x: Vec, F_vec: Vec
) -> None:
"""Assemble the residual."""
self.update_solution(x)
with F_vec.localForm() as F_vec_local:
F_vec_local.set(0.0)
dolfinx.fem.petsc.assemble_vector(F_vec, self._F)
dolfinx.fem.apply_lifting(F_vec, [self._J], [self._bcs], x0=[x], scale=-1.0)
F_vec.ghostUpdate(addv=InsertMode.ADD, mode=ScatterMode.REVERSE)
dolfinx.fem.set_bc(F_vec, self._bcs, x, -1.0)
def J( # type: ignore[no-any-unimported]
self, snes: SNES, x: Vec, J_mat: Mat,
P_mat: Mat
) -> None:
"""Assemble the jacobian."""
J_mat.zeroEntries()
dolfinx.fem.petsc.assemble_matrix( # type: ignore[misc]
J_mat, self._J, self._bcs, diagonal=1.0) # type: ignore[arg-type]
J_mat.assemble()
if self._P is not None:
P_mat.zeroEntries()
dolfinx.fem.petsc.assemble_matrix( # type: ignore[misc]
P_mat, self._P, self._bcs, diagonal=1.0) # type: ignore[arg-type]
P_mat.assemble()
# Create problem
problem = NavierStokesProblem(F, J, up, bc)
F_vec = dolfinx.fem.petsc.create_vector(problem._F)
J_mat = dolfinx.fem.petsc.create_matrix(problem._J)
# Solve
snes = SNES().create(mesh.comm)
snes.setTolerances(max_it=20)
snes.getKSP().setType("preonly")
snes.getKSP().getPC().setType("lu")
snes.getKSP().getPC().setFactorSolverType("mumps")
snes.setObjective(problem.obj)
snes.setFunction(problem.F, F_vec)
snes.setJacobian(problem.J, J=J_mat, P=None)
snes.setMonitor(lambda _, it, residual: print(it, residual))
up_copy = problem.create_snes_solution()
snes.solve(None, up_copy)
problem.update_solution(up_copy) # TODO can this be safely removed?
up_copy.destroy()
snes.destroy()
return up
up_m = run_monolithic()
(u_m, p_m) = (up_m.sub(0).collapse(), up_m.sub(1).collapse())
from dolfinx import fem, io
import ufl
Wplot = fem.FunctionSpace(mesh, ufl.VectorElement("CG", mesh.ufl_cell(), 1 + 1, mesh.geometry.dim))
uv = fem.Function(Wplot)
uv.name = "uv"
uv.interpolate(u_m)
uv.x.scatter_forward()
with io.VTXWriter(mesh.comm, "solution_u.bp", [uv], engine="BP4") as vtx:
vtx.write(0.0)