jNO for LLMs

jNO ships a single self-contained briefing file written for language models: the house conventions, the full API vocabulary, the measured gotchas, and a set of verified end-to-end examples — everything a coding agent needs to write correct, idiomatic jNO without crawling the whole site.

Two raw files, stable URLs, plain markdown:

File	Contents
`llms.txt`	Short index — what jNO is, where the key pages are
`llms-full.txt`	The full guide (~1300 lines): conventions, API, traps, verified examples

Use it by pasting llms-full.txt into a system prompt, pointing a coding agent at the URL, or dropping it into an agent-skill directory (it began life as one). Every example in its §12 was executed against a released jNO and exits 0, and the file is maintained alongside the code — when the API and this file disagree, one of them has a bug, and we treat it that way.

The same content, rendered for humans:

jNO — jax Numerical Operators (guide for LLMs)

jNO is a JAX-native library for differentiable numerical methods. Classical solvers (finite elements, finite differences, spectral/RCWA) and scientific machine learning (PINNs, neural operators, Bayesian inference) sit on one substrate: you write the math — a weak form, a strong-form stencil, a PDE residual, a data loss — and it lowers to a single GPU-ready, end-to-end reverse-mode-differentiable, jit-compiled graph.

Because every solve is differentiable, things that are normally separate frameworks are the same tool here: an inverse problem, a PDE-constrained optimisation, and a neural-network coefficient are one composition away from a forward solve.

Every example in §12 was executed against jNO 0.3.0 and exited 0. This file records what actually runs, not what reads plausibly — the measured traps live in §11.

1. House convention — write jNO this way

Four rules. They are not stylistic; they keep PINN, FEM and FDM code reading identically, and three of them route around API spellings that no longer exist.

1.1 Derivatives: `.bind` first, then attributes

u  = (x * (1 - x) * net(x)).scalar.bind(x=x)     # PINN: .scalar.bind, then u.x / u.xx / u.t / u.tt
ui = u.bind(x=xi, y=yi, t=ti)                     # FEM & FDM: bind the symbol, then ui.x / ui.xx / ui.t
pde = ui.xx + ui.yy + f                           # ✓

Not u.d2(x) + u.d2(y), not u.laplacian(x, y). Those work, but the bound-attribute form is the one convention across all three solvers, and .bind is what makes attribute derivatives available even when the coordinate is not the network's own input. Costs nothing: u.xx + u.yy folds to a single Laplacian node in the trace.

Available on a bound view: .x .y .z .t, .xx .yy .zz .tt, .xy and other mixed partials.

1.2 Domains: `jno.Shape` / `jno.Path` only

d = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain()                    # 2-D
d = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain(time=(0, 0.5, 4))    # 2-D + time
d = jno.Path(0, 0).line_to(1, 0).curve(size=0.02).domain()            # 1-D, ends named left/right
d = jno.Shape.box(0, 0, 0, 1, 1, 1, size=0.1).domain()                # 3-D

Never shapely.geometry.box, never jno.domain.rect(...) / .line(...) / .cube(...) / .l_shape(...). The jno.domain.* constructors have been removed — they are the single largest source of stale example code (see §11).

1.3 `jno.setup` always, first line after the imports

run = jno.setup(__file__)     # -> "./runs/<script-stem>"; also initialises logging

Everything written during the run goes under run: stats.plot(f"{run}/training.png"), jno.save(crux, f"{run}/model.pkl").

1.4 Inline and short where it stays readable

One line per idea. Chain the controls that return self. Build the constraint list inside jno.core([...]) when the expressions are short enough to read there.

run = jno.setup(__file__)
d = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain()
x, y, _ = d.variable("interior")

net = jno.nn(foundax.mlp(in_features=2, hidden_dims=64, num_layers=4, key=jax.random.PRNGKey(0)))
net.optimizer(optax.adam(1e-3))

u = (net(x, y) * x * (1 - x) * y * (1 - y)).scalar.bind(x=x, y=y)      # hard BC
pde = u.xx + u.yy + 2 * pi**2 * sin(pi * x) * sin(pi * y)

jno.core([pde.mse]).solve(5000).plot(f"{run}/training.png")

Also settled: jno.nn(...) over jno.nn.wrap(...) (identical, shorter), jno.np over import jno.numpy as jnn (one namespace, no alias), and splat a split coordinate tuple into a field call — u(*a) - u(*b) to glue two faces, u(*bb) - 0.0 for a Dirichlet edge, u(*ci) - g for an initial condition — never u(a[0], a[1], a[2]). The call accepts the whole tuple including the trailing time coordinate (measured on Dirichlet, ties and ICs). Indexing stays only where you genuinely want one coordinate (ci[0] inside a formula) or a keyword bind (.bind(x=sv[0], ...)).

2. The one idea: everything is a trace

You describe the problem as a symbolic expression, not as a training loop. Domain points, network calls, derivatives, PDE residuals, weak forms, integrals, noise, and trainable parameters are all nodes in one graph. jno.core(...) JIT-compiles that graph once into a JAX function reused by both crux.solve() (training) and crux.eval() (evaluation) — which is why the same expression can be a residual loss during training and a quantity of interest afterwards.

Four normally-separate workflows are just different nodes in that graph, so they compose and differentiate uniformly:

Workflow	Trial	Loss / operator
PINN	a network	strong-form residual `u.xx + f`, derivatives by autodiff
Plain NN / operator learning	a network	supervised `(pred - data).mse`
FEM	FE basis	weak form as a term list → `jno.fem([...])`, differentiable `fem.solve()`
FDM	nodal unknown	same term list, no test function → `jno.fdm([...])`

Mix them in one jno.core(...) and inverse problems fall out for free: put a trainable jno.np.parameter (or a whole network) anywhere in the graph, compare to data, and the gradient flows back through the derivatives, the solve, and the network in one pass.

Term	What it is
Placeholder	base symbolic node — a coordinate, a network call, an operation, a residual
Constraint	any expression reduced to a scalar (`expr.mse`) and handed to `jno.core`
Crux	what `jno.core(...)` returns — compiled step, optimiser state, history
Model	what `jno.nn(module)` returns — a wrapped Equinox module with per-model controls
Tag	a string label on a domain mapping to a point set (`"interior"`, `"boundary"`, `"left"`)

3. Install and environments

pip install jax-numerical-operators             # FEM, FDM, solvers, PINNs, SciML — CPU out of the box
pip install "jax-numerical-operators[cuda]"     # NVIDIA GPU (CUDA-capable JAX build, same pin)

Optional extras: [cuda] (GPU JAX), [fem] (meta-extra = [mesh] remeshing + [pardiso] + [cudss] sparse-direct backends), [rcwa] (Fourier solver, pulls fmmax), [amg] (GPU algebraic multigrid), [iree]; combinable as [cuda,fem]. There is deliberately no [fdm] extra — FDM is core. Models come from foundax — jNO ships no model zoo of its own.

In the jNO repo itself everything runs through pixi:

pixi run -e dev  python script.py     # dev adds matplotlib
pixi run -e rcwa python script.py
pixi run test                         # pytest -x --tb=short
pixi run ci-test                      # skips @pytest.mark.slow

4. Domains and geometry

variable always returns a trailing time coordinate (constant on a steady domain): a 1-D domain unpacks as x, t, 2-D as x, y, t, 3-D as x, y, z, t. Forgetting the trailing _ is the most common beginner error.

x, y, _ = d.variable("interior")                          # PINN collocation coords
xi, yi, _ = d.variable("interior", split=True)            # FEM/FDM quadrature coords
xb, yb, _, nx, ny = d.variable("top", normals=True, split=True)
x, y, _ = d.variable("interior", sample=(500, None))      # 500 sampled points
x0, y0, t0 = d.variable("initial")                        # the t = t0 slice

The `Shape` DSL

A Shape is an immutable build-plan (gmsh-OpenCASCADE); every call returns a new shape.

Primitive	Auto-named boundaries
`Shape.rect(x0, y0, x1, y1)`	`left` `right` `top` `bottom`
`Shape.disk(cx, cy, r)`	`arc`
`Shape.polygon(points)`	`e0 e1 … eN`
`Shape.box(x0, y0, z0, x1, y1, z1)`	`left right top bottom front back`
`Shape.cylinder(x, y, z, dx, dy, dz, r)`	`side` `top` `bottom`
`Shape.sphere(cx, cy, cz, r)`	`surface`
`Path(…).face()`	per-segment, `name=` as you draw

interior and boundary always exist. Combine with a - b (cut), a | b (fuse), a & b (intersect). Transforms: .extrude(h), .revolve(pt, dir, angle), .sweep(path), .array(n, …), .translate(v), .rotate(...), .fillet(r, where=…), .sized(size).

Path chains line_to / arc_to; .face() closes it into a 2-D shape, .curve(size=) makes it a 1-D domain, and an open Path is a sweep trajectory. A sharp line_to→line_to corner is rejected — round it with arc_to.

Mesh density lives on the shape, not the domain — a small size= on the shape covering the region you want refined, or .sized(lambda x, y, z: 0.03 + 0.10*y) for grading.

Naming, selection, materials

d.tag("inlet",  lambda x, y, z: x < 1e-6)                                # coordinate predicate
d.tag("outlet", lambda x, n, name: (n[:, 0] > 0.9) & (name != "left"))   # coords + normal + name
xl, yl, _ = d.variable("left", where=lambda x, y, z: x < 1e-6)           # tag AND bind in one call

d = (core.name("inclusion") + plate.name("matrix")).sized(0.05).domain()  # conforming multi-material
d.variable("inclusion|matrix")     # the interface facets, auto-named by the sorted region pair

+ keeps distinct materials with a conforming interface; | fuses them into one. Interfaces are listed by d.interface_tags() and kept out of d.boundary_tags().

Other domain sources

d = jno.domain("part.msh")                                        # physical groups become tags
d = jno.domain.from_array({"interior": pts, "boundary": bpts})    # point cloud (no mesh ⇒ no integrate/FD)
d = jno.Shape.rect(0, 0, 1, 1, size=0.02).domain(structured=True) # uniform grid → fast FD stencils
dom = 500 * jno.Shape.rect(0, 0, 2, 1, size=0.05).domain()        # replicate for operator learning

d.plot("domain.png") renders mesh, regions and normals. Node coordinates: d.built_mesh.points, or d.mesh_connectivity["points"] on an FDM domain.

5. Expressions

Integration

.integrate() collapses a field to a scalar; the region (volume vs boundary) is auto-detected from the Variable tags inside the expression, so you pass no region argument.

vol  = u.integrate()               # ∫_Ω u dV
bnd  = u_b.integrate()             # ∮_∂Ω u ds
heat = (u * v).integrate(ti)       # time-window integral (trapezoid)
Ku   = (kernel(x, y) * u_of(y)).integrate(var=x)   # non-local / Fredholm — one value per point

xb, yb, _, nx, ny = d.variable("boundary", normals=True, split=True)
flux = (u_b.x * nx + u_b.y * ny).integrate()       # ∮ ∂u/∂n ds — flux is written explicitly

Needs compute_mesh_connectivity=True. A temporal .integrate(t) needs min_consecutive=None (or >= 2) in solve().

Reductions, math, views

u.mse  u.mae  u.mean  u.sum  u.min  u.max  u.std       # → scalar nodes
u.shape  u.T  u.real  u.imag       a.equal(b)   a > b
u.scalar  u.vector  u.complex  u.matrix  u.voigt  u.field    # typed views, each with .bind

jno.np carries the elementwise library: sin cos tan exp log sqrt cbrt abs sign square power floor ceil round, where maximum minimum, concat stack reshape squeeze expand_dims transpose, dot matmul cross norm inner, inv det eigvalsh logm expm sqrtm, symgrad trace vector, constants pi e inf nan. Vector calculus: jno.np.jacobian/divergence/curl_2d/curl_3d, and on a VectorView directly — jno.np.vector(a, b).div(x, y).

Differentiation schemes

The scheme rides on the call: u.d(x, scheme="finite_difference"). (Attribute derivatives use the default; drop to the method form only when you need a non-default scheme.)

Scheme string	Grad	Lap/Hess	Notes
`"automatic_differentiation"` (default)	✅	✅	exact, any domain
`":forward"` / `":reverse"`	✅	—	`jacfwd` / `jacrev`
`":fwd-over-rev"` (default Hessian), `":fwd-over-fwd"`, `":rev-over-rev"`, `":rev-over-fwd"`	—	✅	2nd-order AD variants
`"finite_difference"`	✅	✅	area-weighted, unstructured meshes
`":lsq"` / `":uniform"` / `":inverse_distance"`	✅	✅
`":cotangent"`	—	✅	cotangent Laplacian, 2-D only

Forward mode is cheaper when input dim ≤ output dim (typical for PINNs); reverse for scalar losses with many inputs. Project-wide default via .jno.toml or jno.setup(__file__, diff_type="forward", hessian_type="fwd-over-rev").

Trap. ":cotangent" and ":lsq" return the whole Laplacian in one shot, so u.d2(x, scheme="finite_difference:cotangent") + u.d2(y, …) doubles it.

Units, escape hatch, trackers

x = x.unit("m").scale(L);  u = net(x, t).unit("K").scale(U)
jno.units.check(res);  jno.units.nondimensionalize(res);  jno.units.rescale(res)

result = jno.fn(lambda x, y: jnp.exp(-x**2) * jnp.sin(y), [x, y])   # any JAX fn, differentiable
val = jno.np.mean(jno.np.abs(u - u_exact)).tracker(100)             # logged, NOT part of the loss
J = u.grad(u_net)                                                    # ∂u/∂θ, (B, N, P) — the NTK's J
J_sg = J.stop_gradient                                               # identity fwd, zero bwd

.name("label") tags an expression for logs/W&B; .print(what="shape") emits a runtime shape and passes the value through. Both chain.

6. Training — `jno.core`

crux = jno.core(constraints=[pde.mse, bc.mse], domain=d, rng_seed=42,
                weights=[1.0, 10.0], mesh=(1, 1))   # mesh = (batch_devices, model_devices)
stats = crux.solve(epochs=10_000, batchsize=32, callbacks=[...], min_consecutive=None)
stats.plot(f"{run}/training.png")

Every non-frozen model must have an optimiser before solve().

from jno import LearningRateSchedule as lrs

net.optimizer(optax.adam(1e-3))                                              # bake the rate in
net.optimizer(optax.adam(1)).scale(lrs.exponential(1e-3, 0.9, 2000, 1e-5))   # or attach a schedule
net.optimizer(optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-3)))

Construct with a placeholder rate of 1 when using .scale(...) — the schedule sets the effective rate and is re-evaluated every step. Factories: lrs.constant, lrs.exponential, lrs.cosine, lrs.warmup_cosine, lrs.piecewise_constant, all taking min_lr/max_lr; any (epoch, losses) -> scalar callable is a schedule. jno.optimizers adds second-order optimisers absent from optax: engd, ssbroyden, ssbfgs, soap, md.

What `solve()` actually returns

stats.total_loss           # FINAL SCALAR total loss (a float, NOT an array)
stats.total_loss_history   # 1-D array, concatenated across every solve() call
stats.training_logs        # list of per-solve()-call dicts
stats.plot(path)           # returns self, so it chains off solve()
stats.summary()

Evaluation and debugging

pred = crux.eval(u)                                # on the training domain
pred = crux.eval(u, domain=fine_domain)            # on any other domain
a, b = crux.eval([a, b])                           # several nodes at once
crux.print_tree("tree.txt");  crux.print_shapes()
crux.solve(5000, profile=True)                     # Perfetto trace in <run>/traces/
crux.checkpoints                                    # every solve() call checkpoints

Adaptive weights, resampling, callbacks

w_pde, w_bc = jno.fn.adaptive.relobralo([pde, bc])          # also softadapt, dwa
crux = jno.core([w_pde * pde, w_bc * bc, w_pde.tracker(), w_bc.tracker()])

x, y, _ = d.variable("interior", sample=(None, None),
                     resampling_strategy=jno.sampler.rad(resample_every=100,
                                                         resample_fraction=0.1, start_epoch=1000))

cb = jno.callbacks.early_stopping(patience=1000, min_delta=1e-6, mode="min")
crux.solve(100_000, callbacks=[cb]);  print(cb.stopped_epoch, cb.best_metric)

Samplers: random, rad, rard, ha, cr3 (time-causal), pinnfluence. Always delay with start_epoch > 0; fractions of 0.1–0.3 are typical. Custom callbacks subclass jno.utils.adaptive.callbacks.Callback and override on_solve_begin, on_before_update, on_epoch_end (return True to stop) or on_training_end.

Model controls

net.freeze()  net.unfreeze()
net.mask(param_mask)                    # one-shot boolean-pytree scope for the NEXT control
net.constrain(jax.nn.softplus)          # reparameterise WEIGHTS before every forward pass
net.lora(rank=4, alpha=1.0)
net.dtype(jnp.float64);  net.initialize("weights.eqx")
net.summary();  net.reset();  net.to_iree(sample_inputs)

.scale(...) is overloaded by receiver: on an expression it declares a characteristic magnitude for non-dimensionalisation; on a model it sets the learning-rate schedule.

constrain() shapes the weights, not the field values — softplus on every weight does not make the output positive. For a positive field use an output transform: k = jno.fn.exp(k_raw(x)).

7. FEM — `jno.fem`

Write the weak form as a plain list of residual terms: volume physics, natural boundary terms, and essential boundary conditions all in the same list. There is no dirichlet(...)/neumann(...) call — jno.fem classifies each term by the region its symbols are bound to.

Symbols

d.fem_symbols(value_shape=(), order=1, space=None, names=("u", "phi"))

value_shape=(2,) → vector unknown (elasticity, velocity)
order=k → degree-k Lagrange (any k ≥ 1)
space="RT" (H(div)), "N1curl" (H(curl) — Maxwell, eddy currents), "Argyris"/"Morley" (C¹, plates/biharmonic), "Hermite" (1-D beam), "P0"
call fem_symbols once per field for coupled systems

Boundary conditions are terms

Condition	Term
Dirichlet `u = g`	`u(xb, yb) - g`
Per-component (roller)	`u(xb, yb)[i] - g`
Neumann flux `∂u/∂n = g`	`-g * phi.bind(x=xb, y=yb)`
Robin `∂u/∂n + a u = g`	`(a * u.bind(x=xb, y=yb) - g) * phi.bind(x=xb, y=yb)`
Vector traction `t`	`-jno.np.inner(t, phi.bind(x=xb, y=yb), n_contract=1)`
Periodic / tie	`u("left") - u("right")` (no test function ⇒ a tie)
Bloch	`u(A) - c * u(B)`

g may be a constant or a coordinate expression. Zero Neumann is the natural default and needs no term. Tag periodic faces with a predicate so both sides include their corner nodes.

Vocabulary inside a term

Closed-form nonlinear physics is already symbolic — + - * / ** and jno.np compose straight in:

D      = k * (1e-3 + ui.x**2 + ui.y**2) ** 0.3     # nonlinear diffusivity D(|∇u|)
react  = ui * ui * vi                               # u² reaction
energy = (ui.x**2 + ui.y**2).integrate()            # ∫|∇u|²
tau    = d.cell_size / (2 * beta.norm())            # SUPG stabilisation from the element size

Geometry symbols: d.variable(tag, normals=True, split=True) → nx, ny; d.cell_size; d.enclosure(tags) (view factors for grey-body radiation). Second derivatives for 4th-order forms: jno.np.laplacian(ui, [xi, yi]), jno.np.hessian(ui, [xi, yi]) — needs order >= 2.

Conformity caveat. Lagrange is C⁰, so ∫Δu·Δv over P2 is non-conforming and does not give a convergent biharmonic discretisation. Use space="Argyris" (C¹), space="Morley", or the mixed Ciarlet–Raviart pair.

Escape hatch: jno.fn(lambda a, b: ..., [ui, vi]) turns any differentiable JAX function of traced arguments into a term.

What `jno.fem` returns

Form	flags	Use
steady, linear	`is_linear=True`, `is_transient=False`	`fem.A`, `fem.b`
steady, nonlinear	`is_linear=False`	`fem.residual(u)`, `fem.jacobian(u)`, `fem.dofs`
has a `ui.t` term	`is_transient=True`	`fem.M`, `fem.operator.A`, `fem.state0`, `fem.dt`, `fem.t0`, `fem.t1`

Always available: fem.dofs, fem.points (the coordinates the DOFs live on — use these, not mesh vertices, for P2+), fem.operator, fem.classification, fem.offsets, fem.blocks. fem.operator.evaluate({"k": values}) materialises a parametric operator at given coefficients.

A second time derivative (ui.tt) is auto-reduced to the first-order system y = [u, v] and integrated by the energy-conserving trapezoidal rule. It needs two initial conditions (displacement and velocity). Do not hand-roll backward Euler on a second-order block — it spuriously damps the wave.

Coefficients — known vs trainable

k = jno.np.parameter(phi).initialize(lambda x, y: 1.0 + 4.0*x).freeze()   # KNOWN → assembles normally
k = jno.np.parameter(phi, name="k")                                        # TRAINABLE → parametric, via crux
k = d.variable("kappa", sample=k0).trainable()                             # promote anything to a parameter

A jno.np.parameter is trainable by default, which makes the system runtime-parametric — so fem.A / fem.b are unavailable until you .freeze() it. A frozen value must be a scalar or a coordinate function; a raw per-node array fails loud. .trainable() on a coordinate makes the mesh vertices a design variable (shape optimisation, r-adaptivity).

Solving — the slot API

Every slot is a configured callable (never a string); every None keeps the default.

u = fem.solve(x0        = u_guess,                 # warm start
              nonlinear = jno.solve.newton(),      # or picard(damping=0.7)
              linear    = jno.solve.gmres(),
              precond   = jno.precond.jacobi())

Structure	Solver
SPD (Poisson, elasticity, mass)	`cg`
non-symmetric (advection, SUPG)	`bicgstab` / `gmres` — `bicgstab + jacobi` is the default
iterative preconditioner inside	`fgmres`
symmetric indefinite (Stokes/Biot saddle, biharmonic)	`minres`
SPD, batched/GPU-heavy	`chebyshev` (inner-product free)
indefinite, single solve	`lu` — no vmap rule, use a Krylov inside batched solves
cuSolver refuses it or is slow	`lu(backend="host")` (SuperLU)
shift-invert eigs, constant-operator transient	`lu(backend="cudss")`
a Newton loop, or no GPU	`lu(backend="pardiso")`
small systems / coarse blocks	`dense`

jno.solve.newton(damping=, rtol=1e-8, atol=1e-8, max_steps=100, line_search=, direct=) — note rtol/atol, there is no tol=. direct=True factorises the assembled tangent each step; it is required, not preferred, when the matrix-free Newton-Krylov has no preconditioner to lean on (stiff tangents, saddle points) — otherwise it diverges to NaN.

Preconditioners: jacobi, chebyshev(degree=), nystrom(rank=) (SPD only), amg(cycles=), gmg(), ams, form([terms], inner=), inner(solver), block_diag/triangular (per-field over fem.blocks), cached.

AMG is for repeated solves against the same operator — build the hierarchy once, or an unbuilt spec re-runs pyamg's host-side setup every solve:

M = jno.precond.amg().build(fem.operator[0])
u = fem.solve(linear=jno.solve.cg(tol=1e-10), precond=M)

Break-even is ~7 solves at 95k DOFs, ~66 at 3k. Below ~100k DOFs for a single solve, Jacobi wins.

Picard for lagged coefficients — when a solution-dependent coefficient's Newton tangent destroys the linearised structure (shear-thinning viscosity in non-Newtonian Stokes), freeze it with jno.lag(...) and drive with jno.solve.picard(). lag is stop_gradient on the traced expression, so the linearisation is the Picard iteration.

Eigenproblems and SVD

K = jno.fem([ui.x*vi.x + ui.y*vi.y, u(xb, yb) - 0.0])   # stiffness, no source term
lam, X = K.eigs(mass=[ui*vi], k=6)                       # dense reduction — exact, O(N²) memory
lam, X = K.eigs(mass=[ui*vi], k=6, precond=jno.precond.amg())   # LOBPCG, never densified
lam, X = K.eigs(mass=[ui*vi], k=4, sigma=60.0)           # the 4 modes nearest λ = 60
lam2, X2 = K2.eigs(mass=mass, k=6, precond=…, X0=X)      # warm-start a sweep
U, s, Vt = jno.solve.svd(snapshots, k=6)                 # POD basis; depth must exceed k

tol/maxiter are rejected without precond= or sigma=, so a tolerance can never be silently ignored by the dense path. Do not set tol near machine precision — an exhausted budget NaN-poisons the result by design.

Build time and memory

A 13.8k-DOF 2-D Poisson issues ~209 XLA compilations (~3.7 s), paid once per distinct mesh shape and cached in-process (~320 ms to rebuild). A remeshing loop pays it per iteration. Across separate processes nothing is reused unless you opt in:

jax.config.update("jax_compilation_cache_dir", "~/.cache/jax")
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0.0)   # REQUIRED — default is 1.0s

The second line is not optional here: every one of these compilations is far below the 1 s default, so the cache directory alone writes nothing. jno.fem(chunk=) bounds the element loop's peak memory (default ~0.15 % of device memory per chunk, floor 8192 cells); chunk=False restores a single vmap.

8. FDM — `jno.fdm`

The strong-form sibling of jno.fem: the same kind of constraint list, but the strong residual is collocated at the mesh nodes with finite-difference stencils. No test function, no mass matrix, no quadrature — and the solve is differentiable through custom_root, so it composes into inverse problems exactly like fem.solve().

d.unknown() is the strong-form counterpart of fem_symbols(): a valued P1 nodal field. Autodiff against a coordinate is meaningless on a discrete field, so the bound view's derivatives are finite differences by default — no scheme= needed.

Flux BCs are written with the edge's own normal; jno.fdm handles any condition affine in ∂u/∂n, and any mix of Dirichlet/Neumann/Robin composes. A non-affine condition raises rather than returning a wrong answer. In 2-D a corner node shared by two flux edges has no single normal and falls back to the interior residual — give it an explicit Dirichlet value if it needs anchoring.

Structured grids. structured=True on an axis-aligned Shape.rect/Shape.box builds a uniform grid and switches the interior operators to direct 5-point (2-D) / 7-point (3-D) stencils by array reshaping. Same answer as the unstructured cotangent operator, cheaper. The inner Krylov defaults to GMRES + geometric-multigrid V-cycle because row-replaced Dirichlet makes the reduced operator nonsymmetric. Pick a size giving an even (ideally power-of-two) cell count per axis or multigrid falls back to plain GMRES. Periodic boundaries and complex fields are not supported in jno.fdm.

9. Inverse problems

Always the same shape: put an unknown in the graph, compare to data, crux.solve.

# scalar
a = jno.np.parameter((1,), key=k1, name="a").optimizer(optax.adam(1e-2))

# through a FEM forward solve — the gradient flows through the assembled solve
k = jno.np.parameter((1,), name="k")
k.dtype(jnp.float64); k.initialize(jax.nn.initializers.constant(2.0)); k.optimizer(optax.adam(5e-2))
fem  = jno.fem([k * (ui.x*vi.x + ui.y*vi.y) - f*vi, u(xb, yb) - 0.0])
crux = jno.core([(fem.solve() - u_obs).mse], domain=obs_domain)

For a transient weak form fem.solve() returns the trajectory, so a rate constant is recovered from a time series.

Regularisation. .regularize(kind, *vars) returns an unreduced pointwise placeholder — apply .mean or .mse yourself. Kinds: "smooth" / "h1seminorm" (H1 seminorm, good default), "tv" (sharp interfaces), "l2", "nonneg", "bounded"(lo=, hi=). For a FEM nodal-parameter field the penalty is assembled FE-exact; for a coordinate field it uses autodiff, so pass the spatial variables.

.regularize lives on the network call or the parameter, not on an arbitrary expression. jno.fn.exp(k_raw(x)).regularize(...) raises AttributeError — regularise k_raw(x) and apply the transform separately (see example 08).

Enforce positivity of the field with an output transform (k = jno.fn.exp(k_raw(x))), not with constrain() on the weights.

10. Models, RCWA, Bayesian, save/load

Models

Every model comes from foundax, wrapped with jno.nn(...). Any Equinox module works.

Family	foundax constructors
Linear / MLP	`foundax.linear`, `foundax.mlp`
DeepONet	`foundax.deeponet`
FNO / CNO / U-Net / MgNO	`foundax.fno1d/2d/3d`, `cno2d`, `unet1d/2d/3d`, `mgno1d/2d`
Geometry-aware	`foundax.geofno`, `pcno`, `pit`, `pointnet`
GNOT family	`foundax.cgptno`, `gnot`, `moegptno`
Transformer	`foundax.transformer`
Foundation models	`foundax.poseidon`, `morph`, `mpp`, `walrus`, `dpot`, `prose`, …

Bayesian

model.bayesian(kernel_factory, **kw) replaces the per-step gradient update with one blackjax MCMC transition. It mirrors .optimizer(...) — each parameter is independently optimised or sampled, mixed freely, and crux.solve() dispatches per model.

Model.bayesian(kernel_factory, *, prior=None, warmup=500, keep=1000, thin=1, **kernel_kwargs). The factory is duck-typed by its first parameter name: logdensity_fn → full-data MCMC (blackjax.nuts, hmc, mala); grad_estimator → stochastic-gradient (sgld, sghmc). step_size= is required except for HMC-family kernels with adapt=True and warmup > 0. Afterwards: p.posterior_samples, p.posterior_diagnostics, crux.eval([u], samples="chain"). .vi(blackjax.meanfield_vi, optimizer=…) is the variational alternative, mutually exclusive.

RCWA

Optional ([rcwa], built on fmmax, imported lazily). Hand it the same constraint list you would give jno.fem — it infers period (from the Floquet ties — absent ⇒ raise), ambients, permittivity, wavelength, and the incident wave. Only orders is genuinely yours.

sol = jno.rcwa(constraints, orders=300).solve()
sol.efficiency("T");  sol.order(+1, 0);  sol.power("up");  sol.extraction("up")

A finite aperture with no ties is rejected rather than silently periodicised; model an isolated scatterer by keeping the ties and adding an in-plane PML frame.

Save, load, configuration

jno.save(crux, f"{run}/model.pkl");   crux = jno.load(f"{run}/model.pkl")

cloudpickle serialises weights, optimiser state, logs, the domain and mesh, the expression tree, checkpoints and the RNG state. After loading, Python variable references no longer point at the models inside the solver — reassign in one call: crux.set_optimizer(optax.adam, scale=lrs(1e-4)).

.jno.toml in the cwd (or ~/.jno/config.toml) sets [jno] seed, [runs] base_dir, [rsa] keys for signed save/load, and the AD defaults diff_type / hessian_type.

11. Gotchas — measured traps

These were found by running code, not by reading. (Earlier revisions of this file listed docs pages that contradicted the measurements; those pages have since been fixed — the API facts below remain the part worth keeping loaded.)

Removed API — these raise AttributeError and appear all over older examples:

Stale	Current
`jno.domain.rect(mesh_size=…, time=…)`	`jno.Shape.rect(…, size=…).domain(time=…)`
`jno.domain.cube(...)`	`jno.Shape.box(...).domain()`
`jno.domain.l_shape(...)`	build it from `Shape` booleans
`jno.AdaptSpec`	gone

(jno.domain.line(mesh_size=…) still exists as a 1-D shorthand, but the house form — and what the docs use everywhere — is jno.Path(0,0).line_to(1,0).curve(size=…).domain().)

API facts that trip people up:

crux.eval on a jno.np.parameter of shape (1,) returns the array. float(_a) raises TypeError: Only scalar arrays can be converted to Python scalars. Use _a[0].
.regularize(...) is not available on the result of jno.fn.exp(...) — AttributeError: 'FunctionCall' object has no attribute 'regularize'. Regularise the network call and apply the transform separately.
solve() returns a statistics whose total_loss is the final scalar, not a history array; the array is total_loss_history. There is no stats.epoch, stats.losses, stats.weights, stats.training_time, stats.trainable_params or stats.total_params — that data lives in the stats.training_logs dicts.
jno.solve.newton takes rtol= / atol=, not tol=. Its default 1e-8 is tighter than some well-posed problems reach — a stationary Allen–Cahn floors at ~3.5e-8 and raises RuntimeError: newton_direct did not converge. Relax to atol=1e-7, rtol=1e-7.

Ordinary traps:

variable returns a trailing time coordinate — 1-D unpacks as x, t, 2-D as x, y, t. Use split=True for FEM/FDM quadrature coordinates.
Every non-frozen model needs an optimiser before solve().
.integrate() and finite-difference schemes need compute_mesh_connectivity=True.
fem.points, not mesh vertices, for P2 and above.
Hard BC enforcement beats a soft BC loss whenever the geometry admits a vanishing factor.
A temporal .integrate(t) needs min_consecutive=None or >= 2 in solve().
jno.solve.lu has no vmap rule — use a Krylov solver inside batched/parametric solves.
A trainable jno.np.parameter makes the FEM system runtime-parametric, so fem.A/fem.b are unavailable until you .freeze() it.
jax.config.update("jax_enable_x64", True) must precede the first JAX array, and jno.fdm wants it.
Optional backends fail at import, not at call: mmgpy (adaptive remeshing), diffrax, fmmax (RCWA), pyamg (AMG). Check before recommending them.

12. Verified examples

Every script below ran to completion, exit code 0, against jNO 0.3.0 with a CUDA JAX build. Each is standalone and self-checking: run it and the final assert tells you whether it still works.

12.1 PINN — 1-D Poisson, hard BC

"""1-D Poisson  -u'' = sin(pi x),  u(0) = u(1) = 0,  hard BC ansatz."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi = jno.np.pi

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.02).domain()
x, _ = d.variable("interior")

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

u = (x * (1 - x) * net(x)).scalar.bind(x=x)          # hard u(0)=u(1)=0
pde = u.xx + jno.np.sin(pi * x)                       # -u'' = sin(pi x)

crux = jno.core([pde.mse])
crux.solve(5000).plot(f"{run}/training.png")

_u, _exact = crux.eval([u, jno.np.sin(pi * x) / pi**2])
rel = float(jax.numpy.linalg.norm(_u - _exact) / jax.numpy.linalg.norm(_exact))
print(f"rel_L2 = {rel:.3e}")
assert rel < 1e-1

12.2 PINN — 2-D Poisson

"""2-D Poisson  -Delta u = 2 pi^2 sin(pi x) sin(pi y)  on the unit square, u = 0 on the boundary."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi, sin = jno.np.pi, jno.np.sin

d = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain()
x, y, _ = d.variable("interior")

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

u = (net(x, y) * x * (1 - x) * y * (1 - y)).scalar.bind(x=x, y=y)   # hard u = 0 on the boundary
pde = u.xx + u.yy + 2 * pi**2 * sin(pi * x) * sin(pi * y)

crux = jno.core([pde.mse])
crux.solve(5000).plot(f"{run}/training.png")

_u, _exact = crux.eval([u, sin(pi * x) * sin(pi * y)])
rel = float(jax.numpy.linalg.norm(_u - _exact) / jax.numpy.linalg.norm(_exact))
print(f"rel_L2 = {rel:.3e}")
assert rel < 1e-1

12.3 PINN — transient heat (`u.t`, IC + BC)

"""1-D heat  u_t = alpha u_xx  on a time-dependent 1-D domain (soft IC + BC)."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi, alpha, T = jno.np.pi, 0.1, 0.5

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.05).domain(time=(0, T, 4))
x, t = d.variable("interior")
x0, t0 = d.variable("initial")          # the t = t0 slice
xb, tb = d.variable("boundary")

net = jno.nn(foundax.deeponet(n_sensors=1, coord_dim=1, n_outputs=1, n_layers=3,
                              basis_functions=48, hidden_dim=32, key=jax.random.PRNGKey(0)))
net.optimizer(optax.adam(optax.exponential_decay(1e-3, 2000, 0.5, end_value=1e-5)))

u = net(t, x).scalar.bind(x=x, t=t)
crux = jno.core([(u.t - alpha * u.xx).mse,                     # PDE
                 (net(t0, x0) - jno.np.sin(pi * x0)).mse,      # initial condition
                 net(tb, xb).mse])                             # u = 0 on the spatial boundary
crux.solve(5000).plot(f"{run}/training.png")

_u, _exact = crux.eval([u, jno.np.exp(-alpha * pi**2 * t) * jno.np.sin(pi * x)])
rel = float(jax.numpy.linalg.norm(_u - _exact) / jax.numpy.linalg.norm(_exact))
print(f"rel_L2 = {rel:.3e}")
assert rel < 2e-1

12.4 PINN — wave equation (`u.tt`, two initial conditions)

"""1-D wave  u_tt = c^2 u_xx  — second order in time needs TWO initial conditions."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi, c, T = jno.np.pi, 1.0, 1.0

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.05).domain(time=(0, T, 8))
x, t = d.variable("interior")
x0, t0 = d.variable("initial")
xb, tb = d.variable("boundary")

net = jno.nn(foundax.deeponet(n_sensors=1, coord_dim=1, n_outputs=1, n_layers=4, basis_functions=64,
                              hidden_dim=48, activation=jax.nn.tanh, key=jax.random.PRNGKey(7)))
net.optimizer(optax.adam(optax.warmup_cosine_decay_schedule(1e-6, 1e-3, 100, 10_000, 1e-6)))

u = net(t, x).scalar.bind(x=x, t=t)
u0 = net(t0, x0).scalar.bind(x=x0, t=t0)

crux = jno.core([(u.tt - c**2 * u.xx).mse,                 # PDE
                 (u0 - jno.np.sin(pi * x0)).mse,           # displacement IC
                 u0.t.mse,                                  # velocity IC (at rest)
                 net(tb, xb).mse])                          # u = 0 on the spatial boundary
crux.solve(10_000).plot(f"{run}/training.png")

_u, _exact = crux.eval([u, jno.np.cos(c * pi * t) * jno.np.sin(pi * x)])
rel = float(jax.numpy.linalg.norm(_u - _exact) / jax.numpy.linalg.norm(_exact))
print(f"rel_L2 = {rel:.3e}")
assert rel < 3e-1

12.5 PINN — coupled system (two networks, one core)

"""Coupled elliptic system: two networks, two residuals, one core."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi, sin = jno.np.pi, jno.np.sin

d = jno.Shape.rect(0, 0, 1, 1, size=0.1).domain()
x, y, _ = d.variable("interior")

f = 2 * pi**2 * sin(pi * x) * sin(pi * y) + sin(2 * pi * x) * sin(pi * y)
g = 5 * pi**2 * sin(2 * pi * x) * sin(pi * y) + sin(pi * x) * sin(pi * y)

k1, k2 = jax.random.split(jax.random.PRNGKey(0))
u_net = jno.nn(foundax.mlp(in_features=2, hidden_dims=48, num_layers=4, key=k1))
v_net = jno.nn(foundax.mlp(in_features=2, hidden_dims=48, num_layers=4, key=k2))
for net in (u_net, v_net):
    net.optimizer(optax.adam(optax.warmup_cosine_decay_schedule(0.0, 1e-3, 50, 5000, 1e-5)))

ansatz = 16.0 * x * (1 - x) * y * (1 - y)          # peaks at 1 — keeps the field O(1)
u = (u_net(x, y) * ansatz).scalar.bind(x=x, y=y)
v = (v_net(x, y) * ansatz).scalar.bind(x=x, y=y)

crux = jno.core([(-(u.xx + u.yy) + v - f).mse,
                 (-(v.xx + v.yy) + u - g).mse])
crux.solve(5000).plot(f"{run}/training.png")

_u, _ue, _v, _ve = crux.eval([u, sin(pi * x) * sin(pi * y), v, sin(2 * pi * x) * sin(pi * y)])
nrm = jax.numpy.linalg.norm
print(f"u rel_L2 = {float(nrm(_u - _ue) / nrm(_ue)):.3e}   v rel_L2 = {float(nrm(_v - _ve) / nrm(_ve)):.3e}")
assert float(nrm(_u - _ue) / nrm(_ue)) < 1.5e-1 and float(nrm(_v - _ve) / nrm(_ve)) < 1.5e-1

12.6 PINN — Fokker–Planck (Path geometry, `.div`, an integral constraint, a noise node)

"""2-D Fokker-Planck on a disc: a VectorView divergence, a normalisation INTEGRAL, and a noise node."""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi = jno.np.pi

d = (jno.Path(3, 0).arc_to(-3, 0, through=(0, 3)).arc_to(3, 0, through=(0, -3))
     .face().sized(0.25).domain())                       # a disc of radius 3, Shape/Path only
x, y, _ = d.variable("interior")
xb, yb, _ = d.variable("boundary")

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

p = net(x, y).scalar.bind(x=x, y=y)
drift = jno.np.vector(x * p, y * p).div(x, y)             # the OU probability flux, div of a VectorView
crux = jno.core([(drift + 0.5 * (p.xx + p.yy)).mse,       # Fokker-Planck residual
                 (p.integrate() - 1.0).mse,                # normalisation over the disc
                 (net(xb, yb) - (jno.np.exp(-(xb**2 + yb**2)) / pi
                                 + jno.noise.gaussian(std=1e-4))).mse])
crux.solve(15_000).plot(f"{run}/training.png")

_p, _exact = crux.eval([p, jno.np.exp(-(x**2 + y**2)) / pi])
rel = float(jax.numpy.linalg.norm(_p - _exact) / jax.numpy.linalg.norm(_exact))
print(f"rel_L2 = {rel:.3e}")
assert rel < 3e-1

12.7 Inverse — scalar parameters

"""Recover three unknown scalars from data: jno.np.parameter optimises like a network."""
import jax, optax, jno

run = jno.setup(__file__)
pi = jno.np.pi
A, B, C = 3.14, -2.71, 42.0

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.01).domain()
x, _ = d.variable("interior")
target = A * jno.np.sin(pi * x) + B * jno.np.cos(pi * x) + C * x * (1 - x)

k1, k2, k3 = jax.random.split(jax.random.PRNGKey(0), 3)
a = jno.np.parameter((1,), key=k1, name="a").optimizer(optax.adam(1e-2))
b = jno.np.parameter((1,), key=k2, name="b").optimizer(optax.adam(1e-2))
c = jno.np.parameter((1,), key=k3, name="c").optimizer(optax.adam(1e-2))

crux = jno.core([((a * jno.np.sin(pi * x) + b * jno.np.cos(pi * x) + c * x * (1 - x)) - target).mse])
crux.solve(30_000).plot(f"{run}/training.png")

_a, _b, _c = crux.eval([a, b, c])
# crux.eval returns the parameter ARRAY, shape (1,) — index it, float() on it raises.
print(f"recovered a={_a[0]:.3f} b={_b[0]:.3f} c={_c[0]:.3f}   truth {A} {B} {C}")
assert abs(_a[0] - A) / abs(A) < 0.1 and abs(_b[0] - B) / abs(B) < 0.1 and abs(_c[0] - C) / abs(C) < 0.1

12.8 Inverse — a coefficient field with a smoothness prior

"""Recover an unknown coefficient FIELD k(x) with a network + a smoothness prior.

k > 0 is enforced by an output transform (jno.fn.exp), not by constraining weights.
"""
import foundax, jax, optax, jno

run = jno.setup(__file__)
pi = jno.np.pi

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.01).domain()
x, _ = d.variable("interior")

f_pde = -(pi**2) * jno.np.sin(pi * x)      # k_true = 1, u_true = sin(pi x)  =>  k u'' = f
u_obs = jno.np.sin(pi * x)

k1, k2 = jax.random.split(jax.random.PRNGKey(0))
k_raw = jno.nn(foundax.mlp(in_features=1, output_dim=1, hidden_dims=16, num_layers=2, key=k1))
u_net = jno.nn(foundax.mlp(in_features=1, output_dim=1, hidden_dims=32, num_layers=3, key=k2))
k_raw.optimizer(optax.adam(1e-3))
u_net.optimizer(optax.adam(1e-3))

k_field = k_raw(x)                                          # regularize lives on the network CALL
k = jno.fn.exp(k_field)                                     # > 0 by construction
u = (u_net(x) * x * (1 - x)).scalar.bind(x=x)               # hard zero Dirichlet

crux = jno.core([(k * u.xx - f_pde).mse, (u - u_obs).mse, k_field.regularize("smooth", x).mean])
crux.solve(5_000).plot(f"{run}/training.png")

_u, _k, _obs = crux.eval([u, k, u_obs])
rel = float(jax.numpy.linalg.norm(_u - _obs) / jax.numpy.linalg.norm(_obs))
print(f"u rel_L2 = {rel:.3e}   k min/mean = {_k.min():.3f}/{_k.mean():.3f}")
assert rel < 1e-1 and _k.min() > 0

12.9 FEM — Poisson

"""FEM Poisson: the weak form IS the term list.  -Delta u = f,  u = 0 on the boundary."""
import jax.numpy as jnp, numpy as np, jno

run = jno.setup(__file__)

d = jno.Shape.rect(0, 0, 1, 1, size=0.18).domain()
u, phi = d.fem_symbols()                                  # trial / test
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))                 # -Delta[x(1-x)y(1-y)]
fem = jno.fem([ui.x * vi.x + ui.y * vi.y - f * vi,        # volume: grad u . grad v - f v
               u(xb, yb) - 0.0],                          # essential BC is just another term
              quad_degree=3)
u_h = jnp.asarray(fem.solve(linear=jno.solve.cg(), precond=jno.precond.jacobi()))   # SPD -> CG

pts = np.asarray(fem.points)                              # the coords the DOFs live on (NOT mesh vertices)
exact = pts[:, 0] * (1 - pts[:, 0]) * pts[:, 1] * (1 - pts[:, 1])
rel = float(jnp.linalg.norm(exact - u_h) / jnp.linalg.norm(exact))
print(f"dofs={fem.dofs} linear={fem.is_linear} rel_L2={rel:.3e}")
assert fem.is_linear and rel < 5e-2

12.10 FEM — mixed Dirichlet + Robin

"""Mixed Dirichlet + Robin.  -Delta u + sigma u = f,  exact u = x sin(pi y) + y.

A Robin condition du/dn + a u = g enters the weak form as the natural surface term
(a u - g) * phi bound to THAT edge — no BC objects, the binding carries the region.
"""
import jax.numpy as jnp, numpy as np, jno

run = jno.setup(__file__)
pi, sigma, a_r, a_t = np.pi, 4.0, 2.0, 3.0
sin = jno.np.sin

d = jno.Shape.rect(0, 0, 1, 1, size=0.22).domain()
u, phi = d.fem_symbols()
xi, yi, _ = d.variable("interior", split=True)
xl, yl, _ = d.variable("left", split=True)
xbo, ybo, _ = d.variable("bottom", split=True)
xr, yr, _ = d.variable("right", split=True)
xt, yt, _ = d.variable("top", split=True)
ui, vi = u.bind(x=xi, y=yi), phi.bind(x=xi, y=yi)

f = xi * pi**2 * sin(pi * yi) + sigma * (xi * sin(pi * yi) + yi)
fem = jno.fem([
    ui.x * vi.x + ui.y * vi.y + sigma * u * vi - f * vi,                                   # volume
    (a_r * u.bind(x=xr, y=yr) - (sin(pi * yr) + a_r * (sin(pi * yr) + yr))) * phi.bind(x=xr, y=yr),
    (a_t * u.bind(x=xt, y=yt) - (1.0 - pi * xt + a_t)) * phi.bind(x=xt, y=yt),              # Robin
    u(xl, yl) - yl,                                                                         # Dirichlet
    u(xbo, ybo) - 0.0,
], quad_degree=3)
u_h = jnp.asarray(fem.solve())                                # default: Jacobi-BiCGStab, matrix-free

pts = np.asarray(fem.points)
exact = pts[:, 0] * np.sin(np.pi * pts[:, 1]) + pts[:, 1]
rel = float(jnp.linalg.norm(exact - u_h) / jnp.linalg.norm(exact))
print(f"dofs={fem.dofs} rel_L2={rel:.3e}")
assert rel < 5e-2

12.11 FEM — vector P2 elasticity

"""Vector P2 elasticity: a cantilever under an end shear, checked against Euler-Bernoulli."""
import jax.numpy as jnp, numpy as np, jno

run = jno.setup(__file__)
E, nu, L, H, q = 1000.0, 0.3, 10.0, 1.0, 0.1
lam, mu = E * nu / (1.0 - nu**2), E / (2.0 * (1.0 + nu))       # plane stress
inner, symgrad, trace = jno.np.inner, jno.np.symgrad, jno.np.trace

d = jno.Shape.rect(0, 0, L, H, size=0.5).domain()
u, phi = d.fem_symbols(value_shape=(2,), order=2)              # P2 vector displacement
xi, yi, _ = d.variable("interior", split=True)
xl, yl, _ = d.variable("left", split=True)                     # clamped root
xr, yr, _ = d.variable("right", split=True)                    # loaded tip
eu, ep = symgrad(u, [xi, yi]), symgrad(phi, [xi, yi])

fem = jno.fem([lam * trace(eu) * trace(ep) + 2.0 * mu * inner(eu, ep, n_contract=2),
               u(xl, yl) - (0.0, 0.0),
               -1.0 * inner(jnp.array([0.0, -q]), phi.bind(x=xr, y=yr), n_contract=1)])
# A slender P2 stiffness is too ill-conditioned for Jacobi-CG in float32 -> sparse direct.
sol = jnp.asarray(fem.solve(linear=jno.solve.lu())).reshape(-1, 2)

tip = np.asarray(fem.points)[:, 0] > L - 1e-6
delta, eb = float(-jnp.mean(sol[tip, 1])), (q * H) * L**3 / (3.0 * E * (H**3 / 12.0))
print(f"dofs={fem.dofs}  FEM tip={delta:.4f}  Euler-Bernoulli={eb:.4f}  ratio={delta / eb:.3f}")
assert fem.is_linear and abs(delta - eb) / eb < 0.05

12.12 FEM — transient (`M`, `A`, `state0`)

"""Transient FEM: time=(t0,t1,n) on the domain + a ui.t term gives the semidiscrete M u' + A u = 0."""
import jax.numpy as jnp, numpy as np, jno

run = jno.setup(__file__)
nu = 1.0
dense = lambda A: jnp.asarray(A.todense()) if hasattr(A, "todense") else jnp.asarray(A)  # noqa: E731

d = jno.Shape.rect(0, 0, 1, 1, size=0.08).domain(time=(0.0, 0.05, 26))
u, phi = d.fem_symbols()
xi, yi, ti = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
xi0, yi0, _ = d.variable("initial", split=True)             # the t = t0 slice
ui, vi = u.bind(x=xi, y=yi, t=ti), phi.bind(x=xi, y=yi, t=ti)

fem = jno.fem([ui.t * vi + nu * (ui.x * vi.x + ui.y * vi.y),
               u(xb, yb) - 0.0,
               u(xi0, yi0) - jno.fn(lambda x, y: jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y), [xi0, yi0])])
assert fem.is_transient

M, A, dt = dense(fem.M), dense(fem.operator.A), float(fem.dt)
w = jnp.asarray(fem.state0)
for _ in range(round((fem.t1 - fem.t0) / dt)):              # backward Euler: (M + dt A) w' = M w
    w = jnp.linalg.solve(M + dt * A, M @ w)

pts = np.asarray(fem.points)
exact = np.exp(-2.0 * nu * np.pi**2 * fem.t1) * np.sin(np.pi * pts[:, 0]) * np.sin(np.pi * pts[:, 1])
rel = float(jnp.linalg.norm(exact - w) / jnp.linalg.norm(exact))
print(f"dofs={fem.dofs} rel_L2 at t={fem.t1}: {rel:.3e}")
assert rel < 5e-2

12.13 FEM — nonlinear (Newton)

"""Nonlinear FEM: a cubic reaction makes the form nonlinear, so jno.fem returns a RESIDUAL operator.

Newton is driven through fem.solve() so the implicit-function gradient survives.
direct=True is required here: the matrix-free Newton-Krylov diverges to NaN from this
deliberately over-wide interface, because the cubic makes the tangent stiff.
"""
import jax.numpy as jnp, numpy as np, jno

run = jno.setup(__file__)
eps = 0.15
exact = lambda x: np.tanh((x - 0.5) / (np.sqrt(2.0) * eps))  # noqa: E731

d = jno.Shape.rect(0, 0, 1, 1, size=0.06).domain()
u, phi = d.fem_symbols()
xi, yi, _ = d.variable("interior", split=True)
xl, yl, _ = d.variable("left", split=True)
xr, yr, _ = d.variable("right", split=True)
ui, vi = u.bind(x=xi, y=yi), phi.bind(x=xi, y=yi)

fem = jno.fem([eps**2 * (ui.x * vi.x + ui.y * vi.y) + (u**3 - u) * vi,
               u(xl, yl) - exact(0.0), u(xr, yr) - exact(1.0)], quad_degree=3)
assert not fem.is_linear

pts = np.asarray(fem.points)
u0 = np.tanh((pts[:, 0] - 0.5) / (np.sqrt(2.0) * 0.30))     # over-wide interface = the Newton start
uh = np.asarray(fem.solve(nonlinear=jno.solve.newton(direct=True, atol=1e-7, rtol=1e-7),
                          x0=jnp.asarray(u0))).reshape(-1)

rel = float(jnp.linalg.norm(exact(pts[:, 0]) - uh) / jnp.linalg.norm(exact(pts[:, 0])))
res = float(jnp.linalg.norm(jnp.asarray(fem.residual(jnp.asarray(uh)))))
print(f"dofs={fem.dofs} |residual|={res:.2e} rel_L2={rel:.3e}")
assert res < 1e-6 and rel < 1e-2

12.14 FEM — generalized eigenproblem

"""Generalized eigenproblem K x = lambda M x through fem.eigs — the Dirichlet-Laplacian spectrum."""
import numpy as np, jno

run = jno.setup(__file__)

d = jno.Shape.rect(0, 0, 1, 1, size=0.06).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)

K = jno.fem([ui.x * vi.x + ui.y * vi.y, u(xb, yb) - 0.0])    # stiffness, no source term
lam, X = K.eigs(mass=[ui * vi], k=5)                          # dense reduction: exact, O(N^2) memory

got = np.asarray(lam)[:5] / np.pi**2
exact = np.array([2, 5, 5, 8, 10])                            # pi^2 (m^2 + n^2)
print(f"lambda/pi^2 computed = {np.round(got, 3)}")
print(f"lambda/pi^2 analytic = {exact}")
assert np.max(np.abs(got - exact) / exact) < 5e-2

12.15 FEM — inverse through a differentiable solve

"""Inverse problem THROUGH a differentiable FEM solve: recover a whole nodal field k(x).

fem.solve() is a trace node, so the gradient of the data misfit flows back through the
assembled solve to the parameter. Field inversion is ill-posed, hence the H1 prior.
"""
import jax, jax.numpy as jnp, numpy as np, optax, jno

run = jno.setup(__file__)

d = jno.Shape.rect(0, 0, 1, 1, size=0.1).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 = 30.0 * (xi * (1 - xi) + yi * (1 - yi))                    # strong source: u is sensitive to k

nodes = np.asarray(d.built_mesh.points)[:, :2]
k_true = 1.0 + 0.8 * np.exp(-((nodes[:, 0] - 0.5) ** 2 + (nodes[:, 1] - 0.5) ** 2) / (2 * 0.12**2))

k = jno.np.parameter(phi, name="k")                           # trainable P1 field on the trial space
fem = jno.fem([k * (ui.x * vi.x + ui.y * vi.y) - f * vi, u(xb, yb) - 0.0], quad_degree=3)

A, b = fem.operator.evaluate({"k": jnp.asarray(k_true)})       # synthesise data at the true k
A = A.todense() if hasattr(A, "todense") else jnp.asarray(A)
u_obs = jnp.linalg.solve(A, jnp.asarray(b).reshape(-1))

k.dtype(jnp.float64)
k.initialize(jax.nn.initializers.constant(1.0))                # start from a uniform field
k.optimizer(optax.adam(2e-2))
crux = jno.core([(fem.solve() - u_obs).mse, 1e-3 * k.regularize("h1seminorm").mean],
                domain=jno.domain.from_array({"_": np.zeros((1, 1))}))   # no spatial domain needed
crux.solve(500).plot(f"{run}/training.png")

rec = np.asarray(crux.eval([k])).reshape(-1)                   # the recovered field (do NOT index [0])
rel = float(np.linalg.norm(rec - k_true) / np.linalg.norm(k_true))
print(f"nodes={k_true.shape[0]}  k(x) rel_L2={rel:.3e}  peak rec/true={rec.max():.3f}/{k_true.max():.3f}")
assert rel < 0.1

12.16 FDM — Poisson

"""FDM Poisson: the strong form, collocated at the mesh nodes.

domain.unknown() is the strong-form counterpart of fem_symbols(). Autodiff is meaningless on
a discrete nodal field, so the bound view's derivatives are finite differences by default —
no scheme= needed. The Dirichlet condition is the same u(region) - g term jno.fem uses.
"""
import jax
jax.config.update("jax_enable_x64", True)      # the strong-form solve accumulates in float64

import numpy as np  # noqa: E402
import jno  # noqa: E402

run = jno.setup(__file__)

d = jno.Shape.rect(0, 0, 1, 1, size=0.06).domain()
x, y, _ = d.variable("interior", split=True)
xb, yb, _ = d.variable("boundary", split=True)
u = d.unknown()                                 # valued P1 nodal field
ui = u.bind(x=x, y=y)                           # bound view: .x/.xx are FD stencils

f = 2.0 * np.pi**2 * jno.np.sin(np.pi * x) * jno.np.sin(np.pi * y)
sol = jno.fdm([-ui.xx - ui.yy - f,              # -Delta u = f
               u(xb, yb) - 0.0]).solve()

p = np.asarray(d.mesh_connectivity["points"])[:, :2]
exact = np.sin(np.pi * p[:, 0]) * np.sin(np.pi * p[:, 1])
rel = float(np.linalg.norm(np.asarray(sol).reshape(-1) - exact) / np.linalg.norm(exact))
print(f"nodes={p.shape[0]} rel_L2={rel:.3e}")
assert rel < 3e-2

12.17 FDM — Dirichlet + Neumann + Robin at once

"""FDM with Dirichlet + Neumann + Robin at once. Exact u = y^2.

A flux BC is written with THAT edge's own tags: bind the field to the edge and take its
normal derivative against the edge's outward normal. Any condition affine in du/dn works.
"""
import jax
jax.config.update("jax_enable_x64", True)

import numpy as np  # noqa: E402
import jno  # noqa: E402

run = jno.setup(__file__)

d = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain()
x, y, _ = d.variable("interior", split=True)
xbo, ybo, _ = d.variable("bottom", split=True)
xl, yl, _ = d.variable("left", split=True)
xr, yr, _ = d.variable("right", split=True)
xt, yt, _ = d.variable("top", split=True)
nl, nr, nt = (d.variable(t, normals=True) for t in ("left", "right", "top"))

u = d.unknown()
ui = u.bind(x=x, y=y)                                    # interior view for the PDE
ul, ur, ut = u.bind(x=xl, y=yl), u.bind(x=xr, y=yr), u.bind(x=xt, y=yt)

sol = jno.fdm([-ui.xx - ui.yy + 2.0,                     # -Delta u = -2
               u(xbo, ybo) - 0.0,                        # Dirichlet: bottom held at 0
               ul.d(nl) - 0.0,                           # Neumann: left insulated
               ur.d(nr) - 0.0,                           # Neumann: right insulated
               ut.d(nt) + 1.0 * (ut - 3.0)]).solve()     # Robin: top convects to u_inf = 3

p = np.asarray(d.mesh_connectivity["points"])[:, :2]
rel = float(np.linalg.norm(np.asarray(sol).reshape(-1) - p[:, 1] ** 2) / np.linalg.norm(p[:, 1] ** 2))
print(f"nodes={p.shape[0]} rel_L2={rel:.3e}")
assert rel < 5e-2

12.18 Operator learning — DeepONet on a PDE residual

"""Operator learning: one DeepONet over 50 random coefficients, trained on the PDE residual.

N * domain replicates one mesh across N operator-learning samples; domain.variable(name, array)
attaches the per-sample coefficient as a tensor tag.
"""
import foundax, jax, optax, jno

run = jno.setup(__file__)
KEY, N, EPOCHS = jax.random.PRNGKey(0), 50, 2_000

d = N * jno.Shape.rect(0, 0, 2, 1, size=0.05).domain()
x, y, _ = d.variable("interior")
k = d.variable("k", jax.random.uniform(KEY, shape=(N, 1, 1), minval=0.5, maxval=1.5))

net = jno.nn(foundax.deeponet(n_sensors=1, coord_dim=2, basis_functions=32, hidden_dim=128,
                              activation=jax.numpy.tanh, key=KEY))
net.optimizer(optax.adam(optax.cosine_decay_schedule(1e-3, EPOCHS, alpha=1e-2)))

u = (net(k, jno.np.concat([x, y], axis=-1)) * x * (2 - x) * y * (1 - y)).scalar.bind(x=x, y=y)
crux = jno.core(constraints=[(k * (u.xx + u.yy) + 1.0).mse])
stats = crux.solve(epochs=EPOCHS, batchsize=32)
stats.plot(f"{run}/training.png")

# stats.total_loss is the FINAL SCALAR; the array is stats.total_loss_history.
hist = stats.total_loss_history
print(f"loss {hist[0]:.3e} -> {stats.total_loss:.3e} over {hist.size} logged epochs")
assert stats.total_loss < hist[0]

12.19 Bayesian — NUTS posterior over unknown coefficients

"""Bayesian inverse problem: .bayesian replaces the gradient update with a blackjax NUTS transition.

Each parameter is independently optimised (optax) or sampled (blackjax) — mix freely.
"""
import blackjax, jax, jax.numpy as jnp, jno

run = jno.setup(__file__)
pi = jno.np.pi
A, B = 3.14, -2.71

d = jno.Path(0, 0).line_to(1, 0).curve(size=0.02).domain()
x, _ = d.variable("interior")
target = A * jno.np.sin(pi * x) + B * jno.np.cos(pi * x)

k1, k2 = jax.random.split(jax.random.PRNGKey(0), 2)
a = jno.np.parameter((1,), key=k1, name="a")
b = jno.np.parameter((1,), key=k2, name="b")
for p in (a, b):
    # adapt=True (default) runs window adaptation for `warmup` steps and tunes step_size
    # + inverse_mass_matrix; the step_size given here is only the adapter's initial guess.
    p.bayesian(blackjax.nuts, step_size=1e-2, warmup=300, keep=500)

crux = jno.core([(a * jno.np.sin(pi * x) + b * jno.np.cos(pi * x) - target).mse])
crux.solve(800)

for name, chain, truth in (("A", a.posterior_samples, A), ("B", b.posterior_samples, B)):
    mean = float(jnp.mean(chain))
    lo, hi = (float(v) for v in jnp.quantile(chain, jnp.array([0.05, 0.95])))
    print(f"{name} = {mean:.3f}  90% CI [{lo:.3f}, {hi:.3f}]   truth {truth}")
    assert abs(mean - truth) / abs(truth) < 0.3