cai-aa/text-to-cae/fenics-fem

# Recommended: use conda with conda-forge (dolfinx + gmsh are complex to compile)
conda create -n fenics-env python=3.11 -y
conda activate fenics-env
conda install -c conda-forge fenics-dolfinx mpich petsc4py gmsh pyvista -y
 
# Verify installation
python -c "import dolfinx; print('dolfinx version:', dolfinx.__version__)"
python -c "import gmsh; print('gmsh version:', gmsh.__version__)"
 
# For Docker users (simplest approach)
docker pull dolfinx/dolfinx:stable
docker run -it --rm -v $(pwd):/work dolfinx/dolfinx:stable bash

"""
Solve −∇²u = f on the unit square [0,1]×[0,1]
with homogeneous Dirichlet BC u=0 on ∂Ω.
Manufactured solution: u_exact = sin(πx)sin(πy)
Source term:           f = 2π² sin(πx)sin(πy)
"""
from mpi4py import MPI
import numpy as np
from dolfinx import mesh, fem, io
from dolfinx.fem.petsc import LinearProblem
import ufl
def solve_poisson_unit_square(n_cells: int = 32, degree: int = 1) -> dict:
    """
    Solve the Poisson equation −∇²u = f on the unit square.
    Args:
        n_cells: Number of cells in each direction (total cells = 2*n_cells²).
        degree:  Polynomial degree of Lagrange finite elements.
    Returns:
        Dictionary with keys: uh (solution), L2_error, H1_error.
    """
    # --- Mesh ---
    domain = mesh.create_unit_square(
        MPI.COMM_WORLD, n_cells, n_cells, mesh.CellType.triangle
    )
    # --- Function space ---
    V = fem.functionspace(domain, ("Lagrange", degree))
    # --- Exact solution for manufactured source term ---
    x = ufl.SpatialCoordinate(domain)
    u_exact_expr = ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
    f_expr       = 2.0 * ufl.pi**2 * u_exact_expr
    # --- Dirichlet BC: u = 0 on all boundaries ---
    def boundary_all(x):
        return (
            np.isclose(x[0], 0.0) | np.isclose(x[0], 1.0) |
            np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0)
        )
    boundary_dofs = fem.locate_dofs_geometrical(V, boundary_all)
    u0 = fem.Function(V)
    u0.x.array[:] = 0.0
    bc = fem.dirichletbc(u0, boundary_dofs)
    # --- Variational problem ---
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
    L = f_expr * v * ufl.dx
    # --- Solve ---
    problem = LinearProblem(a, L, bcs=[bc],
                            petsc_options={"ksp_type": "cg", "pc_type": "hypre"})
    uh = problem.solve()
    # --- Compute errors ---
    diff = uh - u_exact_expr
    L2_error = float(fem.assemble_scalar(fem.form(ufl.inner(diff, diff) * ufl.dx)) ** 0.5)
    H1_error = float(fem.assemble_scalar(
        fem.form(ufl.inner(ufl.grad(diff), ufl.grad(diff)) * ufl.dx)
    ) ** 0.5)
    print(f"n_cells={n_cells}, degree={degree}: "
          f"L2={L2_error:.2e}, H1={H1_error:.2e}")
    return {"uh": uh, "L2_error": L2_error, "H1_error": H1_error}
def convergence_study() -> None:
    """Run convergence study with successive mesh refinement."""
    print("Convergence study for Poisson on unit square:")
    print(f"{'N':>6}  {'L2 error':>12}  {'H1 error':>12}  {'L2 rate':>8}")
    prev_L2 = None
    for n in [4, 8, 16, 32, 64]:
        result = solve_poisson_unit_square(n_cells=n, degree=1)
        L2 = result["L2_error"]
        rate = (np.log(prev_L2 / L2) / np.log(2.0)) if prev_L2 else float("nan")
        print(f"{n:>6}  {L2:>12.4e}  {result['H1_error']:>12.4e}  {rate:>8.2f}")
        prev_L2 = L2

"""
Solve linear elasticity on a 2D beam under body force (gravity).
Strong form: −div(σ(u)) = f in Ω
σ(u) = λ tr(ε(u)) I + 2μ ε(u)   (Hooke's law, Lamé form)
ε(u) = ½(∇u + ∇uᵀ)             (small-strain tensor)
"""
from mpi4py import MPI
import numpy as np
from dolfinx import mesh, fem, io
from dolfinx.fem.petsc import LinearProblem
import ufl
def solve_linear_elasticity(
    nx: int = 40,
    ny: int = 10,
    E: float = 210e9,     # Young's modulus (Pa), steel
    nu: float = 0.3,      # Poisson's ratio
    rho: float = 7850.0,  # Density (kg/m³)
    g: float = 9.81,      # Gravitational acceleration (m/s²)
    output_xdmf: str = "elasticity.xdmf",
) -> None:
    """
    Solve 2D linear elasticity (plane stress) on a rectangular beam.
    The left end is clamped (u=0), a body force f=(0,-ρg) is applied.
    Exports solution to XDMF for ParaView.
    Args:
        nx:          Cells in x direction.
        ny:          Cells in y direction.
        E:           Young's modulus in Pa.
        nu:          Poisson's ratio.
        rho:         Material density in kg/m³.
        g:           Gravity magnitude in m/s².
        output_xdmf: Output file path.
    """
    # Lamé parameters
    lam = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
    mu  = E / (2.0 * (1.0 + nu))
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [np.array([0.0, 0.0]), np.array([1.0, 0.25])],
        [nx, ny],
        cell_type=mesh.CellType.triangle,
    )
    # Vector function space (displacement u ∈ R²)
    V = fem.functionspace(domain, ("Lagrange", 1, (2,)))
    # Strain and stress tensors
    def epsilon(u):
        return ufl.sym(ufl.nabla_grad(u))
    def sigma(u):
        return lam * ufl.nabla_div(u) * ufl.Identity(2) + 2 * mu * epsilon(u)
    # Body force: gravity
    f = fem.Constant(domain, np.array([0.0, -rho * g]))
    # Clamped BC on left boundary x=0
    def left_boundary(x):
        return np.isclose(x[0], 0.0)
    boundary_dofs = fem.locate_dofs_geometrical(V, left_boundary)
    u_D = fem.Function(V)
    u_D.x.array[:] = 0.0
    bc = fem.dirichletbc(u_D, boundary_dofs)
    # Variational formulation
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    a = ufl.inner(sigma(u), epsilon(v)) * ufl.dx
    L = ufl.inner(f, v) * ufl.dx
    # Solve
    problem = LinearProblem(a, L, bcs=[bc],
                            petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
    uh = problem.solve()
    uh.name = "Displacement"
    # Compute von Mises stress for visualization
    s = sigma(uh) - (1.0 / 3) * ufl.tr(sigma(uh)) * ufl.Identity(2)
    von_mises = ufl.sqrt(1.5 * ufl.inner(s, s))
    W = fem.functionspace(domain, ("DG", 0))
    vm_expr = fem.Expression(von_mises, W.element.interpolation_points())
    vm_func = fem.Function(W)
    vm_func.interpolate(vm_expr)
    vm_func.name = "VonMises"
    # Export to XDMF
    with io.XDMFFile(MPI.COMM_WORLD, output_xdmf, "w") as xdmf:
        xdmf.write_mesh(domain)
        xdmf.write_function(uh)
        xdmf.write_function(vm_func)
    max_disp = np.max(np.abs(uh.x.array))
    print(f"Max displacement: {max_disp:.4e} m")
    print(f"Solution written to {output_xdmf}")
if __name__ == "__main__":
    solve_linear_elasticity()

"""
Use gmsh to generate a structured mesh of a disk with a hole,
then import into dolfinx for FEM analysis.
"""
import gmsh
import numpy as np
from mpi4py import MPI
from dolfinx.io.gmshio import model_to_mesh
from dolfinx import fem, io
import ufl
from dolfinx.fem.petsc import LinearProblem
def create_annular_mesh(
    r_inner: float = 0.2,
    r_outer: float = 1.0,
    mesh_size: float = 0.05,
    output_msh: str = "annulus.msh",
) -> None:
    """
    Create a 2D annular mesh (disk with circular hole) using gmsh.
    Args:
        r_inner:   Inner radius (hole).
        r_outer:   Outer radius.
        mesh_size: Target mesh element size.
        output_msh: Output .msh file path.
    """
    gmsh.initialize()
    gmsh.model.add("annulus")
    # Outer disk
    outer = gmsh.model.occ.addDisk(0, 0, 0, r_outer, r_outer)
    # Inner disk (hole)
    inner = gmsh.model.occ.addDisk(0, 0, 0, r_inner, r_inner)
    # Boolean cut: outer minus inner
    gmsh.model.occ.cut([(2, outer)], [(2, inner)])
    gmsh.model.occ.synchronize()
    # Mesh size field
    gmsh.model.mesh.setSize(gmsh.model.getEntities(0), mesh_size)
    # Physical groups (needed for dolfinx boundary conditions)
    surfaces = gmsh.model.getEntities(2)
    for _, tag in surfaces:
        gmsh.model.addPhysicalGroup(2, [tag], tag)
        gmsh.model.setPhysicalName(2, tag, f"domain_{tag}")
    curves = gmsh.model.getBoundary(surfaces, oriented=False)
    for dim, tag in curves:
        gmsh.model.addPhysicalGroup(1, [abs(tag)], abs(tag))
        gmsh.model.setPhysicalName(1, abs(tag), f"boundary_{abs(tag)}")
    gmsh.model.mesh.generate(2)
    gmsh.model.mesh.optimize("Netgen")
    gmsh.write(output_msh)
    gmsh.finalize()
    print(f"Annular mesh written to {output_msh}")
def solve_poisson_annulus(msh_file: str = "annulus.msh") -> None:
    """
    Import gmsh mesh and solve Poisson equation on the annular domain.
    BC: u=1 on inner boundary, u=0 on outer boundary.
    """
    gmsh.initialize()
    gmsh.open(msh_file)
    domain, cell_tags, facet_tags = model_to_mesh(
        gmsh.model, MPI.COMM_WORLD, 0, gdim=2
    )
    gmsh.finalize()
    V = fem.functionspace(domain, ("Lagrange", 2))
    # Locate facets for BCs by geometric criterion
    def inner_boundary(x):
        return np.sqrt(x[0]**2 + x[1]**2) < 0.25
    def outer_boundary(x):
        return np.sqrt(x[0]**2 + x[1]**2) > 0.9
    dofs_inner = fem.locate_dofs_geometrical(V, inner_boundary)
    dofs_outer = fem.locate_dofs_geometrical(V, outer_boundary)
    u_inner = fem.Function(V); u_inner.x.array[:] = 1.0
    u_outer = fem.Function(V); u_outer.x.array[:] = 0.0
    bcs = [
        fem.dirichletbc(u_inner, dofs_inner),
        fem.dirichletbc(u_outer, dofs_outer),
    ]
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
    L = fem.Constant(domain, 0.0) * v * ufl.dx  # Laplace: f=0
    problem = LinearProblem(a, L, bcs=bcs)
    uh = problem.solve()
    uh.name = "u"