Hi everyone,
I implemented a FEniCS code for Boussinesq flow using mixed elements ([P2, P1, P1]), and it worked well. However, after transitioning to FEniCSx, the nonlinear Newton solver fails to converge, with the residual going to -nan.
The problem is the differentially heated square cavity for natural convection. The boundary conditions are:
Velocity: No-slip at all walls
Temperature: High temperature on the left wall, lower temperature on the right wall
Pressure: A constraint at the bottom right corner for better convergence
I have carefully verified my weak forms, which worked correctly in FEniCS. I suspect the issue may be related to how boundary conditions are applied in FEniCSx.
Below is my FEniCSx implementation:
import numpy as np
from mpi4py import MPI
from petsc4py import PETSc
from dolfinx import mesh, fem, plot
from dolfinx.plot import vtk_mesh
from dolfinx.io import XDMFFile
from dolfinx.mesh import create_unit_square, CellType, locate_entities_boundary
import basix
from basix.ufl import element, mixed_element
from ufl import split, grad, inner, dx, div, sym, dot, Constant, TestFunctions, TrialFunctions, as_vector, derivative
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver
from dolfinx.log import set_log_level, LogLevel
import pyvista as pv
import pyvistaqt as pvqt
from pyvista.utilities import xvfb
set_log_level(LogLevel.INFO) # Enables logging to monitor solver progress
def plot_vector_field(field, function_space, mesh, field_name, output_filename, scale_factor=0.2):
pv.OFF_SCREEN = True
pv.start_xvfb()
field_collapsed = field.collapse()
topology, cell_types, geometry = plot.vtk_mesh(function_space)
values = np.zeros((geometry.shape[0], 3), dtype=np.float64)
values[:, :len(field_collapsed)] = field_collapsed.x.array.real.reshape((geometry.shape[0], len(field_collapsed)))
function_grid = pv.UnstructuredGrid(topology, cell_types, geometry)
function_grid[field_name] = values
glyphs = function_grid.glyph(orient=field_name, factor=scale_factor)
mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim)
grid = pv.UnstructuredGrid(*plot.vtk_mesh(mesh, mesh.topology.dim))
plotter = pv.Plotter(off_screen=True)
plotter.add_mesh(grid, style="wireframe", color="k")
plotter.add_mesh(glyphs, show_scalar_bar=True)
plotter.view_xy()
plotter.screenshot(output_filename)
print(f"{field_name} visualization saved as {output_filename}")
def plot_scalar_field(field, function_space, field_name, output_filename):
pv.OFF_SCREEN = True
pv.start_xvfb()
u_n = field.collapse()
grid = pv.UnstructuredGrid(*plot.vtk_mesh(function_space))
grid.point_data[field_name] = u_n.x.array.real
grid.set_active_scalars(field_name)
plotter = pv.Plotter(off_screen=True)
plotter.add_mesh(grid, show_edges=True, show_scalar_bar=True)
plotter.view_xy()
plotter.screenshot(output_filename)
print(f"{field_name} field visualization saved as {output_filename}")
# Create a mesh
nx, ny = 32, 32
domain = create_unit_square(MPI.COMM_WORLD, nx, ny)
# Define vector and scalar elements using basix.element
P2 = element("P", domain.basix_cell(), degree=2, shape=(domain.geometry.dim,)) # Velocity element
P1 = element("P", domain.basix_cell(), degree=1) # Pressure and temperature elements
# Create a mixed function space
W = fem.functionspace(domain, mixed_element([P2, P1, P1]))
# Define trial and test functions
w = fem.Function(W)
w0 = fem.Function(W) # Previous solution
(v, p, T) = split(w) # Trial functions
(v_test, p_test, T_test) = TestFunctions(W) # Test functions
w0.x.array[:] = 0.0
# Collapse the subspaces
V_vel, _ = W.sub(0).collapse()
T_space, _ = W.sub(2).collapse()
p_space, _ = W.sub(1).collapse()
# Define initialization functions
def T_init(x):
values = 1.0 - x[0] # Linearly decreasing with respect to x-axis
return values.reshape(1, -1)
def V_init(x):
return np.zeros((domain.geometry.dim, x.shape[1]))
def P_init(x):
return np.zeros((1, x.shape[1]))
# Assign the interpolated values back to the mixed function space
w0.sub(0).interpolate(V_init)
w0.sub(1).interpolate(P_init)
w0.sub(2).interpolate(T_init)
# w.x.scatter_forward()
# w.sub(2).interpolate(T_init) # Small perturbation
print(w0.sub(0).x.array[:5])
print(w0.sub(1).x.array[:5])
print(w0.sub(2).x.array[:5])
print("Initial velocity max:", np.max(w0.sub(0).x.array))
print("Initial pressure max:", np.max(w0.sub(1).x.array))
print("Initial temperature max:", np.max(w0.sub(2).x.array))
(v0, p0, T0) = split(w0) # Previous solution
# Time-stepping parameters
dt = 0.005
T_end = 0.5
t = 0.0
num_steps = int(T_end / dt)
# Define material properties
rho_0 = 1.0 # Density
mu = 1.0 # Dynamic viscosity
K = 1.0 # Thermal conductivity
C_p = 1.0 # Heat capacity
alpha = K / (rho_0 * C_p) # Thermal diffusivity
Ra = 1.e3 # Rayleigh number
# # Define buoyancy force
# def f_expr(x):
# """Expression for the applied buoyancy force"""
# force = Ra * T * C_p * mu / K
# return np.vstack((np.zeros_like(x[0]), force)) # Force applied in y-direction
# f_B = fem.Function(V_vel)
# f_B.interpolate(f_expr)
f_B = as_vector([0, Ra * T * C_p * mu / K])
v_t = (v - v0) / dt
T_t = (T - T0) / dt
# Define variational forms
dx = dx(metadata={"quadrature_degree": 3 * 2})
F_mass = -div(v) * p_test * dx
F_momentum = (rho_0 * dot(v_t, v_test) * dx
+ rho_0 * dot(dot(grad(v), v), v_test) * dx
+ 2 * mu * inner(sym(grad(v)), sym(grad(v_test))) * dx
- p * div(v_test) * dx
- rho_0 * dot(f_B, v_test) * dx)
F_energy = (dot(T_t, T_test) * dx
- dot(v * T, grad(T_test)) * dx
+ alpha * dot(grad(T), grad(T_test)) * dx)
# Combine all forms
F = F_mass + F_momentum + F_energy
# Define boundary markers
def walls(x):
return np.logical_or.reduce((
np.isclose(x[0], 0.0), # Left wall
np.isclose(x[0], 1.0), # Right wall
np.isclose(x[1], 0.0), # Bottom wall
np.isclose(x[1], 1.0) # Top wall
))
def hot_wall(x):
return np.isclose(x[0], 0.0)
def cold_wall(x):
return np.isclose(x[0], 1.0)
def corner_point(x):
return np.logical_and(np.isclose(x[0], 1.0), np.isclose(x[1], 0.0))
boundary_facets = locate_entities_boundary(domain, domain.topology.dim - 1, walls)
hot_wall_facets = locate_entities_boundary(domain, domain.topology.dim - 1, hot_wall)
cold_wall_facets = locate_entities_boundary(domain, domain.topology.dim - 1, cold_wall)
corner_facets = locate_entities_boundary(domain, domain.topology.dim - 1, corner_point)
# Collapse the subspaces
V_vel, _ = W.sub(0).collapse()
T_space, _ = W.sub(2).collapse()
p_space, _ = W.sub(1).collapse()
# Define boundary function values
u_noslip_func = fem.Function(V_vel)
u_noslip_func.interpolate(lambda x: np.zeros((domain.geometry.dim, x.shape[1])))
T_hot_func = fem.Function(T_space)
T_hot_func.interpolate(lambda x: np.ones((1, x.shape[1])))
T_cold_func = fem.Function(T_space)
T_cold_func.interpolate(lambda x: np.zeros((1, x.shape[1])))
p_corner_func = fem.Function(p_space)
p_corner_func.interpolate(lambda x: np.zeros((1, x.shape[1])))
# Locate boundary DOFs using topology
boundary_dofs_v = fem.locate_dofs_topological(W.sub(0), domain.topology.dim - 1, boundary_facets)
boundary_dofs_T_hot = fem.locate_dofs_topological(W.sub(2), domain.topology.dim - 1, hot_wall_facets)
boundary_dofs_T_cold = fem.locate_dofs_topological(W.sub(2), domain.topology.dim - 1, cold_wall_facets)
corner_dofs = fem.locate_dofs_geometrical(p_space, corner_point) #pointwise
# Define boundary conditions for velocity (no-slip condition)
u_noslip = np.array((0,) * domain.geometry.dim, dtype=PETSc.ScalarType)
bc_v = fem.dirichletbc(u_noslip_func, boundary_dofs_v)#, W.sub(0))
# Define boundary conditions for temperature
bc_T_hot = fem.dirichletbc(T_hot_func, boundary_dofs_T_hot) #, W.sub(2))
bc_T_cold = fem.dirichletbc(T_cold_func, boundary_dofs_T_cold) #, W.sub(2))
# Define the Dirichlet condition for pressure at the corner
bc_p_corner = fem.dirichletbc(p_corner_func, corner_dofs) #, W.sub(1))
# Apply boundary conditions
bcs = [bc_v, bc_T_hot, bc_T_cold, bc_p_corner]
if np.any(np.isnan(w.x.array)):
raise RuntimeError(f"NaN detected in solution at step {step}, time {t:.3f}")
# Create nonlinear problem and solver
problem = NonlinearProblem(F, w, bcs)
solver = NewtonSolver(domain.comm, problem)
solver.rtol = 1e-6 # tolerance
solver.atol = 1e-6 # absolute tolerance
solver.max_it = 50 # max iterations (default is 10)
solver.convergence_criterion = "residual"
solver.relaxation_parameter = 0.7 # Add relaxation
solver.report = True
# Configure PETSc linear solver options
ksp = solver.krylov_solver
opts = PETSc.Options()
option_prefix = ksp.getOptionsPrefix()
# opts[f"{option_prefix}ksp_type"] = "gmres"
# opts[f"{option_prefix}pc_type"] = "ilu"
opts[f"{option_prefix}ksp_type"] = "preonly" # Use direct solver
opts[f"{option_prefix}pc_type"] = "lu" # Use LU factorization
opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps" # Set MUMPS as the solver
ksp.setFromOptions()
# Prepare output file
output_file = XDMFFile(MPI.COMM_WORLD, "output.xdmf", "w")
output_file.write_mesh(domain)
# Visualization setup
xvfb.start_xvfb()
topology, cell_types, x = plot.vtk_mesh(T_space)
grid = pv.UnstructuredGrid(topology, cell_types, x)
grid.point_data["T"] = w.sub(2).collapse().x.array
plotter = pv.Plotter()
plotter.open_gif("temperature_evolution.gif", fps=5)
plotter.add_mesh(grid, scalars="T", show_edges=True)
plotter.add_text(f"time: {t:.3f}", font_size=12, name="timelabel")
# Time-stepping loop
for step in range(num_steps):
print("Max velocity:", np.max(w.sub(0).x.array))
print("Max pressure:", np.max(w.sub(1).x.array))
print("Max temperature:", np.max(w.sub(2).x.array))
if np.any(np.isnan(w.x.array)):
raise RuntimeError(f"NaN detected in solution at step {step}, time {t:.3f}")
t += dt
num_iterations, converged = solver.solve(w)
if not converged:
raise RuntimeError(f"Solver failed to converge at step {step}, time {t:.3f}")
w0.x.array[:] = w.x.array # Update previous solution
print(f"Time: {t:.3f}, Newton iterations: {num_iterations}")
# Save results
output_file.write_function(w.sub(2), t)
# Update visualization
grid.point_data["T"][:] = w.sub(2).collapse().x.array
plotter.update_scalars(grid.point_data["T"])
plotter.write_frame()
# Finalize visualization
plotter.close()
output_file.close()
print("Simulation complete!")
# Plot the final results for velocity and pressure
plot_vector_field(w.sub(0), V_vel, domain, "Velocity", "velocity_final.png", scale_factor=0.5)
plot_scalar_field(w.sub(2), T_space, "Temperature", "temperature_final.png")
plot_scalar_field(w.sub(1), p_space, "Pressure", "pressure_final.png")
print("Simulation complete!")
Has anyone encountered similar issues when moving from FEniCS to FEniCSx? Any insights or suggestions would be greatly appreciated.