Mixed-Boundary Poisson 2D

This example shows how to combine different boundary conditions on different parts of the same domain.

Problem Setup

-Delta u = f(x,y),   (x,y) in [0,1]^2
u = 0 on x = 0 and x = 1
du/dy = 0 on y = 0 and y = 1

with exact solution u(x,y) = sin(pi x) cos(pi y).

Step 1: Sample Boundary Segments Separately

The script requests top and bottom variables from the domain in addition to interior points. This lets it apply separate Neumann terms to those boundaries.

domain = jno.domain.rect(mesh_size=0.05)
x,  y,  _ = domain.variable("interior")
xt, yt, _ = domain.variable("top")
xb, yb, _ = domain.variable("bottom")

u_exact = jno.np.sin(pi * x) * jno.np.cos(pi * y)
forcing  = 2 * pi**2 * u_exact

Step 2: Hard-Enforce the Dirichlet Part

The model is multiplied by x(1-x), which automatically satisfies the zero-value condition on the left and right sides.

net = jno.nn.wrap(foundax.mlp(in_features=2, hidden_dims=80, num_layers=5, key=jax.random.PRNGKey(14)))
net.optimizer(optax.adam(optax.exponential_decay(1e-3, 80, 0.5, end_value=1e-5)))

u        = net(x,  y)  * x  * (1 - x)
u_top    = net(xt, yt) * xt * (1 - xt)
u_bottom = net(xb, yb) * xb * (1 - xb)

Step 3: Add Neumann Residuals

The top and bottom boundary losses are built by differentiating the boundary-evaluated field with respect to y.

pde           = -jno.np.laplacian(u, [x, y]) - forcing
neumann_top   = u_top.d(yt)
neumann_bottom = u_bottom.d(yb)

Step 4: Solve With Multiple Loss Terms

The core includes the PDE residual and both Neumann losses.

crux    = jno.core([pde.mse, neumann_top.mse, neumann_bottom.mse])
history = crux.solve(40_000)

What To Notice

Mixed boundary problems are common in practice.
Tagged boundary variables make it straightforward to isolate different edges.
You do not need to choose between all-hard or all-soft boundary handling.

Download full script Back to 02 Elliptic

Script Snippet

"""02 — 2-D Poisson equation with mixed boundary conditions"""

from pathlib import Path

import foundax
import jax
import optax

import jno

π = jno.np.pi
# jno.domain.rect gives left/right/top/bottom edge tags out of the box.
domain = jno.domain.rect(mesh_size=0.05)
x, y, _ = domain.variable("interior")
xt, yt, _ = domain.variable("top")
xb, yb, _ = domain.variable("bottom")

u_exact = jno.np.sin(π * x) * jno.np.cos(π * y)
forcing = 2 * π**2 * u_exact

net = jno.nn.wrap(foundax.mlp(in_features=2, hidden_dims=48, num_layers=4, key=jax.random.PRNGKey(14)))
net.optimizer(optax.adam(optax.exponential_decay(1e-3, 1000, 0.5, end_value=1e-5)))

# Hard Dirichlet on x = 0, 1 via the x(1-x) factor; Neumann on y boundaries is soft.
u = (net(x, y) * x * (1 - x)).scalar.bind(x=x, y=y)
u_top = (net(xt, yt) * xt * (1 - xt)).scalar.bind(x=xt, y=yt)
u_bot = (net(xb, yb) * xb * (1 - xb)).scalar.bind(x=xb, y=yb)

pde = -(u.xx + u.yy) - forcing
neumann_top = u_top.y  # ∂u/∂y at y = 1
neumann_bot = u_bot.y  # ∂u/∂y at y = 0

crux = jno.core([pde.mse, neumann_top.mse, neumann_bot.mse])
crux.solve(5000)

_u, _u_exact = crux.eval([u, u_exact])
rel_l2 = float(jax.numpy.linalg.norm(_u - _u_exact) / (jax.numpy.linalg.norm(_u_exact) + 1e-8))
print(f"Relative L2 error: {rel_l2:.4e}")

results_file = Path(__file__).parent.parent.parent / "tutorial_results.txt"
with open(results_file, "a") as f:
    f.write(f"02_elliptic/mixed_boundary_poisson_2d.py | epochs=5000 | rel_L2={rel_l2:.6e}\n")

assert rel_l2 < 1e-1, f"relative L2 error too large: {rel_l2:.3e}"