Solve the Poisson and linearised elasticity equations using pyamg
Download sources
The demo illustrates solving the Poisson and linearised elasticity equations with using algebraic multigrid from pyamg. pyamg is not MPI-parallel, therefore this demo runs in serial only.
import sys
from mpi4py import MPI
import numpy as np
import numpy.typing as npt
import ufl
from dolfinx import fem, io
from dolfinx.fem import (
apply_lifting,
assemble_matrix,
assemble_vector,
dirichletbc,
form,
functionspace,
locate_dofs_topological,
)
from dolfinx.mesh import CellType, create_box, locate_entities_boundary
try:
import pyamg
except ImportError:
print("This demo requires pyamg.")
exit(0)
if MPI.COMM_WORLD.size > 1:
print("This demo works only in serial.")
exit(0)
def poisson_problem(dtype: npt.DTypeLike, solver_type: str) -> None:
"""Solve a 3D Poisson problem using Ruge-Stuben algebraic multigrid.
Args:
dtype: Scalar type to use.
solver_type: pyamg solver type, either "ruge_stuben" or
"smoothed_aggregation"
"""
real_type = np.real(dtype(0)).dtype
mesh = create_box(
comm=MPI.COMM_WORLD,
points=[(0.0, 0.0, 0.0), (3.0, 2.0, 1.0)],
n=(30, 20, 10),
cell_type=CellType.tetrahedron,
dtype=real_type,
)
V = functionspace(mesh, ("Lagrange", 1))
tdim = mesh.topology.dim
fdim = tdim - 1
facets = locate_entities_boundary(
mesh,
dim=fdim,
marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 3.0),
)
dofs = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets)
bc = dirichletbc(value=dtype(0), dofs=dofs, V=V)
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
x = ufl.SpatialCoordinate(mesh)
f = 10 * ufl.exp(-((x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2) / 0.02)
g = ufl.sin(5 * x[0])
a = form(ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx, dtype=dtype)
L = form(ufl.inner(f, v) * ufl.dx + ufl.inner(g, v) * ufl.ds, dtype=dtype)
A = assemble_matrix(a, [bc]).to_scipy()
b = assemble_vector(L)
apply_lifting(b.array, [a], bcs=[[bc]])
bc.set(b.array)
# Create solution function and create AMG solver
uh = fem.Function(V, dtype=dtype)
if solver_type == "ruge_stuben":
ml = pyamg.ruge_stuben_solver(A)
elif solver_type == "smoothed_aggregation":
ml = pyamg.smoothed_aggregation_solver(A)
else:
raise ValueError(f"Invalid multigrid type: {solver_type}")
print(ml)
# Solve linear systems
print(f"\nSolve Poisson equation: {dtype.__name__}")
res: list[float] = []
tol = 1e-10 if real_type == np.float64 else 1e-6
uh.x.array[:] = ml.solve(b.array, tol=tol, residuals=res, accel="cg")
for i, q in enumerate(res):
print(f"Convergence history: iteration {i}, residual= {q}")
with io.XDMFFile(mesh.comm, f"out_pyamg/poisson_{dtype.__name__}.xdmf", "w") as file:
file.write_mesh(mesh)
file.write_function(uh)
def nullspace_elasticty(Q: fem.FunctionSpace) -> list[np.ndarray]:
"""Create the elasticity (near)nulspace.
Args:
Q: Displacement field function space.
Returns:
Nullspace for the unconstrained problem.
"""
B = np.zeros((Q.dofmap.index_map.size_local * Q.dofmap.bs, 6))
# Get dof indices for each subspace (x, y and z dofs)
dofs = [Q.sub(i).dofmap.list.flatten() for i in range(3)]
# Set the three translational rigid body modes
for i in range(3):
B[dofs[i], i] = 1.0
# Set the three rotational rigid body modes
x = Q.tabulate_dof_coordinates()
dofs_block = Q.dofmap.list.flatten()
x0, x1, x2 = x[dofs_block, 0], x[dofs_block, 1], x[dofs_block, 2]
B[dofs[0], 3] = -x1
B[dofs[1], 3] = x0
B[dofs[0], 4] = x2
B[dofs[2], 4] = -x0
B[dofs[2], 5] = x1
B[dofs[1], 5] = -x2
return B
def elasticity_problem(dtype) -> None:
"""Solve a 3D linearised elasticity problem using AMG.
Uses a smoothed aggregation algebraic multigrid method.
Args:
dtype: Scalar type to use.
"""
mesh = create_box(
comm=MPI.COMM_WORLD,
points=[(0.0, 0.0, 0.0), (3.0, 2.0, 1.0)],
n=(30, 20, 10),
cell_type=CellType.tetrahedron,
dtype=dtype,
)
tdim = mesh.topology.dim
fdim = tdim - 1
facets = locate_entities_boundary(
mesh,
dim=fdim,
marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 3.0),
)
ω, ρ = 300.0, 10.0
x = ufl.SpatialCoordinate(mesh)
f = ufl.as_vector((ρ * ω**2 * x[0], ρ * ω**2 * x[1], 0.0))
# Define the elasticity parameters and create a function that
# computes an expression for the stress given a displacement field.
E = 1.0e9
ν = 0.3
μ = E / (2.0 * (1.0 + ν))
λ = E * ν / ((1.0 + ν) * (1.0 - 2.0 * ν))
def σ(v):
"""Expression for the stress σ given a displacement field."""
return 2.0 * μ * ufl.sym(ufl.grad(v)) + λ * ufl.tr(ufl.sym(ufl.grad(v))) * ufl.Identity(
len(v)
)
V = functionspace(mesh, ("Lagrange", 1, (mesh.geometry.dim,)))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
a = form(ufl.inner(σ(u), ufl.grad(v)) * ufl.dx, dtype=dtype)
L = form(ufl.inner(f, v) * ufl.dx, dtype=dtype)
dofs = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets)
bc = dirichletbc(np.zeros(3, dtype=dtype), dofs, V=V)
A = assemble_matrix(a, bcs=[bc]).to_scipy()
b = assemble_vector(L)
apply_lifting(b.array, [a], bcs=[[bc]])
bc.set(b.array)
uh = fem.Function(V, dtype=dtype)
B = nullspace_elasticty(V)
ml = pyamg.smoothed_aggregation_solver(A, B=B)
print(ml)
print(f"\nLinearised elasticity: {dtype.__name__}")
res_e: list[float] = []
tol = 1e-10 if dtype == np.float64 else 1e-6
uh.x.array[:] = ml.solve(b.array, tol=tol, residuals=res_e, accel="cg")
for i, q in enumerate(res_e):
print(f"Convergence history: iteration {i}, residual= {q}")
with io.XDMFFile(mesh.comm, f"out_pyamg/elasticity_{dtype.__name__}.xdmf", "w") as file:
file.write_mesh(mesh)
file.write_function(uh)
Solve Poission problem with different scalar types
poisson_problem(np.float32, "ruge_stuben")
poisson_problem(np.float64, "ruge_stuben")
For complex, pyamg requires smoothed aggregation multigrid
if not sys.platform.startswith("win32"):
poisson_problem(np.complex64, "smoothed_aggregation")
poisson_problem(np.complex128, "smoothed_aggregation")
Solve elasticity problem with different scalar types
elasticity_problem(np.float32)
elasticity_problem(np.float64)