3D Helmholtz on an F-shaped Domain (FEM)
A 3D screened-Helmholtz solve \(-\Delta u + \sigma u = f\) on a non-trivial geometry — an
extruded "F" prism meshed with gmsh into TET4 elements — with a Dirichlet bottom, a Neumann
top, and natural side walls.
Same API, one more coordinate
3D uses the identical workflow with a z axis: bind z=zi, take ui.z, and write the
Dirichlet condition over three coordinates. The geometry is any jNO/gmsh constructor:
d = jno.domain(constructor=letter_F_3d(mesh_size=0.4), compute_mesh_connectivity=True)
ui, vi = u.bind(x=xi, y=yi, z=zi), phi.bind(x=xi, y=yi, z=zi)
volume = ui.x * vi.x + ui.y * vi.y + ui.z * vi.z + sigma * u * vi - f * vi
fem = jno.fem([volume, top_neumann, u(cb[0], cb[1], cb[2]) - 0.0], element_type="TET4", quad_degree=2)
What to notice
- 3D is the same
jno.femAPI withzadded — no special path. - The extrusion uses
num_layers=8so the through-thickness mode is resolved. - Recovers \(u^\*=z+\alpha\sin(\pi z)\) to rel-\(L^2 \approx 2\times10^{-3}\).
Full script
"""05 - 3D screened-Helmholtz solve on an extruded F-shaped domain via ``jno.fem``.
-Delta u + sigma u = f on a 3D "F" prism (gmsh extrusion) (sigma = 4)
u = 0 on the bottom face (Dirichlet)
du/dn = 1 - alpha pi on the top face (Neumann)
du/dn = 0 on the side walls (natural)
Shows TET4 assembly on a non-trivial 3-D geometry with mixed Dirichlet / Neumann faces.
Manufactured solution u(x, y, z) = z + alpha sin(pi z).
"""
import os
os.environ["JAX_PLATFORMS"] = "cpu" # FEM surface arrays are CPU-pinned
import jax # noqa: E402
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp # noqa: E402
import numpy as np # noqa: E402
import jno # noqa: E402
alpha, sigma, MS = 0.20, 4.0, 0.40
exact = lambda z: z + alpha * jnp.sin(jnp.pi * z) # noqa: E731
dense = lambda A: jnp.asarray(A.todense()) if hasattr(A, "todense") else jnp.asarray(A) # noqa: E731
def letter_F_3d(depth=1.0, mesh_size=MS):
"""A jNO/gmsh constructor for a 3-D extruded "F" (tags: interior, bottom, top, wall)."""
outline = [
(0, 0),
(0.35, 0),
(0.35, 0.9),
(0.9, 0.9),
(0.9, 1.2),
(0.35, 1.2),
(0.35, 1.65),
(1.2, 1.65),
(1.2, 2.0),
(0, 2.0),
]
def construct(geo):
pts = [geo.add_point([x, y, 0.0], mesh_size=mesh_size) for x, y in outline]
loop = geo.add_curve_loop([geo.add_line(pts[i], pts[(i + 1) % len(pts)]) for i in range(len(pts))])
base = geo.add_plane_surface(loop)
flat = [] # flatten the extrusion result (volume + side/top surfaces)
def _flat(xs):
[(_flat(x) if isinstance(x, (list, tuple)) else flat.append(x)) for x in xs]
_flat(geo.extrude(base, [0.0, 0.0, depth], num_layers=8)) # 8 layers resolve sin(pi z) in z
surfaces = [e for e in flat if getattr(e, "dim", None) == 2]
volume = next(e for e in flat if getattr(e, "dim", None) == 3)
geo.add_physical(volume, "interior")
geo.add_physical([base] + surfaces, "boundary")
geo.add_physical([base], "bottom")
geo.add_physical([surfaces[0]], "top")
geo.add_physical(surfaces[1:], "wall")
return geo, 3, mesh_size
return construct
d = jno.domain(constructor=letter_F_3d(mesh_size=MS), compute_mesh_connectivity=True)
u, phi = d.fem_symbols()
ci, cb, ct = d.variable("interior", split=True), d.variable("bottom", split=True), d.variable("top", split=True)
xi, yi, zi = ci[0], ci[1], ci[2]
ui, vi = u.bind(x=xi, y=yi, z=zi), phi.bind(x=xi, y=yi, z=zi)
f = alpha * np.pi**2 * jno.np.sin(np.pi * zi) + sigma * (zi + alpha * jno.np.sin(np.pi * zi)) # -Delta u* + sigma u*
volume = ui.x * vi.x + ui.y * vi.y + ui.z * vi.z + sigma * u * vi - f * vi
top_neumann = -(1.0 - alpha * np.pi) * phi.bind(x=ct[0], y=ct[1], z=ct[2]) # du/dn given on the top face
fem = jno.fem([volume, top_neumann, u(cb[0], cb[1], cb[2]) - 0.0], element_type="TET4", quad_degree=2)
u_fem = jnp.linalg.solve(dense(fem.A), jnp.asarray(fem.b).reshape(-1))
pts = np.asarray(fem.points)
rel_l2 = float(jnp.linalg.norm(exact(pts[:, 2]) - u_fem) / jnp.linalg.norm(exact(pts[:, 2])))
print(f"\n3D F-domain screened Helmholtz via jno.fem: dofs={fem.dofs} rel_L2={rel_l2:.3e}")
assert fem.is_linear and rel_l2 < 1.0e-2