Fenics Fem

"""
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"

    with io.XDMFFile(MPI.COMM_WORLD, "annulus_solution.xdmf", "w") as xdmf:
        xdmf.write_mesh(domain)
        xdmf.write_function(uh)
    print("Annulus Laplace solution saved to annulus_solution.xdmf")

"""
Solve the steady Stokes equations (viscous flow at Re→0) on a channel.

−μ ∇²u + ∇p = f
∇·u = 0

Uses Taylor-Hood P2/P1 elements (LBB-stable mixed formulation).
"""

from mpi4py import MPI
import numpy as np
from dolfinx import mesh, fem
from dolfinx.fem.petsc import LinearProblem
import ufl


def solve_stokes_channel(
    nx: int = 64,
    ny: int = 16,
    mu: float = 1.0,      # dynamic viscosity
    U_max: float = 1.0,   # max inlet velocity
) -> None:
    """
    Solve steady Stokes flow in a 2D channel.
    Parabolic inlet profile, no-slip walls, stress-free outlet.
    """
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [np.array([0.0, 0.0]), np.array([4.0, 1.0])],
        [nx, ny],
        cell_type=mesh.CellType.triangle,
    )

    # Taylor-Hood P2/P1 elements
    P2 = fem.functionspace(domain, ("Lagrange", 2, (2,)))
    P1 = fem.functionspace(domain, ("Lagrange", 1))

    # Mixed space
    V_el = ufl.VectorElement("Lagrange", domain.ufl_cell(), 2)
    Q_el = ufl.FiniteElement("Lagrange", domain.ufl_cell(), 1)
    W = fem.functionspace(domain, ufl.MixedElement([V_el, Q_el]))

    (u, p) = ufl.TrialFunctions(W)
    (v, q) = ufl.TestFunctions(W)

    f = fem.Constant(domain, np.array([0.0, 0.0]))

    a = (mu * ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
         - ufl.div(v) * p * ufl.dx
         + q * ufl.div(u) * ufl.dx)
    L = ufl.inner(f, v) * ufl.dx

    # No-slip on top and bottom walls
    def walls(x):
        return np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0)

    # Parabolic inlet profile: u_x = 4*U_max*y*(1-y), u_y=0
    def inlet_velocity(x):
        vals = np.zeros((2, x.shape[1]))
        vals[0] = 4.0 * U_max * x[1] * (1.0 - x[1])
        return vals

    def inlet(x):
        return np.isclose(x[0], 0.0)

    W0, _ = W.sub(0).collapse()
    dofs_walls = fem.locate_dofs_geometrical((W.sub(0), W0), walls)
    dofs_inlet = fem.locate_dofs_geometrical((W.sub(0), W0), inlet)

    u_no_slip = fem.Function(W0); u_no_slip.x.array[:] = 0.0
    u_inflow  = fem.Function(W0); u_inflow.interpolate(inlet_velocity)

    bcs = [
        fem.dirichletbc(u_no_slip, dofs_walls,  W.sub(0)),
        fem.dirichletbc(u_inflow,  dofs_inlet,  W.sub(0)),
    ]

    problem = LinearProblem(a, L, bcs=bcs,
                            petsc_options={"ksp_type": "minres", "pc_type": "hypre"})
    wh = problem.solve()
    print("Stokes channel flow solved. Extract wh.sub(0) for velocity, wh.sub(1) for pressure.")

"""
Solve the unsteady heat equation with backward Euler time integration.
∂u/∂t − α∇²u = 0  (α = thermal diffusivity)
"""

from mpi4py import MPI
import numpy as np
from dolfinx import mesh, fem
from dolfinx.fem.petsc import LinearProblem
import ufl


def solve_heat_equation(
    nx: int = 40,
    alpha: float = 0.01,
    T_final: float = 1.0,
    dt: float = 0.01,
) -> None:
    """Solve 2D heat equation with Gaussian initial condition."""
    domain = mesh.create_unit_square(MPI.COMM_WORLD, nx, nx)
    V = fem.functionspace(domain, ("Lagrange", 1))

    u_n = fem.Function(V)
    x = fem.Expression(
        ufl.exp(-50.0 * ((ufl.SpatialCoordinate(domain)[0] - 0.5)**2
                         + (ufl.SpatialCoordinate(domain)[1] - 0.5)**2)),
        V.element.interpolation_points()
    )
    u_n.interpolate(x)

    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    dt_const = fem.Constant(domain, dt)

    a = (u * v + dt_const * alpha * ufl.inner(ufl.grad(u), ufl.grad(v))) * ufl.dx
    L = u_n * v * ufl.dx

    problem = LinearProblem(a, L, bcs=[],
                            petsc_options={"ksp_type": "cg", "pc_type": "hypre"})

    t = 0.0
    n_steps = int(T_final / dt)
    for step in range(n_steps):
        t += dt
        uh = problem.solve()
        u_n.x.array[:] = uh.x.array[:]
        if step % 20 == 0:
            max_u = np.max(np.abs(uh.x.array))
            print(f"t={t:.3f}: max(u)={max_u:.4f}")
    print("Heat equation time integration complete.")

if __name__ == "__main__":
    print("=== Poisson Convergence Study ===")
    convergence_study()
    print()
    print("=== Poisson on Unit Square (n=64) ===")
    result = solve_poisson_unit_square(n_cells=64, degree=2)
    print(f"L2 error with P2 elements: {result['L2_error']:.2e}")

if __name__ == "__main__":
    print("=== Linear Elasticity: Steel Cantilever Beam ===")
    solve_linear_elasticity(
        nx=80, ny=20,
        E=210e9,    # Steel Young's modulus
        nu=0.3,
        rho=7850.0,
        g=9.81,
        output_xdmf="steel_beam.xdmf",
    )
    print("Open steel_beam.xdmf in ParaView to visualize displacement and von Mises stress.")

if __name__ == "__main__":
    print("=== gmsh Annular Mesh + Laplace Equation ===")
    create_annular_mesh(r_inner=0.2, r_outer=1.0, mesh_size=0.05)
    solve_poisson_annulus("annulus.msh")
    print("Open annulus_solution.xdmf in ParaView.")