Poisson 2D (FEM)

Download .py All tutorials

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.asarray(fem.solve(linear=jno.solve.cg(), precond=jno.precond.jacobi()))  # SPD -> CG

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.
• fem.solve() alone would also work (matrix-free Jacobi-BiCGStab on the sparse operator); the solver slots pick a structure-appropriate method — CG for this SPD system — without ever densifying. See Choosing the solver for the full jno.solve / jno.precond namespaces.
• The solution is audited against the manufactured field \(u^\*=x(1-x)y(1-y)\) (rel-\(L^2 \approx 8\times10^{-3}\)).

Result

The computed field (left) matches the manufactured \(u^\*=x(1-x)y(1-y)\) to rel-\(L^2\approx8.6\times10^{-3}\) (middle, signed error). Re-solving on five successively finer meshes (right) shows the error falling as the mesh is refined.

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 with ``fem.solve()``. The default is a matrix-free Jacobi-preconditioned BiCGStab on
the sparse operator (never densifies); the solver slots choose anything else, e.g. CG for this
SPD system: ``fem.solve(linear=jno.solve.cg(), precond=jno.precond.jacobi())`` -- see
``docs/fem.md`` "Choosing the solver".
"""

import jax.numpy as jnp
import numpy as np

import jno

exact = lambda x, y: x * (1 - x) * y * (1 - y)  # noqa: E731

d = jno.Shape.rect(0, 0, 1, 1, size=0.18).domain()
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.asarray(fem.solve(linear=jno.solve.cg(), precond=jno.precond.jacobi()))  # SPD -> CG

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