Skip to content

Poisson 2D (FEM)

Download .py Back to chapter

A pure finite-element solve of \(-\Delta u = f\) on the unit square with \(u = 0\) on the boundary, assembled through jno.fem. It is the primer for the API: write the weak form as a list of residual terms — the volume physics plus the essential boundary condition as u(region) - g — and solve the assembled A u = b.

The weak form

The trial/test functions come from d.fem_symbols(); binding them to the quadrature coordinates gives the .x / .y derivatives, and the Dirichlet condition is simply the term u(xb, yb) - 0.0:

ui, vi = u.bind(x=xi, y=yi), phi.bind(x=xi, y=yi)
f = 2.0 * (xi * (1 - xi) + yi * (1 - yi))
fem = jno.fem([ui.x * vi.x + ui.y * vi.y - f * vi, u(xb, yb) - 0.0], quad_degree=3)
u_fem = jnp.linalg.solve(dense(fem.A), jnp.asarray(fem.b).reshape(-1))

What to notice

  • One call, jno.fem([...]), assembles the stiffness fem.A and load fem.b.
  • Boundary conditions are terms in the same list — no separate BC objects.
  • The solution is audited against the manufactured field \(u^\*=x(1-x)y(1-y)\) (rel-\(L^2 \approx 8\times10^{-3}\)).

Full script

"""01 - 2D Poisson equation assembled + solved through ``jno.fem``.

    -Delta u = f on the unit square, u = 0 on the boundary.
    Manufactured  u*(x, y) = x(1-x) y(1-y),   f = -Delta u* = 2[x(1-x) + y(1-y)].

A pure finite-element solve: write the weak form as a list of residual terms (volume physics
+ the essential condition ``u(region) - g``), hand it to ``jno.fem`` -- the single FEM entry --
and solve the assembled ``A u = b``.
"""

import jax.numpy as jnp
import numpy as np
from shapely.geometry import box

import jno

exact = lambda x, y: x * (1 - x) * y * (1 - y)  # noqa: E731
dense = lambda A: jnp.asarray(A.todense()) if hasattr(A, "todense") else jnp.asarray(A)  # noqa: E731

d = jno.domain(box(0.0, 0.0, 1.0, 1.0), mesh_size=0.18)
u, phi = d.fem_symbols()
xi, yi, _ = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
ui, vi = u.bind(x=xi, y=yi), phi.bind(x=xi, y=yi)

f = 2.0 * (xi * (1 - xi) + yi * (1 - yi))
fem = jno.fem([ui.x * vi.x + ui.y * vi.y - f * vi, u(xb, yb) - 0.0], quad_degree=3)  # weak form + u = 0
u_fem = jnp.linalg.solve(dense(fem.A), jnp.asarray(fem.b).reshape(-1))  # solve A u = b

pts = np.asarray(fem.points)  # coordinates the DOFs live on
rel_l2 = float(jnp.linalg.norm(exact(pts[:, 0], pts[:, 1]) - u_fem) / jnp.linalg.norm(exact(pts[:, 0], pts[:, 1])))
print(f"\nPoisson via jno.fem: dofs={fem.dofs}  rel_L2={rel_l2:.3e}")
assert fem.is_linear and rel_l2 < 5e-2