1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
|
# Copyright (C) 2026 Stein K.F. Stoter
#
# This file is part of DOLFINX_MPC
#
# SPDX-License-Identifier: MIT
from __future__ import annotations
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import pytest
import ufl
from basix.ufl import element
from dolfinx import default_real_type, default_scalar_type, fem, la, mesh
from dolfinx.mesh import CellType, create_unit_cube
import dolfinx_mpc
@pytest.mark.parametrize("cell_type", [CellType.tetrahedron, CellType.hexahedron])
@pytest.mark.parametrize("els", [2, 4, 8])
@pytest.mark.parametrize("order", [2, 3])
def test_stokes_channelflow(cell_type, els, order):
"""Test 3D Stokes channel flow with periodic BCs and nested assembly and Taylor-Hood elements.
The analytical solution for Poiseuille flow between parallel plates is:
u_x = (dp/dx) * (1/(2*mu)) * y * (H - y)
where H is the channel height, dp/dx is the pressure gradient enforced as a body force, mu is viscosity.
"""
# Cube mesh
rtype = default_real_type
domain = create_unit_cube(MPI.COMM_WORLD, els, els, els, cell_type=cell_type, dtype=rtype)
# Function spaces: Taylor-Hood
V = fem.functionspace(domain, element("Lagrange", domain.basix_cell(), order, shape=(3,), dtype=rtype))
Q = fem.functionspace(domain, element("Lagrange", domain.basix_cell(), order - 1, dtype=rtype))
# Wall BCs (top/bottom y={0,1})
wall_facets = mesh.locate_entities_boundary(
domain, 2, lambda x: np.logical_or(np.isclose(x[1], 0), np.isclose(x[1], 1))
)
wall_dofs = fem.locate_dofs_topological(V, 2, wall_facets)
bc = fem.dirichletbc(np.zeros(3, dtype=default_scalar_type), wall_dofs, V)
# Periodicity definition (x and z directions)
def periodic_boundary(x):
return np.logical_or(np.isclose(x[0], 1), np.isclose(x[2], 1))
def periodic_map(x):
out = x.copy()
out[0][np.isclose(x[0], 1).nonzero()] -= 1
out[2][np.isclose(x[2], 1).nonzero()] -= 1
return out
# Create MPCs for velocity and pressure
mpc_u = dolfinx_mpc.MultiPointConstraint(V)
mpc_u.create_periodic_constraint_geometrical(V, periodic_boundary, periodic_map, [bc])
mpc_u.finalize()
mpc_p = dolfinx_mpc.MultiPointConstraint(Q)
mpc_p.create_periodic_constraint_geometrical(Q, periodic_boundary, periodic_map, [])
mpc_p.finalize()
# Stokes weak form
u = ufl.TrialFunction(V)
p = ufl.TrialFunction(Q)
v = ufl.TestFunction(V)
q = ufl.TestFunction(Q)
# Viscosity and forcing
mu = fem.Constant(domain, default_scalar_type(1.0))
dp_dx = fem.Constant(domain, default_scalar_type(1.0)) # Pressure gradient in x-direction
f = ufl.as_vector([dp_dx, 0.0, 0.0])
# Weak statement
a00 = mu * ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
a01 = -ufl.inner(p, ufl.div(v)) * ufl.dx
a10 = -ufl.inner(ufl.div(u), q) * ufl.dx
L0 = ufl.inner(f, v) * ufl.dx
L1 = ufl.ZeroBaseForm((q,))
forms = fem.form([[a00, a01], [a10, None]])
rhs_forms = fem.form([L0, L1])
# Assemble - THIS IS WHERE THE BUG OCCURS with MPI and hexahedral cells
A = dolfinx_mpc.create_matrix_nest(forms, constraints=[mpc_u, mpc_p])
dolfinx_mpc.assemble_matrix_nest(A, forms, constraints=[mpc_u, mpc_p], bcs=[bc])
A.assemble()
b = dolfinx_mpc.create_vector_nest(rhs_forms, constraints=[mpc_u, mpc_p])
dolfinx_mpc.assemble_vector_nest(b, rhs_forms, constraints=[mpc_u, mpc_p])
# Apply lifting for boundary conditions
bcs_block = fem.bcs_by_block(fem.extract_function_spaces(forms, 1), [bc])
dolfinx_mpc.apply_lifting(b, forms, bcs_block, constraint=[mpc_u, mpc_p])
# Ghost update after lifting
for b_sub in b.getNestSubVecs():
b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
# Set BC values in RHS
bcs0 = fem.bcs_by_block(fem.extract_function_spaces(rhs_forms), [bc])
fem.petsc.set_bc(b, bcs0)
# Setup Minres solver
ksp = PETSc.KSP().create(domain.comm)
ksp.setOperators(A)
ksp.setErrorIfNotConverged(True)
ksp.setType("minres")
rtol = 100 * np.finfo(rtype).eps
atol = 2 * np.finfo(rtype).eps
ksp.setTolerances(rtol=rtol, atol=atol)
# Setup preconditioning
pc = ksp.getPC()
pc.setType("fieldsplit")
pc.setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)
nested_IS = A.getNestISs()
pc.setFieldSplitIS(("u", nested_IS[0][0]), ("p", nested_IS[0][1]))
ksp_u, ksp_p = pc.getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.getPC().setType("lu")
ksp_u.getPC().setFactorSolverType("mumps")
ksp_p.setType("preonly")
ksp_p.getPC().setType("jacobi")
# Finalize KSP setup
ksp.setFromOptions()
# Create solution functions
uh = fem.Function(mpc_u.function_space)
ph = fem.Function(mpc_p.function_space)
# Solve the system
x = b.duplicate()
ksp.solve(b, x)
la.petsc._ghost_update(
x,
insert_mode=PETSc.InsertMode.INSERT, # type: ignore
scatter_mode=PETSc.ScatterMode.FORWARD, # type: ignore
) # type: ignore
fem.petsc.assign(x, [uh, ph])
# Apply backsubstitution for MPCs
mpc_u.backsubstitution(uh)
mpc_p.backsubstitution(ph)
uh.x.scatter_forward()
ph.x.scatter_forward()
# Verify solution: Poiseuille flow u_x = (dp/dx) * (1/(2*mu)) * y * (H - y)
# For our parameters: mu=1, dp/dx=-1, H=1
# u_x = -(-1) * (1/2) * y * (1 - y) = 0.5 * y * (1 - y)
# Maximum velocity at y=0.5: u_max = 0.5 * 0.5 * 0.5 = 0.125
# Create analytical solution function
def u_exact_expr(x):
"""Exact solution for Poiseuille flow"""
values = np.zeros((3, x.shape[1]))
values[0] = 0.5 * x[1] * (1.0 - x[1]) # u_x component
return values
u_exact = fem.Function(mpc_u.function_space)
u_exact.interpolate(u_exact_expr)
u_exact.x.scatter_forward()
# Compute error
error_u = uh.x.array - u_exact.x.array
error_norm = np.linalg.norm(error_u)
error_norm_global = domain.comm.allreduce(error_norm**2, op=MPI.SUM) ** 0.5
assert error_norm_global < np.sqrt(rtol), f"Velocity error too high: {error_norm_global}"
# Cleanup
ksp.destroy()
A.destroy()
b.destroy()
x.destroy()
|
